308 lines
13 KiB
Python
308 lines
13 KiB
Python
 
|
"""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}
|