fix(llm): lineage sources prefer the question's own codes; element-diff anchors; cap 8 (refs #691 #692)

src/llm/config.py

@@ -53,7 +53,7 @@ class LlmConfig:
     code_cited_max: int = 12
     lineage_max_rows: int = 25
     lineage_on_demand_max: int = 3
-    lineage_sources_max: int = 6
+    lineage_sources_max: int = 8
 
 
 def parse_hosts(spec: str) -> tuple[tuple[str, float], ...]:
@@ -129,7 +129,7 @@ def load() -> LlmConfig:
         code_cited_max=int(_opt(section, "code_cited_max", 12)),
         lineage_max_rows=int(_opt(section, "lineage_max_rows", 25)),
         lineage_on_demand_max=int(_opt(section, "lineage_on_demand_max", 3)),
-        lineage_sources_max=int(_opt(section, "lineage_sources_max", 6)),
+        lineage_sources_max=int(_opt(section, "lineage_sources_max", 8)),
         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/lineage.py

@@ -203,6 +203,12 @@ class LineageEvidence:
     #: ``cfg.lineage_max_rows`` at construction time — ``prompt_block``
     #: is self-contained (no cfg argument), like ``ValuationEvidence``.
     max_prompt_rows: int = 25
+    #: Ruling B10: codes literally named in the question
+    #: (``pfs.families.Detection.explicit``) — ``lineage_sources`` sorts
+    #: paragraphs for these codes ahead of a detected family's other
+    #: codes, so "what replaced G2058" surfaces G2058's own paragraphs
+    #: before its family-mates'.
+    explicit: tuple[str, ...] = ()
 
     def _prompt_events(self) -> list[LineageEvent]:
         return _select_for_prompt(self.events, self.max_prompt_rows)
@@ -229,6 +235,8 @@ class LineageEvidence:
         }
         if self.elements_note:
             out["elements_note"] = self.elements_note
+        if self.explicit:
+            out["explicit"] = list(self.explicit)
         return out
 
 
@@ -612,23 +620,26 @@ def _lineage_paragraph(store: Any, item_key: str, p_id: int) -> dict | None:
     return result
 
 
