PlantAI-rag-new/vstore/flow.py

"""flow.py —— 生信流程知识库（flow_docs 独立集合，只新增不动主库）。

与主库（公文/资料库）完全隔离：独立目录 + 独立 RetrievalIndex。
不加密：语料为团队内部公开分析文档，检索 meta 与正文均需在对话内引用原文。

- FlowStore(flow_dir)：add_chunk / search / delete_source / stats
- 幂等：以 source_id + chunk_key 为唯一键，重跑不产生重复块

数据目录布局：
  flow_dir/
  ├── index.json            # 检索索引（复用 RetrievalIndex，含向量）
  └── chunks/               # 块正文 json（chunk_key.json）

检索三档降级复用 RetrievalIndex 语义（向量 → meta 关键词 → 最新）。
"""
from __future__ import annotations

import hashlib
import json
import logging
import os
import re
from pathlib import Path
from typing import Any

from .embed import embed_batch
from .retrieval import RetrievalIndex, RetrievalError
from .settings import DEFAULT_TOP_K, MAX_TOP_K

logger = logging.getLogger(__name__)

INDEX_FILE = "index.json"
CHUNKS_DIR = "chunks"

# 各语料源稳定 ID（入库脚本与检索共用）
SOURCE_CORPUS = "corpus"
SOURCE_STAGE_CONTRACT = "stage_contract"
SOURCE_SKILL = "skill"

DOC_KIND_CORPUS = "corpus"
DOC_KIND_CONTRACT = "contract"
DOC_KIND_SKILL = "skill"

DOC_KIND_LABELS = {
    DOC_KIND_CORPUS: "全流程文档",
    DOC_KIND_CONTRACT: "阶段契约",
    DOC_KIND_SKILL: "技能说明",
}

# 引用块对 UI 的跳转约定（A3 对齐）：
#  - doc_kind=corpus     : path 为「仓库相对路径」，anchor 为「## 章标题」；
#                          UI 跳转 仓库根 /{path} 并定位 {anchor}（源文件与 A3 一并交付）
#  - doc_kind=contract   : path 为 dev 仓库内相对路径（gfa_core/schemas/stages/…）；
#                          UI 以文件卡片形式打开；行号锚点直接支持时用 line_start/line_end
#  - doc_kind=skill      : path 为 dev 仓库内相对路径（skills/…/SKILL.md），同 contract 渲染
CITE_ANCHOR = "anchor"
CITE_LINE = "line"

# 嵌入调用函数（测试可注入假 embed，镜像 vault.py 的约定）
_EMBED_FN = embed_batch


def set_embed_fn(fn) -> None:
    global _EMBED_FN
    _EMBED_FN = fn


def get_embed_fn():
    return _EMBED_FN

_META_KEYWORDS_FIELDS = ("name", "title", "title_path", "stage_name", "source", "doc_kind_label")


_CJK_SEG = re.compile(r"[一-鿿]{2,}")


def _cjk_2grams(text: str) -> list[str]:
    """中文连续段 2-gram 切词（检索词面通道用）。"""
    out: list[str] = []
    for m in _CJK_SEG.finditer(text):
        seg = m.group(0)
        out.append(seg[:2])
        if len(seg) > 2:
            for i in range(len(seg) - 1):
                out.append(seg[i:i + 2])
    return out


def _ascii_keywords(query: str) -> list[str]:
    """查询的 ASCII 词（词面通道专用；≥2 字符，允许 _ . - 数字）。"""
    ascii_part = _CJK_SEG.sub(" ", query)
    out: list[str] = []
    for w in re.split(r"[\s,，。；;:：/|?？!！()（）]+", ascii_part):
        if re.fullmatch(r"[A-Za-z0-9_.\-]+", w) and len(w) >= 3:
            out.append(w.lower())
    return out


def _meta_to_keywords(meta: dict[str, Any]) -> list[str]:
    """flow meta → 关键词列表（检索二档/词面通道）。

    规则：ASCII 按空白/标点切词（≥2 字符）；中文连续段取 2-gram（含段首双字）。
    """
    words: list[str] = []
    for key in _META_KEYWORDS_FIELDS:
        value = str(meta.get(key) or "")
        if not value:
            continue
        ascii_part = _CJK_SEG.sub(" ", value)
        words.extend([w for w in re.split(r"[\s,，。；;:：/|]+", ascii_part) if len(w) >= 2])
        words.extend(_cjk_2grams(value))
    return words


