diff --git a/src/basic_memory/api/v2/utils.py b/src/basic_memory/api/v2/utils.py index 5c2b101f..61aac74e 100644 --- a/src/basic_memory/api/v2/utils.py +++ b/src/basic_memory/api/v2/utils.py @@ -18,7 +18,7 @@ from basic_memory.services.context_service import ( class EntityBatchLookup(Protocol): - async def find_by_ids(self, ids: List[int]) -> Sequence[Any]: ... + async def find_by_ids_for_hydration(self, ids: List[int]) -> Sequence[Any]: ... class EntityServiceBatchLookup(Protocol): @@ -76,7 +76,7 @@ async def to_graph_context( if item.to_id: entity_ids_needed.add(item.to_id) - # Batch fetch all entities at once - get both title and external_id + # Batch fetch just the entity fields needed to shape the response. entity_title_lookup: dict[int, str] = {} entity_external_id_lookup: dict[int, str] = {} if entity_ids_needed: @@ -87,7 +87,9 @@ async def to_graph_context( phase="lookup_entities", result_count=len(entity_ids_needed), ): - entities = await entity_repository.find_by_ids(list(entity_ids_needed)) + entities = await entity_repository.find_by_ids_for_hydration( + list(entity_ids_needed) + ) for e in entities: entity_title_lookup[e.id] = e.title entity_external_id_lookup[e.id] = e.external_id diff --git a/src/basic_memory/repository/entity_repository.py b/src/basic_memory/repository/entity_repository.py index a07caa14..62adeddd 100644 --- a/src/basic_memory/repository/entity_repository.py +++ b/src/basic_memory/repository/entity_repository.py @@ -8,7 +8,7 @@ from loguru import logger from sqlalchemy import select, func from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker -from sqlalchemy.orm import selectinload +from sqlalchemy.orm import load_only, selectinload from sqlalchemy.orm.interfaces import LoaderOption from sqlalchemy.engine import Row @@ -178,6 +178,24 @@ class EntityRepository(Repository[Entity]): result = await self.execute_query(query, use_query_options=False) return list(result.scalars().all()) + async def find_by_ids_for_hydration(self, ids: List[int]) -> Sequence[Entity]: + """Fetch minimal entity fields needed for context hydration. + + Context hydration only needs an entity's primary key, title, and external + UUID. Keeping this separate from find_by_ids avoids the relationship eager + loads that are useful for full entity reads but expensive for response shaping. + """ + if not ids: + return [] + + query = ( + self.select() + .where(Entity.id.in_(ids)) + .options(load_only(Entity.id, Entity.title, Entity.external_id)) + ) + result = await self.execute_query(query, use_query_options=False) + return list(result.scalars().all()) + async def get_permalink_to_file_path_map(self) -> dict[str, str]: """Get a mapping of permalink -> file_path for all entities. diff --git a/tests/api/v2/test_memory_hydration.py b/tests/api/v2/test_memory_hydration.py index 5154c394..e5275982 100644 --- a/tests/api/v2/test_memory_hydration.py +++ b/tests/api/v2/test_memory_hydration.py @@ -55,6 +55,25 @@ class SpyEntityRepository: self.calls.append(ids) return [self.entities_by_id[i] for i in ids if i in self.entities_by_id] + async def find_by_ids_for_hydration(self, ids: list[int]): + self.calls.append(ids) + return [self.entities_by_id[i] for i in ids if i in self.entities_by_id] + + +class LightweightOnlyEntityRepository: + """Raises if graph hydration uses the eager-loading repository method.""" + + def __init__(self, entities_by_id: dict[int, SimpleNamespace]): + self.entities_by_id = entities_by_id + self.hydration_calls: list[list[int]] = [] + + async def find_by_ids(self, ids: list[int]): + raise AssertionError("graph hydration must use the lightweight hydration lookup") + + async def find_by_ids_for_hydration(self, ids: list[int]): + self.hydration_calls.append(ids) + return [self.entities_by_id[i] for i in ids if i in self.entities_by_id] + # --- Single batch fetch (N+1 elimination) --- @@ -198,3 +217,63 @@ async def test_to_graph_context_empty_results_skip_entity_lookup(): assert repo.calls == [] assert list(graph.results) == [] + + +@pytest.mark.asyncio +async def test_to_graph_context_uses_lightweight_hydration_lookup(): + """Hydration should not load observations/relations when only entity fields are needed.""" + repo = LightweightOnlyEntityRepository( + { + 1: _make_entity(1, "Root", "ext-root"), + 2: _make_entity(2, "Child", "ext-child"), + } + ) + now = datetime.now(timezone.utc) + + context = ServiceContextResult( + results=[ + ContextResultItem( + primary_result=_make_row( + type="entity", + id=1, + root_id=1, + title="Root", + permalink="notes/root", + file_path="notes/root.md", + created_at=now, + ), + observations=[], + related_results=[ + _make_row( + type="relation", + id=20, + root_id=1, + title="links_to: Child", + permalink="notes/root", + file_path="notes/root.md", + relation_type="links_to", + from_id=1, + to_id=2, + depth=1, + created_at=now, + ) + ], + ) + ], + metadata=ContextMetadata( + types=[SearchItemType.ENTITY, SearchItemType.RELATION], + depth=1, + primary_count=1, + related_count=1, + total_relations=1, + ), + ) + + graph = await to_graph_context(context, entity_repository=repo) + + assert len(repo.hydration_calls) == 1 + assert set(repo.hydration_calls[0]) == {1, 2} + relation = graph.results[0].related_results[0] + assert isinstance(relation, RelationSummary) + assert relation.from_entity_external_id == "ext-root" + assert relation.to_entity_external_id == "ext-child" diff --git a/tests/repository/test_entity_repository.py b/tests/repository/test_entity_repository.py index 8a47d76d..1351e6c8 100644 --- a/tests/repository/test_entity_repository.py +++ b/tests/repository/test_entity_repository.py @@ -1050,6 +1050,25 @@ async def test_get_all_permalinks(entity_repository: EntityRepository, session_m assert isinstance(permalink, str) +@pytest.mark.asyncio +async def test_find_by_ids_for_hydration_skips_eager_load_options( + entity_repository: EntityRepository, sample_entity: Entity, monkeypatch: pytest.MonkeyPatch +): + """Context hydration should bypass relationship loader options.""" + + def fail_get_load_options(): + raise AssertionError("hydration lookup must not eager load entity relationships") + + monkeypatch.setattr(entity_repository, "get_load_options", fail_get_load_options) + + found = await entity_repository.find_by_ids_for_hydration([sample_entity.id]) + + assert len(found) == 1 + assert found[0].id == sample_entity.id + assert found[0].title == sample_entity.title + assert found[0].external_id == sample_entity.external_id + + @pytest.mark.asyncio async def test_get_permalink_to_file_path_map(entity_repository: EntityRepository, session_maker): """Test getting permalink -> file_path mapping for bulk operations."""