Skip to content

fix(recall): inherit observation entities through source_memory_ids#1397

Merged
nicoloboschi merged 2 commits intovectorize-io:mainfrom
youchi1:fix/recall-observation-entities
May 5, 2026
Merged

fix(recall): inherit observation entities through source_memory_ids#1397
nicoloboschi merged 2 commits intovectorize-io:mainfrom
youchi1:fix/recall-observation-entities

Conversation

@youchi1
Copy link
Copy Markdown
Contributor

@youchi1 youchi1 commented May 3, 2026

Summary

Recall returns entities: null for every observation result, even when include_entities=true is set, while GET /v1/default/banks/{bank_id}/memories/{id} returns the same observation with its inherited entities populated. The asymmetry strips entity context (e.g. URLs, named projects) from observation-only recall and silently drops those entities from the top-level entities aggregate map.

Related: #1374 fixes the same conceptual gap on the document detail view (per-document graph + Observations tab were empty because the observation→source edge wasn't traversed). This PR is the recall-projection sibling: observation results lose their entities for the same underlying reason — code paths that need to traverse the observation→source edge weren't doing it.

Root cause

get_memory_unit (hindsight_api/engine/memory_engine.py) already handles observations explicitly: if a unit has no rows in unit_entities, it inherits entities from its source memories. The recall path's entity fetch only queried unit_entities directly; observations have no direct rows there, so every observation in the recall result lost its entities.

The inheritance edge is dialect-shaped: PG stores it on memory_units.source_memory_ids, Oracle keeps it in the observation_sources junction (same setup #1374 navigates with _observations_via_source_match_sql).

Fix

Two commits:

  1. fix(recall): inherit observation entities through source_memory_ids — mirror get_memory_unit's fallback inside recall's include_entities block: for observation result IDs with no direct unit_entities rows, look up source_memory_ids, fetch entities for the union of source IDs, project back. Downstream code derives the per-result entities field and the top-level aggregate map from the same map, so inheritance flows through both automatically.

  2. refactor(recall): consolidate observation entity inheritance in one SQL helper — replace the procedural fallback with _entity_rows_for_units_sql, a private engine helper that returns one dialect-correct UNION SELECT covering both direct and inherited entity rows. Mirrors the shape of _observations_via_source_match_sql from fix(api): include observations in per-document graph and counts #1374. Also fixes silent Oracle-parity gap (the procedural patch reached for source_memory_ids directly, which is PG-only). get_memory_unit is refactored to use the same helper, so the two call sites that hand-rolled the same inheritance now share one source of truth.

Reproduction (before fix)

# Same observation via memory get returns entities
curl -sS "$API/v1/default/banks/$BANK/memories/$OBS_ID" | jq '.entities'
# => ["HeadClaw","Reddit","waitlist users"]

# But via recall they are stripped, even with include.entities forced on
curl -sS -X POST "$API/v1/default/banks/$BANK/memories/recall" \
  -H 'Content-Type: application/json' \
  -d '{"query":"...","types":["observation"],"include":{"entities":{"max_tokens":1000}}}' \
  | jq '.results[0].entities, .entities'
# => null
# => null

A mixed-type recall (["observation","world","experience"]) populates the top-level entities map but only via world/experience contributions; observation results still come back with entities: null.

Test plan

  • New regression test (tests/test_recall_observation_entities.py) seeds:
    • A world fact with two entities linked via unit_entities,
    • One observation referencing that fact via source_memory_ids and zero direct entity links,
    • One observation with its own direct unit_entities link,
      and asserts that:
    • both observations come back from recall_async with the expected entities (inherited and direct),
    • the top-level entities aggregate map is populated with all of them,
    • get_memory_unit on the same observations also surfaces inherited and direct entities (guards against drift between the two call sites that share the helper).
  • Test fails on main with AssertionError: Observation entities must inherit through source_memory_ids; got [] and passes after the fix.
  • tests/test_observations.py and tests/test_graph_filtering.py continue to pass.
  • ruff check and ruff format clean.

youchi1 added 2 commits May 3, 2026 15:12
`include_entities=True` returns `entities: null` for every observation in
the recall response, even when those observations are linked through
`source_memory_ids` to facts whose entities are populated. The
per-memory endpoint (`get_memory_unit`) already handles this case: if an
observation has no rows in `unit_entities`, it inherits the union of
entities from its source memories. The recall path queried
`unit_entities` directly and stopped there, so observation results lost
both their per-result `entities` field and their contribution to the
top-level aggregate map.

The asymmetry made observation-only recall hostile to clients that
needed entity context (URL recovery, entity-aware ranking). The
documented workaround was to add `world` and `experience` to the
`types` filter and rely on those facts to carry the entity payload.

Mirror `get_memory_unit`'s fallback inside the recall entity-fetching
block: for observation result IDs that produced no direct
`unit_entities` rows, look up their `source_memory_ids`, fetch entities
for the union of source IDs in a single batched query, and project the
results back onto the original observation IDs (deduped by entity_id,
preserving source-memory order). The downstream code that derives
per-result `entities` and the top-level aggregate map both consume
`fact_entity_map`, so the inheritance flows through both paths
automatically.

Add a regression test that seeds an observation linked via
`source_memory_ids` to a fact carrying two entities, plus a second
observation with its own direct `unit_entities` link, then asserts
recall projects both per-result entity lists and the top-level map.
…QL helper

The first commit on this branch fixed the recall projection by mirroring
get_memory_unit's procedural fallback in Python: query unit_entities,
detect observations that came back empty, separately fetch
source_memory_ids, separately fetch entities for the union of source
IDs, then dedupe and merge in Python. That worked but had two issues
worth fixing before the PR lands.

First, the inheritance edge ("observation linked through its source
memories") is dialect-shaped: PG stores it on `memory_units.source_memory_ids`,
Oracle keeps it in the `observation_sources` junction table. The
procedural patch reached for `source_memory_ids` directly, which made
recall observation-entity inheritance silently PG-only.

Second, the same fallback already existed inline in get_memory_unit, so
shipping a second copy in recall left two places that had to stay in
sync forever, by hand.

Introduce `_entity_rows_for_units_sql`, a private engine helper that
returns a single dialect-correct UNION SELECT producing
`(unit_id, entity_id, canonical_name)` rows. Direct rows come from
`unit_entities`; observations that have no direct row inherit through
`source_memory_ids` (PG) or `observation_sources` (Oracle), guarded by
NOT EXISTS so the inheritance only fires when the direct path is empty.
This is the same conceptual shape as `_observations_via_source_match_sql`
on the document view fix branch — both are SQL primitives over the
observation-source edge.

Use the helper in two places that previously hand-rolled the same
inheritance logic:

- The recall entity-fetch block collapses from three queries plus a
  Python dedupe loop to one fetch into the same `fact_entity_map`.
- get_memory_unit's two-query "fetch direct, fall back to sources"
  pattern collapses to one fetch, with identical observable behavior.

Add a get_memory_unit assertion to the existing regression test so the
shared helper is exercised through both call sites and any future drift
between recall and the per-memory endpoint trips a test, not a
production report.
@nicoloboschi nicoloboschi merged commit 8507095 into vectorize-io:main May 5, 2026
55 of 56 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants