fix(recall): fill the source_facts budget in rank order and flag truncation (#3221)

hindsight-api-slim/hindsight_api/api/http.py

@@ -598,6 +598,14 @@ class RecallResponse(BaseModel):
     source_facts: dict[str, RecallResult] | None = Field(
         default=None, description="Source facts for observation-type results, keyed by fact ID"
     )
+    source_facts_truncated: bool | None = Field(
+        default=None,
+        description=(
+            "Whether the source_facts map was cut short by the token budget. When true, some IDs in "
+            "results[].source_fact_ids have no entry in source_facts — the budget ran out, the "
+            "references are not dangling. Only set when source facts were requested."
+        ),
+    )
 
 
 class EntityInput(BaseModel):
@@ -4651,6 +4659,7 @@ def _fact_to_result(fact: "MemoryFact") -> RecallResult:
                 entities=entities_response,
                 chunks=chunks_response,
                 source_facts=source_facts_response,
+                source_facts_truncated=core_result.source_facts_truncated,
             )
 
             handler_duration = time.time() - handler_start

hindsight-api-slim/hindsight_api/engine/memory_engine.py

@@ -456,6 +456,7 @@ def validate_sql_schema(sql: str) -> None:
 from .search.reranking import CrossEncoderReranker, apply_combined_scoring
 from .search.tags import TagGroup, TagsMatch, build_tag_groups_where_clause, build_tags_where_clause
 from .search.types import ScoredResult
+from .source_facts import select_source_facts_within_budget
 from .task_backend import TaskBackend
 
 # Recall ranking strategy: how the per-arm (semantic/bm25/graph/temporal) results are
@@ -6569,6 +6570,7 @@ def to_tuple_format(results):
             source_fact_start = time.time()
             source_fact_ids_by_obs: dict[str, list[str]] = {}  # obs_id -> [source_id, ...]
             source_facts_dict: dict[str, MemoryFact] | None = None
+            source_facts_truncated = False
             if include_source_facts:
                 observation_ids = [uuid.UUID(sr.id) for sr in top_scored if sr.retrieval.fact_type == "observation"]
                 if observation_ids:
