语料 106 文件/724 块(全流程 79 章 118 + 契约 60 份 429 + 技能 45 份 177);QA 树与 vstore 既有集合零改动 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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}
|