RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat: consolidate db calls in session context (#380) 

* feat: consolidate db calls in session context

* fix: guard against embedding failures

* Parallelize Async Calls (#383)

* fix: parallelize db calls in context()

* fix: (test) use session factory to further isolate tests

---------

Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
This commit is contained in:
Rajat Ahuja 2026-02-13 11:54:41 -05:00 committed by GitHub
parent 33ef0f8eab
commit 97df0a80cd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 216 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
CLAUDE.md
View File

@ -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

46
src/crud/representation.py
View File

@ -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,

182
src/routers/sessions.py
View File

@ -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(

22
tests/conftest.py
View File

@ -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),

7
tests/deriver/test_queue_processing.py
View File

@ -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",

1
tests/sdk_typescript/conftest.py
View File

@ -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),