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())
|