51a5abf3bc
The code-cited query joined every embedding to langchain_pg_collection before filtering, and matched codes with no index to help it; it now filters on a scalar collection_id subquery, and ensure_hnsw builds a GIN index over the codes array plus one on collection_id. The DuckDB replica is opened once and kept, reopened only when publish_replica swaps the file (mtime_ns). Prompt size: at most two vintages once a question touches more than three codes, and code-cited snippets cap at 250 chars.
53 lines
2.1 KiB
Python
53 lines
2.1 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
|
|
# 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_collection" in executed
|
|
and "langchain_pg_embedding (collection_id)" in executed
|
|
)
|
|
|
|
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
|