diff --git a/CLAUDE.md b/CLAUDE.md index f71650f9..7ace3529 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -113,6 +113,7 @@ cd sdks/typescript && bun run tsc --noEmit - Line length: 88 chars (Black compatible) - Explicit error handling with appropriate exception types - Docstrings: Use Google style docstrings +- **Never hold a DB session during external calls** (LLM, embedding, HTTP). If a function needs both a DB session and an external call result, compute the external result first and pass it as a parameter. This avoids tying up DB connections during slow network I/O. Use `tracked_db` for short-lived, DB-only operations; pass a shared session when multiple DB-only calls can reuse one connection. ### Agent Architecture diff --git a/src/crud/representation.py b/src/crud/representation.py index 91407cd8..83f0617a 100644 --- a/src/crud/representation.py +++ b/src/crud/representation.py @@ -3,6 +3,7 @@ from __future__ import annotations import datetime import logging import time +from contextlib import suppress from typing import Any from sqlalchemy import select @@ -181,8 +182,10 @@ class RepresentationManager: async def get_working_representation( self, *, + db: AsyncSession | None = None, session_name: str | None = None, include_semantic_query: str | None = None, + embedding: list[float] | None = None, semantic_search_top_k: int | None = None, semantic_search_max_distance: float | None = None, include_most_derived: bool = False, @@ -192,8 +195,11 @@ class RepresentationManager: Get working representation with flexible query options. Args: + db: Optional database session. If provided, uses it directly; + otherwise creates a new session via tracked_db. session_name: Optional session to filter by include_semantic_query: Query for semantic search + embedding: Pre-computed embedding for the semantic query. semantic_search_top_k: Number of semantic results semantic_search_max_distance: Maximum distance for semantic search include_most_derived: Include most derived observations @@ -202,13 +208,31 @@ class RepresentationManager: Returns: Representation combining various query strategies """ - async with tracked_db( - "representation_manager.get_working_representation" - ) as db: + if include_semantic_query and embedding is None: + with suppress(Exception): + # Best-effort precompute + embedding = await embedding_client.embed(include_semantic_query) + + if db is not None: return await self._get_working_representation_internal( db, session_name=session_name, include_semantic_query=include_semantic_query, + embedding=embedding, + semantic_search_top_k=semantic_search_top_k, + semantic_search_max_distance=semantic_search_max_distance, + include_most_derived=include_most_derived, + max_observations=max_observations, + ) + + async with tracked_db( + "representation_manager.get_working_representation" + ) as new_db: + return await self._get_working_representation_internal( + new_db, + session_name=session_name, + include_semantic_query=include_semantic_query, + embedding=embedding, semantic_search_top_k=semantic_search_top_k, semantic_search_max_distance=semantic_search_max_distance, include_most_derived=include_most_derived, @@ -223,6 +247,7 @@ class RepresentationManager: *, session_name: str | None = None, include_semantic_query: str | None = None, + embedding: list[float] | None = None, semantic_search_top_k: int | None = None, semantic_search_max_distance: float | None = None, include_most_derived: bool = False, @@ -268,6 +293,7 @@ class RepresentationManager: query=include_semantic_query, top_k=semantic_observations, max_distance=semantic_search_max_distance, + embedding=embedding, ) representation.merge_representation( Representation.from_documents(semantic_docs) @@ -298,6 +324,7 @@ class RepresentationManager: top_k: int, max_distance: float | None = None, level: str | None = None, + embedding: list[float] | None = None, ) -> list[models.Document]: """Query documents by semantic similarity.""" try: @@ -308,6 +335,7 @@ class RepresentationManager: level, top_k, max_distance, + embedding=embedding, ) else: documents = await crud.query_documents( @@ -318,6 +346,7 @@ class RepresentationManager: query=query, max_distance=max_distance, top_k=top_k, + embedding=embedding, ) db.expunge_all() return list(documents) @@ -391,6 +420,7 @@ class RepresentationManager: level: str, count: int, max_distance: float | None = None, + embedding: list[float] | None = None, ) -> list[models.Document]: """Query documents for a specific level.""" documents = await crud.query_documents( @@ -402,6 +432,7 @@ class RepresentationManager: max_distance=max_distance, top_k=count, filters=self._build_filter_conditions(level), + embedding=embedding, ) # Sort by creation time @@ -433,10 +464,12 @@ class RepresentationManager: async def get_working_representation( workspace_name: str, *, + db: AsyncSession | None = None, observer: str, observed: str, session_name: str | None = None, include_semantic_query: str | None = None, + embedding: list[float] | None = None, semantic_search_top_k: int | None = None, semantic_search_max_distance: float | None = None, include_most_derived: bool = False, @@ -447,6 +480,11 @@ async def get_working_representation( This is a convenience function that creates a RepresentationManager and calls get_working_representation on it. + + Args: + db: Optional database session. If provided, uses it directly; + otherwise creates a new session via tracked_db. + embedding: Pre-computed embedding for the semantic query. """ manager = RepresentationManager( workspace_name=workspace_name, @@ -454,8 +492,10 @@ async def get_working_representation( observed=observed, ) return await manager.get_working_representation( + db=db, session_name=session_name, include_semantic_query=include_semantic_query, + embedding=embedding, semantic_search_top_k=semantic_search_top_k, semantic_search_max_distance=semantic_search_max_distance, include_most_derived=include_most_derived, diff --git a/src/routers/sessions.py b/src/routers/sessions.py index a7be8943..f68f93ff 100644 --- a/src/routers/sessions.py +++ b/src/routers/sessions.py @@ -1,4 +1,5 @@ import logging +from contextlib import suppress from fastapi import APIRouter, Body, Depends, Path, Query, Response from fastapi_pagination import Page @@ -8,8 +9,9 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import config, crud, schemas from src.cache.client import safe_cache_delete from src.crud.session import session_cache_key -from src.dependencies import db, tracked_db +from src.dependencies import db from src.deriver.enqueue import enqueue_deletion +from src.embedding_client import embedding_client from src.exceptions import ( AuthenticationException, ResourceNotFoundException, @@ -30,6 +32,7 @@ router = APIRouter( async def _get_working_representation_task( + db: AsyncSession, workspace_id: str, last_message: str | None, *, @@ -40,11 +43,13 @@ async def _get_working_representation_task( search_max_distance: float | None, include_most_derived: bool, max_observations: int | None, + embedding: list[float] | None = None, ) -> Representation: """ - Atomic task to get working representation using tracked_db. + Get working representation using an externally-provided DB session. Args: + db: Database session to use for queries workspace_id: The workspace identifier last_message: Optional last message for semantic query observer: Name of the observer peer @@ -54,12 +59,14 @@ async def _get_working_representation_task( search_max_distance: Maximum distance to search for semantically relevant observations include_most_derived: Whether to include the most derived observations in the representation max_observations: Maximum number of observations to include in the representation + embedding: Pre-computed embedding for the semantic query Returns: The working representation """ return await crud.get_working_representation( workspace_name=workspace_id, + db=db, observer=observer, observed=observed, session_name=session_name, @@ -67,6 +74,7 @@ async def _get_working_representation_task( semantic_search_top_k=search_top_k, semantic_search_max_distance=search_max_distance, include_most_derived=include_most_derived, + embedding=embedding, max_observations=max_observations if max_observations is not None else config.settings.DERIVER.WORKING_REPRESENTATION_MAX_OBSERVATIONS, @@ -74,15 +82,17 @@ async def _get_working_representation_task( async def _get_peer_card_task( + db: AsyncSession, workspace_id: str, *, observer: str, observed: str, ) -> list[str] | None: """ - Atomic task to get peer card using tracked_db. + Get peer card using an externally-provided DB session. Args: + db: Database session to use for queries workspace_id: The workspace identifier observer: Name of the observer peer observed: Name of the observed peer @@ -90,25 +100,26 @@ async def _get_peer_card_task( Returns: The peer card or None if not found """ - async with tracked_db("get_peer_card") as db: - return await crud.get_peer_card( - db, - workspace_name=workspace_id, - observer=observer, - observed=observed, - ) + return await crud.get_peer_card( + db, + workspace_name=workspace_id, + observer=observer, + observed=observed, + ) async def _get_session_context_task( + db: AsyncSession, workspace_id: str, session_id: str, token_limit: int, include_summary: bool, ) -> tuple[schemas.Summary | None, list[schemas.Message]]: """ - Atomic task to get session context using tracked_db. + Get session context Args: + db: Database session to use for queries workspace_id: The workspace identifier session_id: The session identifier token_limit: Maximum tokens for the context @@ -117,17 +128,109 @@ async def _get_session_context_task( Returns: Tuple of (summary, messages) """ - async with tracked_db("get_session_context") as db: - summary, messages = await summarizer.get_session_context( - db, - workspace_name=workspace_id, - session_name=session_id, - token_limit=token_limit, - include_summary=include_summary, + summary, messages = await summarizer.get_session_context( + db, + workspace_name=workspace_id, + session_name=session_id, + token_limit=token_limit, + include_summary=include_summary, + ) + # Convert SQLAlchemy models to Pydantic schemas while session is active + message_schemas = [schemas.Message.model_validate(msg) for msg in messages] + return summary, message_schemas + + +async def _get_both_summaries_task( + db: AsyncSession, + workspace_id: str, + session_id: str, +) -> tuple[schemas.Summary | None, schemas.Summary | None]: + """ + Fetch both short and long summaries. + + Returns: + Tuple of (short_summary, long_summary) as Pydantic schemas. + """ + short_raw, long_raw = await summarizer.get_both_summaries( + db, workspace_name=workspace_id, session_name=session_id + ) + short = summarizer.to_schema_summary(short_raw) if short_raw else None + long = summarizer.to_schema_summary(long_raw) if long_raw else None + return short, long + + +async def _get_messages_for_context_task( + db: AsyncSession, + workspace_id: str, + session_id: str, + start_id: int, + token_limit: int, +) -> list[schemas.Message]: + """ + Fetch messages for context. + + Args: + db: Database session to use for queries + workspace_id: The workspace identifier + session_id: The session identifier + start_id: Internal message PK to start from (messages after summary coverage) + token_limit: Maximum tokens for the messages + + Returns: + List of messages as Pydantic schemas + """ + if token_limit <= 0: + return [] + messages = await crud.get_messages_id_range( + db, + workspace_id, + session_id, + start_id=start_id, + token_limit=token_limit, + ) + return [schemas.Message.model_validate(msg) for msg in messages] + + +def _select_summary_for_context( + short_summary: schemas.Summary | None, + long_summary: schemas.Summary | None, + token_limit: int, + include_summary: bool, +) -> tuple[schemas.Summary | None, int, int]: + """ + Pick the best summary that fits within the token budget using 40/60 allocation. + + Args: + short_summary: The short summary, or None + long_summary: The long summary, or None + token_limit: Total token budget for summary + messages + include_summary: Whether summaries should be considered + + Returns: + Tuple of (chosen_summary, messages_start_id, messages_token_budget) + """ + if not include_summary or token_limit <= 0: + return None, 0, max(token_limit, 0) + + summary_budget = int(token_limit * 0.4) + + long_len = long_summary.token_count if long_summary else 0 + short_len = short_summary.token_count if short_summary else 0 + + if long_summary and long_len <= summary_budget and long_len > short_len: + return ( + long_summary, + long_summary.message_id, + token_limit - long_len, ) - # Convert SQLAlchemy models to Pydantic schemas while session is active - message_schemas = [schemas.Message.model_validate(msg) for msg in messages] - return summary, message_schemas + if short_summary and short_len <= summary_budget and short_len > 0: + return ( + short_summary, + short_summary.message_id, + token_limit - short_len, + ) + + return None, 0, token_limit @router.post( @@ -506,6 +609,7 @@ async def get_session_peers( async def get_session_context( workspace_id: str = Path(...), session_id: str = Path(...), + db: AsyncSession = db, tokens: int | None = Query( None, le=config.settings.GET_CONTEXT_MAX_TOKENS, @@ -574,7 +678,7 @@ async def get_session_context( if not peer_target: # No representation or card needed summary, messages = await _get_session_context_task( - workspace_id, session_id, token_limit, include_summary + db, workspace_id, session_id, token_limit, include_summary ) return schemas.SessionContext( name=session_id, @@ -585,9 +689,15 @@ async def get_session_context( observer = peer_perspective or peer_target observed = peer_target - # Run representation and card tasks sequentially to avoid event loop issues - # with tracked_db creating separate database sessions + # Pre-compute embedding outside the DB session (best-effort) + embedding: list[float] | None = None + if search_query: + with suppress(Exception): + embedding = await embedding_client.embed(search_query) + + # Sequential calls on shared DB session representation = await _get_working_representation_task( + db, workspace_id, search_query, observer=observer, @@ -597,18 +707,28 @@ async def get_session_context( search_max_distance=search_max_distance, include_most_derived=include_most_frequent, max_observations=max_conclusions, + embedding=embedding, + ) + card = await _get_peer_card_task( + db, workspace_id, observer=observer, observed=observed + ) + short_summary, long_summary = await _get_both_summaries_task( + db, workspace_id, session_id ) - card = await _get_peer_card_task(workspace_id, observer=observer, observed=observed) - # adjust token limit downward to account for approximate token count of representation and card - # TODO determine if this impacts performance too much - adjusted_token_limit = ( + # Adjust token budget after accounting for representation + card tokens + adjusted_limit = ( token_limit - estimate_tokens(str(representation)) - estimate_tokens(card) ) - # Get the session context with the adjusted limit - summary, messages = await _get_session_context_task( - workspace_id, session_id, adjusted_token_limit, include_summary + # Pick best summary with 40/60 allocation against the adjusted budget + summary, messages_start_id, messages_budget = _select_summary_for_context( + short_summary, long_summary, adjusted_limit, include_summary + ) + + # Fetch messages with the correct start_id and budget + messages = await _get_messages_for_context_task( + db, workspace_id, session_id, messages_start_id, messages_budget ) return schemas.SessionContext( diff --git a/tests/conftest.py b/tests/conftest.py index 46cb67ca..ebb2db68 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -345,19 +345,19 @@ async def sample_data( # Create test app test_workspace = models.Workspace(name=str(generate_nanoid())) db_session.add(test_workspace) - await db_session.flush() # Create test user test_peer = models.Peer( name=str(generate_nanoid()), workspace_name=test_workspace.name ) db_session.add(test_peer) - await db_session.flush() + + # Commit so data is visible to independent tracked_db sessions. + # _truncate_all_tables handles cleanup between tests. + await db_session.commit() yield test_workspace, test_peer - await db_session.rollback() - @pytest.fixture(autouse=True) def mock_langfuse(): @@ -687,20 +687,26 @@ def mock_honcho_llm_call(): @pytest.fixture(autouse=True) -def mock_tracked_db(db_session: AsyncSession): - """Mock tracked_db to use the test database session""" +def mock_tracked_db(db_engine: AsyncEngine): + """Mock tracked_db to create fresh sessions per call. + + Using a session factory instead of a shared session avoids asyncio lock + errors when multiple tracked_db calls run concurrently via asyncio.gather. + """ from contextlib import asynccontextmanager + session_factory = async_sessionmaker(bind=db_engine, expire_on_commit=False) + @asynccontextmanager async def mock_tracked_db_context(_: str | None = None): - yield db_session + async with session_factory() as session: + yield session with ( patch("src.dependencies.tracked_db", mock_tracked_db_context), patch("src.deriver.queue_manager.tracked_db", mock_tracked_db_context), patch("src.deriver.consumer.tracked_db", mock_tracked_db_context), patch("src.deriver.enqueue.tracked_db", mock_tracked_db_context), - patch("src.routers.sessions.tracked_db", mock_tracked_db_context), patch("src.routers.peers.tracked_db", mock_tracked_db_context), patch("src.crud.representation.tracked_db", mock_tracked_db_context), patch("src.dreamer.orchestrator.tracked_db", mock_tracked_db_context), diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py index 51c8af07..409630bd 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -367,6 +367,7 @@ class TestQueueProcessing: qm.worker_ownership[worker_id] = WorkerOwnership( work_unit_key=work_unit_key, aqs_id=aqs_id ) + await db_session.commit() with patch( "src.deriver.queue_manager.process_representation_batch", @@ -790,6 +791,7 @@ class TestQueueProcessing: qm.worker_ownership[worker_id] = WorkerOwnership( work_unit_key=work_unit_key, aqs_id=aqs_id ) + await db_session.commit() with patch( "src.deriver.queue_manager.process_item", @@ -802,6 +804,9 @@ class TestQueueProcessing: assert all(batch["task_type"] == "summary" for batch in processed_batches) assert all(batch["payload_count"] == 1 for batch in processed_batches) + # Expire cached objects so we see updates made by tracked_db sessions + db_session.expire_all() + # Query for the summary queue items that were processed processed_items = ( ( @@ -929,6 +934,7 @@ class TestQueueProcessing: qm.worker_ownership[worker_id] = WorkerOwnership( work_unit_key=work_unit_key, aqs_id=aqs_id ) + await db_session.commit() with patch( "src.deriver.queue_manager.process_representation_batch", @@ -1048,6 +1054,7 @@ class TestQueueProcessing: qm.worker_ownership[worker_id] = WorkerOwnership( work_unit_key=work_unit_key, aqs_id=aqs_id ) + await db_session.commit() with patch( "src.deriver.queue_manager.process_representation_batch", diff --git a/tests/sdk_typescript/conftest.py b/tests/sdk_typescript/conftest.py index 292b559a..15f2c98b 100644 --- a/tests/sdk_typescript/conftest.py +++ b/tests/sdk_typescript/conftest.py @@ -129,7 +129,6 @@ def mock_tracked_db(ts_db_session: async_sessionmaker[AsyncSession]): patch("src.deriver.queue_manager.tracked_db", ts_tracked_db), patch("src.deriver.consumer.tracked_db", ts_tracked_db), patch("src.deriver.enqueue.tracked_db", ts_tracked_db), - patch("src.routers.sessions.tracked_db", ts_tracked_db), patch("src.routers.peers.tracked_db", ts_tracked_db), patch("src.crud.representation.tracked_db", ts_tracked_db), patch("src.dreamer.dream_scheduler.tracked_db", ts_tracked_db),