basic-memory 0.7.0__py3-none-any.whl → 0.17.4__py3-none-any.whl

basic_memory/services/context_service.py

@@ -1,16 +1,20 @@
 """Service for building rich context from the knowledge graph."""
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from typing import List, Optional, Tuple
 
+
 from loguru import logger
 from sqlalchemy import text
 
 from basic_memory.repository.entity_repository import EntityRepository
-from basic_memory.repository.search_repository import SearchRepository
+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
 
 
 @dataclass
@@ -31,6 +35,38 @@ class ContextResultRow:
     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
+
+
+@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.
 
@@ -44,9 +80,11 @@ class ContextService:
         self,
         search_repository: SearchRepository,
         entity_repository: EntityRepository,
+        observation_repository: ObservationRepository,
     ):
         self.search_repository = search_repository
         self.entity_repository = entity_repository
+        self.observation_repository = observation_repository
 
     async def build_context(
         self,
@@ -57,31 +95,42 @@ class ContextService:
         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}'"
         )
 
+        normalized_path: Optional[str] = None
         if memory_url:
             path = memory_url_path(memory_url)
-            # Pattern matching - use search
-            if "*" in path:
-                logger.debug(f"Pattern search for '{path}'")
+            # 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=path, limit=limit, offset=offset
+                    permalink_match=normalized_path, limit=limit, offset=offset
                 )
-
-            # Direct lookup for exact path
             else:
-                logger.debug(f"Direct lookup for '{path}'")
+                # 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=path, limit=limit, offset=offset
+                    permalink=normalized_path, limit=limit, offset=offset
                 )
         else:
             logger.debug(f"Build context for '{types}'")
             primary = await self.search_repository.search(
-                types=types, after_date=since, limit=limit, offset=offset
+                search_item_types=types, after_date=since, limit=limit, offset=offset
             )
 
         # Get type_id pairs for traversal
@@ -94,24 +143,78 @@ class ContextService:
             type_id_pairs, max_depth=depth, since=since, max_results=max_related
         )
         logger.debug(f"Found {len(related)} related results")
-        for r in related:
-            logger.debug(f"Found related {r.type}: {r.permalink}")
-
-        # Build response
-        return {
-            "primary_results": primary,
-            "related_results": related,
-            "metadata": {
-                "uri": memory_url_path(memory_url) if memory_url else None,
-                "types": types if types else None,
-                "depth": depth,
-                "timeframe": since.isoformat() if since else None,
-                "generated_at": datetime.now(timezone.utc).isoformat(),
-                "matched_results": len(primary),
-                "total_results": len(primary) + len(related),
-                "total_relations": sum(1 for r in related if r.type == SearchItemType.RELATION),
-            },
-        }
+
+        # 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),
+        )
+
+        # 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,
@@ -124,7 +227,6 @@ class ContextService:
 
         Uses recursive CTE to find:
         - Connected entities
-        - Their observations
         - Relations that connect them
 
         Note on depth:
@@ -138,107 +240,344 @@ class ContextService:
         if not type_id_pairs:
             return []
 
-        logger.debug(f"Finding connected items for {type_id_pairs} with depth {max_depth}")
+        # 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])
 
-        # Build the VALUES clause directly since SQLite doesn't handle parameterized IN well
-        values = ", ".join([f"('{t}', {i})" for t, i in type_id_pairs])
+        # 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,
+        }
 
-        # Parameters for bindings
-        params = {"max_depth": max_depth, "max_results": max_results}
+        # Build date and timeframe filters conditionally based on since parameter
         if since:
-            params["since_date"] = since.isoformat()  # pyright: ignore
-
-        # Build date filter
-        date_filter = "AND base.created_at >= :since_date" if since else ""
-        r1_date_filter = "AND r.created_at >= :since_date" if since else ""
-        related_date_filter = "AND e.created_at >= :since_date" if since else ""
-
-        query = text(f"""
-        WITH RECURSIVE context_graph AS (
-            -- Base case: seed items 
-            SELECT 
-                id,
-                type,
-                title, 
-                permalink,
-                file_path,
-                from_id,
-                to_id,
-                relation_type,
-                content,
-                category,
-                entity_id,
+            # 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,
-                id as root_id,
-                created_at,
-                created_at as relation_date,
+                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 search_index base
-            WHERE (base.type, base.id) IN ({values})
+            FROM entity e
+            WHERE e.id IN ({entity_id_values})
             {date_filter}
+            {project_filter}
 
-            UNION ALL  -- Allow same paths at different depths
+            UNION ALL
 
-            -- Get relations from current entities 
-            SELECT DISTINCT
+            -- Get relations from current entities
+            SELECT
                 r.id,
-                r.type,
-                r.title,
-                r.permalink,
-                r.file_path,
+                '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,
-                r.content,
-                r.category,
-                r.entity_id,
-                cg.depth + 1,
-                cg.root_id,
-                r.created_at,
-                r.created_at as relation_date,
-                CASE WHEN r.from_id = cg.id THEN 0 ELSE 1 END as is_incoming
-            FROM context_graph cg
-            JOIN search_index r ON (
-                cg.type = 'entity' AND
-                r.type = 'relation' AND 
-                (r.from_id = cg.id OR r.to_id = cg.id)
-                {r1_date_filter}
+                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 cg.depth < :max_depth
+            LEFT JOIN entity e_to ON (r.to_id = e_to.id)
+            WHERE eg.depth < :max_depth
+            AND (r.to_id IS NULL OR e_to.project_id = :project_id)
 
             UNION ALL
 
             -- Get entities connected by relations
-            SELECT DISTINCT
+            SELECT
                 e.id,
-                e.type,
+                'entity' as type,
                 e.title,
-                e.permalink,
+                CASE
+                    WHEN e.permalink IS NULL THEN ''
+                    ELSE e.permalink
+                END as permalink,
                 e.file_path,
-                e.from_id,
-                e.to_id,
-                e.relation_type,
-                e.content,
-                e.category,
-                e.entity_id,
-                cg.depth + 1,  -- Increment depth for entities
-                cg.root_id,
+                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,
-                cg.relation_date,
-                cg.is_incoming
-            FROM context_graph cg
-            JOIN search_index e ON (
-                cg.type = 'relation' AND 
-                e.type = 'entity' AND
-                e.id = CASE 
-                    WHEN cg.is_incoming = 0 THEN cg.to_id  -- Fixed entity lookup
-                    ELSE cg.from_id                             
+                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
-                {related_date_filter}
+                {date_filter}
+                {project_filter}
             )
-            WHERE cg.depth < :max_depth
+            WHERE eg.depth < :max_depth
+            {timeframe_condition}
         )
-        SELECT DISTINCT 
+        SELECT DISTINCT
             type,
             id,
             title,
@@ -253,36 +592,10 @@ class ContextService:
             MIN(depth) as depth,
             root_id,
             created_at
-        FROM context_graph
-        WHERE (type, id) NOT IN ({values})
-        GROUP BY
-            type, id, title, permalink, from_id, to_id,
-            relation_type, category, entity_id,
-            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
        """)
-
-        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