From 392c35d7d681e6e69185bff19cc1af4abe4953a1 Mon Sep 17 00:00:00 2001 From: kert Date: Tue, 8 Sep 2026 22:36:04 -0400 Subject: [PATCH] feat(llm): stamp codes metadata on every chunk; links.as_source shared source dict (refs P48) --- src/llm/chunk.py | 4 ++++ src/llm/links.py | 19 +++++++++++++++++++ src/llm/rag.py | 19 ++----------------- tests/llm/test_chunk.py | 24 ++++++++++++++++++++++++ tests/llm/test_links.py | 31 +++++++++++++++++++++++++++++++ 5 files changed, 80 insertions(+), 17 deletions(-) diff --git a/src/llm/chunk.py b/src/llm/chunk.py index 89d4a4a..7721866 100644 --- a/src/llm/chunk.py +++ b/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), diff --git a/src/llm/links.py b/src/llm/links.py index a702468..f611ddc 100644 --- a/src/llm/links.py +++ b/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), + } diff --git a/src/llm/rag.py b/src/llm/rag.py index 4ecc272..dc7712c 100644 --- a/src/llm/rag.py +++ b/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( diff --git a/tests/llm/test_chunk.py b/tests/llm/test_chunk.py index 88b770f..ff1355e 100644 --- a/tests/llm/test_chunk.py +++ b/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" diff --git a/tests/llm/test_links.py b/tests/llm/test_links.py index 7b0f5cb..1727cec 100644 --- a/tests/llm/test_links.py +++ b/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