From cb6d1c469c35625e6cab1b522833bf5538fe03d3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Thu, 5 Mar 2026 14:12:21 +0100 Subject: [PATCH] fix: resolve chunks for observation results via source_memory_ids (#496) * fix: resolve chunks for observation results via source_memory_ids Observations have no direct chunk_id (they are synthesized from source memories). When include_chunks=True and fact_type includes 'observation', chunks were silently returned as None. Fix collects source chunk_ids via a single JOIN on source_memory_ids, using array_position to preserve observation rank order so observation source chunks are interleaved at the correct position rather than appended after all direct-fact chunks. * fix: use correct run_consolidation method name in test --- .../hindsight_api/engine/memory_engine.py | 48 ++++++++++++++-- .../tests/test_recall_chunks_independence.py | 55 +++++++++++++++++++ 2 files changed, 98 insertions(+), 5 deletions(-) diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index 5c1bcb10..cc39f645 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -2891,15 +2891,53 @@ class MemoryEngine(MemoryEngineInterface): if include_chunks and top_scored: from .response_models import ChunkInfo - # Collect chunk_ids in order of fact relevance (preserving order from top_scored) - # Use a list to maintain order, but track seen chunks to avoid duplicates - chunk_ids_ordered = [] - seen_chunk_ids = set() + # Collect chunk_ids in order of fact relevance (preserving order from top_scored). + # Observations have no direct chunk_id — use a placeholder so their source + # chunks end up at the observation's rank position, not appended at the end. + # ordered_items: list of ('chunk', chunk_id) | ('obs', sr.id) + ordered_items: list[tuple[str, str]] = [] + seen_chunk_ids: set[str] = set() + observation_ids_ordered: list[uuid.UUID] = [] for sr in top_scored: chunk_id = sr.retrieval.chunk_id if chunk_id and chunk_id not in seen_chunk_ids: - chunk_ids_ordered.append(chunk_id) + ordered_items.append(("chunk", chunk_id)) seen_chunk_ids.add(chunk_id) + elif not chunk_id and sr.retrieval.fact_type == "observation": + ordered_items.append(("obs", sr.id)) + observation_ids_ordered.append(uuid.UUID(sr.id)) + + # Resolve source chunk_ids for all observations in a single query, + # ordered by observation rank so per-observation results stay grouped correctly. + obs_chunk_ids: dict[str, list[str]] = {} + if observation_ids_ordered: + async with acquire_with_retry(pool) as obs_conn: + obs_source_rows = await obs_conn.fetch( + f""" + SELECT obs.id AS obs_id, mu.chunk_id + FROM {fq_table("memory_units")} obs + JOIN {fq_table("memory_units")} mu + ON mu.id = ANY(obs.source_memory_ids) + WHERE obs.id = ANY($1::uuid[]) + AND mu.chunk_id IS NOT NULL + ORDER BY array_position($1::uuid[], obs.id) + """, + observation_ids_ordered, + ) + for row in obs_source_rows: + obs_id = str(row["obs_id"]) + cid = row["chunk_id"] + if cid not in seen_chunk_ids: + obs_chunk_ids.setdefault(obs_id, []).append(cid) + seen_chunk_ids.add(cid) + + # Flatten ordered_items into chunk_ids_ordered, expanding obs placeholders + chunk_ids_ordered = [] + for item_type, item_id in ordered_items: + if item_type == "chunk": + chunk_ids_ordered.append(item_id) + else: + chunk_ids_ordered.extend(obs_chunk_ids.get(item_id, [])) if chunk_ids_ordered: chunks_dict = {} diff --git a/hindsight-api/tests/test_recall_chunks_independence.py b/hindsight-api/tests/test_recall_chunks_independence.py index e972d791..a3d5738b 100644 --- a/hindsight-api/tests/test_recall_chunks_independence.py +++ b/hindsight-api/tests/test_recall_chunks_independence.py @@ -231,6 +231,61 @@ async def test_recall_chunks_ordering_by_relevance(memory, request_context): await memory.delete_bank(bank_id, request_context=request_context) +@pytest.mark.asyncio +async def test_recall_chunks_for_observations(memory, request_context): + """ + Test that chunks are returned when recalling only observations. + + Observations have no direct chunk_id (they are synthesized from source memories). + When include_chunks=True, chunks should be resolved via source_memory_ids. + """ + bank_id = "test-chunks-observations" + + try: + # Retain content that will generate observations via consolidation + test_content = """ + Alice is a senior software engineer at a large technology company. + She specializes in distributed systems and has 10 years of experience. + Alice leads a team of 8 engineers working on cloud infrastructure. + She holds a PhD in computer science from Stanford University. + Alice has published several papers on fault-tolerant distributed systems. + """ * 8 + + await memory.retain_async( + bank_id=bank_id, + content=test_content, + context="profile notes", + request_context=request_context, + ) + + # Trigger consolidation explicitly to ensure observations exist + await memory.run_consolidation(bank_id=bank_id, request_context=request_context) + + # Recall observations only with chunks enabled + result = await memory.recall_async( + bank_id=bank_id, + query="Alice software engineer", + fact_type=["observation"], + max_tokens=4096, + include_chunks=True, + max_chunk_tokens=2000, + budget=Budget.MID, + request_context=request_context, + ) + + # If observations were created, chunks should be resolved from source memories + if len(result.results) > 0: + assert result.chunks is not None, "Should include chunks dict when observations are found" + assert len(result.chunks) > 0, "Should return chunks resolved from observation source memories" + + for chunk_id, chunk_info in result.chunks.items(): + assert len(chunk_info.chunk_text) > 0, "Chunks should contain text" + assert chunk_info.chunk_index >= 0, "Chunk should have valid index" + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + @pytest.mark.asyncio async def test_recall_chunks_without_include_flag(memory, request_context): """