def _dump_json(path: Path, payload: dict[str, Any]) -> None:
    tmp = path.with_name(path.name + ".tmp")
    tmp.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
    os.replace(tmp, path)


class FlowStore:
    """flow_docs 集合：chunks/ 正文 + index.json 检索索引。"""

    def __init__(self, flow_dir: str | Path, dim: int | None = None) -> None:
        self.dir = Path(flow_dir)
        self.chunks_dir = self.dir / CHUNKS_DIR
        index_kwargs: dict[str, Any] = {"index_path": self.dir} if dim is None else {"index_path": self.dir, "dim": dim}
        self.index = RetrievalIndex(**index_kwargs)
        self.index.load()

    # ---- 写入 ----

    def add_chunk(self, chunk: dict[str, Any], text: str,
                  vector: list[float] | None = None) -> int:
        """追加/覆盖一个块（幂等键 source_id + chunk_key）。

        返回 row_id；键已存在时覆盖正文、更新向量并保持 row_id 稳定。
        vector 不传时内部嵌入「标题路径 + 标题 + 正文」拼接文本
        （标题路径并入嵌入文本，保证阶段名参与向量召回）。
        """
        if not (chunk.get("source_id") and chunk.get("chunk_key")):
            raise ValueError("chunk 必须含非空 source_id 与 chunk_key")
        self.chunks_dir.mkdir(parents=True, exist_ok=True)
        chunk_path = self.chunks_dir / f"{chunk["chunk_key"]}.json"
        _dump_json(chunk_path, {"meta": chunk, "text": text})
        if vector is None:
            embed_text = " ".join(str(chunk.get(k) or "") for k in ("title_path", "heading", "title")) + " " + text
            vector = _EMBED_FN([embed_text])[0]

        row = {k: v for k, v in chunk.items() if k != "chunk_key"}
        row.update({
            "name": chunk.get("title") or chunk.get("heading") or chunk.get("chunk_key", ""),
            "description": chunk.get("title_path") or "",
            "id": chunk["chunk_key"],
            "chunk_key": chunk["chunk_key"],
            "text_chars": len(text),
            "keywords": _meta_to_keywords(chunk),
        })
        # 幂等：同键先删旧行再追加（faiss 逐行删需重建，块规模下可接受）
        existing = [r for r in self.index.rows() if str(r.get("id")) == chunk["chunk_key"]]
        if existing:
            self.index.remove_item(chunk["chunk_key"])
        return self.index.add_item(row, vector)

    # ---- 检索 ----

    def search(self, query: str, top_k: int = DEFAULT_TOP_K,
               source: str | None = None, doc_kind: str | None = None) -> dict[str, Any]:
        """检索 flow_docs。query 为空走 latest。

        混合检索（语义 + 词面双通道）：
          - 语义通道：向量 top（top_k 条）
          - 词面通道：标题路径/关键词字面命中（排除向量已覆盖的块）
        各占一半配额（top_k 为偶数时对半；奇数时向量多 1），保持向量优先。

        source 过滤：corpus / stage_contract / skill；
        doc_kind 与 source 当前一一对应（corpus→corpus，contract/skill→自身）。
        """
        top_k = max(1, min(int(top_k), MAX_TOP_K))
        keywords = _meta_to_keywords({"title": query}) if query else None
        vector = _EMBED_FN([query or ""])[0] if query else [0.0] * self.index.dim

        result = self.index.search(
            vector, top_k=max(1, min(top_k * 4, MAX_TOP_K)),
            category=None, as_template=None, keywords=keywords,
        )
        items = result.get("items", [])
        vector_quota = (top_k + 1) // 2
        merged: list[dict[str, Any]] = []

        def _pass(row: dict[str, Any]) -> bool:
            if source and str(row.get("source") or "") != source:
                return False
            if doc_kind and str(row.get("doc_kind") or "") != doc_kind:
                return False
            return True

        # 1) 语义通道
        for row in items:
            if not _pass(row):
                continue
            merged.append(row)
            if len(merged) >= vector_quota:
                break

        # 2) 词面通道：全索引标题路径/heading 字面命中（不依赖向量通道候选项，
        #    精确词匹配兜底语义盲区——阶段名/文件名是用户提问的高频字面词）。
        #    候选按「命中词数 desc」排序，命中更全的块优先。
        if query:
            qwords = _ascii_keywords(query)
            if qwords:
                merged_ids = {str(r.get("id")) for r in merged}
                lex_hits: list[tuple[int, int, dict[str, Any]]] = []
                for row in self.index.rows():
                    if str(row.get("id")) in merged_ids:
                        continue
                    if not _pass(row):
                        continue
                    hay = " ".join(
                        str(row.get(k) or "") for k in ("title_path", "heading", "stage_name", "name")
                    ).lower()
                    matched = sum(1 for w in qwords if w in hay)
                    # 阶段名精确命中加权（stage_name == 查询词时 +3）
                    sn = str(row.get("stage_name") or "").lower()
                    if sn and sn in qwords:
                        matched += 3
                    if matched > 0:
                        lex_hits.append((matched, len(lex_hits), row))
                lex_hits.sort(key=lambda t: (-t[0], t[1]))
                # 去重：契约/技能按文件粒度，corpus 按章粒度（title_path），
                # 避免同文件多节块或同章多块霸占配额
                seen_keys: set[str] = set()
                for _, _, row in lex_hits:
                    if str(row.get("source") or "") == SOURCE_CORPUS:
                        key_path = str(row.get("title_path") or str(row.get("id")))
                    else:
                        key_path = str(row.get("path") or "")
                    if key_path and key_path in seen_keys:
                        continue
                    seen_keys.add(key_path)
                    merged.append(row)
                    if len(merged) >= top_k:
                        break
                # 2b) 正文字面通道：ASCII 词扫描块正文（标题未中但正文含软件/阶段名时兜底）
                if len(merged) < top_k:
                    for row in self.index.rows():
                        if str(row.get("id")) in merged_ids:
                            continue
                        if not _pass(row):
                            continue
                        doc = self.get_chunk(str(row.get("chunk_key") or row.get("id") or ""))
                        if not doc:
                            continue
                        body = str(doc.get("text") or "").lower()
                        if any(w in body for w in qwords):
                            merged.append(row)
                            if len(merged) >= top_k:
                                break
                # 3) 若词面/正文通道命中不足且向量通道未满，用向量余量补足
                if len(merged) < top_k:
                    for row in items:
                        if str(row.get("id")) in merged_ids:
                            continue
                        if not _pass(row):
                            continue
                        merged.append(row)
                        if len(merged) >= top_k:
                            break

        return {"mode": result.get("mode", "empty"), "items": merged[:top_k]}

    def get_chunk(self, chunk_key: str) -> dict[str, Any] | None:
        """取块正文（meta + text），供答案引用原文。"""
        path = self.chunks_dir / f"{chunk_key}.json"
        if not path.exists():
            return None
        try:
            return json.loads(path.read_text(encoding="utf-8"))
        except (ValueError, OSError) as exc:
            logger.warning("flow: 块读取失败 %s：%s", chunk_key, exc)
            return None

    # ---- 维护 ----

    def delete_source(self, source_id: str) -> int:
        """删除一个源的所有块（含正文文件）。返回删除数。"""
        removed = 0
        for row in self.index.rows():
            if str(row.get("source_id")) == source_id:
                if self.index.remove_item(str(row.get("id"))):
                    removed += 1
        for f in self.chunks_dir.glob(f"{source_id}__*.json"):
            f.unlink(missing_ok=True)
        return removed

    def stats(self) -> dict[str, Any]:
        """集合统计：总块数、按 source/doc_kind 分布、字符总量。"""
        rows = self.index.rows()
        per_source: dict[str, int] = {}
        per_kind: dict[str, int] = {}
        total_chars = 0
        for row in rows:
            per_source[str(row.get("source") or "?")] = per_source.get(str(row.get("source") or "?"), 0) + 1
            per_kind[str(row.get("doc_kind") or "?")] = per_kind.get(str(row.get("doc_kind") or "?"), 0) + 1
            total_chars += int(row.get("text_chars") or 0)
        return {"total_chunks": len(rows), "per_source": per_source,
                "per_doc_kind": per_kind, "total_chars": total_chars}