homelab/stack
1
0
Fork 0
Code Issues 127 Pull Requests 1 Actions Packages Projects 1 Releases Wiki Activity

feat(llm): fingerprint-first index_refs; PDFs opened only when embedding; docket completion (refs #615)

Browse Source
This commit is contained in:
kert
2026-09-08 14:00:51 -04:00
parent b9c353c41a
commit 414b416bd6
3 changed files with 281 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

130
src/llm/index.py
View File

@@ -31,6 +31,7 @@ from llm.config import LlmConfig, pg_url
from llm.migrate import ensure_hnsw, migrate
from llm.pages import enrich_pdf_pages
from llm.pool import HostPool, PoolEmbeddings, embed_texts
from llm.source import DocRef
log = logging.getLogger(__name__)
@@ -51,15 +52,28 @@ def vectorstore(collection: str, cfg: LlmConfig, pool: HostPool):
)
def _state(engine: Engine, collection: str) -> dict[str, str]:
def _state(engine: Engine, collection: str) -> dict[str, tuple[str, str]]:
with engine.begin() as conn:
rows = conn.execute(
text(
"SELECT item_key, content_hash FROM index_state WHERE collection = :c"
"SELECT item_key, content_hash, fingerprint FROM index_state "
"WHERE collection = :c"
),
{"c": collection},
).fetchall()
return dict(rows)
return {r[0]: (r[1], r[2] or "") for r in rows}
def docket_complete(engine: Engine, collection: str) -> dict[str, str]:
"""docket → sealed_at for dockets fully indexed under that seal."""
with engine.begin() as conn:
rows = conn.execute(
text(
"SELECT docket, sealed_at FROM index_docket_state WHERE collection = :c"
),
{"c": collection},
).fetchall()
return {r[0]: r[1] for r in rows}
def _delete_old_chunks(engine: Engine, item_key: str, collection: str) -> None:
@@ -87,32 +101,70 @@ def _delete_old_chunks(engine: Engine, item_key: str, collection: str) -> None:
log.debug("langchain_pg_embedding not present yet; skipping delete")
def index_docs(
docs: Iterable[Doc],
def index_refs(
refs: Iterable[DocRef],
*,
collection: str,
cfg: LlmConfig,
pool: HostPool,
force: bool = False,
sealed: dict[str, str] | None = None,
mark_complete: bool = True,
) -> dict:
"""Fingerprint-first incremental indexing.
Per ref: (1) fingerprint equal to the stored one → skip without
loading; (2) load, hash the text; hash equal → record the new
fingerprint, skip embedding; (3) chunk, locate PDF pages, embed,
write. PDFs are opened only on path (3). After the loop, every
docket in *sealed* that was iterated is recorded in
``index_docket_state`` so the next run does not list it at all —
unless *mark_complete* is False (a ``--limit`` run is never complete).
"""
engine = _engine(cfg)
migrate(engine)
pool.check(cfg.embed_model)
store = vectorstore(collection, cfg, pool)
seen = _state(engine, collection)
stats = {"indexed": 0, "skipped": 0, "chunks": 0}
stats = {
"indexed": 0,
"skipped": 0,
"chunks": 0,
"fingerprint_skipped": 0,
"hash_skipped": 0,
"docket_complete": 0,
}
dockets_seen: set[str] = set()
for doc in docs:
for ref in refs:
if ref.docket:
dockets_seen.add(ref.docket)
prev_hash, prev_fp = seen.get(ref.key, ("", ""))
if not force and ref.fingerprint and ref.fingerprint == prev_fp:
stats["fingerprint_skipped"] += 1
continue
doc = ref.load()
if doc is None or not doc.text.strip():
stats["skipped"] += 1
continue
h = content_hash(doc.text)
if not force and prev_hash == h:
with engine.begin() as conn:
conn.execute(
text(
"UPDATE index_state SET fingerprint = :f "
"WHERE item_key = :k AND collection = :c"
),
{"f": ref.fingerprint, "k": ref.key, "c": collection},
)
stats["hash_skipped"] += 1
continue
chunks = enrich_pdf_pages(doc, chunk_doc(doc))
if not chunks:
stats["skipped"] += 1
continue
h = content_hash(doc.text)
if not force and seen.get(doc.key) == h:
stats["skipped"] += 1
continue
vectors = embed_texts(pool, cfg.embed_model, [c.text for c in chunks])
_delete_old_chunks(engine, doc.key, collection)
_delete_old_chunks(engine, ref.key, collection)
store.add_embeddings(
texts=[c.text for c in chunks],
embeddings=vectors,
@@ -123,12 +175,19 @@ def index_docs(
conn.execute(
text(
"INSERT INTO index_state "
"(item_key, collection, content_hash, chunk_count) "
"VALUES (:k, :c, :h, :n) "
"(item_key, collection, content_hash, chunk_count, fingerprint) "
"VALUES (:k, :c, :h, :n, :f) "
"ON CONFLICT (item_key, collection) DO UPDATE SET "
"content_hash = :h, chunk_count = :n, indexed_at = now()"
"content_hash = :h, chunk_count = :n, fingerprint = :f, "
"indexed_at = now()"
),
{"k": doc.key, "c": collection, "h": h, "n": len(chunks)},
{
"k": ref.key,
"c": collection,
"h": h,
"n": len(chunks),
"f": ref.fingerprint,
},
)
stats["indexed"] += 1
stats["chunks"] += len(chunks)
@@ -136,8 +195,47 @@ def index_docs(
if stats["indexed"] % 100 == 0:
log.info("indexed %(indexed)s (+%(chunks)s)", stats) # pragma: no cover
if mark_complete and sealed:
for d in sorted(dockets_seen & set(sealed)):
with engine.begin() as conn:
conn.execute(
text(
"INSERT INTO index_docket_state (collection, docket, sealed_at) "
"VALUES (:c, :d, :s) "
"ON CONFLICT (collection, docket) DO UPDATE SET "
"sealed_at = :s, indexed_at = now()"
),
{"c": collection, "d": d, "s": sealed[d]},
)
stats["docket_complete"] += 1
if cfg.build_ann_index:
ensure_hnsw(engine, cfg.embed_dim)
else:
log.info("ANN index skipped (build_ann_index=false); using exact search")
return stats
def index_docs(
docs: Iterable[Doc],
*,
collection: str,
cfg: LlmConfig,
pool: HostPool,
force: bool = False,
) -> dict:
"""Back-compat wrapper: Docs without fingerprints (never fingerprint-skipped)."""
refs = (
DocRef(
key=d.key,
collection=collection,
docket=d.metadata.get("docket") or None,
fingerprint="",
load=(lambda d=d: d),
)
for d in docs
)
stats = index_refs(
refs, collection=collection, cfg=cfg, pool=pool, force=force, sealed=None
)
return {k: stats[k] for k in ("indexed", "skipped", "chunks")}

34
tests/llm/test_index.py
View File

@@ -56,13 +56,18 @@ class TestIndexDocs:
assert kwargs["ids"][0].startswith("K1:")
def test_unchanged_doc_skipped(self):
"""A hash match lands in the internal ``hash_skipped`` bucket, which
index_docs (indexed/skipped/chunks only) doesn't surface — but it
must not be embedded or counted as indexed."""
h = content_hash(DOC.text)
stats, store, _, _ensure_hnsw = _run([DOC], state_rows=[("K1", h)])
assert stats["skipped"] == 1
stats, store, _, _ensure_hnsw = _run([DOC], state_rows=[("K1", h, "")])
assert stats["indexed"] == 0
store.add_embeddings.assert_not_called()
def test_changed_doc_deletes_old_chunks_first(self):
stats, store, conn, _ensure_hnsw = _run([DOC], state_rows=[("K1", "stalehash")])
stats, store, conn, _ensure_hnsw = _run(
[DOC], state_rows=[("K1", "stalehash", "")]
)
assert stats["indexed"] == 1
deletes = [
c
@@ -81,7 +86,9 @@ class TestIndexDocs:
def test_force_reembeds_unchanged(self):
h = content_hash(DOC.text)
stats, store, _, _ensure_hnsw = _run([DOC], state_rows=[("K1", h)], force=True)
stats, store, _, _ensure_hnsw = _run(
[DOC], state_rows=[("K1", h, "")], force=True
)
assert stats["indexed"] == 1
def test_empty_doc_counts_skipped(self):
@@ -91,19 +98,22 @@ class TestIndexDocs:
class TestPageEnrichmentHook:
def test_enriches_chunks_before_add(self, monkeypatch):
"""index_docs runs llm.pages.enrich_pdf_pages on every doc's chunks."""
def test_enriches_only_docs_that_embed(self, monkeypatch):
"""index_docs runs llm.pages.enrich_pdf_pages only on docs that reach
the embed path — a hash match must not open the doc's PDFs."""
from llm import index as index_mod
seen = []
def fake_enrich(doc, chunks):
seen.append(doc.key)
return chunks
monkeypatch.setattr(index_mod, "enrich_pdf_pages", fake_enrich)
monkeypatch.setattr(
index_mod,
"enrich_pdf_pages",
lambda doc, chunks: seen.append(doc.key) or chunks,
)
_run([DOC], [])
assert seen == ["K1"]
seen.clear()
_run([DOC], [("K1", content_hash(DOC.text), "")])
assert seen == []
class TestDeleteOldChunksFallback:

145
tests/llm/test_index_refs.py Normal file
View File

@@ -0,0 +1,145 @@
from __future__ import annotations
from unittest.mock import MagicMock, patch
from llm.chunk import Doc, content_hash
from llm.config import LlmConfig
from llm.index import index_refs
from llm.source import DocRef
CFG = LlmConfig(
ollama_hosts=("http://h1:11434",),
embed_model="m",
instruct_model="g",
embed_dim=768,
build_ann_index=False,
pg_host="x",
pg_port=5432,
pg_db="llm",
pg_user="llm",
)
DOC = Doc(key="K1", text="Some body text.", metadata={"docket": "D"})
def _ref(fp="fp1", docket="D", loads=None, doc=DOC):
def _load():
if loads is not None:
loads.append(doc.key)
return doc
return DocRef(
key=doc.key, collection="comments", docket=docket, fingerprint=fp, load=_load
)
def _run(
refs, state_rows, *, force=False, sealed=None, mark_complete=True, complete_rows=()
):
store = MagicMock()
engine = MagicMock()
conn = engine.begin.return_value.__enter__.return_value
def fake_execute(clause, *a, **k):
sql = str(clause)
r = MagicMock()
if "FROM index_docket_state" in sql:
r.fetchall.return_value = list(complete_rows)
elif "FROM index_state" in sql:
r.fetchall.return_value = state_rows
else:
r.fetchall.return_value = []
return r
conn.execute.side_effect = fake_execute
with (
patch("llm.index._engine", return_value=engine),
patch("llm.index.vectorstore", return_value=store),
patch("llm.index.embed_texts", return_value=[[0.0] * 3]),
patch("llm.index.ensure_hnsw"),
patch("llm.index.enrich_pdf_pages", side_effect=lambda d, c: c) as enrich,
patch("llm.index.HostPool") as MockPool,
):
MockPool.return_value.check.return_value = ["http://h1:11434"]
stats = index_refs(
refs,
collection="comments",
cfg=CFG,
pool=MockPool.return_value,
force=force,
sealed=sealed,
mark_complete=mark_complete,
)
return stats, store, conn, enrich
def test_fingerprint_match_skips_without_load():
loads = []
stats, store, _, enrich = _run([_ref("fp1", loads=loads)], [("K1", "h-old", "fp1")])
assert stats["fingerprint_skipped"] == 1 and stats["indexed"] == 0
assert loads == []
store.add_embeddings.assert_not_called()
enrich.assert_not_called()
def test_hash_match_updates_fingerprint_without_embedding():
h = content_hash(DOC.text)
stats, store, conn, enrich = _run([_ref("fp2")], [("K1", h, "fp1")])
assert stats["hash_skipped"] == 1 and stats["indexed"] == 0
store.add_embeddings.assert_not_called()
enrich.assert_not_called()
upd = [
c
for c in conn.execute.call_args_list
if "UPDATE index_state SET fingerprint" in str(c.args[0])
]
assert len(upd) == 1 and upd[0].args[1]["f"] == "fp2"
def test_changed_doc_embeds_and_records_fingerprint():
stats, store, conn, enrich = _run([_ref("fp2")], [("K1", "stale", "fp1")])
assert stats["indexed"] == 1 and stats["chunks"] == 1
store.add_embeddings.assert_called_once()
enrich.assert_called_once()
ins = [
c
for c in conn.execute.call_args_list
if "INSERT INTO index_state" in str(c.args[0])
]
assert ins[0].args[1]["f"] == "fp2"
def test_empty_fingerprint_never_matches():
stats, store, _, _ = _run([_ref("")], [("K1", "stale", "")])
assert stats["indexed"] == 1
def test_force_ignores_fingerprint_and_hash():
h = content_hash(DOC.text)
stats, store, _, _ = _run([_ref("fp1")], [("K1", h, "fp1")], force=True)
assert stats["indexed"] == 1
def test_load_none_counts_skipped():
ref = DocRef(
key="K9", collection="comments", docket="D", fingerprint="x", load=lambda: None
)
stats, *_ = _run([ref], [])
assert stats["skipped"] == 1 and stats["indexed"] == 0
def test_sealed_docket_marked_complete_after_clean_run():
stats, _, conn, _ = _run([_ref("fp1")], [], sealed={"D": "2026-10-20T00:00:00Z"})
assert stats["docket_complete"] == 1
ins = [
c
for c in conn.execute.call_args_list
if "INSERT INTO index_docket_state" in str(c.args[0])
]
assert ins[0].args[1] == {"c": "comments", "d": "D", "s": "2026-10-20T00:00:00Z"}
def test_not_marked_when_mark_complete_false_or_unsealed():
stats, _, conn, _ = _run([_ref("fp1")], [], sealed={"D": "s"}, mark_complete=False)
assert stats["docket_complete"] == 0
stats, _, conn, _ = _run([_ref("fp1")], [], sealed={})
assert stats["docket_complete"] == 0