holetron-lab/fleet-memory
1
0
Fork 0
Code Releases Packages Activity Actions

fix: resolve chunks for observation results via source_memory_ids (#496)

Browse source 
* 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
This commit is contained in:
Nicolò Boschi 2026-03-05 14:12:21 +01:00 committed by GitHub
parent 4c058b4b98
commit cb6d1c469c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 98 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
hindsight-api/hindsight_api/engine/memory_engine.py
View file

@ -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 = {}

55
hindsight-api/tests/test_recall_chunks_independence.py
View file

@ -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):
"""