fix: build_context related_results schema validation failure (#631)

phernandez · claude · web-flow · commit e97eafa55a29 · 2026-02-27T12:30:44.000-06:00
Signed-off-by: phernandez &lt;paul@basicmachines.co&gt;
Co-authored-by: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/src/basic_memory/alembic/versions/j3d4e5f6g7h8_rename_entity_type_to_note_type.py b/src/basic_memory/alembic/versions/j3d4e5f6g7h8_rename_entity_type_to_note_type.py
@@ -22,10 +22,7 @@ def table_exists(connection, table_name: str) -> bool:
     """Check if a table exists (idempotent migration support)."""
     if connection.dialect.name == "postgresql":
         result = connection.execute(
-            text(
-                "SELECT 1 FROM information_schema.tables "
-                "WHERE table_name = :table_name"
-            ),
+            text("SELECT 1 FROM information_schema.tables WHERE table_name = :table_name"),
             {"table_name": table_name},
         )
         return result.fetchone() is not None
diff --git a/src/basic_memory/mcp/prompts/continue_conversation.py b/src/basic_memory/mcp/prompts/continue_conversation.py
@@ -125,7 +125,7 @@ def _format_continuation_results(result: SearchResponse, topic: str) -> str:
         lines.append(f"### {title}")
         if permalink:
             lines.append(f"permalink: {permalink}")
-            lines.append(f"Read with: `read_note(\"{permalink}\")`")
+            lines.append(f'Read with: `read_note("{permalink}")`')
         if item.content:
             content = item.content[:300] + "..." if len(item.content) > 300 else item.content
             lines.append(f"\n{content}")
diff --git a/src/basic_memory/mcp/tools/build_context.py b/src/basic_memory/mcp/tools/build_context.py
@@ -22,74 +22,6 @@
     RelationSummary,
 )
 
-# --- Fields to strip from each model (redundant with parent entity) ---
-
-_OBSERVATION_STRIP = {
-    "observation_id",
-    "entity_id",
-    "entity_external_id",
-    "title",
-    "file_path",
-    "created_at",
-}
-_RELATION_STRIP = {
-    "relation_id",
-    "entity_id",
-    "from_entity_id",
-    "from_entity_external_id",
-    "to_entity_id",
-    "to_entity_external_id",
-    "title",
-    "file_path",
-    "created_at",
-}
-_ENTITY_STRIP = {"entity_id", "created_at"}
-_METADATA_STRIP = {"total_results", "generated_at"}
-
-
-def _slim_summary(summary: EntitySummary | RelationSummary | ObservationSummary) -> dict:
-    """Strip redundant fields from a summary model based on its type."""
-    if isinstance(summary, ObservationSummary):
-        strip = _OBSERVATION_STRIP
-    elif isinstance(summary, RelationSummary):
-        strip = _RELATION_STRIP
-    else:
-        strip = _ENTITY_STRIP
-
-    data = summary.model_dump()
-    for key in strip:
-        data.pop(key, None)
-    return data
-
-
-def _slim_context(graph: GraphContext) -> dict:
-    """Transform GraphContext into a slimmed dict, stripping redundant fields.
-
-    Reduces payload size ~40% by removing fields on nested objects that
-    duplicate information already present on the parent entity (IDs,
-    timestamps, file paths).
-    """
-    slimmed_results = []
-    for result in graph.results:
-        slimmed_results.append(
-            {
-                "primary_result": _slim_summary(result.primary_result),
-                "observations": [_slim_summary(obs) for obs in result.observations],
-                "related_results": [_slim_summary(rel) for rel in result.related_results],
-            }
-        )
-
-    metadata = graph.metadata.model_dump()
-    for key in _METADATA_STRIP:
-        metadata.pop(key, None)
-
-    return {
-        "results": slimmed_results,
-        "metadata": metadata,
-        "page": graph.page,
-        "page_size": graph.page_size,
-    }
-
 
 def _format_entity_block(result: ContextResult) -> str:
     """Format a single context result as a markdown block."""
@@ -194,7 +126,7 @@ def _format_context_markdown(graph: GraphContext, project: str) -> str:
     - Or standard formats like "7d", "24h"
 
     Format options:
