feat(llm): code-cited and family-cited sources from rules, comments and corpus (refs #688)

src/llm/config.py

@@ -47,6 +47,7 @@ class LlmConfig:
     duckdb_replica: str = "data/replica/aco.ro.duckdb"
     valuation_years: int = 4
     code_cited_per_code: int = 3
+    code_cited_collections: tuple[str, ...] = ("rules", "comments", "corpus")
 
 
 def parse_hosts(spec: str) -> tuple[tuple[str, float], ...]:
@@ -105,6 +106,9 @@ def load() -> LlmConfig:
         or str(_opt(section, "duckdb_replica", "data/replica/aco.ro.duckdb")),
         valuation_years=int(_opt(section, "valuation_years", 4)),
         code_cited_per_code=int(_opt(section, "code_cited_per_code", 3)),
+        code_cited_collections=tuple(
+            _opt(section, "code_cited_collections", ["rules", "comments", "corpus"])
+        ),
         embed_dim=int(section.embed_dim),
         build_ann_index=build_ann_index,
         pg_host=os.environ.get("LLM_PG_HOST", str(section.pg_host)),

src/llm/evidence.py

@@ -3,9 +3,12 @@
 ``valuation_evidence`` detects HCPCS/CPT codes (and registered families)
 in the question, reads their RVUs from the DuckDB read replica and
 attaches conversion factors from ``pfs.rules``. ``code_cited_sources``
-pulls rule chunks whose ``codes`` metadata mention those codes straight
-from pgvector so the answer can cite where CMS valued them. Both degrade
-to "no evidence" with a warning — the chat never fails because of them.
+pulls chunks — from rules, comments and the reference corpus — whose
+``codes`` metadata mention those codes, and (when a family was detected)
+chunks whose ``families`` metadata mentions the family key, straight
+from pgvector so the answer can cite where CMS valued them and where the
+library discusses them. Both degrade to "no evidence" with a warning —
+the chat never fails because of them.
 """
 
 from __future__ import annotations
@@ -181,46 +184,120 @@ _CITED_SQL = text(
     "ORDER BY cmetadata->>'date' DESC NULLS LAST, cmetadata->>'item_key', cmetadata->>'p_id'"
 )
 
+#: Same shape as ``_CITED_SQL`` but matches on the chunk's ``families``
+#: metadata via array overlap — this is how a CCM question pulls the CPT
+#: manual's "Chronic Care Management Services" guideline chunk even though
+#: that chunk never spells out 99490.
+_FAMILY_CITED_SQL = text(
+    "SELECT document, cmetadata FROM ( "
+    "SELECT e.document, e.cmetadata, "
+    "row_number() OVER ( "
+    "PARTITION BY w.family "
+    "ORDER BY e.cmetadata->>'date' DESC NULLS LAST, e.cmetadata->>'item_key', e.cmetadata->>'p_id' "
+    ") AS rn "
+    "FROM langchain_pg_embedding e "
+    "JOIN unnest(CAST(:families AS text[])) AS w(family) "
+    "ON string_to_array(COALESCE(e.cmetadata->>'families', ''), ' ') && ARRAY[w.family] "
+    "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'"
+)
 
-def code_cited_sources(
-    engine: Any, codes: Sequence[str], *, per_code: int, collection: str = "rules"
+
+def _dedupe_key(md: dict[str, str]) -> tuple[str, str]:
+    """(item_key, p_id) for rule chunks — several paragraphs share an
+    item_key — else (item_key, seq), since comment/corpus chunks carry a
+    sequence number instead of a paragraph id."""
+    second = md.get("p_id", "") if md.get("kind") == "rule" else md.get("seq", "")
+    return (md.get("item_key", ""), second)
+
+
+def _collect(
+    engine: Any,
+    sql: Any,
+    key: str,
+    wanted: Sequence[str],
+    *,
+    per: int,
+    collections: Sequence[str],
+    seen: set[tuple[str, str]],
 ) -> list[dict]:
-    """Chunks in *collection* whose ``codes`` metadata mention any of
-    *codes* — at most *per_code* per code, newest rule first, deduped on
-    (item_key, p_id). Scores are 0.0: these are additive, not ranked."""
-    wanted = [c.upper() for c in codes]
-    if not wanted:
-        return []
-    try:
-        with engine.begin() as conn:
-            rows = conn.execute(
-                _CITED_SQL,
-                {
-                    "collection": collection,
-                    "codes": wanted,
-                    "window": per_code * 2,
-                },
-            ).fetchall()
-    except Exception as e:  # noqa: BLE001
-        log.warning("code-cited sources skipped: %s", e)
+    """Run *sql* once per collection in *collections*, capping hits per
+    *key* value (a code or a family) per collection, deduping across the
+    whole call via the shared *seen* set. *key* names both the query
+    parameter (``:codes``/``:families``) and the chunk metadata field."""
+    upper = [w.upper() for w in wanted]
+    if not upper:
         return []
     out: list[dict] = []
+    for collection in collections:
+        try:
+            with engine.begin() as conn:
+                rows = conn.execute(
+                    sql,
+                    {"collection": collection, key: upper, "window": per * 2},
+                ).fetchall()
+        except Exception as e:  # noqa: BLE001
+            log.warning("code-cited sources skipped (%s/%s): %s", collection, key, e)
+            continue
+        count: dict[str, int] = {w: 0 for w in upper}
+        for document, md in rows:
+            md = {k: str(v) for k, v in (md or {}).items()}
+            dkey = _dedupe_key(md)
+            if dkey in seen:
+                continue
+            hits = [w for w in md.get(key, "").split() if w in count and count[w] < per]
+            if not hits:
+                continue
+            seen.add(dkey)
+            for w in hits:
+                count[w] += 1
+            out.append(as_source(md, document[:250], 0.0))
+    return out
+
+
+def code_cited_sources(
+    engine: Any,
+    codes: Sequence[str],
+    *,
+    per_code: int,
+    collections: Sequence[str] = ("rules",),
+    families: Sequence[str] = (),
+) -> list[dict]:
+    """Chunks across *collections* whose ``codes`` metadata mention any of
+    *codes* — at most *per_code* per code per collection, newest first,
+    deduped on (item_key, p_id) for rules or (item_key, seq) otherwise.
+    Concatenated in the order *collections* is given.
+
+    When *families* is non-empty, chunks whose ``families`` metadata
+    mentions one of them are matched too (same per-collection cap) and
+    appended after the code-cited rows — skipped entirely, no query, when
+    *families* is empty. Scores are 0.0: these are additive, not ranked.
+    """
+    if not codes and not families:
+        return []
     seen: set[tuple[str, str]] = set()
-    count: dict[str, int] = {c: 0 for c in wanted}
-    for document, md in rows:
-        md = {k: str(v) for k, v in (md or {}).items()}
-        key = (md.get("item_key", ""), md.get("p_id", ""))
-        if key in seen:
-            continue
-        hit_codes = [
-            c for c in md.get("codes", "").split() if c in count and count[c] < per_code
-        ]
-        if not hit_codes:
-            continue
-        seen.add(key)
-        for c in hit_codes:
-            count[c] += 1
-        out.append(as_source(md, document[:250], 0.0))
+    out = _collect(
+        engine,
+        _CITED_SQL,
+        "codes",
+        codes,
+        per=per_code,
+        collections=collections,
+        seen=seen,
+    )
+    if families:
+        out += _collect(
+            engine,
+            _FAMILY_CITED_SQL,
+            "families",
+            families,
+            per=per_code,
+            collections=collections,
+            seen=seen,
+        )
     return out
 
 

src/llm/rag.py

@@ -9,9 +9,9 @@ source gets a deep link (``llm.links``). Generation streams from the
 largest live Ollama host (``HostPool.acquire_generation``) with the
 model tier that host can hold (``pick_model``). When the question
 cites HCPCS/CPT codes, ``llm.evidence`` supplies authoritative
-valuation rows and rule chunks that cite those codes; the SSE event
-order is ``valuation`` (only when evidence is found) → ``token``* →
-``sources`` → ``done``.
+valuation rows and code- or family-cited chunks from rules, comments and
+the reference corpus; the SSE event order is ``valuation`` (only when
+evidence is found) → ``token``* → ``sources`` → ``done``.
 """
 
 from __future__ import annotations
@@ -139,8 +139,10 @@ def stream_answer(
     """Retrieve, then stream a grounded answer from the largest live host.
 
     When the question cites HCPCS/CPT codes, yields a ``valuation`` event
-    first (evidence's ``payload()``) and merges rule chunks that cite
-    those codes into ``sources``. Then yields ``{"type":"token","text":…}``
+    first (evidence's ``payload()``) and merges chunks from rules,
+    comments and the corpus that cite those codes — or, for a detected
+    family, its family key — into ``sources``. Then yields
+    ``{"type":"token","text":…}``
     events as the model generates, then one
     ``{"type":"sources", "sources": […], "model": …, "host": …}`` and a
     final ``{"type":"done"}``.
@@ -149,7 +151,11 @@ def stream_answer(
     evidence = valuation_evidence(question, cfg)
     if evidence is not None:
         cited = code_cited_sources(
-            _engine(cfg), evidence.codes, per_code=cfg.code_cited_per_code
+            _engine(cfg),
+            evidence.codes,
+            per_code=cfg.code_cited_per_code,
+            collections=tuple(cfg.code_cited_collections),
+            families=evidence.families,
         )
         sources = merge_sources(sources, cited)
         yield evidence.payload()

stack.toml

@@ -126,7 +126,8 @@ recency_weight = 0.3
 top_n = 8
 duckdb_replica = "data/replica/aco.ro.duckdb"   # read-only DuckDB replica the chat reads valuations from
 valuation_years = 4                      # final-rule vintages shown per code (plus the newest NPRM)
-code_cited_per_code = 3                  # rule paragraphs literally citing each detected code
+code_cited_per_code = 3                  # excerpts literally citing each detected code (or family)
+code_cited_collections = ["rules", "comments", "corpus"]  # searched in this order
 
 [llm.k_per_kind]                       # over-fetched ×3 per kind, then re-ranked
 comment = 8

tests/llm/test_config.py

@@ -94,6 +94,7 @@ class TestNewKnobs:
         assert cfg.duckdb_replica == "data/replica/aco.ro.duckdb"
         assert cfg.valuation_years == 4
         assert cfg.code_cited_per_code == 3
+        assert cfg.code_cited_collections == ("rules", "comments", "corpus")
 
     def test_duckdb_replica_env_override(self, monkeypatch):
         monkeypatch.setenv("LLM_DUCKDB_REPLICA", "/app/data/replica/aco.ro.duckdb")

tests/llm/test_evidence.py

@@ -348,6 +348,51 @@ def _rule_md(p_id, codes, date="2024-12-09"):
     }
 
 
+def _comment_md(seq, codes, item_key="C1", date="2024-12-09"):
+    return {
+        "kind": "comment",
+        "item_key": item_key,
+        "seq": str(seq),
+        "comment_id": "CMS-2026-2377-3438",
+        "date": date,
+        "codes": codes,
+        "families": "",
+    }
+
+
+def _corpus_md(seq, codes="", families="", item_key="CPT-ED", date="2024-12-09"):
+    return {
+        "kind": "corpus",
+        "item_key": item_key,
+        "seq": str(seq),
+        "url": "https://example.org/cpt.pdf",
+        "title": "CPT Manual",
+        "date": date,
+        "codes": codes,
+        "families": families,
+    }
+
+
+def _multi_engine(by_code=None, by_family=None):
+    """A fake engine whose rows depend on which collection *and* which of
+    the two SQL statements (codes vs. families) is being executed —
+    distinguished by which bind parameter the call carries."""
+    engine = MagicMock()
+    conn = engine.begin.return_value.__enter__.return_value
+    by_code = by_code or {}
+    by_family = by_family or {}
+
+    def _execute(_sql, params):
+        collection = params["collection"]
+        table = by_code if "codes" in params else by_family
+        result = MagicMock()
+        result.fetchall.return_value = table.get(collection, [])
+        return result
+
+    conn.execute.side_effect = _execute
+    return engine, conn
+
+
 class TestCodeCitedSources:
     def test_sql_and_per_code_cap_and_dedupe(self):
         engine, conn = _engine(
@@ -396,6 +441,108 @@ class TestCodeCitedSources:
         engine.begin.assert_not_called()
 
 
+class TestCodeCitedSourcesMultiCollection:
+    def test_concatenated_rules_then_comments_then_corpus(self):
+        engine, _conn = _multi_engine(
+            by_code={
+                "rules": [("rule chunk", _rule_md(1, "G0556"))],
+                "comments": [("comment chunk", _comment_md(1, "G0556"))],
+                "corpus": [("corpus chunk", _corpus_md(1, "G0556"))],
+            }
+        )
+        out = code_cited_sources(
+            engine, ["G0556"], per_code=2, collections=("rules", "comments", "corpus")
+        )
+        assert [s["kind"] for s in out] == ["rule", "comment", "corpus"]
+
+    def test_per_collection_cap_applies_independently(self):
+        engine, _conn = _multi_engine(
+            by_code={
+                "rules": [
+                    ("r1", _rule_md(1, "G0556")),
+                    ("r2", _rule_md(2, "G0556")),
+                    ("r3", _rule_md(3, "G0556")),  # over the per-collection cap
+                ],
+                "comments": [
+                    ("c1", _comment_md(1, "G0556")),
+                    ("c2", _comment_md(2, "G0556")),
+                ],
+            }
+        )
+        out = code_cited_sources(
+            engine, ["G0556"], per_code=2, collections=("rules", "comments")
+        )
+        assert [s["snippet"] for s in out] == ["r1", "r2", "c1", "c2"]
+
+    def test_dedupe_uses_seq_for_non_rule_kinds(self):
+        engine, _conn = _multi_engine(
+            by_code={
+                "comments": [
+                    ("c1", _comment_md(1, "G0556", item_key="C1")),
+                    ("c1 dup", _comment_md(1, "G0556", item_key="C1")),  # same seq
+                    ("c2", _comment_md(2, "G0556", item_key="C1")),  # different seq
+                ]
+            }
+        )
+        out = code_cited_sources(
+            engine, ["G0556"], per_code=5, collections=("comments",)
+        )
+        assert [s["snippet"] for s in out] == ["c1", "c2"]
+
+
+class TestFamilyCitedSources:
+    def test_skipped_entirely_when_no_families(self):
+        engine, conn = _multi_engine(by_code={"rules": [("r1", _rule_md(1, "G0556"))]})
+        code_cited_sources(engine, ["G0556"], per_code=2, collections=("rules",))
+        assert conn.execute.call_count == 1  # codes query only, no family round trip
+
+    def test_sql_shape_and_params(self):
+        engine, conn = _multi_engine(by_family={"rules": []})
+        code_cited_sources(
+            engine, [], per_code=2, collections=("rules",), families=["CCM"]
+        )
+        sql = str(conn.execute.call_args.args[0])
+        assert "&&" in sql and "PARTITION BY w.family" in sql
+        assert conn.execute.call_args.args[1] == {
+            "collection": "rules",
+            "families": ["CCM"],
+            "window": 4,
+        }
+
+    def test_family_rows_appended_after_code_rows(self):
+        engine, _conn = _multi_engine(
+            by_code={"rules": [("code hit", _rule_md(1, "G0556"))]},
+            by_family={
+                "corpus": [
+                    (
+                        "CCM guideline",
+                        _corpus_md(1, codes="", families="CCM", item_key="CPT-ED"),
+                    )
+                ]
+            },
+        )
+        out = code_cited_sources(
+            engine,
+            ["G0556"],
+            per_code=2,
+            collections=("rules", "corpus"),
+            families=["CCM"],
+        )
+        assert [s["snippet"] for s in out] == ["code hit", "CCM guideline"]
+        assert out[1]["kind"] == "corpus"
+
+    def test_family_dedupes_against_a_code_cited_chunk(self):
+        md = _rule_md(1, "G0556")
+        md["families"] = "APCM"
+        engine, _conn = _multi_engine(
+            by_code={"rules": [("shared", md)]}, by_family={"rules": [("shared", md)]}
+        )
+        out = code_cited_sources(
+            engine, ["G0556"], per_code=2, collections=("rules",), families=["APCM"]
+        )
+        assert [s["snippet"] for s in out] == ["shared"]
+
+
 class TestMergeSources:
     def test_keeps_order_and_dedupes_by_label(self):
         a = [{"label": "X", "score": 0.9}, {"label": "Y", "score": 0.8}]

tests/llm/test_rag.py

@@ -413,7 +413,11 @@ class TestStreamAnswer:
         assert events[1] == {"type": "token", "text": "x"}
         assert events[-2]["sources"] == [src, cited]
         mock_cited.assert_called_once_with(
-            mock_engine.return_value, ("G0556",), per_code=CFG.code_cited_per_code
+            mock_engine.return_value,
+            ("G0556",),
+            per_code=CFG.code_cited_per_code,
+            collections=CFG.code_cited_collections,
+            families=("APCM",),
         )
         body = client.stream.call_args.kwargs["json"]
         assert "Valuation (authoritative" in body["messages"][1]["content"]