From 62725fa2c3b16c348b8d2db67bcd7a1da65e18a4 Mon Sep 17 00:00:00 2001
From: phernandez <phernandez@basicmachines.co>
Date: Wed, 15 Jan 2025 13:40:15 -0600
Subject: [PATCH] fix tests

---
 .../repository/search_repository.py           | 24 ++++++++++---
 src/basic_memory/services/context_service.py  | 35 +++++++++++--------
 2 files changed, 39 insertions(+), 20 deletions(-)

diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py
index e3db2a3a..6a22f87f 100644
--- a/src/basic_memory/repository/search_repository.py
+++ b/src/basic_memory/repository/search_repository.py
@@ -1,7 +1,7 @@
 """Repository for search operations."""
 
 import json
-from typing import List, Optional, Any
+from typing import List, Optional, Any, Dict
 
 from loguru import logger
 from sqlalchemy import text, Executable, Result
@@ -169,10 +169,24 @@ class SearchRepository:
             )
             await session.commit()
 
-    async def execute_query(self, query: Executable, use_query_options:bool = True) -> Result[Any]:
-        """Execute a query asynchronously."""
+    async def execute_query(
+        self, 
+        query: Executable, 
+        params: Optional[Dict[str, Any]] = None, 
+        use_query_options:bool = True
+    ) -> Result[Any]:
+        """Execute a query asynchronously.
+        
+        Args:
+            query: The query to execute
+            params: Optional parameters to bind to the query
+            use_query_options: Whether to apply query options
+        """
         logger.debug(f"Executing query: {query}")
         async with db.scoped_session(self.session_maker) as session:
-            result = await session.execute(query)
+            if params:
+                result = await session.execute(query, params)
+            else:
+                result = await session.execute(query)
             logger.debug("Query executed successfully")
-            return result
+            return result
\ No newline at end of file
diff --git a/src/basic_memory/services/context_service.py b/src/basic_memory/services/context_service.py
index 06baf8b0..5b714329 100644
--- a/src/basic_memory/services/context_service.py
+++ b/src/basic_memory/services/context_service.py
@@ -101,14 +101,9 @@ class ContextService:
             
         logger.debug(f"Finding connected items for {type_id_pairs} with depth {max_depth}")
 
-        # Build the VALUES clause for our seed items
+        # 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])
 
-        # Build date condition for timeframe filtering
-        date_filter = f""
-        if since:
-            date_filter = f"AND created_at >= '{since.isoformat()}'"
-
         # Debug: Check what's in the search index
         debug_query = text("""
             SELECT type, id, from_id, to_id, relation_type 
@@ -123,6 +118,16 @@ class ContextService:
             else:
                 logger.debug(f"{r.type} {r.id}")
 
+        # Parameters for bindings
+        params = {"max_depth": max_depth}
+        if since:
+            params["since_date"] = since.isoformat()
+
+        # Build date filter
+        date_filter = f"AND base.created_at >= :since_date" if since else ""
+        r1_date_filter = f"AND r1.created_at >= :since_date" if since else ""
+        related_date_filter = f"AND related.created_at >= :since_date" if since else ""
+
         query = text(f"""
             WITH RECURSIVE context_graph AS (
                 -- Base case: seed items
@@ -141,12 +146,12 @@ class ContextService:
                     id as root_id,
                     created_at
                 FROM search_index base
-                WHERE (type, id) IN (VALUES {values})
+                WHERE (base.type, base.id) IN ({values})
                 {date_filter}
 
                 UNION 
 
-                -- Find all connected items (relations + entities) at each depth
+                -- Find relations and their connected items at each depth
                 SELECT
                     related.id,
                     related.type,
@@ -162,15 +167,15 @@ class ContextService:
                     cg.root_id,
                     related.created_at
                 FROM context_graph cg
-                INNER JOIN search_index r1 ON (
+                JOIN search_index r1 ON (
                     -- First find the relations
                     cg.type = 'entity' AND 
                     r1.type = 'relation' AND 
                     (r1.from_id = cg.id OR r1.to_id = cg.id)
-                    {date_filter.replace('created_at', 'r1.created_at')}
+                    {r1_date_filter}
                 )
                 -- Then join to ALL related items at the same depth 
-                LEFT JOIN search_index related ON (
+                JOIN search_index related ON (
                     -- The found relation 
                     related.id = r1.id
                     OR
@@ -181,9 +186,9 @@ class ContextService:
                     -- Any observations that are relevant
                     (related.type = 'observation' AND 
                      (related.entity_id = r1.from_id OR related.entity_id = r1.to_id))
+                    {related_date_filter}
                 )
-                WHERE cg.depth < {max_depth}
-                {date_filter.replace('created_at', 'related.created_at')}
+                WHERE cg.depth < :max_depth
             )
             -- Select shortest path to each item
             SELECT DISTINCT 
@@ -205,8 +210,8 @@ class ContextService:
                 type, id, title, permalink, from_id, to_id,
                 relation_type, category, entity_id, content,
                 root_id, created_at
-            ORDER BY depth, type
+            ORDER BY depth, type, id
         """)
 
-        results = await self.search_repository.execute_query(query)
+        results = await self.search_repository.execute_query(query, params=params)
         return results.all()
\ No newline at end of file