Preserve decision sequence in recall spans

benchmarks/recall_span_adapter.py

@@ -7,13 +7,39 @@
 
 from __future__ import annotations
 
+import re
 from typing import Iterable, Mapping, Sequence
 
 from agentmem_os.benchmarks.dated_event_adapter import DatedEventContextAssembler
 from agentmem_os.benchmarks.dated_event_reserve import prepend_reserve
 from agentmem_os.benchmarks.real_code_utils import TfIdfChromaAdapter
 
 
+_DECISION_RECALL_RE = re.compile(
+    r"\b(?:decid(?:e|ed)|chos(?:e|en)|choose|sett(?:le|led)|finally|"
+    r"end(?:ed)?\s+up|went\s+with)\b",
+    re.IGNORECASE,
+)
+_POSITIVE_ACCEPTANCE_RE = re.compile(
+    r"\b(?:love|cool|good|great|perfect|exactly|works|favorite|favourite|"
+    r"go\s+with|settle\s+on|choose|chose)\b",
+    re.IGNORECASE,
+)
+_NAMING_RECALL_RE = re.compile(r"\b(?:name|named|call|called)\b", re.IGNORECASE)
+_NAMING_EVIDENCE_RE = re.compile(
+    r"\b(?:name|names|named|call|called)\b", re.IGNORECASE)
+_MORE_OPTIONS_RE = re.compile(
+    r"\b(?:any\s+other|more|another|few)\b[^.!?\n]{0,48}\b"
+    r"(?:name|option|idea|suggestion)s?\b",
+    re.IGNORECASE,
+)
+_WORD_RE = re.compile(r"[A-Za-z][A-Za-z0-9_-]{4,}")
+_DECISION_STOPWORDS = {
+    "about", "after", "again", "could", "finally", "great", "really",
+    "their", "there", "these", "thing", "think", "those", "would",
+}
+
+
 def _field(turn, name: str):
     if isinstance(turn, dict):
         return turn.get(name)
@@ -42,6 +68,53 @@ def _best_window(text: str, query: str, max_chars: int) -> str:
     return windows[max(range(len(windows)), key=lambda index: (scores[index], -index))]
 
 