@@ -6610,25 +6612,35 @@ def _source_fact_dict(
                         # Resolve each observation's sources. This is a recall hot path, so the SQL
                         # store reads only the two columns it needs rather than a full memory row; a
                         # store that owns its rows answers from its own objects via one addressed read.
+                        #
+                        # Both branches keep observation-rank order: the token budget below is filled
+                        # in this order, so an unordered read would let a low-ranked observation
+                        # spend the budget the top-ranked one needs (issue #3221).
                         if store.writes_memory_rows_in_sql_for(bank_id):
                             obs_rows = [
                                 {"id": str(r["id"]), "source_memory_ids": r["source_memory_ids"]}
                                 for r in await sf_conn.fetch(
                                     f"SELECT id, source_memory_ids FROM {fq_table('memory_units')} "
-                                    f"WHERE id = ANY($1::uuid[]) AND fact_type = 'observation'",
+                                    f"WHERE id = ANY($1::uuid[]) AND fact_type = 'observation' "
+                                    f"ORDER BY array_position($1::uuid[], id)",
                                     observation_ids,
                                 )
                             ]
                         else:
-                            obs_rows = [
-                                {"id": m.unit_id, "source_memory_ids": m.source_memory_ids}
+                            obs_by_id = {
+                                m.unit_id: m
                                 for m in await store.get_memories(
                                     conn=sf_conn,
                                     fq_table=fq_table,
                                     bank_id=bank_id,
                                     unit_ids=[str(o) for o in observation_ids],
                                 )
                                 if m.fact_type == "observation"
+                            }
+                            obs_rows = [
+                                {"id": m.unit_id, "source_memory_ids": m.source_memory_ids}
+                                for m in (obs_by_id.get(str(o)) for o in observation_ids)
+                                if m is not None
                             ]
 
                         # Collect unique source IDs in order of first appearance
@@ -6691,7 +6703,6 @@ def _source_fact_dict(
                                 }
 
                             encoding = _get_tiktoken_encoding()
-                            source_facts_dict = {}
 
                             def _make_source_fact(sid: str, r: Any) -> MemoryFact:
                                 return MemoryFact(
@@ -6708,36 +6719,18 @@ def _make_source_fact(sid: str, r: Any) -> MemoryFact:
                                     tags=r["tags"] or None,
                                 )
 
-                            if max_source_facts_tokens_per_observation >= 0:
-                                # Per-observation capping: each observation independently selects
-                                # source facts up to its token budget.
-                                for obs_id, sids in source_fact_ids_by_obs.items():
-                                    obs_tokens = 0
-                                    for sid in sids:
-                                        if sid not in source_row_by_id:
-                                            continue
-                                        r = source_row_by_id[sid]
-                                        fact_tokens = len(encoding.encode(r["text"]))
-                                        if obs_tokens + fact_tokens > max_source_facts_tokens_per_observation:
-                                            break
-                                        obs_tokens += fact_tokens
-                                        if sid not in source_facts_dict:
-                                            source_facts_dict[sid] = _make_source_fact(sid, r)
-                            else:
-                                # Global budget: fill in order of first appearance until exhausted.
-                                total_source_tokens = 0
-                                for sid in source_ids_ordered:
-                                    if sid not in source_row_by_id:
-                                        continue
-                                    r = source_row_by_id[sid]
-                                    fact_tokens = len(encoding.encode(r["text"]))
-                                    if (
-                                        max_source_facts_tokens >= 0
-                                        and total_source_tokens + fact_tokens > max_source_facts_tokens
-                                    ):
-                                        break
-                                    source_facts_dict[sid] = _make_source_fact(sid, r)
-                                    total_source_tokens += fact_tokens
+                            selection = select_source_facts_within_budget(
+                                source_ids_ordered=source_ids_ordered,
+                                source_fact_ids_by_obs=source_fact_ids_by_obs,
+                                text_by_id={sid: r["text"] for sid, r in source_row_by_id.items()},
+                                max_total_tokens=max_source_facts_tokens,
+                                max_tokens_per_observation=max_source_facts_tokens_per_observation,
+                                count_tokens=lambda text: len(encoding.encode(text)),
+                            )
+                            source_facts_truncated = selection.truncated
+                            source_facts_dict = {
+                                sid: _make_source_fact(sid, source_row_by_id[sid]) for sid in selection.ids
+                            }
 
             # Source-fact enrichment is two SQL passes + tiktoken encoding; record it
             # only when requested (issue #2361).
@@ -6901,6 +6894,7 @@ def _make_source_fact(sid: str, r: Any) -> MemoryFact:
                 entities=entities_dict,
                 chunks=chunks_dict,
                 source_facts=source_facts_dict,
+                source_facts_truncated=source_facts_truncated if include_source_facts else None,
             )
 
         except OperationCancelledError:

hindsight-api-slim/hindsight_api/engine/response_models.py

@@ -353,6 +353,14 @@ class RecallResult(BaseModel):
     source_facts: dict[str, MemoryFact] | None = Field(
         None, description="Source facts for observation-type results, keyed by fact ID"
     )
+    source_facts_truncated: bool | None = Field(
+        None,
+        description=(
+            "Whether the source_facts map was cut short by the token budget. When true, some IDs in "
+            "results[].source_fact_ids have no entry in source_facts — the budget ran out, the "
+            "references are not dangling. Only set when source facts were requested."
+        ),
+    )
 
 
 class ReflectResult(BaseModel):

hindsight-api-slim/hindsight_api/engine/source_facts.py

@@ -0,0 +1,94 @@
+"""Token-budget selection for the recall ``source_facts`` map.
+
+Recall returns each observation's ``source_fact_ids`` in full and resolves them
+through a ``source_facts`` map that is filled up to a token budget. Which facts
+make it into that map is what this module decides.
+
+Two properties matter to callers resolving provenance (issue #3221):
+
+* the budget is spent in **observation-rank order**, so truncation hits the tail
+  of the result list and the top-ranked results keep their provenance;
+* one oversized fact skips itself only — it does not evict every shorter fact
+  behind it — and any skip is reported so the caller can tell a budget-truncated
+  map from a dangling reference.
+"""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class SourceFactSelection:
+    """The source facts that fit the budget, plus whether anything was dropped."""
+
+    ids: list[str]
+    """Selected source fact IDs, in the order they were considered."""
+
+    truncated: bool
+    """True when at least one resolvable source fact was skipped for budget."""
+
+
+def select_source_facts_within_budget(
+    *,
+    source_ids_ordered: list[str],
+    source_fact_ids_by_obs: dict[str, list[str]],
+    text_by_id: dict[str, str],
+    max_total_tokens: int,
+    max_tokens_per_observation: int,
+    count_tokens: Callable[[str], int],
+) -> SourceFactSelection:
+    """Pick the source facts that fit the requested token budget.
+
+    ``source_ids_ordered`` and ``source_fact_ids_by_obs`` must both be in
+    observation-rank order — the budget is spent front to back, so their order
+    decides which results keep their provenance when the budget runs out.
+
+    A non-negative ``max_tokens_per_observation`` gives every observation its own
+    budget (the global one is then unused); otherwise a single ``max_total_tokens``
+    budget is shared, with a negative value meaning unlimited. IDs absent from
+    ``text_by_id`` are skipped without counting as truncation: those rows did not
+    resolve at all, which is a different condition from running out of budget.
+    """
+    selected: list[str] = []
+    seen: set[str] = set()
+    # Skipped for budget, not necessarily lost: under a per-observation cap the same
+    # source can be dropped by one observation's budget and still fit another's. Only
+    # a skip that leaves the fact out of the final map counts as truncation, so the
+    # flag means exactly what it says — some advertised ID has no entry.
+    skipped: set[str] = set()
+
+    def _take(sid: str) -> None:
+        if sid not in seen:
+            seen.add(sid)
+            selected.append(sid)
+
+    if max_tokens_per_observation >= 0:
+        # Per-observation capping: each observation independently selects source
+        # facts up to its own budget.
+        for sids in source_fact_ids_by_obs.values():
+            obs_tokens = 0
+            for sid in sids:
+                text = text_by_id.get(sid)
+                if text is None:
+                    continue
+                fact_tokens = count_tokens(text)
+                if obs_tokens + fact_tokens > max_tokens_per_observation:
+                    skipped.add(sid)
+                    continue
+                obs_tokens += fact_tokens
+                _take(sid)
+        return SourceFactSelection(ids=selected, truncated=bool(skipped - seen))
+
+    total_tokens = 0
+    for sid in source_ids_ordered:
+        text = text_by_id.get(sid)
+        if text is None:
+            continue
+        fact_tokens = count_tokens(text)
+        if max_total_tokens >= 0 and total_tokens + fact_tokens > max_total_tokens:
+            skipped.add(sid)
+            continue
+        total_tokens += fact_tokens
+        _take(sid)
+
+    return SourceFactSelection(ids=selected, truncated=bool(skipped - seen))

hindsight-api-slim/tests/test_source_facts_selection.py

@@ -0,0 +1,143 @@
+"""Tests for the recall source_facts token-budget selection (issue #3221).
+
+The budget must be spent in observation-rank order, an oversized fact must skip
+only itself, and any budget skip must be reported so a caller can tell a
+truncated map from a dangling reference.
+"""
+
+from hindsight_api.engine.source_facts import select_source_facts_within_budget
+
+
+def _count_words(text: str) -> int:
+    """Stand-in tokenizer: one token per word, so budgets read literally."""
+    return len(text.split())
+
+
+def _select(
+    *,
+    by_obs: dict[str, list[str]],
+    texts: dict[str, str],
+    max_total_tokens: int = -1,
+    max_tokens_per_observation: int = -1,
+):
+    ordered: list[str] = []
+    seen: set[str] = set()
+    for sids in by_obs.values():
+        for sid in sids:
+            if sid not in seen:
+                seen.add(sid)
+                ordered.append(sid)
+    return select_source_facts_within_budget(
+        source_ids_ordered=ordered,
+        source_fact_ids_by_obs=by_obs,
+        text_by_id=texts,
+        max_total_tokens=max_total_tokens,
+        max_tokens_per_observation=max_tokens_per_observation,
+        count_tokens=_count_words,
+    )
+
+
+class TestGlobalBudget:
+    def test_budget_is_spent_in_rank_order(self):
+        """The top-ranked observation keeps its provenance; the tail loses it."""
+        by_obs = {"obs-rank-1": ["s1"], "obs-rank-2": ["s2"], "obs-rank-3": ["s3"]}
+        texts = {"s1": "one two", "s2": "three four", "s3": "five six"}
+
+        selection = _select(by_obs=by_obs, texts=texts, max_total_tokens=2)
+
+        assert selection.ids == ["s1"]
+        assert selection.truncated is True
+
+    def test_unlimited_budget_keeps_everything(self):
+        by_obs = {"obs-1": ["s1", "s2"], "obs-2": ["s3"]}
+        texts = {"s1": "a b c", "s2": "d e f", "s3": "g h i"}
+
+        selection = _select(by_obs=by_obs, texts=texts, max_total_tokens=-1)
+
+        assert selection.ids == ["s1", "s2", "s3"]
+        assert selection.truncated is False
+
+    def test_oversized_fact_does_not_evict_the_facts_behind_it(self):
+        """One long fact skips itself only — shorter facts behind it still fit."""
+        by_obs = {"obs-1": ["long"], "obs-2": ["short-a"], "obs-3": ["short-b"]}
+        texts = {"long": "w " * 50, "short-a": "a", "short-b": "b"}
+
+        selection = _select(by_obs=by_obs, texts=texts, max_total_tokens=2)
+
+        assert selection.ids == ["short-a", "short-b"]
+        assert selection.truncated is True
+
+    def test_shared_source_counted_once(self):
+        """A source cited by two observations is selected (and charged) once."""
+        by_obs = {"obs-1": ["shared"], "obs-2": ["shared", "s2"]}
+        texts = {"shared": "a b", "s2": "c d"}
+
+        selection = _select(by_obs=by_obs, texts=texts, max_total_tokens=4)
+
+        assert selection.ids == ["shared", "s2"]
+        assert selection.truncated is False
+
+    def test_unresolvable_id_is_not_truncation(self):
+        """A source row that did not resolve is skipped without flagging truncation."""
+        by_obs = {"obs-1": ["s1", "missing"]}
+        texts = {"s1": "a b"}
+
+        selection = _select(by_obs=by_obs, texts=texts, max_total_tokens=-1)
+
+        assert selection.ids == ["s1"]
+        assert selection.truncated is False
+
+
+class TestPerObservationBudget:
+    def test_each_observation_gets_its_own_budget(self):
+        by_obs = {"obs-1": ["s1"], "obs-2": ["s2"]}
+        texts = {"s1": "a b", "s2": "c d"}
+
+        selection = _select(by_obs=by_obs, texts=texts, max_tokens_per_observation=2)
+
+        assert selection.ids == ["s1", "s2"]
+        assert selection.truncated is False
+
+    def test_oversized_fact_does_not_evict_the_facts_behind_it(self):
+        by_obs = {"obs-1": ["long", "short"]}
+        texts = {"long": "w " * 50, "short": "a"}
+
+        selection = _select(by_obs=by_obs, texts=texts, max_tokens_per_observation=1)
+
+        assert selection.ids == ["short"]
+        assert selection.truncated is True
+
+    def test_per_observation_cap_takes_precedence_over_global(self):
+        """With a per-observation cap set, the global budget does not apply."""
+        by_obs = {"obs-1": ["s1"], "obs-2": ["s2"]}
+        texts = {"s1": "a b", "s2": "c d"}
+
+        selection = _select(
+            by_obs=by_obs,
+            texts=texts,
+            max_total_tokens=2,
+            max_tokens_per_observation=2,
+        )
+
+        assert selection.ids == ["s1", "s2"]
+        assert selection.truncated is False
+
+    def test_fact_dropped_by_one_observation_but_kept_by_another_is_not_truncation(self):
+        """A source that still lands in the map was not lost, so nothing is flagged."""
+        by_obs = {"obs-1": ["filler", "shared"], "obs-2": ["shared"]}
+        texts = {"filler": "a", "shared": "b"}
+
+        # obs-1 spends its whole budget on "filler" and skips "shared"; obs-2 has room.
+        selection = _select(by_obs=by_obs, texts=texts, max_tokens_per_observation=1)
+
+        assert set(selection.ids) == {"filler", "shared"}
+        assert selection.truncated is False
+
+    def test_zero_cap_drops_every_fact_and_reports_truncation(self):
+        by_obs = {"obs-1": ["s1"]}
+        texts = {"s1": "a"}
+
+        selection = _select(by_obs=by_obs, texts=texts, max_tokens_per_observation=0)
+
+        assert selection.ids == []
+        assert selection.truncated is True