speed up context query

src/basic_memory/api/routers/memory_router.py

@@ -1,12 +1,15 @@
 """Routes for memory:// URI operations."""
+
 from dataclasses import asdict
-from typing import List, Optional
 from datetime import datetime, timedelta
+from typing import Optional
+
 from fastapi import APIRouter
+from loguru import logger
 
 from basic_memory.config import config
-from basic_memory.schemas.memory import MemoryUrl, GraphContext
 from basic_memory.deps import ContextServiceDep
+from basic_memory.schemas.memory import MemoryUrl, GraphContext
 from basic_memory.schemas.search import SearchResult, RelatedResult
 
 router = APIRouter(prefix="/memory", tags=["memory"])
@@ -36,22 +39,26 @@ async def get_memory_context(
     uri: str,
     depth: int = 1,
     timeframe: str = "7d",
-    max_results: int = 10
+    max_results: int = 10,
 ) -> GraphContext:
     """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}`")
     memory_url = MemoryUrl(f"memory://{config.project}/{uri}")
 
     # Parse timeframe
     since = parse_timeframe(timeframe)
 
     # Build context
-    context = await context_service.build_context(memory_url, depth=depth, since=since, max_results=max_results)
+    context = await context_service.build_context(
+        memory_url, depth=depth, since=since, max_results=max_results
+    )
 
-    primary_entities = [SearchResult(**asdict(r)) for r in context["primary_entities"]]
-    related_entities = [RelatedResult(**asdict(r)) for r in context["related_entities"]]
+    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"]
     # Transform to GraphContext
-    return GraphContext(primary_entities=primary_entities, related_entities=related_entities, metadata=metadata)
-
+    return GraphContext(
+        primary_results=primary_results, related_results=related_results, metadata=metadata
+    )