fix tests

This commit is contained in:
phernandez
2025-01-15 13:40:15 -06:00
parent 65cc945b28
commit 62725fa2c3
2 changed files with 39 additions and 20 deletions
@@ -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
+20 -15
View File
@@ -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()