feat(reflect): resolve entity names on reflect sub-recalls (#3334)

recall_async only populates each result's entities field when
include_entities=True, and it defaults to False — so reflect's recall and
search_observations tools never surfaced them. Canonical entity names are
semantic signal the surface text may lack ("Bob" in the text vs canonical
"Robert Smith"): they give the agent resolved names to cite and to pivot
follow-up queries on, for the cost of one extra lookup query per recall.

The top-level EntityState dict recall also builds is not serialized into
tool results; only the per-fact names reach the agent.
This commit is contained in:
Nicolò Boschi
2026-08-10 13:59:02 +02:00
committed by GitHub
parent e9a14da690
commit e852acc7c0
2 changed files with 90 additions and 0 deletions
@@ -257,6 +257,11 @@ async def tool_search_observations(
tags_match=tags_match,
tag_groups=tag_groups,
include_source_facts=include_source_facts,
# Canonical entity names are semantic signal the surface text may lack
# ("Bob" in the text vs canonical "Robert Smith"): they populate each
# result's `entities` field, giving the agent resolved names to cite
# and to pivot follow-up queries on.
include_entities=True,
created_after=created_after,
created_before=created_before,
_connection_budget=1,
@@ -337,6 +342,9 @@ async def tool_recall(
tag_groups=tag_groups,
created_after=created_after,
created_before=created_before,
# See tool_search_observations: resolved entity names on each result
# are worth the one extra lookup query.
include_entities=True,
_connection_budget=connection_budget,
_quiet=True, # Suppress logging for internal operations
include_chunks=include_chunks,
@@ -0,0 +1,82 @@
"""Reflect sub-recalls must request entity resolution.
Canonical entity names are semantic signal the surface text may lack ("Bob"
in the text vs canonical "Robert Smith"). `recall_async` only populates each
result's `entities` field when `include_entities=True`, and it defaults to
False — so both reflect retrieval tools must pass it explicitly, and the
names must survive the serialization into the tool result the agent reads.
"""
from __future__ import annotations
from dataclasses import dataclass
from unittest.mock import AsyncMock, MagicMock
import pytest
from hindsight_api.engine.reflect.tools import tool_recall, tool_search_observations
from hindsight_api.engine.response_models import MemoryFact, RecallResult
@dataclass
class _FakeRequestContext:
"""Dataclass stand-in matching the fields used by ``dataclasses.replace``."""
api_key: str | None = None
api_key_id: str | None = None
tenant_id: str | None = None
internal: bool = False
mcp_authenticated: bool = False
user_initiated: bool = False
allowed_bank_ids: list[str] | None = None
def _fact_with_entities() -> MemoryFact:
return MemoryFact(
id="123e4567-e89b-12d3-a456-426614174000",
text="Bob moved the deploy to 09:00 UTC.",
fact_type="world",
entities=["Robert Smith"],
)
def _mock_engine(results: list[MemoryFact] | None = None):
engine = MagicMock()
engine.recall_async = AsyncMock(return_value=RecallResult(results=results or [], source_facts={}))
return engine
class TestReflectRecallRequestsEntities:
@pytest.mark.asyncio
async def test_recall_passes_include_entities(self):
engine = _mock_engine()
await tool_recall(engine, "bank-1", "query", _FakeRequestContext())
assert engine.recall_async.call_args.kwargs["include_entities"] is True
@pytest.mark.asyncio
async def test_search_observations_passes_include_entities(self):
engine = _mock_engine()
await tool_search_observations(engine, "bank-1", "query", _FakeRequestContext())
assert engine.recall_async.call_args.kwargs["include_entities"] is True
@pytest.mark.asyncio
async def test_entity_names_reach_the_agent(self):
"""End-to-end through the tool's serialization (null-pruning, field
trimming): the canonical names land in the payload the agent reads."""
engine = _mock_engine(results=[_fact_with_entities()])
result = await tool_recall(engine, "bank-1", "query", _FakeRequestContext())
assert result["memories"][0]["entities"] == ["Robert Smith"]
@pytest.mark.asyncio
async def test_observation_entity_names_reach_the_agent(self):
engine = _mock_engine(results=[_fact_with_entities()])
result = await tool_search_observations(engine, "bank-1", "query", _FakeRequestContext())
assert result["observations"][0]["entities"] == ["Robert Smith"]