Disambiguate recall spans with session context

benchmarks/recall_span_adapter.py

@@ -48,16 +48,20 @@ def select_recall_spans(
     *,
     session_limit: int = 1,
     min_similarity: float = 0.18,
+    group_similarity_weight: float = 1.0,
     max_chars_per_turn: int = 3_200,
 ) -> list[str]:
     """Select matched user prompts and their immediately paired replies.
 
     A character-within-word score is used because the measured recall failures
-    include inflection, punctuation and compound-word differences.  Ranking is
-    over user turns only.  Once a user request is selected, its next assistant
-    turn is reserved as provenance-preserving source text.
+    include inflection, punctuation and compound-word differences.  Each user
+    turn receives a bounded score from its complete source group so distinctive
+    entities elsewhere in the same conversation can disambiguate generic recall
+    wording.  Once a user request is selected, its next assistant turn is
+    reserved as provenance-preserving source text.
     """
-    if session_limit <= 0 or max_chars_per_turn <= 0:
+    if (session_limit <= 0 or max_chars_per_turn <= 0
+            or group_similarity_weight < 0):
         return []
     groups = [list(group) for group in turn_groups]
     candidates = []
@@ -77,14 +81,31 @@ def select_recall_spans(
         sublinear_tf=True, min_df=1)
     matrix = vectorizer.fit_transform([row[2] for row in candidates])
     scores = cosine_similarity(vectorizer.transform([query]), matrix)[0]
+    group_texts = [
+        "\n".join(str(_field(turn, "content") or "") for turn in turns)
+        for turns in groups
+    ]
+    group_vectorizer = TfidfVectorizer(
+        analyzer="char_wb", ngram_range=(3, 5), max_features=20_000,
+        sublinear_tf=True, min_df=1)
+    group_matrix = group_vectorizer.fit_transform(group_texts)
+    group_scores = cosine_similarity(
+        group_vectorizer.transform([query]), group_matrix)[0]
+    # Keep the established turn-level admission boundary. Group context may
+    # rerank an already-qualified recall query, but it must not broaden which
+    # queries receive a reserve.
+    if max(scores) < min_similarity:
+        return []
+    ranking_scores = [
+        scores[index] + group_similarity_weight * group_scores[row[0]]
+        for index, row in enumerate(candidates)
+    ]
     ranked = sorted(
-        range(len(candidates)), key=lambda index: (-scores[index], index))
+        range(len(candidates)), key=lambda index: (-ranking_scores[index], index))
 
     selected_groups = set()
     reserve = []
     for index in ranked:
-        if scores[index] < min_similarity:
-            break
         group_index, turn_index, user_content = candidates[index]
         if group_index in selected_groups:
             continue
@@ -112,6 +133,7 @@ def __init__(
         *,
         session_limit: int = 1,
         min_similarity: float = 0.18,
+        group_similarity_weight: float = 1.0,
         max_chars_per_turn: int = 3_200,
         base=None,
     ):
@@ -121,6 +143,7 @@ def __init__(
         }
         self.session_limit = session_limit
         self.min_similarity = min_similarity
+        self.group_similarity_weight = group_similarity_weight
         self.max_chars_per_turn = max_chars_per_turn
         self.base = base or TfIdfChromaAdapter()
         self.last_receipt = None
@@ -155,6 +178,7 @@ def no_recall(reason: str) -> list:
             groups,
             session_limit=self.session_limit,
             min_similarity=self.min_similarity,
+            group_similarity_weight=self.group_similarity_weight,
             max_chars_per_turn=self.max_chars_per_turn,
         )
         if not recall_reserve:

tests/test_recall_span_adapter.py

@@ -27,6 +27,53 @@ def test_selects_user_request_and_immediate_assistant_reply():
     assert selected == [groups[0][0]["content"], groups[0][1]["content"]]
 
 
+def test_group_context_disambiguates_generic_recall_wording():
+    groups = [[
+        {"role": "assistant", "content": "Ideas for the Radiation Amplified zombie."},
+        {"role": "user", "content": "Can you suggest some one-word names?"},
+        {"role": "assistant", "content": "Radik, Irradon, and Fissionator."},
+        {"role": "user", "content": "Fissionator is a really cool one."},
+        {"role": "assistant", "content": "The Fissionator design can use a protective suit."},
+    ], [
+        {"role": "user", "content": "What was the word you were not supposed to remember?"},
+        {"role": "assistant", "content": "I cannot provide that word."},
+    ]]
+    selected = select_recall_spans(
+        "What did we finally decide to name the Radiation Amplified zombie?",
+        groups,
+    )
+    assert selected[0] in {groups[0][1]["content"], groups[0][3]["content"]}
+    assert selected[1] in {groups[0][2]["content"], groups[0][4]["content"]}
+
+
+def test_group_context_cannot_bypass_similarity_floor_when_unrelated():
+    groups = [[
+        {"role": "user", "content": "Tell me about garden soil."},
+        {"role": "assistant", "content": "Use compost."},
+    ]]
+    assert select_recall_spans(
+        "What did we decide to name the Radiation Amplified zombie?",
+        groups,
+        min_similarity=0.95,
+    ) == []
+
+
+def test_group_context_reranks_without_broadening_admission():
+    groups = [[
+        {"role": "assistant", "content": "Radiation Amplified zombie naming discussion."},
+        {"role": "user", "content": "Any other options?"},
+        {"role": "assistant", "content": "Fissionator."},
+    ], [
+        {"role": "user", "content": "What was the word to remember?"},
+        {"role": "assistant", "content": "Another topic."},
+    ]]
+    assert select_recall_spans(
+        "Which name did we choose for the Radiation Amplified zombie?",
+        groups,
+        min_similarity=0.99,
+    ) == []
+
+
 def test_oversized_turn_returns_source_only_relevant_window():
     text = "A" * 3500 + " Construction of the house began in 2014. " + "B" * 3500
     window = _best_window(text, "When did construction of the house begin?", 3200)