tps-dashboard/backend/app/services/knowledge_service.py

"""
知识库服务 — 写入向量 + 语义检索 + md 文件解析
使用 pgvector 原生 SQL 向量检索（<=> 余弦距离算子），不在 Python 侧计算
"""

from typing import Optional
from datetime import date

import frontmatter
from sqlalchemy import text
from sqlmodel import Session, select
from pgvector.sqlalchemy import Vector

from app.models.knowledge import KnowledgeItem, KnowledgeEmbedding
from app.services.embedding_service import EmbeddingService
from app.db.session import engine


class KnowledgeService:
    """知识库 CRUD + 语义检索 + md 解析"""

    # yaml 类型字段 → source_type 枚举映射
    SOURCE_TYPE_MAP = {
        "杂志文章": "military_report",
        "军报": "military_report",
        "节目文稿": "manuscript",
        "报题单": "baoti",
    }

    def __init__(self):
        self.embedder = EmbeddingService()

    def parse_md_file(self, file_content: bytes, file_name: str) -> dict:
        """
        解析一个 .md 文件的 yaml frontmatter + 正文，返回入库用的字典。
        严格按真实样本的字段名映射，不猜测。

        Returns:
            dict 含 keys: title, content_md, source_type, author, publish_date,
            source_detail, metadata(JSONB), related_entities(JSONB)
        """
        content = file_content.decode("utf-8", errors="replace")
        parsed = frontmatter.loads(content)
        fm = parsed.metadata or {}

        # —— 类型 → source_type（硬映射，不猜测）——
        raw_type = str(fm.get("类型", "")).strip()
        source_type = self.SOURCE_TYPE_MAP.get(raw_type, "manual")

        # —— 标题：名称 或 标题——
        title = str(fm.get("名称", "") or fm.get("标题", "")).strip()
        if not title:
            # fallback: 用正文第一行或文件名
            lines = [l.strip() for l in content.split("\n") if l.strip() and not l.strip().startswith("---")]
            title = lines[0] if lines else file_name

        # —— 作者：作者 或 编导——
        author = str(fm.get("作者", "") or fm.get("编导", "") or "").strip() or None

        # —— 出处详情：期刊 + 期号（拼在一起存进 JSONB 的 source_detail）——
        journal = str(fm.get("期刊", "") or "").strip()
        issue = str(fm.get("期号", "") or "").strip()
        if journal or issue:
            source_detail = f"{journal} {issue}".strip()
        else:
            source_detail = None

        # —— 播出日期：容错 "待补充" 等非日期文本——
        raw_date = str(fm.get("播出日期", "") or "").strip()
        publish_date = None
        if raw_date and raw_date not in ("待补充", "待确认", ""):
            try:
                publish_date = date.fromisoformat(raw_date)
            except ValueError:
                # 非 ISO 格式，尝试 common 格式
                for fmt in ("%Y-%m-%d", "%Y年%m月%d日", "%Y/%m/%d"):
                    try:
                        publish_date = date.fromisoformat(raw_date.replace("年", "-").replace("月", "-").replace("日", ""))
                        break
                    except ValueError:
                        continue

        # —— 权重（不展示，存 JSONB 备 Phase 4）——
        weight = str(fm.get("权重", "") or "").strip() or None

        # —— 相关实体（涉及装备/涉及技术/涉及厂商/主题）——
        related_entities = []
        for key in ("涉及装备", "涉及技术", "涉及厂商", "主题"):
            val = fm.get(key)
            if val:
                if isinstance(val, list):
                    related_entities.extend(val)
                elif isinstance(val, str):
                    # 可能是 "山东舰, 福建舰" 这样的逗号分隔字符串
                    for item in val.replace("，", ",").split(","):
                        item = item.strip()
                        if item:
                            related_entities.append(item)

        # —— metadata JSONB：权重、出处详情、双链预留——
        metadata = {}
        if weight:
            metadata["weight"] = weight
        if source_detail:
            metadata["source_detail"] = source_detail
        # related_concepts 字段预留给双链解析（Phase 4），本 Task 原样存入
        metadata["double_bracket_links"] = self._extract_double_brackets(parsed.content)

        # —— 正文（去掉 frontmatter 的纯内容）——
        content_md = parsed.content

        return {
            "title": title,
            "content_md": content_md,
            "source_type": source_type,
            "author": author,
            "publish_date": publish_date,
            "metadata": metadata if metadata else None,
            "related_entities": related_entities if related_entities else None,
            "source_file_name": file_name,
        }

    def _extract_double_brackets(self, text: str) -> list[str]:
        """提取 [[...]] 双链标记，原样返回列表，不解析成图谱（本 Task 留门）。"""
        import re
        return re.findall(r"\[\[([^\]]+)\]\]", text)

    def store_md_file(self, file_content: bytes, file_name: str) -> KnowledgeItem:
        """
        读取一篇 md 内容，调用 embo-01 拿到向量，写入 knowledge_items + knowledge_embeddings
        """
        parsed = self.parse_md_file(file_content, file_name)

        # 调用 embedding（type="db" 表示存入知识库）
        embedding_list = self.embedder.embed_single(parsed["content_md"], embed_type="db")

        with Session(engine) as session:
            item = KnowledgeItem(
                title=parsed["title"],
                content_md=parsed["content_md"],
                source_type=parsed["source_type"],
                source_file_name=parsed["source_file_name"],
                author=parsed["author"],
                publish_date=parsed["publish_date"],
                tags=parsed["metadata"],
                related_entities=parsed["related_entities"],
            )
            session.add(item)
            session.flush()

            emb = KnowledgeEmbedding(
                knowledge_id=item.id,
                chunk_index=0,
                chunk_text=parsed["content_md"],
                embedding=embedding_list,
            )
            session.add(emb)
            session.commit()
            session.refresh(item)
            return item

    def delete_item(self, knowledge_id: int) -> bool:
        """删除知识库条目及其向量（CASCADE 已由 DB 层配置）。"""
        with Session(engine) as session:
            item = session.get(KnowledgeItem, knowledge_id)
            if item is None:
                return False
            session.delete(item)
            session.commit()
            return True

    def list_items(self, source_type: Optional[str] = None) -> list[dict]:
        """返回知识库条目列表（含 source_detail 从 metadata 解压）。"""
        with Session(engine) as session:
            statement = select(KnowledgeItem).order_by(KnowledgeItem.created_at.desc())
            if source_type:
                statement = statement.where(KnowledgeItem.source_type == source_type)
            items = session.exec(statement).all()

            results = []
            for item in items:
                # 从 tags(JSONB) 取 source_detail
                tags = item.tags or {}
                source_detail = tags.get("source_detail") if isinstance(tags, dict) else None
                results.append({
                    "id": item.id,
                    "title": item.title,
                    "author": item.author,
                    "publish_date": item.publish_date,
                    "source_type": item.source_type,
                    "source_file_name": item.source_file_name,
                    "source_detail": source_detail,
                    "created_at": item.created_at,
                })
            return results

    def get_distinct_sources(self) -> list[str]:
        """返回库里所有不重复的 source_detail（动态从 JSONB 提取），供筛选下拉用。"""
        with Session(engine) as session:
            items = session.exec(select(KnowledgeItem)).all()
            sources = set()
            for item in items:
                tags = item.tags or {}
                if isinstance(tags, dict) and tags.get("source_detail"):
                    sources.add(tags["source_detail"])
            return sorted(list(sources))

    def search_similar(self, query_text: str, top_k: int = 5) -> list[dict]:
        """
        语义检索：查询句转为向量，用 SQL 余弦距离（<=>）在数据库层检索
        返回 top_k 条相似笔记，含相似度分数
        """
        query_vector = self.embedder.embed_single(query_text, embed_type="query")
        vec_str = "[" + ",".join(str(v) for v in query_vector) + "]"

        with Session(engine) as session:
            sql = f"""
                SELECT
                    ki.id,
                    ki.title,
                    ki.source_type,
                    1 - (ke.embedding <=> '{vec_str}'::vector) AS similarity
                FROM knowledge_embeddings ke
                JOIN knowledge_items ki ON ke.knowledge_id = ki.id
                WHERE ke.chunk_index = 0
                ORDER BY ke.embedding <=> '{vec_str}'::vector
                LIMIT {top_k}
            """
            stmt = text(sql)
            rows = session.execute(stmt).all()
            return [
                {"id": r.id, "title": r.title, "source_type": r.source_type, "similarity": round(r.similarity, 4)}
                for r in rows
            ]

    def get_item_count(self) -> int:
        with Session(engine) as session:
            return len(session.exec(select(KnowledgeItem)).all())

    def get_embedding_count(self) -> int:
        with Session(engine) as session:
            return len(session.exec(select(KnowledgeEmbedding)).all())