mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 06:50:53 +00:00
feat: PR 8 — Object Vector Retrieval (paperforge_objects collection)
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.
This commit is contained in:
LLLin000
parent
e551f79a77
commit
0dfbf2a68d
10 changed files with 570 additions and 36 deletions
363
docs/plans/pr8-object-vector-retrieval-plan.md
Normal file
363
docs/plans/pr8-object-vector-retrieval-plan.md
Normal file
|
|
@ -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 |
|
||||
|
|
@ -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="",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Reference in a new issue