improve build_context results

src/basic_memory/api/routers/memory_router.py

@@ -1,6 +1,5 @@
 """Routes for memory:// URI operations."""
 
-from dataclasses import asdict
 from datetime import datetime, timedelta
 from typing import Optional
 
@@ -8,9 +7,18 @@ from fastapi import APIRouter
 from loguru import logger
 
 from basic_memory.config import config
-from basic_memory.deps import ContextServiceDep
-from basic_memory.schemas.memory import MemoryUrl, GraphContext
-from basic_memory.schemas.search import SearchResult, RelatedResult
+from basic_memory.deps import ContextServiceDep, EntityRepositoryDep
+from basic_memory.repository.search_repository import SearchIndexRow
+from basic_memory.schemas.memory import (
+    MemoryUrl,
+    GraphContext,
+    RelationSummary,
+    EntitySummary,
+    ObservationSummary,
+    MemoryMetadata,
+)
+from basic_memory.schemas.search import SearchItemType
+from basic_memory.services.context_service import ContextResultRow
 
 router = APIRouter(prefix="/memory", tags=["memory"])
 
@@ -36,6 +44,7 @@ def parse_timeframe(timeframe: str) -> Optional[datetime]:
 @router.get("/{uri:path}", response_model=GraphContext)
 async def get_memory_context(
     context_service: ContextServiceDep,
+    entity_repository: EntityRepositoryDep,
     uri: str,
     depth: int = 1,
     timeframe: str = "7d",
@@ -44,7 +53,9 @@ async def get_memory_context(
     """Get rich context from memory:// URI."""
     # add the project name from the config to the url as the "host
     # Parse URI
-    logger.debug(f"Getting context for URI: `{uri}` depth: `{depth}` timeframe: `{timeframe}` max_results: `{max_results}`")
+    logger.debug(
+        f"Getting context for URI: `{uri}` depth: `{depth}` timeframe: `{timeframe}` max_results: `{max_results}`"
+    )
     memory_url = MemoryUrl(f"memory://{config.project}/{uri}")
 
     # Parse timeframe
@@ -55,10 +66,39 @@ async def get_memory_context(
         memory_url, depth=depth, since=since, max_results=max_results
     )
 
-    primary_results = [SearchResult(**asdict(r)) for r in context["primary_results"]]
-    related_results = [RelatedResult(**asdict(r)) for r in context["related_results"]]
-    metadata = context["metadata"]
+    # return results
+    async def to_summary(item: SearchIndexRow | ContextResultRow):
+        match item.type:
+            case SearchItemType.ENTITY:
+                return EntitySummary(
+                    title=item.title,
+                    permalink=item.permalink,
+                    file_path=item.file_path,
+                    created_at=item.created_at,
+                )
+            case SearchItemType.OBSERVATION:
+                return ObservationSummary(
+                    category=item.category, content=item.content, permalink=item.permalink
+                )
+            case SearchItemType.RELATION:
+                
+                from_entity = await entity_repository.find_by_id(item.from_id)
+                to_entity = await entity_repository.find_by_id(item.to_id)
+                
+                return RelationSummary(
+                    permalink=item.permalink,
+                    type=item.type,
+                    from_id=from_entity.permalink,
+                    to_id=to_entity.permalink,
+                    created_at=item.created_at,
+                )
+
+    primary_results = [await to_summary(r) for r in context["primary_results"]]
+    related_results = [await to_summary(r) for r in context["related_results"]]
+    metadata = MemoryMetadata.model_validate(context["metadata"])
+    
     # Transform to GraphContext
     return GraphContext(
         primary_results=primary_results, related_results=related_results, metadata=metadata
     )
+

src/basic_memory/repository/search_repository.py

@@ -137,6 +137,7 @@ class SearchRepository:
         sql = f"""
             SELECT 
                 id, 
+                title, 
                 permalink,
                 file_path,
                 type,
@@ -162,6 +163,7 @@ class SearchRepository:
         results = [
             SearchIndexRow(
                 id=row.id,
+                title=row.title,
                 permalink=row.permalink,
                 file_path=row.file_path,
                 type=row.type,

src/basic_memory/schemas/memory.py

@@ -1,11 +1,11 @@
 """Schemas for memory context."""
 
+from datetime import datetime
 from typing import Dict, List, Any
 
 from pydantic import AnyUrl, Field, BaseModel
 
 from basic_memory.config import config
-from basic_memory.schemas.search import SearchResult, RelatedResult
 
 """Memory URL schema for knowledge addressing.
 
@@ -47,25 +47,56 @@ class MemoryUrl(AnyUrl):
         return f"memory://{self.host}{self.path}"
 
 
+class EntitySummary(BaseModel):
+    """Simplified entity representation."""
+
+    permalink: str
+    title: str
+    file_path: str
+    created_at: datetime
+
+
+class RelationSummary(BaseModel):
+    """Simplified relation representation."""
+
+    permalink: str
+    type: str
+    from_id: str
+    to_id: str
+    created_at: datetime
+
+
+class ObservationSummary(BaseModel):
+    """Simplified observation representation."""
+
+    permalink: str
+    category: str
+    content: str
+
+
+class MemoryMetadata(BaseModel):
+    """Simplified response metadata."""
+
+    url: str
+    depth: int
+    timeframe: str
+    generated_at: datetime
+    total_results: int
+    total_relations: int
+
+
 class GraphContext(BaseModel):
     """Complete context response."""
 
     # Direct matches
-    primary_results: List[SearchResult] = Field(description="Entities directly matching URI")
+    primary_results: List[EntitySummary | RelationSummary | ObservationSummary] = Field(
+        description="results directly matching URI"
+    )
 
     # Related entities
-    related_results: List[RelatedResult] = Field(description="Entities found via relations")
+    related_results: List[EntitySummary | RelationSummary | ObservationSummary] = Field(
+        description="related results"
+    )
 
     # Context metadata
-    metadata: Dict[str, Any] = Field(
-        default_factory=dict,
-        example={
-            "uri": "memory://specs/search/*",
-            "depth": 2,
-            "timeframe": "7d",
-            "generated_at": "2024-01-14T12:00:00Z",
-            "matched_results": 3,
-            "total_results": 8,
-            "total_relations": 12,
-        },
-    )
+    metadata: MemoryMetadata

src/basic_memory/services/context_service.py

@@ -19,6 +19,7 @@ class ContextResultRow:
     id: int
     title: str
     permalink: str
+    file_path: str
     depth: int
     root_id: int
     created_at: datetime
@@ -88,7 +89,7 @@ class ContextService:
             "primary_results": primary,
             "related_results": related,
             "metadata": {
-                "uri": memory_url.relative_path(),
+                "url": memory_url.relative_path(),
                 "depth": depth,
                 "timeframe": since.isoformat() if since else None,
                 "generated_at": datetime.now(timezone.utc).isoformat(),
@@ -104,7 +105,7 @@ class ContextService:
         max_depth: int = 1,
         since: Optional[datetime] = None,
         max_results: int = 10,
-    ):
+    ) -> List[ContextResultRow]:
         """Find items connected through relations.
 
         Uses recursive CTE to find:
@@ -138,6 +139,7 @@ WITH RECURSIVE context_graph AS (
         type,
         title, 
         permalink,
+        file_path,
         from_id,
         to_id,
         relation_type,
@@ -160,6 +162,7 @@ WITH RECURSIVE context_graph AS (
         r.type,
         r.title,
         r.permalink,
+        r.file_path,
         r.from_id,
         r.to_id,
         r.relation_type,
@@ -187,6 +190,7 @@ WITH RECURSIVE context_graph AS (
         e.type,
         e.title,
         e.permalink,
+        e.file_path,
         e.from_id,
         e.to_id,
         e.relation_type,
@@ -214,6 +218,7 @@ SELECT DISTINCT
     id,
     title,
     permalink,
+    file_path,
     from_id,
     to_id,
     relation_type,
@@ -240,6 +245,7 @@ LIMIT :max_results
                 id=row.id,
                 title=row.title,
                 permalink=row.permalink,
+                file_path=row.file_path,
                 from_id=row.from_id,
                 to_id=row.to_id,
                 relation_type=row.relation_type,