#!/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())