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",
        "军报文章": "military_report",
        "节目文稿": "manuscript",
        "报题单": "baoti",
        "装备": "manual",
        "技术": "manual",
        "动态": "manual",
        "术语": "manual",
        "厂商": "manual",
        "索引": "manual",
    }

    # 来源大类固定显示顺序（制片人 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

        # fallback：如果 source_detail 仍为空，取"来源"字段
        if not source_detail:
            raw_source = str(fm.get("来源", "") or "").strip()
            if raw_source:
                source_detail = raw_source

        # —— 播出日期：容错 "待补充" 等非日期文本——
        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

        # —— 相关实体（涉及装备/涉及技术/涉及厂商/主题/相关装备/应用领域）——
        import re as _re
        related_entities = []
        for key in ("涉及装备", "涉及技术", "涉及厂商", "主题", "相关装备", "应用领域"):
            val = fm.get(key)
            if val:
                if isinstance(val, list):
                    for item in val:
                        # 去掉 Obsidian 双链格式 [[xxx]]
                        cleaned = _re.sub(r"\[\[|\]\]", "", str(item)).strip()
                        if cleaned:
                            related_entities.append(cleaned)
                elif isinstance(val, str):
                    # 可能是 "山东舰, 福建舰" 这样的逗号分隔字符串
                    for item in val.replace("，", ",").split(","):
                        item = _re.sub(r"\[\[|\]\]", "", 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)

        # 保留原始 Obsidian 类别（映射到 manual 的类型，如装备/技术/动态/术语/厂商等）
        if raw_type and source_type == "manual":
            metadata["obsidian_category"] = raw_type

        # 源文件路径保留（预埋 PDF 原文链接）
        raw_source_files = fm.get("源文件")
        if raw_source_files:
            cleaned_files = []
            if isinstance(raw_source_files, list):
                for sf in raw_source_files:
                    cleaned = _re.sub(r'[\[\]"]', "", str(sf)).strip()
                    if cleaned:
                        cleaned_files.append(cleaned)
            elif isinstance(raw_source_files, str):
                cleaned = _re.sub(r'[\[\]"]', "", raw_source_files).strip()
                if cleaned:
                    cleaned_files.append(cleaned)
            if cleaned_files:
                metadata["source_files"] = cleaned_files

        # —— 正文（去掉 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 = 10, min_similarity: float = 0.3) -> list[dict]:
        """
        语义检索：查询句转为向量，用 SQL 余弦距离（<=>）在数据库层检索
        返回 top_k 条相似笔记，含相似度分数 + 智能摘要片段。
        过滤掉 similarity < min_similarity 的条目。
        """
        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,
                    ki.author,
                    ki.tags,
                    ki.content_md,
                    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()
            results = []
            for r in rows:
                similarity = round(r.similarity, 4)
                # 过滤低于阈值的条目
                if similarity < min_similarity:
                    continue
                tags = r.tags or {}
                source_detail = tags.get("source_detail") if isinstance(tags, dict) else None
                snippet = self._extract_smart_snippet(r.content_md, query_text)
                results.append({
                    "id": r.id,
                    "title": r.title,
                    "source_type": r.source_type,
                    "author": r.author,
                    "source_detail": source_detail,
                    "snippet": snippet,
                    "content_md": r.content_md,
                    "similarity": similarity,
                })
            return results

    def _extract_smart_snippet(self, content_md: str, query_text: str, max_len: int = 300) -> str:
        """
        智能摘要提取：根据搜索词定位最相关段落，截取摘要并加粗关键词。
        中文连续字符串会被拆成 2-gram 子串以提高段落匹配命中率。
        """
        import re as _re

        if not content_md:
            return ""

        # 1. 切分关键词：先按空格/标点拆，再把长中文词拆成 2-gram
        raw_parts = _re.split(r'[\s，。！？、；：""''（）\(\)\[\]\-]+', query_text)
        raw_parts = [p.strip() for p in raw_parts if len(p.strip()) > 1]

        keywords = []
        for part in raw_parts:
            keywords.append(part)
            if len(part) > 2 and _re.fullmatch(r'[一-鿿]+', part):
                for i in range(len(part) - 1):
                    bigram = part[i:i+2]
                    if bigram not in keywords:
                        keywords.append(bigram)

        if not keywords:
            return content_md[:max_len]

        # 2. 按段落分割（跳过纯 markdown 标记行如 # ## --- 等）
        paragraphs = content_md.split("\n\n")
        if len(paragraphs) < 3:
            paragraphs = content_md.split("\n")
        paragraphs = [p.strip() for p in paragraphs
                       if p.strip() and not _re.fullmatch(r'[#\-=\s>]+', p.strip())]

        # 3. 计算每段关键词命中数（完整词权重 3，bigram 权重 1）
        best_para = ""
        best_score = 0
        for para in paragraphs:
            score = 0
            for kw in keywords:
                if kw in para:
                    score += 3 if kw in raw_parts else 1
            if score > best_score:
                best_score = score
                best_para = para

        # 4. fallback：无关键词命中时用正文前 max_len 字
        if best_score == 0:
            snippet = content_md[:max_len]
        else:
            snippet = best_para[:max_len]

        # 5. 加粗命中关键词（优先标记完整词，避免 bigram 重复标记）
        marked = set()
        for kw in sorted(keywords, key=len, reverse=True):
            if kw in snippet and kw not in marked:
                snippet = snippet.replace(kw, f"**{kw}**", 1)
                marked.add(kw)
                for sub_kw in keywords:
                    if sub_kw != kw and sub_kw in kw:
                        marked.add(sub_kw)

        return snippet

    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 固定顺序排列，仅显示有数据的大类（count > 0）。

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

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

        total_count = len(items)

        # 按 source_type 分组：仅收集有数据的类别，再按固定顺序排列
        type_groups: dict = {}
        for item in items:
            st = item.source_type or "manual"
            if st not in type_groups:
                type_groups[st] = []
            type_groups[st].append(item)

        # 只遍历有数据的类别，按 SOURCE_TYPE_ORDER 顺序
        children = []
        for st in self.SOURCE_TYPE_ORDER:
            if st not in type_groups:
                continue
            st_items = type_groups[st]

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

            if secondary_field is not None:
                # 按二级字段分组
                detail_groups: dict = {}
                for item in st_items:
                    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,
        }]