fix(llm): code-cited query filtered by collection + GIN index; cached replica handle; prompt caps (refs P48)

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.

src/llm/evidence.py

@@ -11,6 +11,7 @@ to "no evidence" with a warning — the chat never fails because of them.
 from __future__ import annotations
 
 import logging
+import os
 from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any, Sequence
@@ -85,6 +86,40 @@ class ValuationEvidence:
         }
 
 
+#: The open replica, keyed by (path, mtime_ns). DuckDB pays a real cost
+#: to open a database, and every chat turn that mentions a code would
+#: otherwise pay it. ``publish_replica`` swaps the file with os.replace,
+#: so the handle stays valid on the old inode until we notice the new
+#: mtime and reopen.
+_REPLICA: tuple[str, int, Any] | None = None
+
+
+def _mtime(path: str) -> int:
+    try:
+        return os.stat(path).st_mtime_ns
+    except OSError:
+        return -1
+
+
+def _connect(path: str) -> Any:
+    """The cached read-only connection to *path*, reopened when the
+    replica file has been republished since we opened it."""
+    global _REPLICA
+    mtime = _mtime(path)
+    if _REPLICA is not None:
+        cached_path, cached_mtime, con = _REPLICA
+        if (cached_path, cached_mtime) == (path, mtime):
+            return con
+        _REPLICA = None
+        try:
+            con.close()
+        except Exception as e:  # noqa: BLE001 — a stale handle must not break the chat
+            log.warning("stale replica handle not closed: %s", e)
+    con = duckdb.connect(path, read_only=True)
+    _REPLICA = (path, mtime, con)
+    return con
+
+
 def _replica_path(cfg: LlmConfig) -> str:
     p = Path(cfg.duckdb_replica)
     if p.is_absolute():
@@ -100,17 +135,18 @@ def valuation_evidence(question: str, cfg: LlmConfig) -> ValuationEvidence | Non
         return None
     path = _replica_path(cfg)
     try:
-        con = duckdb.connect(path, read_only=True)
+        con = _connect(path)
     except Exception as e:  # noqa: BLE001 — duckdb raises several types
         log.warning("valuation evidence skipped (%s): %s", path, e)
         return None
+    # A whole family is six codes; showing four vintages of each would
+    # crowd the excerpts out of the prompt.
+    years = cfg.valuation_years if len(det.codes) <= 3 else 2
     try:
-        rows, unpriced = valuation(con, list(det.codes), years=cfg.valuation_years)
+        rows, unpriced = valuation(con, list(det.codes), years=years)
     except Exception as e:  # noqa: BLE001
         log.warning("valuation evidence skipped (query): %s", e)
         return None
-    finally:
-        con.close()
     return ValuationEvidence(det.codes, det.families, tuple(rows), tuple(unpriced))
 
 