-    - "json" (default): Slimmed JSON with redundant fields removed
+    - "json" (default): Structured JSON with internal fields excluded
     - "text": Compact markdown text for LLM consumption
     """,
     annotations={"readOnlyHint": True, "openWorldHint": False},
@@ -231,12 +163,12 @@ async def build_context(
         page: Page number of results to return (default: 1)
         page_size: Number of results to return per page (default: 10)
         max_related: Maximum number of related results to return (default: 10)
-        output_format: Response format - "json" for slimmed JSON dict,
+        output_format: Response format - "json" for structured JSON dict,
             "text" for compact markdown text
         context: Optional FastMCP context for performance caching.
 
     Returns:
-        dict (output_format="json"): Slimmed JSON with redundant fields removed
+        dict (output_format="json"): Structured JSON with internal fields excluded
         str (output_format="text"): Compact markdown representation
 
     Examples:
@@ -292,4 +224,4 @@ async def build_context(
         if output_format == "text":
             return _format_context_markdown(graph, active_project.name)
 
-        return _slim_context(graph)
+        return graph.model_dump()
diff --git a/src/basic_memory/mcp/tools/search.py b/src/basic_memory/mcp/tools/search.py
@@ -466,7 +466,12 @@ async def search_notes(
             effective_query = (query or "").strip()
             if effective_query:
                 valid_search_types = {
-                    "text", "title", "permalink", "vector", "semantic", "hybrid",
+                    "text",
+                    "title",
+                    "permalink",
+                    "vector",
+                    "semantic",
+                    "hybrid",
                 }
                 if effective_search_type == "text":
                     search_query.text = effective_query
diff --git a/src/basic_memory/repository/observation_repository.py b/src/basic_memory/repository/observation_repository.py
@@ -2,9 +2,10 @@
 
 from typing import Dict, List, Sequence
 
-
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import async_sessionmaker
+from sqlalchemy.orm import selectinload
+from sqlalchemy.orm.interfaces import LoaderOption
 
 from basic_memory.models import Observation
 from basic_memory.repository.repository import Repository
@@ -22,6 +23,10 @@ def __init__(self, session_maker: async_sessionmaker, project_id: int):
         """
         super().__init__(session_maker, Observation, project_id=project_id)
 
+    def get_load_options(self) -> List[LoaderOption]:
+        """Eager-load parent entity to prevent N+1 if obs.entity is accessed."""
+        return [selectinload(Observation.entity)]
+
     async def find_by_entity(self, entity_id: int) -> Sequence[Observation]:
         """Find all observations for a specific entity."""
         query = select(Observation).filter(Observation.entity_id == entity_id)
