adminandClaude Fable 5
adb97f49a8
语料 106 文件/724 块(全流程 79 章 118 + 契约 60 份 429 + 技能 45 份 177)；QA 树与 vstore 既有集合零改动 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
311 lines
12 KiB
Python
311 lines
12 KiB
Python
#!/usr/bin/env python3
|
||
"""ingest_flow_docs.py —— 生信流程语料切块入库 flow_docs（幂等可重跑）。
|
||
|
||
语料源（按任务书优先级）：
|
||
1) 仓库 docs/bioflow_corpus/PlantAI_现有分析全流程_基因家族_RNAseq_WGS_GWAS.md（79 章）
|
||
--corpus-file 覆盖路径；切块策略：按文件实际 #（一级/组）与 ##（二级/章）标题切块
|
||
2) dev gfa_core/schemas/stages/*_contract_v1.md（阶段契约，按 ## 节切块）
|
||
--stages-dir 覆盖路径
|
||
3) dev skills/*/SKILL.md（技能说明，按 ## 节切块）
|
||
--skills-dir 覆盖路径
|
||
|
||
输出：flow_dir/{index.json, chunks/*.json}（独立目录，与主库 vault/index 隔离）
|
||
幂等：块键 = source_id + '__' + md5(标题路径 + 源文件相对路径)[:12]，
|
||
同键重跑覆盖，源文件删除后 --prune 清理孤儿块。
|
||
|
||
依赖：vstore.embed（ollama bge-m3）；单元测试经 vstore.vault.set_embed_fn 注入假嵌入，
|
||
不依赖真服务。运行：/home/xzq/5090/teamWorkspace/zh/plantai-rag-new-venv/bin/python
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import hashlib
|
||
import json
|
||
import sys
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
# 脚本放 vstore/scripts/ 下；允许包内直跑（python -m）与路径执行
|
||
if __package__ in (None, ""):
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from vstore.flow import (
|
||
CITE_ANCHOR,
|
||
CITE_LINE,
|
||
DOC_KIND_CONTRACT,
|
||
DOC_KIND_CORPUS,
|
||
DOC_KIND_LABELS,
|
||
DOC_KIND_SKILL,
|
||
SOURCE_CORPUS,
|
||
SOURCE_SKILL,
|
||
SOURCE_STAGE_CONTRACT,
|
||
FlowStore,
|
||
)
|
||
from vstore.vault import get_embed_fn
|
||
|
||
REPO = Path(__file__).resolve().parents[2] # PlantAI-rag-new 仓库根
|
||
DEFAULT_FLOW_DIR = REPO / "vstore" / "data" / "flow_docs"
|
||
DEFAULT_CORPUS_FILE = REPO / "docs" / "bioflow_corpus" / "PlantAI_现有分析全流程_基因家族_RNAseq_WGS_GWAS.md"
|
||
# 只读语料源（dev 仓库）
|
||
DEFAULT_STAGES_DIR = Path("/home/xzq/5090/teamWorkspace/zh/PlantAI-dev/gfa_core/schemas/stages")
|
||
DEFAULT_SKILLS_DIR = Path("/home/xzq/5090/teamWorkspace/zh/PlantAI-dev/skills")
|
||
|
||
CHUNK_TEXT_CAP = 6000 # 单块嵌入/展示上限（字符，超长块只保留前部并在 meta 注明截断）
|
||
|
||
|
||
# ---- 文本工具 ----
|
||
|
||
def _split_headings(lines: list[str]) -> list[tuple[int, str, int]]:
|
||
"""提取标题切点：[(行号0基, 标题原文, 级别1/2), ...]，按出现顺序。"""
|
||
out: list[tuple[int, str, int]] = []
|
||
for i, line in enumerate(lines):
|
||
stripped = line.strip()
|
||
if stripped.startswith("# ") and not stripped.startswith("## "):
|
||
out.append((i, stripped, 1))
|
||
elif stripped.startswith("## "):
|
||
out.append((i, stripped, 2))
|
||
return out
|
||
|
||
|
||
def _chunk_key(source_id: str, title_path: str, rel_path: str) -> str:
|
||
digest = hashlib.md5((title_path + "|" + rel_path).encode("utf-8")).hexdigest()[:12]
|
||
return f"{source_id}__{digest}"
|
||
|
||
|
||
def _store_text(store: FlowStore, meta: dict[str, Any], key: str,
|
||
text: str, line_start: int, line_end: int) -> None:
|
||
meta["line_start"] = int(line_start)
|
||
meta["line_end"] = int(line_end)
|
||
cite = meta.get("cite")
|
||
if isinstance(cite, dict) and cite.get("line") == 1:
|
||
cite["line"] = int(line_start)
|
||
truncated = len(text) > CHUNK_TEXT_CAP
|
||
meta["truncated"] = truncated
|
||
store.add_chunk(
|
||
{**meta, "chunk_key": key, "created_at": _now()},
|
||
text[:CHUNK_TEXT_CAP],
|
||
)
|
||
|
||
|
||
def _now() -> str:
|
||
from datetime import datetime, timezone
|
||
return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
||
|
||
|
||
def _embed(texts: list[str]) -> list[list[float]]:
|
||
return get_embed_fn()(texts)
|
||
|
||
|
||
# ---- 语料 1：全流程文档 ----
|
||
|
||
def ingest_corpus(store: FlowStore, corpus_file: Path, corpus_rel: str) -> int:
|
||
"""按文件实际 #（一级/组）与 ##（二级/章）标题切块。
|
||
|
||
分组规则：一级标题定义组；组内二级标题各成一块；组标题到首个二级标题
|
||
之间的内容（含无二级标题的章正文，如第 79 章）并入组块；首个一级标题
|
||
之前的前导区（文档头 + 前导二级章）分别成块。三级标题不切分，留在父块内。
|
||
"""
|
||
if not corpus_file.exists():
|
||
raise FileNotFoundError(f"corpus 文件不存在：{corpus_file}")
|
||
lines = corpus_file.read_text(encoding="utf-8").splitlines()
|
||
cuts = _split_headings(lines)
|
||
if not cuts:
|
||
raise ValueError(f"corpus 无标题切点：{corpus_file}")
|
||
|
||
doc_title = lines[0].replace("# ", "").strip()
|
||
added = 0
|
||
|
||
def _add(path: list[str], text: str, start: int, end: int, note: str) -> None:
|
||
nonlocal added
|
||
if not text.strip():
|
||
return
|
||
key = _chunk_key(SOURCE_CORPUS, " | ".join(path), corpus_rel)
|
||
meta = _corpus_meta(doc_title, path, corpus_rel, text, note)
|
||
_store_text(store, meta, key, text, start + 1, end)
|
||
added += 1
|
||
|
||
l1_idx = [k for k, c in enumerate(cuts) if c[2] == 1]
|
||
|
||
# 前导区：文档头（标题 + 引言 + 结构图），以及首个 L1 之前的 ## 章（如 ## 1 文档阅读说明）
|
||
first_l1_pos = cuts[l1_idx[0]][0] if l1_idx else len(lines)
|
||
pre_children = [c for c in cuts if c[2] == 2 and c[0] < first_l1_pos]
|
||
header_end = pre_children[0][0] if pre_children else first_l1_pos
|
||
_add([doc_title], "\n".join(lines[0:header_end]).strip(), 0, header_end,
|
||
"文档头（引言/结构图）")
|
||
for ci, (cstart, ctitle, _) in enumerate(pre_children):
|
||
cend = pre_children[ci + 1][0] if ci + 1 < len(pre_children) else first_l1_pos
|
||
_add([doc_title, ctitle], "\n".join(lines[cstart:cend]).strip(),
|
||
cstart, cend, "前导章（二级标题）")
|
||
|
||
# 组内：L1 组块（标题 + 组首内容）+ 各 ## 块
|
||
for gi, k in enumerate(l1_idx):
|
||
start, title, _ = cuts[k]
|
||
end = cuts[l1_idx[gi + 1]][0] if gi + 1 < len(l1_idx) else len(lines)
|
||
children = [c for c in cuts if c[2] == 2 and start < c[0] < end]
|
||
first_child = children[0][0] if children else end
|
||
_add([title], "\n".join(lines[start:first_child]).strip(),
|
||
start, first_child, "一级标题组块（组首内容，无二级标题章正文并入）")
|
||
for ci, (cstart, ctitle, _) in enumerate(children):
|
||
cend = children[ci + 1][0] if ci + 1 < len(children) else end
|
||
_add([title, ctitle], "\n".join(lines[cstart:cend]).strip(),
|
||
cstart, cend, "二级标题（章节）切块")
|
||
return added
|
||
|
||
|
||
def _corpus_meta(doc_title: str, full_path: list[str], corpus_rel: str,
|
||
text: str, chunking_note: str) -> dict[str, Any]:
|
||
return {
|
||
"source": SOURCE_CORPUS,
|
||
"source_id": SOURCE_CORPUS,
|
||
"doc_kind": DOC_KIND_CORPUS,
|
||
"doc_kind_label": DOC_KIND_LABELS[DOC_KIND_CORPUS],
|
||
"title": doc_title,
|
||
"title_path": " | ".join(full_path),
|
||
"heading": full_path[-1],
|
||
"path": corpus_rel,
|
||
"cite": {"anchor": full_path[-1]},
|
||
"line_start": 1, # _store_text 会覆盖
|
||
"line_end": 1,
|
||
"chunking": chunking_note,
|
||
"truncated": False,
|
||
}
|
||
|
||
|
||
# ---- 语料 2/3：阶段契约与技能说明（按 ## 节切块） ----
|
||
|
||
def ingest_contracts(store: FlowStore, stages_dir: Path, stages_rel: str) -> int:
|
||
added = 0
|
||
files = sorted(stages_dir.glob("*_contract_v1.md"))
|
||
if not files:
|
||
raise FileNotFoundError(f"契约目录无 *_contract_v1.md：{stages_dir}")
|
||
for f in files:
|
||
added += _ingest_md_sections(
|
||
store, f, f"{stages_rel}/{f.name}",
|
||
kind=DOC_KIND_CONTRACT, source=SOURCE_STAGE_CONTRACT,
|
||
)
|
||
return added
|
||
|
||
|
||
def ingest_skills(store: FlowStore, skills_dir: Path, skills_rel: str) -> int:
|
||
added = 0
|
||
files = sorted(skills_dir.glob("*/SKILL.md"))
|
||
if not files:
|
||
raise FileNotFoundError(f"技能目录无 */SKILL.md：{skills_dir}")
|
||
for f in files:
|
||
added += _ingest_md_sections(
|
||
store, f, f"{skills_rel}/{f.name}",
|
||
kind=DOC_KIND_SKILL, source=SOURCE_SKILL,
|
||
)
|
||
return added
|
||
|
||
|
||
def _ingest_md_sections(store: FlowStore, file: Path, rel_path: str,
|
||
kind: str, source: str) -> int:
|
||
"""按 ## 节切块（含标题路径）；无 ## 时整文件一块。"""
|
||
lines = file.read_text(encoding="utf-8").splitlines()
|
||
h1 = next((l.strip() for l in lines if l.strip().startswith("# ")), "")
|
||
h1 = h1.replace("# ", "").strip() or file.stem
|
||
cuts = [(i, t) for i, t, lv in _split_headings(lines) if lv == 2]
|
||
stage_name = file.stem.removesuffix("_contract_v1")
|
||
if kind == DOC_KIND_SKILL:
|
||
stage_name = file.parent.name
|
||
|
||
def _meta(path: list[str], text: str) -> dict[str, Any]:
|
||
heading = path[-1]
|
||
return {
|
||
"source": source,
|
||
"source_id": source,
|
||
"doc_kind": kind,
|
||
"doc_kind_label": DOC_KIND_LABELS[kind],
|
||
"title": h1,
|
||
"title_path": " | ".join(path),
|
||
"heading": heading,
|
||
"stage_name": stage_name,
|
||
"path": rel_path,
|
||
"cite": {"anchor": heading, "line": 1},
|
||
"line_start": 1,
|
||
"line_end": 1,
|
||
"chunking": "二级标题（节）切块" if len(path) > 1 else "整文件（无二级标题）",
|
||
"truncated": False,
|
||
}
|
||
|
||
added = 0
|
||
if not cuts:
|
||
text = "\n".join(lines).strip()
|
||
if not text:
|
||
return 0
|
||
key = _chunk_key(source, h1, rel_path)
|
||
_store_text(store, _meta([h1], text), key, text, 1, len(lines))
|
||
return 1
|
||
for i, (start, heading) in enumerate(cuts):
|
||
end = cuts[i + 1][0] if i + 1 < len(cuts) else len(lines)
|
||
text = "\n".join(lines[start:end]).strip()
|
||
if not text:
|
||
continue
|
||
key = _chunk_key(source, f"{h1} | {heading}", rel_path)
|
||
_store_text(store, _meta([h1, heading], text), key, text, start + 1, end)
|
||
added += 1
|
||
# 文件级全文块：单节向量分散时由全文兜底（标题路径含文件级主题词）
|
||
full_text = "\n".join(lines).strip()
|
||
if full_text and len(full_text) > 200:
|
||
key = _chunk_key(source, f"{h1} | 全文", rel_path)
|
||
_store_text(store, _meta([h1, "全文"], full_text), key, full_text, 1, len(lines))
|
||
added += 1
|
||
return added
|
||
|
||
|
||
# ---- 主流程 ----
|
||
|
||
def main() -> int:
|
||
ap = argparse.ArgumentParser(description="生信流程语料切块入库 flow_docs（幂等）")
|
||
ap.add_argument("--flow-dir", default=str(DEFAULT_FLOW_DIR), help="flow_docs 数据目录")
|
||
ap.add_argument("--corpus-file", default=str(DEFAULT_CORPUS_FILE), help="全流程文档路径")
|
||
ap.add_argument("--stages-dir", default=str(DEFAULT_STAGES_DIR), help="阶段契约目录（dev 只读）")
|
||
ap.add_argument("--skills-dir", default=str(DEFAULT_SKILLS_DIR), help="技能说明目录（dev 只读）")
|
||
ap.add_argument("--skip-corpus", action="store_true", help="跳过全流程文档")
|
||
ap.add_argument("--skip-stages", action="store_true", help="跳过阶段契约")
|
||
ap.add_argument("--skip-skills", action="store_true", help="跳过技能说明")
|
||
ap.add_argument("--prune", action="store_true", help="清理索引中源文件已不存在的孤儿块")
|
||
args = ap.parse_args()
|
||
|
||
corpus_file = Path(args.corpus_file)
|
||
stages_dir = Path(args.stages_dir)
|
||
skills_dir = Path(args.skills_dir)
|
||
flow_dir = Path(args.flow_dir)
|
||
|
||
corpus_rel = f"docs/bioflow_corpus/{corpus_file.name}"
|
||
stages_rel = "gfa_core/schemas/stages"
|
||
skills_rel = "skills"
|
||
|
||
store = FlowStore(flow_dir)
|
||
counts: dict[str, int] = {}
|
||
if not args.skip_corpus:
|
||
counts["corpus"] = ingest_corpus(store, corpus_file, corpus_rel)
|
||
if not args.skip_stages:
|
||
counts["stages"] = ingest_contracts(store, stages_dir, stages_rel)
|
||
if not args.skip_skills:
|
||
counts["skills"] = ingest_skills(store, skills_dir, skills_rel)
|
||
|
||
pruned = 0
|
||
if args.prune:
|
||
# 保留：当前仍存在于三个源目录/文件中的所有块
|
||
known_files = {corpus_file.name}
|
||
known_files |= {f.name for f in stages_dir.glob("*_contract_v1.md")}
|
||
known_files |= {f.name for f in skills_dir.glob("*/SKILL.md")}
|
||
for row in store.index.rows():
|
||
fname = str(row.get("path") or "").rsplit("/", 1)[-1]
|
||
if fname and fname not in known_files:
|
||
store.index.remove_item(str(row.get("id")))
|
||
pruned += 1
|
||
|
||
stats = store.stats()
|
||
print(json.dumps({
|
||
"added_by_source": counts,
|
||
"pruned": pruned,
|
||
"stats": stats,
|
||
}, ensure_ascii=False, indent=2))
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|