diff --git a/src/basic_memory/api/routers/memory_router.py b/src/basic_memory/api/routers/memory_router.py index 56b9821a..ce8c1299 100644 --- a/src/basic_memory/api/routers/memory_router.py +++ b/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 ) + diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py index 4fa6e8e5..bc7b0551 100644 --- a/src/basic_memory/repository/search_repository.py +++ b/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, diff --git a/src/basic_memory/schemas/memory.py b/src/basic_memory/schemas/memory.py index abb4b9a6..1cbf2753 100644 --- a/src/basic_memory/schemas/memory.py +++ b/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 diff --git a/src/basic_memory/services/context_service.py b/src/basic_memory/services/context_service.py index db4c0e7e..0284bc48 100644 --- a/src/basic_memory/services/context_service.py +++ b/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, diff --git a/tests/api/test_memory_router.py b/tests/api/test_memory_router.py index cefff054..d1170cb3 100644 --- a/tests/api/test_memory_router.py +++ b/tests/api/test_memory_router.py @@ -1,4 +1,5 @@ """Tests for memory router endpoints.""" +from datetime import datetime import pytest @@ -17,11 +18,11 @@ async def test_get_memory_context(client, test_graph): assert len(context.related_results) > 0 # Verify metadata - assert context.metadata["uri"] == "test/root" - assert context.metadata["depth"] == 1 # default depth + assert context.metadata.url == "test/root" + assert context.metadata.depth == 1 # default depth # assert context.metadata["timeframe"] == "7d" # default timeframe - assert isinstance(context.metadata["generated_at"], str) - assert context.metadata["matched_results"] == 1 + assert isinstance(context.metadata.generated_at, datetime) + assert context.metadata.total_results == 3 @pytest.mark.asyncio