Files
stack/tests/llm/test_migrate.py
kert 26f01bd0ba
Some checks failed
CI / lint (push) Successful in 58s
CI / notebooks-smoke (push) Successful in 1m48s
Deploy / notebooks (push) Has been skipped
Deploy / zotero (push) Has been skipped
CI / test (push) Failing after 2m58s
Deploy / docs (push) Has been skipped
Deploy / api (push) Has been skipped
Deploy / llm (push) Has been skipped
Deploy / mc (push) Has been skipped
Infra CI / zotero (push) Successful in 25s
Infra CI / notebooks (push) Successful in 58s
Infra CI / docs (push) Successful in 1m32s
Infra CI / api (push) Successful in 1m11s
Infra CI / mc (push) Failing after 41s
Infra CI / llm (push) Successful in 1m9s
Deploy / report (push) Successful in 19s
fix(llm): filtered vector search returns results — hnsw.iterative_scan = relaxed_order on every connection and at the database level
pgvector's HNSW scan collects ef_search (40) nearest candidates and
only then applies the WHERE clause, so a selective metadata filter
(one theme inside one doctype; one docket's comments) could drop every
candidate and answer nothing. pgvector 0.8's iterative scan keeps
walking the graph until the LIMIT is met. The engine now SETs
hnsw.iterative_scan = relaxed_order on each pooled connection (a
placeholder SET needs no privilege) and PGVector is handed that engine
instead of a URL so its sessions inherit it; migrate() also tries
ALTER DATABASE (best effort — the llm role may not; applied once as
the superuser on the live database).
2026-09-24 20:07:38 -04:00

108 lines
4.2 KiB
Python

"""llm.migrate — idempotent DDL for the llm database."""
from unittest.mock import MagicMock
from llm import migrate
class TestDdl:
def test_index_state_ddl_is_idempotent_sql(self):
assert "CREATE TABLE IF NOT EXISTS index_state" in migrate.INDEX_STATE_DDL
assert "PRIMARY KEY (item_key, collection)" in migrate.INDEX_STATE_DDL
def test_migrate_executes_ddl(self):
engine = MagicMock()
conn = engine.begin.return_value.__enter__.return_value
migrate.migrate(engine)
executed = " ".join(str(call.args[0]) for call in conn.execute.call_args_list)
assert "index_state" in executed
def test_ensure_hnsw_types_column_then_indexes(self):
engine = MagicMock()
conn = engine.begin.return_value.__enter__.return_value
migrate.ensure_hnsw(engine, 768)
executed = " ".join(str(call.args[0]) for call in conn.execute.call_args_list)
assert "vector(768)" in executed
assert "USING hnsw" in executed
assert (
"ix_embedding_collection" in executed
and "langchain_pg_embedding (collection_id)" in executed
)
# F6: the metadata GIN indexes are ensure_metadata_indexes' job now,
# not ensure_hnsw's — this must never take the ACCESS EXCLUSIVE lock.
assert "ix_embedding_codes" not in executed
assert "ix_embedding_families" not in executed
assert "ix_embedding_elements" not in executed
def test_ensure_metadata_indexes_issues_exactly_the_three_gin_statements(self):
engine = MagicMock()
conn = engine.begin.return_value.__enter__.return_value
migrate.ensure_metadata_indexes(engine)
calls = [str(call.args[0]) for call in conn.execute.call_args_list]
assert len(calls) == 3
executed = " ".join(calls)
# code-cited lookups scan cmetadata->>'codes' inside one collection
assert (
"ix_embedding_codes" in executed
and "USING gin (string_to_array" in executed
)
assert (
"ix_embedding_families" in executed and "cmetadata->>'families'" in executed
)
assert (
"ix_embedding_elements" in executed and "cmetadata->>'elements'" in executed
)
# No vector-column ALTER, no HNSW build — this must be safe on a
# live table.
assert "ALTER TABLE" not in executed
assert "hnsw" not in executed.lower()
def test_fingerprint_column_and_docket_state(self):
assert "fingerprint" in migrate.INDEX_STATE_DDL
assert "ADD COLUMN IF NOT EXISTS fingerprint" in migrate.INDEX_STATE_ALTER
assert (
"CREATE TABLE IF NOT EXISTS index_docket_state"
in migrate.INDEX_DOCKET_STATE_DDL
)
assert "PRIMARY KEY (collection, docket)" in migrate.INDEX_DOCKET_STATE_DDL
def test_migrate_runs_alter_and_docket_state(self):
engine = MagicMock()
conn = engine.begin.return_value.__enter__.return_value
migrate.migrate(engine)
executed = " ".join(str(call.args[0]) for call in conn.execute.call_args_list)
assert "ADD COLUMN IF NOT EXISTS fingerprint" in executed
assert "index_docket_state" in executed
class TestIterativeScan:
def test_migrate_sets_iterative_scan_on_the_database(self):
from unittest.mock import MagicMock
from llm.migrate import migrate
engine = MagicMock()
conn = engine.begin.return_value.__enter__.return_value
conn.execute.return_value.scalar.return_value = "llm"
migrate(engine)
sql = [str(c.args[0]) for c in conn.execute.call_args_list]
assert any(
"ALTER DATABASE llm SET hnsw.iterative_scan = 'relaxed_order'" in q
for q in sql
)
def test_failure_is_a_warning_not_an_error(self, caplog):
from unittest.mock import MagicMock
from llm.migrate import ensure_iterative_scan
engine = MagicMock()
engine.begin.return_value.__enter__.return_value.execute.side_effect = (
RuntimeError("permission denied")
)
import logging
caplog.set_level(logging.INFO, logger="llm.migrate")
ensure_iterative_scan(engine) # no raise
assert "hnsw.iterative_scan not set" in caplog.text