diff --git a/src/basic_memory/repository/postgres_search_repository.py b/src/basic_memory/repository/postgres_search_repository.py
@@ -689,9 +689,7 @@ async def search(
             for idx, note_type in enumerate(note_types):
                 param_name = f"note_type_{idx}"
                 params[param_name] = json.dumps({"note_type": note_type})
-                type_conditions.append(
-                    f"search_index.metadata @> CAST(:{param_name} AS jsonb)"
-                )
+                type_conditions.append(f"search_index.metadata @> CAST(:{param_name} AS jsonb)")
             conditions.append(f"({' OR '.join(type_conditions)})")
 
         # Handle date filter
diff --git a/src/basic_memory/schemas/memory.py b/src/basic_memory/schemas/memory.py
@@ -125,7 +125,7 @@ class EntitySummary(BaseModel):
 
     type: Literal["entity"] = "entity"
     external_id: str  # UUID for v2 API routing
-    entity_id: int  # Database ID for v2 API consistency
+    entity_id: Optional[int] = Field(None, exclude=True)  # Internal DB ID
     permalink: Optional[str]
     title: str
     content: Optional[str] = None
@@ -143,18 +143,18 @@ class RelationSummary(BaseModel):
     """Simplified relation representation."""
 
     type: Literal["relation"] = "relation"
-    relation_id: int  # Database ID for v2 API consistency
-    entity_id: Optional[int] = None  # ID of the entity this relation belongs to
+    relation_id: Optional[int] = Field(None, exclude=True)  # Internal DB ID
+    entity_id: Optional[int] = Field(None, exclude=True)  # Internal FK
     title: str
     file_path: str
     permalink: str
     relation_type: str
     from_entity: Optional[str] = None
-    from_entity_id: Optional[int] = None  # ID of source entity
-    from_entity_external_id: Optional[str] = None  # UUID of source entity for v2 API routing
+    from_entity_id: Optional[int] = Field(None, exclude=True)  # Internal FK
+    from_entity_external_id: Optional[str] = Field(None, exclude=True)  # Internal routing ID
     to_entity: Optional[str] = None
-    to_entity_id: Optional[int] = None  # ID of target entity
-    to_entity_external_id: Optional[str] = None  # UUID of target entity for v2 API routing
+    to_entity_id: Optional[int] = Field(None, exclude=True)  # Internal FK
+    to_entity_external_id: Optional[str] = Field(None, exclude=True)  # Internal routing ID
     created_at: Annotated[
         datetime, Field(json_schema_extra={"type": "string", "format": "date-time"})
     ]
@@ -168,10 +168,10 @@ class ObservationSummary(BaseModel):
     """Simplified observation representation."""
 
     type: Literal["observation"] = "observation"
-    observation_id: int  # Database ID for v2 API consistency
-    entity_id: Optional[int] = None  # ID of the entity this observation belongs to
-    entity_external_id: Optional[str] = None  # UUID of parent entity for v2 API routing
-    title: str
+    observation_id: Optional[int] = Field(None, exclude=True)  # Internal DB ID
+    entity_id: Optional[int] = Field(None, exclude=True)  # Internal FK
+    entity_external_id: Optional[str] = Field(None, exclude=True)  # Internal routing ID
+    title: Optional[str] = Field(None, exclude=True)  # Redundant with parent entity
     file_path: str
     permalink: str
     category: str
@@ -192,19 +192,13 @@ class MemoryMetadata(BaseModel):
     types: Optional[List[SearchItemType]] = None
     depth: int
     timeframe: Optional[str] = None
-    generated_at: Annotated[
-        datetime, Field(json_schema_extra={"type": "string", "format": "date-time"})
-    ]
+    generated_at: Optional[datetime] = Field(None, exclude=True)  # Internal timing
     primary_count: Optional[int] = None  # Changed field name
     related_count: Optional[int] = None  # Changed field name
-    total_results: Optional[int] = None  # For backward compatibility
+    total_results: Optional[int] = Field(None, exclude=True)  # Internal counter
     total_relations: Optional[int] = None
     total_observations: Optional[int] = None
 
-    @field_serializer("generated_at")
-    def serialize_generated_at(self, dt: datetime) -> str:
-        return dt.isoformat()
-
 
 class ContextResult(BaseModel):
     """Context result containing a primary item with its observations and related items."""
diff --git a/src/basic_memory/services/project_service.py b/src/basic_memory/services/project_service.py
@@ -977,16 +977,13 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus:
                 total_indexed_entities=total_indexed_entities,
                 vector_tables_exist=False,
                 reindex_recommended=True,
-                reindex_reason=(
-                    "Vector tables not initialized — run: bm reindex --embeddings"
-                ),
+                reindex_reason=("Vector tables not initialized — run: bm reindex --embeddings"),
             )
 
         # --- Count queries (tables exist) ---
         si_result = await self.repository.execute_query(
             text(
-                "SELECT COUNT(DISTINCT entity_id) FROM search_index "
-                "WHERE project_id = :project_id"
+                "SELECT COUNT(DISTINCT entity_id) FROM search_index WHERE project_id = :project_id"
             ),
             {"project_id": project_id},
         )
@@ -1040,9 +1037,7 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus:
                 "WHERE c.project_id = :project_id AND e.rowid IS NULL"
             )
 
-        orphan_result = await self.repository.execute_query(
-            orphan_sql, {"project_id": project_id}
-        )
+        orphan_result = await self.repository.execute_query(orphan_sql, {"project_id": project_id})
         orphaned_chunks = orphan_result.scalar() or 0
 
         # --- Reindex recommendation logic (priority order) ---
@@ -1051,9 +1046,7 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus:
 
         if total_indexed_entities > 0 and total_chunks == 0:
             reindex_recommended = True
