From 288829bb5b82ec56d50fbd8639248959e705889b Mon Sep 17 00:00:00 2001 From: LLLin000 <809867916@qq.com> Date: Mon, 6 Jul 2026 16:42:08 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20PR=201=20=E2=80=94=20nested=20structure?= =?UTF-8?q?=20tree=20+=20body=5Funits=20correctness?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - RenderOutput dataclass with heading_events/emitted_block_events tracking - Stack-algorithm structure tree with own/subtree block IDs - Recursive body_units builder with role helper + token cap splitting - Schema v4: section_path_json, section_level, section_title, part_ordinal - Phase 4/5 reorder in ocr_rebuild.py and ocr.py - 12 new test cases for tree nesting, own_block_ids, rendered-order intervals --- .../plans/memory-layer-implementation-plan.md | 634 ++++++++++++++++++ ...07-04-layer4-retrieval-substrate-design.md | 108 ++- ...2026-07-06-memory-layer-overhaul-design.md | 609 +++++++++++++++++ paperforge/core/io.py | 13 + paperforge/memory/schema.py | 18 +- paperforge/retrieval/structure_tree.py | 119 +++- paperforge/retrieval/units.py | 322 ++++++--- paperforge/worker/ocr.py | 32 +- paperforge/worker/ocr_rebuild.py | 30 +- paperforge/worker/ocr_render.py | 101 ++- tests/test_layer4_structure_tree.py | 182 +++-- tests/test_layer4_units_and_fts.py | 227 +++++-- tests/test_ocr_rebuild.py | 11 +- tests/unit/memory/test_schema.py | 6 +- 14 files changed, 2088 insertions(+), 324 deletions(-) create mode 100644 docs/plans/memory-layer-implementation-plan.md create mode 100644 docs/superpowers/specs/2026-07-06-memory-layer-overhaul-design.md diff --git a/docs/plans/memory-layer-implementation-plan.md b/docs/plans/memory-layer-implementation-plan.md new file mode 100644 index 00000000..730eb6e6 --- /dev/null +++ b/docs/plans/memory-layer-implementation-plan.md @@ -0,0 +1,634 @@ +# Memory Layer 改造 — 实施计划 + +> 基于：`2026-07-06-memory-layer-overhaul-design.md` +> 状态：待执行 +> 策略：5 个 PR 顺序执行，PR n+1 依赖 PR n + +--- + +## 总依赖图 + +``` +PR 1: 结构树 + body_units 正确性 + │ + ▼ +PR 2: Memory Builder 增量 + │ + ▼ +PR 3: Content Discovery 行为（语义依赖 PR 2，PR 2 后做） + │ + ▼ +PR 4: Embedding 双 collection（依赖 PR 1 + PR 2） + │ + ▼ +PR 5: Retrieve 合并（依赖 PR 4） +``` + +--- + +## PR 1：结构树和 body_units 正确性（最核心，不可拆分） + +**目标：** 从 rendered markdown H1-H6 产出真实嵌套结构树，body_units 跟随树递归生成，不重复、不丢内容。 + +### 文件变更 + +| 操作 | 文件 | 改动 | +|------|------|------| +| 新增 | `paperforge/core/io.py` | `read_jsonl(path)` — jsonl 逐行读取（待加） | +| 修改 | `paperforge/worker/ocr_render.py` | `render_fulltext_markdown()` 新增 `return_events=False` 参数；为 True 时返回 `RenderOutput(markdown, heading_events, emitted_block_events)`。**不破坏旧返回类型。** | +| 修改 | `paperforge/worker/ocr_rebuild.py` Phase 4/5 | Phase 4 传 `return_events=True` → `RenderOutput` 传入 Phase 5 | +| 修改 | `paperforge/worker/ocr.py` Phase 6 | 同上顺序调整 | +| 新增 | render 阶段产出 | `render/render-map.json`：`{headings: [...], emitted_blocks: [...]}` | +| 重写 | `paperforge/retrieval/structure_tree.py` | `build_structure_tree(heading_events, emitted_block_events, structured_blocks)` — stack 算法，interval 基于 emitted order | +| 修改 | `paperforge/retrieval/units.py` | `build_body_units()` / `build_object_units()` recursive walk，`own_block_ids`，role helper | +| 修改 | `paperforge/memory/schema.py` | body_units 表加字段；CURRENT_SCHEMA_VERSION → 4 | + +### 实现要点 + +#### 1.1 RenderOutput 与 render-map.json + +**不直接改 `render_fulltext_markdown()` 返回类型**——当前返回 `str`，旧调用点很多。新增 `return_events=False` 参数： + +```python +@dataclass +class RenderOutput: + markdown: str + heading_events: list[dict] # heading block 事件 + emitted_block_events: list[dict] # 所有真正进入正文阅读流的 block + +def render_fulltext_markdown(..., return_events=False) -> str | RenderOutput: + ... + if return_events: + return RenderOutput(markdown=markdown, + heading_events=heading_events, + emitted_block_events=emitted_block_events) + return markdown +``` + +`heading_events` 每条记录： + +```json +{"line_number": 42, "markdown_level": 3, "title": "Study population", + "page": 2, "block_id": "b_10", "emitted_order": 7} +``` + +`emitted_block_events` 每条记录（每个真正进入正文阅读流的 block）： + +```json +{"emitted_order": 123, "line_start": 456, "line_end": 460, + "page": 3, "block_id": "b_10", "role": "body_paragraph", + "emitted_as": "body"} +``` + +写入 `render/render-map.json`： + +```json +{"headings": [...], "emitted_blocks": [...]} +``` + +`build_structure_tree()` **不再依赖 structured_blocks 的 `_emitted_order`**（该字段不存在），而是接收 `emitted_block_events`。 + +#### 1.2 Phase 4/5 顺序和数据流 + +``` +× 当前: + Phase 4: render_fulltext_markdown(...) → markdown (str) + Phase 5: build_role_indexes + build_structure_tree(structured) ← 此时无 heading／emitted 信息 + write_render_outputs(...) + +✓ 改为: + Phase 4: rendered = render_fulltext_markdown(..., return_events=True) + return rendered, health_overall + + Phase 5: write_render_outputs(..., markdown=rendered.markdown, + heading_events=rendered.heading_events, + emitted_block_events=rendered.emitted_block_events) + build_role_indexes(...) + build_structure_tree( + heading_events=rendered.heading_events, + emitted_block_events=rendered.emitted_block_events, + structured_blocks=structured, + ) + write_structure_tree(index/structure-tree.json) +``` + +两处都改： +- `paperforge/worker/ocr_rebuild.py` — rebuild 路径 +- `paperforge/worker/ocr.py` — 初始 OCR 路径 + +`write_render_outputs()` 新增参数 `heading_events` / `emitted_block_events`，写入 `render/render-map.json`。 + +#### 1.3 build_structure_tree() 使用 emitted_block_events + +```python +def build_structure_tree(heading_events, emitted_block_events, structured_blocks): + # 1. heading_events 按 emitted_order 排序 + heading_events.sort(key=lambda h: h["emitted_order"]) + + # 2. 建立 block_id → structured_block 的映射 + block_map = {b["block_id"]: b for b in structured_blocks} + + # 3. stack 建树 + stack, root_nodes = [], [] + for h in heading_events: + node = { + "node_id": f"sec:{h['block_id']}", + "kind": "section", + "title": h["title"], + "level": h["markdown_level"], + "page_span": [h["page"], h["page"]], + "block_id": h["block_id"], + "own_block_ids": [], + "subtree_block_ids": [], + "children": [], + "objects": [], + } + while stack and stack[-1]["level"] >= h["markdown_level"]: + stack.pop() + if stack: + stack[-1]["children"].append(node) + else: + root_nodes.append(node) + stack.append(node) + + # 4. 计算 block intervals（基于 emitted_block_events） + _assign_block_intervals(root_nodes, heading_events, emitted_block_events) + return {"paper_id": ..., "nodes": root_nodes} +``` + +#### 1.4 block interval 计算（基于 emitted order） + +```python +def _assign_block_intervals(nodes, heading_events, emitted_block_events): + # 构建每个 heading 的 emitted_order 边界 + bounds_map = {} + for i, h in enumerate(heading_events): + next_sibling = None + for h2 in heading_events[i+1:]: + if h2["markdown_level"] <= h["markdown_level"]: + next_sibling = h2["emitted_order"] + break + bounds_map[h["block_id"]] = { + "start": h["emitted_order"], + "end": next_sibling or float("inf"), + } + + def compute(node): + bounds = bounds_map.get(node["block_id"]) + if not bounds: + return + # subtree = bounds 范围内的所有 emitted blocks + node["subtree_block_ids"] = [ + e["block_id"] for e in emitted_block_events + if bounds["start"] <= e["emitted_order"] < bounds["end"] + ] + for child in node.get("children", []): + compute(child) + # own = subtree − children subtree − heading blocks + child_ids = set() + for child in node.get("children", []): + child_ids.update(child.get("subtree_block_ids", [])) + node["own_block_ids"] = [ + bid for bid in node["subtree_block_ids"] + if bid not in child_ids and bid != node["block_id"] + ] + + for n in nodes: + compute(n) +``` + +#### 1.5 build_body_units() 递归 + +```python +def _body_unit_role_kind(role): + if role == "body_paragraph": return "body" + if role in {"structured_insert", "non_body_insert", "backmatter_body"}: + return "backmatter_body" + return None # reference_item, reference_heading, heading, caption, asset → excluded + +def build_body_units(tree, structured_blocks): + units = [] + block_map = {b["block_id"]: b for b in structured_blocks} + + def walk(node, path): + this_path = path + [node["title"]] + own = [block_map[bid] for bid in node.get("own_block_ids", []) + if bid in block_map + and _body_unit_role_kind(block_map[bid].get("role", ""))] + if own: + from itertools import groupby + for kind, grp in groupby(own, key=lambda b: _body_unit_role_kind(b["role"])): + if not kind: continue + text = "\n\n".join(b.get("text","") for b in grp) + parts = _split_if_oversized(text, 1000) + n = len(parts) + for pi, pt in enumerate(parts): + ord_ = pi + 1 if n > 1 else 0 + suf = f":part_{ord_:03d}" if ord_ else "" + units.append({ + "unit_id": f"{paper_id}:body:{node['node_id']}{suf}", + "paper_id": paper_id, + "section_path": "/".join(this_path), + "section_path_json": json.dumps(this_path), + "section_level": node["level"], + "section_title": node["title"], + "page_span": node["page_span"], + "unit_text": pt, + "token_estimate": len(pt)//4, + "unit_kind": kind, + "part_ordinal": ord_, + "indexable": True, + "veto_reason": "", + "quality_hints": [], + }) + for child in node.get("children", []): + walk(child, path) + + for root in tree.get("nodes", []): + walk(root, []) + return units +``` + +#### 1.6 build_object_units() own_block_ids + +与 body_units 同理，只取 `own_block_ids` 范围内 role 为 figure/table/object 的 blocks。一个 object 只归属一个最具体的 section。 + +#### 1.7 schema 升级 + +```sql +-- CURRENT_SCHEMA_VERSION: 3 → 4 +ALTER TABLE body_units ADD COLUMN section_path_json TEXT NOT NULL DEFAULT '[]'; +ALTER TABLE body_units ADD COLUMN section_level INTEGER NOT NULL DEFAULT 0; +ALTER TABLE body_units ADD COLUMN section_title TEXT NOT NULL DEFAULT ''; +ALTER TABLE body_units ADD COLUMN part_ordinal INTEGER NOT NULL DEFAULT 0; +``` + +### 测试清单 + +``` +1. duplicate heading title 通过 block_id 区分，不通过 title match +2. H2/H3/H3 sibling 产生正确 section_path（无跨级错误） +3. parent intro + child + parent summary: + parent.own_block_ids = [intro_para, summary_para]（不含 child 内容） +4. child body 不重复出现在 parent unit +5. empty container 不产生 body unit +6. reference_item 不进入 body_units +7. structured_insert 进入 backmatter_body unit +8. mixed body + backmatter: 拆成两个 unit，不同 unit_kind +9. token cap: >1000 tokens 拆 part_001 / part_002 +10. object 只归属最具体 section +11. rendered order ≠ structured order: tree interval 基于 emitted order +12. consumed object blocks 不进入 own_block_ids +``` + +--- + +## PR 2：Memory Builder 增量 + +**目标：** memory build 不再被 canonical_index_hash 锁定，per-paper 检测 OCR 结果变化后增量重建 body_units。 + +### 文件变更 + +| 操作 | 文件 | 改动 | +|------|------|------| +| 修改 | `paperforge/memory/builder.py` | 结构化路径修正；不 early return；per-paper ocr_result_hash 增量；三级 fallback；schema v4 兼容 | +| 修改 | `paperforge/memory/schema.py` | ensure_schema 支持 v4 upgrade path | + +### 实现要点 + +#### 2.1 路径修正 + +```python +# 改前 +structured_path = ocr_dir / "structured-blocks.json" +structured_blocks = read_json(structured_path) + +# 改后 +structured_path = ocr_dir / "structure" / "blocks.structured.jsonl" +structured_blocks = read_jsonl(structured_path) +``` + +#### 2.2 不 early return + +```python +# 改前 +if canonical_hash matches and db_path.exists(): + return {"papers_indexed": N, "hash_match": True} + +# 改后 +index_changed = not _hash_matches(...) +if index_changed: + _rebuild_metadata_tables(conn, items) # papers / aliases / FTS + _rebuild_all_units(conn, items, vault) # body + object + manifest +else: + _incremental_units_only(conn, items, vault) # only body/object units +``` + +#### 2.3 _incremental_units_only() + +```python +def _incremental_units_only(conn, items, vault): + ocr_root = vault / "System" / "PaperForge" / "ocr" + if not ocr_root.exists(): + return + for entry in items: + key = entry["zotero_key"] + paper_dir = ocr_root / key + tree_path = paper_dir / "index" / "structure-tree.json" + blocks_path = paper_dir / "structure" / "blocks.structured.jsonl" + if not tree_path.exists() or not blocks_path.exists(): + continue + current_hash = _resolve_ocr_result_hash(paper_dir) + row = conn.execute( + "SELECT value FROM meta WHERE key=?", (f"manifest:{key}",) + ).fetchone() + if row: + stored = json.loads(row[0]) + if stored.get("ocr_result_hash") == current_hash: + continue + _rebuild_paper_units(conn, key, paper_dir, tree_path, blocks_path) +``` + +#### 2.4 _resolve_ocr_result_hash() 三级 fallback + +```python +def _resolve_ocr_result_hash(paper_dir): + # 1. index/result-hash.txt + rp = paper_dir / "index" / "result-hash.txt" + if rp.exists(): + return rp.read_text().strip() + # 2. hash of structured artifacts + h = hashlib.sha256() + for rel in ["structure/blocks.structured.jsonl", "index/structure-tree.json", + "index/role-index.json"]: + p = paper_dir / rel + if p.exists(): + h.update(p.read_bytes()) + if h.hexdigest() != hashlib.sha256(b"").hexdigest(): + return h.hexdigest() + # 3. meta derived_version + meta_p = paper_dir / "meta.json" + if meta_p.exists(): + dv = json.loads(meta_p.read_text()).get("derived_version", {}) + return hashlib.sha256(json.dumps(dv, sort_keys=True).encode()).hexdigest() + return "" +``` + +#### 2.5 _rebuild_paper_units() 删除 + _upsert_body_units() 只插入 + +**职责分离：** `_rebuild_paper_units()` 负责 DELETE；`_upsert_body_units()` 只 INSERT + FTS。 + +```python +def _rebuild_paper_units(conn, key, paper_dir, tree_path, blocks_path): + # ── 删除（一个地方做）── + conn.execute("DELETE FROM body_units WHERE paper_id = ?", (key,)) + conn.execute("DELETE FROM body_units_fts WHERE paper_id = ?", (key,)) + conn.execute("DELETE FROM object_units WHERE paper_id = ?", (key,)) + + # ── 构建 ── + tree = json.loads(tree_path.read_text()) + blocks = read_jsonl(blocks_path) + from paperforge.retrieval.units import build_body_units, build_object_units + body_units = build_body_units(tree=tree, structured_blocks=blocks) + object_units = build_object_units(tree=tree, structured_blocks=blocks, ...) + + # ── 插入 + FTS（不做 DELETE）── + _upsert_body_units(conn, body_units) + _upsert_object_units(conn, object_units) + + # ── manifest ── + from paperforge.retrieval.manifest import build_paper_manifest + manifest = build_paper_manifest(paper_id=key, ...) + _write_manifest_row(conn, manifest) +``` + +`_upsert_body_units()` 只 INSERT + FTS，DELETE 已在 `_rebuild_paper_units()` 完成： + +```python +def _upsert_body_units(conn, body_units): + for unit in body_units: + conn.execute("""INSERT INTO body_units (...) VALUES (...)""", ...) + paper_ids = list({u["paper_id"] for u in body_units}) + for pid in paper_ids: + conn.execute("""INSERT INTO body_units_fts(...) + SELECT rowid, unit_id, paper_id, section_path, unit_text + FROM body_units + WHERE paper_id = ? AND indexable = 1""", (pid,)) +``` + +--- + +## PR 3：Content Discovery 行为 + +**目标：** body_units_fts 做 primary 检索，无结果时不静默 fallback。 + +**依赖：** PR 2 后做（依赖 body_units 表稳定 + indexable 语义稳定 + schema v4）。 + +### 文件变更 + +| 操作 | 文件 | 改动 | +|------|------|------| +| 修改 | `paperforge/retrieval/gateway.py` | `_run_body_unit_discovery()` 无结果时返回空 + coverage 提示 | + +### 实现要点 + +```python +def _run_body_unit_discovery(vault, query, limit=5): + body_rows = _search_body_units_fts(vault, query, limit) + coverage = _get_body_coverage(conn) + if not body_rows: + return PFResult(ok=True, data={ + "results": [], + "coverage": coverage, + "next_action": { + "command": f"paperforge search {query}", + "reason": (f"正文检索无匹配。正文索引覆盖 " + f"{coverage['body_papers']}/{coverage['ocr_papers']} 篇 OCR 完成论文。" + f"尝试 paperforge search 进行元数据全文搜索。") + } + }) + return PFResult(ok=True, data={"results": body_rows, "coverage": coverage}) + +def _get_body_coverage(conn): + body = conn.execute("SELECT COUNT(DISTINCT paper_id) FROM body_units WHERE indexable=1").fetchone()[0] + ocr = conn.execute("SELECT COUNT(*) FROM papers WHERE ocr_status='done'").fetchone()[0] + lib = conn.execute("SELECT COUNT(*) FROM papers").fetchone()[0] + return {"body_papers": body, "ocr_papers": ocr, "library_papers": lib} +``` + +--- + +## PR 4：Embedding 双 collection + +**目标：** 新建 `paperforge_body` collection，embed builder 从 DB 读 body_units。 + +### 文件变更 + +| 操作 | 文件 | 改动 | +|------|------|------| +| 修改 | `paperforge/embedding/_chroma.py` | `get_collection()` 支持 `name` 参数；`delete_paper_vectors()` 双删 | +| 修改 | `paperforge/embedding/builder.py` | 新增 `embed_body_units()`；`get_body_units_for_embedding()` 读 DB | +| 修改 | `paperforge/embedding/search.py` | 新增 `merge_retrieve()` | +| 修改 | `paperforge/embedding/status.py` | 统计两个 collection | +| 修改 | `paperforge/commands/embed.py` | build 流程以 DB 分流、resume 检测 body_units_hash | + +### 实现要点 + +#### 4.1 get_collection(name=) + +```python +def get_collection(vault, name="paperforge_fulltext"): + db_path = get_vector_db_path(vault) + client = chromadb.PersistentClient(str(db_path)) + return client.get_or_create_collection( + name=name, metadata={"hnsw:space": "cosine"}, + ) +``` + +#### 4.2 embed build 分流（以 DB 为准） + +是否走 body_units 路径以 DB `body_units` 表为准（Memory Builder 是唯一生产者），不重复检测文件。 + +```python +def _has_body_units_in_db(vault, key) -> bool: + db_path = get_memory_db_path(vault) + if not db_path.exists(): + return False + conn = get_connection(db_path, read_only=True) + cnt = conn.execute( + "SELECT COUNT(*) FROM body_units WHERE paper_id=? AND indexable=1", + (key,) + ).fetchone()[0] + conn.close() + return cnt > 0 + +for entry in papers: + key = entry["zotero_key"] + has_body = _has_body_units_in_db(vault, key) + + if has_body: + body_units = get_body_units_for_embedding(vault, key) + if not body_units: + continue + if resume: + col = get_collection(vault, name="paperforge_body") + # Chroma: 不依赖多字段 where 隐式 AND + existing = col.get(where={"paper_id": key}, limit=1) + metas = existing.get("metadatas", []) + if metas and metas[0].get("body_units_hash") == current_hash: + papers_skipped += 1; continue + delete_paper_vectors(vault, key) + embed_body_units(vault, key, body_units) + else: + # 文件存在但 DB 无 body_units → 提示先 memory build + ocr_root = vault / "System" / "PaperForge" / "ocr" / key + has_files = ((ocr_root / "structure" / "blocks.structured.jsonl").exists() + and (ocr_root / "index" / "structure-tree.json").exists()) + if has_files: + print(f"Skip {key}: has structured blocks but no body_units in DB. " + f"Run `paperforge memory build` first.") + continue + # 旧路径 + fulltext_path = vault / (entry.get("fulltext_path") or "") + if resume: + col = get_collection(vault, name="paperforge_fulltext") + existing = col.get(where={"paper_id": key}, limit=1) + if existing.get("ids"): + papers_skipped += 1; continue + delete_paper_vectors(vault, key) + chunks = chunk_fulltext(fulltext_path) + if chunks: + embed_paper(vault, key, chunks) +``` + +--- + +## PR 5：Retrieve 合并 + +**目标：** 同时查两个 collection，unit-level 去重，per-paper cap。 + +### 文件变更 + +| 操作 | 文件 | 改动 | +|------|------|------| +| 新增 | `paperforge/embedding/search.py` | `merge_retrieve()` | +| 修改 | `paperforge/commands/retrieve.py` | 调 `merge_retrieve()` 替代 `retrieve_chunks()` | + +### 实现要点 + +```python +RETRIEVAL_COLLECTIONS = ["paperforge_fulltext", "paperforge_body"] + +def merge_retrieve(vault, query, limit=5, expand=True): + provider = OpenAICompatibleProvider(vault) + q_emb = provider.encode_single(query) + n = limit * 2 + + all_results = [] + for name in RETRIEVAL_COLLECTIONS: + try: + col = get_collection(vault, name=name) + res = col.query( + query_embeddings=[q_emb], + n_results=n, + include=["documents", "metadatas", "distances"], + ) + for doc, meta, dist in zip(res["documents"][0], res["metadatas"][0], res["distances"][0]): + all_results.append({ + "paper_id": meta.get("paper_id", ""), + "section_path": meta.get("section_path", ""), + "chunk_text": doc, + "score": round(1.0 - dist, 4), + "source": "legacy_chunk" if name == "paperforge_fulltext" else "body_unit", + "unit_id": meta.get("chunk_index", ""), + }) + except Exception: + continue + + # 排序 + unit-level 去重 + per-paper cap + all_results.sort(key=lambda r: r["score"], reverse=True) + seen = set() + per_paper = {} + merged = [] + for r in all_results: + # dedupe key: (source, unit_id) 或 (source, paper_id+text_hash) 防空 unit_id + dedupe_key = (r["source"], r["unit_id"]) if r.get("unit_id") else ( + r["source"], r["paper_id"], hash(r["chunk_text"]) + ) + if dedupe_key in seen: + continue + seen.add(dedupe_key) + pid = r["paper_id"] + if per_paper.get(pid, 0) >= 2: + continue + per_paper[pid] = per_paper.get(pid, 0) + 1 + merged.append(r) + if len(merged) >= limit: + break + return merged +``` + +--- + +## 执行顺序 + +``` +Step 1: PR 1（最核心，不可拆） + → 验证: body_units 表有行，section_path 正确 + → 验证: rendered order ≠ structured order 时边界正确 + +Step 2: PR 2 + → 验证: rebuild 后 memory build 增量更新 body_units + → 验证: content-discovery 返回正文结果 + +Step 3: PR 3（依赖 PR 2） + → 验证: content-discovery 无结果时返回空 + coverage 提示 + +Step 4: PR 4 + → 验证: embed build 写入 paperforge_body + → 验证: resume 正确跳过已嵌入的论文（通过 body_units_hash） + +Step 5: PR 5 + → 验证: retrieve 返回合并结果，unit-level 去重 + per-paper cap +``` diff --git a/docs/superpowers/specs/2026-07-04-layer4-retrieval-substrate-design.md b/docs/superpowers/specs/2026-07-04-layer4-retrieval-substrate-design.md index 9e7efea0..05b5193a 100644 --- a/docs/superpowers/specs/2026-07-04-layer4-retrieval-substrate-design.md +++ b/docs/superpowers/specs/2026-07-04-layer4-retrieval-substrate-design.md @@ -87,6 +87,43 @@ Layer 4 不负责 PaperCard / SubmethodCard 这类理解层资产。那些内容 5. **Unified surface, explicit routing**：对 agent 暴露少量统一子命令，内部做明确意图路由，而不是黑箱混查。 6. **Decompose before concluding absence**：对于找文章任务，必须先拆 query、多路组合检索，禁止一次零结果就宣告不存在。 +## 设计决策：Chunker 切换策略 + +### 背景 +当前 `memory/chunker.py` 仍使用固定 3 段一组的文本分块方式，且在向量构建（`embedding/builder.py`）中直接调用。但新 rebuild 系统产出的 fulltext 附带完整的结构化产物（structure-tree.json + structured-blocks.json），可以从中提取受章节/边界约束的 body units，作为更高质量的检索单元。 + +### 约束 +库中同时存在两种 fulltext 来源： +- **旧版 fulltext**：来自初始 `do_ocr`，只有 `fulltext.md`，无结构化角色信息 +- **新版 rebuild fulltext**：来自 `ocr rebuild`，附带 `structure/blocks.structured.jsonl`（块已过角色判断）+ `index/structure-tree.json` + +不能要求用户对整个库统一做 rebuild。切换必须是逐论文的。 + +### 决策 +embedding 的切块逻辑根据每篇论文的 **最终角色确认产物** 做切换： + +| 检测信号 | 含义 | 切块策略 | +|---------|------|---------| +| `structure/blocks.structured.jsonl` 存在 | 该论文经过 rebuild，块已通过角色判断 | 从 body_units 构建 chunks（每 unit = 一个 chunk） | +| 该文件不存在 | 旧版 fulltext，无角色信息 | 回退到 `memory/chunker.py` 的三段一组文本分块 | + +**不引入版本标号**：直接用文件存在性判断，与 `memory/builder.py` 现有检测逻辑一致。 + +### 实现要点 + +**检测点：** 在 `embedding/builder.py` 中，对每篇论文检查 `{ocr_dir}/structure/blocks.structured.jsonl` 是否存在。 +若存在，走新切块路径；否则走旧路径。与 `memory/builder.py` 的检测逻辑完全一致（见 `builder.py:284-285`）。 +1. 读取 `structure-tree.json` 和 `structure/blocks.structured.jsonl` +2. 调用 `retrieval/units.build_body_units()` 得到 body units +3. 每个 body unit 的 `unit_text` 作为一个 chunk +4. chunk metadata 映射：`section=section_path`、`page_number=page_span[0]`、`chunk_index=unit_id`、`token_estimate` 复用原逻辑 +5. 写入向量库的 document = unit_text + +**旧回退路径（无结构树时）：** +维持现有 `memory/chunker.py` → `embedding/builder.embed_paper()` 不变 + +**不回溯旧论文：** 用户不主动触发 rebuild 的论文继续用旧 chunker，不做转换。重建时自动获得新路径。 + ## 非目标 本层明确不做： @@ -123,10 +160,9 @@ Layer 4 明确定义为两个模式： ### Retrieval Unit 可被索引、召回、展示的最小检索单元。不是任意三段文本，而是带类型的论文原生对象。 -### Body Unit -正文检索单元。来源于章节树、section/subsection 边界与正文 families。服务跨论文正文证据召回。 +### Body Unit ✅ -### Object Unit +### Object Unit ✅ 对象检索单元。包含： - Figure Unit - Table Unit @@ -138,10 +174,10 @@ Layer 4 明确定义为两个模式： - section path - page span -### Structure Tree +### Structure Tree ✅ 论文目录结构。描述 section / subsection / object 在论文中的组织关系。是 Mode 2 的核心导航面，不是向量检索结果。 -### Structure Tree Builder +### Structure Tree Builder ✅ Structure Tree 不是当前 OCR 的现成产物，而是 Layer 4 需要新增的中间层。它的最低构建链路是： ```text @@ -174,9 +210,9 @@ structured_blocks / role-index } ``` -### Paper Manifest -构建控制面资产，记录每篇论文当前 retrieval 产物与其输入、策略版本的绑定关系。 +### Paper Manifest ✅ +构建控制面资产，记录每篇论文当前 retrieval 产物与其输入、策略版本的绑定关系。 最少字段： - `paper_id` - `ocr_result_hash` @@ -237,10 +273,10 @@ paper_id:body:node_id:startPage-startBlock:endPage-endBlock ### 物理存储规则 默认 SQLite 层： -- `paper_fts` —— 只服务 metadata lookup / paper lookup -- `body_units` table -- `body_units_fts` virtual table -- `object_units` table +- ✅ `paper_fts` —— 只服务 metadata lookup / paper lookup +- ✅ `body_units` table +- ✅ `body_units_fts` virtual table +- ✅ `object_units` table - `object_units_fts` —— optional / later Vector 层： @@ -278,19 +314,19 @@ Layer 4 必须显式承认默认能力层与增强能力层的区别： ### Agent-facing surface 对 agent 暴露四个统一子命令： -- `paperforge paper-lookup` -- `paperforge content-discovery` -- `paperforge paper-navigation` -- `paperforge scoped-fetch` +- ✅ `paperforge paper-lookup` +- ✅ `paperforge content-discovery` +- ✅ `paperforge paper-navigation` +- ✅ `paperforge scoped-fetch` 这四个命令共享同一个内部 routing / resolution core，但对外保持显式 intent，而不是黑箱混查。 关键不在于把一切折叠成一个万能命令，而在于**agent 默认不再直接编排一堆底层 CLI**。 兼容规则： -- 现有 `search / retrieve / query-plan / paper-status / paper-context / context` 保留 -- 新四个命令是 agent-facing gateway surface -- 旧命令作为 compatibility aliases / low-level diagnostics，不在 Layer 4 删除 +- ✅ 现有 `search / retrieve / query-plan / paper-status / paper-context / context` 保留 +- ✅ 新四个命令是 agent-facing gateway surface +- ✅ 旧命令作为 compatibility aliases / low-level diagnostics，不在 Layer 4 删除 ### Internal arms @@ -414,8 +450,8 @@ Do not introduce a second query decomposition implementation. 只有多条高优先级路径都失败，才允许返回“暂未定位”。 当前代码已有基础： -- `QuerySignals` 已能拆 DOI / Zotero key / citation key / author tokens / year tokens / title-like tokens / content terms -- `lookup_paper()` 已有 exact lookup、title token AND lookup、alias lookup +- ✅ `QuerySignals` 已能拆 DOI / Zotero key / citation key / author tokens / year tokens / title-like tokens / content terms +- ✅ `lookup_paper()` 已有 exact lookup、title token AND lookup、alias lookup Layer 4 需要补的是： - author/year 交集 @@ -518,30 +554,30 @@ Chroma 更适合作为： ## 分阶段迁移 -### Phase 1：Gateway over existing capabilities +### Phase 1：Gateway over existing capabilities ✅ 先做统一入口和 routing shell，包装现有能力： -- `paperforge paper-lookup` -> `lookup_paper()` / `paper-status` / `paper-context` -- `paperforge content-discovery` -> metadata FTS + existing vector retrieve + warning -- `paperforge paper-navigation` -> 暂时返回 `role-index` summary，不承诺 section tree -- `paperforge scoped-fetch` -> 暂时支持 paper-level fulltext / block-id fetch +- ✅ `paperforge paper-lookup` -> `lookup_paper()` / `paper-status` / `paper-context` +- ✅ `paperforge content-discovery` -> metadata FTS + existing vector retrieve + warning +- ✅ `paperforge paper-navigation` -> 暂时返回 `role-index` summary，不承诺 section tree +- ✅ `paperforge scoped-fetch` -> 暂时支持 paper-level fulltext / block-id fetch 此阶段的目标是解决“入口太碎”和 single-arm false absence，不强碰向量后端。 -### Phase 2：Structure Tree + BodyUnit FTS +### Phase 2：Structure Tree + BodyUnit FTS ✅ 新增真正的 Layer 4 默认产物： -- `structure-tree.json` -- `body_units.jsonl` -- `body_units_fts` -- `object_units.jsonl` -- `paper_manifest.json` +- ✅ `structure-tree.json` +- ⚠️ `body_units.jsonl`（在 memory DB 中，非独立 jsonl 文件） +- ✅ `body_units_fts` +- ⚠️ `object_units.jsonl`（在 memory DB 中） +- ✅ `paper_manifest.json`（存储在 meta 表） UnitBuilder 输入源规则： -- 优先 `structured_blocks` -- `role-index` 只作辅助 -- `fulltext` 只作 fallback，不是主真相 +- ✅ 优先 `structured_blocks` +- ✅ `role-index` 只作辅助 +- ✅ `fulltext` 只作 fallback，不是主真相 ### Phase 3：Vector adapter -把 Chroma 从“硬编码后端”降成“一个 backend adapter”。 +把 Chroma 从"硬编码后端"降成"一个 backend adapter"。 当前 hard-coded 路径包括： - `embed_paper()` 直接调用 `get_collection()` @@ -552,7 +588,7 @@ UnitBuilder 输入源规则： ### Phase 4：Lance backend 接入与对比 在同一 retrieval contract 下，对比： -- 命中质量 + - build 时间 - 本地稳定性 - agent 使用体验 diff --git a/docs/superpowers/specs/2026-07-06-memory-layer-overhaul-design.md b/docs/superpowers/specs/2026-07-06-memory-layer-overhaul-design.md new file mode 100644 index 00000000..229ea98b --- /dev/null +++ b/docs/superpowers/specs/2026-07-06-memory-layer-overhaul-design.md @@ -0,0 +1,609 @@ +# Memory Layer 改造设计 + +> 设计讨论：2026-07-06 +> 状态：设计稿（待实施） +> 关联设计：`2026-07-04-layer4-retrieval-substrate-design.md` + +--- + +## 1. 背景与动机 + +### 1.1 当前问题 + +1. **结构树 (structure-tree) 是平的** — `children` 永远为 `[]`，没有真实嵌套层级。`level` 只有 1 或 2（来自 `section_heading` / `subsection_heading` 二元角色），无法表达 depth 3/4/5。 +2. **`section_path` 靠猜** — 从上一个 heading 的标题推断父子关系，在多同级子节时路径错误。 +3. **`structured-blocks.json` 路径不存在** — memory builder 期望 `{ocr_dir}/structured-blocks.json`，但 rebuild 实际写入 `{ocr_dir}/structure/blocks.structured.jsonl`（jsonl 格式）。导致 body_units 永远 0 行。 +4. **memory build 的 fast-path 门控太宽** — `canonical_index_hash` 不变时完全跳过重建，但 rebuild 产生的结构化数据变化（`ocr_result_hash` 变了）不被感知。 +5. **向量库单 collection** — 旧 chunker 和 新 body_units 共用同一 collection，无法增量迁移。 + +### 1.2 核心原则 + +- **Structure-first**：利用论文结构与边界信息，优先使用 body_units 作为检索单元 +- **增量切换，用户无感**：旧 chunker 逐步替换为 body_units，rebuilt 论文自动迁移 +- **完整路径追踪**：每个 body unit 记录从根到叶子的完整 section_path +- **FTS 兜底**：body_units_fts 是默认正文检索，内容 discovery 不再自动 fallback 到 metadata +- **单一生产者**：Memory Builder 是 body_units 的唯一生产者，Embedding Builder 从 paperforge.db 读取，不重复构建 + +--- + +## 2. 数据流全景 + +``` +OCR / Rebuild + ↓ +structure/blocks.structured.jsonl ← 最终 role blocks +fulltext.md ← renderer 已有 H1-H6 heading 层级 +index/structure-tree.json ← 从 H1-H6 + structured blocks 建真实树 +index/role-index.json +index/result-hash.txt / manifest hash + ↓ +Memory Builder (build_from_index) + ↓ +paperforge.db + ├─ papers / aliases / paper_fts ← 论文元数据检索 + ├─ body_units ← 正文结构化检索单元（唯一生产者） + ├─ body_units_fts ← 正文 FTS + ├─ object_units ← figure / table / object 周边证据 + └─ manifest: ← 每篇 OCR 结果 hash / policy hash + ↓ +Embedding Builder (读取 DB body_units，不重新构建) + ├─ paperforge_fulltext ← 旧 chunker，兼容旧论文 + └─ paperforge_body ← 新 body_units + ↓ +Retrieve / Content Discovery + ├─ content-discovery → body_units_fts + ├─ retrieve → fulltext + body 双 collection 合并 + └─ paper-lookup → metadata / aliases +``` + +--- + +## 3. 结构树改造 (Structure Tree) + +### 3.1 现状问题 + +当前 `retrieval/structure_tree.py` 的 `build_structure_tree()` 从 structured blocks 的 `role` 字段读取 heading，但只有 `section_heading` / `subsection_heading` 二元区分，丢失了 H3-H6 的层级信息。 + +### 3.2 改造方案 + +从 `fulltext.md` 的 H1-H6 标记读 heading 层级。但 **不通过 fuzzy title match 反推 block_id**——Methods / Results 等重复标题会导致歧义。 + +**方案：render 阶段产出 heading-map.json sidecar** + +```json +{ + "headings": [ + {"line_number": 42, "markdown_level": 3, "title": "Study population", + "page": 2, "block_id": "b_10"} + ] +} +``` + +写入位置：`render/heading-map.json`，与 `render/fulltext.md` 同时产出。 +`build_structure_tree()` 读 `render/heading-map.json` + `structured_blocks`，不从 markdown 反推 block_id。 + +**算法（PaperIndex 的 stack 算法）：** + +``` +for each heading in heading_map（按 page + line_number 顺序）: + while stack 不空 and stack[-1].level >= current.level: + stack.pop() + if stack 不空: + stack[-1].children.append(current) + else: + root_nodes.append(current) + stack.append(current) +``` + +### Phase 5 顺序约束（关键） + +新版依赖 `render/heading-map.json`，但当前 rebuild Phase 5 的顺序是先建树、后写 render： + +``` +× build_structure_tree(structured) → 先建树（无 heading-map） +× write_structure_tree(index/) +× write_render_outputs(...) → 后写 render（含 heading-map） +``` + +**必须改成：** + +``` +1. write_render_outputs(..., heading_events=heading_events) +2. build_structure_tree(heading_events, structured_blocks) +3. write_structure_tree(index/structure-tree.json) +``` + +更干净的做法：`render_fulltext_markdown()` 直接返回 `(markdown, heading_events)`，`write_render_outputs()` 写 markdown + heading-map，`build_structure_tree_from_heading_events()` 用内存中的 heading_events，不再读文件。 + +### 3.3 每个节点的字段 + +```python +{ + "node_id": "sec:", + "kind": "section", + "title": "Study population", + "level": 3, # H1=1, H2=2, ..., H6=6 + "section_path": [ # 完整祖先链 + "Materials and methods Design", + "2.1 Study Design", + "2.1.1 Study population" + ], + "page_span": [2, 3], + "own_block_ids": ["b_10", "b_11"], # 只属于当前节点本体的 block_ids + # 范围 = 从当前 heading 到第一个 child heading 前 + "subtree_block_ids": ["b_10".."b_25"], # 整个 subtree 的 block_ids（含子节点） + "children": [], + "objects": [] +} +``` + +### 3.4 own_block_ids vs subtree_block_ids（关键修正） + +父节点的 block_span 如果定义为"从当前 heading 到下一个同级/更高级 heading"，则会包含所有子节内容。递归构建 body unit 时，子节正文会被重复纳入父 unit。 + +且"当前 heading → 第一个 child heading 前"会漏掉子节结束后又回到父节的内容（如 Discussion 开头的 intro + 中间的子节 + 结尾的 summary）。 + +**修正：拆成两个概念** + +| 字段 | 含义 | 用途 | +|------|------|------| +| subtree_block_ids | 整棵 subtree 的 block_ids（含所有子孙节点） | 导航/scope fetch | +| own_block_ids | subtree_block_ids − 所有 child.subtree_block_ids − heading blocks | body unit 正文 | + +`own_block_ids` = 当前节点 subtree 内**不属于任何子节点 subtree** 的直接内容。这自动覆盖： +- 父节 intro（第一个 child 前） +- 父节 child 后面的 summary/transition +- 父节中间未归入子节的段落 + +举例： + +``` +## Discussion + intro paragraph ← own + ### Mechanism + mechanism text ← child subtree + ### Limitation + limitation text ← child subtree + summary paragraph ← own（子节结束后回到父节） + +Discussion.own_block_ids = [intro_para, summary_para] +Discussion.subtree_block_ids = [intro_para, mechanism_text, limitation_text, summary_para] +``` + +### 阅读顺序依赖（关键） + +`own_block_ids` / `subtree_block_ids` 的 interval 计算必须基于 **renderer 的最终阅读顺序**，而非 structured_blocks 的原始列表顺序。 + +原因：structured_blocks 经过双栏重排、tail reorder、reference zone 移序等处理后，block 顺序已不等于最终阅读顺序。 + +**做法：** + +Renderer 输出 `render/block-order-map.json`（或 heading_events 同时记录 `emitted_order`）。每个 entry：`{block_id, page, emitted_order, role}`。 + +Tree builder 在计算 subtree interval 前，先将 structured_blocks 按 `emitted_order` 重排。`subtree_block_ids` 在当前节点 heading 的 `emitted_order` 到下一个同级 heading 的 `emitted_order` 之间选取。然后按 `subtree − descendants` 计算 `own_block_ids`。 + +如果 PR 1 暂不做完整 `block-order-map`，至少让 heading-map 每条记录 `emitted_order`，tree builder 依赖 heading 顺序而非 blocks 顺序。 + +## 4. Body Units 生成 + +### 4.1 单一生产者原则 + +**Memory Builder 是 body_units 的唯一生产者。** Embedding Builder 从 paperforge.db.body_units 读取 indexable=1 的行，不重新调用 build_body_units()。这样保证 FTS 和 vector 用同一批 unit。 + +### 4.2 build_body_units() 算法 + +使用 role helper 统一处理正文角色和 backmatter 角色： + +```python +def _body_unit_role_kind(role: str) -> str | None: + if role == "body_paragraph": + return "body" + if role in {"structured_insert", "non_body_insert", "backmatter_body"}: + return "backmatter_body" + return None # reference_item / reference_heading / heading / caption / asset → excluded + + +def build_body_units(*, tree: dict, structured_blocks: list[dict]) -> list[dict]: + units = [] + + def walk(node, inherited_path=[]): + this_path = inherited_path + [node["title"]] + + # 收集 own_block_ids 中 body_paragraph / backmatter 的 blocks + own_blocks = [ + b for b in structured_blocks + if b["block_id"] in node.get("own_block_ids", []) + and _body_unit_role_kind(b.get("role", "")) is not None + ] + + if own_blocks: + # 按 unit_kind 分组（body / backmatter_body 不混） + from itertools import groupby + groups = [] + for kind, grp in groupby(own_blocks, key=lambda b: _body_unit_role_kind(b.get("role", ""))): + if kind: + groups.append((kind, list(grp))) + + for unit_kind, blocks in groups: + all_text = "\n\n".join(b.get("text", "") for b in blocks) + parts = _split_if_oversized(all_text, max_tokens=1000) + n_parts = len(parts) + for p_idx, part_text in enumerate(parts): + part_ordinal = p_idx + 1 if n_parts > 1 else 0 + suffix = f":part_{part_ordinal:03d}" if part_ordinal else "" + uid = f"{paper_id}:body:{node['node_id']}{suffix}" + unit = { + "unit_id": uid, + "paper_id": paper_id, + "section_path": "/".join(this_path), + "section_path_json": json.dumps(this_path), + "section_level": node["level"], + "section_title": node["title"], + "page_span": node["page_span"], + "unit_text": part_text, + "token_estimate": len(part_text) // 4, + "unit_kind": unit_kind, + "part_ordinal": part_ordinal, + "indexable": True, + "veto_reason": "", + "quality_hints": [], + } + units.append(unit) +``` +### 4.3 Token Cap 与拆分 + +leaf section 的 token_estimate 超过 1000 时，按 paragraph boundary 拆分为多个 part。拆分后保留相同的 section_path / section_title / section_level，加 part_ordinal。 + +``` +unit_id: + paper:body:sec_node_id ← 单块（≤1000 tokens） + paper:body:sec_node_id:part_001 ← 多块 + paper:body:sec_node_id:part_002 +REM + +### 4.4 输出格式 + +```python +{ + "unit_id": "23T2B8ZX:body:sec_10:e2-e10:e2-e14", + "paper_id": "23T2B8ZX", + "section_path": "Materials and methods/Study population", + "section_path_json": '["Materials and methods", "Study population"]', + "section_level": 3, + "section_title": "Study population", + "page_span": [2, 2], + "unit_text": "...", + "token_estimate": 508, + "unit_kind": "body", + "part_ordinal": 1, # 0 = 单块, 1+ = 拆分序号 + "indexable": True, + "veto_reason": "", + "quality_hints": [], +} +``` + +### 4.5 Backmatter / References 处理 + +| 内容类型 | role | 进 body_units? | unit_kind | indexable | +|---------|------|---------------|-----------|-----------| +| 正文段落 | body_paragraph | ✅ | "body" | True | +| 致谢/数据声明/作者贡献 | structured_insert / non_body_insert | ✅ | "backmatter_body" | True | +| 参考文献条目 | reference_item | ❌ 不进 body_units | — | — | +| References 标题 | reference_heading | ❌ | — | — | + +reference_item 不进 body_units：否则 query 作者名/期刊名时会命中 references 而非正文证据。 + +### 4.6 _upsert_body_units() 增量安全 + +per-paper 精确刷新，避免全库 FTS 重建的重复/残留风险： + +```python +def _upsert_body_units(conn, body_units: list[dict]): + paper_ids = list({u["paper_id"] for u in body_units}) + for pid in paper_ids: + conn.execute("DELETE FROM body_units WHERE paper_id = ?", (pid,)) + conn.execute("DELETE FROM body_units_fts WHERE paper_id = ?", (pid,)) + for unit in body_units: + conn.execute("""INSERT INTO body_units (...) VALUES (...)""", ...) + for pid in paper_ids: + conn.execute("""INSERT INTO body_units_fts(...) + SELECT rowid, unit_id, paper_id, section_path, unit_text + FROM body_units + WHERE paper_id = ? AND indexable = 1""", (pid,)) +``` + +--- + +## 5. Object Units 适配嵌套树 + +当前 `build_object_units()` 也是 flat node 遍历。树变 nested 后，父节点 span 包含 child object，同样出现重复归属。 + +**修正：** object ownership 只匹配 own_block_ids 范围。 + +一个 object 只归属一个最具体的 section（own_block_ids 匹配的最深节点）。 + +--- + +## 6. Memory Builder 增量逻辑 + +### 6.1 改造后流程 + +``` +build_from_index(vault): + 1. 计算 canonical_index_hash + 2. 检查 canonical_index_hash 是否匹配 + ├─ 匹配（index 未变）： + │ a. 跳过 papers / aliases / FTS 重建 + │ b. 继续到步骤 3 + │ + └─ 不匹配（index 变了）： + a. 全量重建 + b. 跳过步骤 3 + + 3. Per-paper body_units 增量更新 + for each paper with tree + blocks.structured.jsonl: + 读取 manifest: 的 ocr_result_hash + 计算当前 ocr_result_hash + ├─ 相同 → skip + └─ 不同或缺失 → 重建 body_units + object_units + 更新 manifest +``` + +### 6.2 ocr_result_hash 三级 fallback + +```python +def _resolve_ocr_result_hash(paper_dir: Path) -> str: + # 1. index/result-hash.txt + rp = paper_dir / "index" / "result-hash.txt" + if rp.exists(): + return rp.read_text().strip() + # 2. hash of 结构化产物 + hasher = hashlib.sha256() + for rel in ["structure/blocks.structured.jsonl", + "index/structure-tree.json", + "index/role-index.json"]: + p = paper_dir / rel + if p.exists(): + hasher.update(p.read_bytes()) + result = hasher.hexdigest() + if result != hashlib.sha256(b"").hexdigest(): + return result + # 3. meta.json derived_version + meta_p = paper_dir / "meta.json" + if meta_p.exists(): + meta = json.loads(meta_p.read_text()) + dv = meta.get("derived_version", {}) + return hashlib.sha256(json.dumps(dv, sort_keys=True).encode()).hexdigest() + return "" +``` + +### 6.3 Manifest 与 Vector 共用 body_units_hash + +manifest 和 vector resume 使用**同一个 canonical hash**，命名统一为 `body_units_hash`。 + +```python +def compute_body_units_hash(units: list[dict], policy_version: str) -> str: + raw = json.dumps([{ + "unit_id": u["unit_id"], + "section_path": u["section_path"], + "section_level": u["section_level"], + "section_title": u["section_title"], + "unit_kind": u["unit_kind"], + "part_ordinal": u["part_ordinal"], + "unit_text": u["unit_text"], + "retrieval_policy_version": policy_version, + } for u in units], sort_keys=True) + return hashlib.sha256(raw.encode()).hexdigest() +``` + +**存储位置：** + +```text +meta manifest: → body_units_hash +vector metadata → body_units_hash（每条向量都存） +embed build-state → 用于 resume 比较 +``` + +这样 manifest 判断和 vector resume 判断始终对齐。 + +### 7.2 Embedding Builder 从 DB 读取 + +```python +def get_body_units_for_embedding(vault, zotero_key) -> list[dict]: + db_path = get_memory_db_path(vault) + conn = get_connection(db_path, read_only=True) + rows = conn.execute( + """SELECT unit_id, paper_id, section_path, section_level, section_title, + unit_kind, part_ordinal, unit_text, token_estimate + FROM body_units + WHERE paper_id = ? AND indexable = 1 + ORDER BY unit_id""", + (zotero_key,) + ).fetchall() + conn.close() + return [dict(r) for r in rows] +``` + +SELECT 必须包含所有 `body_units_hash` 计算字段，否则 hash 会用 None 做 silent input。 + +### 7.3 embed_body_units() + +```python +RETRIEVAL_POLICY_VERSION = "l4.body.v1" + +def embed_body_units(vault, zotero_key, body_units): + collection = get_collection(vault, name="paperforge_body") + provider = OpenAICompatibleProvider(vault) + + body_units_hash = compute_body_units_hash(body_units, RETRIEVAL_POLICY_VERSION) + + texts = [u["unit_text"] for u in body_units] + ids = [u["unit_id"] for u in body_units] + metadatas = [{ + "paper_id": zotero_key, + "section_path": u["section_path"], + "section_level": u["section_level"], + "section_title": u["section_title"], + "chunk_index": u["unit_id"], + "token_estimate": u["token_estimate"], + "unit_kind": u["unit_kind"], # "body" or "backmatter_body" + "part_ordinal": u["part_ordinal"], + "retrieval_policy_version": RETRIEVAL_POLICY_VERSION, + "body_units_hash": body_units_hash, + } for u in body_units] + + embeddings = provider.encode(texts) + collection.add(ids=ids, embeddings=embeddings, documents=texts, metadatas=metadatas) + return len(body_units) +``` + +### 7.4 delete_paper_vectors() 双删 + +```python +def delete_paper_vectors(vault, zotero_key): + for name in ["paperforge_fulltext", "paperforge_body"]: + try: + collection = get_collection(vault, name=name) + ids = collection.get(where={"paper_id": zotero_key}).get("ids", []) + if ids: + collection.delete(ids=ids) + except Exception: + pass +``` + +### 7.5 Resume 检测（body_units_hash + policy_version） + +仅比较 `retrieval_policy_version` 不够——文本变化时不会触发重嵌。 + +**判断逻辑：** + +``` +if paperforge_body 中存在该 paper_id + 且 stored.body_units_hash == current.body_units_hash + 且 stored.retrieval_policy_version == current.retrieval_policy_version + → skip + +否则： + delete + re-embed +``` + +`current.body_units_hash` 由 `compute_body_units_hash(body_units, policy_version)` 计算，与 manifest 使用同一函数。 + +### 7.6 embed status 内部分 collection 统计 + +对外总数，对内明细： + +```json +{ + "chunk_count": 61258, + "collections": { + "paperforge_fulltext": 60225, + "paperforge_body": 1033 + } +} +``` +REM +REM + +--- + +## 8. Retrieve 双 Collection 合并 + +### 8.1 merge_retrieve() 签名 + +```python +def merge_retrieve(vault, query, limit=5, expand=True) -> list[dict]: +``` + +### 8.2 合并逻辑 + +1. 查 paperforge_fulltext + paperforge_body +2. unit-level 去重（legacy: collection+id, body: unit_id） +3. per-paper cap：每篇最多 2-3 条（不按 paper_id 一条去重——同一篇论文的不同 section 都可能命中） +4. 标注 source: legacy_chunk / body_unit +5. score 降序，截取 top limit + +--- + +## 9. Content Discovery 行为 + +### 9.1 核心规则 + +- body_units_fts 为 primary 检索源 +- 无结果时不静默 fallback 到 paper_fts（metadata 搜索） +- 返回 coverage 信息（denominator 固定）： + body_units_papers = COUNT(DISTINCT paper_id) FROM body_units WHERE indexable=1 + ocr_done_papers = COUNT(*) FROM papers WHERE ocr_status='done' + library_papers = COUNT(*) FROM papers +- UI 建议显示：`正文索引覆盖 body_units_papers / ocr_done_papers 篇 OCR 完成论文` +- metadata 搜索只作为显式 next_action + +### 9.2 混合迁移 coverage 保护 + +22 篇有 body_units 时 body_units_fts.count > 0，content-discovery 会自动切到 body_units_fts 路径。此时 846 篇旧论文不可见。不做静默 fallback，但 coverage 提示让用户知道搜索范围有限。 + +--- + +## 10. SQLite Schema 调整 + +body_units 表补充字段（CURRENT_SCHEMA_VERSION → 4）： + +```sql +section_path_json TEXT NOT NULL DEFAULT '[]', +section_level INTEGER NOT NULL DEFAULT 0, +section_title TEXT NOT NULL DEFAULT '', +part_ordinal INTEGER NOT NULL DEFAULT 0, +``` + +--- + +## 11. PR / Wave 拆分 + +### PR 1：结构树和 body_units 正确性（最核心） +- core.io.read_jsonl（待加） +- render/heading-map.json sidecar（renderer 同步产出） +- build_structure_tree() — stack 算法，own_block_ids +- build_body_units() — recursive，own_block_ids，token cap +- build_object_units() — recursive，own_block_ids +- schema 字段扩展 + version 4 +- tests: 嵌套路径、parent 不重复、空壳跳过、token cap + +### PR 2：Memory Builder 增量 +- 路径修正（待改） +- 不 early return on canonical_hash +- per-paper ocr_result_hash 增量 +- _resolve_ocr_result_hash 三级 fallback +- _upsert_body_units per-paper FTS 精确刷新 + +### PR 3：Content Discovery 行为 +- body_units_fts primary +- 无结果时不 fallback +- coverage 提示 + +### PR 4：Embedding 双 collection +- get_collection(name=) +- embed_body_units() +- delete_paper_vectors() 双删 +- embed build 从 DB 读 body_units +- resume 检测 body_unit_hash + policy_version +- status 分 collection 统计 + +### PR 5：Retrieve 合并 +- merge_retrieve() +- unit-level dedup + per-paper cap +- retrieve.py 切到 merge_retrieve() +- 输出 source 标注 + +--- + +## 12. 边界情况与风险 + +| 场景 | 处理方式 | 风险 | +|------|---------|------| +| 0 篇有 blocks.structured.jsonl | 全部走旧 chunker | 🟢 无 | +| 部分论文有（22/868） | 分流处理 | 🟡 注意 content-discovery coverage | +| 父节点 own 包含子内容 | 已拆为 own / subtree | 🟢 已处理 | +| body_units > 1000 tokens | token cap 拆分 | 🟢 已处理 | +| References 污染正文检索 | reference_item 不进 body_units | 🟢 已处理 | +| result-hash.txt 不存在 | 三级 fallback hash | 🟡 已处理 | +| resume 命中旧策略版本或文本变了 | vector metadata 存 body_unit_hash + policy_version | 🟢 已处理 | +| content-discovery 只搜 22 篇 | coverage 提示 | 🟡 用户可见 | diff --git a/paperforge/core/io.py b/paperforge/core/io.py index 389ac90f..190d9d99 100644 --- a/paperforge/core/io.py +++ b/paperforge/core/io.py @@ -20,3 +20,16 @@ def write_json_atomic(path: Path, data) -> None: def write_json(path: Path, data) -> None: """Write data as JSON, creating parent directories as needed.""" write_json_atomic(path, data) + + +def read_jsonl(path: Path) -> list[dict]: + """Read a newline-delimited JSON file (one JSON object per line).""" + if not path.exists(): + return [] + rows: list[dict] = [] + with open(path, encoding="utf-8") as f: + for line in f: + line = line.strip() + if line: + rows.append(json.loads(line)) + return rows diff --git a/paperforge/memory/schema.py b/paperforge/memory/schema.py index 50324d90..4d6e3e7f 100644 --- a/paperforge/memory/schema.py +++ b/paperforge/memory/schema.py @@ -5,7 +5,7 @@ import sqlite3 logger = logging.getLogger(__name__) -CURRENT_SCHEMA_VERSION = 3 # Bump from 2 for unit_kind + object_units schema changes +CURRENT_SCHEMA_VERSION = 4 # Bump from 3 for section_path_json, section_level, section_title, part_ordinal CREATE_META = """ CREATE TABLE IF NOT EXISTS meta ( @@ -180,8 +180,12 @@ CREATE TABLE IF NOT EXISTS body_units ( unit_id TEXT PRIMARY KEY, paper_id TEXT NOT NULL, section_path TEXT NOT NULL, + section_path_json TEXT NOT NULL DEFAULT '[]', + section_level INTEGER NOT NULL DEFAULT 0, + section_title TEXT NOT NULL DEFAULT '', unit_text TEXT NOT NULL, unit_kind TEXT NOT NULL DEFAULT 'body', + part_ordinal INTEGER NOT NULL DEFAULT 0, page_span_json TEXT NOT NULL, block_span_json TEXT NOT NULL, token_estimate INTEGER NOT NULL, @@ -241,6 +245,18 @@ def ensure_schema(conn: sqlite3.Connection) -> None: logger.info("Migrating schema v%s -> v3: rebuilding body_units, body_units_fts, object_units", current_version) for table in ("body_units", "body_units_fts", "object_units"): conn.execute(f"DROP TABLE IF EXISTS {table};") + if current_version < 4: + logger.info("Migrating schema v%s -> v4: adding body_units columns", current_version) + for col_sql in [ + "ALTER TABLE body_units ADD COLUMN section_path_json TEXT NOT NULL DEFAULT '[]'", + "ALTER TABLE body_units ADD COLUMN section_level INTEGER NOT NULL DEFAULT 0", + "ALTER TABLE body_units ADD COLUMN section_title TEXT NOT NULL DEFAULT ''", + "ALTER TABLE body_units ADD COLUMN part_ordinal INTEGER NOT NULL DEFAULT 0", + ]: + try: + conn.execute(col_sql) + except Exception: + pass # column may already exist conn.execute(CREATE_BODY_UNITS) conn.execute(CREATE_BODY_UNITS_FTS) diff --git a/paperforge/retrieval/structure_tree.py b/paperforge/retrieval/structure_tree.py index c0e00d17..485d38a4 100644 --- a/paperforge/retrieval/structure_tree.py +++ b/paperforge/retrieval/structure_tree.py @@ -6,38 +6,109 @@ from typing import Any from paperforge.core.io import write_json -def build_structure_tree(structured_blocks: list[dict[str, Any]]) -> dict[str, Any]: - nodes: list[dict[str, Any]] = [] - current_section: dict[str, Any] | None = None - for block in structured_blocks: - role = block.get("role") - text = str(block.get("text", "")).strip() - if role in {"section_heading", "subsection_heading", "introduction_heading", "abstract_heading"} and text: - parent_title = nodes[-1]["title"] if nodes else "" - current_section = { - "node_id": f"sec:{block.get('block_id')}", - "kind": "section", - "title": text, - "level": 1 if role != "subsection_heading" else 2, - "section_path": [text] if role != "subsection_heading" else [parent_title, text], - "page_span": [block.get("page", 0), block.get("page", 0)], - "block_span": [[block.get("page", 0), block.get("block_id", "")]], - "children": [], - "objects": [], - } - nodes.append(current_section) - elif current_section is not None: - current_section["page_span"][1] = block.get("page", current_section["page_span"][1]) - current_section["block_span"].append([block.get("page", 0), block.get("block_id", "")]) - return {"paper_id": structured_blocks[0].get("paper_id", "") if structured_blocks else "", "nodes": nodes} +def build_structure_tree( + heading_events: list[dict[str, Any]], + emitted_block_events: list[dict[str, Any]], + structured_blocks: list[dict[str, Any]], +) -> dict[str, Any]: + """Build nested structure tree from heading events and emitted block events. + + Uses PaperIndex stack algorithm for nesting, then assigns own_block_ids + and subtree_block_ids via interval computation based on emitted_order. + """ + if not heading_events: + paper_id = structured_blocks[0].get("paper_id", "") if structured_blocks else "" + return {"paper_id": paper_id, "nodes": []} + + heading_events = sorted(heading_events, key=lambda h: h["emitted_order"]) + + # ── stack algorithm ── + root_nodes: list[dict[str, Any]] = [] + stack: list[dict[str, Any]] = [] + + for h in heading_events: + node = { + "node_id": f"sec:{h['block_id']}", + "kind": "section", + "title": h["title"], + "level": h["markdown_level"], + "block_id": h["block_id"], + "page_span": [h["page"], h["page"]], + "own_block_ids": [], + "subtree_block_ids": [], + "children": [], + "objects": [], + } + while stack and stack[-1]["level"] >= h["markdown_level"]: + stack.pop() + if stack: + stack[-1]["children"].append(node) + else: + root_nodes.append(node) + stack.append(node) + + # ── build bounds map ── + bounds_map: dict[str | int, dict[str, int | float]] = {} + for i, h in enumerate(heading_events): + next_sibling: int | float | None = None + for h2 in heading_events[i + 1 :]: + if h2["markdown_level"] <= h["markdown_level"]: + next_sibling = h2["emitted_order"] + break + bounds_map[h["block_id"]] = { + "start": h["emitted_order"], + "end": next_sibling or float("inf"), + } + + # ── interval assignment (recursive) ── + def _assign_intervals(node: dict[str, Any]) -> None: + bounds = bounds_map.get(node["block_id"]) + if bounds is None: + return + + start = bounds["start"] + end = bounds["end"] + + node["subtree_block_ids"] = [ + str(e["block_id"]) for e in emitted_block_events if start <= e["emitted_order"] < end + ] + + # Extend page_span from emitted blocks + for e in emitted_block_events: + if start <= e["emitted_order"] < end: + page = e.get("page") + if page is not None: + if page < node["page_span"][0]: + node["page_span"][0] = page + if page > node["page_span"][1]: + node["page_span"][1] = page + + for child in node.get("children", []): + _assign_intervals(child) + + # own = subtree - all children subtrees - heading block + child_ids: set[str] = set() + for child in node.get("children", []): + child_ids.update(child.get("subtree_block_ids", [])) + child_ids.add(str(node["block_id"])) + + node["own_block_ids"] = [bid for bid in node["subtree_block_ids"] if bid not in child_ids] + + for n in root_nodes: + _assign_intervals(n) + + paper_id = structured_blocks[0].get("paper_id", "") if structured_blocks else "" + return {"paper_id": paper_id, "nodes": root_nodes} def write_structure_tree(index_root: Path, tree: dict[str, Any]) -> None: + """Write structure tree to index/structure-tree.json.""" index_root.mkdir(parents=True, exist_ok=True) write_json(index_root / "structure-tree.json", tree) def summarize_role_index(role_index: dict[str, Any]) -> dict[str, Any]: + """Summarize role index by counting entries per role.""" role_counts: dict[str, int] = {} for key, entries in role_index.items(): if isinstance(entries, list): diff --git a/paperforge/retrieval/units.py b/paperforge/retrieval/units.py index 083fcc3d..e1fb26e2 100644 --- a/paperforge/retrieval/units.py +++ b/paperforge/retrieval/units.py @@ -1,15 +1,58 @@ """BodyUnit / ObjectUnit builders for the retrieval substrate. Produces atomic retrieval units from a paper's structure tree, structured -blocks, and role index. Each unit is a self-contained chunk that can be +blocks, and role index. Each unit is a self-contained chunk that can be indexed by FTS and retrieved independently. """ from __future__ import annotations +import json +from itertools import groupby from typing import Any +# ── Role helper ── + +def _body_unit_role_kind(role: str) -> str | None: + """Map a block's role to body unit kind, or None if excluded.""" + if role == "body_paragraph": + return "body" + if role in {"structured_insert", "non_body_insert", "backmatter_body"}: + return "backmatter_body" + return None # reference_item, reference_heading, heading, caption, asset, etc. + + +def _split_if_oversized(text: str, max_tokens: int = 1000) -> list[str]: + """Split text by paragraph if token estimate exceeds max_tokens.""" + if len(text) // 4 <= max_tokens: + return [text] + paragraphs = text.split("\n\n") + parts: list[str] = [] + current: list[str] = [] + current_tokens = 0 + for para in paragraphs: + para_tokens = len(para) // 4 + if current_tokens + para_tokens > max_tokens and current: + parts.append("\n\n".join(current)) + current = [para] + current_tokens = para_tokens + else: + current.append(para) + current_tokens += para_tokens + if current: + parts.append("\n\n".join(current)) + # Fallback: if a single oversized paragraph wasn't split, halve it + if len(parts) == 1 and len(parts[0]) // 4 > max_tokens: + mid = len(parts[0]) // 2 + break_at = parts[0].rfind(". ", 0, mid) + if break_at < mid // 2: + break_at = parts[0].rfind(" ", 0, mid) + if break_at > 0: + return [parts[0][:break_at + 1].strip(), parts[0][break_at + 1:].strip()] + return parts if parts else [text] + + def build_unit_id( paper_id: str, kind: str, @@ -23,49 +66,108 @@ def build_unit_id( return f"{paper_id}:{kind}:{node_id}:{start_page}-{start_block}:{end_page}-{end_block}" +def build_unit_id_v2( + paper_id: str, + kind: str, + node_id: str, + part_suffix: str = "", +) -> str: + """Simplified unit identifier without block span.""" + return f"{paper_id}:{kind}:{node_id}{part_suffix}" + + def build_body_units( *, tree: dict[str, Any], structured_blocks: list[dict[str, Any]], ) -> list[dict[str, Any]]: - """Build body-text retrieval units for every section node in the tree. + """Build body-text retrieval units recursively from the structure tree. - Each body unit groups all ``body_paragraph`` blocks inside the section's - ``block_span``, joined by double newlines. + Walks each node in the tree, collects own_block_ids that match body + roles, groups by unit_kind (body / backmatter_body), splits oversized + units, and produces schema v4 compatible output. """ units: list[dict[str, Any]] = [] paper_id = tree.get("paper_id", "") - for node in tree.get("nodes", []): - block_ids = {block_id for _, block_id in node.get("block_span", [])} - texts = [ - b.get("text", "") - for b in structured_blocks - if b.get("block_id") in block_ids and b.get("role") == "body_paragraph" - ] - unit_text = "\n\n".join(t for t in texts if t) - span = node.get("block_span", []) - unit = { - "unit_id": build_unit_id( - paper_id, - "body", - node["node_id"], - node["page_span"][0], - span[0][1] if span else "", - node["page_span"][1], - span[-1][1] if span else "", - ), - "paper_id": paper_id, - "section_path": "/".join(node.get("section_path", [])), - "page_span": node.get("page_span", []), - "block_span": span, - "unit_text": unit_text, - "token_estimate": len(unit_text) // 4, - "unit_kind": "body", - "indexable": bool(unit_text.strip()), - "veto_reason": "" if unit_text.strip() else "empty", - "quality_hints": [], - } - units.append(unit) + if not paper_id or not tree.get("nodes"): + return units + + block_map: dict[str | int, dict[str, Any]] = {} + for b in structured_blocks: + bid = b.get("block_id") + if bid is not None: + block_map[str(bid)] = b + block_map[bid] = b + + def walk(node: dict[str, Any], inherited_path: list[str]) -> None: + this_path = inherited_path + [node["title"]] + + # Collect own blocks that match body/backmatter roles + own_body_blocks: list[dict[str, Any]] = [] + for bid in node.get("own_block_ids", []): + b = block_map.get(bid) or block_map.get(str(bid)) + if b is None: + continue + kind = _body_unit_role_kind(b.get("role", "")) + if kind is not None: + own_body_blocks.append(b) + + if own_body_blocks: + # Group by unit_kind + groups: list[tuple[str, list[dict[str, Any]]]] = [] + sorted_blocks = sorted( + own_body_blocks, + key=lambda b: _body_unit_role_kind(b.get("role", "")), + ) + for kind, grp in groupby( + sorted_blocks, + key=lambda b: _body_unit_role_kind(b.get("role", "")), + ): + if kind: + groups.append((kind, list(grp))) + + for unit_kind, blocks in groups: + all_text = "\n\n".join( + b.get("text", "") for b in blocks if b.get("text") + ) + if not all_text.strip(): + continue + + parts = _split_if_oversized(all_text, max_tokens=1000) + n_parts = len(parts) + + for p_idx, part_text in enumerate(parts): + part_ordinal = p_idx + 1 if n_parts > 1 else 0 + suffix = f":part_{part_ordinal:03d}" if part_ordinal else "" + uid = build_unit_id_v2( + paper_id, "body", node["node_id"], suffix, + ) + + unit = { + "unit_id": uid, + "paper_id": paper_id, + "section_path": "/".join(this_path), + "section_path_json": json.dumps(this_path), + "section_level": node.get("level", 0), + "section_title": node.get("title", ""), + "page_span": node.get("page_span", []), + "block_span": [], + "unit_text": part_text, + "token_estimate": len(part_text) // 4, + "unit_kind": unit_kind, + "part_ordinal": part_ordinal, + "indexable": True, + "veto_reason": "", + "quality_hints": [], + } + units.append(unit) + + for child in node.get("children", []): + walk(child, this_path) + + for root in tree.get("nodes", []): + walk(root, []) + return units @@ -77,69 +179,109 @@ def build_object_units( ) -> list[dict[str, Any]]: """Build object retrieval units from the role index scoped by section. - For each section node, objects (figures, tables, equations, …) whose - ``block_id`` falls within the node's ``block_span`` are collected from - the role index and emitted as individual object units with structured - metadata fields (``object_kind``, ``object_label``, ``caption_text``, - ``nearby_body_text``). + Each object unit captures a figure/table/object with its caption and + nearby body text, scoped to the most specific section that owns it. """ import re units: list[dict[str, Any]] = [] paper_id = tree.get("paper_id", "") - for node in tree.get("nodes", []): - block_ids = {block_id for _, block_id in node.get("block_span", [])} - span = node.get("block_span", []) - # Pre-collect body text from this section for nearby_body_text - body_texts = [ - b.get("text", "") - for b in structured_blocks - if b.get("block_id") in block_ids and b.get("role") == "body_paragraph" - ] - nearby_body = "\n\n".join(t for t in body_texts if t) + # Build block_map for text resolution + block_map: dict[str | int, dict[str, Any]] = {} + for b in structured_blocks: + bid = b.get("block_id") + if bid is not None: + block_map[str(bid)] = b + block_map[bid] = b - for role, entries in role_index.items(): - for entry in entries: - if not isinstance(entry, dict): - continue - if entry.get("block_id") not in block_ids: - continue - text = str(entry.get("text", "")) - uid = build_unit_id( - paper_id, "object", node["node_id"], - node["page_span"][0], entry.get("block_id", ""), - node["page_span"][1], entry.get("block_id", ""), - ) + # Collect all objects from role_index (figures, tables, etc.) + objects = role_index.get("figure_captions", []) + role_index.get("table_captions", []) - # Map role to object_kind - role_lower = role.lower() - if role_lower.startswith("figure") or "figure" in role_lower: - object_kind = "figure" - elif role_lower.startswith("table") or "table" in role_lower: - object_kind = "table" - else: - object_kind = role + def find_owning_node(node: dict[str, Any], object_block_id: str) -> dict[str, Any] | None: + """Find the most specific section that owns an object block.""" + # Check children first (most specific) + for child in node.get("children", []): + result = find_owning_node(child, object_block_id) + if result: + return result + # If object is in this node's subtree but not in any child's subtree + str_bid = str(object_block_id) + if str_bid in node.get("subtree_block_ids", []): + return node + return None - # Extract object_label from caption prefix - label_match = re.match(r'(Figure|Fig\.?\s*|Table|Tab\.?\s*)\s*\d+', text, re.IGNORECASE) - object_label = label_match.group(0) if label_match else entry.get("block_id", "") + for obj in objects: + obj_id = obj.get("figure_id") or obj.get("table_id") or "" + obj_type = "figure" if "figure_id" in obj else "table" + caption = str(obj.get("text", "") or "") + owning_node: dict[str, Any] | None = None + + # Try to find a caption block or asset block in this object + caption_bid = obj.get("caption_block_id") + if caption_bid: + for root in tree.get("nodes", []): + owning_node = find_owning_node(root, str(caption_bid)) + if owning_node: + break + + if not owning_node: + # Fallback: iterate all roots to find the first that contains object block_ids + consumed_ids = [ + str(x) if not isinstance(x, dict) else str(x.get("block_id", "")) + for x in obj.get("consumed_block_ids", []) + ] + for cid in consumed_ids: + if cid: + for root in tree.get("nodes", []): + owning_node = find_owning_node(root, cid) + if owning_node: + break + if owning_node: + break + + if not owning_node: + owning_node = tree.get("nodes", [{}])[0] if tree.get("nodes") else None + + if not owning_node: + continue + + # Build section path + section_path_parts: list[str] = [] + + def _build_path(n: dict[str, Any], target: dict[str, Any]) -> list[str]: + if n is target: + return [n.get("title", "")] + for child in n.get("children", []): + result = _build_path(child, target) + if result: + return [n.get("title", "")] + result + return [] + + for root in tree.get("nodes", []): + section_path_parts = _build_path(root, owning_node) + if section_path_parts: + break + + uid = f"{paper_id}:obj:{obj_id}" + unit = { + "unit_id": uid, + "paper_id": paper_id, + "section_path": "/".join(section_path_parts), + "object_kind": obj_type, + "object_label": obj_id, + "caption_text": caption, + "nearby_body_text": "", # Can be populated from own_block_ids text + "page_span": [ + owning_node.get("page_span", [0, 0])[0], + owning_node.get("page_span", [0, 0])[1], + ], + "block_span": [], + "token_estimate": len(caption) // 4, + "indexable": bool(caption.strip()), + "veto_reason": "" if caption.strip() else "empty_caption", + "quality_hints": [], + } + units.append(unit) - unit = { - "unit_id": uid, - "paper_id": paper_id, - "section_path": "/".join(node.get("section_path", [])), - "page_span": node.get("page_span", []), - "block_span": span, - "object_kind": object_kind, - "object_label": object_label, - "caption_text": text, - "nearby_body_text": nearby_body, - "token_estimate": len(text) // 4, - "unit_kind": "object", - "indexable": bool(text.strip()), - "veto_reason": "" if text.strip() else "empty", - "quality_hints": [], - } - units.append(unit) return units diff --git a/paperforge/worker/ocr.py b/paperforge/worker/ocr.py index fe8c324d..ec3252c7 100644 --- a/paperforge/worker/ocr.py +++ b/paperforge/worker/ocr.py @@ -1977,9 +1977,9 @@ def postprocess_ocr_result(vault: Path, key: str, all_results: list[dict]) -> tu ) # --- Phase 3: structured renderer --- - from paperforge.worker.ocr_render import render_fulltext_markdown, write_render_outputs + from paperforge.worker.ocr_render import render_fulltext_markdown, write_render_outputs, RenderOutput - markdown = render_fulltext_markdown( + rendered = render_fulltext_markdown( structured_blocks=structured, resolved_metadata=resolved, figure_inventory=figure_inventory, @@ -1987,10 +1987,11 @@ def postprocess_ocr_result(vault: Path, key: str, all_results: list[dict]) -> tu page_count=page_num, document_structure=doc_structure, reader_payload=reader_payload, + return_events=True, ) - # --- Phase 3: OCR health report --- + # --- health report --- from paperforge.worker.ocr_health import ( build_ocr_health, build_ocr_raw_integrity_health, @@ -2006,7 +2007,7 @@ def postprocess_ocr_result(vault: Path, key: str, all_results: list[dict]) -> tu table_inventory=table_inventory, doc_structure=doc_structure, reader_payload=reader_payload, - rendered_markdown=markdown, + rendered_markdown=rendered.markdown, ) health_report["ocr_raw_integrity"] = ocr_raw_integrity write_ocr_health(ocr_root / "health", health_report) @@ -2016,6 +2017,16 @@ def postprocess_ocr_result(vault: Path, key: str, all_results: list[dict]) -> tu from paperforge.worker.ocr_decisions import collect_decisions, write_decision_log write_decision_log(ocr_root / "health" / "decision_log.jsonl", collect_decisions(structured)) + # --- Phase 4: render output --- + meta = write_render_outputs( + render_root=ocr_root / "render", + user_fulltext=ocr_root / "fulltext.md", + markdown=rendered.markdown, + heading_events=rendered.heading_events, + emitted_block_events=rendered.emitted_block_events, + meta=meta, + rebuild_increment=False, + ) # --- Phase 5: role-based OCR index --- from paperforge.worker.ocr_index import build_role_indexes, write_role_index @@ -2028,7 +2039,11 @@ def postprocess_ocr_result(vault: Path, key: str, all_results: list[dict]) -> tu # --- Phase 6: structure tree --- from paperforge.retrieval.structure_tree import build_structure_tree, write_structure_tree - structure_tree = build_structure_tree(structured) + structure_tree = build_structure_tree( + heading_events=rendered.heading_events, + emitted_block_events=rendered.emitted_block_events, + structured_blocks=structured, + ) write_structure_tree(ocr_root / "index", structure_tree) # Update meta.json with version payloads @@ -2056,13 +2071,6 @@ def postprocess_ocr_result(vault: Path, key: str, all_results: list[dict]) -> tu except ValueError: meta["legacy_images_path"] = str(images_dir) meta["path_map"] = {"structured_truth": "assets/", "legacy_compat": "images/"} - meta = write_render_outputs( - render_root=ocr_root / "render", - user_fulltext=ocr_root / "fulltext.md", - markdown=markdown, - meta=meta, - rebuild_increment=False, - ) write_json(meta_path, meta) fulltext_path = ocr_root / "fulltext.md" diff --git a/paperforge/worker/ocr_rebuild.py b/paperforge/worker/ocr_rebuild.py index 1877594b..f9720bea 100644 --- a/paperforge/worker/ocr_rebuild.py +++ b/paperforge/worker/ocr_rebuild.py @@ -439,9 +439,9 @@ def _rebuild_one_paper(vault: Path, key: str) -> dict: doc_structure: dict, ocr_meta: dict, source_pdf_path: Path | None, - ) -> str: + ) -> tuple["RenderOutput", str]: """Extract object artifacts, render fulltext markdown, build health report. - Returns markdown string.""" + Returns (RenderOutput, health_overall) tuple.""" from paperforge.worker.ocr_objects import ( extract_and_write_objects, _resolve_object_crop_pdf_path, @@ -468,13 +468,13 @@ def _rebuild_one_paper(vault: Path, key: str) -> dict: use_disk_page_cache=False, ) - from paperforge.worker.ocr_render import render_fulltext_markdown, write_render_outputs + from paperforge.worker.ocr_render import RenderOutput, render_fulltext_markdown, write_render_outputs rebuild_page_count = ocr_meta.get("page_count", 0) or 0 if not rebuild_page_count: all_rebuild_pages = {int(b["page"]) for b in structured if b.get("page")} rebuild_page_count = max(all_rebuild_pages) if all_rebuild_pages else 0 - markdown = render_fulltext_markdown( + rendered = render_fulltext_markdown( structured_blocks=structured, resolved_metadata=resolved, figure_inventory=figure_inventory, @@ -482,8 +482,8 @@ def _rebuild_one_paper(vault: Path, key: str) -> dict: page_count=rebuild_page_count, document_structure=doc_structure, reader_payload=reader_payload, + return_events=True, ) - from paperforge.worker.ocr_health import build_ocr_health, build_ocr_raw_integrity_health, write_ocr_health health_report = build_ocr_health( @@ -494,7 +494,7 @@ def _rebuild_one_paper(vault: Path, key: str) -> dict: table_inventory=table_inventory, doc_structure=doc_structure, reader_payload=reader_payload, - rendered_markdown=markdown, + rendered_markdown=rendered.markdown, ) health_report["ocr_raw_integrity"] = build_ocr_raw_integrity_health(all_raw_blocks) write_ocr_health(paper_root / "health", health_report) @@ -503,13 +503,13 @@ def _rebuild_one_paper(vault: Path, key: str) -> dict: write_decision_log(paper_root / "health" / "decision_log.jsonl", collect_decisions(structured)) - return markdown, health_report.get("overall", "unknown") + return rendered, health_report.get("overall", "unknown") # ── Phase 5: indexes, version flags, write meta ── def _phase5_finalize( resolved: dict, structured: list[dict], - markdown: str, + rendered: "RenderOutput", span_meta_patch: dict, health_overall: str = "unknown", ) -> None: @@ -524,7 +524,11 @@ def _rebuild_one_paper(vault: Path, key: str) -> dict: from paperforge.retrieval.structure_tree import build_structure_tree, write_structure_tree - structure_tree = build_structure_tree(structured) + structure_tree = build_structure_tree( + heading_events=rendered.heading_events, + emitted_block_events=rendered.emitted_block_events, + structured_blocks=structured, + ) write_structure_tree(paper_root / "index", structure_tree) meta = ocr_meta @@ -537,7 +541,9 @@ def _rebuild_one_paper(vault: Path, key: str) -> dict: meta = write_render_outputs( render_root=paper_root / "render", user_fulltext=artifacts.compat_fulltext, - markdown=markdown, + markdown=rendered.markdown, + heading_events=rendered.heading_events, + emitted_block_events=rendered.emitted_block_events, meta=meta, rebuild_increment=True, ) @@ -565,12 +571,12 @@ def _rebuild_one_paper(vault: Path, key: str) -> dict: table_inventory = phase3_result["table_inventory"] reader_payload = phase3_result["reader_payload"] - markdown, health_overall = _phase4_render_health( + rendered, health_overall = _phase4_render_health( structured, resolved, figure_inventory, table_inventory, reader_payload, doc_structure, ocr_meta, source_pdf_path, ) - _phase5_finalize(resolved, structured, markdown, span_meta_patch, health_overall=health_overall) + _phase5_finalize(resolved, structured, rendered, span_meta_patch, health_overall=health_overall) return {"key": key, "status": "ok"} def _run_parallel_rebuild(vault: Path, keys: list[str], workers: int, checkpoint_dir: Path | None) -> list[dict]: """Run rebuild in parallel using a process pool.""" diff --git a/paperforge/worker/ocr_render.py b/paperforge/worker/ocr_render.py index c0018a64..138833ae 100644 --- a/paperforge/worker/ocr_render.py +++ b/paperforge/worker/ocr_render.py @@ -3,6 +3,8 @@ from __future__ import annotations import re from html import unescape from pathlib import Path +from dataclasses import dataclass +from typing import Any from paperforge.worker.ocr_document import ( _TAIL_ROLES, @@ -1247,6 +1249,26 @@ def _collect_frontmatter_fallback_fields( if fallback_doi: result["doi"] = fallback_doi return result +@dataclass +class RenderOutput: + markdown: str + heading_events: list[dict] + emitted_block_events: list[dict] + + +def _determine_emitted_as(role: str) -> str: + emitted_map = { + "body_paragraph": "body", + "backmatter_body": "backmatter", + "tail_candidate_body": "body", + "structured_insert": "structured_insert", + "non_body_insert": "structured_insert", + "backmatter_boundary_heading": "backmatter_heading", + "backmatter_heading": "backmatter_heading", + } + return emitted_map.get(role, "other") + + def render_fulltext_markdown( @@ -1258,8 +1280,12 @@ def render_fulltext_markdown( page_count: int | None = None, document_structure: DocumentStructure | None = None, reader_payload: dict | None = None, -) -> str: + return_events: bool = False, +) -> str | RenderOutput: lines: list[str] = [] + heading_events: list[dict] = [] + emitted_block_events: list[dict] = [] + emitted_order_counter: int = 0 emitted_figure_captions: set[str] = set() @@ -1795,9 +1821,21 @@ def render_fulltext_markdown( if role == "backmatter_boundary_heading" or role == "backmatter_heading": last_structured_insert_page = None last_structured_insert_bbox = None + bm_start_line = len(lines) if text and not _should_suppress_frontmatter_heading(text): lines.append(f"**{text}**") lines.append("") + if len(lines) > bm_start_line: + emitted_block_events.append({ + "emitted_order": emitted_order_counter, + "line_start": bm_start_line, + "line_end": len(lines), + "page": block_page, + "block_id": block.get("block_id"), + "role": role, + "emitted_as": _determine_emitted_as(role), + }) + emitted_order_counter += 1 elif role == "reference_heading": last_structured_insert_page = None last_structured_insert_bbox = None @@ -1842,7 +1880,18 @@ def render_fulltext_markdown( prefix = "##" if depth <= 1 else "###" lines.append(f"{prefix} {text}") lines.append("") + h_level = len(prefix) # "##" → 2, "###" → 3, "####" → 4 + heading_events.append({ + "line_number": len(lines) - 1, + "markdown_level": h_level, + "title": text, + "page": block.get("page"), + "block_id": block.get("block_id"), + "emitted_order": emitted_order_counter, + }) + emitted_order_counter += 1 elif role == "structured_insert": + si_start_line = len(lines) container_text = block.get("_container_text") if container_text and isinstance(container_text, str): source_text = normalize_ocr_math_text(" ".join(container_text.replace("\n", " ").split())) @@ -1886,6 +1935,17 @@ def render_fulltext_markdown( lines.append("") last_structured_insert_page = block_page last_structured_insert_bbox = bbox if len(bbox) >= 4 else None + if len(lines) > si_start_line: + emitted_block_events.append({ + "emitted_order": emitted_order_counter, + "line_start": si_start_line, + "line_end": len(lines), + "page": block.get("page"), + "block_id": block.get("block_id"), + "role": role, + "emitted_as": _determine_emitted_as(role), + }) + emitted_order_counter += 1 elif role == "table_caption": tbl_ids_for_page = tables_by_page.get(block_page, []) if tbl_ids_for_page: @@ -1900,6 +1960,7 @@ def render_fulltext_markdown( lines.append(text) lines.append("") elif role in ("backmatter_body", "tail_candidate_body", "body_paragraph"): + body_start_line = len(lines) if last_structured_insert_page is not None: lines.append("") last_structured_insert_page = None @@ -1915,7 +1976,19 @@ def render_fulltext_markdown( _emitted_body_text_by_page.setdefault(block_page if block_page is not None else -1, []).append(text) lines.append(text) lines.append("") + if text and len(lines) > body_start_line: + emitted_block_events.append({ + "emitted_order": emitted_order_counter, + "line_start": body_start_line, + "line_end": len(lines), + "page": block.get("page"), + "block_id": block.get("block_id"), + "role": role, + "emitted_as": _determine_emitted_as(role), + }) + emitted_order_counter += 1 else: + else_start_line = len(lines) if last_structured_insert_page is not None: lines.append("") last_structured_insert_page = None @@ -1923,6 +1996,17 @@ def render_fulltext_markdown( if text: lines.append(text) lines.append("") + if text and len(lines) > else_start_line: + emitted_block_events.append({ + "emitted_order": emitted_order_counter, + "line_start": else_start_line, + "line_end": len(lines), + "page": block.get("page"), + "block_id": block.get("block_id"), + "role": role, + "emitted_as": _determine_emitted_as(role), + }) + emitted_order_counter += 1 # Emit objects for the last rendered page if current_page is not None: @@ -1959,7 +2043,14 @@ def render_fulltext_markdown( emitted_figure_captions=emitted_figure_captions, ) - return "\n".join(lines).strip() + "\n" + markdown = "\n".join(lines).strip() + "\n" + if return_events: + return RenderOutput( + markdown=markdown, + heading_events=heading_events, + emitted_block_events=emitted_block_events, + ) + return markdown import datetime as _dt @@ -1979,10 +2070,16 @@ def write_render_outputs( meta: dict, rebuild_increment: bool, now_utc: _dt.datetime | None = None, + heading_events: list[dict] | None = None, + emitted_block_events: list[dict] | None = None, ) -> dict: now_utc = now_utc or _dt.datetime.now(_dt.timezone.utc) render_root.mkdir(parents=True, exist_ok=True) (render_root / "fulltext.md").write_text(markdown, encoding="utf-8") + if heading_events is not None and emitted_block_events is not None: + from paperforge.core.io import write_json + render_map = {"headings": heading_events, "emitted_blocks": emitted_block_events} + write_json(render_root / "render-map.json", render_map) backup_info = create_pre_rebuild_backup(user_fulltext, now_utc) atomic_replace_text(user_fulltext, markdown) diff --git a/tests/test_layer4_structure_tree.py b/tests/test_layer4_structure_tree.py index 85222d95..9b9b13c7 100644 --- a/tests/test_layer4_structure_tree.py +++ b/tests/test_layer4_structure_tree.py @@ -1,10 +1,7 @@ -"""Tests for Layer 4 structure tree builder and paper navigation.""" +"""Tests for Layer 4 structure tree builder using heading_events + emitted_block_events.""" from __future__ import annotations -from pathlib import Path -from unittest.mock import patch - from paperforge.retrieval.structure_tree import ( build_structure_tree, summarize_role_index, @@ -12,71 +9,90 @@ from paperforge.retrieval.structure_tree import ( ) -def test_build_structure_tree_creates_section_nodes_from_headings(): - structured_blocks = [ - {"paper_id": "ABCD1234", "page": 1, "block_id": "b1", "role": "section_heading", "text": "Methods"}, - {"paper_id": "ABCD1234", "page": 1, "block_id": "b2", "role": "body_paragraph", "text": "We recruited 30 patients."}, - ] - tree = build_structure_tree(structured_blocks) +def _h(emitted_order, markdown_level, title, page=1, block_id=None): + return { + "line_number": emitted_order * 10, + "markdown_level": markdown_level, + "title": title, + "page": page, + "block_id": block_id or f"h_{emitted_order}", + "emitted_order": emitted_order, + } + + +def _e(emitted_order, block_id, role="body_paragraph", page=1, emitted_as="body"): + return { + "emitted_order": emitted_order, + "line_start": emitted_order * 10, + "line_end": emitted_order * 10 + 5, + "page": page, + "block_id": block_id, + "role": role, + "emitted_as": emitted_as, + } + + +def test_single_heading_with_body(): + heading_events = [_h(0, 2, "Methods", block_id="b1")] + emitted = [_e(0, "b1"), _e(1, "b2")] + structured = [{"paper_id": "ABCD1234"}] + tree = build_structure_tree(heading_events, emitted, structured) assert tree["paper_id"] == "ABCD1234" assert len(tree["nodes"]) == 1 - assert tree["nodes"][0]["title"] == "Methods" - assert tree["nodes"][0]["section_path"] == ["Methods"] - assert tree["nodes"][0]["level"] == 1 - assert tree["nodes"][0]["page_span"] == [1, 1] + n = tree["nodes"][0] + assert n["title"] == "Methods" + assert n["level"] == 2 + assert n["block_id"] == "b1" + assert set(n["subtree_block_ids"]) == {"b1", "b2"} + assert n["own_block_ids"] == ["b2"] -def test_build_structure_tree_handles_subsection(): - structured_blocks = [ - {"paper_id": "ABCD1234", "page": 1, "block_id": "b1", "role": "section_heading", "text": "Methods"}, - {"paper_id": "ABCD1234", "page": 1, "block_id": "b2", "role": "body_paragraph", "text": "We recruited 30 patients."}, - {"paper_id": "ABCD1234", "page": 2, "block_id": "b3", "role": "subsection_heading", "text": "Statistical Analysis"}, - {"paper_id": "ABCD1234", "page": 2, "block_id": "b4", "role": "body_paragraph", "text": "We used t-tests."}, +def test_h2_h3_nesting(): + heading_events = [ + _h(0, 2, "Methods", block_id="b1"), + _h(2, 3, "Statistics", block_id="b3"), ] - tree = build_structure_tree(structured_blocks) - assert tree["paper_id"] == "ABCD1234" + emitted = [_e(0, "b1"), _e(1, "b2"), _e(2, "b3"), _e(3, "b4")] + structured = [{"paper_id": "P001"}] + tree = build_structure_tree(heading_events, emitted, structured) + assert len(tree["nodes"]) == 1 + parent = tree["nodes"][0] + child = parent["children"][0] + assert parent["title"] == "Methods" + assert child["title"] == "Statistics" + assert parent["own_block_ids"] == ["b2"] + assert child["own_block_ids"] == ["b4"] + + +def test_h2_h3_h2_sibling(): + heading_events = [ + _h(0, 2, "Methods", block_id="b1"), + _h(2, 3, "Statistics", block_id="b3"), + _h(4, 2, "Results", block_id="b5"), + ] + emitted = [ + _e(0, "b1"), _e(1, "b2"), + _e(2, "b3"), _e(3, "b4"), + _e(4, "b5"), _e(5, "b6"), + ] + structured = [{"paper_id": "P002"}] + tree = build_structure_tree(heading_events, emitted, structured) assert len(tree["nodes"]) == 2 - assert tree["nodes"][1]["title"] == "Statistical Analysis" - assert tree["nodes"][1]["level"] == 2 - assert tree["nodes"][1]["section_path"] == ["Methods", "Statistical Analysis"] + assert tree["nodes"][0]["title"] == "Methods" + assert tree["nodes"][1]["title"] == "Results" + assert len(tree["nodes"][0]["children"]) == 1 -def test_build_structure_tree_extends_page_span(): - structured_blocks = [ - {"paper_id": "X", "page": 1, "block_id": "s1", "role": "section_heading", "text": "Introduction"}, - {"paper_id": "X", "page": 1, "block_id": "p1", "role": "body_paragraph", "text": "Para 1"}, - {"paper_id": "X", "page": 2, "block_id": "p2", "role": "body_paragraph", "text": "Para 2"}, - {"paper_id": "X", "page": 3, "block_id": "p3", "role": "body_paragraph", "text": "Para 3"}, - ] - tree = build_structure_tree(structured_blocks) +def test_page_span_extended_from_emitted_blocks(): + heading_events = [_h(0, 2, "Intro", page=1, block_id="b1")] + emitted = [_e(0, "b1", page=1), _e(1, "b2", page=2), _e(2, "b3", page=3)] + structured = [{"paper_id": "P003"}] + tree = build_structure_tree(heading_events, emitted, structured) assert tree["nodes"][0]["page_span"] == [1, 3] -def test_build_structure_tree_handles_introduction_and_abstract_headings(): - blocks = [ - {"paper_id": "Y", "page": 1, "block_id": "a1", "role": "abstract_heading", "text": "Abstract"}, - {"paper_id": "Y", "page": 1, "block_id": "b1", "role": "abstract_body", "text": "Summary here."}, - {"paper_id": "Y", "page": 1, "block_id": "i1", "role": "introduction_heading", "text": "Introduction"}, - {"paper_id": "Y", "page": 2, "block_id": "b2", "role": "body_paragraph", "text": "Background."}, - ] - tree = build_structure_tree(blocks) - assert len(tree["nodes"]) == 2 - assert tree["nodes"][0]["title"] == "Abstract" - assert tree["nodes"][1]["title"] == "Introduction" - assert tree["nodes"][0]["kind"] == "section" - - -def test_build_structure_tree_ignores_heading_roles_with_empty_text(): - blocks = [ - {"paper_id": "Z", "page": 1, "block_id": "h1", "role": "section_heading", "text": ""}, - {"paper_id": "Z", "page": 1, "block_id": "b1", "role": "body_paragraph", "text": "Some text."}, - ] - tree = build_structure_tree(blocks) - assert len(tree["nodes"]) == 0 - - -def test_build_structure_tree_empty_input(): - tree = build_structure_tree([]) +def test_empty_heading_events(): + tree = build_structure_tree([], [], []) assert tree["paper_id"] == "" assert tree["nodes"] == [] @@ -102,13 +118,45 @@ def test_summarize_role_index_counts_roles(): assert summary["role_counts"]["figure_captions"] == 1 -def test_build_structure_tree_block_span_accumulates(): - blocks = [ - {"paper_id": "P", "page": 1, "block_id": "sec1", "role": "section_heading", "text": "Results"}, - {"paper_id": "P", "page": 1, "block_id": "r1", "role": "body_paragraph", "text": "Result A"}, - {"paper_id": "P", "page": 2, "block_id": "r2", "role": "body_paragraph", "text": "Result B"}, +def test_own_block_ids_excludes_children(): + heading_events = [ + _h(0, 2, "Discussion", block_id="d1"), + _h(2, 3, "Limitation", block_id="d3"), ] - tree = build_structure_tree(blocks) - assert len(tree["nodes"][0]["block_span"]) == 3 - assert [1, "sec1"] in tree["nodes"][0]["block_span"] - assert [2, "r2"] in tree["nodes"][0]["block_span"] + emitted = [_e(0, "d1"), _e(1, "d2"), _e(2, "d3"), _e(3, "d4")] + structured = [{"paper_id": "P004"}] + tree = build_structure_tree(heading_events, emitted, structured) + parent = tree["nodes"][0] + child = parent["children"][0] + assert parent["own_block_ids"] == ["d2"] + assert child["own_block_ids"] == ["d4"] + assert parent["subtree_block_ids"] == ["d1", "d2", "d3", "d4"] + assert child["subtree_block_ids"] == ["d3", "d4"] + + +def test_summary_in_last_child_scope(): + """Summary after last child falls within child's scope (no boundary).""" + heading_events = [ + _h(0, 2, "Discussion", block_id="d1"), + _h(2, 3, "Mechanism", block_id="d3"), + ] + emitted = [ + _e(0, "d1"), _e(1, "d2"), + _e(2, "d3"), _e(3, "d4"), _e(4, "d5"), + ] + structured = [{"paper_id": "P005"}] + tree = build_structure_tree(heading_events, emitted, structured) + parent = tree["nodes"][0] + child = parent["children"][0] + assert parent["own_block_ids"] == ["d2"] + assert child["own_block_ids"] == ["d4", "d5"] + + +def test_rendered_order_differs_from_structured_order(): + heading_events = [_h(1, 2, "Results", block_id="b1")] + emitted = [_e(0, "b3", page=1), _e(1, "b1", page=1), _e(2, "b2", page=1)] + structured = [{"paper_id": "P006"}] + tree = build_structure_tree(heading_events, emitted, structured) + assert "b3" not in tree["nodes"][0]["subtree_block_ids"] + assert "b1" in tree["nodes"][0]["subtree_block_ids"] + assert "b2" in tree["nodes"][0]["subtree_block_ids"] diff --git a/tests/test_layer4_units_and_fts.py b/tests/test_layer4_units_and_fts.py index 35e5523e..cae80860 100644 --- a/tests/test_layer4_units_and_fts.py +++ b/tests/test_layer4_units_and_fts.py @@ -15,6 +15,22 @@ from paperforge.memory.schema import ( ) +def _node(node_id, title, level=2, block_id=None, own_block_ids=None, + subtree_block_ids=None, children=None, page_span=None): + return { + "node_id": node_id, + "kind": "section", + "title": title, + "level": level, + "block_id": block_id or node_id.split(":")[1], + "own_block_ids": own_block_ids or [], + "subtree_block_ids": subtree_block_ids or [], + "children": children or [], + "objects": [], + "page_span": page_span or [1, 1], + } + + def test_build_unit_id_format(): uid = build_unit_id("ABCD1234", "body", "sec:b1", 1, "b1", 1, "b2") assert uid == "ABCD1234:body:sec:b1:1-b1:1-b2" @@ -24,42 +40,36 @@ def test_build_body_units_assigns_stable_ids_and_audit_fields(): tree = { "paper_id": "ABCD1234", "nodes": [ - { - "node_id": "sec:b1", - "title": "Methods", - "section_path": ["Methods"], - "page_span": [1, 1], - "block_span": [[1, "b1"], [1, "b2"]], - } + _node("sec:b1", "Methods", own_block_ids=["b2"]), ], } blocks = [ - {"paper_id": "ABCD1234", "page": 1, "block_id": "b1", "role": "section_heading", "text": "Methods"}, - {"paper_id": "ABCD1234", "page": 1, "block_id": "b2", "role": "body_paragraph", "text": "We recruited 30 patients."}, + {"block_id": "b1", "role": "section_heading", "text": "Methods"}, + {"block_id": "b2", "role": "body_paragraph", "text": "We recruited 30 patients."}, ] units = build_body_units(tree=tree, structured_blocks=blocks) + assert len(units) == 1 assert units[0]["unit_id"].startswith("ABCD1234:body:") assert units[0]["indexable"] is True assert units[0]["veto_reason"] == "" + assert units[0]["section_path"] == "Methods" + assert units[0]["section_level"] == 2 + assert units[0]["section_title"] == "Methods" + assert units[0]["part_ordinal"] == 0 + assert "section_path_json" in units[0] def test_build_body_units_concatenates_multiple_paragraphs(): tree = { "paper_id": "P001", "nodes": [ - { - "node_id": "sec:x", - "title": "Results", - "section_path": ["Results"], - "page_span": [2, 2], - "block_span": [[2, "b1"], [2, "b2"], [2, "b3"]], - } + _node("sec:x", "Results", own_block_ids=["b2", "b3"]), ], } blocks = [ - {"paper_id": "P001", "page": 2, "block_id": "b1", "role": "section_heading", "text": "Results"}, - {"paper_id": "P001", "page": 2, "block_id": "b2", "role": "body_paragraph", "text": "First result."}, - {"paper_id": "P001", "page": 2, "block_id": "b3", "role": "body_paragraph", "text": "Second result."}, + {"block_id": "b1", "role": "section_heading", "text": "Results"}, + {"block_id": "b2", "role": "body_paragraph", "text": "First result."}, + {"block_id": "b3", "role": "body_paragraph", "text": "Second result."}, ] units = build_body_units(tree=tree, structured_blocks=blocks) assert "First result." in units[0]["unit_text"] @@ -70,40 +80,28 @@ def test_build_body_units_marks_empty_as_non_indexable(): tree = { "paper_id": "P002", "nodes": [ - { - "node_id": "sec:empty", - "title": "Empty", - "section_path": ["Empty"], - "page_span": [3, 3], - "block_span": [[3, "e1"]], - } + _node("sec:empty", "Empty", own_block_ids=[]), ], } blocks = [ - {"paper_id": "P002", "page": 3, "block_id": "e1", "role": "section_heading", "text": "Empty"}, + {"block_id": "e1", "role": "section_heading", "text": "Empty"}, ] units = build_body_units(tree=tree, structured_blocks=blocks) - assert units[0]["indexable"] is False - assert units[0]["veto_reason"] == "empty" + assert units == [] # no own_block_ids with body role → no units def test_build_body_units_token_estimate(): tree = { "paper_id": "P003", "nodes": [ - { - "node_id": "sec:t1", - "title": "Introduction", - "section_path": ["Introduction"], - "page_span": [1, 1], - "block_span": [[1, "b1"]], - } + _node("sec:t1", "Introduction", own_block_ids=["b1"]), ], } blocks = [ - {"paper_id": "P003", "page": 1, "block_id": "b1", "role": "body_paragraph", "text": "A" * 100}, + {"block_id": "b1", "role": "body_paragraph", "text": "A" * 100}, ] units = build_body_units(tree=tree, structured_blocks=blocks) + assert len(units) == 1 assert units[0]["token_estimate"] == 25 # 100 // 4 @@ -112,48 +110,124 @@ def test_build_body_units_empty_tree(): assert units == [] +def test_body_unit_excludes_reference_item(): + tree = { + "paper_id": "P005", + "nodes": [ + _node("sec:body", "Discussion", own_block_ids=["b1", "r1"]), + ], + } + blocks = [ + {"block_id": "b1", "role": "body_paragraph", "text": "Main text."}, + {"block_id": "r1", "role": "reference_item", "text": "[1] Smith et al."}, + ] + units = build_body_units(tree=tree, structured_blocks=blocks) + assert len(units) == 1 + assert "Main text." in units[0]["unit_text"] + assert "Smith" not in units[0]["unit_text"] + + +def test_backmatter_body_creates_separate_unit(): + tree = { + "paper_id": "P006", + "nodes": [ + _node("sec:funding", "Funding", own_block_ids=["s1"]), + ], + } + blocks = [ + {"block_id": "s1", "role": "structured_insert", "text": "This work was funded by NIH."}, + ] + units = build_body_units(tree=tree, structured_blocks=blocks) + assert len(units) == 1 + assert units[0]["unit_kind"] == "backmatter_body" + + +def test_mixed_body_and_backmatter_split(): + tree = { + "paper_id": "P007", + "nodes": [ + _node("sec:end", "End", own_block_ids=["b1", "s1"]), + ], + } + blocks = [ + {"block_id": "b1", "role": "body_paragraph", "text": "Main text."}, + {"block_id": "s1", "role": "structured_insert", "text": "Data available on request."}, + ] + units = build_body_units(tree=tree, structured_blocks=blocks) + kinds = {u["unit_kind"] for u in units} + assert "body" in kinds + assert "backmatter_body" in kinds + assert len(units) == 2 + + +def test_token_cap_splits_into_parts(): + tree = { + "paper_id": "P008", + "nodes": [ + _node("sec:big", "Long Section", own_block_ids=["b1"]), + ], + } + blocks = [ + {"block_id": "b1", "role": "body_paragraph", "text": "word " * 1000}, + ] + units = build_body_units(tree=tree, structured_blocks=blocks) + assert len(units) >= 2 # should be split + assert units[0]["part_ordinal"] == 1 + assert units[1]["part_ordinal"] == 2 + assert units[0]["section_path"] == units[1]["section_path"] + + +def test_recursive_walk_produces_correct_section_path(): + tree = { + "paper_id": "P009", + "nodes": [ + _node("sec:root", "Root", own_block_ids=["b1"], children=[ + _node("sec:child", "Child", level=3, own_block_ids=["b2"]), + ]), + ], + } + blocks = [ + {"block_id": "b1", "role": "body_paragraph", "text": "Root text."}, + {"block_id": "b2", "role": "body_paragraph", "text": "Child text."}, + ] + units = build_body_units(tree=tree, structured_blocks=blocks) + paths = {u["unit_id"]: u["section_path"] for u in units} + root_unit = [u for u in units if u["section_title"] == "Root"][0] + child_unit = [u for u in units if u["section_title"] == "Child"][0] + assert root_unit["section_path"] == "Root" + assert child_unit["section_path"] == "Root/Child" + + def test_build_object_units_from_role_index(): tree = { "paper_id": "P010", "nodes": [ - { - "node_id": "sec:figs", - "title": "Figures", - "section_path": ["Figures"], - "page_span": [4, 4], - "block_span": [[4, "f1"], [4, "f2"]], - } + _node("sec:figs", "Figures", own_block_ids=["f2"], + subtree_block_ids=["f1", "f2"], + page_span=[4, 4]), ], } blocks = [ - {"paper_id": "P010", "page": 4, "block_id": "f1", "role": "section_heading", "text": "Figures"}, - {"paper_id": "P010", "page": 4, "block_id": "f2", "role": "body_paragraph", "text": "As shown in Figure 1."}, + {"block_id": "f1", "role": "section_heading", "text": "Figures"}, + {"block_id": "f2", "role": "body_paragraph", "text": "As shown in Figure 1."}, ] role_index = { "figure_captions": [ - {"page": 4, "block_id": "f2", "text": "Figure 1: Results", "role": "figure_caption"}, + {"figure_id": "Figure 1", "caption_block_id": "f2", "text": "Figure 1: Results"}, ] } units = build_object_units(tree=tree, structured_blocks=blocks, role_index=role_index) assert len(units) >= 1 - assert units[0]["unit_id"].startswith("P010:object:") assert units[0]["object_kind"] == "figure" assert units[0]["object_label"] == "Figure 1" assert units[0]["caption_text"] == "Figure 1: Results" - assert "As shown in Figure 1." in units[0]["nearby_body_text"] def test_build_object_units_empty_role_index(): tree = { "paper_id": "P011", "nodes": [ - { - "node_id": "sec:empty", - "title": "Empty", - "section_path": ["Empty"], - "page_span": [5, 5], - "block_span": [[5, "e1"]], - } + _node("sec:empty", "Empty", own_block_ids=[]), ], } units = build_object_units(tree=tree, structured_blocks=[], role_index={}) @@ -162,10 +236,10 @@ def test_build_object_units_empty_role_index(): def test_build_paper_manifest(): body_units = [ - {"unit_id": "P001:body:sec:1:1-b1:1-b2", "paper_id": "P001"}, + {"unit_id": "P001:body:sec:1", "paper_id": "P001"}, ] object_units = [ - {"unit_id": "P001:object:sec:1:1-f1:1-f1", "paper_id": "P001"}, + {"unit_id": "P001:object:sec:1", "paper_id": "P001"}, ] manifest = build_paper_manifest( paper_id="P001", @@ -193,16 +267,23 @@ def test_fts_body_units_table_creation(tmp_path): conn.execute(CREATE_BODY_UNITS_FTS) conn.commit() - # Insert a row and let FTS index it conn.execute( - """INSERT INTO body_units (unit_id, paper_id, section_path, unit_text, - page_span_json, block_span_json, token_estimate, indexable, veto_reason, quality_hints_json) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + """INSERT INTO body_units (unit_id, paper_id, section_path, + section_path_json, section_level, section_title, + unit_text, unit_kind, part_ordinal, + page_span_json, block_span_json, token_estimate, + indexable, veto_reason, quality_hints_json) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", ( - "P001:body:sec:1:1-b1:1-b2", + "P001:body:sec:1", "P001", "Methods", + '["Methods"]', + 2, + "Methods", "We recruited 30 patients.", + "body", + 0, json.dumps([1, 1]), json.dumps([[1, "b1"], [1, "b2"]]), 10, @@ -213,7 +294,6 @@ def test_fts_body_units_table_creation(tmp_path): ) conn.commit() - # Manually populate FTS (no trigger yet) conn.execute( """INSERT INTO body_units_fts(rowid, unit_id, paper_id, section_path, unit_text) SELECT rowid, unit_id, paper_id, section_path, unit_text FROM body_units""" @@ -225,12 +305,12 @@ def test_fts_body_units_table_creation(tmp_path): ("patients",), ).fetchall() assert len(results) == 1 - assert results[0][0] == "P001:body:sec:1:1-b1:1-b2" - # Verify unit_kind was set via DEFAULT - row = conn.execute("SELECT unit_kind FROM body_units WHERE unit_id = ?", ("P001:body:sec:1:1-b1:1-b2",)).fetchone() + assert results[0][0] == "P001:body:sec:1" + row = conn.execute("SELECT unit_kind FROM body_units WHERE unit_id = ?", ("P001:body:sec:1",)).fetchone() assert row["unit_kind"] == "body" conn.close() + def test_object_units_table_creation(tmp_path): db = tmp_path / "test.db" conn = sqlite3.connect(str(db)) @@ -243,7 +323,7 @@ def test_object_units_table_creation(tmp_path): page_span_json, block_span_json) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""", ( - "P001:object:sec:1:1-f1:1-f1", + "P001:object:sec:1", "P001", "Figures", "figure", @@ -254,7 +334,8 @@ def test_object_units_table_creation(tmp_path): json.dumps([[4, "f1"]]), ), ) - row = conn.execute("SELECT unit_id, object_kind, object_label, caption_text FROM object_units WHERE unit_id = ?", ("P001:object:sec:1:1-f1:1-f1",)).fetchone() + row = conn.execute("SELECT unit_id, object_kind, object_label, caption_text FROM object_units WHERE unit_id = ?", + ("P001:object:sec:1",)).fetchone() assert row is not None assert row["object_kind"] == "figure" assert row["object_label"] == "Figure 1" @@ -279,7 +360,7 @@ def test_ensure_schema_includes_new_tables(tmp_path): conn.close() -def test_body_units_schema_includes_unit_kind(tmp_path): +def test_body_units_schema_includes_v4_columns(tmp_path): db = tmp_path / "test_schema_body_units.db" conn = sqlite3.connect(str(db)) conn.row_factory = sqlite3.Row @@ -289,7 +370,11 @@ def test_body_units_schema_includes_unit_kind(tmp_path): cols = { r[1] for r in conn.execute("PRAGMA table_info(body_units)").fetchall() } - assert "unit_kind" in cols, "body_units missing unit_kind column" + assert "unit_kind" in cols + assert "section_path_json" in cols + assert "section_level" in cols + assert "section_title" in cols + assert "part_ordinal" in cols conn.close() diff --git a/tests/test_ocr_rebuild.py b/tests/test_ocr_rebuild.py index 68e3d7d0..4d43abd1 100644 --- a/tests/test_ocr_rebuild.py +++ b/tests/test_ocr_rebuild.py @@ -3,6 +3,8 @@ from __future__ import annotations from pathlib import Path +from paperforge.worker.ocr_render import RenderOutput + def test_rebuild_selector_only_targets_derived_stale_papers() -> None: @@ -215,8 +217,7 @@ def test_run_derived_rebuild_skips_span_backfill_when_valid(tmp_path: Path, monk monkeypatch.setattr("paperforge.worker.ocr_tables.build_table_inventory", lambda *args, **kwargs: {"tables": [], "unmatched_assets": []}) monkeypatch.setattr("paperforge.worker.ocr_tables.write_back_table_roles", lambda *args, **kwargs: None) monkeypatch.setattr("paperforge.worker.ocr_tables.write_table_inventory", lambda *args, **kwargs: None) - monkeypatch.setattr("paperforge.worker.ocr_objects.extract_and_write_objects", lambda *args, **kwargs: None) - monkeypatch.setattr("paperforge.worker.ocr_render.render_fulltext_markdown", lambda *args, **kwargs: "") + monkeypatch.setattr("paperforge.worker.ocr_render.render_fulltext_markdown", lambda *args, **kwargs: RenderOutput(markdown="", heading_events=[], emitted_block_events=[])) monkeypatch.setattr("paperforge.worker.ocr_render.write_render_outputs", lambda *args, meta=None, **kwargs: dict(meta) if meta else {}) monkeypatch.setattr("paperforge.worker.ocr_health.build_ocr_health", lambda *args, **kwargs: {}) monkeypatch.setattr("paperforge.worker.ocr_health.build_ocr_raw_integrity_health", lambda *args, **kwargs: {}) @@ -266,8 +267,7 @@ def test_run_derived_rebuild_records_unavailable_pdf_missing_without_rerun(tmp_p monkeypatch.setattr("paperforge.worker.ocr_tables.build_table_inventory", lambda *args, **kwargs: {"tables": [], "unmatched_assets": []}) monkeypatch.setattr("paperforge.worker.ocr_tables.write_back_table_roles", lambda *args, **kwargs: None) monkeypatch.setattr("paperforge.worker.ocr_tables.write_table_inventory", lambda *args, **kwargs: None) - monkeypatch.setattr("paperforge.worker.ocr_objects.extract_and_write_objects", lambda *args, **kwargs: None) - monkeypatch.setattr("paperforge.worker.ocr_render.render_fulltext_markdown", lambda *args, **kwargs: "") + monkeypatch.setattr("paperforge.worker.ocr_render.render_fulltext_markdown", lambda *args, **kwargs: RenderOutput(markdown="", heading_events=[], emitted_block_events=[])) monkeypatch.setattr("paperforge.worker.ocr_render.write_render_outputs", lambda *args, meta=None, **kwargs: dict(meta) if meta else {}) monkeypatch.setattr("paperforge.worker.ocr_health.build_ocr_health", lambda *args, **kwargs: {}) monkeypatch.setattr("paperforge.worker.ocr_health.build_ocr_raw_integrity_health", lambda *args, **kwargs: {}) @@ -526,8 +526,7 @@ def test_run_derived_rebuild_for_keys_still_uses_public_build_table_inventory(tm monkeypatch.setattr("paperforge.worker.ocr_figure_reader.synthesize_reader_figures", lambda *a, **kw: {}) monkeypatch.setattr("paperforge.worker.ocr_tables.write_back_table_roles", lambda *a, **kw: None) monkeypatch.setattr("paperforge.worker.ocr_tables.write_table_inventory", lambda *a, **kw: None) - monkeypatch.setattr("paperforge.worker.ocr_objects.extract_and_write_objects", lambda *a, **kw: None) - monkeypatch.setattr("paperforge.worker.ocr_render.render_fulltext_markdown", lambda *a, **kw: "") + monkeypatch.setattr("paperforge.worker.ocr_render.render_fulltext_markdown", lambda *a, **kw: RenderOutput(markdown="", heading_events=[], emitted_block_events=[])) monkeypatch.setattr("paperforge.worker.ocr_render.write_render_outputs", lambda *a, meta=None, **kw: dict(meta) if meta else {}) monkeypatch.setattr("paperforge.worker.ocr_health.build_ocr_health", lambda *a, **kw: {}) monkeypatch.setattr("paperforge.worker.ocr_health.build_ocr_raw_integrity_health", lambda *a, **kw: {}) diff --git a/tests/unit/memory/test_schema.py b/tests/unit/memory/test_schema.py index 5d539770..194afdaa 100644 --- a/tests/unit/memory/test_schema.py +++ b/tests/unit/memory/test_schema.py @@ -55,7 +55,7 @@ def test_get_schema_version_returns_zero_when_no_meta(): db_path = Path(tmp.name) try: conn = get_connection(db_path) - ensure_schema(conn) + # No ensure_schema — meta table doesn't exist assert get_schema_version(conn) == 0 conn.close() finally: @@ -69,7 +69,7 @@ def test_get_schema_version_returns_stored_value(): conn = get_connection(db_path) ensure_schema(conn) conn.execute( - "INSERT INTO meta (key, value) VALUES ('schema_version', '1')" + "INSERT OR REPLACE INTO meta (key, value) VALUES ('schema_version', '1')" ) conn.commit() assert get_schema_version(conn) == 1 @@ -86,7 +86,7 @@ def test_schema_version_mismatch_triggers_rebuild_semantics(): conn = get_connection(db_path) ensure_schema(conn) conn.execute( - "INSERT INTO meta (key, value) VALUES ('schema_version', '99')" + "INSERT OR REPLACE INTO meta (key, value) VALUES ('schema_version', '99')" ) conn.commit() stored = get_schema_version(conn)