From 33ef0f8eab8f1e11650279f6cf6b2d04edff6774 Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Fri, 13 Feb 2026 00:53:46 -0500 Subject: [PATCH] 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> --- src/cache/client.py | 43 ++++- src/crud/__init__.py | 7 +- src/crud/collection.py | 99 +++++++++- src/crud/message.py | 8 +- src/crud/peer.py | 6 +- src/crud/peer_card.py | 4 +- src/crud/session.py | 15 +- src/crud/workspace.py | 11 +- src/deriver/enqueue.py | 35 ++-- src/dialectic/core.py | 6 + src/routers/sessions.py | 4 + src/schemas.py | 36 +++- src/utils/agent_tools.py | 321 +++++++++++++++++++------------- tests/utils/test_agent_tools.py | 269 +++++++++++++++++++++++++- 14 files changed, 683 insertions(+), 181 deletions(-) diff --git a/src/cache/client.py b/src/cache/client.py index d4005fb1..319fab6e 100644 --- a/src/cache/client.py +++ b/src/cache/client.py @@ -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", ] diff --git a/src/crud/__init__.py b/src/crud/__init__.py index 9d5a01a4..fdff810d 100644 --- a/src/crud/__init__.py +++ b/src/crud/__init__.py @@ -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", diff --git a/src/crud/collection.py b/src/crud/collection.py index c1759b9c..fadbd145 100644 --- a/src/crud/collection.py +++ b/src/crud/collection.py @@ -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)) diff --git a/src/crud/message.py b/src/crud/message.py index 18c1c053..b73dc414 100644 --- a/src/crud/message.py +++ b/src/crud/message.py @@ -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 = ( diff --git a/src/crud/peer.py b/src/crud/peer.py index 8d3a761e..8276a93e 100644 --- a/src/crud/peer.py +++ b/src/crud/peer.py @@ -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 diff --git a/src/crud/peer_card.py b/src/crud/peer_card.py index 37c2749c..f386400b 100644 --- a/src/crud/peer_card.py +++ b/src/crud/peer_card.py @@ -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: diff --git a/src/crud/session.py b/src/crud/session.py index 7008725c..9657413e 100644 --- a/src/crud/session.py +++ b/src/crud/session.py @@ -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) diff --git a/src/crud/workspace.py b/src/crud/workspace.py index 91a79705..f759543e 100644 --- a/src/crud/workspace.py +++ b/src/crud/workspace.py @@ -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 diff --git a/src/deriver/enqueue.py b/src/deriver/enqueue.py index fdfbce84..71b00a2b 100644 --- a/src/deriver/enqueue.py +++ b/src/deriver/enqueue.py @@ -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)", diff --git a/src/dialectic/core.py b/src/dialectic/core.py index 8a3943ac..6dd215b5 100644 --- a/src/dialectic/core.py +++ b/src/dialectic/core.py @@ -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(): diff --git a/src/routers/sessions.py b/src/routers/sessions.py index 7cad530e..a7be8943 100644 --- a/src/routers/sessions.py +++ b/src/routers/sessions.py @@ -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: diff --git a/src/schemas.py b/src/schemas.py index a476ee92..07c24576 100644 --- a/src/schemas.py +++ b/src/schemas.py @@ -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.""" diff --git a/src/utils/agent_tools.py b/src/utils/agent_tools.py index 2e8cf8e5..0d89fa8a 100644 --- a/src/utils/agent_tools.py +++ b/src/utils/agent_tools.py @@ -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}" diff --git a/tests/utils/test_agent_tools.py b/tests/utils/test_agent_tools.py index 7f135bcb..4fa7e160 100644 --- a/tests/utils/test_agent_tools.py +++ b/tests/utils/test_agent_tools.py @@ -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: