fce24fec6b
llm.chunk.chunk_doc scanned the import-time pfs.families.FAMILIES (five hand families) at chunk time, while `stack llm restamp` refreshed from the DuckDB replica first — an indexed chunk's `families` metadata could disagree with what a restamp of the same text would compute. chunk_doc now takes an optional `code_index` (pfs.anchors.code_family_ index(families)) and threads it into every anchor_metadata call, on both the section-chunking and paragraph-chunking paths. llm.index. index_refs builds one such index per run — hand families plus pfs.families.refresh_from against the DuckDB replica when it exists, the same mechanism `stack llm restamp` already used — and passes it to every chunk_doc call for the run. Corrected pfs/anchors.py's docstring, which claimed the two paths already wrote identical values.
297 lines
9.3 KiB
Python
297 lines
9.3 KiB
Python
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 _maybe_chunk_doc(fn):
|
|
"""Patch llm.index.chunk_doc only when a test asks for it."""
|
|
from contextlib import nullcontext
|
|
|
|
return nullcontext() if fn is None else patch("llm.index.chunk_doc", side_effect=fn)
|
|
|
|
|
|
def _run(
|
|
refs,
|
|
state_rows,
|
|
*,
|
|
force=False,
|
|
sealed=None,
|
|
mark_complete=True,
|
|
complete_rows=(),
|
|
chunk_doc=None,
|
|
code_index=None,
|
|
):
|
|
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,
|
|
# Real `_code_index` opens the live DuckDB replica — never in a
|
|
# unit test; every test here gets a fixed (default empty) index
|
|
# unless it asks for one.
|
|
patch("llm.index._code_index", return_value=code_index or {}),
|
|
_maybe_chunk_doc(chunk_doc),
|
|
):
|
|
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
|
|
|
|
|
|
def _state_inserts(conn):
|
|
return [
|
|
c
|
|
for c in conn.execute.call_args_list
|
|
if "INSERT INTO index_state" in str(c.args[0])
|
|
]
|
|
|
|
|
|
def _docket_inserts(conn):
|
|
return [
|
|
c
|
|
for c in conn.execute.call_args_list
|
|
if "INSERT INTO index_docket_state" in str(c.args[0])
|
|
]
|
|
|
|
|
|
def test_empty_text_ref_is_stamped_final_and_counts_toward_completion():
|
|
"""No text is a finished state, not a failure: stamp it so the next
|
|
run fingerprint-skips it, and let the docket still seal complete."""
|
|
blank = DocRef(
|
|
key="K9", collection="comments", docket="D", fingerprint="x", load=lambda: None
|
|
)
|
|
stats, store, conn, _ = _run([blank], [], sealed={"D": "s"})
|
|
assert stats["skipped"] == 1 and stats["indexed"] == 0
|
|
store.add_embeddings.assert_not_called()
|
|
ins = _state_inserts(conn)
|
|
assert len(ins) == 1
|
|
assert ins[0].args[1] == {"k": "K9", "c": "comments", "h": "", "n": 0, "f": "x"}
|
|
assert stats["docket_complete"] == 1
|
|
assert _docket_inserts(conn)
|
|
|
|
|
|
def test_stamped_empty_row_is_fingerprint_skipped_on_the_next_run():
|
|
loads = []
|
|
blank = DocRef(
|
|
key="K9",
|
|
collection="comments",
|
|
docket="D",
|
|
fingerprint="x",
|
|
load=lambda: loads.append("K9"),
|
|
)
|
|
stats, _, conn, _ = _run([blank], [("K9", "", "x")])
|
|
assert stats["fingerprint_skipped"] == 1 and stats["skipped"] == 0
|
|
assert loads == [] # stamped-empty rows are never re-loaded
|
|
assert _state_inserts(conn) == []
|
|
|
|
|
|
def test_sealed_docket_not_marked_when_a_ref_yields_no_chunks():
|
|
"""Text that chunks to nothing is a real failure — it blocks the seal
|
|
and is not stamped."""
|
|
stats, _, conn, _ = _run(
|
|
[_ref("fp1")], [], sealed={"D": "s"}, chunk_doc=lambda d, **k: []
|
|
)
|
|
assert stats["skipped"] == 1 and stats["docket_complete"] == 0
|
|
assert _state_inserts(conn) == []
|
|
assert _docket_inserts(conn) == []
|
|
|
|
|
|
# ── F1: indexer and restamp must write the same `families` ──────────
|
|
|
|
|
|
def test_index_refs_threads_code_index_into_chunk_doc():
|
|
"""`index_refs` passes the same `code_index` it built to every
|
|
`chunk_doc` call — a code that maps to a CPT-derived family in the
|
|
index built for this run must show up in the chunk's metadata."""
|
|
captured = []
|
|
|
|
def fake_chunk_doc(doc, **kwargs):
|
|
from llm.chunk import Chunk
|
|
|
|
captured.append(kwargs.get("code_index"))
|
|
return [Chunk(id="c1", text=doc.text, metadata={})]
|
|
|
|
fixed_index = {"99490": ("CCM",)}
|
|
stats, store, conn, enrich = _run(
|
|
[_ref("fp2")],
|
|
[("K1", "stale", "fp1")],
|
|
chunk_doc=fake_chunk_doc,
|
|
code_index=fixed_index,
|
|
)
|
|
assert stats["indexed"] == 1
|
|
assert captured == [fixed_index]
|
|
|
|
|
|
def test_code_index_builds_once_per_run_via_refresh_from_replica():
|
|
"""`llm.index._code_index` refreshes from the DuckDB replica when it
|
|
exists — the same mechanism `stack llm restamp` uses — before
|
|
building `code_family_index`, so both paths compute `families` from
|
|
the same snapshot (Ruling F1)."""
|
|
import llm.index as index_mod
|
|
|
|
calls = []
|
|
|
|
class FakeCon:
|
|
def close(self):
|
|
calls.append("closed")
|
|
|
|
def fake_connect(path, read_only):
|
|
calls.append(("connect", path, read_only))
|
|
return FakeCon()
|
|
|
|
def fake_refresh_from(con):
|
|
calls.append("refreshed")
|
|
return 0
|
|
|
|
with (
|
|
patch("llm.index.os.path.exists", return_value=True),
|
|
patch("duckdb.connect", side_effect=fake_connect),
|
|
patch("pfs.families.refresh_from", side_effect=fake_refresh_from),
|
|
):
|
|
idx = index_mod._code_index(CFG)
|
|
|
|
assert calls == [("connect", CFG.duckdb_replica, True), "refreshed", "closed"]
|
|
assert isinstance(idx, dict)
|
|
|
|
|
|
def test_code_index_skips_refresh_when_no_replica_file():
|
|
"""No replica on disk (a fresh checkout, or replica=False) — the
|
|
index falls back to hand families only, no DuckDB touched."""
|
|
import llm.index as index_mod
|
|
|
|
with (
|
|
patch("llm.index.os.path.exists", return_value=False),
|
|
patch("pfs.families.refresh_from") as refresh,
|
|
):
|
|
idx = index_mod._code_index(CFG)
|
|
|
|
refresh.assert_not_called()
|
|
assert isinstance(idx, dict)
|