-def _lineage_source(store: Any, event: LineageEvent) -> dict | None:
+def _labeled_rule_source(
-    """One rule-kind source dict for *event*'s FR paragraph, built the
+    store: Any, item_key: str, p_id: int, page_hint: int, label: str
-    same way the indexer builds a rule chunk (``links._rule`` reads
+) -> dict | None:
-    ``item_key``/``p_id``/``page``/``ordinal``/``fr_volume``/``html_url``
+    """One rule-kind source dict for the FR paragraph at *(item_key,
-    from the metadata) so ``as_source`` produces the same kind of
+    p_id)*, labelled *label* — built the same way the indexer builds a
-    (url, label) a retrieved excerpt for this paragraph would — except
+    rule chunk (``links._rule`` reads ``item_key``/``p_id``/``page``/
-    the label/id are then overridden with *event*'s own lineage label,
+    ``ordinal``/``fr_volume``/``html_url`` from the metadata) so
-    so the prompt block's ``[label]`` and the sources drawer's entry are
+    ``as_source`` produces the same kind of (url, label) a retrieved
-    the exact same string. ``None`` when the paragraph can't be found."""
+    excerpt for this paragraph would, before the label/id are
-    para = _lineage_paragraph(store, event.item_key, event.p_id)
+    overridden with *label*. Shared by ``_lineage_source`` (a lineage
+    event's own paragraph) and ``lineage_sources``'s element-diff
+    anchors (Ruling B10). ``None`` when the paragraph can't be found."""
+    para = _lineage_paragraph(store, item_key, p_id)
     if para is None:
         return None
     md = {
         "kind": "rule",
-        "item_key": event.item_key,
+        "item_key": item_key,
-        "p_id": str(event.p_id),
+        "p_id": str(p_id),
-        "page": str(para.get("page") or event.page or ""),
+        "page": str(para.get("page") or page_hint or ""),
         "ordinal": str(para.get("ordinal") or ""),
         "fr_volume": str(para.get("fr_volume") or ""),
         "html_url": para.get("html_url") or "",
@@ -637,13 +648,34 @@ def _lineage_source(store: Any, event: LineageEvent) -> dict | None:
         "date": (para.get("date_published") or "")[:10],
     }
     src = as_source(md, para.get("text") or "", 0.0)
-    src["label"] = event.label
+    src["label"] = label
-    src["id"] = event.label
+    src["id"] = label
     return src
 
 
+def _lineage_source(store: Any, event: LineageEvent) -> dict | None:
+    """*event*'s own FR paragraph as a source, labelled with *event*'s
+    own lineage label so the prompt block's ``[label]`` and the sources
+    drawer's entry are the exact same string."""
+    return _labeled_rule_source(
+        store, event.item_key, event.p_id, event.page, event.label
+    )
+
+
+#: Ruling B10: element-diff anchors are attempted regardless of whether
+#: the event-paragraph cap (*max_items*) has room left — up to this
+#: many, subject only to ``_LINEAGE_SOURCES_HARD_CAP``.
+_ELEMENT_ANCHOR_MAX = 2
+
+#: Ruling B10: absolute ceiling on ``lineage_sources``' output —
+#: *max_items* event paragraphs plus up to ``_ELEMENT_ANCHOR_MAX``
+#: element-diff anchors never exceeds this even when the event loop
+#: already used every *max_items* slot.
+_LINEAGE_SOURCES_HARD_CAP = 10
+
+
 def lineage_sources(
-    store: Any, evidence: LineageEvidence, *, max_items: int = 6
+    store: Any, evidence: LineageEvidence, *, max_items: int = 8
 ) -> list[dict]:
     """FR paragraphs for the events selected into the lineage prompt
     block, as ``rule``-kind source dicts — so the model can actually
@@ -652,21 +684,40 @@ def lineage_sources(
 
     Distinct ``(item_key, p_id)`` pairs among ``source == "fr"`` prompt
     events (the same selection ``LineageEvidence.prompt_block`` uses),
-    priority-kind events first (``_PRIORITY_KINDS`` — created, adopted,
+    sorted (Ruling B10) with events whose ``code`` is one of
-    replaces, replaced_by, deleted, disappeared), then by year; capped
+    ``evidence.explicit`` (codes literally named in the question) first,
-    at *max_items*. Each source's ``label``/``id`` are overridden with
+    then priority-kind events (``_PRIORITY_KINDS`` — created, adopted,
-    the lineage event's own label (``_lineage_source``), so the same
+    replaces, replaced_by, deleted, disappeared), then by year — a
+    question naming one code out of a larger detected family gets that
+    code's own paragraphs ahead of its family-mates'. Capped at
+    *max_items*. Each source's ``label``/``id`` are overridden with the
+    lineage event's own label (``_lineage_source``), so the same
     paragraph never shows two different bracketed labels to the model.
 
+    Ruling B10: when ``evidence.element_diffs`` is non-empty, up to
+    ``_ELEMENT_ANCHOR_MAX`` distinct element-diff paragraphs (most-
+    shared value first — the order ``_element_diffs`` already emits)
+    are appended after the event paragraphs, labelled with the diff's
+    own label. These are attempted even when *max_items* has no room
+    left — the only cap on them is ``_LINEAGE_SOURCES_HARD_CAP`` on the
+    total.
+
     Never raises: a failure fetching any one paragraph is logged and
     skipped; a failure before the loop starts degrades to ``[]``, the
     same contract as ``lineage_evidence`` itself."""
     try:
+        explicit = set(evidence.explicit)
         prompt_events = _select_for_prompt(evidence.events, evidence.max_prompt_rows)
         candidates = [
             e for e in prompt_events if e.source == "fr" and e.item_key and e.p_id
         ]
-        candidates.sort(key=lambda e: (0 if e.kind in _PRIORITY_KINDS else 1, e.year))
+        candidates.sort(
+            key=lambda e: (
+                0 if e.code in explicit else 1,
+                0 if e.kind in _PRIORITY_KINDS else 1,
+                e.year,
+            )
+        )
     except Exception as e:  # noqa: BLE001
         log.warning("lineage sources skipped (select): %s", e)
         return []
@@ -686,6 +737,27 @@ def lineage_sources(
             out.append(src)
         if len(out) >= max_items:
             break
+
+    added = 0
+    for d in evidence.element_diffs:
+        if added >= _ELEMENT_ANCHOR_MAX or len(out) >= _LINEAGE_SOURCES_HARD_CAP:
+            break
+        if not d.item_key or not d.p_id:
+            continue
+        key = (d.item_key, d.p_id)
+        if key in seen:
+            continue
+        seen.add(key)
+        try:
+            src = _labeled_rule_source(store, d.item_key, d.p_id, 0, d.label)
+        except Exception as ex:  # noqa: BLE001
+            log.warning(
+                "element-diff source skipped (%s p-%s): %s", d.item_key, d.p_id, ex
+            )
+            continue
+        if src is not None:
+            out.append(src)
+            added += 1
     return out
 
 
@@ -770,4 +842,5 @@ def lineage_evidence(question: str, cfg: LlmConfig) -> LineageEvidence | None:
         guidance=guidance,
         elements_note=elements_note,
         max_prompt_rows=cfg.lineage_max_rows,
+        explicit=det.explicit,
     )

stack.toml

@@ -133,7 +133,7 @@ code_cited_collections = ["rules", "comments", "corpus"]  # searched in this ord
 code_cited_max = 12                      # cited excerpts kept after round-robin interleave across collections; <= 0 = unlimited
 lineage_max_rows = 25                    # collapsed lineage events kept in the prompt block (the SSE payload always carries every collapsed row)
 lineage_on_demand_max = 3                # detected codes per turn allowed to fall back to pfs.lineage.lineage() when pfs.code_event has no rows for them
-lineage_sources_max = 6                  # FR paragraphs fetched as sources for events selected into the lineage prompt block, outside code_cited_max
+lineage_sources_max = 8                  # FR paragraphs fetched as sources for events selected into the lineage prompt block, outside code_cited_max (Ruling B10: + up to 2 element-diff anchors, hard cap 10 total)
 
 [llm.k_per_kind]                       # over-fetched ×3 per kind, then re-ranked
 comment = 8

tests/llm/test_lineage.py

@@ -949,14 +949,17 @@ class TestCachingAndScopedUrlResolution:
         assert len(seen_keys) == len(set(seen_keys))  # never re-fetched
 
 
-def _lineage_evidence(events, max_prompt_rows=25) -> LineageEvidence:
+def _lineage_evidence(
+    events, max_prompt_rows=25, explicit=(), element_diffs=()
+) -> LineageEvidence:
     return LineageEvidence(
         codes=(),
         families=(),
         events=tuple(events),
-        element_diffs=(),
+        element_diffs=tuple(element_diffs),
         guidance=(),
         max_prompt_rows=max_prompt_rows,
+        explicit=tuple(explicit),
     )
 
 
@@ -1124,6 +1127,166 @@ class TestLineageSources:
         assert lineage_sources(store, ev, max_items=6) == []
 
 
+def _insert_extra_anchor(store, key, year, p_id, text="extra text"):
+    store.con.execute(
+        "INSERT INTO items VALUES (?,?,?,?)",
+        (
+            key,
+            f"Medicare Program; CY {year} PFS Final Rule",
+            f"{year - 1}-11-01",
+            f"https://www.federalregister.gov/d/{year}-doc-{key}",
+        ),
+    )
+    store.con.execute(
+        "INSERT INTO fr_anchors VALUES (?,?,?,?,?)",
+        (key, p_id, 1000 + p_id, 1, text),
+    )
+    store.con.commit()
+
+
+def _diff(item_key, p_id, label, type_="element", value="v") -> ElementDiff:
+    return ElementDiff(
+        type=type_,
+        value=value,
+        in_codes=("A",),
+        not_in_codes=("B",),
+        label=label,
+        item_key=item_key,
+        p_id=p_id,
+    )
+
+
+class TestLineageSourcesRulingB10:
+    """Ruling B10 — explicit-code-first ordering and element-diff
+    anchors appended to ``lineage_sources``' output."""
+
+    def test_explicit_code_sorts_first_even_when_oldest_and_non_priority(self, store):
+        # A (explicit, literally named in the question) has the oldest,
+        # non-priority-kind event; B and C (family-mates, not named) have
+        # newer, priority-kind events. Explicit-code membership must
+        # outrank both the priority-kind tier and the year — A still
+        # sorts first.
+        a = lineage._to_lineage_event(
+            store,
+            _ev("A", 2015, "revalued", item_key="DE2VH9PD", p_id=1251),
+            mtime=0,
+        )
+        b = lineage._to_lineage_event(
+            store,
+            _ev("B", 2021, "created", item_key="YBM4IZUS", p_id=686),
+            mtime=0,
+        )
+        c = lineage._to_lineage_event(
+            store,
+            _ev("C", 2021, "created", item_key="YBM4IZUS", p_id=1578),
+            mtime=0,
+        )
+        ev = _lineage_evidence([a, b, c], explicit=("A",))
+        out = lineage_sources(store, ev, max_items=8)
+        assert [s["label"] for s in out] == [a.label, b.label, c.label]
+
+    def test_explicit_code_survives_a_cap_that_would_otherwise_drop_it(self, store):
+        a = lineage._to_lineage_event(
+            store,
+            _ev("A", 2015, "revalued", item_key="DE2VH9PD", p_id=1251),
+            mtime=0,
+        )
+        b = lineage._to_lineage_event(
+            store,
+            _ev("B", 2021, "created", item_key="YBM4IZUS", p_id=686),
+            mtime=0,
+        )
+        ev = _lineage_evidence([a, b], explicit=("A",))
+        (out,) = lineage_sources(store, ev, max_items=1)
+        assert out["label"] == a.label
+
+    def test_element_diff_anchors_appended_after_events(self, store):
+        e = lineage._to_lineage_event(
+            store,
+            _ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251),
+            mtime=0,
+        )
+        diff = _diff("YBM4IZUS", 686, "CY2021 PFS final ¶686")
+        ev = _lineage_evidence([e], element_diffs=[diff])
+        out = lineage_sources(store, ev, max_items=8)
+        assert [s["label"] for s in out] == [e.label, diff.label]
+        assert out[1]["kind"] == "rule"
+        assert out[1]["item_key"] == "YBM4IZUS"
+        assert out[1]["p_id"] == "686"
+
+    def test_element_diff_anchors_capped_at_two_most_shared_first(self, store):
+        _insert_extra_anchor(store, "EX1", 2016, 10)
+        _insert_extra_anchor(store, "EX2", 2017, 11)
+        diffs = [
+            _diff("DE2VH9PD", 1251, "CY2015 PFS final ¶1251"),  # most shared
+            _diff("EX1", 10, "CY2016 PFS final ¶10"),
+            _diff("EX2", 11, "CY2017 PFS final ¶11"),  # third — dropped by the cap
+        ]
+        ev = _lineage_evidence([], element_diffs=diffs)
+        out = lineage_sources(store, ev, max_items=8)
+        assert [s["label"] for s in out] == [diffs[0].label, diffs[1].label]
+
+    def test_element_diffs_still_added_past_max_items_up_to_the_hard_cap(self, store):
+        # 8 distinct event paragraphs exactly fill max_items=8; two more
+        # distinct element-diff paragraphs are still appended (the "add
+        # up to 2 even when max_items has no room" rule), reaching the
+        # absolute hard cap of 10, not 11.
+        events = []
+        for i in range(8):
+            key = f"EV{i}"
+            _insert_extra_anchor(store, key, 2000 + i, 100 + i)
+            events.append(
+                lineage._to_lineage_event(
+                    store,
+                    _ev("Z", 2000 + i, "created", item_key=key, p_id=100 + i),
+                    mtime=0,
+                )
+            )
+        _insert_extra_anchor(store, "DX1", 2050, 500)
+        _insert_extra_anchor(store, "DX2", 2051, 501)
+        _insert_extra_anchor(store, "DX3", 2052, 502)
+        diffs = [
+            _diff("DX1", 500, "d1"),
+            _diff("DX2", 501, "d2"),
+            _diff("DX3", 502, "d3"),  # third diff — dropped by the element cap of 2
+        ]
+        ev = _lineage_evidence(events, explicit=("Z",), element_diffs=diffs)
+        out = lineage_sources(store, ev, max_items=8)
+        assert len(out) == 10
+        assert [s["label"] for s in out[8:]] == ["d1", "d2"]
+
+    def test_element_diff_with_no_anchor_is_skipped(self, store):
+        e = lineage._to_lineage_event(
+            store,
+            _ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251),
+            mtime=0,
+        )
+        rvu_only_diff = ElementDiff(
+            type="t",
+            value="v",
+            in_codes=("A",),
+            not_in_codes=("B",),
+            label="l",
+            item_key="",
+            p_id=0,
+        )
+        ev = _lineage_evidence([e], element_diffs=[rvu_only_diff])
+        out = lineage_sources(store, ev, max_items=8)
+        assert [s["label"] for s in out] == [e.label]
+
+    def test_element_diff_deduped_against_an_event_paragraph(self, store):
+        e = lineage._to_lineage_event(
+            store,
+            _ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251),
+            mtime=0,
+        )
+        same_paragraph_diff = _diff("DE2VH9PD", 1251, "some other label")
+        ev = _lineage_evidence([e], element_diffs=[same_paragraph_diff])
+        out = lineage_sources(store, ev, max_items=8)
+        assert len(out) == 1
+        assert out[0]["label"] == e.label
+
+
 class TestReconcileLabels:
     """``_reconcile_labels`` — a retrieved/cited source for the same FR
     paragraph as a lineage row wins the slot, renamed to the lineage