feat(llm): stamp codes metadata on every chunk; links.as_source shared source dict (refs P48)

src/llm/chunk.py

@@ -14,6 +14,8 @@ import hashlib
 import re
 from dataclasses import dataclass
 
+from pfs.families import find_codes
+
 
 @dataclass(frozen=True)
 class Paragraph:
@@ -143,6 +145,7 @@ def chunk_doc(
                 "item_key": doc.key,
                 "seq": str(seq),
                 "section": heading,
+                "codes": " ".join(find_codes(piece)),
             },
         )
         for seq, (heading, piece) in enumerate(pieces)
@@ -184,6 +187,7 @@ def _chunk_paragraphs(doc: Doc, target: int, overlap: int) -> list[Chunk]:
                 "item_key": doc.key,
                 "seq": str(seq),
                 "section": "",
+                "codes": " ".join(find_codes(text)),
                 "p_id": str(first.p_id),
                 "p_id_last": str(last.p_id),
                 "page": str(first.page),

src/llm/links.py

@@ -19,6 +19,7 @@ from bib.frlink import text_fragment
 _COMMENT = "https://www.regulations.gov/comment/{cid}"
 _DOWNLOAD = "https://downloads.regulations.gov/{cid}/{name}"
 _TITLE_MAX = 60
+_SNIPPET_CHARS = 500
 
 
 def _short(title: str) -> str:
@@ -84,3 +85,21 @@ def for_source(md: dict[str, str], snippet: str) -> tuple[str, str]:
     if kind == "rule":
         return _rule(md, snippet)
     return _corpus(md)
+
+
+def as_source(md: dict[str, str], text: str, score: float) -> dict:
+    """The source dict the chat sends to the prompt and the UI."""
+    snippet = text[:_SNIPPET_CHARS].strip()
+    url, label = for_source(md, snippet)
+    return {
+        "id": label,
+        "label": label,
+        "kind": md.get("kind", ""),
+        "url": url,
+        "title": md.get("title", ""),
+        "date": md.get("date", ""),
+        "docket": md.get("docket", ""),
+        "comment_id": md.get("comment_id", ""),
+        "snippet": snippet,
+        "score": round(score, 4),
+    }

src/llm/rag.py

@@ -19,14 +19,13 @@ from typing import Iterator
 import httpx
 
 from llm.config import LlmConfig
-from llm.links import for_source
+from llm.links import as_source
 from llm.pool import HostPool, PoolEmbeddings, pick_model
 from llm.rerank import Hit, blend, filter_since
 
 _TIMEOUT = httpx.Timeout(300.0, connect=5.0)
 _COLLECTIONS = {"comment": "comments", "rule": "rules", "corpus": "corpus"}
 _OVERFETCH = 3
-_SNIPPET_CHARS = 500
 
 _SYSTEM = (
     "You answer questions about CMS rulemaking using ONLY the excerpts "
@@ -65,21 +64,7 @@ def _hits(question_vec: list[float], *, cfg: LlmConfig, pool: HostPool) -> list[
 
 
 def _source(hit: Hit) -> dict:
-    md = hit.metadata
-    snippet = hit.text[:_SNIPPET_CHARS].strip()
-    url, label = for_source(md, snippet)
-    return {
-        "id": label,
-        "label": label,
-        "kind": md.get("kind", ""),
-        "url": url,
-        "title": md.get("title", ""),
-        "date": md.get("date", ""),
-        "docket": md.get("docket", ""),
-        "comment_id": md.get("comment_id", ""),
-        "snippet": snippet,
-        "score": round(hit.score, 4),
-    }
+    return as_source(hit.metadata, hit.text, hit.score)
 
 
 def retrieve(

tests/llm/test_chunk.py

@@ -151,3 +151,27 @@ class TestParagraphChunks:
         a = chunk_doc(Doc(key="R", text="same", metadata={}, paragraphs=paras))
         b = chunk_doc(Doc(key="R", text="same", metadata={}, paragraphs=paras))
         assert [c.id for c in a] == [c.id for c in b]
+
+
+class TestCodesMetadata:
+    def test_codes_stamped_per_chunk(self):
+        text = "We propose G0556 and G0557.\n\n## Other\n\nNo codes here.\n\n## More\n\n99490 applies."
+        chunks = chunk_doc(_doc(text), target_chars=60, overlap_chars=10)
+        found = {c.metadata["codes"] for c in chunks}
+        assert "G0556 G0557" in found
+        assert "99490" in found
+        assert "" in found
+
+    def test_rule_paragraph_chunks_get_codes(self):
+        from llm.chunk import Paragraph
+
+        doc = Doc(
+            key="R1",
+            text="x",
+            metadata={"kind": "rule"},
+            paragraphs=(
+                Paragraph(1, 100, 1, "APCM code G0556 is valued at 0.25 work RVUs."),
+            ),
+        )
+        (chunk,) = chunk_doc(doc)
+        assert chunk.metadata["codes"] == "G0556"

tests/llm/test_links.py

@@ -127,3 +127,34 @@ class TestCorpus:
 def test_unknown_kind_treated_as_corpus():
     url, label = for_source({"url": "https://u.test", "title": "T"}, "s")
     assert (url, label) == ("https://u.test", "T")
+
+
+class TestAsSource:
+    def test_shape_matches_rag_source(self):
+        from llm.links import as_source
+
+        s = as_source(
+            {
+                "kind": "comment",
+                "comment_id": "CMS-2026-2377-1",
+                "date": "2026-08-19",
+                "title": "t",
+                "docket": "CMS-2026-2377",
+            },
+            "body " * 200,
+            0.123456,
+        )
+        assert set(s) == {
+            "id",
+            "label",
+            "kind",
+            "url",
+            "title",
+            "date",
+            "docket",
+            "comment_id",
+            "snippet",
+            "score",
+        }
+        assert s["label"] == "CMS-2026-2377-1" and s["kind"] == "comment"
+        assert len(s["snippet"]) <= 500 and s["score"] == 0.1235