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

perf: add redis caching and reduce embedding API calls (#372) 

* feat: add redis to get collection + get session peer config

* fix: reduce embedding calls

* fix: delete session caches when soft-deleting session

* fix: batch embeddings in create observatiosn tool

* fix: batch embeddings call in extract_preferences

* fix: double embed in fallback for _handle_search_memory

* fix: claude comments

* fix: return ObservationResult

* fix: make cache delete/set retryable. remove session peer config cache

* chore: claude nitpicks

---------

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

43
src/cache/client.py vendored
View File

@ -2,7 +2,7 @@ from __future__ import annotations
import asyncio
import logging
from typing import cast
from typing import Any, cast
import sentry_sdk
from cashews import cache
@ -11,6 +11,7 @@ from redis import exceptions as redis_exc
from tenacity import (
AsyncRetrying,
retry_if_exception_type,
stop_after_attempt,
stop_after_delay,
wait_exponential_jitter,
)
@ -112,6 +113,44 @@ async def init_cache() -> None:
cache.setup("mem://", pickle_type=PicklerType.SQLALCHEMY)
_TRANSIENT_CACHE_ERRORS = (
redis_exc.TimeoutError,
redis_exc.ConnectionError,
asyncio.TimeoutError,
TimeoutError,
)
async def safe_cache_set(key: str, value: Any, expire: int | float) -> None:
"""Best-effort cache set with retries on transient errors. Failures are logged but never propagate."""
try:
async for attempt in AsyncRetrying(
stop=stop_after_attempt(3),
wait=wait_exponential_jitter(initial=0.1, max=0.5),
retry=retry_if_exception_type(_TRANSIENT_CACHE_ERRORS),
reraise=True,
):
with attempt:
await cache.set(key, value, expire=expire)
except Exception:
logger.warning("Cache set failed for key %s", key, exc_info=True)
async def safe_cache_delete(key: str) -> None:
"""Best-effort cache delete with retries on transient errors. Failures are logged but never propagate."""
try:
async for attempt in AsyncRetrying(
stop=stop_after_attempt(3),
wait=wait_exponential_jitter(initial=0.1, max=0.5),
retry=retry_if_exception_type(_TRANSIENT_CACHE_ERRORS),
reraise=True,
):
with attempt:
await cache.delete(key)
except Exception:
logger.warning("Cache delete failed for key %s", key, exc_info=True)
async def close_cache() -> None:
await cache.close()
@ -120,4 +159,6 @@ __all__ = [
"init_cache",
"close_cache",
"cache",
"safe_cache_delete",
"safe_cache_set",
]

7
src/crud/__init__.py
View File

@ -1,4 +1,8 @@
from .collection import get_collection, get_or_create_collection
from .collection import (
get_collection,
get_or_create_collection,
update_collection_internal_metadata,
)
from .deriver import get_deriver_status, get_queue_status
from .document import (
create_documents,
@ -71,6 +75,7 @@ __all__ = [
# Collection
"get_collection",
"get_or_create_collection",
"update_collection_internal_metadata",
# Deriver
"get_deriver_status",
"get_queue_status",

99
src/crud/collection.py
View File

@ -1,14 +1,67 @@
from logging import getLogger
from typing import Any
from sqlalchemy import select
from cashews import NOT_NONE
from sqlalchemy import select, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from src import models
from src.cache.client import (
cache,
get_cache_namespace,
safe_cache_delete,
safe_cache_set,
)
from src.config import settings
from src.exceptions import ConflictException, ResourceNotFoundException
logger = getLogger(__name__)
COLLECTION_CACHE_KEY_TEMPLATE = (
"workspace:{workspace_name}:collection:{observer}:{observed}"
)
COLLECTION_LOCK_PREFIX = f"{get_cache_namespace()}:lock"
def collection_cache_key(workspace_name: str, observer: str, observed: str) -> str:
"""Generate cache key for collection."""
return (
get_cache_namespace()
+ ":"
+ COLLECTION_CACHE_KEY_TEMPLATE.format(
workspace_name=workspace_name,
observer=observer,
observed=observed,
)
)
@cache(
key=COLLECTION_CACHE_KEY_TEMPLATE,
ttl=f"{settings.CACHE.DEFAULT_TTL_SECONDS}s",
prefix=get_cache_namespace(),
condition=NOT_NONE,
)
@cache.locked(
key=COLLECTION_CACHE_KEY_TEMPLATE,
ttl=f"{settings.CACHE.DEFAULT_LOCK_TTL_SECONDS}s",
prefix=COLLECTION_LOCK_PREFIX,
)
async def _fetch_collection(
db: AsyncSession,
workspace_name: str,
observer: str,
observed: str,
) -> models.Collection | None:
"""Fetch a collection from the database."""
return await db.scalar(
select(models.Collection)
.where(models.Collection.workspace_name == workspace_name)
.where(models.Collection.observer == observer)
.where(models.Collection.observed == observed)
)
async def get_collection(
db: AsyncSession,
@ -32,16 +85,11 @@ async def get_collection(
Raises:
ResourceNotFoundException: If the collection does not exist
"""
stmt = (
select(models.Collection)
.where(models.Collection.workspace_name == workspace_name)
.where(models.Collection.observer == observer)
.where(models.Collection.observed == observed)
)
result = await db.execute(stmt)
collection = result.scalar_one_or_none()
collection = await _fetch_collection(db, workspace_name, observer, observed)
if collection is None:
raise ResourceNotFoundException("Collection not found")
# Merge cached object into session (cached objects are detached)
collection = await db.merge(collection, load=False)
return collection
@ -66,6 +114,14 @@ async def get_or_create_collection(
)
db.add(honcho_collection)
await db.commit()
key = collection_cache_key(workspace_name, observer, observed)
await safe_cache_set(
key,
honcho_collection,
expire=settings.CACHE.DEFAULT_TTL_SECONDS,
)
return honcho_collection
except IntegrityError:
await db.rollback()
@ -76,3 +132,28 @@ async def get_or_create_collection(
return await get_or_create_collection(
db, workspace_name, observer=observer, observed=observed, _retry=True
)
async def update_collection_internal_metadata(
db: AsyncSession,
workspace_name: str,
observer: str,
observed: str,
update_data: dict[str, Any],
) -> None:
"""Merge a patch into a collection's internal_metadata (JSONB ||) and invalidate cache."""
stmt = (
update(models.Collection)
.where(
models.Collection.workspace_name == workspace_name,
models.Collection.observer == observer,
models.Collection.observed == observed,
)
.values(
internal_metadata=models.Collection.internal_metadata.op("||")(update_data)
)
)
await db.execute(stmt)
await db.commit()
await safe_cache_delete(collection_cache_key(workspace_name, observer, observed))

8
src/crud/message.py
View File

@ -585,6 +585,7 @@ async def search_messages(
query: str,
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 and return conversation snippets.
@ -599,14 +600,17 @@ async def search_messages(
query: Search query text
limit: Maximum number of matching messages to return
context_window: Number of messages before/after each match to include
embedding: Optional pre-computed embedding
Returns:
List of tuples: (matched_messages, context_messages)
Each snippet may contain multiple matches if they were close together.
Context messages are ordered chronologically and include the matched messages.
"""
# 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)
)
# First, find the top matching messages
match_stmt = (

6
src/crud/peer.py
View File

@ -7,7 +7,7 @@ from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from src import models, schemas
from src.cache.client import cache, get_cache_namespace
from src.cache.client import cache, get_cache_namespace, safe_cache_delete
from src.config import settings
from src.crud.workspace import get_or_create_workspace
from src.exceptions import ConflictException, ResourceNotFoundException
@ -125,7 +125,7 @@ async def get_or_create_peers(
# Only invalidate cache for changed/new peers - read-through pattern
for peer_obj in changed_peers + new_peers:
cache_key = peer_cache_key(workspace_name, peer_obj.name)
await cache.delete(cache_key)
await safe_cache_delete(cache_key)
logger.debug(
"Peer %s cache invalidated in workspace %s (changed or new)",
peer_obj.name,
@ -254,7 +254,7 @@ async def update_peer(
await db.refresh(honcho_peer)
cache_key = peer_cache_key(workspace_name, honcho_peer.name)
await cache.delete(cache_key)
await safe_cache_delete(cache_key)
logger.debug("Peer %s updated successfully", peer_name)
return honcho_peer

4
src/crud/peer_card.py
View File

@ -7,7 +7,7 @@ from sqlalchemy import update
from sqlalchemy.ext.asyncio import AsyncSession
from src import exceptions, models, schemas
from src.cache.client import cache
from src.cache.client import safe_cache_delete
from src.crud.peer import get_or_create_peers, get_peer, peer_cache_key
logger = logging.getLogger(__name__)
@ -93,7 +93,7 @@ async def set_peer_card(
# Invalidate cache - read-through pattern
cache_key = peer_cache_key(workspace_name, observer)
await cache.delete(cache_key)
await safe_cache_delete(cache_key)
def construct_peer_card_label(*, observer: str, observed: str) -> str:

15
src/crud/session.py
View File

@ -11,7 +11,12 @@ from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.types import BigInteger, Boolean
from src import models, schemas
from src.cache.client import cache, get_cache_namespace
from src.cache.client import (
cache,
get_cache_namespace,
safe_cache_delete,
safe_cache_set,
)
from src.config import settings
from src.exceptions import (
ConflictException,
@ -231,7 +236,7 @@ async def get_or_create_session(
# Only update cache if session data changed or was newly created
if needs_cache_update:
cache_key = session_cache_key(workspace_name, session.name)
await cache.set(
await safe_cache_set(
cache_key, honcho_session, expire=settings.CACHE.DEFAULT_TTL_SECONDS
)
logger.debug(
@ -341,7 +346,7 @@ async def update_session(
# Only invalidate if we actually updated
cache_key = session_cache_key(workspace_name, session_name)
await cache.delete(cache_key)
await safe_cache_delete(cache_key)
logger.debug("Session %s updated successfully", session_name)
return honcho_session
@ -560,6 +565,10 @@ async def delete_session(
# Finally, delete the session itself
await db.delete(honcho_session)
await db.commit()
# Invalidate session cache
await safe_cache_delete(session_cache_key(workspace_name, session_name))
logger.debug("Session %s and all associated data deleted", session_name)
except Exception as e:
logger.error("Failed to delete session %s: %s", session_name, e)

11
src/crud/workspace.py
View File

@ -8,7 +8,12 @@ from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from src import models, schemas
from src.cache.client import cache, get_cache_namespace
from src.cache.client import (
cache,
get_cache_namespace,
safe_cache_delete,
safe_cache_set,
)
from src.config import settings
from src.exceptions import ConflictException, ResourceNotFoundException
from src.utils.filter import apply_filter
@ -108,7 +113,7 @@ async def get_or_create_workspace(
logger.debug("Workspace created successfully: %s", workspace.name)
cache_key = workspace_cache_key(workspace.name)
await cache.set(
await safe_cache_set(
cache_key, honcho_workspace, expire=settings.CACHE.DEFAULT_TTL_SECONDS
)
return GetOrCreateResult(honcho_workspace, created=True)
@ -221,7 +226,7 @@ async def update_workspace(
# Only invalidate if we actually updated
cache_key = workspace_cache_key(workspace_name)
await cache.delete(cache_key)
await safe_cache_delete(cache_key)
logger.debug("Workspace with id %s updated successfully", honcho_workspace.id)
return honcho_workspace

35
src/deriver/enqueue.py
View File

@ -2,7 +2,7 @@ import logging
from datetime import datetime, timezone
from typing import Any, Literal
from sqlalchemy import exists, insert, select, update
from sqlalchemy import exists, insert, select
from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, models, schemas
@ -516,28 +516,21 @@ async def enqueue_dream(
stmt = insert(QueueItem).returning(QueueItem)
await db_session.execute(stmt, [dream_record])
# Update collection metadata
# Update collection metadata (CRUD handles cache invalidation)
now_iso = datetime.now(timezone.utc).isoformat()
update_stmt = (
update(models.Collection)
.where(
models.Collection.workspace_name == workspace_name,
models.Collection.observer == observer,
models.Collection.observed == observed,
)
.values(
internal_metadata=models.Collection.internal_metadata.op("||")(
{
"dream": {
"last_dream_document_count": document_count,
"last_dream_at": now_iso,
}
}
)
)
await crud.update_collection_internal_metadata(
db_session,
workspace_name,
observer,
observed,
update_data={
"dream": {
"last_dream_document_count": document_count,
"last_dream_at": now_iso,
}
},
)
await db_session.execute(update_stmt)
await db_session.commit()
# update_collection_internal_metadata commits already
logger.info(
"Enqueued dream task for %s/%s/%s (type: %s)",

6
src/dialectic/core.py
View File

@ -16,6 +16,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
from src import crud
from src.config import ReasoningLevel, settings
from src.dialectic import prompts
from src.embedding_client import embedding_client
from src.telemetry import prometheus_metrics
from src.telemetry.events import DialecticCompletedEvent, emit
from src.telemetry.logging import (
@ -168,6 +169,9 @@ class DialecticAgent:
prefetch_limit = 10 if self.reasoning_level == "minimal" else 25
try:
# Pre-compute embedding once for both searches
query_embedding = await embedding_client.embed(query)
# Search explicit observations separately
explicit_repr = await search_memory(
db=self.db,
@ -177,6 +181,7 @@ class DialecticAgent:
query=query,
limit=prefetch_limit,
levels=["explicit"],
embedding=query_embedding,
)
# Search derived observations separately
@ -188,6 +193,7 @@ class DialecticAgent:
query=query,
limit=prefetch_limit,
levels=["deductive", "inductive", "contradiction"],
embedding=query_embedding,
)
if explicit_repr.is_empty() and derived_repr.is_empty():

4
src/routers/sessions.py
View File

@ -6,6 +6,8 @@ from fastapi_pagination.ext.sqlalchemy import apaginate
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.deriver.enqueue import enqueue_deletion
from src.exceptions import (
@ -265,6 +267,8 @@ async def delete_session(
await db.commit()
await safe_cache_delete(session_cache_key(workspace_id, session_id))
logger.debug("Session %s marked as inactive, deletion enqueued", session_id)
return {"message": "Session deleted successfully"}
except ValueError as e:

36
src/schemas.py
View File

@ -1,7 +1,7 @@
import datetime
import ipaddress
from enum import Enum
from typing import Annotated, Any, Self, cast
from typing import Annotated, Any, Literal, Self, cast
from urllib.parse import urlparse
import tiktoken
@ -562,6 +562,40 @@ class DocumentCreate(DocumentBase):
)
class ObservationInput(BaseModel):
"""Validated observation input from LLM tool calls."""
content: str = Field(min_length=1)
level: DocumentLevel = "explicit"
source_ids: list[str] | None = None
premises: list[str] | None = None
sources: list[str] | None = None
pattern_type: (
Literal["preference", "behavior", "personality", "tendency", "correlation"]
| None
) = None
confidence: Literal["high", "medium", "low"] | None = None
@model_validator(mode="after")
def validate_level_fields(self) -> Self:
"""Validate that level-specific fields are present when required."""
if self.level == "deductive" and not self.source_ids:
raise ValueError(
"deductive observations require 'source_ids' field with document IDs of premises"
)
if self.level == "inductive" and not self.source_ids:
raise ValueError(
"inductive observations require 'source_ids' field with document IDs of sources"
)
if self.level == "contradiction" and (
not self.source_ids or len(self.source_ids) < 2
):
raise ValueError(
"contradiction observations require 'source_ids' field with at least 2 IDs of contradicting observations"
)
return self
class ConclusionGet(BaseModel):
"""Schema for listing conclusions with optional filters."""

321
src/utils/agent_tools.py
View File

@ -5,6 +5,7 @@ from dataclasses import dataclass
from datetime import datetime
from typing import Any, cast
from pydantic import ValidationError
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
@ -22,7 +23,7 @@ from src.telemetry.events import (
from src.utils import summarizer
from src.utils.formatting import format_new_turn_with_timestamp, utc_now_iso
from src.utils.representation import Representation
from src.utils.types import DocumentLevel, get_current_iteration
from src.utils.types import get_current_iteration
logger = logging.getLogger(__name__)
@ -58,6 +59,23 @@ async def get_observation_lock(
return _observation_locks[key]
@dataclass
class ObservationFailure:
"""Records a single observation that failed during creation."""
content_preview: str
error: str
@dataclass
class ObservationsCreatedResult:
"""Result of a batch create_observations call."""
created_count: int
created_levels: list[str]
failed: list[ObservationFailure]
def _truncate_tool_output(output: str, max_chars: int | None = None) -> str:
"""Truncate tool output to prevent token explosion."""
if max_chars is None:
@ -561,20 +579,20 @@ INDUCTION_SPECIALIST_TOOLS: list[dict[str, Any]] = [
async def create_observations(
db: AsyncSession,
observations: list[dict[str, Any]],
observations: list[schemas.ObservationInput],
observer: str,
observed: str,
session_name: str | None,
workspace_name: str,
message_ids: list[int],
message_created_at: str,
) -> None:
) -> ObservationsCreatedResult:
"""
Create multiple observations (documents) in the memory system in a single call.
Args:
db: Database session
observations: List of observations, each with 'content', 'level', and level-specific fields
observations: List of validated observation inputs
observer: The peer making the observation
observed: The peer being observed
session_name: Session identifier
@ -582,13 +600,12 @@ async def create_observations(
message_ids: List of message IDs these observations are based on
message_created_at: Timestamp of the message that triggered these observations
Level-specific fields:
- deductive: 'premises' (list of strings)
- inductive: 'sources' (list of strings), 'pattern_type', 'confidence'
Returns:
ObservationsCreatedResult with created count and any per-observation failures
"""
if not observations:
logger.warning("create_observations called with empty list")
return
return ObservationsCreatedResult(created_count=0, created_levels=[], failed=[])
# Get or create collection
await crud.get_or_create_collection(
@ -598,58 +615,71 @@ async def create_observations(
observed=observed,
)
# Generate embeddings and create document objects for all observations
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)
embeddings_by_index = dict(
zip(range(len(observations)), embeddings, strict=True)
)
except Exception as e:
logger.warning(
"Batch embedding failed for create_observations; falling back to per-observation embedding: %s",
e,
)
# Build document objects with pre-computed embeddings
documents: list[schemas.DocumentCreate] = []
for obs in observations:
content = obs.get("content", "")
level_str = obs.get("level", "explicit")
if not content:
logger.warning("Skipping observation with empty content")
continue
# Validate and cast level
level: DocumentLevel
if level_str == "inductive":
level = "inductive"
elif level_str == "deductive":
level = "deductive"
elif level_str == "contradiction":
level = "contradiction"
failed: list[ObservationFailure] = []
for i, obs in enumerate(observations):
embedding: list[float]
if embeddings_by_index is not None:
embedding = embeddings_by_index[i]
else:
level = "explicit"
# Generate embedding for the observation
embedding = await embedding_client.embed(content)
try:
embedding = await embedding_client.embed(obs.content)
except Exception as e:
logger.warning(
"Error embedding observation content for level '%s': %s",
obs.level,
e,
)
failed.append(
ObservationFailure(
content_preview=obs.content[:50],
error=f"Embedding failed: {e}",
)
)
continue
# Build metadata with level-specific fields
metadata = schemas.DocumentMetadata(
message_ids=message_ids,
message_created_at=message_created_at,
source_ids=obs.get("source_ids")
if level in ("deductive", "inductive", "contradiction")
source_ids=obs.source_ids
if obs.level in ("deductive", "inductive", "contradiction")
else None,
# Deductive-specific (human-readable premises)
premises=obs.get("premises") if level == "deductive" else None,
# Inductive/Contradiction-specific (human-readable sources)
sources=obs.get("sources")
if level in ("inductive", "contradiction")
premises=obs.premises if obs.level == "deductive" else None,
sources=obs.sources
if obs.level in ("inductive", "contradiction")
else None,
pattern_type=obs.get("pattern_type") if level == "inductive" else None,
confidence=obs.get("confidence", "medium")
if level == "inductive"
pattern_type=obs.pattern_type if obs.level == "inductive" else None,
confidence=(obs.confidence or "medium")
if obs.level == "inductive"
else None,
)
# Create document with tree linkage at top level
doc = schemas.DocumentCreate(
content=content,
content=obs.content,
session_name=session_name,
level=level,
level=obs.level,
metadata=metadata,
embedding=embedding,
source_ids=obs.get("source_ids")
if level in ("deductive", "inductive", "contradiction")
source_ids=obs.source_ids
if obs.level in ("deductive", "inductive", "contradiction")
else None,
)
documents.append(doc)
@ -665,9 +695,19 @@ async def create_observations(
deduplicate=True,
)
logger.info(
f"Created {len(documents)} observations in {workspace_name}/{observer}/{observed}"
"Created %d observations in %s/%s/%s",
len(documents),
workspace_name,
observer,
observed,
)
return ObservationsCreatedResult(
created_count=len(documents),
created_levels=[doc.level for doc in documents],
failed=failed,
)
async def get_recent_history(
db: AsyncSession,
@ -731,6 +771,7 @@ async def search_memory(
query: str,
limit: int,
levels: list[str] | None = None,
embedding: list[float] | None = None,
) -> Representation:
"""
Search for observations in memory using semantic similarity.
@ -744,6 +785,7 @@ async def search_memory(
limit: Maximum number of results
levels: Optional list of observation levels to filter by
(e.g., ["explicit"], ["deductive", "inductive", "contradiction"])
embedding: Optional pre-computed embedding to avoid redundant API calls
Returns:
Representation object containing relevant observations
@ -761,6 +803,7 @@ async def search_memory(
query=query,
top_k=limit,
filters=filters,
embedding=embedding,
)
return Representation.from_documents(documents)
@ -862,6 +905,20 @@ async def extract_preferences(
"things user wants or does not want",
]
# Batch embed all queries in a single API call.
# If batching fails, each search call will generate its own embedding.
query_embeddings_by_query: dict[str, list[float]] | None = None
try:
query_embeddings = await embedding_client.simple_batch_embed(semantic_queries)
query_embeddings_by_query = dict(
zip(semantic_queries, query_embeddings, strict=True)
)
except Exception as e:
logger.warning(
"Batch embedding failed for extract_preferences; falling back to per-query embedding in search_messages: %s",
e,
)
for query in semantic_queries:
try:
snippets = await crud.search_messages(
@ -871,6 +928,11 @@ async def extract_preferences(
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:
@ -880,7 +942,7 @@ async def extract_preferences(
seen_content.add(content_key)
messages.append(f"'{msg.content.strip()}'")
except Exception as e:
logger.warning(f"Error in semantic search for '{query}': {e}")
logger.warning("Error in semantic search for '%s': %s", query, e)
return {
"instructions": [], # Deprecated - LLM will categorize
@ -917,106 +979,77 @@ async def _handle_create_observations(
ctx: ToolContext, tool_input: dict[str, Any]
) -> str:
"""Handle create_observations tool."""
observations = tool_input.get("observations", [])
raw_observations = tool_input.get("observations", [])
if not observations:
if not raw_observations:
return "ERROR: observations list is empty"
valid_levels = ["explicit", "deductive", "inductive", "contradiction"]
valid_pattern_types = [
"preference",
"behavior",
"personality",
"tendency",
"correlation",
]
valid_confidence = ["high", "medium", "low"]
# Set context-specific default level before Pydantic validation
default_level = "explicit" if ctx.current_messages else "deductive"
for obs in raw_observations:
obs.setdefault("level", default_level)
# Determine message context based on whether we have current_messages
# Validate observations individually so valid ones are still processed
observations: list[schemas.ObservationInput] = []
validation_failures: list[ObservationFailure] = []
for obs in raw_observations:
try:
validated = schemas.ObservationInput.model_validate(obs)
except ValidationError as e:
validation_failures.append(
ObservationFailure(
content_preview=str(obs.get("content", ""))[:50],
error=f"Validation failed: {e}",
)
)
continue
# Deriver can only create explicit observations
if ctx.current_messages and validated.level != "explicit":
validation_failures.append(
ObservationFailure(
content_preview=validated.content[:50],
error=f"Deriver can only create 'explicit' observations, got '{validated.level}'",
)
)
continue
observations.append(validated)
if not observations:
failure_details = "; ".join(
f"'{f.content_preview}': {f.error}" for f in validation_failures
)
return f"ERROR: All observations failed validation: {failure_details}"
# Determine message context
if ctx.current_messages:
# Deriver agent: uses simplified schema, default all to explicit
for i, obs in enumerate(observations):
if "content" not in obs:
return f"ERROR: observation {i} missing 'content' field"
# Default to explicit - the simplified deriver schema doesn't include level
if "level" not in obs:
obs["level"] = "explicit"
# Enforce explicit-only for deriver
if obs["level"] != "explicit":
return f"ERROR: Deriver can only create 'explicit' observations, got '{obs['level']}' at index {i}"
message_ids = [msg.id for msg in ctx.current_messages]
message_created_at = str(ctx.current_messages[-1].created_at)
obs_session_name = ctx.session_name
else:
# Dreamer/Dialectic agent: allow deductive and inductive, no source messages
for i, obs in enumerate(observations):
if "content" not in obs:
return f"ERROR: observation {i} missing 'content' field"
# Default to deductive for backwards compatibility
if "level" not in obs:
obs["level"] = "deductive"
if obs["level"] not in valid_levels:
return f"ERROR: observation {i} has invalid level '{obs['level']}'"
# Validate deductive-specific fields (tree linkage required)
if obs["level"] == "deductive":
if not obs.get("source_ids"):
return f"ERROR: deductive observation {i} requires 'source_ids' field with document IDs of premises"
# Validate source_ids are strings
for sid in obs.get("source_ids", []):
if not isinstance(sid, str):
return f"ERROR: observation {i} source_ids must be strings, got {type(sid)}"
# Validate inductive-specific fields (tree linkage required)
if obs["level"] == "inductive":
if not obs.get("source_ids"):
return f"ERROR: inductive observation {i} requires 'source_ids' field with document IDs of sources"
# Validate source_ids are strings
for sid in obs.get("source_ids", []):
if not isinstance(sid, str):
return f"ERROR: observation {i} source_ids must be strings, got {type(sid)}"
if (
obs.get("pattern_type")
and obs["pattern_type"] not in valid_pattern_types
):
return f"ERROR: observation {i} has invalid pattern_type '{obs['pattern_type']}'"
if obs.get("confidence") and obs["confidence"] not in valid_confidence:
return f"ERROR: observation {i} has invalid confidence '{obs['confidence']}'"
# Validate contradiction-specific fields (need source_ids for the two contradicting obs)
if obs["level"] == "contradiction":
if not obs.get("source_ids"):
return f"ERROR: contradiction observation {i} requires 'source_ids' field with IDs of contradicting observations"
if len(obs.get("source_ids", [])) < 2:
return f"ERROR: contradiction observation {i} requires at least 2 source_ids (the contradicting observations)"
for sid in obs.get("source_ids", []):
if not isinstance(sid, str):
return f"ERROR: observation {i} source_ids must be strings, got {type(sid)}"
message_ids = []
message_created_at = utc_now_iso()
obs_session_name = ctx.session_name
# Use lock to serialize database writes (prevents concurrent commit issues)
async with ctx.db_lock:
await create_observations(
result = await create_observations(
ctx.db,
observations=observations,
observer=ctx.observer,
observed=ctx.observed,
session_name=obs_session_name,
session_name=ctx.session_name,
workspace_name=ctx.workspace_name,
message_ids=message_ids,
message_created_at=message_created_at,
)
explicit_count = sum(1 for o in observations if o.get("level") == "explicit")
deductive_count = sum(1 for o in observations if o.get("level") == "deductive")
inductive_count = sum(1 for o in observations if o.get("level") == "inductive")
contradiction_count = sum(
1 for o in observations if o.get("level") == "contradiction"
)
# Merge validation and embedding failures
all_failures = validation_failures + result.failed
# Count levels from actually-created observations
levels = result.created_levels
explicit_count = levels.count("explicit")
deductive_count = levels.count("deductive")
inductive_count = levels.count("inductive")
contradiction_count = levels.count("contradiction")
# Emit telemetry event if context is available
if ctx.run_id and ctx.agent_type and ctx.parent_category:
@ -1029,12 +1062,24 @@ async def _handle_create_observations(
workspace_name=ctx.workspace_name,
observer=ctx.observer,
observed=ctx.observed,
conclusion_count=len(observations),
levels=[o.get("level", "explicit") for o in observations],
conclusion_count=result.created_count,
levels=levels,
)
)
return f"Created {len(observations)} observations for {ctx.observed} by {ctx.observer} ({explicit_count} explicit, {deductive_count} deductive, {inductive_count} inductive, {contradiction_count} contradiction)"
response = (
f"Created {result.created_count} observations for {ctx.observed} by {ctx.observer} "
f"({explicit_count} explicit, {deductive_count} deductive, "
f"{inductive_count} inductive, {contradiction_count} contradiction)"
)
if all_failures:
failure_details = "; ".join(
f"'{f.content_preview}': {f.error}" for f in all_failures
)
response += f"\nFailed {len(all_failures)}: {failure_details}"
return response
async def _handle_update_peer_card(ctx: ToolContext, tool_input: dict[str, Any]) -> str:
@ -1042,7 +1087,8 @@ async def _handle_update_peer_card(ctx: ToolContext, tool_input: dict[str, Any])
# Check if peer card creation is disabled via configuration
if ctx.configuration is not None and not ctx.configuration.peer_card.create:
logger.info(
f"Peer card creation disabled for {ctx.workspace_name}, skipping update"
"Peer card creation disabled for %s, skipping update",
ctx.workspace_name,
)
return (
"Peer card creation is disabled for this workspace/session configuration."
@ -1118,13 +1164,20 @@ async def _handle_get_recent_history(
async def _handle_search_memory(ctx: ToolContext, tool_input: dict[str, Any]) -> str:
"""Handle search_memory tool."""
top_k = min(tool_input.get("top_k", 20), 40)
query = tool_input["query"]
try:
query_embedding = await embedding_client.embed(query)
except ValueError:
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,
workspace_name=ctx.workspace_name,
observer=ctx.observer,
observed=ctx.observed,
query=tool_input["query"],
query=query,
top_k=top_k,
embedding=query_embedding,
)
mem = Representation.from_documents(documents)
total_count = mem.len()
@ -1134,7 +1187,6 @@ async def _handle_search_memory(ctx: ToolContext, tool_input: dict[str, Any]) ->
# this stage, and be efficient with tool calls, and make sure the model
# doesn't short-circuit and think there's nothing here, we
# automatically search the message history for relevant information.
query = tool_input["query"]
if ctx.agent_type == "dialectic":
limit = min(tool_input.get("top_k", 20), 20)
snippets = await crud.search_messages(
@ -1144,6 +1196,7 @@ async def _handle_search_memory(ctx: ToolContext, tool_input: dict[str, Any]) ->
query=query,
limit=limit,
context_window=0,
embedding=query_embedding,
)
if snippets:
message_output = _format_message_snippets(
@ -1158,7 +1211,7 @@ async def _handle_search_memory(ctx: ToolContext, tool_input: dict[str, Any]) ->
)
return f"No observations found for query '{query}'"
mem_str = mem.str_with_ids() if ctx.include_observation_ids else str(mem)
return f"Found {total_count} observations for query '{tool_input['query']}':\n\n{mem_str}"
return f"Found {total_count} observations for query '{query}':\n\n{mem_str}"
async def _handle_get_observation_context(
@ -1461,7 +1514,7 @@ async def _handle_delete_observations(
)
deleted_count += 1
except Exception as e:
logger.warning(f"Failed to delete observation {obs_id}: {e}")
logger.warning("Failed to delete observation %s: %s", obs_id, e)
# Emit telemetry event if context is available
if deleted_count > 0 and ctx.run_id and ctx.agent_type and ctx.parent_category:
@ -1723,13 +1776,13 @@ async def create_tool_executor(
Returns:
String result describing what was done
"""
logger.info(f"[tool call] {tool_name} {tool_input}")
logger.info("[tool call] %s %s", tool_name, tool_input)
try:
handler = _TOOL_HANDLERS.get(tool_name)
if handler:
result = await handler(ctx, tool_input)
logger.info(f"[tool result] {tool_name} {result}")
logger.info("[tool result] %s %s", tool_name, result)
return result
return f"Unknown tool: {tool_name}"

269
tests/utils/test_agent_tools.py
View File

@ -10,9 +10,10 @@ from nanoid import generate as generate_nanoid
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, models
from src import crud, models, schemas
from src.config import settings
from src.utils.agent_tools import (
ObservationsCreatedResult,
ToolContext,
_handle_create_observations, # pyright: ignore[reportPrivateUsage]
_handle_delete_observations, # pyright: ignore[reportPrivateUsage]
@ -28,7 +29,9 @@ from src.utils.agent_tools import (
_handle_search_memory, # pyright: ignore[reportPrivateUsage]
_handle_search_messages, # pyright: ignore[reportPrivateUsage]
_handle_update_peer_card, # pyright: ignore[reportPrivateUsage]
create_observations,
create_tool_executor,
extract_preferences,
)
# =============================================================================
@ -242,6 +245,133 @@ class TestCreateObservations:
assert "ERROR" in result
assert "empty" in result.lower()
async def test_batch_embedding_failure_falls_back_to_individual_embeds(
self,
db_session: AsyncSession,
tool_test_data: Any,
monkeypatch: pytest.MonkeyPatch,
):
"""If batch embedding fails but individual embeds succeed, all observations are created."""
workspace, peer1, peer2, session, _, _ = tool_test_data
async def fail_batch_embed(_texts: list[str]) -> list[list[float]]:
raise RuntimeError("embedding provider timeout")
async def succeed_single_embed(_content: str) -> list[float]:
return [0.1, 0.2, 0.3]
created_documents: list[Any] = []
async def fake_create_documents(
_db: AsyncSession,
documents: list[Any],
workspace_name: str,
*,
observer: str,
observed: str,
deduplicate: bool = False,
) -> int:
_ = (workspace_name, observer, observed, deduplicate)
created_documents.extend(documents)
return len(documents)
monkeypatch.setattr(
"src.utils.agent_tools.embedding_client.simple_batch_embed",
fail_batch_embed,
)
monkeypatch.setattr(
"src.utils.agent_tools.embedding_client.embed",
succeed_single_embed,
)
monkeypatch.setattr(
"src.utils.agent_tools.crud.create_documents", fake_create_documents
)
result = await create_observations(
db_session,
observations=[
schemas.ObservationInput(content="First obs", level="explicit"),
schemas.ObservationInput(content="Second obs", level="explicit"),
],
observer=peer1.name,
observed=peer2.name,
session_name=session.name,
workspace_name=workspace.name,
message_ids=[],
message_created_at=str(datetime.now(timezone.utc)),
)
assert isinstance(result, ObservationsCreatedResult)
assert result.created_count == 2
assert len(result.failed) == 0
assert len(created_documents) == 2
async def test_batch_embedding_failure_individual_embed_partial_failure(
self,
db_session: AsyncSession,
tool_test_data: Any,
monkeypatch: pytest.MonkeyPatch,
):
"""If batch embedding fails and some individual embeds also fail, only successful ones are created."""
workspace, peer1, peer2, session, _, _ = tool_test_data
async def fail_batch_embed(_texts: list[str]) -> list[list[float]]:
raise RuntimeError("embedding provider timeout")
async def embed_per_observation(content: str) -> list[float]:
if content == "Fails embed":
raise RuntimeError("single-item embed failure")
return [0.1, 0.2, 0.3]
created_documents: list[Any] = []
async def fake_create_documents(
_db: AsyncSession,
documents: list[Any],
workspace_name: str,
*,
observer: str,
observed: str,
deduplicate: bool = False,
) -> int:
_ = (workspace_name, observer, observed, deduplicate)
created_documents.extend(documents)
return len(documents)
monkeypatch.setattr(
"src.utils.agent_tools.embedding_client.simple_batch_embed",
fail_batch_embed,
)
monkeypatch.setattr(
"src.utils.agent_tools.embedding_client.embed",
embed_per_observation,
)
monkeypatch.setattr(
"src.utils.agent_tools.crud.create_documents", fake_create_documents
)
result = await create_observations(
db_session,
observations=[
schemas.ObservationInput(content="Embeds fine", level="explicit"),
schemas.ObservationInput(content="Fails embed", level="explicit"),
],
observer=peer1.name,
observed=peer2.name,
session_name=session.name,
workspace_name=workspace.name,
message_ids=[],
message_created_at=str(datetime.now(timezone.utc)),
)
assert isinstance(result, ObservationsCreatedResult)
assert result.created_count == 1
assert len(result.failed) == 1
assert result.failed[0].content_preview == "Fails embed"
assert "Embedding failed" in result.failed[0].error
assert len(created_documents) == 1
assert created_documents[0].content == "Embeds fine"
@pytest.mark.asyncio
class TestDeleteObservations:
@ -366,6 +496,77 @@ class TestSearchMemory:
assert "No observations found" in result
async def test_reuses_single_embedding_for_dialectic_fallback(
self,
make_tool_context: Callable[..., ToolContext],
monkeypatch: pytest.MonkeyPatch,
):
"""Uses one embedding for query_documents and search_messages fallback."""
ctx = make_tool_context()
ctx.agent_type = "dialectic"
embed_calls: list[str] = []
query_embeddings: list[list[float] | None] = []
fallback_embeddings: list[list[float] | None] = []
async def fake_embed(query: str) -> list[float]:
embed_calls.append(query)
return [0.1, 0.2, 0.3]
async def fake_query_documents(
db: AsyncSession,
workspace_name: str,
query: str,
*,
observer: str,
observed: str,
top_k: int = 5,
embedding: list[float] | None = None,
**_kwargs: Any,
) -> list[models.Document]:
_ = (db, workspace_name, query, observer, observed, top_k)
query_embeddings.append(embedding)
return []
async def fake_search_messages(
db: AsyncSession,
workspace_name: str,
session_name: str | None,
query: str,
limit: int = 10,
context_window: int = 2,
embedding: list[float] | None = None,
) -> list[tuple[list[models.Message], list[models.Message]]]:
_ = (db, workspace_name, session_name, query, limit, context_window)
fallback_embeddings.append(embedding)
msg = models.Message(
workspace_name=ctx.workspace_name,
session_name=ctx.session_name,
peer_name=ctx.observed,
content="Relevant fallback message",
seq_in_session=1,
token_count=5,
created_at=datetime.now(timezone.utc),
)
return [([msg], [msg])]
monkeypatch.setattr("src.utils.agent_tools.embedding_client.embed", fake_embed)
monkeypatch.setattr(
"src.utils.agent_tools.crud.query_documents", fake_query_documents
)
monkeypatch.setattr(
"src.utils.agent_tools.crud.search_messages",
fake_search_messages,
)
result = await _handle_search_memory(ctx, {"query": "coffee preferences"})
assert "No observations yet. Message search results:" in result
assert embed_calls == ["coffee preferences"]
assert len(query_embeddings) == 1
assert len(fallback_embeddings) == 1
assert query_embeddings[0] == fallback_embeddings[0]
@pytest.mark.asyncio
class TestSearchMessages:
@ -684,6 +885,72 @@ class TestExtractPreferences:
# Should return some result about preferences
assert isinstance(result, str)
async def test_falls_back_to_per_query_embedding_when_batch_fails(
self,
db_session: AsyncSession,
tool_test_data: Any,
monkeypatch: pytest.MonkeyPatch,
):
"""Batch embedding failure should not abort preference extraction."""
workspace, _, observed_peer, session, _, _ = tool_test_data
async def fail_batch_embed(_texts: list[str]) -> list[list[float]]:
raise RuntimeError("embedding provider timeout")
async def unexpected_embed_call(_query: str) -> list[float]:
raise AssertionError(
"extract_preferences should not call embedding_client.embed "
+ "when batch embedding fails"
)
embedding_args: list[list[float] | None] = []
async def fake_search_messages(
_db: AsyncSession,
workspace_name: str,
session_name: str | None,
query: str,
limit: int,
context_window: int,
embedding: list[float] | None,
) -> list[tuple[list[models.Message], list[models.Message]]]:
_ = (limit, context_window)
embedding_args.append(embedding)
msg = models.Message(
workspace_name=workspace_name,
session_name=session_name,
peer_name=observed_peer.name,
content=f"Relevant from {query}",
seq_in_session=1,
token_count=5,
created_at=datetime.now(timezone.utc),
)
return [([msg], [])]
monkeypatch.setattr(
"src.utils.agent_tools.embedding_client.simple_batch_embed",
fail_batch_embed,
)
monkeypatch.setattr(
"src.utils.agent_tools.embedding_client.embed",
unexpected_embed_call,
)
monkeypatch.setattr(
"src.utils.agent_tools.crud.search_messages", fake_search_messages
)
result = await extract_preferences(
db_session,
workspace_name=workspace.name,
session_name=session.name,
observed=observed_peer.name,
)
# We still get partial results despite one per-query failure.
assert result["messages"]
assert len(embedding_args) == 5
assert all(embedding is None for embedding in embedding_args)
@pytest.mark.asyncio
class TestFinishConsolidation: