151 lines
5.4 KiB
Python
151 lines
5.4 KiB
Python
"""llm.index — incremental, resumable embedding indexer."""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from sqlalchemy.exc import ProgrammingError
|
|
|
|
from llm.chunk import Doc, content_hash
|
|
from llm.config import LlmConfig
|
|
from llm.index import index_docs
|
|
|
|
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 _run(docs, state_rows, force=False, cfg=CFG):
|
|
"""Run index_docs with everything external mocked; return mocks."""
|
|
store = MagicMock()
|
|
engine = MagicMock()
|
|
conn = engine.begin.return_value.__enter__.return_value
|
|
conn.execute.return_value.fetchall.return_value = state_rows
|
|
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") as mock_ensure_hnsw,
|
|
patch("llm.index.HostPool") as MockPool,
|
|
):
|
|
MockPool.return_value.check.return_value = ["http://h1:11434"]
|
|
stats = index_docs(
|
|
docs,
|
|
collection="comments",
|
|
cfg=cfg,
|
|
pool=MockPool.return_value,
|
|
force=force,
|
|
)
|
|
return stats, store, conn, mock_ensure_hnsw
|
|
|
|
|
|
class TestIndexDocs:
|
|
def test_new_doc_embedded_and_recorded(self):
|
|
stats, store, conn, _ensure_hnsw = _run([DOC], state_rows=[])
|
|
assert stats == {"indexed": 1, "skipped": 0, "chunks": 1}
|
|
store.add_embeddings.assert_called_once()
|
|
kwargs = store.add_embeddings.call_args.kwargs
|
|
assert kwargs["ids"][0].startswith("K1:")
|
|
|
|
def test_unchanged_doc_skipped(self):
|
|
h = content_hash(DOC.text)
|
|
stats, store, _, _ensure_hnsw = _run([DOC], state_rows=[("K1", h)])
|
|
assert stats["skipped"] == 1
|
|
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")])
|
|
assert stats["indexed"] == 1
|
|
deletes = [
|
|
c
|
|
for c in conn.execute.call_args_list
|
|
if "DELETE FROM langchain_pg_embedding" in str(c.args[0])
|
|
]
|
|
assert len(deletes) == 1
|
|
sql = str(deletes[0].args[0])
|
|
assert "item_key" in sql # DELETE ... cmetadata->>'item_key'
|
|
# Scoped to the target collection, not a global delete by item_key.
|
|
assert "collection_id" in sql
|
|
assert "langchain_pg_collection" in sql
|
|
params = deletes[0].args[1]
|
|
assert params["k"] == "K1"
|
|
assert params["c"] == "comments"
|
|
|
|
def test_force_reembeds_unchanged(self):
|
|
h = content_hash(DOC.text)
|
|
stats, store, _, _ensure_hnsw = _run([DOC], state_rows=[("K1", h)], force=True)
|
|
assert stats["indexed"] == 1
|
|
|
|
def test_empty_doc_counts_skipped(self):
|
|
empty = Doc(key="K2", text=" ", metadata={})
|
|
stats, store, _, _ensure_hnsw = _run([empty], state_rows=[])
|
|
assert stats == {"indexed": 0, "skipped": 1, "chunks": 0}
|
|
|
|
|
|
class TestDeleteOldChunksFallback:
|
|
def test_programming_error_swallowed_and_indexing_continues(self):
|
|
"""A fresh DB where langchain_pg_embedding doesn't exist yet raises
|
|
ProgrammingError on DELETE; _delete_old_chunks must swallow it and
|
|
indexing must still proceed to add_embeddings + record state."""
|
|
store = MagicMock()
|
|
engine = MagicMock()
|
|
conn = engine.begin.return_value.__enter__.return_value
|
|
|
|
def fake_execute(clause, *args, **kwargs):
|
|
sql = str(clause)
|
|
if "DELETE FROM langchain_pg_embedding" in sql:
|
|
raise ProgrammingError("stmt", {}, Exception("relation missing"))
|
|
result = MagicMock()
|
|
result.fetchall.return_value = [] # no prior state
|
|
return result
|
|
|
|
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.HostPool") as MockPool,
|
|
):
|
|
MockPool.return_value.check.return_value = ["http://h1:11434"]
|
|
stats = index_docs(
|
|
[DOC],
|
|
collection="comments",
|
|
cfg=CFG,
|
|
pool=MockPool.return_value,
|
|
force=False,
|
|
)
|
|
|
|
assert stats == {"indexed": 1, "skipped": 0, "chunks": 1}
|
|
store.add_embeddings.assert_called_once()
|
|
|
|
|
|
class TestAnnIndexGating:
|
|
def test_build_ann_index_false_skips_ensure_hnsw(self):
|
|
_stats, _store, _conn, mock_ensure_hnsw = _run([DOC], state_rows=[], cfg=CFG)
|
|
mock_ensure_hnsw.assert_not_called()
|
|
|
|
def test_build_ann_index_true_calls_ensure_hnsw(self):
|
|
cfg = LlmConfig(
|
|
ollama_hosts=("http://h1:11434",),
|
|
embed_model="m",
|
|
instruct_model="g",
|
|
embed_dim=768,
|
|
build_ann_index=True,
|
|
pg_host="x",
|
|
pg_port=5432,
|
|
pg_db="llm",
|
|
pg_user="llm",
|
|
)
|
|
_stats, _store, _conn, mock_ensure_hnsw = _run([DOC], state_rows=[], cfg=cfg)
|
|
assert mock_ensure_hnsw.called
|