From 0533c6dd26d2fb4928eae0c837efae608e6351d3 Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Thu, 2 Apr 2026 11:13:09 -0400 Subject: [PATCH] fix dialectic held connection (#477) * fix: dialectic held connection * fix: (agent) pre-compute embeddings for agent tools * fix: (tests) refactor tests to use smaller test db connections * fix: Embedding client to branch depending on vector store * fix: reflect dedup-skipped observations in created counts and isolate DB sessions in extract_preferences * fix: (tests) update tests to match changes * fix: expunge docs + don't pass in db to query_documents --------- Co-authored-by: Rajat Ahuja --- src/crud/__init__.py | 4 + src/crud/document.py | 266 +++++++---- src/crud/message.py | 8 +- src/crud/representation.py | 4 +- src/dialectic/chat.py | 62 ++- src/dialectic/core.py | 38 +- src/dreamer/orchestrator.py | 43 +- src/dreamer/specialists.py | 42 +- src/dreamer/surprisal.py | 55 ++- src/routers/peers.py | 1 - src/utils/agent_tools.py | 575 +++++++++++++----------- tests/conftest.py | 5 + tests/crud/test_document.py | 2 +- tests/integration/test_token_metrics.py | 3 - tests/utils/test_agent_tools.py | 64 ++- 15 files changed, 658 insertions(+), 514 deletions(-) diff --git a/src/crud/__init__.py b/src/crud/__init__.py index fdff810d..ff8f87ac 100644 --- a/src/crud/__init__.py +++ b/src/crud/__init__.py @@ -9,6 +9,7 @@ from .document import ( create_observations, delete_document, delete_document_by_id, + fetch_documents_by_ids, get_all_documents, get_child_observations, get_documents_by_ids, @@ -16,6 +17,7 @@ from .document import ( query_documents, query_documents_most_derived, query_documents_recent, + query_external_vector_document_ids, ) from .message import ( create_messages, @@ -82,6 +84,7 @@ __all__ = [ # Document "create_documents", "create_observations", + "fetch_documents_by_ids", "get_all_documents", "get_child_observations", "get_documents_by_ids", @@ -89,6 +92,7 @@ __all__ = [ "query_documents", "query_documents_most_derived", "query_documents_recent", + "query_external_vector_document_ids", "delete_document", "delete_document_by_id", # Message diff --git a/src/crud/document.py b/src/crud/document.py index 4e5d8751..688ed2f0 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -15,6 +15,7 @@ from src.config import settings from src.crud.collection import get_or_create_collection from src.crud.peer import get_peer from src.crud.session import get_session +from src.dependencies import tracked_db from src.embedding_client import embedding_client from src.exceptions import ResourceNotFoundException, ValidationException from src.utils.filter import apply_filter @@ -190,9 +191,127 @@ async def query_documents_most_derived( return result.scalars().all() -async def query_documents( +def _uses_pgvector() -> bool: + """Check whether queries should go through pgvector (DB-only) path.""" + return ( + settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED + ) + + +async def query_external_vector_document_ids( + workspace_name: str, + observer: str, + observed: str, + embedding: list[float], + top_k: int = 5, + max_distance: float | None = None, + filters: dict[str, Any] | None = None, +) -> list[str] | None: + """Query external vector store for document IDs sorted by similarity. + + No DB session needed — safe to call outside a tracked_db scope. + + Returns: + Ordered list of document IDs on the external-store path, + empty list when the external store has no results, + or None when the pgvector (DB-only) path should be used instead. + """ + if _uses_pgvector(): + return None + + external_vector_store = get_external_vector_store() + if external_vector_store is None: + return [] + + namespace = external_vector_store.get_vector_namespace( + "document", workspace_name, observer, observed + ) + + vector_filters: dict[str, Any] = {} + if filters: + for key in ["level", "session_name"]: + if key in filters: + vector_filters[key] = filters[key] + + vector_results = await external_vector_store.query( + namespace, + embedding, + top_k=top_k, + max_distance=max_distance, + filters=vector_filters if vector_filters else None, + ) + + if not vector_results: + return [] + + return [result.id for result in vector_results] + + +async def fetch_documents_by_ids( db: AsyncSession, workspace_name: str, + observer: str, + observed: str, + document_ids: list[str], + filters: dict[str, Any] | None = None, +) -> list[models.Document]: + """Fetch documents by IDs, preserving input order. DB-only operation.""" + if not document_ids: + return [] + + stmt = ( + select(models.Document) + .where(models.Document.workspace_name == workspace_name) + .where(models.Document.observer == observer) + .where(models.Document.observed == observed) + .where(models.Document.deleted_at.is_(None)) + .where(models.Document.id.in_(document_ids)) + ) + stmt = apply_filter(stmt, models.Document, filters) + + result = await db.execute(stmt) + documents = {doc.id: doc for doc in result.scalars().all()} + + return [documents[doc_id] for doc_id in document_ids if doc_id in documents] + + +async def _query_documents_pgvector( + db: AsyncSession, + workspace_name: str, + observer: str, + observed: str, + embedding: list[float], + filters: dict[str, Any] | None, + max_distance: float | None, + top_k: int, +) -> list[models.Document]: + """pgvector similarity search — pure DB operation.""" + stmt = ( + select(models.Document) + .where(models.Document.workspace_name == workspace_name) + .where(models.Document.observer == observer) + .where(models.Document.observed == observed) + .where(models.Document.embedding.isnot(None)) + .where(models.Document.deleted_at.is_(None)) + ) + + if max_distance is not None: + stmt = stmt.where( + models.Document.embedding.cosine_distance(embedding) <= max_distance + ) + + stmt = apply_filter(stmt, models.Document, filters) + stmt = stmt.order_by(models.Document.embedding.cosine_distance(embedding)).limit( + top_k + ) + + result = await db.execute(stmt) + return list(result.scalars().all()) + + +async def query_documents( + db: AsyncSession | None, + workspace_name: str, query: str, *, observer: str, @@ -205,8 +324,12 @@ async def query_documents( """ Query documents using semantic similarity. + When *db* is provided the caller owns the session lifetime. When *db* is + ``None`` the function opens (and closes) its own short-lived session so that + no DB connection is held during external vector-store calls. + Args: - db: Database session + db: Database session, or None to let the function manage its own workspace_name: Name of the workspace query: Search query text observer: Name of the observing peer @@ -228,89 +351,69 @@ async def query_documents( f"Query exceeds maximum token limit of {settings.MAX_EMBEDDING_TOKENS}." ) from e - # Query Postgres directly when using pgvector OR during migration (not yet migrated) - # This ensures we use pgvector as source of truth until migration is complete - if settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED: - stmt = ( - select(models.Document) - .where(models.Document.workspace_name == workspace_name) - .where(models.Document.observer == observer) - .where(models.Document.observed == observed) - .where(models.Document.embedding.isnot(None)) - .where(models.Document.deleted_at.is_(None)) - ) - - if max_distance is not None: - stmt = stmt.where( - models.Document.embedding.cosine_distance(embedding) <= max_distance + if _uses_pgvector(): + # pgvector path — pure DB, open a short session if none provided + if db is not None: + return await _query_documents_pgvector( + db, + workspace_name, + observer, + observed, + embedding, + filters, + max_distance, + top_k, ) + async with tracked_db("query_documents.pgvector") as managed_db: + docs = await _query_documents_pgvector( + managed_db, + workspace_name, + observer, + observed, + embedding, + filters, + max_distance, + top_k, + ) + for doc in docs: + managed_db.expunge(doc) + return docs - stmt = apply_filter(stmt, models.Document, filters) - stmt = stmt.order_by( - models.Document.embedding.cosine_distance(embedding) - ).limit(top_k) - - result = await db.execute(stmt) - return list(result.scalars().all()) - - # FALLBACK: Use external vector store (Turbopuffer, LanceDB) - external_vector_store = get_external_vector_store() - if external_vector_store is None: - return [] - - namespace = external_vector_store.get_vector_namespace( - "document", workspace_name, observer, observed - ) - - # Build vector store filters - # Convert filter dict to vector store format (handles level, session_name, etc.) - vector_filters: dict[str, Any] = {} - if filters: - # Direct pass-through for simple equality filters - # The filters dict can contain: level, session_name, or other document fields - # We can push level and session_name to vector store since they're in metadata - for key in ["level", "session_name"]: - if key in filters: - vector_filters[key] = filters[key] - - # Query external vector store for similar documents with filters applied - vector_results = await external_vector_store.query( - namespace, - embedding, + # External vector store — network call first, DB only for the ID fetch + document_ids = await query_external_vector_document_ids( + workspace_name=workspace_name, + observer=observer, + observed=observed, + embedding=embedding, top_k=top_k, max_distance=max_distance, - filters=vector_filters if vector_filters else None, + filters=filters, ) - if not vector_results: + if not document_ids: return [] - # Get document IDs from vector results (vector ID = document ID for documents) - document_ids = [result.id for result in vector_results] - - # Fetch documents from database - stmt = ( - select(models.Document) - .where(models.Document.workspace_name == workspace_name) - .where(models.Document.observer == observer) - .where(models.Document.observed == observed) - .where(models.Document.deleted_at.is_(None)) - .where(models.Document.id.in_(document_ids)) - ) - # Re-apply all filters at the database layer to catch any constraints - # that aren't supported by the vector store metadata. - stmt = apply_filter(stmt, models.Document, filters) - - result = await db.execute(stmt) - documents = {doc.id: doc for doc in result.scalars().all()} - - # Return documents in order of similarity (preserving vector store order) - ordered_docs: list[models.Document] = [] - for vr in vector_results: - if vr.id in documents: - ordered_docs.append(documents[vr.id]) - - return ordered_docs + if db is not None: + return await fetch_documents_by_ids( + db=db, + workspace_name=workspace_name, + observer=observer, + observed=observed, + document_ids=document_ids, + filters=filters, + ) + async with tracked_db("query_documents.fetch") as managed_db: + docs = await fetch_documents_by_ids( + db=managed_db, + workspace_name=workspace_name, + observer=observer, + observed=observed, + document_ids=document_ids, + filters=filters, + ) + for doc in docs: + managed_db.expunge(doc) + return docs async def create_documents( @@ -321,7 +424,7 @@ async def create_documents( observer: str, observed: str, deduplicate: bool = False, -) -> int: +) -> list[schemas.DocumentCreate]: """ Create multiple documents with optional duplicate detection. @@ -333,9 +436,11 @@ async def create_documents( observed: Name of the observed peer Returns: - Count of new documents + List of DocumentCreate schemas that were actually inserted (excludes + duplicates and failures). """ honcho_documents: list[models.Document] = [] + accepted_documents: list[schemas.DocumentCreate] = [] # Store (document_model, embedding) pairs - IDs aren't available until after commit docs_with_embeddings: list[tuple[models.Document, list[float]]] = [] @@ -391,6 +496,7 @@ async def create_documents( if doc.embedding: new_doc.sync_state = "pending" honcho_documents.append(new_doc) + accepted_documents.append(doc) # Track embedding for vector store (ID will be available after commit) if doc.embedding: @@ -489,7 +595,7 @@ async def create_documents( "Failed to create documents due to integrity constraint violation" ) from e - return len(honcho_documents) + return accepted_documents async def delete_document( diff --git a/src/crud/message.py b/src/crud/message.py index b73dc414..fb9733af 100644 --- a/src/crud/message.py +++ b/src/crud/message.py @@ -739,6 +739,7 @@ async def search_messages_temporal( before_date: datetime | None = None, limit: int = 10, context_window: int = 2, + embedding: list[float] | None = None, ) -> list[tuple[list[models.Message], list[models.Message]]]: """ Search for messages using semantic similarity with optional date filtering. @@ -755,13 +756,16 @@ async def search_messages_temporal( before_date: Only return messages before this datetime limit: Maximum number of matching messages to return context_window: Number of messages before/after each match to include + embedding: Optional pre-computed embedding for the query Returns: List of tuples: (matched_messages, context_messages) Each snippet may contain multiple matches if they were close together. """ - # Generate embedding for the search query - query_embedding = await embedding_client.embed(query) + # Use provided embedding or generate one + query_embedding = ( + embedding if embedding is not None else await embedding_client.embed(query) + ) # Build query with date filters match_stmt = ( diff --git a/src/crud/representation.py b/src/crud/representation.py index ca246bef..616a17d1 100644 --- a/src/crud/representation.py +++ b/src/crud/representation.py @@ -162,7 +162,7 @@ class RepresentationManager: ) # Use bulk creation with optional duplicate detection - new_documents = await crud.create_documents( + accepted_documents = await crud.create_documents( db, documents_to_create, self.workspace_name, @@ -177,7 +177,7 @@ class RepresentationManager: except Exception as e: logger.warning(f"Failed to check dream scheduling: {e}") - return new_documents + return len(accepted_documents) async def get_working_representation( self, diff --git a/src/dialectic/chat.py b/src/dialectic/chat.py index 5ea0e11f..e9659a21 100644 --- a/src/dialectic/chat.py +++ b/src/dialectic/chat.py @@ -39,13 +39,12 @@ async def agentic_chat( Returns: The synthesized answer string """ - async with tracked_db("dialectic.agentic_chat") as db: - # Validate that the peers exist before proceeding + # Short-lived DB session for validation + config + async with tracked_db("dialectic.preflight") as db: await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observer)) if observer != observed: await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observed)) - # Resolve configuration to check if peer cards should be used session = None if session_name: session = await crud.get_session( @@ -54,7 +53,6 @@ async def agentic_chat( workspace = await crud.get_workspace(db, workspace_name=workspace_name) configuration = get_configuration(None, session, workspace) - # Get peer cards for context (if enabled) observer_peer_card = None observed_peer_card = None if configuration.peer_card.use: @@ -65,22 +63,19 @@ async def agentic_chat( observed_peer_card = await crud.get_peer_card( db, workspace_name, observer=observer, observed=observed ) + # DB session closed — agent runs without holding a connection - # Create and run the dialectic agent - agent = DialecticAgent( - db=db, - workspace_name=workspace_name, - session_name=session_name, - observer=observer, - observed=observed, - observer_peer_card=observer_peer_card, - observed_peer_card=observed_peer_card, - reasoning_level=reasoning_level, - ) + agent = DialecticAgent( + workspace_name=workspace_name, + session_name=session_name, + observer=observer, + observed=observed, + observer_peer_card=observer_peer_card, + observed_peer_card=observed_peer_card, + reasoning_level=reasoning_level, + ) - response = await agent.answer(query) - - return response + return await agent.answer(query) async def agentic_chat_stream( @@ -105,13 +100,12 @@ async def agentic_chat_stream( Yields: Chunks of the response text as they are generated """ - async with tracked_db("dialectic.agentic_chat_stream") as db: - # Validate that the peers exist before proceeding + # Short-lived DB session for validation + config + async with tracked_db("dialectic.preflight") as db: await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observer)) if observer != observed: await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observed)) - # Resolve configuration to check if peer cards should be used session = None if session_name: session = await crud.get_session( @@ -120,7 +114,6 @@ async def agentic_chat_stream( workspace = await crud.get_workspace(db, workspace_name=workspace_name) configuration = get_configuration(None, session, workspace) - # Get peer cards for context (if enabled) observer_peer_card = None observed_peer_card = None if configuration.peer_card.use: @@ -131,18 +124,17 @@ async def agentic_chat_stream( observed_peer_card = await crud.get_peer_card( db, workspace_name, observer=observer, observed=observed ) + # DB session closed — agent streams without holding a connection - # Create and run the dialectic agent - agent = DialecticAgent( - db=db, - workspace_name=workspace_name, - session_name=session_name, - observer=observer, - observed=observed, - observer_peer_card=observer_peer_card, - observed_peer_card=observed_peer_card, - reasoning_level=reasoning_level, - ) + agent = DialecticAgent( + workspace_name=workspace_name, + session_name=session_name, + observer=observer, + observed=observed, + observer_peer_card=observer_peer_card, + observed_peer_card=observed_peer_card, + reasoning_level=reasoning_level, + ) - async for chunk in agent.answer_stream(query): - yield chunk + async for chunk in agent.answer_stream(query): + yield chunk diff --git a/src/dialectic/core.py b/src/dialectic/core.py index 6dd215b5..f2dee426 100644 --- a/src/dialectic/core.py +++ b/src/dialectic/core.py @@ -11,10 +11,9 @@ import uuid from collections.abc import AsyncIterator, Callable from typing import Any, cast -from sqlalchemy.ext.asyncio import AsyncSession - from src import crud from src.config import ReasoningLevel, settings +from src.dependencies import tracked_db from src.dialectic import prompts from src.embedding_client import embedding_client from src.telemetry import prometheus_metrics @@ -52,7 +51,6 @@ class DialecticAgent: def __init__( self, - db: AsyncSession, workspace_name: str, session_name: str | None, observer: str, @@ -66,7 +64,6 @@ class DialecticAgent: Initialize the dialectic agent. Args: - db: Database session workspace_name: Workspace identifier session_name: Session identifier (may be None for global queries) observer: The peer making the query @@ -76,7 +73,6 @@ class DialecticAgent: metric_key: Optional key for logging metrics (if provided, agent won't log separately) reasoning_level: Level of reasoning to apply """ - self.db: AsyncSession = db self.workspace_name: str = workspace_name self.session_name: str | None = session_name self.observer: str = observer @@ -118,19 +114,20 @@ class DialecticAgent: token_limit=max_tokens, reverse=False, # chronological order ) - result = await self.db.execute(stmt) - messages = result.scalars().all() + async with tracked_db("dialectic.session_history") as db: + result = await db.execute(stmt) + messages = result.scalars().all() - if not messages: - return + if not messages: + return - # Format messages for injection - formatted_messages: list[str] = [] - for msg in messages: - formatted = format_new_turn_with_timestamp( - msg.content, msg.created_at, msg.peer_name - ) - formatted_messages.append(formatted) + # Format messages for injection (must access ORM attrs before session closes) + formatted_messages: list[str] = [] + for msg in messages: + formatted = format_new_turn_with_timestamp( + msg.content, msg.created_at, msg.peer_name + ) + formatted_messages.append(formatted) session_history_section = ( "\n\n## SESSION HISTORY\n\n" @@ -169,12 +166,12 @@ class DialecticAgent: prefetch_limit = 10 if self.reasoning_level == "minimal" else 25 try: - # Pre-compute embedding once for both searches + # Pre-compute embedding once for both searches (no DB needed) query_embedding = await embedding_client.embed(query) - # Search explicit observations separately + # search_memory manages its own short-lived DB sessions so no + # connection is held during external vector-store calls. explicit_repr = await search_memory( - db=self.db, workspace_name=self.workspace_name, observer=self.observer, observed=self.observed, @@ -184,9 +181,7 @@ class DialecticAgent: embedding=query_embedding, ) - # Search derived observations separately derived_repr = await search_memory( - db=self.db, workspace_name=self.workspace_name, observer=self.observer, observed=self.observed, @@ -280,7 +275,6 @@ class DialecticAgent: tool_executor: Callable[ [str, dict[str, Any]], Any ] = await create_tool_executor( - db=self.db, workspace_name=self.workspace_name, session_name=self.session_name, observer=self.observer, diff --git a/src/dreamer/orchestrator.py b/src/dreamer/orchestrator.py index adb8329b..000d45a6 100644 --- a/src/dreamer/orchestrator.py +++ b/src/dreamer/orchestrator.py @@ -20,7 +20,6 @@ from dataclasses import dataclass from typing import Any import sentry_sdk -from sqlalchemy.ext.asyncio import AsyncSession from src import crud from src.config import settings @@ -64,7 +63,6 @@ class DreamResult: async def run_dream( - db: AsyncSession, workspace_name: str, observer: str, observed: str, @@ -78,8 +76,9 @@ async def run_dream( 1. Deduction specialist: Creates deductive observations from explicit facts 2. Induction specialist: Creates inductive observations from patterns + Uses short-lived DB sessions to avoid holding connections during LLM calls. + Args: - db: Database session workspace_name: Workspace identifier observer: Observer peer name observed: Observed peer name @@ -96,16 +95,17 @@ async def run_dream( f"[{run_id}] Starting dream cycle for {workspace_name}/{observer}/{observed}" ) - if session_name is not None: - session = await crud.get_session( - db, workspace_name=workspace_name, session_name=session_name - ) - else: - session = None + # Short-lived DB session for config resolution + async with tracked_db("dream.config") as db: + if session_name is not None: + session = await crud.get_session( + db, workspace_name=workspace_name, session_name=session_name + ) + else: + session = None - workspace = await crud.get_workspace(db, workspace_name=workspace_name) - - configuration = get_configuration(None, session, workspace) + workspace = await crud.get_workspace(db, workspace_name=workspace_name) + configuration = get_configuration(None, session, workspace) if not configuration.dream.enabled: logger.info( f"[{run_id}] Dreams disabled for {workspace_name}/{session_name}, skipping dream" @@ -129,7 +129,6 @@ async def run_dream( from src.dreamer.surprisal import sample_observations_with_surprisal high_surprisal_obs = await sample_observations_with_surprisal( - db=db, workspace_name=workspace_name, observer=observer, observed=observed, @@ -164,12 +163,11 @@ async def run_dream( accumulate_metric(task_name, "surprisal_error", str(e), "blob") # Specialists will explore freely without hints - # Phase 1: Run deduction specialist + # Phase 1: Run deduction specialist (manages its own DB sessions) logger.info(f"[{run_id}] Phase 1: Running deduction specialist") deduction_specialist = SPECIALISTS["deduction"] try: deduction_result = await deduction_specialist.run( - db=db, workspace_name=workspace_name, observer=observer, observed=observed, @@ -194,7 +192,6 @@ async def run_dream( induction_specialist = SPECIALISTS["induction"] try: induction_result = await induction_specialist.run( - db=db, workspace_name=workspace_name, observer=observer, observed=observed, @@ -311,14 +308,12 @@ DREAM: {payload.dream_type} documents for {workspace_name}/{payload.observer}/{p try: match payload.dream_type: case DreamType.OMNI: - async with tracked_db("dream_orchestrator") as db: - result = await run_dream( - db=db, - workspace_name=workspace_name, - observer=payload.observer, - observed=payload.observed, - session_name=payload.session_name, - ) + result = await run_dream( + workspace_name=workspace_name, + observer=payload.observer, + observed=payload.observed, + session_name=payload.session_name, + ) # Log completion (telemetry event already emitted in run_dream) if result is not None: diff --git a/src/dreamer/specialists.py b/src/dreamer/specialists.py index 70d51c53..608db3b0 100644 --- a/src/dreamer/specialists.py +++ b/src/dreamer/specialists.py @@ -18,10 +18,9 @@ from collections.abc import Callable from dataclasses import dataclass from typing import Any -from sqlalchemy.ext.asyncio import AsyncSession - from src import crud, schemas from src.config import settings +from src.dependencies import tracked_db from src.schemas import ResolvedConfiguration from src.telemetry import prometheus_metrics from src.telemetry.events import DreamSpecialistEvent, emit @@ -116,7 +115,6 @@ If you update it, send the full deduplicated list and remove stale entries. async def run( self, - db: AsyncSession, workspace_name: str, observer: str, observed: str, @@ -128,8 +126,9 @@ If you update it, send the full deduplicated list and remove stale entries. """ Run the specialist agent. + Uses short-lived DB sessions to avoid holding connections during LLM calls. + Args: - db: Database session workspace_name: Workspace identifier observer: The observing peer observed: The peer being observed @@ -145,23 +144,27 @@ If you update it, send the full deduplicated list and remove stale entries. task_name = f"dreamer_{self.name}_{run_id}" start_time = time.perf_counter() - # Validate that the peers exist before proceeding - await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observer)) - if observer != observed: - await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observed)) + # Short-lived DB session for preflight operations + async with tracked_db("dream.specialist.preflight") as db: + await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observer)) + if observer != observed: + await crud.get_peer( + db, workspace_name, schemas.PeerCreate(name=observed) + ) - # Determine if peer card tools should be included - peer_card_enabled = configuration is None or configuration.peer_card.create + # Determine if peer card tools should be included + peer_card_enabled = configuration is None or configuration.peer_card.create - # Fetch current peer card to inject into prompt (saves a tool call) - current_peer_card: list[str] | None = None - if peer_card_enabled: - current_peer_card = await crud.get_peer_card( - db, - workspace_name=workspace_name, - observer=observer, - observed=observed, - ) + # Fetch current peer card to inject into prompt (saves a tool call) + current_peer_card: list[str] | None = None + if peer_card_enabled: + current_peer_card = await crud.get_peer_card( + db, + workspace_name=workspace_name, + observer=observer, + observed=observed, + ) + # DB session closed — LLM calls happen without holding a connection # Build messages messages: list[dict[str, str]] = [ @@ -181,7 +184,6 @@ If you update it, send the full deduplicated list and remove stale entries. tool_executor: Callable[ [str, dict[str, Any]], Any ] = await create_tool_executor( - db=db, workspace_name=workspace_name, observer=observer, observed=observed, diff --git a/src/dreamer/surprisal.py b/src/dreamer/surprisal.py index faaf7c7e..093b2d7b 100644 --- a/src/dreamer/surprisal.py +++ b/src/dreamer/surprisal.py @@ -18,22 +18,32 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import models from src.config import settings from src.crud.document import get_all_documents +from src.dependencies import tracked_db from src.dreamer.trees import SurprisalTree, create_tree logger = logging.getLogger(__name__) +@dataclass +class ObservationData: + """Plain data extracted from a Document ORM object (session-safe).""" + + id: str + content: str + level: str | None + embedding: np.ndarray | list[float] + + @dataclass class SurprisalScore: """Container for observation with surprisal score.""" - observation: models.Document + observation: ObservationData surprisal: float embedding: np.ndarray async def sample_observations_with_surprisal( - db: AsyncSession, workspace_name: str, observer: str, observed: str, @@ -42,15 +52,14 @@ async def sample_observations_with_surprisal( Sample observations and compute surprisal scores. Workflow: - 1. Fetch observations based on SAMPLING_STRATEGY - 2. Extract embeddings from DB (already stored) + 1. Fetch observations based on SAMPLING_STRATEGY (short DB scope) + 2. Extract embeddings (already stored on objects, no DB needed) 3. Build tree structure using trees.create_tree() 4. Compute surprisal for each observation 5. Rank by surprisal (highest first) 6. Filter by threshold and take top N Args: - db: Database session workspace_name: Workspace identifier observer: Observer peer name observed: Observed peer name @@ -59,13 +68,29 @@ async def sample_observations_with_surprisal( List of SurprisalScore objects, ranked by surprisal (highest first) """ try: - # 1. Fetch observations - observations = await _fetch_observations( - db=db, - workspace_name=workspace_name, - observer=observer, - observed=observed, - ) + # 1. Fetch observations (short DB scope, closed before compute) + async with tracked_db("dream.surprisal.fetch") as db: + raw_observations = await _fetch_observations( + db=db, + workspace_name=workspace_name, + observer=observer, + observed=observed, + ) + # Extract plain data from ORM objects before session closes, + # skipping any with null embeddings (can't compute surprisal). + observations = [ + ObservationData( + id=obs.id, + content=obs.content, + level=obs.level, + embedding=obs.embedding, + ) + for obs in raw_observations + if obs.embedding is not None + ] + skipped = len(raw_observations) - len(observations) + if skipped: + logger.warning(f"Skipped {skipped} observations with null embeddings") # Edge case: No observations if not observations: @@ -326,7 +351,7 @@ async def _fetch_all_observations( return list(result.scalars().all()) -def _extract_embeddings(observations: list[models.Document]) -> np.ndarray: +def _extract_embeddings(observations: list[ObservationData]) -> np.ndarray: """ Extract embeddings from observations as numpy array. @@ -370,7 +395,7 @@ def _build_tree(embeddings: np.ndarray) -> SurprisalTree: def _compute_surprisal_scores( - observations: list[models.Document], + observations: list[ObservationData], embeddings: np.ndarray, tree: SurprisalTree, ) -> list[SurprisalScore]: @@ -378,7 +403,7 @@ def _compute_surprisal_scores( Compute surprisal score for each observation. Args: - observations: List of Document objects + observations: List of ObservationData objects embeddings: np.ndarray of embeddings matching observations tree: Built SurprisalTree diff --git a/src/routers/peers.py b/src/routers/peers.py index ba30b959..01aa0f71 100644 --- a/src/routers/peers.py +++ b/src/routers/peers.py @@ -143,7 +143,6 @@ async def get_sessions_for_peer( @router.post( "/{peer_id}/chat", - summary="Query a Peer's representation using natural language", responses={ 200: { "content": { diff --git a/src/utils/agent_tools.py b/src/utils/agent_tools.py index 646087e5..a6c3009d 100644 --- a/src/utils/agent_tools.py +++ b/src/utils/agent_tools.py @@ -12,6 +12,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import crud, models, schemas from src.config import settings +from src.dependencies import tracked_db from src.embedding_client import embedding_client from src.models import Document from src.schemas import ResolvedConfiguration @@ -615,7 +616,6 @@ INDUCTION_SPECIALIST_TOOLS: list[dict[str, Any]] = [ async def create_observations( - db: AsyncSession, observations: list[schemas.ObservationInput], observer: str, observed: str, @@ -627,8 +627,9 @@ async def create_observations( """ Create multiple observations (documents) in the memory system in a single call. + Uses short-lived DB sessions to avoid holding connections during embedding API calls. + Args: - db: Database session observations: List of validated observation inputs observer: The peer making the observation observed: The peer being observed @@ -644,18 +645,17 @@ async def create_observations( logger.warning("create_observations called with empty list") return ObservationsCreatedResult(created_count=0, created_levels=[], failed=[]) - # Get or create collection - await crud.get_or_create_collection( - db, - workspace_name, - observer=observer, - observed=observed, - ) + # Phase 1: Ensure collection exists (short DB scope) + async with tracked_db("create_observations.collection") as db: + await crud.get_or_create_collection( + db, + workspace_name, + observer=observer, + observed=observed, + ) + # Phase 2: Compute embeddings (no DB needed) contents = [obs.content for obs in observations] - - # Batch embed all observation contents. - # If batching fails, fall back to per-observation embedding. embeddings_by_index: dict[int, list[float]] | None = None try: embeddings = await embedding_client.simple_batch_embed(contents) @@ -721,27 +721,29 @@ async def create_observations( ) documents.append(doc) - # Bulk create all documents + # Phase 3: Bulk create all documents (short DB scope) + accepted: list[schemas.DocumentCreate] = [] if documents: - await crud.create_documents( - db, - documents=documents, - workspace_name=workspace_name, - observer=observer, - observed=observed, - deduplicate=True, - ) + async with tracked_db("create_observations.save") as db: + accepted = await crud.create_documents( + db, + documents=documents, + workspace_name=workspace_name, + observer=observer, + observed=observed, + deduplicate=True, + ) logger.info( "Created %d observations in %s/%s/%s", - len(documents), + len(accepted), workspace_name, observer, observed, ) return ObservationsCreatedResult( - created_count=len(documents), - created_levels=[doc.level for doc in documents], + created_count=len(accepted), + created_levels=[doc.level for doc in accepted], failed=failed, ) @@ -801,7 +803,6 @@ async def get_recent_history( async def search_memory( - db: AsyncSession, workspace_name: str, observer: str, observed: str, @@ -813,8 +814,10 @@ async def search_memory( """ Search for observations in memory using semantic similarity. + Does not require a DB session — ``query_documents`` manages its own + short-lived sessions so no connection is held during external calls. + Args: - db: Database session workspace_name: Workspace identifier observer: The peer who made the observations observed: The peer who was observed @@ -833,7 +836,7 @@ async def search_memory( filters = {"level": {"in": levels}} documents = await crud.query_documents( - db=db, + db=None, workspace_name=workspace_name, observer=observer, observed=observed, @@ -910,7 +913,6 @@ async def get_observation_context( async def extract_preferences( - db: AsyncSession, workspace_name: str, session_name: str | None, observed: str, @@ -922,7 +924,6 @@ async def extract_preferences( This is language-agnostic and doesn't rely on keyword matching. Args: - db: Database session workspace_name: Workspace identifier session_name: Session identifier (optional) observed: The peer whose preferences to extract @@ -942,7 +943,7 @@ async def extract_preferences( "things user wants or does not want", ] - # Batch embed all queries in a single API call. + # Batch embed all queries in a single API call (no DB needed). # If batching fails, each search call will generate its own embedding. query_embeddings_by_query: dict[str, list[float]] | None = None try: @@ -958,26 +959,27 @@ async def extract_preferences( for query in semantic_queries: try: - snippets = await crud.search_messages( - db, - workspace_name=workspace_name, - session_name=session_name, - query=query, - limit=10, - context_window=0, - embedding=( - query_embeddings_by_query.get(query) - if query_embeddings_by_query is not None - else None - ), - ) - for matches, _ in snippets: - for msg in matches: - if msg.peer_name == observed: - content_key = msg.content[:100].lower() - if content_key not in seen_content: - seen_content.add(content_key) - messages.append(f"'{msg.content.strip()}'") + async with tracked_db("extract_preferences") as db: + snippets = await crud.search_messages( + db, + workspace_name=workspace_name, + session_name=session_name, + query=query, + limit=10, + context_window=0, + embedding=( + query_embeddings_by_query.get(query) + if query_embeddings_by_query is not None + else None + ), + ) + for matches, _ in snippets: + for msg in matches: + if msg.peer_name == observed: + content_key = msg.content[:100].lower() + if content_key not in seen_content: + seen_content.add(content_key) + messages.append(f"'{msg.content.strip()}'") except Exception as e: logger.warning("Error in semantic search for '%s': %s", query, e) @@ -992,7 +994,6 @@ async def extract_preferences( class ToolContext: """Context object passed to tool handlers.""" - db: AsyncSession workspace_name: str observer: str observed: str @@ -1068,7 +1069,6 @@ async def _handle_create_observations( # Use lock to serialize database writes (prevents concurrent commit issues) async with ctx.db_lock: result = await create_observations( - ctx.db, observations=observations, observer=ctx.observer, observed=ctx.observed, @@ -1178,9 +1178,9 @@ async def _handle_update_peer_card(ctx: ToolContext, tool_input: dict[str, Any]) ) normalized_peer_card = normalized_peer_card[:MAX_PEER_CARD_FACTS] - async with ctx.db_lock: + async with ctx.db_lock, tracked_db("tool.update_peer_card") as db: await crud.set_peer_card( - ctx.db, + db, workspace_name=ctx.workspace_name, peer_card=normalized_peer_card, observer=ctx.observer, @@ -1214,18 +1214,19 @@ async def _handle_get_recent_history( ) -> str: """Handle get_recent_history tool.""" _ = tool_input - history: list[models.Message] = await get_recent_history( - ctx.db, - workspace_name=ctx.workspace_name, - session_name=ctx.session_name, - observed=ctx.observed, - token_limit=ctx.history_token_limit, - ) - if not history: - return "No conversation history available" - history_text = "\n".join( - [f"{m.peer_name}: {_truncate_message_content(m.content)}" for m in history] - ) + async with tracked_db("tool.get_recent_history") as db: + history: list[models.Message] = await get_recent_history( + db, + workspace_name=ctx.workspace_name, + session_name=ctx.session_name, + observed=ctx.observed, + token_limit=ctx.history_token_limit, + ) + if not history: + return "No conversation history available" + history_text = "\n".join( + [f"{m.peer_name}: {_truncate_message_content(m.content)}" for m in history] + ) scope = ( f"from session {ctx.session_name}" if ctx.session_name @@ -1245,7 +1246,7 @@ async def _handle_search_memory(ctx: ToolContext, tool_input: dict[str, Any]) -> return f"ERROR: Query exceeds maximum token limit of {settings.MAX_EMBEDDING_TOKENS}. Please use a shorter query." documents = await crud.query_documents( - db=ctx.db, + db=None, workspace_name=ctx.workspace_name, observer=ctx.observer, observed=ctx.observed, @@ -1263,19 +1264,22 @@ async def _handle_search_memory(ctx: ToolContext, tool_input: dict[str, Any]) -> # automatically search the message history for relevant information. if ctx.agent_type == "dialectic": limit = min(_safe_int(tool_input.get("top_k"), 20), 20) - snippets = await crud.search_messages( - ctx.db, - workspace_name=ctx.workspace_name, - session_name=ctx.session_name, - query=query, - limit=limit, - context_window=0, - embedding=query_embedding, - ) - if snippets: - message_output = _format_message_snippets( - snippets, f"for query '{query}'" + message_output = None + async with tracked_db("tool.search_memory.fallback") as db: + snippets = await crud.search_messages( + db, + workspace_name=ctx.workspace_name, + session_name=ctx.session_name, + query=query, + limit=limit, + context_window=0, + embedding=query_embedding, ) + if snippets: + message_output = _format_message_snippets( + snippets, f"for query '{query}'" + ) + if message_output: return ( f"No observations yet. Message search results:\n\n{message_output}" ) @@ -1292,24 +1296,25 @@ async def _handle_get_observation_context( ctx: ToolContext, tool_input: dict[str, Any] ) -> str: """Handle get_observation_context tool.""" - messages = await get_observation_context( - ctx.db, - workspace_name=ctx.workspace_name, - session_name=ctx.session_name, - message_ids=tool_input["message_ids"], - ) - if not messages: - return f"No messages found for IDs {tool_input['message_ids']}" - messages_text = "\n".join( - [ - format_new_turn_with_timestamp( - _truncate_message_content(m.content), - m.created_at, - m.peer_name, - ) - for m in messages - ] - ) + async with tracked_db("tool.get_observation_context") as db: + messages = await get_observation_context( + db, + workspace_name=ctx.workspace_name, + session_name=ctx.session_name, + message_ids=tool_input["message_ids"], + ) + if not messages: + return f"No messages found for IDs {tool_input['message_ids']}" + messages_text = "\n".join( + [ + format_new_turn_with_timestamp( + _truncate_message_content(m.content), + m.created_at, + m.peer_name, + ) + for m in messages + ] + ) output = f"Retrieved {len(messages)} messages with context:\n{messages_text}" return _truncate_tool_output(output) @@ -1318,18 +1323,23 @@ async def _handle_search_messages(ctx: ToolContext, tool_input: dict[str, Any]) """Handle search_messages tool.""" query = tool_input["query"] limit = min(_safe_int(tool_input.get("limit"), 10), 20) # Cap at 20 - snippets = await crud.search_messages( - ctx.db, - workspace_name=ctx.workspace_name, - session_name=ctx.session_name, - query=query, - limit=limit, - context_window=2, - ) - if not snippets: - return f"No messages found for query '{query}'" - - return _format_message_snippets(snippets, f"for query '{query}'") + # Pre-compute embedding outside DB session to avoid holding a connection + # during the external API call (same pattern as _handle_search_memory). + query_embedding = await embedding_client.embed(query) + async with tracked_db("tool.search_messages") as db: + snippets = await crud.search_messages( + db, + workspace_name=ctx.workspace_name, + session_name=ctx.session_name, + query=query, + limit=limit, + context_window=2, + embedding=query_embedding, + ) + if not snippets: + return f"No messages found for query '{query}'" + formatted = _format_message_snippets(snippets, f"for query '{query}'") + return formatted async def _handle_grep_messages(ctx: ToolContext, tool_input: dict[str, Any]) -> str: @@ -1342,32 +1352,35 @@ async def _handle_grep_messages(ctx: ToolContext, tool_input: dict[str, Any]) -> _safe_int(tool_input.get("context_window"), 2), 2 ) # Cap context - snippets = await crud.grep_messages( - ctx.db, - workspace_name=ctx.workspace_name, - session_name=ctx.session_name, - text=text, - limit=limit, - context_window=context_window, - ) - if not snippets: - return f"No messages found containing '{text}'" - - # Format with pattern-based snippet extraction - snippet_texts: list[str] = [] - total_matches = sum(len(matches) for matches, _ in snippets) - for i, (matches, context) in enumerate(snippets, 1): - lines: list[str] = [] - for msg in context: - truncated = _extract_pattern_snippet(msg.content, text) - lines.append( - format_new_turn_with_timestamp(truncated, msg.created_at, msg.peer_name) - ) - sess = context[0].session_name if context else "unknown" - snippet_texts.append( - f"--- Snippet {i} (session: {sess}, {len(matches)} match(es)) ---\n" - + "\n".join(lines) + async with tracked_db("tool.grep_messages") as db: + snippets = await crud.grep_messages( + db, + workspace_name=ctx.workspace_name, + session_name=ctx.session_name, + text=text, + limit=limit, + context_window=context_window, ) + if not snippets: + return f"No messages found containing '{text}'" + + # Format with pattern-based snippet extraction + snippet_texts: list[str] = [] + total_matches = sum(len(matches) for matches, _ in snippets) + for i, (matches, context) in enumerate(snippets, 1): + lines: list[str] = [] + for msg in context: + truncated = _extract_pattern_snippet(msg.content, text) + lines.append( + format_new_turn_with_timestamp( + truncated, msg.created_at, msg.peer_name + ) + ) + sess = context[0].session_name if context else "unknown" + snippet_texts.append( + f"--- Snippet {i} (session: {sess}, {len(matches)} match(es)) ---\n" + + "\n".join(lines) + ) output = ( f"Found {total_matches} messages containing '{text}' in {len(snippets)} conversation snippets:\n\n" @@ -1403,15 +1416,29 @@ async def _handle_get_messages_by_date_range( if isinstance(before_date, str): return before_date # Error message - messages = await crud.get_messages_by_date_range( - ctx.db, - workspace_name=ctx.workspace_name, - session_name=ctx.session_name, - after_date=after_date, - before_date=before_date, - limit=limit, - order=order, - ) + async with tracked_db("tool.get_messages_by_date_range") as db: + messages = await crud.get_messages_by_date_range( + db, + workspace_name=ctx.workspace_name, + session_name=ctx.session_name, + after_date=after_date, + before_date=before_date, + limit=limit, + order=order, + ) + msg_count = len(messages) + messages_text = ( + "\n".join( + [ + format_new_turn_with_timestamp( + _truncate_message_content(m.content), m.created_at, m.peer_name + ) + for m in messages + ] + ) + if messages + else "" + ) date_range: list[str] = [] if after_date_str: @@ -1419,23 +1446,16 @@ async def _handle_get_messages_by_date_range( if before_date_str: date_range.append(f"before {before_date_str}") - if not messages: + if not msg_count: range_desc = " and ".join(date_range) if date_range else "specified range" return f"No messages found {range_desc}" - messages_text = "\n".join( - [ - format_new_turn_with_timestamp( - _truncate_message_content(m.content), m.created_at, m.peer_name - ) - for m in messages - ] - ) - range_desc = " and ".join(date_range) if date_range else "all time" order_desc = "oldest first" if order == "asc" else "newest first" - output = f"Found {len(messages)} messages ({range_desc}, {order_desc}):\n\n{messages_text}" + output = ( + f"Found {msg_count} messages ({range_desc}, {order_desc}):\n\n{messages_text}" + ) return _truncate_tool_output(output) @@ -1460,28 +1480,35 @@ async def _handle_search_messages_temporal( if isinstance(before_date, str): return before_date - snippets = await crud.search_messages_temporal( - ctx.db, - workspace_name=ctx.workspace_name, - session_name=ctx.session_name, - query=query, - after_date=after_date, - before_date=before_date, - limit=limit, - context_window=context_window, - ) + # Pre-compute embedding outside DB session to avoid holding a connection + # during the external API call. + query_embedding = await embedding_client.embed(query) + async with tracked_db("tool.search_messages_temporal") as db: + snippets = await crud.search_messages_temporal( + db, + workspace_name=ctx.workspace_name, + session_name=ctx.session_name, + query=query, + after_date=after_date, + before_date=before_date, + limit=limit, + context_window=context_window, + embedding=query_embedding, + ) + date_filter: list[str] = [] + if after_date_str: + date_filter.append(f"after {after_date_str}") + if before_date_str: + date_filter.append(f"before {before_date_str}") + filter_desc = f" ({' and '.join(date_filter)})" if date_filter else "" - date_filter: list[str] = [] - if after_date_str: - date_filter.append(f"after {after_date_str}") - if before_date_str: - date_filter.append(f"before {before_date_str}") - filter_desc = f" ({' and '.join(date_filter)})" if date_filter else "" + if not snippets: + return f"No messages found for query '{query}'{filter_desc}" - if not snippets: - return f"No messages found for query '{query}'{filter_desc}" - - return _format_message_snippets(snippets, f"for query '{query}'{filter_desc}") + formatted = _format_message_snippets( + snippets, f"for query '{query}'{filter_desc}" + ) + return formatted async def _handle_get_recent_observations( @@ -1489,15 +1516,16 @@ async def _handle_get_recent_observations( ) -> str: """Handle get_recent_observations tool.""" session_only = tool_input.get("session_only", False) - documents = await crud.query_documents_recent( - db=ctx.db, - workspace_name=ctx.workspace_name, - observer=ctx.observer, - observed=ctx.observed, - limit=min(_safe_int(tool_input.get("limit"), 10), 100), - session_name=ctx.session_name if session_only else None, - ) - representation = Representation.from_documents(documents) + async with tracked_db("tool.get_recent_observations") as db: + documents = await crud.query_documents_recent( + db=db, + workspace_name=ctx.workspace_name, + observer=ctx.observer, + observed=ctx.observed, + limit=min(_safe_int(tool_input.get("limit"), 10), 100), + session_name=ctx.session_name if session_only else None, + ) + representation = Representation.from_documents(documents) total_count = representation.len() if total_count == 0: return "No recent observations found" @@ -1514,14 +1542,15 @@ async def _handle_get_most_derived_observations( ctx: ToolContext, tool_input: dict[str, Any] ) -> str: """Handle get_most_derived_observations tool.""" - documents = await crud.query_documents_most_derived( - db=ctx.db, - workspace_name=ctx.workspace_name, - observer=ctx.observer, - observed=ctx.observed, - limit=min(_safe_int(tool_input.get("limit"), 10), 100), - ) - representation = Representation.from_documents(documents) + async with tracked_db("tool.get_most_derived_observations") as db: + documents = await crud.query_documents_most_derived( + db=db, + workspace_name=ctx.workspace_name, + observer=ctx.observer, + observed=ctx.observed, + limit=min(_safe_int(tool_input.get("limit"), 10), 100), + ) + representation = Representation.from_documents(documents) total_count = representation.len() if total_count == 0: return "No established observations found" @@ -1545,9 +1574,10 @@ async def _handle_get_session_summary( if summary_type == "long" else summarizer.SummaryType.SHORT ) - summary = await summarizer.get_summary( - ctx.db, ctx.workspace_name, ctx.session_name, st - ) + async with tracked_db("tool.get_session_summary") as db: + summary = await summarizer.get_summary( + db, ctx.workspace_name, ctx.session_name, st + ) if not summary: return "No session summary available yet" return f"Session summary ({summary['summary_type']}):\n{summary['content']}" @@ -1556,12 +1586,13 @@ async def _handle_get_session_summary( async def _handle_get_peer_card(ctx: ToolContext, tool_input: dict[str, Any]) -> str: """Handle get_peer_card tool.""" _ = tool_input - peer_card = await crud.get_peer_card( - ctx.db, - workspace_name=ctx.workspace_name, - observer=ctx.observer, - observed=ctx.observed, - ) + async with tracked_db("tool.get_peer_card") as db: + peer_card = await crud.get_peer_card( + db, + workspace_name=ctx.workspace_name, + observer=ctx.observer, + observed=ctx.observed, + ) if not peer_card: return f"No peer card available for {ctx.observed}" return f"Peer card for {ctx.observed}:\n" + "\n".join( @@ -1578,11 +1609,11 @@ async def _handle_delete_observations( return "ERROR: observation_ids list is empty" deleted_count = 0 - async with ctx.db_lock: + async with ctx.db_lock, tracked_db("tool.delete_observations") as db: for obs_id in observation_ids: try: await crud.delete_document( - ctx.db, + db, workspace_name=ctx.workspace_name, document_id=obs_id, observer=ctx.observer, @@ -1625,7 +1656,6 @@ async def _handle_extract_preferences( """Handle extract_preferences tool.""" _ = tool_input results = await extract_preferences( - ctx.db, workspace_name=ctx.workspace_name, session_name=ctx.session_name, observed=ctx.observed, @@ -1685,79 +1715,83 @@ async def _handle_get_reasoning_chain( return f"ERROR: Invalid direction '{direction}'. Must be 'premises', 'conclusions', or 'both'" # Get the observation itself - docs = await crud.get_documents_by_ids(ctx.db, ctx.workspace_name, [observation_id]) - if not docs or not docs[0]: - return f"ERROR: Observation '{observation_id}' not found" + async with tracked_db("tool.get_reasoning_chain") as db: + docs = await crud.get_documents_by_ids(db, ctx.workspace_name, [observation_id]) + if not docs or not docs[0]: + return f"ERROR: Observation '{observation_id}' not found" - doc: Document = docs[0] + doc: Document = docs[0] - output_parts: list[str] = [] + output_parts: list[str] = [] - # Format the main observation - level = doc.level or "explicit" - output_parts.append(f"**Observation [id:{doc.id}] ({level}):**\n{doc.content}") + # Format the main observation + level = doc.level or "explicit" + output_parts.append(f"**Observation [id:{doc.id}] ({level}):**\n{doc.content}") - # Get premises/sources if requested - if direction in ("premises", "both"): - if level == "deductive" and doc.source_ids: - premises = await crud.get_documents_by_ids( - ctx.db, ctx.workspace_name, doc.source_ids - ) - if premises: - premise_lines: list[Any] = [] - for p in premises: - p_level = p.level or "explicit" - premise_lines.append(f" - [id:{p.id}] ({p_level}): {p.content}") + # Get premises/sources if requested + if direction in ("premises", "both"): + if level == "deductive" and doc.source_ids: + premises = await crud.get_documents_by_ids( + db, ctx.workspace_name, doc.source_ids + ) + if premises: + premise_lines: list[Any] = [] + for p in premises: + p_level = p.level or "explicit" + premise_lines.append( + f" - [id:{p.id}] ({p_level}): {p.content}" + ) + output_parts.append( + f"\n**Premises ({len(premises)}):**\n" + + "\n".join(premise_lines) + ) + else: + output_parts.append( + f"\n**Premises:** Referenced {len(doc.source_ids)} premise IDs but none found in database" + ) + elif level == "inductive" and doc.source_ids: + sources = await crud.get_documents_by_ids( + db, ctx.workspace_name, doc.source_ids + ) + if sources: + source_lines: list[Any] = [] + for s in sources: + s_level = s.level or "explicit" + source_lines.append(f" - [id:{s.id}] ({s_level}): {s.content}") + output_parts.append( + f"\n**Sources ({len(sources)}):**\n" + "\n".join(source_lines) + ) + else: + output_parts.append( + f"\n**Sources:** Referenced {len(doc.source_ids)} source IDs but none found in database" + ) + elif level == "explicit": output_parts.append( - f"\n**Premises ({len(premises)}):**\n" + "\n".join(premise_lines) + "\n**Premises/Sources:** N/A (explicit observations have no premises)" ) else: - output_parts.append( - f"\n**Premises:** Referenced {len(doc.source_ids)} premise IDs but none found in database" - ) - elif level == "inductive" and doc.source_ids: - sources = await crud.get_documents_by_ids( - ctx.db, ctx.workspace_name, doc.source_ids + output_parts.append("\n**Premises/Sources:** None recorded") + + # Get conclusions if requested + if direction in ("conclusions", "both"): + children = await crud.get_child_observations( + db, + ctx.workspace_name, + observation_id, + observer=ctx.observer, + observed=ctx.observed, ) - if sources: - source_lines: list[Any] = [] - for s in sources: - s_level = s.level or "explicit" - source_lines.append(f" - [id:{s.id}] ({s_level}): {s.content}") + if children: + child_lines: list[Any] = [] + for c in children: + c_level = c.level or "explicit" + child_lines.append(f" - [id:{c.id}] ({c_level}): {c.content}") output_parts.append( - f"\n**Sources ({len(sources)}):**\n" + "\n".join(source_lines) + f"\n**Derived Conclusions ({len(children)}):**\n" + + "\n".join(child_lines) ) else: - output_parts.append( - f"\n**Sources:** Referenced {len(doc.source_ids)} source IDs but none found in database" - ) - elif level == "explicit": - output_parts.append( - "\n**Premises/Sources:** N/A (explicit observations have no premises)" - ) - else: - output_parts.append("\n**Premises/Sources:** None recorded") - - # Get conclusions if requested - if direction in ("conclusions", "both"): - children = await crud.get_child_observations( - ctx.db, - ctx.workspace_name, - observation_id, - observer=ctx.observer, - observed=ctx.observed, - ) - if children: - child_lines: list[Any] = [] - for c in children: - c_level = c.level or "explicit" - child_lines.append(f" - [id:{c.id}] ({c_level}): {c.content}") - output_parts.append( - f"\n**Derived Conclusions ({len(children)}):**\n" - + "\n".join(child_lines) - ) - else: - output_parts.append("\n**Derived Conclusions:** None found") + output_parts.append("\n**Derived Conclusions:** None found") return "\n".join(output_parts) @@ -1785,7 +1819,6 @@ _TOOL_HANDLERS: dict[str, Callable[[ToolContext, dict[str, Any]], Any]] = { async def create_tool_executor( - db: AsyncSession, workspace_name: str, observer: str, observed: str, @@ -1804,8 +1837,10 @@ async def create_tool_executor( This factory function captures the agent's context and returns an async callable that can execute any tool from AGENT_TOOLS or DIALECTIC_AGENT_TOOLS. + Each tool handler manages its own short-lived DB sessions via tracked_db(), + so no long-lived database session is needed. + Args: - db: Database session workspace_name: Workspace identifier observer: The peer making observations/queries observed: The peer being observed/queried about @@ -1826,7 +1861,6 @@ async def create_tool_executor( shared_lock = await get_observation_lock(workspace_name, observer, observed) ctx = ToolContext( - db=db, workspace_name=workspace_name, observer=observer, observed=observed, @@ -1877,9 +1911,8 @@ async def create_tool_executor( # We don't re-raise because the LLM should be able to continue with other tools error_msg = f"Tool {tool_name} failed unexpectedly: {type(e).__name__}: {e}" logger.error(error_msg, exc_info=True) - # Rollback the transaction to clear any failed state - # This is critical for PostgreSQL which blocks subsequent queries on failed transactions - await ctx.db.rollback() + # No explicit rollback needed — each handler uses tracked_db() which + # handles rollback in its finally block return error_msg return execute_tool diff --git a/tests/conftest.py b/tests/conftest.py index 8ff205ff..bd426e89 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -752,6 +752,11 @@ def mock_tracked_db(db_engine: AsyncEngine, request: pytest.FixtureRequest): patch("src.dialectic.chat.tracked_db", mock_tracked_db_context), patch("src.utils.summarizer.tracked_db", mock_tracked_db_context), patch("src.webhooks.events.tracked_db", mock_tracked_db_context), + patch("src.utils.agent_tools.tracked_db", mock_tracked_db_context), + patch("src.crud.document.tracked_db", mock_tracked_db_context), + patch("src.dialectic.core.tracked_db", mock_tracked_db_context), + patch("src.dreamer.specialists.tracked_db", mock_tracked_db_context), + patch("src.dreamer.surprisal.tracked_db", mock_tracked_db_context), ): yield diff --git a/tests/crud/test_document.py b/tests/crud/test_document.py index 6bfc0c47..f25b75d2 100644 --- a/tests/crud/test_document.py +++ b/tests/crud/test_document.py @@ -387,7 +387,7 @@ class TestDocumentCRUD: observed=test_peer2.name, ) - assert count == 2 + assert len(count) == 2 # Verify documents were created stmt = select(models.Document).where( diff --git a/tests/integration/test_token_metrics.py b/tests/integration/test_token_metrics.py index 31163474..7c24223a 100644 --- a/tests/integration/test_token_metrics.py +++ b/tests/integration/test_token_metrics.py @@ -707,7 +707,6 @@ class TestDialecticTokenMetrics: before = metric_checker.capture(dialectic_tokens_processed_counter, labels) agent = DialecticAgent( - db=db_session, workspace_name=workspace.name, session_name=session.name, observer=peer.name, @@ -755,7 +754,6 @@ class TestDialecticTokenMetrics: before = metric_checker.capture(dialectic_tokens_processed_counter, labels) agent = DialecticAgent( - db=db_session, workspace_name=workspace.name, session_name=session.name, observer=peer.name, @@ -819,7 +817,6 @@ class TestDialecticTokenMetrics: ) agent = DialecticAgent( - db=db_session, workspace_name=workspace.name, session_name=session.name, observer=peer.name, diff --git a/tests/utils/test_agent_tools.py b/tests/utils/test_agent_tools.py index 76c8b450..1387f681 100644 --- a/tests/utils/test_agent_tools.py +++ b/tests/utils/test_agent_tools.py @@ -118,15 +118,17 @@ async def tool_test_data( for doc in documents: await db_session.refresh(doc) - yield workspace, peer1, peer2, session, messages, documents + # Commit so data is visible to independent tracked_db sessions. + # Tool handlers no longer share the test's db_session — they open + # their own short-lived sessions via tracked_db. + # _truncate_all_tables handles cleanup between tests. + await db_session.commit() - await db_session.rollback() + yield workspace, peer1, peer2, session, messages, documents @pytest.fixture -def make_tool_context( - db_session: AsyncSession, tool_test_data: Any -) -> Callable[..., ToolContext]: +def make_tool_context(tool_test_data: Any) -> Callable[..., ToolContext]: """Factory fixture to create ToolContext with custom parameters.""" workspace, peer1, peer2, session, _messages, _ = tool_test_data shared_lock = asyncio.Lock() @@ -139,7 +141,6 @@ def make_tool_context( session_name: str | None = None, ) -> ToolContext: return ToolContext( - db=db_session, workspace_name=workspace.name, observer=peer1.name, observed=peer2.name, @@ -248,7 +249,6 @@ class TestCreateObservations: async def test_batch_embedding_failure_falls_back_to_individual_embeds( self, - db_session: AsyncSession, tool_test_data: Any, monkeypatch: pytest.MonkeyPatch, ): @@ -271,10 +271,10 @@ class TestCreateObservations: observer: str, observed: str, deduplicate: bool = False, - ) -> int: + ) -> list[Any]: _ = (workspace_name, observer, observed, deduplicate) created_documents.extend(documents) - return len(documents) + return documents monkeypatch.setattr( "src.utils.agent_tools.embedding_client.simple_batch_embed", @@ -289,7 +289,6 @@ class TestCreateObservations: ) result = await create_observations( - db_session, observations=[ schemas.ObservationInput(content="First obs", level="explicit"), schemas.ObservationInput(content="Second obs", level="explicit"), @@ -309,7 +308,6 @@ class TestCreateObservations: async def test_batch_embedding_failure_individual_embed_partial_failure( self, - db_session: AsyncSession, tool_test_data: Any, monkeypatch: pytest.MonkeyPatch, ): @@ -334,10 +332,10 @@ class TestCreateObservations: observer: str, observed: str, deduplicate: bool = False, - ) -> int: + ) -> list[Any]: _ = (workspace_name, observer, observed, deduplicate) created_documents.extend(documents) - return len(documents) + return documents monkeypatch.setattr( "src.utils.agent_tools.embedding_client.simple_batch_embed", @@ -352,7 +350,6 @@ class TestCreateObservations: ) result = await create_observations( - db_session, observations=[ schemas.ObservationInput(content="Embeds fine", level="explicit"), schemas.ObservationInput(content="Fails embed", level="explicit"), @@ -393,6 +390,8 @@ class TestDeleteObservations: assert "Deleted 1 observations" in result + # Expire the document so the identity map picks up the committed soft-delete + db_session.expire(documents[0]) # Verify soft-deletion (document still exists but has deleted_at timestamp) stmt = select(models.Document).where(models.Document.id == doc_id) doc = (await db_session.execute(stmt)).scalar_one_or_none() @@ -482,7 +481,6 @@ class TestSearchMemory: await db_session.flush() ctx = ToolContext( - db=db_session, workspace_name=workspace.name, observer=peer1.name, observed=peer2.name, @@ -655,14 +653,12 @@ class TestGetRecentHistory: async def test_without_session_uses_observed( self, - db_session: AsyncSession, tool_test_data: Any, ): """Without session, retrieves messages from observed peer.""" workspace, peer1, peer2, _, _, _ = tool_test_data ctx = ToolContext( - db=db_session, workspace_name=workspace.name, observer=peer1.name, observed=peer2.name, @@ -779,6 +775,8 @@ class TestUpdatePeerCard: assert "Updated peer card" in result + # Refresh the observer so the identity map picks up the committed update + await db_session.refresh(peer1) # Verify DB state peer_card = await crud.get_peer_card( db_session, @@ -804,6 +802,8 @@ class TestUpdatePeerCard: await _handle_update_peer_card(ctx, {"content": oversized}) + # Refresh the observer so the identity map picks up the committed update + await db_session.refresh(peer1) peer_card = await crud.get_peer_card( db_session, workspace_name=workspace.name, @@ -834,6 +834,8 @@ class TestUpdatePeerCard: result = await _handle_update_peer_card(ctx, {"content": None}) assert "empty" in result.lower() + # Refresh the observer so the identity map picks up the committed update + await db_session.refresh(peer1) # Verify original card is preserved peer_card = await crud.get_peer_card( db_session, @@ -861,6 +863,8 @@ class TestUpdatePeerCard: result = await _handle_update_peer_card(ctx, {"content": []}) assert "empty" in result.lower() + # Refresh the observer so the identity map picks up the committed update + await db_session.refresh(peer1) # Verify original card is preserved peer_card = await crud.get_peer_card( db_session, @@ -915,7 +919,6 @@ class TestGetPeerCard: await db_session.flush() ctx = ToolContext( - db=db_session, workspace_name=workspace.name, observer=peer1.name, observed=peer2.name, @@ -970,7 +973,6 @@ class TestExtractPreferences: async def test_falls_back_to_per_query_embedding_when_batch_fails( self, - db_session: AsyncSession, tool_test_data: Any, monkeypatch: pytest.MonkeyPatch, ): @@ -1023,7 +1025,6 @@ class TestExtractPreferences: ) result = await extract_preferences( - db_session, workspace_name=workspace.name, session_name=session.name, observed=observed_peer.name, @@ -1062,14 +1063,11 @@ class TestFinishConsolidation: class TestToolExecutor: """Tests for create_tool_executor and the executor function.""" - async def test_create_tool_executor_returns_callable( - self, db_session: AsyncSession, tool_test_data: Any - ): + async def test_create_tool_executor_returns_callable(self, tool_test_data: Any): """create_tool_executor returns an async callable.""" workspace, peer1, peer2, session, _, _ = tool_test_data executor = await create_tool_executor( - db=db_session, workspace_name=workspace.name, observer=peer1.name, observed=peer2.name, @@ -1078,14 +1076,11 @@ class TestToolExecutor: assert callable(executor) - async def test_executor_routes_to_correct_handler( - self, db_session: AsyncSession, tool_test_data: Any - ): + async def test_executor_routes_to_correct_handler(self, tool_test_data: Any): """Executor routes tool calls to correct handlers.""" workspace, peer1, peer2, session, _, _ = tool_test_data executor = await create_tool_executor( - db=db_session, workspace_name=workspace.name, observer=peer1.name, observed=peer2.name, @@ -1098,14 +1093,11 @@ class TestToolExecutor: # Should be from get_peer_card handler assert "peer card" in result.lower() or "No peer card" in result - async def test_executor_unknown_tool_returns_error( - self, db_session: AsyncSession, tool_test_data: Any - ): + async def test_executor_unknown_tool_returns_error(self, tool_test_data: Any): """Unknown tool name returns error message.""" workspace, peer1, peer2, session, _, _ = tool_test_data executor = await create_tool_executor( - db=db_session, workspace_name=workspace.name, observer=peer1.name, observed=peer2.name, @@ -1116,14 +1108,11 @@ class TestToolExecutor: assert "Unknown tool" in result - async def test_executor_handles_exceptions_gracefully( - self, db_session: AsyncSession, tool_test_data: Any - ): + async def test_executor_handles_exceptions_gracefully(self, tool_test_data: Any): """Executor converts exceptions to error strings instead of raising.""" workspace, peer1, peer2, session, _, _ = tool_test_data executor = await create_tool_executor( - db=db_session, workspace_name=workspace.name, observer=peer1.name, observed=peer2.name, @@ -1137,13 +1126,12 @@ class TestToolExecutor: # Should contain error info, not raise exception async def test_executor_dreamer_context_includes_observation_ids( - self, db_session: AsyncSession, tool_test_data: Any + self, tool_test_data: Any ): """Dreamer context (include_observation_ids=True) shows IDs in output.""" workspace, peer1, peer2, session, _, _ = tool_test_data executor = await create_tool_executor( - db=db_session, workspace_name=workspace.name, observer=peer1.name, observed=peer2.name,