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

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 <rahuja445@gmail.com>
This commit is contained in:
Vineeth Voruganti 2026-04-02 11:13:09 -04:00 committed by GitHub
parent cc3483bfaa
commit 0533c6dd26
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
15 changed files with 658 additions and 514 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/crud/__init__.py
View File

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

266
src/crud/document.py
View File

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

8
src/crud/message.py
View File

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

4
src/crud/representation.py
View File

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

62
src/dialectic/chat.py
View File

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

38
src/dialectic/core.py
View File

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

43
src/dreamer/orchestrator.py
View File

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

42
src/dreamer/specialists.py
View File

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

55
src/dreamer/surprisal.py
View File

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

1
src/routers/peers.py
View File

@ -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": {

575
src/utils/agent_tools.py
View File

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

5
tests/conftest.py
View File

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

2
tests/crud/test_document.py
View File

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

3
tests/integration/test_token_metrics.py
View File

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

64
tests/utils/test_agent_tools.py
View File

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