fix(llm): family-cited window is one chunk per item, newest items first (refs #688)

This commit is contained in:
kert
2026-09-09 22:58:54 -04:00
parent 48f4519ca8
commit ef79b4e435
2 changed files with 51 additions and 9 deletions

View File

@@ -184,25 +184,42 @@ _CITED_SQL = text(
"ORDER BY cmetadata->>'date' DESC NULLS LAST, cmetadata->>'item_key', cmetadata->>'p_id'" "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`` #: Matches on the chunk's ``families`` metadata via array overlap — this is
#: metadata via array overlap — this is how a CCM question pulls the CPT #: how a CCM question pulls the CPT manual's "Chronic Care Management
#: manual's "Chronic Care Management Services" guideline chunk even though #: Services" guideline chunk even though that chunk never spells out 99490.
#: that chunk never spells out 99490. #: Two-stage windowing so the result is breadth across ITEMS, not several
#: chunks of whichever single document is newest: stage one keeps one chunk
#: per (family, item_key) — the earliest chunk in the item, by paragraph id
#: for rules or sequence number for comments/corpus; stage two then ranks
#: the surviving (one-per-item) rows per family by date and keeps the
#: newest :window.
_FAMILY_CITED_SQL = text( _FAMILY_CITED_SQL = text(
"SELECT document, cmetadata FROM ( " "SELECT document, cmetadata FROM ( "
"SELECT e.document, e.cmetadata, " "SELECT document, cmetadata, family, item_key, ordkey, "
"row_number() OVER ( " "row_number() OVER ( "
"PARTITION BY w.family " "PARTITION BY family "
"ORDER BY e.cmetadata->>'date' DESC NULLS LAST, e.cmetadata->>'item_key', e.cmetadata->>'p_id' " "ORDER BY cmetadata->>'date' DESC NULLS LAST, item_key, ordkey "
") AS rn " ") AS rn "
"FROM ( "
"SELECT e.document, e.cmetadata, w.family AS family, "
"e.cmetadata->>'item_key' AS item_key, "
"COALESCE(NULLIF(e.cmetadata->>'p_id', '')::int, "
"NULLIF(e.cmetadata->>'seq', '')::int) AS ordkey, "
"row_number() OVER ( "
"PARTITION BY w.family, e.cmetadata->>'item_key' "
"ORDER BY COALESCE(NULLIF(e.cmetadata->>'p_id', '')::int, "
"NULLIF(e.cmetadata->>'seq', '')::int) ASC NULLS LAST "
") AS item_rn "
"FROM langchain_pg_embedding e " "FROM langchain_pg_embedding e "
"JOIN unnest(CAST(:families AS text[])) AS w(family) " "JOIN unnest(CAST(:families AS text[])) AS w(family) "
"ON string_to_array(COALESCE(e.cmetadata->>'families', ''), ' ') && ARRAY[w.family] " "ON string_to_array(COALESCE(e.cmetadata->>'families', ''), ' ') && ARRAY[w.family] "
"WHERE e.collection_id = (SELECT uuid FROM langchain_pg_collection " "WHERE e.collection_id = (SELECT uuid FROM langchain_pg_collection "
"WHERE name = :collection) " "WHERE name = :collection) "
") one_per_item "
"WHERE item_rn = 1 "
") t " ") t "
"WHERE rn <= :window " "WHERE rn <= :window "
"ORDER BY cmetadata->>'date' DESC NULLS LAST, cmetadata->>'item_key', cmetadata->>'p_id'" "ORDER BY cmetadata->>'date' DESC NULLS LAST, item_key, ordkey"
) )

View File

@@ -502,13 +502,38 @@ class TestFamilyCitedSources:
engine, [], per_code=2, collections=("rules",), families=["CCM"] engine, [], per_code=2, collections=("rules",), families=["CCM"]
) )
sql = str(conn.execute.call_args.args[0]) sql = str(conn.execute.call_args.args[0])
assert "&&" in sql and "PARTITION BY w.family" in sql assert "&&" in sql
# one chunk per (family, item_key), then the survivors ranked per family
assert sql.count("PARTITION BY") == 2
assert "PARTITION BY w.family, e.cmetadata->>'item_key'" in sql
assert "PARTITION BY family" in sql
assert conn.execute.call_args.args[1] == { assert conn.execute.call_args.args[1] == {
"collection": "rules", "collection": "rules",
"families": ["CCM"], "families": ["CCM"],
"window": 4, "window": 4,
} }
def test_called_once_per_collection_with_same_params(self):
engine, conn = _multi_engine(
by_family={"rules": [], "comments": [], "corpus": []}
)
code_cited_sources(
engine,
[],
per_code=3,
collections=("rules", "comments", "corpus"),
families=["CCM"],
)
assert conn.execute.call_count == 3
for call, collection in zip(
conn.execute.call_args_list, ("rules", "comments", "corpus")
):
assert call.args[1] == {
"collection": collection,
"families": ["CCM"],
"window": 6,
}
def test_family_rows_appended_after_code_rows(self): def test_family_rows_appended_after_code_rows(self):
engine, _conn = _multi_engine( engine, _conn = _multi_engine(
by_code={"rules": [("code hit", _rule_md(1, "G0556"))]}, by_code={"rules": [("code hit", _rule_md(1, "G0556"))]},