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",
    }

    # 来源大类固定显示顺序（制片人 Obsidian 习惯）
    SOURCE_TYPE_ORDER = [
        "manuscript",      # 节目文稿
        "military_report", # 杂志文章
        "baoti",           # 报题单
        "manual",          # 其他
    ]

    # 二级分组维度映射（与前端 useKnowledgeGrouping.js 的 SECONDARY_GROUP_FIELD 保持一致）
    # key = source_type, value = 用来做二级分组的字段名，None = 不建二级节点
    SECONDARY_GROUP_FIELD = {
        "manuscript": "author",            # 节目文稿 → 按作者(编导)
        "military_report": "source_detail", # 杂志文章 → 按出处
        "baoti": None,                     # 报题单 → 不分组
        "manual": None,                    # 其他 → 不分组
    }

    SOURCE_TYPE_LABEL = {
        "military_report": "杂志文章",
        "manuscript": "节目文稿",
        "baoti": "报题单",
        "manual": "其他",
    }

    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())

    def get_grouped_items(self) -> list[dict]:
        """
        按 source_type → 二级字段（author / source_detail）两层聚合，返回树形结构数据。
        按 SOURCE_TYPE_ORDER 固定顺序排列。

        二级节点 key 格式：`{source_type}|{二级字段名}|{字段值}`
        例：manuscript|author|左鑫
            military_report|source_detail|航空知识 2026年第1期

        二级字段值为 null / 空字串 → 归入对应大类，不造空节点。
        """
        with Session(engine) as session:
            items = session.exec(select(KnowledgeItem)).all()

        total_count = len(items)

        # 按 source_type 分组，初始化所有已知类别（保证空类别也按顺序出现）
        type_groups: dict = {st: [] for st in self.SOURCE_TYPE_ORDER}
        for item in items:
            st = item.source_type or "manual"
            if st not in type_groups:
                type_groups[st] = []
            type_groups[st].append(item)

        children = []
        for st in self.SOURCE_TYPE_ORDER:
            st_items = type_groups.get(st, [])
            if not st_items:
                # 空类别也按顺序出现（0条）
                children.append({
                    "key": st,
                    "label": f"{self.SOURCE_TYPE_LABEL.get(st, st)}（0条）",
                    "count": 0,
                    "children": [],
                })
                continue

            secondary_field = self.SECONDARY_GROUP_FIELD.get(st)
            grandchildren = []

            if secondary_field is not None:
                # 按二级字段分组
                detail_groups: dict = {}
                for item in st_items:
                    # 取二级字段值：author 直接取，source_detail 从 tags(JSONB) 取
                    if secondary_field == "source_detail":
                        tags = item.tags or {}
                        sd = tags.get("source_detail") if isinstance(tags, dict) else None
                        field_val = sd
                    else:
                        field_val = (getattr(item, secondary_field, None) or "").strip() or None

                    if field_val not in detail_groups:
                        detail_groups[field_val] = []
                    detail_groups[field_val].append(item)

                for sd, sd_items in detail_groups.items():
                    if sd is not None:
                        grandchildren.append({
                            "key": f"{st}|{secondary_field}|{sd}",
                            "label": f"{sd}（{len(sd_items)}条）",
                            "count": len(sd_items),
                        })

            children.append({
                "key": st,
                "label": f"{self.SOURCE_TYPE_LABEL.get(st, st)}（{len(st_items)}条）",
                "count": len(st_items),
                "children": grandchildren,
            })

        return [{
            "key": "all",
            "label": f"全部（{total_count}条）",
            "count": total_count,
            "children": children,
        }]