+def _local_window(turns: Sequence, turn_index: int) -> list:
+    """Keep a source-ordered proposal, user response, and continuation."""
+    start = turn_index
+    if (turn_index > 0
+            and str(_field(turns[turn_index - 1], "role") or "").lower()
+            == "assistant"):
+        start -= 1
+    end = turn_index + 1
+    if (end < len(turns)
+            and str(_field(turns[end], "role") or "").lower()
+            == "assistant"):
+        end += 1
+    return list(turns[start:end])
+
+
+def _decision_bonus(
+    turns: Sequence,
+    turn_index: int,
+    *,
+    naming_recall: bool,
+) -> float:
+    """Favor accepted proposals over requests for additional options."""
+    user = str(_field(turns[turn_index], "content") or "")
+    if _MORE_OPTIONS_RE.search(user):
+        return -0.30
+    if not _POSITIVE_ACCEPTANCE_RE.search(user):
+        return 0.0
+    window = _local_window(turns, turn_index)
+    window_text = "\n".join(
+        str(_field(turn, "content") or "") for turn in window)
+    if naming_recall and not _NAMING_EVIDENCE_RE.search(window_text):
+        return 0.0
+    if len(window) != 3:
+        return 0.25
+    token_sets = []
+    for turn in window:
+        token_sets.append({
+            token.lower() for token in _WORD_RE.findall(
+                str(_field(turn, "content") or ""))
+            if token.lower() not in _DECISION_STOPWORDS
+        })
+    # Continued use of the same distinctive term after positive user language
+    # is source-only evidence that a proposal was adopted.
+    repeated = set.intersection(*token_sets)
+    return 0.65 if repeated else 0.40
+
+
 def select_recall_spans(
     query: str,
     turn_groups: Iterable[Sequence],
@@ -96,29 +169,77 @@ def select_recall_spans(
     # queries receive a reserve.
     if max(scores) < min_similarity:
         return []
-    ranking_scores = [
+    decision_recall = bool(_DECISION_RECALL_RE.search(query or ""))
+    naming_recall = bool(_NAMING_RECALL_RE.search(query or ""))
+    if decision_recall:
+        local_texts = [
+            "\n".join(str(_field(turn, "content") or "")
+                      for turn in _local_window(groups[row[0]], row[1]))
+            for row in candidates
+        ]
+        local_vectorizer = TfidfVectorizer(
+            analyzer="char_wb", ngram_range=(3, 5), max_features=20_000,
+            sublinear_tf=True, min_df=1)
+        local_matrix = local_vectorizer.fit_transform(local_texts)
+        local_scores = cosine_similarity(
+            local_vectorizer.transform([query]), local_matrix)[0]
+    else:
+        local_scores = [0.0] * len(candidates)
+    base_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: (-ranking_scores[index], index))
+    within_group_scores = []
+    for index, row in enumerate(candidates):
+        group_index, turn_index, _ = row
+        decision_bonus = (
+            _decision_bonus(
+                groups[group_index], turn_index,
+                naming_recall=naming_recall,
+            )
+            if decision_recall else 0.0
+        )
+        within_group_scores.append(
+            base_ranking_scores[index] + local_scores[index] + decision_bonus)
+    # Select source sessions with the already-verified group-aware score. Apply
+    # decision evidence only after that boundary so generic acceptance language
+    # in an unrelated session cannot change which session is admitted.
+    candidates_by_group = {}
+    for index, row in enumerate(candidates):
+        candidates_by_group.setdefault(row[0], []).append(index)
+    ranked_groups = sorted(
+        candidates_by_group,
+        key=lambda group_index: (
+            -max(base_ranking_scores[index]
+                 for index in candidates_by_group[group_index]),
+            group_index,
+        ),
+    )
 
     selected_groups = set()
     reserve = []
-    for index in ranked:
-        group_index, turn_index, user_content = candidates[index]
-        if group_index in selected_groups:
-            continue
+    for group_index in ranked_groups:
+        index = max(
+            candidates_by_group[group_index],
+            key=lambda candidate_index: (
+                within_group_scores[candidate_index], -candidate_index),
+        )
+        _, turn_index, _ = candidates[index]
         selected_groups.add(group_index)
-        reserve.append(_best_window(user_content, query, max_chars_per_turn))
         turns = groups[group_index]
-        if turn_index + 1 < len(turns):
+        selected_turns = (
+            _local_window(turns, turn_index)
+            if decision_recall else [turns[turn_index]]
+        )
+        if not decision_recall and turn_index + 1 < len(turns):
             reply = turns[turn_index + 1]
-            reply_content = str(_field(reply, "content") or "")
-            if (str(_field(reply, "role") or "").lower() == "assistant"
-                    and reply_content):
+            if str(_field(reply, "role") or "").lower() == "assistant":
+                selected_turns.append(reply)
+        for selected_turn in selected_turns:
+            selected_content = str(_field(selected_turn, "content") or "")
+            if selected_content:
                 reserve.append(_best_window(
-                    reply_content, query, max_chars_per_turn))
+                    selected_content, query, max_chars_per_turn))
         if len(selected_groups) >= session_limit:
             break
     return reserve

tests/test_recall_span_adapter.py

@@ -31,7 +31,7 @@ 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": "assistant", "content": "Names: Radik, Irradon, and Fissionator."},
         {"role": "user", "content": "Fissionator is a really cool one."},
         {"role": "assistant", "content": "The Fissionator design can use a protective suit."},
     ], [
@@ -42,8 +42,26 @@ def test_group_context_disambiguates_generic_recall_wording():
         "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"]}
+    assert selected == [turn["content"] for turn in groups[0][2:5]]
+
+
+def test_decision_recall_reserves_proposal_acceptance_and_continuation():
+    group = [
+        {"role": "assistant", "content": "How about Radialisk?"},
+        {"role": "user", "content": "Any other name ideas?"},
+        {"role": "assistant", "content": "Contaminated Colossus or Irradiated Behemoth."},
+        {"role": "user", "content": "How about a few one-word names?"},
+        {"role": "assistant", "content": "Names for the Radiation Amplified: Radik, Fissionator, Radiatron."},
+        {"role": "user", "content": "Fissionator is a REALLY cool one, especially with a mechanical design."},
+        {"role": "assistant", "content": "The Fissionator could wear a protective radiation suit."},
+        {"role": "user", "content": "Can you give me ideas for what the Fissionator could look like?"},
+        {"role": "assistant", "content": "The Fissionator could be a robotic construct."},
+    ]
+    selected = select_recall_spans(
+        "What did we finally decide to name the Radiation Amplified zombie?",
+        [group],
+    )
+    assert selected == [turn["content"] for turn in group[4:7]]
 
 
 def test_group_context_cannot_bypass_similarity_floor_when_unrelated():