"""Service for building rich context from the knowledge graph.""" from __future__ import annotations from dataclasses import dataclass, field from datetime import datetime, timezone from typing import List, Optional, Tuple, TYPE_CHECKING from loguru import logger from sqlalchemy import text from basic_memory.repository.entity_repository import EntityRepository from basic_memory.repository.observation_repository import ObservationRepository from basic_memory.repository.postgres_search_repository import PostgresSearchRepository from basic_memory.repository.search_repository import SearchRepository, SearchIndexRow from basic_memory.schemas.memory import MemoryUrl, memory_url_path from basic_memory.schemas.search import SearchItemType from basic_memory.utils import generate_permalink if TYPE_CHECKING: from basic_memory.services.link_resolver import LinkResolver @dataclass class ContextResultRow: type: str id: int title: str permalink: str file_path: str depth: int root_id: int created_at: datetime from_id: Optional[int] = None to_id: Optional[int] = None relation_type: Optional[str] = None content: Optional[str] = None category: Optional[str] = None entity_id: Optional[int] = None @dataclass class ContextResultItem: """A hierarchical result containing a primary item with its observations and related items.""" primary_result: ContextResultRow | SearchIndexRow observations: List[ContextResultRow] = field(default_factory=list) related_results: List[ContextResultRow] = field(default_factory=list) @dataclass class ContextMetadata: """Metadata about a context result.""" uri: Optional[str] = None types: Optional[List[SearchItemType]] = None depth: int = 1 timeframe: Optional[str] = None generated_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) primary_count: int = 0 related_count: int = 0 total_observations: int = 0 total_relations: int = 0 has_more: bool = False @dataclass class ContextResult: """Complete context result with metadata.""" results: List[ContextResultItem] = field(default_factory=list) metadata: ContextMetadata = field(default_factory=ContextMetadata) class ContextService: """Service for building rich context from memory:// URIs. Handles three types of context building: 1. Direct permalink lookup - exact match on path 2. Pattern matching - using * wildcards 3. Special modes via params (e.g., 'related') """ def __init__( self, search_repository: SearchRepository, entity_repository: EntityRepository, observation_repository: ObservationRepository, link_resolver: Optional[LinkResolver] = None, ): self.search_repository = search_repository self.entity_repository = entity_repository self.observation_repository = observation_repository self.link_resolver = link_resolver async def build_context( self, memory_url: Optional[MemoryUrl] = None, types: Optional[List[SearchItemType]] = None, depth: int = 1, since: Optional[datetime] = None, limit=10, offset=0, max_related: int = 10, include_observations: bool = True, ) -> ContextResult: """Build rich context from a memory:// URI.""" logger.debug( f"Building context for URI: '{memory_url}' depth: '{depth}' since: '{since}' limit: '{limit}' offset: '{offset}' max_related: '{max_related}'" ) # Fetch one extra item to detect whether more pages exist (N+1 trick) fetch_limit = limit + 1 normalized_path: Optional[str] = None if memory_url: path = memory_url_path(memory_url) # Check for wildcards before normalization has_wildcard = "*" in path if has_wildcard: # For wildcard patterns, normalize each segment separately to preserve the * parts = path.split("*") normalized_parts = [ generate_permalink(part, split_extension=False) if part else "" for part in parts ] normalized_path = "*".join(normalized_parts) logger.debug(f"Pattern search for '{normalized_path}'") primary = await self.search_repository.search( permalink_match=normalized_path, limit=fetch_limit, offset=offset ) else: # For exact paths, normalize the whole thing normalized_path = generate_permalink(path, split_extension=False) logger.debug(f"Direct lookup for '{normalized_path}'") primary = await self.search_repository.search( permalink=normalized_path, limit=fetch_limit, offset=offset ) # Trigger: exact permalink lookup returned no results # Why: the identifier may be valid but not an exact permalink match # (e.g., missing project prefix, title instead of permalink) # Outcome: use LinkResolver's multi-strategy resolution to find the entity, # then retry search with its actual permalink if not primary and self.link_resolver: entity = await self.link_resolver.resolve_link( path, use_search=True, strict=False ) if entity: logger.debug( f"LinkResolver resolved '{path}' to permalink '{entity.permalink}'" ) normalized_path = entity.permalink primary = await self.search_repository.search( permalink=entity.permalink, limit=fetch_limit, offset=offset ) else: logger.debug(f"Build context for '{types}'") primary = await self.search_repository.search( search_item_types=types, after_date=since, limit=fetch_limit, offset=offset ) # Trim to requested limit and set has_more flag has_more = len(primary) > limit if has_more: primary = primary[:limit] # Get type_id pairs for traversal type_id_pairs = [(r.type, r.id) for r in primary] if primary else [] logger.debug(f"found primary type_id_pairs: {len(type_id_pairs)}") # Find related content related = await self.find_related( type_id_pairs, max_depth=depth, since=since, max_results=max_related ) logger.debug(f"Found {len(related)} related results") # Collect entity IDs from primary and related results entity_ids = [] for result in primary: if result.type == SearchItemType.ENTITY.value: entity_ids.append(result.id) for result in related: if result.type == SearchItemType.ENTITY.value: entity_ids.append(result.id) # Fetch observations for all entities if requested observations_by_entity = {} if include_observations and entity_ids: # Use our observation repository to get observations for all entities at once observations_by_entity = await self.observation_repository.find_by_entities(entity_ids) logger.debug(f"Found observations for {len(observations_by_entity)} entities") # Create metadata dataclass metadata = ContextMetadata( uri=normalized_path if memory_url else None, types=types, depth=depth, timeframe=since.isoformat() if since else None, primary_count=len(primary), related_count=len(related), total_observations=sum(len(obs) for obs in observations_by_entity.values()), total_relations=sum(1 for r in related if r.type == SearchItemType.RELATION), has_more=has_more, ) # Build context results list directly with ContextResultItem objects context_results = [] # For each primary result for primary_item in primary: # Find all related items with this primary item as root related_to_primary = [r for r in related if r.root_id == primary_item.id] # Get observations for this item if it's an entity item_observations = [] if primary_item.type == SearchItemType.ENTITY.value and include_observations: # Convert Observation models to ContextResultRows for obs in observations_by_entity.get(primary_item.id, []): item_observations.append( ContextResultRow( type="observation", id=obs.id, title=f"{obs.category}: {obs.content[:50]}...", permalink=generate_permalink( f"{primary_item.permalink}/observations/{obs.category}/{obs.content}" ), file_path=primary_item.file_path, content=obs.content, category=obs.category, entity_id=primary_item.id, depth=0, root_id=primary_item.id, created_at=primary_item.created_at, # created_at time from entity ) ) # Create ContextResultItem directly context_item = ContextResultItem( primary_result=primary_item, observations=item_observations, related_results=related_to_primary, ) context_results.append(context_item) # Return the structured ContextResult return ContextResult(results=context_results, metadata=metadata) async def find_related( self, type_id_pairs: List[Tuple[str, int]], max_depth: int = 1, since: Optional[datetime] = None, max_results: int = 10, ) -> List[ContextResultRow]: """Find items connected through relations. Uses recursive CTE to find: - Connected entities - Relations that connect them Note on depth: Each traversal step requires two depth levels - one to find the relation, and another to follow that relation to an entity. So a max_depth of 4 allows traversal through two entities (relation->entity->relation->entity), while reaching an entity three steps away requires max_depth=6 (relation->entity->relation->entity->relation->entity). """ max_depth = max_depth * 2 if not type_id_pairs: return [] # Extract entity IDs from type_id_pairs for the optimized query entity_ids = [i for t, i in type_id_pairs if t == "entity"] if not entity_ids: logger.debug("No entity IDs found in type_id_pairs") return [] logger.debug( f"Finding connected items for {len(entity_ids)} entities with depth {max_depth}" ) # Build the VALUES clause for entity IDs entity_id_values = ", ".join([str(i) for i in entity_ids]) # Parameters for bindings - include project_id for security filtering params = { "max_depth": max_depth, "max_results": max_results, "project_id": self.search_repository.project_id, } # Build date and timeframe filters conditionally based on since parameter if since: # SQLite accepts ISO strings, but Postgres/asyncpg requires datetime objects if isinstance(self.search_repository, PostgresSearchRepository): # pragma: no cover # asyncpg expects timezone-NAIVE datetime in UTC for DateTime(timezone=True) columns # even though the column stores timezone-aware values since_utc = ( since.astimezone(timezone.utc) if since.tzinfo else since ) # pragma: no cover params["since_date"] = since_utc.replace(tzinfo=None) # pyright: ignore # pragma: no cover else: params["since_date"] = since.isoformat() # pyright: ignore date_filter = "AND e.created_at >= :since_date" relation_date_filter = "AND e_from.created_at >= :since_date" timeframe_condition = "AND eg.relation_date >= :since_date" else: date_filter = "" relation_date_filter = "" timeframe_condition = "" # Add project filtering for security - ensure all entities and relations belong to the same project project_filter = "AND e.project_id = :project_id" relation_project_filter = "AND e_from.project_id = :project_id" # Use a CTE that operates directly on entity and relation tables # This avoids the overhead of the search_index virtual table # Note: Postgres and SQLite have different CTE limitations: # - Postgres: doesn't allow multiple UNION ALL branches referencing the CTE # - SQLite: doesn't support LATERAL joins # So we need different queries for each database backend # Detect database backend is_postgres = isinstance(self.search_repository, PostgresSearchRepository) if is_postgres: # pragma: no cover query = self._build_postgres_query( entity_id_values, date_filter, project_filter, relation_date_filter, relation_project_filter, timeframe_condition, ) else: # SQLite needs VALUES clause for exclusion (not needed for Postgres) values = ", ".join([f"('{t}', {i})" for t, i in type_id_pairs]) query = self._build_sqlite_query( entity_id_values, date_filter, project_filter, relation_date_filter, relation_project_filter, timeframe_condition, values, ) result = await self.search_repository.execute_query(query, params=params) rows = result.all() context_rows = [ ContextResultRow( type=row.type, 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, content=row.content, category=row.category, entity_id=row.entity_id, depth=row.depth, root_id=row.root_id, created_at=row.created_at, ) for row in rows ] return context_rows def _build_postgres_query( # pragma: no cover self, entity_id_values: str, date_filter: str, project_filter: str, relation_date_filter: str, relation_project_filter: str, timeframe_condition: str, ): """Build Postgres-specific CTE query using LATERAL joins.""" return text(f""" WITH RECURSIVE entity_graph AS ( -- Base case: seed entities SELECT e.id, 'entity' as type, e.title, e.permalink, e.file_path, CAST(NULL AS INTEGER) as from_id, CAST(NULL AS INTEGER) as to_id, CAST(NULL AS TEXT) as relation_type, CAST(NULL AS TEXT) as content, CAST(NULL AS TEXT) as category, CAST(NULL AS INTEGER) as entity_id, 0 as depth, e.id as root_id, e.created_at, e.created_at as relation_date FROM entity e WHERE e.id IN ({entity_id_values}) {date_filter} {project_filter} UNION ALL -- Fetch BOTH relations AND connected entities in a single recursive step -- Postgres only allows ONE reference to the recursive CTE in the recursive term -- We use CROSS JOIN LATERAL to generate two rows (relation + entity) from each traversal SELECT CASE WHEN step_type = 1 THEN r.id ELSE e.id END as id, CASE WHEN step_type = 1 THEN 'relation' ELSE 'entity' END as type, CASE WHEN step_type = 1 THEN r.relation_type || ': ' || r.to_name ELSE e.title END as title, CASE WHEN step_type = 1 THEN '' ELSE COALESCE(e.permalink, '') END as permalink, CASE WHEN step_type = 1 THEN e_from.file_path ELSE e.file_path END as file_path, CASE WHEN step_type = 1 THEN r.from_id ELSE NULL END as from_id, CASE WHEN step_type = 1 THEN r.to_id ELSE NULL END as to_id, CASE WHEN step_type = 1 THEN r.relation_type ELSE NULL END as relation_type, CAST(NULL AS TEXT) as content, CAST(NULL AS TEXT) as category, CAST(NULL AS INTEGER) as entity_id, eg.depth + step_type as depth, eg.root_id, CASE WHEN step_type = 1 THEN e_from.created_at ELSE e.created_at END as created_at, CASE WHEN step_type = 1 THEN e_from.created_at ELSE eg.relation_date END as relation_date FROM entity_graph eg CROSS JOIN LATERAL (VALUES (1), (2)) AS steps(step_type) JOIN relation r ON ( eg.type = 'entity' AND (r.from_id = eg.id OR r.to_id = eg.id) ) JOIN entity e_from ON ( r.from_id = e_from.id {relation_project_filter} ) LEFT JOIN entity e ON ( step_type = 2 AND e.id = CASE WHEN r.from_id = eg.id THEN r.to_id ELSE r.from_id END {date_filter} {project_filter} ) WHERE eg.depth < :max_depth AND (step_type = 1 OR (step_type = 2 AND e.id IS NOT NULL AND e.id != eg.id)) {timeframe_condition} ) -- Materialize and filter SELECT DISTINCT type, id, title, permalink, file_path, from_id, to_id, relation_type, content, category, entity_id, MIN(depth) as depth, root_id, created_at FROM entity_graph WHERE depth > 0 GROUP BY type, id, title, permalink, file_path, from_id, to_id, relation_type, content, category, entity_id, root_id, created_at ORDER BY depth, type, id LIMIT :max_results """) def _build_sqlite_query( self, entity_id_values: str, date_filter: str, project_filter: str, relation_date_filter: str, relation_project_filter: str, timeframe_condition: str, values: str, ): """Build SQLite-specific CTE query using multiple UNION ALL branches.""" return text(f""" WITH RECURSIVE entity_graph AS ( -- Base case: seed entities SELECT e.id, 'entity' as type, e.title, e.permalink, e.file_path, NULL as from_id, NULL as to_id, NULL as relation_type, NULL as content, NULL as category, NULL as entity_id, 0 as depth, e.id as root_id, e.created_at, e.created_at as relation_date, 0 as is_incoming FROM entity e WHERE e.id IN ({entity_id_values}) {date_filter} {project_filter} UNION ALL -- Get relations from current entities SELECT r.id, 'relation' as type, r.relation_type || ': ' || r.to_name as title, '' as permalink, e_from.file_path, r.from_id, r.to_id, r.relation_type, NULL as content, NULL as category, NULL as entity_id, eg.depth + 1, eg.root_id, e_from.created_at, e_from.created_at as relation_date, CASE WHEN r.from_id = eg.id THEN 0 ELSE 1 END as is_incoming FROM entity_graph eg JOIN relation r ON ( eg.type = 'entity' AND (r.from_id = eg.id OR r.to_id = eg.id) ) JOIN entity e_from ON ( r.from_id = e_from.id {relation_date_filter} {relation_project_filter} ) WHERE eg.depth < :max_depth UNION ALL -- Get entities connected by relations SELECT e.id, 'entity' as type, e.title, CASE WHEN e.permalink IS NULL THEN '' ELSE e.permalink END as permalink, e.file_path, NULL as from_id, NULL as to_id, NULL as relation_type, NULL as content, NULL as category, NULL as entity_id, eg.depth + 1, eg.root_id, e.created_at, eg.relation_date, eg.is_incoming FROM entity_graph eg JOIN entity e ON ( eg.type = 'relation' AND e.id = CASE WHEN eg.is_incoming = 0 THEN eg.to_id ELSE eg.from_id END {date_filter} {project_filter} ) WHERE eg.depth < :max_depth {timeframe_condition} ) SELECT DISTINCT type, id, title, permalink, file_path, from_id, to_id, relation_type, content, category, entity_id, MIN(depth) as depth, root_id, created_at FROM entity_graph WHERE depth > 0 GROUP BY type, id, title, permalink, file_path, from_id, to_id, relation_type, content, category, entity_id, root_id, created_at ORDER BY depth, type, id LIMIT :max_results """)