From 26f01bd0ba8dd206be38e777f4d0a90b2c8bd447 Mon Sep 17 00:00:00 2001 From: kert Date: Thu, 24 Sep 2026 20:07:38 -0400 Subject: [PATCH] =?UTF-8?q?fix(llm):=20filtered=20vector=20search=20return?= =?UTF-8?q?s=20results=20=E2=80=94=20hnsw.iterative=5Fscan=20=3D=20relaxed?= =?UTF-8?q?=5Forder=20on=20every=20connection=20and=20at=20the=20database?= =?UTF-8?q?=20level?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- src/llm/index.py | 30 ++++++++++++++++++++++++++++-- src/llm/migrate.py | 29 +++++++++++++++++++++++++++++ tests/llm/test_index.py | 23 +++++++++++++++++++++++ tests/llm/test_migrate.py | 32 ++++++++++++++++++++++++++++++++ 4 files changed, 112 insertions(+), 2 deletions(-) diff --git a/src/llm/index.py b/src/llm/index.py index cee9b91..70d7bdf 100644 --- a/src/llm/index.py +++ b/src/llm/index.py @@ -60,17 +60,43 @@ def _code_index(cfg: LlmConfig) -> dict[str, tuple[str, ...]]: return code_family_index(FAMILIES) +ITERATIVE_SCAN_SQL = "SET hnsw.iterative_scan = relaxed_order" + + +def set_iterative_scan(dbapi_conn: Any, _record: Any = None) -> None: + """Per-connection ``hnsw.iterative_scan = relaxed_order`` (pgvector ≥ + 0.8): an HNSW scan otherwise stops after ef_search candidates and a + selective metadata filter (one theme in one doctype, one docket) can + drop every one of them and answer nothing. A SET of the placeholder + needs no privilege, unlike ALTER DATABASE; an older pgvector that + rejects it only costs a debug line.""" + try: + cur = dbapi_conn.cursor() + try: + cur.execute(ITERATIVE_SCAN_SQL) + finally: + cur.close() + except Exception as exc: # noqa: BLE001 — never block a connection on a tuning knob + log.debug("hnsw.iterative_scan not set on this connection: %s", exc) + + def _engine(cfg: LlmConfig) -> Engine: # pragma: no cover — needs a live DB - return create_engine(pg_url(cfg)) + from sqlalchemy import event + + engine = create_engine(pg_url(cfg)) + event.listen(engine, "connect", set_iterative_scan) + return engine def vectorstore(collection: str, cfg: LlmConfig, pool: HostPool): from langchain_postgres import PGVector # pragma: no cover — needs live DB + # Hand PGVector our engine (not the URL) so its sessions come from + # the pool that sets hnsw.iterative_scan on connect. return PGVector( # pragma: no cover — needs a live pgvector DB embeddings=PoolEmbeddings(pool, cfg.embed_model), collection_name=collection, - connection=pg_url(cfg), + connection=_engine(cfg), embedding_length=cfg.embed_dim, use_jsonb=True, ) diff --git a/src/llm/migrate.py b/src/llm/migrate.py index b5a76e0..5662527 100644 --- a/src/llm/migrate.py +++ b/src/llm/migrate.py @@ -76,12 +76,41 @@ _METADATA_INDEX_DDL = [ ] +# pgvector's HNSW scan collects ef_search (40) nearest candidates and only +# then applies a WHERE clause, so a selective metadata filter (one theme +# inside one doctype, one docket's comments) can drop every candidate and +# return nothing. pgvector ≥ 0.8 keeps scanning until the LIMIT is met +# when iterative_scan is on; set at the database level so every session +# — the API's pooled connections included — inherits it without having +# to load the extension's GUCs first. +ITERATIVE_SCAN_DDL = "ALTER DATABASE {db} SET hnsw.iterative_scan = 'relaxed_order'" + + def migrate(engine: Engine) -> None: """Create llm-owned tables/columns. Safe to run on every start.""" with engine.begin() as conn: conn.execute(text(INDEX_STATE_DDL)) conn.execute(text(INDEX_STATE_ALTER)) conn.execute(text(INDEX_DOCKET_STATE_DDL)) + ensure_iterative_scan(engine) + + +def ensure_iterative_scan(engine: Engine) -> None: + """Database-level ``hnsw.iterative_scan = relaxed_order``; a warning, + not a failure, when the role may not alter the database or the + pgvector on hand predates the setting.""" + import logging + + try: + with engine.begin() as conn: + db = conn.execute(text("SELECT current_database()")).scalar() + conn.execute(text(ITERATIVE_SCAN_DDL.format(db=db))) + except Exception as exc: # noqa: BLE001 — the llm role usually may not ALTER DATABASE + # every connection also SETs it itself (llm.index.set_iterative_scan), + # so this is belt-and-braces for sessions opened outside our engine + logging.getLogger(__name__).info( + "hnsw.iterative_scan not set database-wide: %s", exc + ) def ensure_hnsw(engine: Engine, dim: int) -> None: diff --git a/tests/llm/test_index.py b/tests/llm/test_index.py index 91beb6f..29be248 100644 --- a/tests/llm/test_index.py +++ b/tests/llm/test_index.py @@ -239,3 +239,26 @@ class TestPruneRulesFromCorpus: assert prune_rules_from_corpus(engine, store)["rule_items"] == 0 engine.begin.assert_not_called() store.close() + + +class TestIterativeScanOnConnect: + def test_sets_the_guc_on_each_new_connection(self): + from unittest.mock import MagicMock + + from llm.index import ITERATIVE_SCAN_SQL, set_iterative_scan + + conn = MagicMock() + set_iterative_scan(conn, None) + conn.cursor.return_value.execute.assert_called_once_with(ITERATIVE_SCAN_SQL) + conn.cursor.return_value.close.assert_called_once() + + def test_failure_never_propagates(self): + from unittest.mock import MagicMock + + from llm.index import set_iterative_scan + + conn = MagicMock() + conn.cursor.return_value.execute.side_effect = RuntimeError( + "unrecognized parameter" + ) + set_iterative_scan(conn, None) # older pgvector: no raise diff --git a/tests/llm/test_migrate.py b/tests/llm/test_migrate.py index 680018c..00435ad 100644 --- a/tests/llm/test_migrate.py +++ b/tests/llm/test_migrate.py @@ -73,3 +73,35 @@ class TestDdl: 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