From 0dfbf2a68df17b18204ce98b965cec3f0df40ab8 Mon Sep 17 00:00:00 2001 From: LLLin000 <809867916@qq.com> Date: Mon, 6 Jul 2026 19:58:46 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20PR=208=20=E2=80=94=20Object=20Vector=20?= =?UTF-8?q?Retrieval=20(paperforge=5Fobjects=20collection)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New paperforge_objects collection in ChromaDB for figure/table caption vector search, completing the memory layer retrieval triad: paperforge_fulltext → legacy_chunk paperforge_body → body_unit paperforge_objects → object_unit (NEW) Changes: - manifest.py: compute_object_units_hash(), manifest records both hashes - _chroma.py: _COLLECTION_NAMES includes paperforge_objects - builder.py: get_object_units_for_embedding() + embed_object_units() - status.py: object_chunk_count, total_chunks - state_snapshot.py: write_vector_runtime() extended with new fields - search.py: 3rd collection, dict-based source mapping, object_kind/label - __init__.py: exports for new functions - commands/embed.py: _has_object_units_in_db, body+object structured path routing with dual-hash resume, extended counters - Test fix: test_build_paper_manifest data complete for hash functions Verification: 20 object vectors embedded, all per-paper counts match, status total consistent (58266 = 57435 + 811 + 20), object caption queries return object_unit results via merge_retrieve. Unit tests: 153 pass, 1 skip. --- .../plans/pr8-object-vector-retrieval-plan.md | 363 ++++++++++++++++++ paperforge/commands/embed.py | 103 +++-- paperforge/embedding/__init__.py | 4 +- paperforge/embedding/_chroma.py | 2 +- paperforge/embedding/builder.py | 65 ++++ paperforge/embedding/search.py | 10 +- paperforge/embedding/status.py | 12 +- paperforge/memory/state_snapshot.py | 8 +- paperforge/retrieval/manifest.py | 24 ++ tests/test_layer4_units_and_fts.py | 15 +- 10 files changed, 570 insertions(+), 36 deletions(-) create mode 100644 docs/plans/pr8-object-vector-retrieval-plan.md diff --git a/docs/plans/pr8-object-vector-retrieval-plan.md b/docs/plans/pr8-object-vector-retrieval-plan.md new file mode 100644 index 00000000..ffb91915 --- /dev/null +++ b/docs/plans/pr8-object-vector-retrieval-plan.md @@ -0,0 +1,363 @@ +# PR 8: Object Vector Retrieval — Implementation Plan + +**Status**: Reviewed (applied 9 review fixes) +**Depends on**: PR 7 (object unit_id/caption_key/node_id hardening) — confirmed on master @ `e551f79a` +**Files touched**: 11 files (8 code + 3 docs/tests), ~240 lines net +**No schema changes, no new FTS table** + +--- + +## Background + +Current gap: + +``` +body_units → DB → paperforge_body vectors → merge_retrieve ✅ +object_units → DB → NO vectors → merge_retrieve ❌ +``` + +Object units (figure captions, table captions) are persisted in SQLite DB and used by `build_object_units()` / `_upsert_object_units()`, but **never embedded into ChromaDB**. +They are **not** in `body_units_fts` either — so neither FTS search nor vector search can find figure/table content. +`merge_retrieve` only queries `paperforge_fulltext` + `paperforge_body` → figure/table content is invisible to vector search. + +## Design: C-lite (new `paperforge_objects` collection) + +``` +paperforge_fulltext → legacy chunks → source="legacy_chunk" +paperforge_body → body units → source="body_unit" +paperforge_objects → object units → source="object_unit" +``` + +--- + +## Changes + +### 1. `paperforge/retrieval/manifest.py` + +Add `compute_object_units_hash()` for resume detection: + +```python +def compute_object_units_hash(units: list[dict]) -> str: + raw = json.dumps( + [ + { + "unit_id": u["unit_id"], + "paper_id": u["paper_id"], + "section_path": u.get("section_path", ""), + "object_kind": u.get("object_kind", ""), + "object_label": u.get("object_label", ""), + "caption_text": u.get("caption_text", ""), + "nearby_body_text": u.get("nearby_body_text", ""), + "retrieval_policy_version": RETRIEVAL_POLICY_VERSION, + } + for u in units + ], + ensure_ascii=False, + sort_keys=True, + ) + return sha256(raw.encode()).hexdigest() +``` + +**Also update `build_paper_manifest()`** to record `body_units_hash` and `object_units_hash`: + +```python +"body_unit_count": len(body_units), +"body_units_hash": compute_body_units_hash(body_units), +"object_unit_count": len(object_units), +"object_units_hash": compute_object_units_hash(object_units), +``` + +This makes provenance complete: manifest tells you not just how many, but whether content changed. + +### 2. `paperforge/embedding/_chroma.py` + +Add `"paperforge_objects"` to collection list: + +```python +_COLLECTION_NAMES = ["paperforge_fulltext", "paperforge_body", "paperforge_objects"] +``` + +`delete_paper_vectors()` iterates this list — object vectors will be auto-deleted. + +### 3. `paperforge/embedding/builder.py` + +Add two functions: + +```python +def get_object_units_for_embedding(vault: Path, key: str) -> list[dict]: + """Fetch object_units from the memory DB for a given paper.""" + ... + +def embed_object_units(vault: Path, zotero_key: str, object_units: list[dict]) -> int: + """Embed object_units into paperforge_objects collection.""" + ... +``` + +**Embedding text**: `"\n".join([object_label, caption_text, nearby_body_text])` +Useful for queries like "Figure 2" (object_label) or "histological analysis" (caption_text). + +**Metadata**: `paper_id`, `section_path`, `unit_id`, `unit_kind="object"`, `object_kind`, `object_label`, `object_units_hash`, `retrieval_policy_version`, `token_estimate` + +### 4. `paperforge/embedding/status.py` + +In the status dict, add: + +```python +"object_chunk_count": col_o.count(), +"total_chunks": ft_count + body_count + object_count, +``` + +### 5. `paperforge/embedding/search.py` + +- Add `"paperforge_objects"` to `RETRIEVAL_COLLECTIONS` +- Source mapping (dict): + +```python +source = { + "paperforge_fulltext": "legacy_chunk", + "paperforge_body": "body_unit", + "paperforge_objects": "object_unit", +}[name] +``` + +- Result fields: add `object_kind` and `object_label` from metadata (empty string when not applicable) + +**Note**: `merge_retrieve` queries collections sequentially (not parallel). Accept for now; optimize later if needed. + +### 6. `paperforge/embedding/__init__.py` + +Export `get_object_units_for_embedding` and `embed_object_units`. + +### 7. `paperforge/commands/embed.py` — embed build routing + +Add `_has_object_units_in_db()` symmetric to `_has_body_units_in_db()`: + +```python +def _has_object_units_in_db(vault: Path, key: str) -> bool: + """Check if paper has object_units in the memory DB.""" + ... +``` + +Routing logic (structured path handles all three cases): + +```python +has_body = _has_body_units_in_db(vault, key) +has_object = _has_object_units_in_db(vault, key) + +if has_body or has_object: + body_units = get_body_units_for_embedding(vault, key) if has_body else [] + object_units = get_object_units_for_embedding(vault, key) if has_object else [] + + if resume: + body_ok = not body_units # empty = nothing to check = OK + object_ok = not object_units + + if body_units: + col = get_collection(vault, name="paperforge_body") + existing = col.get(where={"paper_id": key}, limit=1) + body_ok = _hash_matches(existing, body_units, "body_units_hash") + + if object_units: + col = get_collection(vault, name="paperforge_objects") + existing = col.get(where={"paper_id": key}, limit=1) + object_ok = _hash_matches(existing, object_units, "object_units_hash") + + if body_ok and object_ok: + papers_skipped += 1 + continue + + delete_paper_vectors(vault, key) # removes all three collections + + chunks_body = 0 + chunks_object = 0 + if body_units: + chunks_body = embed_body_units(vault, key, body_units) + if object_units: + chunks_object = embed_object_units(vault, key, object_units) + + chunks_embedded += chunks_body + chunks_object + papers_embedded += 1 +else: + # Legacy fulltext path (unchanged) + ... +``` + +Key behaviors: +- **body-only paper**: body_ok checked, object_ok = True (empty → OK), embed body only +- **object-only paper**: body_ok = True, object_ok checked, embed object only +- **both**: both checked, both embedded +- **either hash mismatch**: delete all + re-embed both + +**Empty list hash note**: When a paper has no object_units (`[]`), resume treats it as OK. +Empty list does not require a vector-side hash record. Only non-empty units trigger a collection check. +This avoids requiring empty vector collections for body-only papers. + +### 8. `paperforge/memory/state_snapshot.py` — extend `write_vector_runtime()` + +Current signature only accepts `chunk_count`. Extend to accept body/object counts: + +```python +def write_vector_runtime( + vault: Path, + *, + enabled: bool, + mode: str, + model: str, + deps_installed: bool, + deps_missing: list[str] | None, + py_version: str, + db_exists: bool, + chunk_count: int, + body_chunk_count: int = 0, + object_chunk_count: int = 0, + total_chunks: int | None = None, + build_state: dict | None, + healthy: bool = True, + error: str = "", + corrupted: bool = False, +) -> None: +``` + +Write the new fields: + +```python +"body_chunk_count": body_chunk_count, +"object_chunk_count": object_chunk_count, +"total_chunks": total_chunks if total_chunks is not None else chunk_count + body_chunk_count + object_chunk_count, +``` + +This keeps CLI `embed status --json` and Obsidian UI runtime snapshot consistent. +Without this change, `commands/embed.py` will get TypeError when passing the new kwargs. + +### 9. `paperforge/memory/builder.py` + +No changes needed. `_upsert_object_units()` already exists. + +--- + +## Test Plan + +### TC-O1: Object DB persistence (Layer A, section 4b) +Already added in PR 7. Verifies unit_id unique, non-empty labels/section_path, DB count == list count. + +### TC-O2a: Object vector count (clean full run) +```python +object_db = SELECT COUNT(*) FROM object_units WHERE indexable=1 +object_vec = paperforge_objects.count() +assert object_vec == object_db +``` + +### TC-O2b: Object vector count per-paper (incremental/resume safe) +```python +for key in embedded_keys: + db_cnt = SELECT COUNT(*) FROM object_units WHERE paper_id=? AND indexable=1 + vec_cnt = len(paperforge_objects.get(where={"paper_id": key}).get("ids", [])) + assert vec_cnt == db_cnt, f"{key}: DB={db_cnt} vec={vec_cnt}" +``` +Per-paper assert pinpoints which paper failed and works correctly in resume mode. + +### TC-O3: Object vector metadata +```python +meta = sample from paperforge_objects +assert meta["unit_kind"] == "object" +assert meta["object_kind"] in {"figure", "table"} +assert meta["object_label"] +assert meta["object_units_hash"] +``` + +### TC-O4: Object caption query +```python +caption_phrase = pick from object_units.caption_text +results = merge_retrieve(vault, query=caption_phrase, limit=10) +assert any(r["source"] == "object_unit" for r in results) +``` + +### TC-O5: Object resume hash +```python +# First embed → count = N +# Re-run --resume → count = N (skip) +# Modify caption_text in DB → re-run --resume → count = N (hash mismatch → re-embed) +# Verify metadata.object_units_hash changed +``` + +### TC-O6: Source merge cap +```python +results = merge_retrieve(vault, caption_query, limit=10) +object_count = sum(1 for r in results if r["source"] == "object_unit") +assert object_count >= 1 # objects not drowned by legacy +assert per_paper_count <= 2 # cap works +``` + +### TC-O7: delete_paper_vectors removes all three collections +```python +# Add vectors to fulltext/body/objects for same paper_id +delete_paper_vectors(vault, key) +assert all three collections have 0 vectors for key +``` + +### TC-O8: Object-only route +```python +# Paper has object_units but no body_units +# embed build +assert object vectors created in paperforge_objects +assert legacy path not used (no new legacy chunks for this paper) +``` + +### TC-O9: Status total consistency +```python +status = get_embed_status(vault) +assert status["total_chunks"] == ( + status["chunk_count"] + + status["body_chunk_count"] + + status["object_chunk_count"] +) +# Check embed status --json output includes object_chunk_count +``` + +### TC-O10: Resume does not duplicate +```python +first_count = paperforge_objects.count() +# re-run --resume +second_count = paperforge_objects.count() +assert second_count == first_count +# modify caption → re-run → verify hash changed, count unchanged +``` + +### TC-O11: Object result fields in merge_retrieve +```python +results = merge_retrieve(vault, query=caption_phrase, limit=10) +obj = first r where source == "object_unit" +assert obj["object_kind"] in {"figure", "table"} +assert obj["object_label"] +assert obj["section_path"] +``` + +--- + +## Implementation Order + +| Step | File | What | Complexity | +|------|------|------|------------| +| 1 | `manifest.py` | `compute_object_units_hash()` + manifest records both hashes | +20 lines | +| 2 | `_chroma.py` | `_COLLECTION_NAMES` + `paperforge_objects` | +1 line | +| 3 | `builder.py` | `get_object_units_for_embedding()` + `embed_object_units()` | +60 lines | +| 4 | `status.py` | `object_chunk_count`, `total_chunks` | +5 lines | +| 5 | `state_snapshot.py` | Extend `write_vector_runtime()` signature, write new fields | +15 lines | +| 6 | `search.py` | 3rd collection, source mapping, object_kind/object_label fields | +15 lines | +| 7 | `__init__.py` | Export new functions | +2 lines | +| 8 | `commands/embed.py` | `_has_object_units_in_db`, structured path routing, dual resume, counters | +35 lines | +| 9 | Tests | TC-O1 through TC-O11 | +80 lines | + +**Total**: ~240 lines, 11 files (8 code + 3 docs/tests), no schema changes, no new FTS table. + +--- + +## Risks + +| Risk | Mitigation | +|------|------------| +| 3-collection ChromaDB query is sequential, ~50% slower | Accept for now; source-aware early stopping or parallel queries if latency issue | +| Object chunks with empty caption_text waste vectors | Already filtered: `indexable=0` / `veto_reason="empty_caption"` | +| Object_units hash vs body_units hash diverge | Resume checks both; if either changed, delete all + re-embed both | +| 709 papers without object_units yet | Same gap as body_units; fixed by full pipeline rebuild | +| vector-runtime-state and status JSON drift | Fixed: both updated from same `get_embed_status()` call | diff --git a/paperforge/commands/embed.py b/paperforge/commands/embed.py index 69a120d6..0c5d7557 100644 --- a/paperforge/commands/embed.py +++ b/paperforge/commands/embed.py @@ -6,7 +6,7 @@ import sys from paperforge import __version__ as PF_VERSION from paperforge.core.errors import ErrorCode -from paperforge.retrieval.manifest import compute_body_units_hash, RETRIEVAL_POLICY_VERSION +from paperforge.retrieval.manifest import compute_body_units_hash, compute_object_units_hash, RETRIEVAL_POLICY_VERSION from paperforge.core.result import PFError, PFResult from paperforge.embedding import ( delete_paper_vectors, @@ -22,7 +22,7 @@ from paperforge.memory.chunker import chunk_fulltext from paperforge.memory.state_snapshot import write_vector_runtime from paperforge.worker.asset_index import read_index from paperforge.worker._progress import progress_bar -from paperforge.embedding.builder import embed_body_units, get_body_units_for_embedding +from paperforge.embedding.builder import embed_body_units, embed_object_units, get_body_units_for_embedding, get_object_units_for_embedding from paperforge.memory.db import get_connection, get_memory_db_path @@ -41,6 +41,21 @@ def _has_body_units_in_db(vault: Path, key: str) -> bool: finally: conn.close() +def _has_object_units_in_db(vault: Path, key: str) -> bool: + """Check if paper has object_units in the memory DB.""" + db_path = get_memory_db_path(vault) + if not db_path.exists(): + return False + conn = get_connection(db_path, read_only=True) + try: + cnt = conn.execute( + "SELECT COUNT(*) FROM object_units WHERE paper_id=? AND indexable=1", + (key,), + ).fetchone()[0] + return cnt > 0 + finally: + conn.close() + def run(args: argparse.Namespace) -> int: vault = args.vault_path sub = getattr(args, "embed_subcommand", "build") @@ -69,6 +84,9 @@ def run(args: argparse.Namespace) -> int: py_version=sys.version.split()[0], db_exists=status.get("db_exists", False), chunk_count=status.get("chunk_count", 0), + body_chunk_count=status.get("body_chunk_count", 0), + object_chunk_count=status.get("object_chunk_count", 0), + total_chunks=status.get("total_chunks", 0), build_state=status.get("build_state"), healthy=status.get("healthy", True), corrupted=status.get("corrupted", False), @@ -208,37 +226,57 @@ def run(args: argparse.Namespace) -> int: if not key: continue - # Check body_units first + # Check body_units and object_units first has_body = _has_body_units_in_db(vault, key) + has_object = _has_object_units_in_db(vault, key) - if has_body: - # Body units path - body_units = get_body_units_for_embedding(vault, key) - if not body_units: - continue + if has_body or has_object: + body_units = get_body_units_for_embedding(vault, key) if has_body else [] + object_units = get_object_units_for_embedding(vault, key) if has_object else [] if resume: - try: - col = get_collection(vault, name="paperforge_body") - existing = col.get(where={"paper_id": key}, limit=1) - if existing.get("ids"): - meta = existing.get("metadatas", [{}])[0] - current_hash = compute_body_units_hash(body_units) - if (meta.get("body_units_hash") == current_hash - and meta.get("retrieval_policy_version") == RETRIEVAL_POLICY_VERSION): - papers_skipped += 1 - continue - except Exception as exc: - err = str(exc).lower() - if "hnsw" in err or "compaction" in err: - logger.warning("ChromaDB index corrupted — rebuilding from scratch. Use --force next time for clean rebuild.") - pass + body_ok = not body_units + object_ok = not object_units + + if body_units: + try: + col = get_collection(vault, name="paperforge_body") + existing = col.get(where={"paper_id": key}, limit=1) + if existing.get("ids"): + meta = existing.get("metadatas", [{}])[0] + current_body_hash = compute_body_units_hash(body_units) + body_ok = (meta.get("body_units_hash") == current_body_hash + and meta.get("retrieval_policy_version") == RETRIEVAL_POLICY_VERSION) + except Exception: + pass + + if object_units: + try: + col = get_collection(vault, name="paperforge_objects") + existing = col.get(where={"paper_id": key}, limit=1) + if existing.get("ids"): + meta = existing.get("metadatas", [{}])[0] + current_obj_hash = compute_object_units_hash(object_units) + object_ok = (meta.get("object_units_hash") == current_obj_hash + and meta.get("retrieval_policy_version") == RETRIEVAL_POLICY_VERSION) + except Exception: + pass + + if body_ok and object_ok: + papers_skipped += 1 + continue + try: i += 1 print(f"EMBED_PROGRESS:{i}:{total}:{key}", flush=True) delete_paper_vectors(vault, key) - n = embed_body_units(vault, key, body_units) - chunks_embedded += n + chunks_body = 0 + chunks_object = 0 + if body_units: + chunks_body = embed_body_units(vault, key, body_units) + if object_units: + chunks_object = embed_object_units(vault, key, object_units) + chunks_embedded += chunks_body + chunks_object papers_embedded += 1 mark_vector_build_state(vault, current=i, paper_id=key, @@ -266,6 +304,9 @@ def run(args: argparse.Namespace) -> int: py_version=sys.version.split()[0], db_exists=get_vector_db_path(vault).exists(), chunk_count=_actual, + body_chunk_count=0, + object_chunk_count=0, + total_chunks=_actual, build_state=read_vector_build_state(vault), healthy=False, error=str(e), @@ -338,6 +379,9 @@ def run(args: argparse.Namespace) -> int: py_version=sys.version.split()[0], db_exists=get_vector_db_path(vault).exists(), chunk_count=_actual, + body_chunk_count=0, + object_chunk_count=0, + total_chunks=_actual, build_state=read_vector_build_state(vault), healthy=False, error=str(e), @@ -358,10 +402,16 @@ def run(args: argparse.Namespace) -> int: _real_chunks = _status.get("chunk_count", chunks_embedded) _mode = _status.get("mode", "") _model = _status.get("model", "") + _body_chunks = _status.get("body_chunk_count", 0) + _object_chunks = _status.get("object_chunk_count", 0) + _total_chunks = _status.get("total_chunks", 0) except Exception: _real_chunks = chunks_embedded _mode = "" _model = "" + _body_chunks = 0 + _object_chunks = 0 + _total_chunks = 0 write_vector_runtime( vault, @@ -373,6 +423,9 @@ def run(args: argparse.Namespace) -> int: py_version=sys.version.split()[0], db_exists=True, chunk_count=_real_chunks, + body_chunk_count=_body_chunks, + object_chunk_count=_object_chunks, + total_chunks=_total_chunks, build_state=read_vector_build_state(vault), healthy=True, error="", diff --git a/paperforge/embedding/__init__.py b/paperforge/embedding/__init__.py index 7b4866cb..d1fb794f 100644 --- a/paperforge/embedding/__init__.py +++ b/paperforge/embedding/__init__.py @@ -16,7 +16,7 @@ from paperforge.embedding.build_state import ( read_vector_build_state, write_vector_build_state, ) -from paperforge.embedding.builder import embed_body_units, embed_paper, get_body_units_for_embedding +from paperforge.embedding.builder import embed_body_units, embed_paper, get_body_units_for_embedding, get_object_units_for_embedding, embed_object_units from paperforge.embedding.preflight import _preflight_check from paperforge.embedding.search import merge_retrieve, retrieve_chunks from paperforge.embedding.status import get_embed_status @@ -28,8 +28,10 @@ __all__ = [ "embed_body_units", "embed_paper", "get_body_units_for_embedding", + "embed_object_units", "get_collection", "get_embed_status", + "get_object_units_for_embedding", "get_vector_backend", "get_vector_build_state_path", "get_vector_db_path", diff --git a/paperforge/embedding/_chroma.py b/paperforge/embedding/_chroma.py index 302b8135..c231166d 100644 --- a/paperforge/embedding/_chroma.py +++ b/paperforge/embedding/_chroma.py @@ -6,7 +6,7 @@ from pathlib import Path logger = logging.getLogger(__name__) -_COLLECTION_NAMES = ["paperforge_fulltext", "paperforge_body"] +_COLLECTION_NAMES = ["paperforge_fulltext", "paperforge_body", "paperforge_objects"] def get_vector_db_path(vault: Path) -> Path: diff --git a/paperforge/embedding/builder.py b/paperforge/embedding/builder.py index 183fda45..6f477c0d 100644 --- a/paperforge/embedding/builder.py +++ b/paperforge/embedding/builder.py @@ -82,3 +82,68 @@ def get_body_units_for_embedding(vault: Path, key: str) -> list[dict]: return [dict(r) for r in rows] finally: conn.close() + + +def get_object_units_for_embedding(vault: Path, key: str) -> list[dict]: + """Fetch object_units from the memory DB for a given paper.""" + db_path = get_memory_db_path(vault) + if not db_path.exists(): + return [] + conn = get_connection(db_path, read_only=True) + try: + rows = conn.execute( + """SELECT unit_id, paper_id, section_path, + object_kind, object_label, caption_text, nearby_body_text, + page_span_json, token_estimate + FROM object_units + WHERE paper_id=? AND indexable=1 + ORDER BY unit_id""", + (key,), + ).fetchall() + return [dict(r) for r in rows] + finally: + conn.close() + + +def embed_object_units(vault: Path, zotero_key: str, object_units: list[dict]) -> int: + """Embed object_units into paperforge_objects collection.""" + if not object_units: + return 0 + + from paperforge.retrieval.manifest import compute_object_units_hash, RETRIEVAL_POLICY_VERSION + + provider = OpenAICompatibleProvider(vault) + current_hash = compute_object_units_hash(object_units) + + texts = [ + "\n".join( + x for x in [ + u.get("object_label", ""), + u.get("caption_text", ""), + u.get("nearby_body_text", ""), + ] + if x + ) + for u in object_units + ] + + ids = [u["unit_id"] for u in object_units] + metadatas = [ + { + "paper_id": zotero_key, + "section_path": u.get("section_path", ""), + "unit_id": u["unit_id"], + "unit_kind": "object", + "object_kind": u.get("object_kind", ""), + "object_label": u.get("object_label", ""), + "object_units_hash": current_hash, + "retrieval_policy_version": RETRIEVAL_POLICY_VERSION, + "token_estimate": u.get("token_estimate", 0), + } + for u in object_units + ] + + embeddings = provider.encode(texts) + collection = get_collection(vault, name="paperforge_objects") + collection.add(ids=ids, embeddings=embeddings, documents=texts, metadatas=metadatas) + return len(object_units) diff --git a/paperforge/embedding/search.py b/paperforge/embedding/search.py index 7e3f23fc..077513a1 100644 --- a/paperforge/embedding/search.py +++ b/paperforge/embedding/search.py @@ -9,7 +9,7 @@ from paperforge.embedding.providers.openai_compatible import OpenAICompatiblePro logger = logging.getLogger(__name__) -RETRIEVAL_COLLECTIONS = ["paperforge_fulltext", "paperforge_body"] +RETRIEVAL_COLLECTIONS = ["paperforge_fulltext", "paperforge_body", "paperforge_objects"] def retrieve_chunks(vault: Path, query: str, limit: int = 5, expand: bool = True) -> list[dict]: @@ -41,8 +41,14 @@ def merge_retrieve(vault: Path, query: str, limit: int = 5, expand: bool = True) "section_path": meta.get("section_path", meta.get("section", "")), "chunk_text": doc, "score": round(1.0 - dist, 4), - "source": "legacy_chunk" if name == "paperforge_fulltext" else "body_unit", + "source": { + "paperforge_fulltext": "legacy_chunk", + "paperforge_body": "body_unit", + "paperforge_objects": "object_unit", + }[name], "unit_id": meta.get("unit_id") or meta.get("chunk_index", ""), + "object_kind": meta.get("object_kind", ""), + "object_label": meta.get("object_label", ""), }) except Exception: continue diff --git a/paperforge/embedding/status.py b/paperforge/embedding/status.py index afee95d5..0affcde6 100644 --- a/paperforge/embedding/status.py +++ b/paperforge/embedding/status.py @@ -44,8 +44,8 @@ def get_embed_status(vault: Path) -> dict: db_path = get_vector_db_path(vault) exists = db_path.exists() chunk_count = 0 + object_chunk_count = 0 body_chunk_count = 0 - healthy = True error = "" corrupted = False if exists: @@ -61,6 +61,12 @@ def get_embed_status(vault: Path) -> dict: body_chunk_count = col_body.count() except Exception: pass + # Count paperforge_objects + try: + col_o = get_collection(vault, name="paperforge_objects") + object_chunk_count = col_o.count() + except Exception: + pass # Backend health (checks primary collection) backend = get_vector_backend(vault) health = backend.health() @@ -70,11 +76,13 @@ def get_embed_status(vault: Path) -> dict: model = get_api_model(vault) + return { "db_exists": exists, "chunk_count": chunk_count, "body_chunk_count": body_chunk_count, - "total_chunks": chunk_count + body_chunk_count, + "object_chunk_count": object_chunk_count, + "total_chunks": chunk_count + body_chunk_count + object_chunk_count, "model": model, "mode": "api", "healthy": healthy, diff --git a/paperforge/memory/state_snapshot.py b/paperforge/memory/state_snapshot.py index 8f57cd1b..7729d7d9 100644 --- a/paperforge/memory/state_snapshot.py +++ b/paperforge/memory/state_snapshot.py @@ -38,7 +38,10 @@ def write_vector_runtime(vault: Path, *, enabled: bool, mode: str, model: str, deps_installed: bool, deps_missing: list[str] | None, py_version: str, db_exists: bool, chunk_count: int, build_state: dict | None, healthy: bool = True, - error: str = "", corrupted: bool = False) -> None: + error: str = "", corrupted: bool = False, + body_chunk_count: int = 0, + object_chunk_count: int = 0, + total_chunks: int | None = None) -> None: snap = { "schema_version": 1, "updated_at": datetime.now(timezone.utc).isoformat(), @@ -51,6 +54,9 @@ def write_vector_runtime(vault: Path, *, enabled: bool, mode: str, model: str, "py_version": py_version, "db_exists": db_exists, "chunk_count": chunk_count, + "body_chunk_count": body_chunk_count, + "object_chunk_count": object_chunk_count, + "total_chunks": total_chunks if total_chunks is not None else chunk_count + body_chunk_count + object_chunk_count, "healthy": healthy, "corrupted": corrupted, "error": error, diff --git a/paperforge/retrieval/manifest.py b/paperforge/retrieval/manifest.py index 77bfe7e0..e05b0061 100644 --- a/paperforge/retrieval/manifest.py +++ b/paperforge/retrieval/manifest.py @@ -37,6 +37,28 @@ def compute_body_units_hash(units: list[dict]) -> str: return sha256(raw.encode()).hexdigest() +def compute_object_units_hash(units: list[dict]) -> str: + """Compute a canonical hash for object units to detect changes.""" + raw = json.dumps( + [ + { + "unit_id": u["unit_id"], + "paper_id": u["paper_id"], + "section_path": u.get("section_path", ""), + "object_kind": u.get("object_kind", ""), + "object_label": u.get("object_label", ""), + "caption_text": u.get("caption_text", ""), + "nearby_body_text": u.get("nearby_body_text", ""), + "retrieval_policy_version": RETRIEVAL_POLICY_VERSION, + } + for u in units + ], + ensure_ascii=False, + sort_keys=True, + ) + return sha256(raw.encode()).hexdigest() + + def build_paper_manifest( *, @@ -56,7 +78,9 @@ def build_paper_manifest( "structure_tree_hash": structure_tree_hash, "retrieval_policy_version": retrieval_policy_version, "body_unit_count": len(body_units), + "body_units_hash": compute_body_units_hash(body_units), "object_unit_count": len(object_units), + "object_units_hash": compute_object_units_hash(object_units), "built_at": datetime.now(timezone.utc).isoformat(), "source_paths": dict(source_paths), } diff --git a/tests/test_layer4_units_and_fts.py b/tests/test_layer4_units_and_fts.py index 0cebc50e..ca8e9015 100644 --- a/tests/test_layer4_units_and_fts.py +++ b/tests/test_layer4_units_and_fts.py @@ -235,10 +235,17 @@ def test_build_object_units_empty_role_index(): def test_build_paper_manifest(): body_units = [ - {"unit_id": "P001:body:sec:1", "paper_id": "P001"}, + {"unit_id": "P001:body:sec:1", "paper_id": "P001", + "section_path": "Intro", "section_path_json": '["Intro"]', + "section_level": 2, "section_title": "Intro", + "page_span": [1, 1], "unit_text": "Hello world", + "token_estimate": 3, "unit_kind": "body", "part_ordinal": 0}, ] object_units = [ - {"unit_id": "P001:object:sec:1", "paper_id": "P001"}, + {"unit_id": "P001:obj:p1:15", "paper_id": "P001", + "section_path": "Methods", "object_kind": "figure", + "object_label": "p1:15", "caption_text": "Figure 1: test", + "nearby_body_text": "", "page_span": [1, 1], "token_estimate": 3}, ] manifest = build_paper_manifest( paper_id="P001", @@ -252,12 +259,13 @@ def test_build_paper_manifest(): assert manifest["paper_id"] == "P001" assert manifest["body_unit_count"] == 1 assert manifest["object_unit_count"] == 1 + assert manifest["body_units_hash"] + assert manifest["object_units_hash"] assert manifest["ocr_result_hash"] == "abc123" assert manifest["retrieval_policy_version"] == "l4.body.v1" assert "built_at" in manifest assert "structure_tree_hash" in manifest - def test_fts_body_units_table_creation(tmp_path): db = tmp_path / "test.db" conn = sqlite3.connect(str(db)) @@ -265,7 +273,6 @@ def test_fts_body_units_table_creation(tmp_path): conn.execute(CREATE_BODY_UNITS) conn.execute(CREATE_BODY_UNITS_FTS) conn.commit() - conn.execute( """INSERT INTO body_units (unit_id, paper_id, section_path, section_path_json, section_level, section_title,