-            reindex_reason = (
-                "Embeddings have never been built — run: bm reindex --embeddings"
-            )
+            reindex_reason = "Embeddings have never been built — run: bm reindex --embeddings"
         elif orphaned_chunks > 0:
             reindex_recommended = True
             reindex_reason = (
@@ -1063,9 +1056,7 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus:
         elif total_indexed_entities > total_entities_with_chunks:
             missing = total_indexed_entities - total_entities_with_chunks
             reindex_recommended = True
-            reindex_reason = (
-                f"{missing} entities missing embeddings — run: bm reindex --embeddings"
-            )
+            reindex_reason = f"{missing} entities missing embeddings — run: bm reindex --embeddings"
 
         return EmbeddingStatus(
             semantic_search_enabled=True,
diff --git a/test-int/cli/test_cli_tool_edit_note_integration.py b/test-int/cli/test_cli_tool_edit_note_integration.py
@@ -262,7 +262,14 @@ def test_edit_note_json_format_contract(app, app_config, test_project, config_ma
 
     assert result.exit_code == 0, result.output
     data = json.loads(result.stdout)
-    assert set(data.keys()) == {"title", "permalink", "file_path", "operation", "checksum", "fileCreated"}
+    assert set(data.keys()) == {
+        "title",
+        "permalink",
+        "file_path",
+        "operation",
+        "checksum",
+        "fileCreated",
+    }
     assert data["operation"] == "append"
     assert data["fileCreated"] is False
     assert data["title"] == "Edit JSON Note"
diff --git a/tests/api/v2/test_knowledge_router.py b/tests/api/v2/test_knowledge_router.py
@@ -836,9 +836,7 @@ async def test_delete_directory_v2_nested_structure(client: AsyncClient, v2_proj
 
 
 @pytest.mark.asyncio
-async def test_entity_response_includes_user_tracking_fields(
-    client: AsyncClient, v2_project_url
-):
+async def test_entity_response_includes_user_tracking_fields(client: AsyncClient, v2_project_url):
     """EntityResponseV2 includes created_by and last_updated_by fields (null for local)."""
     entity_data = {
         "title": "UserTrackingTest",
diff --git a/tests/mcp/test_tool_build_context.py b/tests/mcp/test_tool_build_context.py
@@ -9,7 +9,7 @@
 
 @pytest.mark.asyncio
 async def test_get_basic_discussion_context(client, test_graph, test_project):
-    """Test getting basic discussion context returns slimmed JSON dict."""
+    """Test getting basic discussion context returns JSON dict with excluded fields removed."""
     result = await build_context(project=test_project.name, url="memory://test/root")
 
     assert isinstance(result, dict)
@@ -19,7 +19,7 @@ async def test_get_basic_discussion_context(client, test_graph, test_project):
     assert primary["permalink"] == f"{test_project.name}/test/root"
     assert len(result["results"][0]["related_results"]) > 0
 
-    # Verify metadata — stripped fields should be absent
+    # Verify metadata — excluded fields should be absent
     meta = result["metadata"]
     assert meta["uri"] == f"{test_project.name}/test/root"
     assert meta["depth"] == 1  # default depth
@@ -28,16 +28,39 @@ async def test_get_basic_discussion_context(client, test_graph, test_project):
     assert "generated_at" not in meta
     assert "total_results" not in meta
 
-    # Verify entity-level stripped fields
+    # Entity: entity_id excluded, created_at kept (needed for related results)
     assert "entity_id" not in primary
-    assert "created_at" not in primary
+    assert "created_at" in primary
 
-    # Verify observation-level stripped fields
+    # Verify observation-level fields: internal IDs excluded, file_path/created_at kept
     if result["results"][0]["observations"]:
         obs = result["results"][0]["observations"][0]
         assert "observation_id" not in obs
         assert "entity_id" not in obs
-        assert "file_path" not in obs
+        assert "title" not in obs
+        # file_path and created_at kept (needed when observation is primary_result)
+        assert "file_path" in obs
+        assert "created_at" in obs
+        # Other kept fields
+        assert "permalink" in obs
+        assert "category" in obs
+        assert "content" in obs
+
+    # Verify related_results item structure — entities have identifying fields
+    for related in result["results"][0]["related_results"]:
+        item_type = related["type"]
+        if item_type == "entity":
+            assert "title" in related
+            assert "file_path" in related
+            assert "created_at" in related
+            assert "entity_id" not in related  # excluded
+        elif item_type == "relation":
+            assert "relation_type" in related
+            assert "title" in related
+            assert "file_path" in related
+            assert "created_at" in related
+            assert "relation_id" not in related  # excluded
+            assert "entity_id" not in related  # excluded
 
 
 @pytest.mark.asyncio
diff --git a/tests/repository/test_hybrid_fusion.py b/tests/repository/test_hybrid_fusion.py
@@ -248,8 +248,11 @@ async def test_dual_source_result_keeps_vector_matched_chunk():
     ]
     vector_results = [
         FakeRow(
-            id=1, score=0.9, title="both",
-            content_snippet=content, matched_chunk_text=vector_chunk,
+            id=1,
+            score=0.9,
+            title="both",
+            content_snippet=content,
+            matched_chunk_text=vector_chunk,
         ),
     ]
 
diff --git a/tests/repository/test_postgres_search_repository.py b/tests/repository/test_postgres_search_repository.py
diff --git a/tests/schemas/test_memory_serialization.py b/tests/schemas/test_memory_serialization.py
diff --git a/tests/services/test_project_service_embedding_status.py b/tests/services/test_project_service_embedding_status.py