@@ -122,10 +158,10 @@ _CITED_SQL = text(
     "ORDER BY e.cmetadata->>'date' DESC NULLS LAST, e.cmetadata->>'item_key', e.cmetadata->>'p_id' "
     ") AS rn "
     "FROM langchain_pg_embedding e "
-    "JOIN langchain_pg_collection c ON c.uuid = e.collection_id "
     "JOIN unnest(CAST(:codes AS text[])) AS w(code) "
     "ON w.code = ANY(string_to_array(COALESCE(e.cmetadata->>'codes', ''), ' ')) "
-    "WHERE c.name = :collection "
+    "WHERE e.collection_id = (SELECT uuid FROM langchain_pg_collection "
+    "WHERE name = :collection) "
     ") t "
     "WHERE rn <= :window "
     "ORDER BY cmetadata->>'date' DESC NULLS LAST, cmetadata->>'item_key', cmetadata->>'p_id'"
@@ -170,7 +206,7 @@ def code_cited_sources(
         seen.add(key)
         for c in hit_codes:
             count[c] += 1
-        out.append(as_source(md, document, 0.0))
+        out.append(as_source(md, document[:250], 0.0))
     return out
 
 

src/llm/migrate.py

@@ -48,6 +48,14 @@ _HNSW_DDL = [
     "ALTER TABLE langchain_pg_embedding ALTER COLUMN embedding TYPE vector({dim})",
     "CREATE INDEX IF NOT EXISTS ix_embedding_hnsw "
     "ON langchain_pg_embedding USING hnsw (embedding vector_cosine_ops)",
+    # Code-cited sources (llm.evidence) match a code against the space-
+    # separated cmetadata->>'codes' inside one collection; without these
+    # two the query is a sequential scan of every embedding.
+    "CREATE INDEX IF NOT EXISTS ix_embedding_codes "
+    "ON langchain_pg_embedding USING gin "
+    "(string_to_array(coalesce(cmetadata->>'codes',''), ' '))",
+    "CREATE INDEX IF NOT EXISTS ix_embedding_collection "
+    "ON langchain_pg_embedding (collection_id)",
 ]
 
 

tests/llm/test_evidence.py

@@ -4,6 +4,9 @@ from __future__ import annotations
 
 from unittest.mock import MagicMock, patch
 
+import pytest
+
+from llm import evidence
 from llm.config import LlmConfig
 from llm.evidence import (
     ValuationEvidence,
@@ -136,6 +139,12 @@ class TestEvidenceText:
 
 
 class TestValuationEvidence:
+    @pytest.fixture(autouse=True)
+    def _clear_replica_cache(self):
+        evidence._REPLICA = None
+        yield
+        evidence._REPLICA = None
+
     def test_none_when_no_codes(self):
         assert (
             valuation_evidence("what did commenters say about telehealth?", CFG) is None
@@ -155,13 +164,42 @@ class TestValuationEvidence:
             and ev.codes == ("G0556", "G0557", "G0558")
             and ev.families == ("APCM",)
         )
-        mock_connect.return_value.close.assert_called_once()
+        mock_connect.return_value.close.assert_not_called()  # cached, not closed
 
     @patch("llm.evidence.duckdb.connect", side_effect=OSError("no replica"))
     def test_unopenable_replica_yields_none(self, _c, caplog):
         assert valuation_evidence("G0556?", CFG) is None
         assert "valuation evidence skipped" in caplog.text
 
+    @patch("llm.evidence.valuation", return_value=([_row()], []))
+    @patch("llm.evidence.duckdb.connect")
+    def test_replica_connection_is_reused(self, mock_connect, _val):
+        valuation_evidence("G0556?", CFG)
+        valuation_evidence("G0557?", CFG)
+        mock_connect.assert_called_once()
+        mock_connect.return_value.close.assert_not_called()
+
+    @patch("llm.evidence.valuation", return_value=([_row()], []))
+    @patch("llm.evidence.duckdb.connect")
+    def test_republished_replica_reopens_and_closes_the_old_handle(
+        self, mock_connect, _val, monkeypatch
+    ):
+        first, second = MagicMock(), MagicMock()
+        mock_connect.side_effect = [first, second]
+        mtimes = iter([1, 2])
+        monkeypatch.setattr(evidence, "_mtime", lambda _p: next(mtimes))
+        valuation_evidence("G0556?", CFG)
+        valuation_evidence("G0557?", CFG)
+        assert mock_connect.call_count == 2
+        first.close.assert_called_once()
+        second.close.assert_not_called()
+
+    @patch("llm.evidence.valuation", return_value=([], []))
+    @patch("llm.evidence.duckdb.connect")
+    def test_many_codes_shorten_the_window(self, _c, mock_val):
+        valuation_evidence("How is CCM valued?", CFG)  # six-code family
+        assert mock_val.call_args.kwargs == {"years": 2}
+
 
 def _engine(rows):
     engine = MagicMock()
@@ -203,6 +241,9 @@ class TestCodeCitedSources:
         sql = str(conn.execute.call_args.args[0])
         assert "string_to_array" in sql and "langchain_pg_collection" in sql
         assert "PARTITION BY" in sql
+        # the collection is a scalar subquery, not a join to filter after
+        assert "JOIN langchain_pg_collection" not in sql
+        assert "e.collection_id = (SELECT uuid FROM langchain_pg_collection" in sql
         assert conn.execute.call_args.args[1]["collection"] == "rules"
         assert conn.execute.call_args.args[1]["codes"] == ["G0556", "G0557"]
         assert conn.execute.call_args.args[1]["window"] == 4
@@ -213,6 +254,11 @@ class TestCodeCitedSources:
         ]
         assert all(s["kind"] == "rule" and s["score"] == 0.0 for s in out)
 
+    def test_snippets_are_capped(self):
+        engine, _ = _engine([("x" * 900, _rule_md(1, "G0556"))])
+        (out,) = code_cited_sources(engine, ["G0556"], per_code=2)
+        assert len(out["snippet"]) == 250
+
     def test_engine_error_yields_empty(self, caplog):
         engine = MagicMock()
         engine.begin.side_effect = RuntimeError("pg down")

tests/llm/test_migrate.py

@@ -24,6 +24,15 @@ class TestDdl:
         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