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

feat: agentic dreamer and agentic dialectic (#309) 

* feat: add better params to working representation fetch in SDKs, return messages when added

* fix: working representation routes now accepting all parameters properly, with tests

* feat: add metadata/config fields to SDK objects where viable

* fix: tests

* feat: refactor SDKs to use representation config; [TEMP STAINLESS BUILD] update API

* feat: add representation object to sdks

* fix: use stainless sdk on branch

* fix: update TypeScript SDK tsconfig to use node16 module resolution

* fix: add isolatedModules = true to tsconfig

* fix: lol

* chore: coderabbit review

* feat: make delete session real

* feat: add observations routes with delete endpoints for documents. make session deletion real.

* chore: type cleanup

* fix: tests

* chore: coderabbit review

* fix: namespace by workspace

* feat: add ability to customize messages_per_summary at both workspace and session level

* chore: tests for summary config

* chore: coderabbit cleanup

* feat: make session and workspace config totally customizeable

* feat: add search by peer knowledge (#250)

* feat: search by peer perspective

* fix: enforce workspace in filters, make messages distinct in join

* fix: batch and merge migration steps

* fix: add refresh, add config to workspace, add refresh function, make fields readonly

* fix: search distinct

* fix: merge migrations

* fix: merge migrations

* fix: batch deletions, improve comments, limit consolidate dream to 100 docs at a time, auth on observations routes

* chore: review

* chore: coderabbit

* chore: review

* chore: broken comment

* feat: add set peer card route to API

* feat: create advanced configuration parameters with message>session>workspace hierarchy

* [wip] build unified testing harness

* chore: lint

* fix: cache invalidation, naming things, etc

* feat: longmem tests

* chore: peer config refactor

* feat: consolidate dream working, refactor representation

* fix: Various CR Comment Fixes

* feat: Allow configurable Redis port for harness instances and update cleanup methods to be asynchronous.

* feat: agentic ingestion task!!!

* feat: agentic deriver

* feat: dialectic agent and dreamer agent

* chore: browbeat tests into passing

* fix: nits

* chore: remove old code, update config files

* fix: simplify deriver

* feat: dialectic agent prompt updates, re-introduce non_agent deriver, eval tweaks

* feat: fast deriver, dreamer, then dialectic

* fix: tweaks across the board

* feat: add baseline tests

* feat: truncation in tools and client, tweaks for evals

* feat: add locomo, fix longmem judge!!!

* fix: locomo f1 is trash, use llm judge

* feat: trace creation

* feat: add first draft of obex benchmark, fix embedding model, fix locomo methodology

* fix: locomo session-optimized, better logging of cache usage and better cache usage

* chore: use openrouter for baselines

* fix: add test for merge migration

* chore: opus-powered cleanup

* fix: add config for vllm, better client

* chore: clean up clients.py a bit

* chore: move magic numbers to config, add tests for agent tools

* fix: wrong mock in dialectic tests, make ToolContext a dataclass

* feat: tweak prompts, make deriver explicit-only

* feat: more prompt & tool tweaks

* chore: more tweaks

* feat: dream with subagents

* fix: make dream trigger override scheduled, play around with dream agents

* chore: cleanup deriver

* chore: cleanup dialectic

* chore: cleanup orchestrator

* chore: comment out dream stuff, WIPing

* fix: inc temp on retry, typechecking

* feat: tweak dreaming

* feat: contradiction obs

* Add dream trees

* chore: preserve reasoning_details from openrouter in client

* fix: get_observation_context correct params

* fix: use correct message id in tool

* chore: cleanup longmem runner

* chore: clean up tests, remove dream tests for now as rearchitecting around trees

* chore: update stainless deps

* Update threholding mechanism

* chore: pre-commit hooks whitespace

* chore: clean up types

* feat: add explicit bench

* fix: address additional basepyright issues

* fix: adding logging as a fixture on honcho_llm_call and supporting dialectic loging. (#305)

* fix: lock on db for tool calls

* chore: clean up experimental derivers

* chore: coderabbit review cleanup

* feat: add streaming support to agentic dialectic

* feat: prometheus token tracking for deriver and dialectic

* fix: self-loops for isolated nodes

* chore: PascalCase for prometheus parameter typing

* feat: add reasoning levels to dialectic agent

* chore: delete old file, add new fake env vars in unittest.yml

* fix: all fields needed for dialectic reasoning level configs

* feat: track dreaming usage in prometheus

* chore: Create backwards compatabile conclusion and queue endpoints

* fix: remove redundant try-catch, add trace label, move .limit to end of statement

* fix: remove vignettes (for now), review fixes, remove merge migration, config cleanup

* chore: code review / cleanup

* chore: merge fixes

* chore: clean up, remove reasoning_focus, reintroduce peer cards in dreamers

* chore: code rabbit nitpicks

* fix: add unique index for pending dreams in queue

* fix: revert removal of surprisal in dreamer config

---------

Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
Co-authored-by: 3un01a <3un01a@plasticlabs.ai>
Co-authored-by: 3un01a <3un01a.labs@gmail.com>
Co-authored-by: ajspig <46900795+ajspig@users.noreply.github.com>
This commit is contained in:
doria 2026-01-12 15:12:17 -05:00 committed by GitHub
parent 49f6d75991
commit 578ef2c665
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
126 changed files with 17451 additions and 3750 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

106
.env.template
View File

@ -21,6 +21,9 @@ LOG_LEVEL=INFO
# LANGFUSE_HOST=
# LANGFUSE_PUBLIC_KEY=
# COLLECT_METRICS_LOCAL=false
# LOCAL_METRICS_FILE=metrics.jsonl
# NAMESPACE="honcho"
# =============================================================================
@ -71,27 +74,35 @@ LLM_ANTHROPIC_API_KEY=your-anthropic-api-key-here
# LLM_OPENAI_COMPATIBLE_BASE_URL=
# LLM_OPENAI_COMPATIBLE_API_KEY=
# Separate vLLM endpoint (for local models)
# LLM_VLLM_API_KEY=
# LLM_VLLM_BASE_URL=
# =============================================================================
# LLM Configuration
# =============================================================================
# Global LLM settings
# LLM_DEFAULT_MAX_TOKENS=2500
# LLM_EMBEDDING_PROVIDER=openai
# LLM_MAX_TOOL_OUTPUT_CHARS=30000 # Max chars for tool output (~7500 tokens)
# LLM_MAX_MESSAGE_CONTENT_CHARS=2000 # Max chars per message in tool results
# =============================================================================
# Deriver (Background Worker) Settings
# =============================================================================
# DERIVER_ENABLED=true
# DERIVER_WORKERS=1
# DERIVER_POLLING_SLEEP_INTERVAL_SECONDS=1.0
# DERIVER_STALE_SESSION_TIMEOUT_MINUTES=5
# DERIVER_QUEUE_ERROR_RETENTION_SECONDS=2592000 # 30 days
# DERIVER_PROVIDER=google
# DERIVER_MODEL=gemini-2.0-flash-lite
# DERIVER_MODEL=gemini-2.5-flash-lite
# DERIVER_DEDUPLICATE=true
# DERIVER_MAX_OUTPUT_TOKENS=2500
# only applied when using Anthropic as provider
# DERIVER_MAX_OUTPUT_TOKENS=4096
# DERIVER_THINKING_BUDGET_TOKENS=1024
# DERIVER_LOG_OBSERVATIONS=false
# DERIVER_WORKING_REPRESENTATION_MAX_OBSERVATIONS=100
# DERIVER_REPRESENTATION_BATCH_MAX_TOKENS=4096
# DERIVER_MAX_INPUT_TOKENS=23000
# DERIVER_BACKUP_PROVIDER=
# DERIVER_BACKUP_MODEL=
@ -99,27 +110,51 @@ LLM_ANTHROPIC_API_KEY=your-anthropic-api-key-here
# Peer Card Configuration
# =============================================================================
# PEER_CARD_ENABLED=true
# PEER_CARD_PROVIDER=openai
# PEER_CARD_MODEL=gpt-5-nano-2025-08-07
# PEER_CARD_MAX_OUTPUT_TOKENS=4000
# PEER_CARD_BACKUP_PROVIDER=
# PEER_CARD_BACKUP_MODEL=
# =============================================================================
# Dialectic Settings
# =============================================================================
# DIALECTIC_PROVIDER=anthropic
# DIALECTIC_MODEL=claude-sonnet-4-20250514
# DIALECTIC_PERFORM_QUERY_GENERATION=false
# DIALECTIC_QUERY_GENERATION_PROVIDER=groq
# DIALECTIC_QUERY_GENERATION_MODEL=llama-3.1-8b-instant
# DIALECTIC_MAX_OUTPUT_TOKENS=2500
# DIALECTIC_SEMANTIC_SEARCH_TOP_K=10
# DIALECTIC_SEMANTIC_SEARCH_MAX_DISTANCE=0.85
# DIALECTIC_THINKING_BUDGET_TOKENS=1024
# DIALECTIC_CONTEXT_WINDOW_SIZE=100000
# DIALECTIC_BACKUP_PROVIDER=
# DIALECTIC_BACKUP_MODEL=
# Global dialectic settings
# DIALECTIC_MAX_OUTPUT_TOKENS=8192
# DIALECTIC_MAX_INPUT_TOKENS=100000
# DIALECTIC_HISTORY_TOKEN_LIMIT=8192
# DIALECTIC_SESSION_HISTORY_MAX_TOKENS=16384
# Per-level settings (reasoning_level parameter in API)
# Each level can have its own provider, model, thinking budget, and tool iterations
# Minimal level
# DIALECTIC_LEVELS__minimal__PROVIDER=google
# DIALECTIC_LEVELS__minimal__MODEL=gemini-2.5-flash-lite
# DIALECTIC_LEVELS__minimal__THINKING_BUDGET_TOKENS=0
# DIALECTIC_LEVELS__minimal__MAX_TOOL_ITERATIONS=2
# Low level
# DIALECTIC_LEVELS__low__PROVIDER=google
# DIALECTIC_LEVELS__low__MODEL=gemini-3-flash
# DIALECTIC_LEVELS__low__THINKING_BUDGET_TOKENS=0
# DIALECTIC_LEVELS__low__MAX_TOOL_ITERATIONS=5
# Medium level
# DIALECTIC_LEVELS__medium__PROVIDER=anthropic
# DIALECTIC_LEVELS__medium__MODEL=claude-haiku-4-5
# DIALECTIC_LEVELS__medium__THINKING_BUDGET_TOKENS=512
# DIALECTIC_LEVELS__medium__MAX_TOOL_ITERATIONS=4
# High level
# DIALECTIC_LEVELS__high__PROVIDER=anthropic
# DIALECTIC_LEVELS__high__MODEL=claude-opus-4-5
# DIALECTIC_LEVELS__high__THINKING_BUDGET_TOKENS=0
# DIALECTIC_LEVELS__high__MAX_TOOL_ITERATIONS=4
# Extra-high level
# DIALECTIC_LEVELS__extra-high__PROVIDER=anthropic
# DIALECTIC_LEVELS__extra-high__MODEL=claude-opus-4-5
# DIALECTIC_LEVELS__extra-high__THINKING_BUDGET_TOKENS=512
# DIALECTIC_LEVELS__extra-high__MAX_TOOL_ITERATIONS=10
# Optional backup per level (must set both or neither):
# DIALECTIC_LEVELS__extra-high__BACKUP_PROVIDER=google
# DIALECTIC_LEVELS__extra-high__BACKUP_MODEL=gemini-2.5-pro
# =============================================================================
# Summary Settings
@ -130,7 +165,7 @@ LLM_ANTHROPIC_API_KEY=your-anthropic-api-key-here
# SUMMARY_PROVIDER=google
# SUMMARY_MODEL=gemini-2.5-flash
# SUMMARY_MAX_TOKENS_SHORT=1000
# SUMMARY_MAX_TOKENS_LONG=2000
# SUMMARY_MAX_TOKENS_LONG=4000
# SUMMARY_THINKING_BUDGET_TOKENS=512
# SUMMARY_BACKUP_PROVIDER=
# SUMMARY_BACKUP_MODEL=
@ -142,13 +177,26 @@ LLM_ANTHROPIC_API_KEY=your-anthropic-api-key-here
# DREAM_DOCUMENT_THRESHOLD=50
# DREAM_IDLE_TIMEOUT_MINUTES=60
# DREAM_MIN_HOURS_BETWEEN_DREAMS=8
# DREAM_ENABLED_TYPES=["consolidate"]
# DREAM_PROVIDER=openai
# DREAM_MODEL=gpt-4o-mini-2024-07-18
# DREAM_MAX_TOKENS=2000
# DREAM_ENABLED_TYPES=["omni"]
# DREAM_PROVIDER=anthropic
# DREAM_MODEL=claude-haiku-4-5
# DREAM_MAX_OUTPUT_TOKENS=4000
# DREAM_THINKING_BUDGET_TOKENS=2048
# DREAM_MAX_TOOL_ITERATIONS=8
# DREAM_HISTORY_TOKEN_LIMIT=8192
# DREAM_BACKUP_PROVIDER=
# DREAM_BACKUP_MODEL=
# Dream Surprisal Settings (Tree-based observation sampling for targeted reasoning)
# DREAM_SURPRISAL__ENABLED=true
# DREAM_SURPRISAL__TREE_TYPE=kdtree # Options: kdtree, balltree, rptree, covertree, lsh, graph, prototype
# DREAM_SURPRISAL__TREE_K=5 # Number of neighbors for kNN-based trees
# DREAM_SURPRISAL__SAMPLING_STRATEGY=recent # Options: recent, random, all
# DREAM_SURPRISAL__SAMPLE_SIZE=200 # Number of observations to sample for tree building
# DREAM_SURPRISAL__TOP_PERCENT_SURPRISAL=0.10 # Top percentage of observations (0.10 = top 10%)
# DREAM_SURPRISAL__MIN_HIGH_SURPRISAL_FOR_REPLACE=10 # Hybrid mode: min observations to replace standard questions
# DREAM_SURPRISAL__INCLUDE_LEVELS=["explicit","deductive"] # Observation levels to include
# =============================================================================
# Webhook Settings
# =============================================================================
@ -170,13 +218,13 @@ LLM_ANTHROPIC_API_KEY=your-anthropic-api-key-here
# Metrics (Optional)
# =============================================================================
# METRICS_ENABLED=false
# METRICS_NAMESPACE=honcho
# METRICS_NAMESPACE=honcho # Inherits from NAMESPACE if not set
# =============================================================================
# Cache
# =============================================================================
# CACHE_ENABLED=false
# CACHE_URL="redis://localhost:6379/0"
# CACHE_NAMESPACE="honcho"
# CACHE_URL="redis://localhost:6379/0?suppress=false"
# CACHE_NAMESPACE="honcho" # Inherits from NAMESPACE if not set
# CACHE_DEFAULT_TTL_SECONDS=300
# CACHE_DEFAULT_LOCK_TTL_SECONDS=5

22
.github/workflows/unittest.yml vendored
View File

@ -105,8 +105,26 @@ jobs:
LLM_OPENAI_COMPATIBLE_BASE_URL: http://localhost:8000
DERIVER_PROVIDER: openai
DERIVER_MODEL: test
DIALECTIC_PROVIDER: openai
DIALECTIC_MODEL: test
DIALECTIC_LEVELS__minimal__PROVIDER: openai
DIALECTIC_LEVELS__minimal__MODEL: test
DIALECTIC_LEVELS__minimal__THINKING_BUDGET_TOKENS: 0
DIALECTIC_LEVELS__minimal__MAX_TOOL_ITERATIONS: 2
DIALECTIC_LEVELS__low__PROVIDER: openai
DIALECTIC_LEVELS__low__MODEL: test
DIALECTIC_LEVELS__low__THINKING_BUDGET_TOKENS: 0
DIALECTIC_LEVELS__low__MAX_TOOL_ITERATIONS: 5
DIALECTIC_LEVELS__medium__PROVIDER: openai
DIALECTIC_LEVELS__medium__MODEL: test
DIALECTIC_LEVELS__medium__THINKING_BUDGET_TOKENS: 0
DIALECTIC_LEVELS__medium__MAX_TOOL_ITERATIONS: 4
DIALECTIC_LEVELS__high__PROVIDER: openai
DIALECTIC_LEVELS__high__MODEL: test
DIALECTIC_LEVELS__high__THINKING_BUDGET_TOKENS: 0
DIALECTIC_LEVELS__high__MAX_TOOL_ITERATIONS: 4
DIALECTIC_LEVELS__extra-high__PROVIDER: openai
DIALECTIC_LEVELS__extra-high__MODEL: test
DIALECTIC_LEVELS__extra-high__THINKING_BUDGET_TOKENS: 0
DIALECTIC_LEVELS__extra-high__MAX_TOOL_ITERATIONS: 10
DIALECTIC_QUERY_GENERATION_PROVIDER: openai
DIALECTIC_QUERY_GENERATION_MODEL: test
SUMMARY_PROVIDER: openai

2
.pre-commit-config.yaml
View File

@ -99,7 +99,7 @@ repos:
stages: [pre-push]
pass_filenames: false
# TypeScript build/test with bun
# # TypeScript build/test with bun
- id: typescript-check
name: TypeScript build and test
entry: bash -c 'if [ -f "sdks/typescript/package.json" ]; then cd sdks/typescript && bun run build && bun run test; fi'

69
CLAUDE.md
View File

@ -93,6 +93,52 @@ All API routes follow the pattern: `/v1/{resource}/{id}/{action}`
- Explicit error handling with appropriate exception types
- Docstrings: Use Google style docstrings
### Agent Architecture
Honcho uses three specialized LLM agents that work together to form memories and answer queries:
#### 1. Deriver Agent (`src/deriver/agent/`)
**Role**: Memory formation through content ingestion
The Deriver processes incoming messages and extracts observations about peers.
- **Trigger**: Messages created via API are enqueued for background processing
- **Tools**: `create_observations`, `update_peer_card`, `get_recent_history`, `search_memory`, `get_observation_context`, `search_messages`
- **Output**: Explicit observations (direct facts) and deductive observations (inferences)
- **Entry point**: `src/deriver/agent/worker.py` → `Agent.run_loop()`
#### 2. Dialectic Agent (`src/dialectic/agent/`)
**Role**: Analysis and recall for answering queries
The Dialectic answers questions about peers by strategically gathering context from memory.
- **Trigger**: API call to `/peers/{peer_id}/chat` with `agentic=true`
- **Tools**: `search_memory`, `get_recent_history`, `get_observation_context`, `search_messages`, `get_recent_observations`, `get_most_derived_observations`, `get_session_summary`, `get_peer_card`, `create_observations` (deductive only)
- **Output**: Natural language response grounded in gathered context
- **Entry point**: `src/dialectic/chat.py` → `agentic_chat()` → `DialecticAgent.answer()`
#### 3. Dreamer Agent (`src/dreamer/agent.py`)
**Role**: Consolidation and self-improvement of memory
The Dreamer explores and consolidates observations to improve memory quality.
- **Trigger**: Scheduled or explicit dream task via queue
- **Tools**: `get_recent_observations`, `get_most_derived_observations`, `search_memory`, `create_observations`, `delete_observations`, `update_peer_card`
- **Strategy**: Random walk exploration - start from recent/high-value observations, search for related content, consolidate redundancies
- **Output**: Consolidated observations, deleted redundancies
- **Entry point**: `src/dreamer/agent.py` → `DreamerAgent.consolidate()`
#### Shared Agent Infrastructure
All agents share common infrastructure in `src/utils/agent_tools.py`:
- **Tool definitions**: Unified tool schemas used by all agents
- **Tool executor**: `create_tool_executor()` factory creates context-aware executors
- **LLM client**: `honcho_llm_call()` handles tool calling loops with configurable iterations
### Project Structure
```
@ -120,9 +166,12 @@ src/
│ └── workspace.py # Workspace CRUD operations
├── dialectic/ # Dialectic API implementation
│ ├── __init__.py
│ ├── chat.py # Chat functionality
│ ├── chat.py # Chat functionality (standard + agentic)
│ ├── prompts.py # Prompt templates
│ └── utils.py # Dialectic utilities
│ └── agent/ # Agentic dialectic implementation
│ ├── __init__.py
│ ├── core.py # DialecticAgent class
│ └── prompts.py # Agent system prompts
├── routers/ # API endpoints
│ ├── workspaces.py
│ ├── peers.py
@ -134,19 +183,25 @@ src/
│ ├── __init__.py
│ ├── __main__.py # Deriver entry point
│ ├── consumer.py # Message consumer
│ ├── deriver.py # Main deriver logic
│ ├── enqueue.py # Queue operations
│ ├── prompts.py # Deriver prompts
│ ├── queue_manager.py # Queue management
│ ├── queue_payload.py # Queue payload schemas
│ └── utils.py # Deriver utilities
│ └── agent/ # Agentic deriver implementation
│ ├── __init__.py
│ ├── core.py # Agent class
│ ├── worker.py # Task processing
│ └── prompts.py # Agent system prompts
├── dreamer/ # Memory consolidation system
│ ├── __init__.py
│ ├── agent.py # DreamerAgent class + process_agent_dream
│ └── dreamer.py # Legacy dreamer (scheduled)
├── utils/ # Utilities
│ ├── __init__.py
│ ├── agent_tools.py # Shared agent tools and executor
│ ├── clients.py # LLM client abstraction
│ ├── files.py # File handling utilities
│ ├── filter.py # Query filtering utilities
│ ├── formatting.py # Message formatting utilities
│ ├── logging.py # Logging configuration
│ ├── logging.py # Logging and metrics (Rich console output)
│ ├── search.py # Search functionality
│ ├── shared_models.py # Shared data models
│ ├── summarizer.py # Session summarization

89
config.toml.example
View File

@ -15,6 +15,8 @@ MAX_EMBEDDING_TOKENS = 8192
MAX_EMBEDDING_TOKENS_PER_REQUEST = 300000
# LANGFUSE_HOST = "https://api.langfuse.com"
# LANGFUSE_PUBLIC_KEY = "your-public-key-here"
# COLLECT_METRICS_LOCAL = false
# LOCAL_METRICS_FILE = "metrics.jsonl"
NAMESPACE="honcho"
# Database settings
@ -48,6 +50,9 @@ PROFILES_SAMPLE_RATE = 0.1
# LLM settings
[llm]
DEFAULT_MAX_TOKENS = 2500
EMBEDDING_PROVIDER = "openai"
MAX_TOOL_OUTPUT_CHARS = 30000 # Max chars for tool output (~7500 tokens)
MAX_MESSAGE_CONTENT_CHARS = 2000 # Max chars per message in tool results
# API Keys for LLM providers
# ANTHROPIC_API_KEY = "your-api-key"
@ -57,46 +62,72 @@ DEFAULT_MAX_TOKENS = 2500
# GROQ_API_KEY = "your-api-key"
# OPENAI_COMPATIBLE_BASE_URL = "your-base-url"
# Separate vLLM endpoint (for local models)
# VLLM_API_KEY = "your-api-key"
# VLLM_BASE_URL = "your-base-url"
# Deriver settings
[deriver]
ENABLED = true
WORKERS = 1
POLLING_SLEEP_INTERVAL_SECONDS = 1.0
STALE_SESSION_TIMEOUT_MINUTES = 5
# QUEUE_ERROR_RETENTION_SECONDS = 2592000 # 30 days
PROVIDER = "google"
MODEL = "gemini-2.5-flash-lite"
# BACKUP_PROVIDER = "anthropic"
# BACKUP_MODEL = "claude-haiku-4-5"
DEDUPLICATE = true
MAX_OUTPUT_TOKENS = 2500
THINKING_BUDGET_TOKENS = 1024 # only applied when using Anthropic
# BACKUP_PROVIDER = "google"
# BACKUP_MODEL = "gemini-2.5-flash-lite"
MAX_OUTPUT_TOKENS = 4096
THINKING_BUDGET_TOKENS = 1024
LOG_OBSERVATIONS = false
WORKING_REPRESENTATION_MAX_OBSERVATIONS = 100
REPRESENTATION_BATCH_MAX_TOKENS = 4096
MAX_INPUT_TOKENS = 23000
# Peer card settings
[peer_card]
ENABLED = true
PROVIDER = "openai"
MODEL = "gpt-5-nano-2025-08-07"
MAX_OUTPUT_TOKENS = 4000
# BACKUP_PROVIDER = "google"
# BACKUP_MODEL = "gemini-2.5-flash-lite"
# Dialectic settings
[dialectic]
MAX_OUTPUT_TOKENS = 8192
MAX_INPUT_TOKENS = 100000
HISTORY_TOKEN_LIMIT = 8192
SESSION_HISTORY_MAX_TOKENS = 16384
# Per-level settings for reasoning levels
[dialectic.levels.minimal]
PROVIDER = "google"
MODEL = "gemini-2.5-flash-lite"
THINKING_BUDGET_TOKENS = 0
MAX_TOOL_ITERATIONS = 2
[dialectic.levels.low]
PROVIDER = "google"
MODEL = "gemini-3-flash"
THINKING_BUDGET_TOKENS = 0
MAX_TOOL_ITERATIONS = 5
[dialectic.levels.medium]
PROVIDER = "anthropic"
MODEL = "claude-sonnet-4-20250514"
PERFORM_QUERY_GENERATION = false
QUERY_GENERATION_PROVIDER = "groq"
QUERY_GENERATION_MODEL = "llama-3.1-8b-instant"
MAX_OUTPUT_TOKENS = 2500
SEMANTIC_SEARCH_TOP_K = 10
SEMANTIC_SEARCH_MAX_DISTANCE = 0.85
THINKING_BUDGET_TOKENS = 1024
CONTEXT_WINDOW_SIZE = 100000
MODEL = "claude-haiku-4-5"
THINKING_BUDGET_TOKENS = 512
MAX_TOOL_ITERATIONS = 4
[dialectic.levels.high]
PROVIDER = "anthropic"
MODEL = "claude-opus-4-5"
THINKING_BUDGET_TOKENS = 0
MAX_TOOL_ITERATIONS = 4
[dialectic.levels.extra-high]
PROVIDER = "anthropic"
MODEL = "claude-opus-4-5"
THINKING_BUDGET_TOKENS = 512
MAX_TOOL_ITERATIONS = 10
# Backup provider example (optional, must set both or neither):
# BACKUP_PROVIDER = "google"
# BACKUP_MODEL = "gemini-2.5-flash"
# BACKUP_MODEL = "gemini-2.5-pro"
# Summary settings
[summary]
@ -106,7 +137,7 @@ MESSAGES_PER_LONG_SUMMARY = 60
PROVIDER = "google"
MODEL = "gemini-2.5-flash"
MAX_TOKENS_SHORT = 1000
MAX_TOKENS_LONG = 2000
MAX_TOKENS_LONG = 4000
THINKING_BUDGET_TOKENS = 512
# BACKUP_PROVIDER = "google"
# BACKUP_MODEL = "gemini-2.5-flash"
@ -117,10 +148,10 @@ ENABLED = true
DOCUMENT_THRESHOLD = 50
IDLE_TIMEOUT_MINUTES = 60
MIN_HOURS_BETWEEN_DREAMS = 8
ENABLED_TYPES = ["consolidate"]
PROVIDER = "openai"
MODEL = "gpt-4o-mini-2024-07-18"
MAX_OUTPUT_TOKENS = 2000
ENABLED_TYPES = ["omni"]
PROVIDER = "anthropic"
MODEL = "claude-haiku-4-5"
MAX_OUTPUT_TOKENS = 4000
# BACKUP_PROVIDER = "google"
# BACKUP_MODEL = "gemini-2.5-flash"
@ -132,12 +163,12 @@ MAX_WORKSPACE_LIMIT = 10
# Metrics settings
[metrics]
ENABLED = false
NAMESPACE = "honcho"
# NAMESPACE = "honcho" # Inherits from app.NAMESPACE if not set
# Cache settings
[cache]
ENABLED = false
URL = "redis://localhost:6379/0"
NAMESPACE="honcho"
URL = "redis://localhost:6379/0?suppress=false"
# NAMESPACE = "honcho" # Inherits from app.NAMESPACE if not set
DEFAULT_TTL_SECONDS = 300
DEFAULT_LOCK_TTL_SECONDS = 5

2
docs/v2.6.0-alpha/contributing/configuration.mdx
View File

@ -379,7 +379,7 @@ DREAM_IDLE_TIMEOUT_MINUTES=60 # Minutes of inactivity before dream can start
DREAM_MIN_HOURS_BETWEEN_DREAMS=8 # Minimum hours between dreams for a peer
# Dream types to enable
DREAM_ENABLED_TYPES=["consolidate"] # Currently supported: consolidate
DREAM_ENABLED_TYPES=["omni"] # Currently supported: omni
# LLM settings for dream processing
DREAM_PROVIDER=openai

44
migrations/versions/7c0d9a4e3b1f_add_unique_index_for_pending_dreams.py Normal file
View File

@ -0,0 +1,44 @@
"""add unique index for pending dream queue deduplication
This ensures `enqueue_dream()` is idempotent under concurrent calls by enforcing
at most one unprocessed dream queue item per `work_unit_key`.
Revision ID: 7c0d9a4e3b1f
Revises: f1a2b3c4d5e6
Create Date: 2026-01-12
"""
from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
from migrations.utils import get_schema
# revision identifiers, used by Alembic.
revision: str = "7c0d9a4e3b1f"
down_revision: str | None = "f1a2b3c4d5e6"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
schema = get_schema()
def upgrade() -> None:
"""Add a partial unique index to prevent duplicate pending dream queue items."""
op.create_index(
"ux_queue_dream_pending_work_unit_key",
"queue",
["work_unit_key"],
unique=True,
schema=schema,
postgresql_where=sa.text("task_type = 'dream' AND processed = false"),
)
def downgrade() -> None:
"""Drop the partial unique index for pending dream queue items."""
op.drop_index(
"ux_queue_dream_pending_work_unit_key", table_name="queue", schema=schema
)

2
migrations/versions/d429de0e5338_adopt_peer_paradigm.py
View File

@ -1165,7 +1165,7 @@ def backfill_token_counts(schema: str) -> None:
# Initialize tokenizer once outside the loop for performance
tokenizer = None
with suppress(Exception):
tokenizer = tiktoken.get_encoding("cl100k_base")
tokenizer = tiktoken.get_encoding("o200k_base")
def _count_tokens(text: str) -> int:
"""Count tokens in a text string using tiktoken."""

75
migrations/versions/f1a2b3c4d5e6_add_reasoning_tree_columns.py Normal file
View File

@ -0,0 +1,75 @@
"""add_reasoning_tree_columns
Add source_ids column to documents table for reasoning tree traversal.
This enables linking observations (deductive, inductive, contradiction)
to their source observations.
Revision ID: f1a2b3c4d5e6
Revises: 110bdf470272
Create Date: 2025-12-11
"""
from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects.postgresql import JSONB
from migrations.utils import column_exists, get_schema, index_exists
# revision identifiers, used by Alembic.
revision: str = "f1a2b3c4d5e6"
down_revision: str | None = "110bdf470272"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
schema = get_schema()
def upgrade() -> None:
"""Add source_ids column with GIN index for tree traversal."""
connection = op.get_bind()
inspector = sa.inspect(connection)
# Add source_ids column (nullable, for linking to parent observations)
if not column_exists("documents", "source_ids", inspector):
op.add_column(
"documents",
sa.Column(
"source_ids",
JSONB,
nullable=True,
server_default=sa.text("NULL"),
),
schema=schema,
)
# Add GIN index on source_ids for efficient child lookups
# (finding all observations that have a given observation as a source)
if not index_exists("documents", "ix_documents_source_ids_gin", inspector):
op.create_index(
"ix_documents_source_ids_gin",
"documents",
["source_ids"],
postgresql_using="gin",
schema=schema,
)
def downgrade() -> None:
"""Remove source_ids column and its index."""
connection = op.get_bind()
inspector = sa.inspect(connection)
# Drop GIN index first
if index_exists("documents", "ix_documents_source_ids_gin", inspector):
op.drop_index(
"ix_documents_source_ids_gin",
table_name="documents",
schema=schema,
)
# Drop column
if column_exists("documents", "source_ids", inspector):
op.drop_column("documents", "source_ids", schema=schema)

1
pyproject.toml
View File

@ -35,6 +35,7 @@ dependencies = [
"json-repair>=0.49.0",
"redis>=7.0.0,<8.0.0",
"cashews[redis]==7.4.4",
"scikit-learn>=1.6.0",
]
[tool.uv]
dev-dependencies = [

2
scripts/generate_message_embeddings.py
View File

@ -91,7 +91,7 @@ async def create_embeddings_for_messages(
return 0
# Initialize tiktoken encoding (same as used in MessageCreate schema)
encoding = tiktoken.get_encoding("cl100k_base")
encoding = tiktoken.get_encoding("o200k_base")
# Prepare data for batch embedding with proper token encoding
id_resource_dict = {

27
scripts/jsonl_to_json.py Executable file
View File

@ -0,0 +1,27 @@
#!/usr/bin/env python3
"""Convert a JSONL file to a JSON array."""
import json
import sys
def main() -> None:
"""Parse command-line arguments and convert JSONL file to JSON array.
Reads a JSONL file where each line is a valid JSON object, aggregates
all records into a list, and outputs as a formatted JSON array to stdout.
Exits with code 1 if arguments are invalid.
"""
if len(sys.argv) != 2:
print(f"Usage: {sys.argv[0]} <input.jsonl>", file=sys.stderr)
sys.exit(1)
with open(sys.argv[1]) as f:
records = [json.loads(line) for line in f if line.strip()]
print(json.dumps(records, indent=2))
if __name__ == "__main__":
main()

125
sdks/python/src/honcho/async_client/client.py
View File

@ -589,6 +589,131 @@ class AsyncHoncho(BaseModel):
await asyncio.sleep(sleep_time)
@validate_call
async def list_observations(
self,
filters: dict[str, object] | None = Field(
None, description="Filters to scope the observations"
),
reverse: bool = Field(
False, description="Whether to reverse the order of results"
),
):
"""
List all observations in the current workspace with optional filtering.
Makes an async API call to retrieve observations that match the specified filters.
Observations can be filtered by session_id, observer_id, and observed_id.
Args:
filters: Optional filter criteria for observations. Supported filters include:
- session_id: Filter observations by session
- observer_id: Filter observations by observer peer
- observed_id: Filter observations by observed peer
reverse: Whether to reverse the order of results (default: False)
Returns:
A paginated list of Observation objects matching the specified criteria
Example:
>>> observations = await client.list_observations(
... filters={"observer_id": "user123", "observed_id": "assistant"}
... )
"""
return await self._client.workspaces.observations.list(
workspace_id=self.workspace_id,
filters=filters,
reverse=reverse,
)
@validate_call
async def query_observations(
self,
query: str = Field(..., min_length=1, description="Semantic search query"),
observer: str = Field(
..., min_length=1, description="Observer peer ID (required)"
),
observed: str = Field(
..., min_length=1, description="Observed peer ID (required)"
),
top_k: int = Field(
default=10, ge=1, le=100, description="Number of results to return"
),
distance: float | None = Field(
default=None,
ge=0.0,
le=1.0,
description="Maximum cosine distance threshold for results",
),
filters: dict[str, object] | None = Field(
None, description="Additional filters to apply"
),
):
"""
Query observations using semantic search.
Performs vector similarity search on observations to find semantically relevant results.
Observer and observed peer IDs are required for semantic search.
Args:
query: The semantic search query
observer: The observer peer ID (required)
observed: The observed peer ID (required)
top_k: Number of results to return (1-100, default: 10)
distance: Maximum cosine distance threshold for results (0.0-1.0)
filters: Optional filters to scope the query
Returns:
A list of Observation objects matching the query
Example:
>>> observations = await client.query_observations(
... query="user preferences about music",
... observer="user123",
... observed="assistant",
... top_k=5,
... distance=0.8
... )
"""
# Merge observer/observed into filters without mutating the input
query_filters: dict[str, object | str] = {
**(filters or {}),
"observer": observer,
"observed": observed,
}
return await self._client.workspaces.observations.query(
workspace_id=self.workspace_id,
query=query,
top_k=top_k,
distance=distance,
filters=query_filters,
)
@validate_call
async def delete_observation(
self,
observation_id: str = Field(
..., min_length=1, description="ID of the observation to delete"
),
) -> None:
"""
Delete a specific observation by ID.
This permanently deletes the observation (document) from the theory-of-mind system.
This action cannot be undone.
Args:
observation_id: The ID of the observation to delete
Example:
>>> await client.delete_observation('obs_123abc')
"""
await self._client.workspaces.observations.delete(
workspace_id=self.workspace_id,
observation_id=observation_id,
)
@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
async def update_message(
self,

5
sdks/python/src/honcho/async_client/session.py
View File

@ -1,11 +1,11 @@
from __future__ import annotations
import asyncio
import json
import logging
import time
from typing import TYPE_CHECKING, Any
import json
from datetime import datetime
from typing import TYPE_CHECKING, Any
from honcho_core import AsyncHoncho as AsyncHonchoCore
from honcho_core._types import omit
@ -915,7 +915,6 @@ class AsyncSession(SessionBase):
if target is None
else (target if isinstance(target, str) else target.id)
)
data = await self._client.workspaces.peers.working_representation(
peer_id,
workspace_id=self.workspace_id,

125
sdks/python/src/honcho/client.py
View File

@ -564,6 +564,131 @@ class Honcho(BaseModel):
time.sleep(sleep_time)
@validate_call
def list_observations(
self,
filters: dict[str, object] | None = Field(
None, description="Filters to scope the observations"
),
reverse: bool = Field(
False, description="Whether to reverse the order of results"
),
):
"""
List all observations in the current workspace with optional filtering.
Makes an API call to retrieve observations that match the specified filters.
Observations can be filtered by session_id, observer_id, and observed_id.
Args:
filters: Optional filter criteria for observations. Supported filters include:
- session_id: Filter observations by session
- observer_id: Filter observations by observer peer
- observed_id: Filter observations by observed peer
reverse: Whether to reverse the order of results (default: False)
Returns:
A paginated list of Observation objects matching the specified criteria
Example:
>>> observations = client.list_observations(
... filters={"observer_id": "user123", "observed_id": "assistant"}
... )
"""
return self._client.workspaces.observations.list(
workspace_id=self.workspace_id,
filters=filters,
reverse=reverse,
)
@validate_call
def query_observations(
self,
query: str = Field(..., min_length=1, description="Semantic search query"),
observer: str = Field(
..., min_length=1, description="Observer peer ID (required)"
),
observed: str = Field(
..., min_length=1, description="Observed peer ID (required)"
),
top_k: int = Field(
default=10, ge=1, le=100, description="Number of results to return"
),
distance: float | None = Field(
default=None,
ge=0.0,
le=1.0,
description="Maximum cosine distance threshold for results",
),
filters: dict[str, object] | None = Field(
None, description="Additional filters to apply"
),
):
"""
Query observations using semantic search.
Performs vector similarity search on observations to find semantically relevant results.
Observer and observed peer IDs are required for semantic search.
Args:
query: The semantic search query
observer: The observer peer ID (required)
observed: The observed peer ID (required)
top_k: Number of results to return (1-100, default: 10)
distance: Maximum cosine distance threshold for results (0.0-1.0)
filters: Optional filters to scope the query
Returns:
A list of Observation objects matching the query
Example:
>>> observations = client.query_observations(
... query="user preferences about music",
... observer="user123",
... observed="assistant",
... top_k=5,
... distance=0.8
... )
"""
# Merge observer/observed into filters without mutating the input
query_filters: dict[str, object | str] = {
**(filters or {}),
"observer": observer,
"observed": observed,
}
return self._client.workspaces.observations.query(
workspace_id=self.workspace_id,
query=query,
top_k=top_k,
distance=distance,
filters=query_filters,
)
@validate_call
def delete_observation(
self,
observation_id: str = Field(
..., min_length=1, description="ID of the observation to delete"
),
) -> None:
"""
Delete a specific observation by ID.
This permanently deletes the observation (document) from the theory-of-mind system.
This action cannot be undone.
Args:
observation_id: The ID of the observation to delete
Example:
>>> client.delete_observation('obs_123abc')
"""
self._client.workspaces.observations.delete(
workspace_id=self.workspace_id,
observation_id=observation_id,
)
@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def update_message(
self,

7
sdks/python/src/honcho/session.py
View File

@ -1,10 +1,10 @@
from __future__ import annotations
import json
import logging
import time
from typing import TYPE_CHECKING, Any
import json
from datetime import datetime
from typing import TYPE_CHECKING, Any
from honcho_core import Honcho as HonchoCore
from honcho_core._types import omit
@ -20,8 +20,8 @@ from .session_context import SessionContext, SessionSummaries, Summary
from .utils import prepare_file_for_upload
if TYPE_CHECKING:
from .types import Representation
from .peer import Peer
from .types import Representation
logger = logging.getLogger(__name__)
@ -886,7 +886,6 @@ class Session(SessionBase):
if target is None
else (target if isinstance(target, str) else target.id)
)
data = self._client.workspaces.peers.working_representation(
peer_id,
workspace_id=self.workspace_id,

2
sdks/python/src/honcho/types.py
View File

@ -5,9 +5,9 @@ from __future__ import annotations
from collections.abc import AsyncIterator, Iterator
from datetime import datetime
from typing import TYPE_CHECKING, Any, cast
from typing_extensions import Required, TypedDict
from pydantic import BaseModel, Field
from typing_extensions import Required, TypedDict
# Re-export observation types from dedicated module
from .observations import AsyncObservationScope, Observation, ObservationScope

249
src/config.py
View File

@ -4,7 +4,7 @@ from typing import Annotated, Any, ClassVar, Literal, Protocol
import tomllib
from dotenv import load_dotenv
from pydantic import Field, field_validator, model_validator
from pydantic import BaseModel, Field, field_validator, model_validator
from pydantic.fields import FieldInfo
from pydantic_settings import (
BaseSettings,
@ -136,8 +136,8 @@ class BackupLLMSettingsMixin:
both fields are set together or both are None.
"""
BACKUP_PROVIDER: SupportedProviders | None = "custom"
BACKUP_MODEL: str | None = "x-ai/grok-4-fast"
BACKUP_PROVIDER: SupportedProviders | None = None
BACKUP_MODEL: str | None = None
@model_validator(mode="after")
def _validate_backup_configuration(self):
@ -206,15 +206,34 @@ class LLMSettings(HonchoSettings):
GROQ_API_KEY: str | None = None
OPENAI_COMPATIBLE_BASE_URL: str | None = None
EMBEDDING_PROVIDER: Literal["openai", "gemini"] = "openai"
# Separate vLLM endpoint (for local models)
VLLM_API_KEY: str | None = None
VLLM_BASE_URL: str | None = None
EMBEDDING_PROVIDER: Literal["openai", "gemini", "openrouter"] = "openai"
# General LLM settings
DEFAULT_MAX_TOKENS: Annotated[int, Field(default=1000, gt=0, le=100_000)] = 2500
# Maximum characters for tool output to prevent token explosion.
# Set to 30,000 chars (~7,500 tokens at 4 chars/token) to stay well under
# typical context limits while providing substantial tool output.
MAX_TOOL_OUTPUT_CHARS: Annotated[int, Field(default=30000, gt=0, le=100_000)] = (
30000
)
# Maximum characters for individual message content in tool results.
# Keeps each message preview concise while preserving key context.
MAX_MESSAGE_CONTENT_CHARS: Annotated[int, Field(default=2000, gt=0, le=10_000)] = (
2000
)
class DeriverSettings(BackupLLMSettingsMixin, HonchoSettings):
model_config = SettingsConfigDict(env_prefix="DERIVER_", extra="ignore") # pyright: ignore
ENABLED: bool = True
WORKERS: Annotated[int, Field(default=1, gt=0, le=100)] = 1
POLLING_SLEEP_INTERVAL_SECONDS: Annotated[
float, Field(default=1.0, gt=0.0, le=60.0)
@ -228,30 +247,29 @@ class DeriverSettings(BackupLLMSettingsMixin, HonchoSettings):
PROVIDER: SupportedProviders = "google"
MODEL: str = "gemini-2.5-flash-lite"
TEMPERATURE: float | None = None
# Whether to deduplicate documents when creating them
DEDUPLICATE: bool = True
MAX_OUTPUT_TOKENS: Annotated[int, Field(default=10_000, gt=0, le=100_000)] = 10_000
# Thinking budget tokens are only applied when using Anthropic as provider
MAX_OUTPUT_TOKENS: Annotated[int, Field(default=4096, gt=0, le=100_000)] = 4096
THINKING_BUDGET_TOKENS: Annotated[int, Field(default=1024, gt=0, le=5000)] = 1024
LOG_OBSERVATIONS: bool = False
MAX_INPUT_TOKENS: Annotated[int, Field(default=23000, gt=0, le=23000)] = 23000
# Maximum number of observations to return in working representation
# This is applied to both explicit and deductive observations
WORKING_REPRESENTATION_MAX_OBSERVATIONS: Annotated[
int, Field(default=50, gt=0, le=500)
] = 50
int, Field(default=100, gt=0, le=1000)
] = 100
REPRESENTATION_BATCH_MAX_TOKENS: Annotated[
int,
Field(
default=4096,
ge=1,
),
Field(default=4096, ge=128, le=16_384),
] = 4096
MAX_INPUT_TOKENS: Annotated[int, Field(default=23000, gt=0, le=23000)] = 23000
@model_validator(mode="after")
def validate_batch_tokens_vs_context_limit(self):
if self.REPRESENTATION_BATCH_MAX_TOKENS > self.MAX_INPUT_TOKENS:
@ -261,39 +279,122 @@ class DeriverSettings(BackupLLMSettingsMixin, HonchoSettings):
return self
class PeerCardSettings(BackupLLMSettingsMixin, HonchoSettings):
class PeerCardSettings(HonchoSettings):
model_config = SettingsConfigDict(env_prefix="PEER_CARD_", extra="ignore") # pyright: ignore
ENABLED: bool = True
PROVIDER: SupportedProviders = "openai"
MODEL: str = "gpt-5-nano-2025-08-07"
# Note: peer cards should be very short, but GPT-5 models need output tokens for thinking which cannot be turned off...
MAX_OUTPUT_TOKENS: Annotated[int, Field(default=4000, gt=1000, le=10_000)] = 4000
# Reasoning levels for dialectic - defined here to avoid circular imports with schemas
ReasoningLevel = Literal["minimal", "low", "medium", "high", "extra-high"]
REASONING_LEVELS: list[ReasoningLevel] = [
"minimal",
"low",
"medium",
"high",
"extra-high",
]
class DialecticSettings(BackupLLMSettingsMixin, HonchoSettings):
model_config = SettingsConfigDict(env_prefix="DIALECTIC_", extra="ignore") # pyright: ignore
class DialecticLevelSettings(BaseModel):
"""Settings for a specific reasoning level in the dialectic."""
PROVIDER: SupportedProviders = "anthropic"
MODEL: str = "claude-sonnet-4-20250514"
model_config = SettingsConfigDict(populate_by_name=True) # pyright: ignore
PERFORM_QUERY_GENERATION: bool = False
QUERY_GENERATION_PROVIDER: SupportedProviders = "groq"
QUERY_GENERATION_MODEL: str = "llama-3.1-8b-instant"
PROVIDER: Annotated[SupportedProviders, Field(validation_alias="provider")]
MODEL: Annotated[str, Field(validation_alias="model")]
BACKUP_PROVIDER: Annotated[
SupportedProviders | None, Field(validation_alias="backup_provider")
] = None
BACKUP_MODEL: Annotated[str | None, Field(validation_alias="backup_model")] = None
THINKING_BUDGET_TOKENS: Annotated[
int, Field(ge=0, le=100_000, validation_alias="thinking_budget_tokens")
]
MAX_TOOL_ITERATIONS: Annotated[
int, Field(ge=0, le=50, validation_alias="max_tool_iterations")
]
MAX_OUTPUT_TOKENS: Annotated[int, Field(default=2500, gt=0, le=100_000)] = 2500
@model_validator(mode="after")
def _validate_backup_configuration(self) -> "DialecticLevelSettings":
"""Ensure both backup fields are set together or both are None."""
if (self.BACKUP_PROVIDER is None) != (self.BACKUP_MODEL is None):
raise ValueError(
"BACKUP_PROVIDER and BACKUP_MODEL must both be set or both be None"
)
return self
SEMANTIC_SEARCH_TOP_K: Annotated[int, Field(default=10, gt=0, le=100)] = 10
SEMANTIC_SEARCH_MAX_DISTANCE: Annotated[
float, Field(default=0.85, ge=0.0, le=1.0)
] = 0.85 # Max distance for semantic search relevance
THINKING_BUDGET_TOKENS: Annotated[int, Field(default=1024, gt=0, le=5000)] = 1024
class DialecticSettings(HonchoSettings):
model_config = SettingsConfigDict( # pyright: ignore
env_prefix="DIALECTIC_", env_nested_delimiter="__", extra="ignore"
)
CONTEXT_WINDOW_SIZE: Annotated[
int, Field(default=100_000, gt=10_000, le=200_000)
] = 100_000
# Per-level settings for provider, model, thinking budget, and tool iterations
# TODO: Fill in appropriate values for each reasoning level
LEVELS: dict[ReasoningLevel, DialecticLevelSettings] = Field(
default_factory=lambda: {
"minimal": DialecticLevelSettings(
PROVIDER="google",
MODEL="gemini-2.5-flash-lite",
THINKING_BUDGET_TOKENS=0,
MAX_TOOL_ITERATIONS=2,
),
"low": DialecticLevelSettings(
PROVIDER="google",
MODEL="gemini-3-flash",
THINKING_BUDGET_TOKENS=0,
MAX_TOOL_ITERATIONS=5,
),
"medium": DialecticLevelSettings(
PROVIDER="anthropic",
MODEL="claude-haiku-4-5",
THINKING_BUDGET_TOKENS=512,
MAX_TOOL_ITERATIONS=4,
),
"high": DialecticLevelSettings(
PROVIDER="anthropic",
MODEL="claude-opus-4-5",
THINKING_BUDGET_TOKENS=0,
MAX_TOOL_ITERATIONS=4,
),
"extra-high": DialecticLevelSettings(
PROVIDER="anthropic",
MODEL="claude-opus-4-5",
THINKING_BUDGET_TOKENS=512,
MAX_TOOL_ITERATIONS=10,
),
}
)
MAX_OUTPUT_TOKENS: Annotated[int, Field(default=8192, gt=0, le=100_000)] = 8192
MAX_INPUT_TOKENS: Annotated[int, Field(default=100_000, gt=0, le=200_000)] = 100_000
# Token limit for get_recent_history tool within the agent
HISTORY_TOKEN_LIMIT: Annotated[int, Field(default=8192, gt=0, le=100_000)] = 8192
# Session history injection: max tokens of recent messages to include when session_id is specified.
# Set to 0 to disable automatic session history injection.
SESSION_HISTORY_MAX_TOKENS: Annotated[
int, Field(default=16_384, ge=0, le=100_000)
] = 16_384
@model_validator(mode="after")
def _validate_token_budgets(self) -> "DialecticSettings":
"""Ensure the output token limit exceeds all thinking budgets."""
for level, level_settings in self.LEVELS.items():
if self.MAX_OUTPUT_TOKENS <= level_settings.THINKING_BUDGET_TOKENS:
raise ValueError(
f"MAX_OUTPUT_TOKENS must be greater than THINKING_BUDGET_TOKENS for level '{level}'"
)
return self
@model_validator(mode="after")
def _validate_all_levels_present(self) -> "DialecticSettings":
"""Ensure all reasoning levels are configured."""
missing = set(REASONING_LEVELS) - set(self.LEVELS.keys())
if missing:
raise ValueError(f"Missing configuration for reasoning levels: {missing}")
return self
class SummarySettings(BackupLLMSettingsMixin, HonchoSettings):
@ -304,8 +405,8 @@ class SummarySettings(BackupLLMSettingsMixin, HonchoSettings):
MESSAGES_PER_SHORT_SUMMARY: Annotated[int, Field(default=20, gt=0, le=100)] = 20
MESSAGES_PER_LONG_SUMMARY: Annotated[int, Field(default=60, gt=0, le=500)] = 60
PROVIDER: SupportedProviders = "openai"
MODEL: str = "gpt-4o-mini-2024-07-18"
PROVIDER: SupportedProviders = "google"
MODEL: str = "gemini-2.5-flash"
MAX_TOKENS_SHORT: Annotated[int, Field(default=1000, gt=0, le=10_000)] = 1000
MAX_TOKENS_LONG: Annotated[int, Field(default=4000, gt=0, le=20_000)] = 4000
@ -341,19 +442,80 @@ class CacheSettings(HonchoSettings):
)
class SurprisalSettings(BaseModel):
"""Settings for tree-based surprisal sampling during dreams."""
ENABLED: bool = False
# Tree configuration
TREE_TYPE: Literal[
"kdtree", "balltree", "rptree", "covertree", "lsh", "graph", "prototype"
] = "kdtree"
TREE_K: Annotated[int, Field(default=5, gt=0, le=20)] = 5 # k for kNN-based trees
# Sampling strategy
SAMPLING_STRATEGY: Literal["recent", "random", "all"] = "recent"
SAMPLE_SIZE: Annotated[int, Field(default=200, gt=0, le=2000)] = 200
# Surprisal filtering (normalized scores: 0.0 = lowest, 1.0 = highest)
TOP_PERCENT_SURPRISAL: Annotated[float, Field(default=0.10, gt=0.0, le=1.0)] = (
0.10 # Top 10% of observations
)
# Hybrid mode: min high-surprisal observations to replace standard questions
MIN_HIGH_SURPRISAL_FOR_REPLACE: Annotated[int, Field(default=10, gt=0)] = 10
# Observation level filtering
INCLUDE_LEVELS: list[str] = ["explicit", "deductive"]
class DreamSettings(BackupLLMSettingsMixin, HonchoSettings):
model_config = SettingsConfigDict(env_prefix="DREAM_", extra="ignore") # pyright: ignore
model_config = SettingsConfigDict( # pyright: ignore
env_prefix="DREAM_", env_nested_delimiter="__", extra="ignore"
)
ENABLED: bool = True
DOCUMENT_THRESHOLD: Annotated[int, Field(default=50, gt=0, le=1000)] = 50
IDLE_TIMEOUT_MINUTES: Annotated[int, Field(default=60, gt=0, le=1440)] = 60
MIN_HOURS_BETWEEN_DREAMS: Annotated[int, Field(default=8, gt=0, le=72)] = 8
ENABLED_TYPES: list[str] = ["consolidate"]
ENABLED_TYPES: list[str] = ["omni"]
# LLM settings for dream processing
PROVIDER: SupportedProviders = "google"
MODEL: str = "gemini-2.5-flash"
MAX_OUTPUT_TOKENS: Annotated[int, Field(default=2000, gt=0, le=10_000)] = 2000
# LLM settings for dream processing - upgraded for extended reasoning
PROVIDER: SupportedProviders = "anthropic"
MODEL: str = "claude-sonnet-4-20250514" # Upgraded from haiku for reasoning
MAX_OUTPUT_TOKENS: Annotated[int, Field(default=16_384, gt=0, le=64_000)] = 16_384
THINKING_BUDGET_TOKENS: Annotated[int, Field(default=8192, gt=0, le=32_000)] = 8192
# Agent iteration limit - increased for extended reasoning workflow
MAX_TOOL_ITERATIONS: Annotated[int, Field(default=20, gt=0, le=50)] = 20
# Token limit for get_recent_history tool within the agent
HISTORY_TOKEN_LIMIT: Annotated[int, Field(default=16_384, gt=0, le=200_000)] = (
16_384
)
# Observation limits for orchestrated dreaming prescan
# Higher = more context for reasoning but slower prescan
PRESCAN_OBSERVATIONS_PER_LEVEL: Annotated[
int, Field(default=200, gt=0, le=1000)
] = 200
# Specialist model settings (OpenRouter format: provider/model)
# DeductionSpecialist: handles logical inference + temporal reasoning
DEDUCTION_MODEL: str = "anthropic/claude-haiku-4.5"
# InductionSpecialist: identifies patterns across observations
INDUCTION_MODEL: str = "anthropic/claude-haiku-4.5"
# Surprisal-based sampling subsystem
SURPRISAL: SurprisalSettings = Field(default_factory=SurprisalSettings)
@model_validator(mode="after")
def _validate_token_budgets(self) -> "DreamSettings":
"""Ensure the output token limit exceeds the thinking budget."""
if self.MAX_OUTPUT_TOKENS <= self.THINKING_BUDGET_TOKENS:
raise ValueError(
"MAX_OUTPUT_TOKENS must be greater than THINKING_BUDGET_TOKENS"
)
return self
class AppSettings(HonchoSettings):
@ -381,6 +543,7 @@ class AppSettings(HonchoSettings):
COLLECT_METRICS_LOCAL: bool = False
LOCAL_METRICS_FILE: str = "metrics.jsonl"
REASONING_TRACES_FILE: str | None = None # Path to JSONL file for reasoning traces
NAMESPACE: str = "honcho" # Top-level namespace for all settings, can be overridden by nested-model settings

19
src/crud/__init__.py
View File

@ -1,21 +1,29 @@
from .collection import get_collection, get_or_create_collection
from .deriver import get_deriver_status
from .deriver import get_deriver_status, get_queue_status
from .document import (
create_documents,
create_observations,
delete_document,
delete_document_by_id,
get_all_documents,
get_child_observations,
get_documents_by_ids,
get_documents_with_filters,
query_documents,
query_documents_most_derived,
query_documents_recent,
)
from .message import (
create_messages,
get_message,
get_message_seq_in_session,
get_messages,
get_messages_by_date_range,
get_messages_by_seq_range,
get_messages_id_range,
grep_messages,
search_messages,
search_messages_temporal,
update_message,
)
from .peer import (
@ -62,21 +70,30 @@ __all__ = [
"get_or_create_collection",
# Deriver
"get_deriver_status",
"get_queue_status",
# Document
"create_documents",
"create_observations",
"get_all_documents",
"get_child_observations",
"get_documents_by_ids",
"get_documents_with_filters",
"query_documents",
"query_documents_most_derived",
"query_documents_recent",
"delete_document",
"delete_document_by_id",
# Message
"create_messages",
"get_messages",
"get_messages_by_date_range",
"get_messages_by_seq_range",
"get_messages_id_range",
"get_message",
"get_message_seq_in_session",
"grep_messages",
"search_messages",
"search_messages_temporal",
"update_message",
# Peer
"get_or_create_peers",

35
src/crud/deriver.py
View File

@ -11,16 +11,16 @@ from src import models, schemas
logger = getLogger(__name__)
async def get_deriver_status(
async def get_queue_status(
db: AsyncSession,
workspace_name: str,
session_name: str | None = None,
*,
observer: str | None = None,
observed: str | None = None,
) -> schemas.DeriverStatus:
) -> schemas.QueueStatus:
"""
Get the deriver processing status, optionally filtered by observer, sender, and/or session.
Get the processing queue status, optionally filtered by observer, sender, and/or session.
Args:
db: Database session
@ -50,6 +50,25 @@ async def get_deriver_status(
)
async def get_deriver_status(
db: AsyncSession,
workspace_name: str,
session_name: str | None = None,
*,
observer: str | None = None,
observed: str | None = None,
) -> schemas.QueueStatus:
"""Deprecated: use get_queue_status."""
return await get_queue_status(
db=db,
workspace_name=workspace_name,
session_name=session_name,
observer=observer,
observed=observed,
)
def _build_queue_status_query(
workspace_name: str,
session_name: str | None,
@ -157,21 +176,21 @@ def _process_queue_rows(rows: Sequence[Row[Any]]) -> schemas.QueueCounts:
def _build_status_response(
session_name: str | None,
counts: schemas.QueueCounts,
) -> schemas.DeriverStatus:
) -> schemas.QueueStatus:
"""Build the final response object."""
if session_name:
return schemas.DeriverStatus(
return schemas.QueueStatus(
total_work_units=counts.total,
completed_work_units=counts.completed,
in_progress_work_units=counts.in_progress,
pending_work_units=counts.pending,
)
sessions: dict[str, schemas.SessionDeriverStatus] = {}
sessions: dict[str, schemas.SessionQueueStatus] = {}
for session_id, data in counts.sessions.items():
total = data.completed + data.in_progress + data.pending
sessions[session_id] = schemas.SessionDeriverStatus(
sessions[session_id] = schemas.SessionQueueStatus(
session_id=session_id,
total_work_units=total,
completed_work_units=data.completed,
@ -179,7 +198,7 @@ def _build_status_response(
pending_work_units=data.pending,
)
return schemas.DeriverStatus(
return schemas.QueueStatus(
sessions=sessions if sessions else None,
total_work_units=counts.total,
completed_work_units=counts.completed,

147
src/crud/document.py
View File

@ -102,6 +102,80 @@ def get_documents_with_filters(
return stmt
async def query_documents_recent(
db: AsyncSession,
workspace_name: str,
*,
observer: str,
observed: str,
limit: int = 10,
session_name: str | None = None,
) -> Sequence[models.Document]:
"""
Query most recent documents.
Args:
db: Database session
workspace_name: Name of the workspace
observer: Name of the observing peer
observed: Name of the observed peer
limit: Maximum number of documents to return
session_name: Optional session name to filter by
Returns:
Sequence of documents ordered by created_at descending
"""
stmt = select(models.Document).where(
models.Document.workspace_name == workspace_name,
models.Document.observer == observer,
models.Document.observed == observed,
)
if session_name is not None:
stmt = stmt.where(models.Document.session_name == session_name)
stmt = stmt.order_by(models.Document.created_at.desc()).limit(limit)
result = await db.execute(stmt)
return result.scalars().all()
async def query_documents_most_derived(
db: AsyncSession,
workspace_name: str,
*,
observer: str,
observed: str,
limit: int = 10,
) -> Sequence[models.Document]:
"""
Query documents sorted by times_derived (most reinforced first).
Args:
db: Database session
workspace_name: Name of the workspace
observer: Name of the observing peer
observed: Name of the observed peer
limit: Maximum number of documents to return
Returns:
Sequence of documents ordered by times_derived descending
"""
stmt = (
select(models.Document)
.where(
models.Document.workspace_name == workspace_name,
models.Document.observer == observer,
models.Document.observed == observed,
)
.order_by(models.Document.times_derived.desc())
.limit(limit)
)
result = await db.execute(stmt)
return result.scalars().all()
async def query_documents(
db: AsyncSession,
workspace_name: str,
@ -205,6 +279,8 @@ async def create_documents(
internal_metadata=metadata_dict,
embedding=doc.embedding,
session_name=doc.session_name,
# Tree linkage column
source_ids=doc.source_ids,
)
)
except Exception as e:
@ -299,7 +375,7 @@ async def delete_document_by_id(
async def create_observations(
db: AsyncSession,
observations: list[schemas.ObservationCreate],
observations: Sequence[schemas.ConclusionCreate],
workspace_name: str,
) -> list[models.Document]:
"""
@ -455,3 +531,72 @@ async def is_rejected_duplicate(
f"[DUPLICATE DETECTION] Rejecting new in favor of existing. new='{doc.content}', existing='{existing_doc.content}'."
)
return True
# =============================================================================
# Tree Traversal Functions - For reasoning chain navigation
# =============================================================================
async def get_documents_by_ids(
db: AsyncSession,
workspace_name: str,
document_ids: list[str],
) -> Sequence[models.Document]:
"""
Get multiple documents by their IDs.
Args:
db: Database session
workspace_name: Workspace identifier
document_ids: List of document IDs to retrieve
Returns:
Sequence of documents found (may be fewer than requested if some IDs don't exist)
"""
if not document_ids:
return []
stmt = select(models.Document).where(
models.Document.workspace_name == workspace_name,
models.Document.id.in_(document_ids),
)
result = await db.execute(stmt)
return result.scalars().all()
async def get_child_observations(
db: AsyncSession,
workspace_name: str,
parent_id: str,
*,
observer: str | None = None,
observed: str | None = None,
) -> Sequence[models.Document]:
"""
Get all observations that have this document as a source/premise.
Useful for traversing the reasoning tree upward (source -> derived observations).
Uses GIN index on source_ids for efficient lookups.
Args:
db: Database session
workspace_name: Workspace identifier
parent_id: Document ID to find children of
observer: Optional filter by observer
observed: Optional filter by observed
Returns:
Sequence of documents that reference this document as a source
"""
# Find documents where source_ids contains the parent_id
stmt = select(models.Document).where(
models.Document.workspace_name == workspace_name,
models.Document.source_ids.contains([parent_id]),
)
if observer:
stmt = stmt.where(models.Document.observer == observer)
if observed:
stmt = stmt.where(models.Document.observed == observed)
result = await db.execute(stmt)
return result.scalars().all()

288
src/crud/message.py
View File

@ -1,3 +1,4 @@
from datetime import datetime
from logging import getLogger
from typing import Any
@ -9,6 +10,7 @@ from src import models, schemas
from src.config import settings
from src.embedding_client import embedding_client
from src.utils.filter import apply_filter
from src.utils.formatting import ILIKE_ESCAPE_CHAR, escape_ilike_pattern
from .session import get_or_create_session
@ -52,6 +54,77 @@ def _apply_token_limit(
)
async def _build_merged_snippets(
db: AsyncSession,
workspace_name: str,
matched_messages: list[models.Message],
context_window: int,
) -> list[tuple[list[models.Message], list[models.Message]]]:
"""
Group matched messages by session, merge overlapping context ranges, and fetch context.
Takes a list of matched messages and builds conversation snippets by:
1. Grouping matches by session name
2. Sorting matches within each session by sequence number
3. Merging overlapping context windows to avoid duplicate context
4. Fetching the full context for each merged range from the database
Args:
db: Database session
workspace_name: Name of the workspace
matched_messages: List of messages that matched a search query
context_window: Number of messages before/after each match to include
Returns:
List of tuples: (matched_messages_in_range, context_messages)
Each tuple represents a snippet where context_messages includes all messages
in the merged range (including the matched messages), ordered chronologically.
"""
if not matched_messages:
return []
session_matches: dict[str, list[models.Message]] = {}
for msg in matched_messages:
session_matches.setdefault(msg.session_name, []).append(msg)
snippets: list[tuple[list[models.Message], list[models.Message]]] = []
for sess_name, matches in session_matches.items():
matches.sort(key=lambda m: m.seq_in_session)
merged_ranges: list[tuple[int, int, list[models.Message]]] = []
for match in matches:
start = match.seq_in_session - context_window
end = match.seq_in_session + context_window
if merged_ranges and start <= merged_ranges[-1][1] + 1:
prev_start, prev_end, prev_matches = merged_ranges[-1]
merged_ranges[-1] = (
prev_start,
max(prev_end, end),
[*prev_matches, match],
)
else:
merged_ranges.append((start, end, [match]))
for start_seq, end_seq, range_matches in merged_ranges:
context_stmt = (
select(models.Message)
.where(models.Message.workspace_name == workspace_name)
.where(models.Message.session_name == sess_name)
.where(models.Message.seq_in_session.between(start_seq, end_seq))
.order_by(models.Message.seq_in_session.asc())
)
context_result = await db.execute(context_stmt)
context_messages = list(context_result.scalars().all())
snippets.append((range_matches, context_messages))
return snippets
async def create_messages(
db: AsyncSession,
messages: list[schemas.MessageCreate],
@ -397,3 +470,218 @@ async def update_message(
await db.commit()
# await db.refresh(honcho_message)
return honcho_message
async def search_messages(
db: AsyncSession,
workspace_name: str,
session_name: str | None,
query: str,
limit: int = 10,
context_window: int = 2,
) -> list[tuple[list[models.Message], list[models.Message]]]:
"""
Search for messages using semantic similarity and return conversation snippets.
Each result includes matched messages plus surrounding context. Overlapping
snippets within the same session are merged to avoid repetition.
Args:
db: Database session
workspace_name: Name of the workspace
session_name: Name of the session (optional)
query: Search query text
limit: Maximum number of matching messages to return
context_window: Number of messages before/after each match to include
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)
# First, find the top matching messages
match_stmt = (
select(models.Message)
.join(
models.MessageEmbedding,
models.Message.public_id == models.MessageEmbedding.message_id,
)
.where(models.MessageEmbedding.workspace_name == workspace_name)
.order_by(models.MessageEmbedding.embedding.cosine_distance(query_embedding))
.limit(limit)
)
if session_name:
match_stmt = match_stmt.where(
models.MessageEmbedding.session_name == session_name
)
result = await db.execute(match_stmt)
matched_messages = list(result.scalars().all())
return await _build_merged_snippets(
db, workspace_name, matched_messages, context_window
)
async def grep_messages(
db: AsyncSession,
workspace_name: str,
session_name: str | None,
text: str,
limit: int = 10,
context_window: int = 2,
) -> list[tuple[list[models.Message], list[models.Message]]]:
"""
Search for messages containing specific text (case-insensitive substring match).
Unlike semantic search, this finds EXACT text matches. Useful for finding
specific names, dates, phrases, or keywords.
Args:
db: Database session
workspace_name: Name of the workspace
session_name: Name of the session (optional - searches all sessions if None)
text: Text to search for (case-insensitive)
limit: Maximum number of matching messages to return
context_window: Number of messages before/after each match to include
Returns:
List of tuples: (matched_messages, context_messages)
Each snippet may contain multiple matches if they were close together.
"""
# Build the base query with ILIKE for case-insensitive text search
escaped_text = escape_ilike_pattern(text)
match_stmt = (
select(models.Message)
.where(models.Message.workspace_name == workspace_name)
.where(
models.Message.content.ilike(f"%{escaped_text}%", escape=ILIKE_ESCAPE_CHAR)
)
.order_by(models.Message.created_at.desc())
.limit(limit)
)
if session_name:
match_stmt = match_stmt.where(models.Message.session_name == session_name)
result = await db.execute(match_stmt)
matched_messages = list(result.scalars().all())
return await _build_merged_snippets(
db, workspace_name, matched_messages, context_window
)
async def get_messages_by_date_range(
db: AsyncSession,
workspace_name: str,
session_name: str | None,
after_date: datetime | None = None,
before_date: datetime | None = None,
limit: int = 20,
order: str = "desc",
) -> list[models.Message]:
"""
Get messages within a date range.
Args:
db: Database session
workspace_name: Name of the workspace
session_name: Name of the session (optional - searches all sessions if None)
after_date: Return messages after this datetime
before_date: Return messages before this datetime
limit: Maximum messages to return
order: Sort order - 'asc' for oldest first, 'desc' for newest first
Returns:
List of messages within the date range
"""
stmt = select(models.Message).where(models.Message.workspace_name == workspace_name)
if session_name:
stmt = stmt.where(models.Message.session_name == session_name)
if after_date:
stmt = stmt.where(models.Message.created_at >= after_date)
if before_date:
stmt = stmt.where(models.Message.created_at <= before_date)
if order == "asc":
stmt = stmt.order_by(models.Message.created_at.asc())
else:
stmt = stmt.order_by(models.Message.created_at.desc())
stmt = stmt.limit(limit)
result = await db.execute(stmt)
return list(result.scalars().all())
async def search_messages_temporal(
db: AsyncSession,
workspace_name: str,
session_name: str | None,
query: str,
after_date: datetime | None = None,
before_date: datetime | None = None,
limit: int = 10,
context_window: int = 2,
) -> list[tuple[list[models.Message], list[models.Message]]]:
"""
Search for messages using semantic similarity with optional date filtering.
Combines the power of semantic search with time constraints. Use after_date
to find recent mentions, or before_date to find what was said before a certain point.
Args:
db: Database session
workspace_name: Name of the workspace
session_name: Name of the session (optional)
query: Search query text
after_date: Only return messages after this datetime
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
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)
# Build query with date filters
match_stmt = (
select(models.Message)
.join(
models.MessageEmbedding,
models.Message.public_id == models.MessageEmbedding.message_id,
)
.where(models.MessageEmbedding.workspace_name == workspace_name)
)
if session_name:
match_stmt = match_stmt.where(
models.MessageEmbedding.session_name == session_name
)
# Apply date filters on the Message table
if after_date:
match_stmt = match_stmt.where(models.Message.created_at >= after_date)
if before_date:
match_stmt = match_stmt.where(models.Message.created_at <= before_date)
# Order by similarity and limit
match_stmt = match_stmt.order_by(
models.MessageEmbedding.embedding.cosine_distance(query_embedding)
).limit(limit)
result = await db.execute(match_stmt)
matched_messages = list(result.scalars().all())
return await _build_merged_snippets(
db, workspace_name, matched_messages, context_window
)

2
src/crud/representation.py
View File

@ -55,7 +55,7 @@ class RepresentationManager:
Args:
representation: Representation object
message_id_range: Message ID range to link with observations
message_ids: Message ID range to link with observations
session_name: Session name to link with existing summary context
message_created_at: Timestamp when the message was created

3
src/crud/workspace.py
View File

@ -266,8 +266,7 @@ async def delete_workspace(db: AsyncSession, workspace_name: str) -> schemas.Wor
# Then delete QueueItem entries
await db.execute(
delete(models.QueueItem).where(
func.split_part(models.QueueItem.work_unit_key, ":", 2)
== workspace_name
models.QueueItem.workspace_name == workspace_name
)
)

13
src/deriver/consumer.py
View File

@ -2,7 +2,6 @@ import logging
import sentry_sdk
from pydantic import ValidationError
from rich.console import Console
from sqlalchemy import select
from src import crud, models
@ -25,8 +24,6 @@ from src.webhooks import webhook_delivery
logger = logging.getLogger(__name__)
logging.getLogger("sqlalchemy.engine.Engine").disabled = True
console = Console(markup=True)
async def process_item(queue_item: models.QueueItem) -> None:
"""Process a single item from the queue."""
@ -133,9 +130,12 @@ async def process_representation_batch(
observer: str | None,
observed: str | None,
) -> None:
"""Prepares and processes a batch of messages for representation tasks.
"""
Prepares and processes a batch of messages for representation tasks.
Args:
messages: List of messages to process
message_level_configuration: Resolved configuration for this batch
observer: The observer of the messages
observed: The observed of the messages
"""
@ -146,11 +146,6 @@ async def process_representation_batch(
if observed is None or observer is None:
raise ValueError("observed and observer are required for representation tasks")
logger.debug(
"process_representation_batch received %s messages",
len(messages),
)
await process_representation_tasks_batch(
messages,
message_level_configuration,

537
src/deriver/deriver.py
View File

@ -1,131 +1,26 @@
import datetime
import logging
import time
import sentry_sdk
from src import crud, exceptions, prometheus
from src import crud, prometheus
from src.config import settings
from src.crud.representation import RepresentationManager
from src.dependencies import tracked_db
from src.models import Message
from src.schemas import ResolvedConfiguration
from src.utils import summarizer
from src.utils.clients import honcho_llm_call
from src.utils.config_helpers import get_configuration
from src.utils.formatting import format_new_turn_with_timestamp
from src.utils.logging import (
accumulate_metric,
conditional_observe,
log_performance_metrics,
# log_representation,
)
from src.utils.peer_card import PeerCardQuery
from src.utils.logging import accumulate_metric, log_performance_metrics
from src.utils.representation import PromptRepresentation, Representation
from src.utils.tokens import estimate_tokens, track_input_tokens
from src.utils.tokens import estimate_tokens, track_deriver_input_tokens
from src.utils.tracing import with_sentry_transaction
from .prompts import (
critical_analysis_prompt,
estimate_critical_analysis_prompt_tokens,
estimate_peer_card_prompt_tokens,
peer_card_prompt,
)
from .prompts import estimate_minimal_deriver_prompt_tokens, minimal_deriver_prompt
logger = logging.getLogger(__name__)
logging.getLogger("sqlalchemy.engine.Engine").disabled = True
@conditional_observe(name="Critical Analysis Call")
async def critical_analysis_call(
peer_id: str,
peer_card: list[str] | None,
message_created_at: datetime.datetime,
working_representation: Representation,
history: str,
new_turns: list[str],
) -> PromptRepresentation:
prompt = critical_analysis_prompt(
peer_id=peer_id,
peer_card=peer_card,
message_created_at=message_created_at,
working_representation=working_representation,
history=history,
new_turns=new_turns,
)
response = await honcho_llm_call(
llm_settings=settings.DERIVER,
prompt=prompt,
max_tokens=settings.DERIVER.MAX_OUTPUT_TOKENS
or settings.LLM.DEFAULT_MAX_TOKENS,
track_name="Critical Analysis Call",
response_model=PromptRepresentation,
json_mode=True,
stop_seqs=[" \n", "\n\n\n\n"],
thinking_budget_tokens=settings.DERIVER.THINKING_BUDGET_TOKENS,
reasoning_effort="minimal",
verbosity="medium",
enable_retry=True,
retry_attempts=3,
)
prometheus.DERIVER_TOKENS_PROCESSED.labels(
task_type="representation",
token_type="output", # nosec B106
component="total",
).inc(response.output_tokens)
return response.content
@conditional_observe(name="Peer Card Call")
async def peer_card_call(
old_peer_card: list[str] | None,
new_observations: Representation,
) -> PeerCardQuery:
"""
Generate peer card prompt, call LLM with response model.
"""
prompt = peer_card_prompt(
old_peer_card=old_peer_card,
new_observations=new_observations.str_no_timestamps(),
)
response = await honcho_llm_call(
llm_settings=settings.PEER_CARD,
prompt=prompt,
max_tokens=settings.PEER_CARD.MAX_OUTPUT_TOKENS
or settings.LLM.DEFAULT_MAX_TOKENS,
track_name="Peer Card Call",
response_model=PeerCardQuery,
json_mode=True,
reasoning_effort="minimal",
enable_retry=True,
retry_attempts=3,
)
# Track input tokens for peer_card task
track_input_tokens(
task_type="peer_card",
components={
"prompt": estimate_peer_card_prompt_tokens(),
"old_peer_card": estimate_tokens(old_peer_card),
"new_observations": estimate_tokens(new_observations.str_no_timestamps()),
},
)
# Track output tokens for peer_card task
prometheus.DERIVER_TOKENS_PROCESSED.labels(
task_type="peer_card",
token_type="output", # nosec B106
component="total",
).inc(response.output_tokens)
return response.content
@with_sentry_transaction("process_representation_tasks_batch", op="deriver")
@with_sentry_transaction("minimal_deriver_batch", op="deriver")
async def process_representation_tasks_batch(
messages: list[Message],
message_level_configuration: ResolvedConfiguration | None,
@ -134,48 +29,28 @@ async def process_representation_tasks_batch(
observed: str,
) -> None:
"""
Process a batch of representation tasks by extracting insights and updating working representations.
Process messages with minimal overhead - single LLM call, no peer card.
Args:
messages: List of messages to process (includes interleaving context).
message_level_configuration: Optional configuration override.
observer: The observer peer ID.
observed: The observed peer ID.
"""
if not messages or len(messages) == 0:
if not messages:
return
messages.sort(key=lambda x: x.id)
overall_start = time.perf_counter()
messages.sort(key=lambda x: x.id)
latest_message = messages[-1]
earliest_message = messages[0]
accumulate_metric(
f"deriver_{latest_message.id}_{observer}",
"starting_message_id",
earliest_message.id,
"id",
)
accumulate_metric(
f"deriver_{latest_message.id}_{observer}",
"ending_message_id",
latest_message.id,
"id",
)
# Start overall timing
overall_start = time.perf_counter()
# Time context preparation
context_prep_start = time.perf_counter()
# Use get_session_context_formatted with configurable token limit
working_representation = await crud.get_working_representation(
latest_message.workspace_name,
observer=observer,
observed=observed,
# include_semantic_query=latest_message.content,
# include_most_derived=False,
)
async with tracked_db("deriver.get_peer_card") as db:
if message_level_configuration is None:
# Get configuration if not provided
# TODO: this appears to be a very rare edge case coming out of `get_queue_item_batch` in queue_manager.py,
# possible that we can remove this and require configuration to come through with the payload.
if message_level_configuration is None:
async with tracked_db("minimal_deriver.get_config") as db:
message_level_configuration = get_configuration(
None,
await crud.get_session(
@ -185,308 +60,152 @@ async def process_representation_tasks_batch(
db, workspace_name=latest_message.workspace_name
),
)
if message_level_configuration.peer_card.use is False:
speaker_peer_card = None
else:
speaker_peer_card = await crud.get_peer_card(
db,
latest_message.workspace_name,
observer=observer,
observed=observed,
)
# Skip if deriver disabled
if message_level_configuration.deriver.enabled is False:
return
# Estimate tokens for deriver input
peer_card_tokens = estimate_tokens(speaker_peer_card)
working_rep_tokens = estimate_tokens(
str(working_representation) if not working_representation.is_empty() else None
accumulate_metric(
f"minimal_deriver_{latest_message.id}_{observer}",
"starting_message_id",
earliest_message.id,
"id",
)
accumulate_metric(
f"minimal_deriver_{latest_message.id}_{observer}",
"ending_message_id",
latest_message.id,
"id",
)
prompt_tokens = estimate_critical_analysis_prompt_tokens()
# Estimate tokens for new conversation turns
new_turns = [
format_new_turn_with_timestamp(m.content, m.created_at, m.peer_name)
for m in messages
]
new_turns_tokens = estimate_tokens(new_turns)
estimated_input_tokens = (
peer_card_tokens + working_rep_tokens + prompt_tokens + new_turns_tokens
# Format messages with timestamps
formatted_messages = "\n".join(
format_new_turn_with_timestamp(msg.content, msg.created_at, msg.peer_name)
for msg in messages
)
# Calculate available tokens for context
safety_buffer = 500
available_context_tokens = max(
0,
settings.DERIVER.MAX_INPUT_TOKENS - estimated_input_tokens - safety_buffer,
)
async with tracked_db("deriver.get_session_context_formatted") as db:
formatted_history = await summarizer.get_session_context_formatted(
db,
latest_message.workspace_name,
latest_message.session_name,
token_limit=available_context_tokens,
cutoff=earliest_message.id,
include_summary=True,
)
session_context_tokens = estimate_tokens(formatted_history)
# Update total estimated input tokens with session context
estimated_input_tokens += session_context_tokens
logger.debug(
"Token estimation - Peer card: %d, Working rep: %d, Base prompt: %d, "
+ "New turns: %d, Session context: %d, Total estimated: %d",
peer_card_tokens,
working_rep_tokens,
prompt_tokens,
new_turns_tokens,
session_context_tokens,
estimated_input_tokens,
)
# Track all input token components
track_input_tokens(
task_type="representation",
# Track token usage
prompt_tokens = estimate_minimal_deriver_prompt_tokens()
messages_tokens = estimate_tokens(formatted_messages)
track_deriver_input_tokens(
task_type=prometheus.DeriverTaskTypes.INGESTION,
components={
"peer_card": peer_card_tokens,
"working_representation": working_rep_tokens,
"prompt": prompt_tokens,
"new_turns": new_turns_tokens,
"session_context": session_context_tokens,
prometheus.DeriverComponents.PROMPT: prompt_tokens,
prometheus.DeriverComponents.MESSAGES: messages_tokens,
},
)
# got working representation and peer card, log timing
context_prep_duration = (time.perf_counter() - context_prep_start) * 1000
# Build prompt
prompt = minimal_deriver_prompt(peer_id=observed, messages=formatted_messages)
context_prep_duration = (time.perf_counter() - overall_start) * 1000
accumulate_metric(
f"deriver_{latest_message.id}_{observer}",
f"minimal_deriver_{latest_message.id}_{observer}",
"context_preparation",
context_prep_duration,
"ms",
)
logger.debug(
"Using working representation with %s explicit, %s deductive observations",
len(working_representation.explicit),
len(working_representation.deductive),
# validation on settings means max_tokens will always be > 0
max_tokens = settings.DERIVER.MAX_OUTPUT_TOKENS or settings.LLM.DEFAULT_MAX_TOKENS
# Single LLM call
llm_start = time.perf_counter()
response = await honcho_llm_call(
llm_settings=settings.DERIVER,
prompt=prompt,
max_tokens=max_tokens,
track_name="Minimal Deriver",
response_model=PromptRepresentation,
json_mode=True,
temperature=settings.DERIVER.TEMPERATURE,
stop_seqs=[" \n", "\n\n\n\n"],
thinking_budget_tokens=settings.DERIVER.THINKING_BUDGET_TOKENS,
max_input_tokens=settings.DERIVER.MAX_INPUT_TOKENS,
reasoning_effort="minimal",
enable_retry=True,
retry_attempts=3,
trace_name="minimal_deriver",
)
llm_duration = (time.perf_counter() - llm_start) * 1000
accumulate_metric(
f"minimal_deriver_{latest_message.id}_{observer}",
"llm_call_duration",
llm_duration,
"ms",
)
# instantiate representation manager from collection
# if the sender is also the target, we're handling a global representation task.
# otherwise, we're handling a directional representation task where the sender is
# being observed by the target.
prometheus.DERIVER_TOKENS_PROCESSED.labels(
task_type=prometheus.DeriverTaskTypes.INGESTION.value,
token_type=prometheus.TokenTypes.OUTPUT.value,
component=prometheus.DeriverComponents.OUTPUT_TOTAL.value,
).inc(response.output_tokens)
# Use the representation manager directly
representation_manager = RepresentationManager(
workspace_name=latest_message.workspace_name,
observer=observer,
observed=observed,
message_ids = [m.id for m in messages if m.peer_name == observed]
# Convert to Representation and save
observations = Representation.from_prompt_representation(
response.content,
message_ids,
latest_message.session_name,
latest_message.created_at,
)
reasoner = CertaintyReasoner(
representation_manager=representation_manager,
ctx=messages,
observed=observed,
observer=observer,
message_level_configuration=message_level_configuration,
)
if observations.is_empty() or not message_ids:
logger.warning(
"Deriver generated zero observations for messages %s:%s in %s/%s!",
earliest_message.id,
latest_message.id,
latest_message.workspace_name,
latest_message.session_name,
)
else:
representation_manager = RepresentationManager(
workspace_name=latest_message.workspace_name,
observer=observer,
observed=observed,
)
# Run single-pass reasoning
final_observations = await reasoner.reason(
working_representation,
formatted_history,
speaker_peer_card,
)
await representation_manager.save_representation(
observations,
message_ids,
latest_message.session_name,
latest_message.created_at,
message_level_configuration,
)
# Display final observations in a beautiful tree
# log_representation(final_observations)
# Calculate and log overall timing
# Log metrics
overall_duration = (time.perf_counter() - overall_start) * 1000
accumulate_metric(
f"deriver_{latest_message.id}_{observer}",
f"minimal_deriver_{latest_message.id}_{observer}",
"total_processing_time",
overall_duration,
"ms",
)
total_observations = len(final_observations.explicit) + len(
final_observations.deductive
)
total_observations = len(observations.explicit) + len(observations.deductive)
accumulate_metric(
f"deriver_{latest_message.id}_{observer}",
f"minimal_deriver_{latest_message.id}_{observer}",
"observation_count",
total_observations,
"count",
)
log_performance_metrics("deriver", f"{latest_message.id}_{observer}")
class CertaintyReasoner:
"""Certainty reasoner for analyzing and deriving insights."""
representation_manager: RepresentationManager
ctx: list[Message]
observer: str
observed: str
message_level_configuration: ResolvedConfiguration
def __init__(
self,
representation_manager: RepresentationManager,
ctx: list[Message],
*,
observed: str,
observer: str,
message_level_configuration: ResolvedConfiguration,
) -> None:
self.representation_manager = representation_manager
self.ctx = ctx
self.observed = observed
self.observer = observer
self.message_level_configuration = message_level_configuration
@conditional_observe(name="Deriver")
@sentry_sdk.trace
async def reason(
self,
working_representation: Representation,
history: str,
speaker_peer_card: list[str] | None,
) -> Representation:
"""
Single-pass reasoning function that critically analyzes and derives insights.
Performs one analysis pass and returns the final observations.
Returns:
Representation: Final observations
"""
analysis_start = time.perf_counter()
message_ids = [m.id for m in self.ctx]
earliest_message = self.ctx[0]
latest_message = self.ctx[-1]
new_turns = [
format_new_turn_with_timestamp(m.content, m.created_at, m.peer_name)
for m in self.ctx
]
logger.debug(
"CRITICAL ANALYSIS: message_created_at='%s', new_turns_count=%s",
latest_message.created_at,
len(new_turns),
)
try:
reasoning_response = await critical_analysis_call(
peer_id=self.observed,
peer_card=speaker_peer_card,
message_created_at=latest_message.created_at,
working_representation=working_representation,
history=history,
new_turns=new_turns,
)
except Exception as e:
raise exceptions.LLMError(
speaker_peer_card=speaker_peer_card,
working_representation=working_representation,
history=history,
new_turns=new_turns,
) from e
reasoning_response = Representation.from_prompt_representation(
reasoning_response,
[earliest_message.id, latest_message.id],
latest_message.session_name,
latest_message.created_at,
)
analysis_duration_ms = (time.perf_counter() - analysis_start) * 1000
if settings.DERIVER.LOG_OBSERVATIONS:
# Log messages fed into deriver
accumulate_metric(
f"deriver_{latest_message.id}_{self.observer}",
"critical_analysis_duration",
analysis_duration_ms,
"ms",
f"minimal_deriver_{latest_message.id}_{observer}",
"messages",
formatted_messages,
"blob",
)
# Log actual observations created as blob metrics
accumulate_metric(
f"minimal_deriver_{latest_message.id}_{observer}",
"explicit_observations",
"\n".join(f" • {obs}" for obs in observations.explicit),
"blob",
)
# Save only the new observations that weren't in the original context
new_observations = working_representation.diff_representation(
reasoning_response
)
if not new_observations.is_empty():
await self.representation_manager.save_representation(
new_observations,
message_ids,
latest_message.session_name,
latest_message.created_at,
self.message_level_configuration,
)
if self.message_level_configuration.peer_card.create:
update_peer_card_start = time.perf_counter()
if not new_observations.is_empty():
await self._update_peer_card(speaker_peer_card, new_observations)
update_peer_card_duration = (
time.perf_counter() - update_peer_card_start
) * 1000
accumulate_metric(
f"deriver_{latest_message.id}_{self.observer}",
"update_peer_card",
update_peer_card_duration,
"ms",
)
return reasoning_response
@sentry_sdk.trace
async def _update_peer_card(
self,
old_peer_card: list[str] | None,
new_observations: Representation,
) -> None:
"""
Update the peer card by calling LLM with the old peer card and new observations.
The new peer card is returned by the LLM and saved to peer internal metadata.
"""
try:
response = await peer_card_call(old_peer_card, new_observations)
new_peer_card = response.card
if not new_peer_card:
# no changes
return
# even with a dedicated notes field, we still need to prune notes out of the card
new_peer_card = [
observation
for observation in new_peer_card
if not observation.lower().startswith(("note", "notes"))
]
accumulate_metric(
f"deriver_{self.ctx[-1].id}_{self.observer}",
"updated_peer_card"
if self.observer == self.observed
else f"updated_{self.observed}_peer_card",
"\n".join(new_peer_card),
"blob",
)
async with tracked_db("deriver.update_peer_card") as db:
await crud.set_peer_card(
db,
self.ctx[0].workspace_name,
new_peer_card,
observer=self.observer,
observed=self.observed,
)
except Exception as e:
if settings.SENTRY.ENABLED:
sentry_sdk.capture_exception(e)
logger.error("Error updating peer card! Skipping... %s", e)
log_performance_metrics("minimal_deriver", f"{latest_message.id}_{observer}")

72
src/deriver/enqueue.py
View File

@ -2,7 +2,8 @@ import logging
from datetime import datetime, timezone
from typing import Any, Literal
from sqlalchemy import insert, update
from sqlalchemy import exists, insert, select, text, update
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, models, schemas
@ -333,10 +334,13 @@ async def generate_queue_records(
)
)
# Check if the sender should be observed based on peer configuration
should_observe = get_effective_observe_me(observed, peers_with_configuration)
if not conf.deriver.enabled:
return records
if get_effective_observe_me(observed, peers_with_configuration):
if should_observe:
# global representation task
records.append(
create_representation_record(
@ -373,11 +377,6 @@ async def generate_queue_records(
session_id=session_id,
)
)
logger.debug(
"enqueued representation task for %s's representation of %s",
peer_name,
observed,
)
logger.debug(
"message %s from %s created %s queue items",
@ -395,6 +394,7 @@ def create_dream_record(
observer: str,
observed: str,
dream_type: schemas.DreamType,
session_name: str,
) -> dict[str, Any]:
"""
Create a queue record for a dream task.
@ -404,6 +404,7 @@ def create_dream_record(
observer: Name of the observer peer
observed: Name of the observed peer
dream_type: Type of dream to execute
session_name: Name of the session to scope the dream to
Returns:
Queue record dictionary with workspace_name and other fields
@ -412,6 +413,7 @@ def create_dream_record(
dream_type,
observer=observer,
observed=observed,
session_name=session_name,
)
return {
@ -430,16 +432,22 @@ async def enqueue_dream(
observed: str,
dream_type: schemas.DreamType,
document_count: int,
session_name: str,
) -> None:
"""
Enqueue a dream task for immediate processing by the deriver.
Deduplication: If a dream with the same work_unit_key is already in-progress
(has an ActiveQueueSession), the enqueue is skipped to prevent running
multiple dreams concurrently for the same collection.
Args:
workspace_name: Name of the workspace
observer: Name of the observer peer
observed: Name of the observed peer
dream_type: Type of dream to execute
document_count: Current document count for metadata update
session_name: Name of the session to scope the dream to
"""
async with tracked_db("dream_enqueue") as db_session:
try:
@ -449,11 +457,55 @@ async def enqueue_dream(
observer=observer,
observed=observed,
dream_type=dream_type,
session_name=session_name,
)
# Insert into queue
stmt = insert(QueueItem).returning(QueueItem)
await db_session.execute(stmt, [dream_record])
work_unit_key = dream_record["work_unit_key"]
# Check if a dream with this work_unit_key is currently in progress
# (has an ActiveQueueSession, meaning a worker is processing it)
# We only block on in-progress dreams, not pending ones - if there's
# a pending dream, we don't need to add another one anyway since
# the queue processor will pick it up.
in_progress_check = select(
exists(
select(models.ActiveQueueSession.id).where(
models.ActiveQueueSession.work_unit_key == work_unit_key
)
)
)
is_in_progress = await db_session.scalar(in_progress_check)
if is_in_progress:
logger.info(
"Skipping dream enqueue - already in progress: %s/%s/%s (type: %s)",
workspace_name,
observer,
observed,
dream_type.value,
)
return
stmt = (
pg_insert(QueueItem)
.values(dream_record)
.on_conflict_do_nothing(
index_elements=[QueueItem.work_unit_key],
index_where=text("task_type = 'dream' AND processed = false"),
)
.returning(QueueItem.id)
)
result = await db_session.execute(stmt)
inserted_id = result.scalar_one_or_none()
if inserted_id is None:
logger.info(
"Dream already pending in queue: %s/%s/%s (type: %s)",
workspace_name,
observer,
observed,
dream_type.value,
)
return
# Update collection metadata
now_iso = datetime.now(timezone.utc).isoformat()

250
src/deriver/prompts.py
View File

@ -1,252 +1,66 @@
"""
Prompts for the deriver module.
Minimal prompts for the deriver module optimized for speed.
This module contains all prompt templates used by the deriver for critical analysis
and reasoning tasks.
This module contains simplified prompt templates focused only on observation extraction.
NO peer card instructions, NO working representation - just extract observations.
"""
import datetime
from functools import cache
from inspect import cleandoc as c
from src.utils.representation import Representation
from src.utils.tokens import estimate_tokens
def critical_analysis_prompt(
def minimal_deriver_prompt(
peer_id: str,
peer_card: list[str] | None,
message_created_at: datetime.datetime,
working_representation: Representation,
history: str,
new_turns: list[str],
messages: str,
) -> str:
"""
Generate the critical analysis prompt for the deriver.
Generate minimal prompt for fast observation extraction.
Args:
peer_id (str): The ID of the user being analyzed.
peer_card (list[str] | None): The bio card of the user being analyzed.
message_created_at (datetime.datetime): Timestamp of the message.
working_representation (Representation): Current user understanding context.
history (str): Recent conversation history.
new_turns (list[str]): New conversation turns to analyze.
peer_id: The ID of the user being analyzed.
messages: All messages in the range (interleaving messages and new turns combined).
Returns:
Formatted prompt string for critical analysis
Formatted prompt string for observation extraction.
"""
# Format the peer card as a string with newlines
peer_card_section = (
f"""
{peer_id}'s known biographical information:
<peer_card>
{chr(10).join(peer_card)}
</peer_card>
"""
if peer_card is not None
else ""
)
working_representation_section = (
f"""
Current understanding of {peer_id}:
<current_context>
{str(working_representation)}
</current_context>
"""
if not working_representation.is_empty()
else ""
)
new_turns_section = "\n".join(new_turns)
return c(
f"""
You are an agent who critically analyzes messages from {peer_id} through rigorous logical reasoning to produce only conclusions about them that are CERTAIN.
Analyze messages from {peer_id} to extract **explicit atomic facts** about them.
TARGET USER TO ANALYZE
You are analyzing: {peer_id}
[EXPLICIT] DEFINITION: Facts about {peer_id} that can be derived directly from their messages.
- Transform statements into one or multiple conclusions
- Each conclusion must be self-contained with enough context
- Use absolute dates/times when possible (e.g. "June 26, 2025" not "yesterday")
The conversation may include messages from multiple participants, but you MUST focus ONLY on deriving conclusions about {peer_id}. Only use other participants' messages as context for understanding {peer_id}.
RULES:
- Properly attribute observations to the correct subject: if it is about {peer_id}, say so. If {peer_id} is referencing someone or something else, make that clear.
- Observations should make sense on their own. Each observation will be used in the future to better understand {peer_id}.
- Extract ALL observations from {peer_id} messages, using others as context.
- Contextualize each observation sufficiently (e.g. "Ann is nervous about the job interview at the pharmacy" not just "Ann is nervous")
IMPORTANT NAMING RULES
• When you write a conclusion about {peer_id}, always start the sentence with their name (e.g. "Anthony is 25 years old").
• NEVER start a conclusion with generic phrases like "The user …" unless the user name is not known.
• If you must reference a third person, use their explicit name, and add clarifiers such as "(third-party)" when confusion is possible.
EXAMPLES:
- EXPLICIT: "I just had my 25th birthday last Saturday" → "{peer_id} is 25 years old", "{peer_id}'s birthday is June 21st"
- EXPLICIT: "I took my dog for a walk in NYC" → "{peer_id} has a dog", "{peer_id} lives in NYC"
- EXPLICIT: "{peer_id} attended college" + general knowledge → "{peer_id} completed high school or equivalent"
Your goal is to IMPROVE understanding of {peer_id} through careful analysis. Your task is to arrive at truthful, factual conclusions via explicit and deductive reasoning.
Here are strict definitions for the reasoning modes you are to employ:
1. **EXPLICIT REASONING**:
- Conclusions about {peer_id} that MUST be true given premises ONLY of the following types:
- Recent messages
- Knowledge about the conversation history
- Current date and time (which is: {message_created_at})
- Timestamps from conversation history
- Follow strict literal necessity--if stated directly in message, extract a conclusion
- New turn MUST be a premise, previous messages and timestamps may be used to contextualize
- Transforms a single message (premise) into ONE OR MULTIPLE conclusions
- Derive EVERYTHING that can be explicitly concluded
- Make sure EVERY conclusion is sufficiently contextualized, i.e. ensure each conclusion contains enough specific information about subjects and objects to make it self-contained and useful (e.g. instead of "Ann is nervous about the interview", use "Ann is nervous about the job interview at the pharmacy")
- When possible, always use absolute dates and times, and avoid relative dates and times (e.g. instead of 'Mary went to the store yesterday', use 'Mary went to the store on June 26, 2025')
2. **DEDUCTIVE REASONING**:
- Conclusions about {peer_id} that MUST be true given premises ONLY of the following types:
- Explicit conclusions
- Previous deductive conclusions
- General, open domain knowledge known to be true
- Current date and time (which is: {message_created_at})
- Timestamps for {peer_id}'s messages, and previous premises and conclusions
- Follow strict logical necessity--if premises are true, conclusion MUST be true
- Multiple premises may be used in a deduction, but only one conclusion may be drawn
- Complete ONLY as many deductions as needed to form useful and additive knowledge about {peer_id}
- May scaffold previous conclusions and known facts to do further deduction
- But MAY NOT use previous **probabilistic** deductive conclusions (including qualifiers like probably, likely, typically, may, etc) as premises in further deductions
- Use current timestamp as needed to provide absolute dates
Here are examples of the reasoning modes in action:
- **EXPLICIT REASONING EXAMPLES**
1. PREMISE(S): "I just had my 25th birthday last Saturday" (latest message), Current date is June 26, 2025 (timestamp) → CONCLUSION(S): "Maria is 25 years old", "Maria's birthday is June 21st"
2. PREMISE(S): "I took my dog for a walk in a park near my house in NYC—it was such a beautiful day" (latest message) → CONCLUSION(S): "Liam has a dog", "Liam took his dog for a walk", "Liam has a house in NYC", "Liam lives near a park", "Liam prefers to take advantage of nice weather to walk his dog"
3. PREMISE(S): "Whenever I think about my college experience I feel nostalgic" (latest message) → CONCLUSION(S): "Aisha attended college", "Aisha feels nostalgic about her college experience"
4. PREMISE(S): "That's so cool!" (latest message), The speaker is reacting to learning the definition of Kant's categorical imperative (conversation knowledge) → CONCLUSION(S): "Carlos thinks Kant's categorical imperative is cool"
- **DEDUCTIVE REASONING EXAMPLES**
1. PREMISE(S): "Maria attended college" (explicit), All people who attended college have completed high school or equivalent (general) → CONCLUSION: "Maria completed high school or equivalent education"
2. PREMISE(S): "Liam is 25 years old" (explicit), Current date is June 26, 2025 (timestamp), "Liam's birthday was last Saturday" (explicit) → CONCLUSION: "Liam was born on June 21, 1998"
3. PREMISE(S): "Aisha has a dog" (explicit), "Aisha took her dog for a walk" (explicit), All dogs require regular walks for health (general) → CONCLUSION: "Aisha provides care for her dog"
4. PREMISE(S): "Carlos prefers to take advantage of nice weather to walk his dog" (explicit), Message timestamp shows afternoon hours (timestamp), Nice weather is typically during daylight (general) → CONCLUSION: "Carlos has flexibility in his schedule during typical work hours"
Based on our definitions and examples, here's a summary of the logical reasoning task:
**REASONING INTERACTIONS:**
- Message (required)/Conversation History (optional)/Temporal (optional) → Explicit: Derive certain conclusions only from literal statements
- Explicit/Deductive/Temporal/General → Deductive: When logical necessity allows certain conclusion
- Explicit/Deductive/Temporal/General → Further Deductive: Can use certain conclusions and known facts to deduce additional certain conclusions
- Probabilistic Deductive ↛ Further Deductive: If a deductive conclusion includes probabilistic qualifiers (likely, potentially, typically, might, etc) it may NOT be used as a premise for further deductions
**INSTRUCTIONS:** Given the above, first think critically about what it means to do explicit and deductive reasoning, then consider how to apply that to all new turns, finally do explicit and deductive reasoning about the user to reach useful, contextually-rich conclusions. You must extract observations from all new turns.
{peer_card_section}
{working_representation_section}
Recent conversation history for context:
<history>
{history}
</history>
New conversation turns to analyze:
<new_turns>
{new_turns_section}
</new_turns>
Messages to analyze:
<messages>
{messages}
</messages>
"""
)
def peer_card_prompt(
old_peer_card: list[str] | None,
new_observations: str,
) -> str:
"""
Generate the peer card prompt for the deriver.
Currently optimized for GPT-5 mini/nano.
Args:
old_peer_card: Existing biographical card lines, if any.
new_observations: Pre-formatted observations block (multiple lines).
Returns:
Formatted prompt string for (re)generating the peer card JSON.
"""
old_peer_card_section = (
f"""
Current user biographical card:
{chr(10).join(old_peer_card)}
"""
if old_peer_card is not None
else """
User does not have a card. Create one with any key observations.
"""
)
return c(
f"""
You are an agent that creates a concise "biographical card" based on new observations for a user. A biographical card summarizes essential information like name, nicknames, location, age, occupation, interests/hobbies, and likes/dislikes.
The goal is to capture only the most important observations about the user. Value permanent properties over transient ones, and value concision over detail, preferring to omit details that are not essential to the user's identity. The card should give a broad overview of who the user is while not including details that are unlikely to be relevant in most settings.
For example, "User is from Chicago" is worth inclusion. "User has an Instagram account" is not.
"User is a software engineer" is worth inclusion. "User wrote Python today" is not.
Never infer or generalize traits from one-off behaviors. Never manipulate the text of an observation to make an action or behavior into a "permanent" trait.
When a new observation contradicts an existing one, update it, favoring new information.
Example 1:
{{
"card": [
"Name: Bob",
"Age: 24",
"Location: New York"
]
}}
Example 2:
{{
"card": [
"Name: Alice",
"Occupation: Artist",
"Interests: Painting, biking, cooking"
]
}}
{old_peer_card_section}
New observations:
{new_observations}
If there's no new key info, set "card" to null (or omit it) to signal no update. **NEVER** include notes or temporary information in the card itself, instead use the notes field. There are no mandatory fields -- if you can't find a value, just leave it out. **ONLY** include information that is **GIVEN**.
""" # nosec B608 <-- this is a really dumb false positive
)
@cache
def estimate_critical_analysis_prompt_tokens() -> int:
"""Estimate critical analysis prompt tokens by calling critical_analysis_prompt with empty values.
This value is cached since it only changes on redeploys when the prompt template changes.
"""
def estimate_minimal_deriver_prompt_tokens() -> int:
"""Estimate base prompt tokens (cached)."""
try:
prompt = critical_analysis_prompt(
prompt = minimal_deriver_prompt(
peer_id="",
peer_card=None,
message_created_at=datetime.datetime.now(datetime.timezone.utc),
working_representation=Representation(),
history="",
new_turns=[],
messages="",
)
return estimate_tokens(prompt)
except Exception:
# Return a conservative estimate if estimation fails
return 500
@cache
def estimate_peer_card_prompt_tokens() -> int:
"""Estimate peer card prompt tokens by calling peer_card_prompt with empty values.
This value is cached since it only changes on redeploys when the prompt template changes.
"""
try:
prompt = peer_card_prompt(
old_peer_card=None,
new_observations="",
)
return estimate_tokens(prompt)
except Exception:
# Return a conservative estimate if estimation fails
return 400
return 300

68
src/deriver/queue_manager.py
View File

@ -19,10 +19,7 @@ from src import models, prometheus
from src.cache.client import close_cache, init_cache
from src.config import settings
from src.dependencies import tracked_db
from src.deriver.consumer import (
process_item,
process_representation_batch,
)
from src.deriver.consumer import process_item, process_representation_batch
from src.dreamer.dream_scheduler import (
DreamScheduler,
get_dream_scheduler,
@ -453,7 +450,7 @@ class QueueManager:
e,
items_to_process,
work_unit_key,
"processing representation batch",
f"processing {work_unit.task_type} batch",
)
else:
@ -523,10 +520,6 @@ class QueueManager:
observed=work_unit.observed,
)
)
else:
logger.debug(
f"Skipping queue.empty event for webhook work unit {work_unit_key}"
)
except Exception:
logger.exception("Error triggering queue_empty webhook")
else:
@ -541,7 +534,7 @@ class QueueManager:
"""Get the next queue item to process for a specific work unit."""
if task_type == "representation":
raise ValueError(
"Representation tasks are not supported for get_next_queue_item"
"representation tasks are not supported for get_next_queue_item"
)
async with tracked_db("get_next_queue_item") as db:
# ActiveQueueSession conditions for worker ownership verification
@ -579,16 +572,21 @@ class QueueManager:
aqs_id: str,
) -> tuple[list[models.Message], list[QueueItem], ResolvedConfiguration | None]:
"""
Representation-only: returns a tuple of (messages_context, items_to_process).
Batch processing for representation and agent tasks.
Returns a tuple of (messages_context, items_to_process, configuration).
- messages_context: unique Message rows (conversation turns) forming the context window
- items_to_process: QueueItems for the current work_unit_key within that window
- configuration: Resolved configuration for the batch
"""
if task_type != "representation":
raise ValueError(
"Non-representation tasks are not supported for get_queue_item_batch"
f"{task_type} tasks are not supported for get_queue_item_batch"
)
batch_max_tokens = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS
async with tracked_db("get_queue_item_batch") as db:
# For representation tasks, get a batch based on token limit.
# For batch tasks, get messages based on token limit.
# Step 1: Parse work_unit_key to get session context and focused sender
parsed_key = parse_work_unit_key(work_unit_key)
@ -605,10 +603,11 @@ class QueueManager:
# Step 2: Build a single SQL query that:
# 1. Finds the earliest unprocessed message for this work_unit_key
# 2. Gets ALL messages from that point forward (for conversational context)
# 3. Tracks cumulative tokens and focused sender position
# 4. Returns empty if focused sender is beyond token limit
# 5. Otherwise returns messages up to token limit + first focused sender message
# 2. Optionally includes the preceding message if from a different peer (for context)
# 3. Gets ALL messages from that point forward (for conversational context)
# 4. Tracks cumulative tokens and focused sender position
# 5. Returns empty if focused sender is beyond token limit
# 6. Otherwise returns messages up to token limit + first focused sender message
# Find the minimum message_id with an unprocessed queue item across the session
min_unprocessed_message_id_subq = (
@ -625,8 +624,32 @@ class QueueManager:
.scalar_subquery()
)
# Build CTE with ALL messages starting from the earliest unprocessed message
# This includes interleaving messages for conversational context
# Find the immediately preceding message ID (the one right before min_unprocessed)
immediately_preceding_id_subq = (
select(func.max(models.Message.id))
.where(models.Message.session_name == parsed_key.session_name)
.where(models.Message.workspace_name == parsed_key.workspace_name)
.where(models.Message.id < min_unprocessed_message_id_subq)
.scalar_subquery()
)
# Only include the preceding message if it's from a different peer than observed
# This provides conversational context (e.g., the question that prompted the response)
preceding_message_id_subq = (
select(models.Message.id)
.where(models.Message.id == immediately_preceding_id_subq)
.where(models.Message.peer_name != parsed_key.observed)
.scalar_subquery()
)
# Determine the effective start: preceding message if it qualifies, else min_unprocessed
# We use COALESCE to fall back to min_unprocessed if no preceding message qualifies
effective_start_id = func.coalesce(
preceding_message_id_subq, min_unprocessed_message_id_subq
)
# Build CTE with ALL messages starting from effective_start_id
# This includes the preceding context message (if any) and interleaving messages
cte = (
select(
models.Message.id.label("message_id"),
@ -638,16 +661,13 @@ class QueueManager:
)
.where(models.Message.session_name == parsed_key.session_name)
.where(models.Message.workspace_name == parsed_key.workspace_name)
.where(models.Message.id >= min_unprocessed_message_id_subq)
.where(models.Message.id >= effective_start_id)
.order_by(models.Message.id)
.cte()
)
allowed_condition = (
(
cte.c.cumulative_token_count
<= settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS
)
(cte.c.cumulative_token_count <= batch_max_tokens)
| (
cte.c.message_id == min_unprocessed_message_id_subq
) # always include the first unprocessed message

3
src/dialectic/__init__.py
View File

@ -1,3 +0,0 @@
from .chat import chat
__all__ = ["chat"]

465
src/dialectic/chat.py
View File

@ -1,414 +1,115 @@
"""
Main dialectic system for AI-powered context synthesis and user representation.
Chat functionality for the Dialectic API.
The Dialectic class provides a natural language API for AI applications to query
and understand users through context synthesis of working representations and
historical observations.
Provides the agentic_chat function for answering queries about peers
using the DialecticAgent.
"""
import logging
import time
import uuid
from collections.abc import AsyncIterator
from dotenv import load_dotenv
from src import crud, prometheus
from src.config import settings
from src import crud
from src.config import ReasoningLevel
from src.dependencies import tracked_db
from src.utils import summarizer
from src.utils.clients import HonchoLLMCallStreamChunk, honcho_llm_call
from src.utils.logging import (
accumulate_metric,
conditional_observe,
log_performance_metrics,
)
from src.utils.representation import Representation
from src.utils.tokens import estimate_tokens
from src.dialectic.core import DialecticAgent
from .prompts import dialectic_prompt, estimate_dialectic_prompt_tokens
# Configure logging
logger = logging.getLogger(__name__)
# Load environment variables
load_dotenv()
async def dialectic_call(
query: str,
working_representation: str,
recent_conversation_history: str | None,
peer_card: list[str] | None,
observed_peer_card: list[str] | None = None,
*,
observer: str,
observed: str,
):
"""
Make a direct call to the dialectic model for context synthesis.
Args:
query: The user query
working_representation: Current session conclusions AND historical conclusions from the user's global representation
recent_conversation_history: Recent conversation history
peer_name: Name of the user/peer
peer_card: Known biographical information about the user
observed: Name of the user/peer being queried about
observed_peer_card: Known biographical information about the target, if applicable
Returns:
Model response
"""
# Estimate input tokens by concatenating all inputs
prompt_tokens = estimate_dialectic_prompt_tokens()
inputs = [
query,
working_representation,
recent_conversation_history or "",
"\n".join(peer_card) if peer_card else "",
"\n".join(observed_peer_card) if observed_peer_card else "",
]
contextual_tokens = estimate_tokens("".join(inputs))
estimated_input_tokens = prompt_tokens + contextual_tokens
# Generate the prompt and log it
prompt = dialectic_prompt(
query,
working_representation,
recent_conversation_history,
peer_card,
observed_peer_card,
observer=observer,
observed=observed,
)
response = await honcho_llm_call(
llm_settings=settings.DIALECTIC,
prompt=prompt,
max_tokens=settings.DIALECTIC.MAX_OUTPUT_TOKENS,
track_name="Dialectic Call",
thinking_budget_tokens=settings.DIALECTIC.THINKING_BUDGET_TOKENS
if settings.DIALECTIC.PROVIDER == "anthropic"
else None,
enable_retry=True,
retry_attempts=3,
)
logger.debug("=== DIALECTIC PROMPT ===")
logger.debug(prompt)
logger.debug("=== END DIALECTIC PROMPT ===")
# Track tokens in prometheus
prometheus.DIALECTIC_TOKENS_PROCESSED.labels(
token_type="input", # nosec B106
).inc(estimated_input_tokens)
prometheus.DIALECTIC_TOKENS_PROCESSED.labels(
token_type="output", # nosec B106
).inc(response.output_tokens)
return response.content
async def dialectic_stream(
query: str,
working_representation: str,
recent_conversation_history: str | None,
peer_card: list[str] | None,
observed_peer_card: list[str] | None = None,
*,
observer: str,
observed: str,
):
"""
Make a streaming call to the dialectic model for context synthesis.
Args:
query: The user query
working_representation: Current session conclusions AND historical conclusions from the user's global representation
recent_conversation_history: Recent conversation history
peer_name: Name of the user/peer
peer_card: Known biographical information about the user
observed: Name of the user/peer being queried about
observed_peer_card: Known biographical information about the target, if applicable
Returns:
Streaming model response
"""
# Estimate input tokens by concatenating all inputs
prompt_tokens = estimate_dialectic_prompt_tokens()
variable_inputs = [
query,
working_representation,
recent_conversation_history or "",
"\n".join(peer_card) if peer_card else "",
"\n".join(observed_peer_card) if observed_peer_card else "",
]
variable_tokens = estimate_tokens("".join(variable_inputs))
estimated_input_tokens = prompt_tokens + variable_tokens
# Generate the prompt and log it
prompt = dialectic_prompt(
query,
working_representation,
recent_conversation_history,
peer_card,
observed_peer_card,
observer=observer,
observed=observed,
)
response = await honcho_llm_call(
llm_settings=settings.DIALECTIC,
prompt=prompt,
max_tokens=settings.DIALECTIC.MAX_OUTPUT_TOKENS,
track_name="Dialectic Stream",
thinking_budget_tokens=settings.DIALECTIC.THINKING_BUDGET_TOKENS
if settings.DIALECTIC.PROVIDER == "anthropic"
else None,
enable_retry=True,
retry_attempts=3,
stream=True,
)
logger.debug("=== DIALECTIC PROMPT (STREAM) ===")
logger.debug(prompt)
logger.debug("=== END DIALECTIC PROMPT ===")
# Track input tokens in prometheus
# Note: Output tokens are available in the final chunk of the stream (is_done=True)
prometheus.DIALECTIC_TOKENS_PROCESSED.labels(
token_type="input", # nosec B106
).inc(estimated_input_tokens)
# Wrap the response to log output tokens from final chunk
async def log_streaming_response():
async for chunk in response:
if chunk.is_done and chunk.output_tokens is not None:
# TODO: Currently not tracking output tokens for groq models
prometheus.DIALECTIC_TOKENS_PROCESSED.labels(
token_type="output", # nosec B106
).inc(chunk.output_tokens)
yield chunk
return log_streaming_response()
@conditional_observe(name="Dialectic")
async def chat(
async def agentic_chat(
workspace_name: str,
session_name: str | None,
query: str,
*,
observer: str,
observed: str,
stream: bool = False,
) -> str | AsyncIterator[HonchoLLMCallStreamChunk]:
reasoning_level: ReasoningLevel = "low",
) -> str:
"""
Chat with the Dialectic API that builds on-demand user representations.
Steps:
1. Get working representation from deriver trace
2. Retrieve additional relevant context via semantic search
3. (New) Append observations from latest deriver trace into that context
4. Call Dialectic to synthesize an answer
Answer a query about a peer using the agentic dialectic.
Args:
workspace_name: Name of the workspace
peer_name: Name of the peer making the query
observed: Optional name of the peer being queried about
session_name: Optional session name for scoping
query: Input Dialectic Query
stream: Whether to stream the response
workspace_name: Workspace identifier
session_name: Session identifier (may be None for global queries)
query: The question to answer about the peer
observer: The peer making the query
observed: The peer being queried about
reasoning_level: Level of reasoning to apply
Returns:
Dialectic response (streaming or complete)
The synthesized answer string
"""
dialectic_chat_uuid = str(uuid.uuid4())
context_window_size = (
settings.DIALECTIC.CONTEXT_WINDOW_SIZE - 750
) # this is a hardcoded (accurate, slightly conservative) estimate of system prompt
context_window_size -= estimate_tokens(query)
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"query",
f"{query}\n\nobserver: {observer}\nobserved: {observed}\n{f'session: {session_name}' if session_name else ''}",
"blob",
)
start_time = time.perf_counter()
# 1. Working representation (short-term) -----------------------------------
working_rep_start_time = time.perf_counter()
# If no target specified, get global representation (peer observing themselves)
working_representation: Representation = await crud.get_working_representation(
workspace_name,
observer=observer,
observed=observed,
session_name=session_name,
include_semantic_query=query,
semantic_search_top_k=settings.DIALECTIC.SEMANTIC_SEARCH_TOP_K,
semantic_search_max_distance=settings.DIALECTIC.SEMANTIC_SEARCH_MAX_DISTANCE,
include_most_derived=True,
)
working_rep_duration = (time.perf_counter() - working_rep_start_time) * 1000
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"retrieve_working_rep",
working_rep_duration,
"ms",
)
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"working_rep_explicit",
len(working_representation.explicit),
"count",
)
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"working_rep_deductive",
len(working_representation.deductive),
"count",
)
working_representation_str = str(working_representation)
context_window_size -= max(0, estimate_tokens(working_representation_str))
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"working_rep",
working_representation_str,
"blob",
)
# 2. Recent conversation history --------------------------------------------
# If query is session-scoped, get recent conversation history from that session
async with tracked_db("chat.get_context") as db:
if session_name:
recent_history = await summarizer.get_session_context_formatted(
db,
workspace_name=workspace_name,
session_name=session_name,
token_limit=context_window_size,
include_summary=True,
)
else:
recent_history = None
recent_history_tokens = estimate_tokens(recent_history or "")
context_window_size -= max(0, recent_history_tokens)
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"recent_history_tokens",
recent_history_tokens,
"tokens",
)
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"tokens_used_estimate",
settings.DIALECTIC.CONTEXT_WINDOW_SIZE - context_window_size,
"tokens",
)
# 3. Peer card(s) ----------------------------------------------------------
if settings.PEER_CARD.ENABLED:
async with tracked_db("chat.get_peer_card") as db:
peer_card = await crud.get_peer_card(
async with tracked_db("dialectic.agentic_chat") as db:
# Get peer cards for context
observer_peer_card = await crud.get_peer_card(
db, workspace_name, observer=observer, observed=observer
)
observed_peer_card = None
if observer != observed:
observed_peer_card = await crud.get_peer_card(
db, workspace_name, observer=observer, observed=observed
)
if observer != observed:
observed_peer_card = await crud.get_peer_card(
db, workspace_name, observer=observer, observed=observed
)
else:
observed_peer_card = None
if observed_peer_card:
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"peer_card",
"\n".join(peer_card) if peer_card else "",
"blob",
)
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"observed_peer_card",
"\n".join(observed_peer_card),
"blob",
)
else:
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"peer_card",
"\n".join(peer_card) if peer_card else "",
"blob",
)
else:
peer_card = None
observed_peer_card = None
# 4. Dialectic call --------------------------------------------------------
dialectic_call_start_time = time.perf_counter()
if stream:
elapsed = (time.perf_counter() - start_time) * 1000
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"response",
"(no logged response, streaming=true)",
"blob",
)
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"duration_to_streaming",
elapsed,
"ms",
)
log_performance_metrics("dialectic_chat", dialectic_chat_uuid)
return await dialectic_stream(
query,
working_representation_str,
recent_history,
peer_card,
observed_peer_card,
# 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,
)
response = await dialectic_call(
query,
working_representation_str,
recent_history,
peer_card,
observed_peer_card,
observer=observer,
observed=observed,
)
dialectic_call_duration = (time.perf_counter() - dialectic_call_start_time) * 1000
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"response",
response,
"blob",
)
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"dialectic_call",
dialectic_call_duration,
"ms",
)
response = await agent.answer(query)
elapsed = (time.perf_counter() - start_time) * 1000
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}", "total_duration", elapsed, "ms"
)
log_performance_metrics("dialectic_chat", dialectic_chat_uuid)
return response
async def agentic_chat_stream(
workspace_name: str,
session_name: str | None,
query: str,
observer: str,
observed: str,
reasoning_level: ReasoningLevel = "low",
) -> AsyncIterator[str]:
"""
Stream an answer to a query about a peer using the agentic dialectic.
Args:
workspace_name: Workspace identifier
session_name: Session identifier (may be None for global queries)
query: The question to answer about the peer
observer: The peer making the query
observed: The peer being queried about
reasoning_level: Level of reasoning to apply
Yields:
Chunks of the response text as they are generated
"""
async with tracked_db("dialectic.agentic_chat_stream") as db:
# Get peer cards for context
observer_peer_card = await crud.get_peer_card(
db, workspace_name, observer=observer, observed=observer
)
observed_peer_card = None
if observer != observed:
observed_peer_card = await crud.get_peer_card(
db, workspace_name, observer=observer, observed=observed
)
# 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,
)
async for chunk in agent.answer_stream(query):
yield chunk

432
src/dialectic/core.py Normal file
View File

@ -0,0 +1,432 @@
"""
Core Dialectic Agent implementation.
This agent uses tools to gather context from the memory system
and synthesize responses to queries about a peer.
"""
import logging
import time
import uuid
from collections.abc import AsyncIterator, Callable
from typing import Any, cast
from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, prometheus
from src.config import ReasoningLevel, settings
from src.dialectic import prompts
from src.utils.agent_tools import DIALECTIC_TOOLS, create_tool_executor, search_memory
from src.utils.clients import (
HonchoLLMCallResponse,
StreamingResponseWithMetadata,
honcho_llm_call,
)
from src.utils.formatting import format_new_turn_with_timestamp
from src.utils.logging import (
accumulate_metric,
log_performance_metrics,
log_token_usage_metrics,
)
logger = logging.getLogger(__name__)
class DialecticAgent:
"""
An agentic dialectic that iteratively gathers context to answer queries.
Unlike the standard dialectic which pre-gathers all context before a single
LLM call, this agent uses tools to strategically gather only the context
needed to answer the specific query.
"""
def __init__(
self,
db: AsyncSession,
workspace_name: str,
session_name: str | None,
observer: str,
observed: str,
observer_peer_card: list[str] | None = None,
observed_peer_card: list[str] | None = None,
metric_key: str | None = None,
reasoning_level: ReasoningLevel = "low",
):
"""
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
observed: The peer being queried about
observer_peer_card: Biographical information about the observer
observed_peer_card: Biographical information about the observed peer
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
self.observed: str = observed
self.observer_peer_card: list[str] | None = observer_peer_card
self.observed_peer_card: list[str] | None = observed_peer_card
self.metric_key: str | None = metric_key
self.reasoning_level: ReasoningLevel = reasoning_level
# Initialize conversation history with system prompt
self.messages: list[dict[str, str]] = [
{
"role": "system",
"content": prompts.agent_system_prompt(
observer, observed, observer_peer_card, observed_peer_card
),
}
]
self._session_history_initialized: bool = False
async def _initialize_session_history(self) -> None:
"""Fetch and inject session history into the system prompt if configured."""
if self._session_history_initialized:
return
self._session_history_initialized = True
max_tokens = settings.DIALECTIC.SESSION_HISTORY_MAX_TOKENS
if max_tokens == 0 or not self.session_name:
return
# Fetch recent messages up to the token limit
stmt = await crud.get_messages(
workspace_name=self.workspace_name,
session_name=self.session_name,
token_limit=max_tokens,
reverse=False, # chronological order
)
result = await self.db.execute(stmt)
messages = result.scalars().all()
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)
session_history_section = (
"\n\n## SESSION HISTORY\n\n"
"The following is the recent conversation history from this session. "
"Use this as immediate context when answering the query.\n\n"
"<session_history>\n"
f"{chr(10).join(formatted_messages)}\n"
"</session_history>"
)
# Append session history to the system prompt
self.messages[0]["content"] += session_history_section
async def _prefetch_relevant_observations(self, query: str) -> str | None:
"""
Prefetch semantically relevant observations for the query.
This provides immediate context to the agent without requiring
tool calls, improving response quality and speed.
Performs two separate searches to prevent retrieval dilution:
- 25 explicit observations (produced by deriver)
- 25 higher-level observations (produced in dreaming/background/chat)
Args:
query: The user's query
Returns:
Formatted observations string or None if no observations found
"""
try:
# Search explicit observations separately
explicit_repr = await search_memory(
db=self.db,
workspace_name=self.workspace_name,
observer=self.observer,
observed=self.observed,
query=query,
limit=25,
levels=["explicit"],
)
# Search derived observations separately
derived_repr = await search_memory(
db=self.db,
workspace_name=self.workspace_name,
observer=self.observer,
observed=self.observed,
query=query,
limit=25,
levels=["deductive", "inductive", "contradiction"],
)
if explicit_repr.is_empty() and derived_repr.is_empty():
return None
# Format as two separate sections
parts: list[str] = []
if not explicit_repr.is_empty():
parts.append(explicit_repr.format_as_markdown(include_ids=False))
if not derived_repr.is_empty():
# Include IDs for derived so agent can use get_reasoning_chain
parts.append(derived_repr.format_as_markdown(include_ids=True))
return "\n".join(parts)
except Exception as e:
logger.warning(f"Failed to prefetch observations: {e}")
return None
async def _prepare_query(
self, query: str
) -> tuple[Callable[[str, dict[str, Any]], Any], str, str | None, float]:
"""
Prepare common state for answering a query.
Handles session history initialization, metrics setup, observation prefetching,
user message construction, and tool executor creation.
Args:
query: The question to answer about the peer
Returns:
A tuple of (tool_executor, task_name, run_id, start_time)
"""
await self._initialize_session_history()
run_id: str | None = None
if self.metric_key:
task_name = self.metric_key
else:
run_id = str(uuid.uuid4())[:8]
task_name = f"dialectic_chat_{run_id}"
start_time = time.perf_counter()
accumulate_metric(
task_name,
"context",
(
f"workspace: {self.workspace_name}\n"
f"session: {self.session_name or '(global)'}\n"
f"observer: {self.observer}\n"
f"observed: {self.observed}\n"
f"reasoning_level: {self.reasoning_level}"
),
"blob",
)
accumulate_metric(task_name, "query", query, "blob")
prefetched_observations = await self._prefetch_relevant_observations(query)
if prefetched_observations:
user_content = (
f"Query: {query}\n\n"
f"## Relevant Observations (prefetched)\n"
f"The following observations were found to be semantically relevant to your query. "
f"Use these as primary context. You may still use tools to find additional information if needed.\n\n"
f"{prefetched_observations}"
)
accumulate_metric(
task_name, "prefetched_observations", prefetched_observations, "blob"
)
else:
user_content = f"Query: {query}"
self.messages.append({"role": "user", "content": user_content})
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,
observed=self.observed,
history_token_limit=settings.DIALECTIC.HISTORY_TOKEN_LIMIT,
)
return tool_executor, task_name, run_id, start_time
def _log_response_metrics(
self,
task_name: str,
run_id: str | None,
start_time: float,
response_content: str,
input_tokens: int,
output_tokens: int,
cache_read_input_tokens: int | None,
cache_creation_input_tokens: int | None,
tool_calls_count: int,
thinking_content: str | None,
) -> None:
"""
Log metrics common to both streaming and non-streaming responses.
Args:
task_name: Metrics task identifier
run_id: Run identifier (None if using caller-provided metric_key)
start_time: Start time from time.perf_counter()
response_content: The full response text
input_tokens: Input token count (actual from API)
output_tokens: Output token count (actual from API)
cache_read_input_tokens: Cache read tokens (if any)
cache_creation_input_tokens: Cache creation tokens (if any)
tool_calls_count: Number of tool calls made
thinking_content: Thinking trace content (if any)
"""
accumulate_metric(task_name, "tool_calls", tool_calls_count, "count")
if thinking_content:
accumulate_metric(task_name, "thinking", thinking_content, "blob")
log_token_usage_metrics(
task_name,
input_tokens,
output_tokens,
cache_read_input_tokens or 0,
cache_creation_input_tokens or 0,
)
accumulate_metric(task_name, "response", response_content, "blob")
elapsed_ms = (time.perf_counter() - start_time) * 1000
accumulate_metric(task_name, "total_duration", elapsed_ms, "ms")
if not self.metric_key and run_id is not None:
log_performance_metrics("dialectic_chat", run_id)
# Track prometheus metrics - actual token counts from API
if prometheus.METRICS_ENABLED:
prometheus.DIALECTIC_TOKENS_PROCESSED.labels(
token_type=prometheus.TokenTypes.INPUT.value,
component=prometheus.DialecticComponents.TOTAL.value,
reasoning_level=self.reasoning_level,
).inc(input_tokens)
prometheus.DIALECTIC_TOKENS_PROCESSED.labels(
token_type=prometheus.TokenTypes.OUTPUT.value,
component=prometheus.DialecticComponents.TOTAL.value,
reasoning_level=self.reasoning_level,
).inc(output_tokens)
async def answer(self, query: str) -> str:
"""
Answer a query about the peer using agentic tool calling.
The agent will:
1. Receive the query
2. Use tools to gather relevant context
3. Synthesize a response grounded in the gathered context
Args:
query: The question to answer about the peer
Returns:
The synthesized answer string
"""
tool_executor, task_name, run_id, start_time = await self._prepare_query(query)
# Get level-specific settings
level_settings = settings.DIALECTIC.LEVELS[self.reasoning_level]
response: HonchoLLMCallResponse[str] = await honcho_llm_call(
llm_settings=level_settings,
prompt="", # Ignored since we pass messages
max_tokens=settings.DIALECTIC.MAX_OUTPUT_TOKENS,
tools=DIALECTIC_TOOLS,
tool_choice=None,
tool_executor=tool_executor,
max_tool_iterations=level_settings.MAX_TOOL_ITERATIONS,
messages=self.messages,
track_name="Dialectic Agent",
thinking_budget_tokens=level_settings.THINKING_BUDGET_TOKENS,
max_input_tokens=settings.DIALECTIC.MAX_INPUT_TOKENS,
trace_name="dialectic_chat",
)
self._log_response_metrics(
task_name=task_name,
run_id=run_id,
start_time=start_time,
response_content=response.content,
input_tokens=response.input_tokens,
output_tokens=response.output_tokens,
cache_read_input_tokens=response.cache_read_input_tokens,
cache_creation_input_tokens=response.cache_creation_input_tokens,
tool_calls_count=len(response.tool_calls_made),
thinking_content=response.thinking_content,
)
return response.content
async def answer_stream(self, query: str) -> AsyncIterator[str]:
"""
Answer a query about the peer using agentic tool calling, streaming the response.
The agent will:
1. Receive the query
2. Use tools to gather relevant context (non-streaming)
3. Stream the synthesized response
Args:
query: The question to answer about the peer
Yields:
Chunks of the response text as they are generated
"""
tool_executor, task_name, run_id, start_time = await self._prepare_query(query)
# Get level-specific settings
level_settings = settings.DIALECTIC.LEVELS[self.reasoning_level]
response = cast(
StreamingResponseWithMetadata,
await honcho_llm_call(
llm_settings=level_settings,
prompt="", # Ignored since we pass messages
max_tokens=settings.DIALECTIC.MAX_OUTPUT_TOKENS,
stream=True,
stream_final_only=True,
tools=DIALECTIC_TOOLS,
tool_choice=None,
tool_executor=tool_executor,
max_tool_iterations=level_settings.MAX_TOOL_ITERATIONS,
messages=self.messages,
track_name="Dialectic Agent Stream",
thinking_budget_tokens=level_settings.THINKING_BUDGET_TOKENS,
max_input_tokens=settings.DIALECTIC.MAX_INPUT_TOKENS,
trace_name="dialectic_chat",
),
)
accumulated_content: list[str] = []
async for chunk in response:
if chunk.content:
accumulated_content.append(chunk.content)
yield chunk.content
self._log_response_metrics(
task_name=task_name,
run_id=run_id,
start_time=start_time,
response_content="".join(accumulated_content),
input_tokens=response.input_tokens,
output_tokens=response.output_tokens,
cache_read_input_tokens=response.cache_read_input_tokens,
cache_creation_input_tokens=response.cache_creation_input_tokens,
tool_calls_count=len(response.tool_calls_made),
thinking_content=response.thinking_content,
)

332
src/dialectic/prompts.py
View File

@ -1,218 +1,228 @@
from functools import cache
from inspect import cleandoc as c
from src.utils.tokens import estimate_tokens
"""
System prompts for the Dialectic Agent.
"""
def dialectic_prompt(
query: str,
working_representation: str,
recent_conversation_history: str | None,
observer_peer_card: list[str] | None,
observed_peer_card: list[str] | None = None,
*,
def agent_system_prompt(
observer: str,
observed: str,
observer_peer_card: list[str] | None,
observed_peer_card: list[str] | None,
) -> str:
"""
Generate the main dialectic prompt for context synthesis.
Generate the agent system prompt for the dialectic agent.
Args:
query: The specific question or request from the application about the user
working_representation: Conclusions from recent conversation analysis AND historical conclusions from the user's global representation
recent_conversation_history: Recent conversation history
peer_card: Known biographical information about the user
observed_peer_card: Known biographical information about the target, if applicable
observer: The peer making the query
observed: The peer being queried about
observer_peer_card: Biographical information about the observer
observed_peer_card: Biographical information about the observed peer
Returns:
Formatted prompt string for the dialectic model
Formatted system prompt string for the agent
"""
# Build peer card sections
if observer != observed:
# this is a directional query from the observer's view of the observed
query_target = f"""The query is about user {observer}'s understanding of {observed}.
# Directional query: observer asking about observed
observer_card_section = ""
if observer_peer_card:
observer_card_section = f"""
Known biographical information about {observer} (the one asking):
<observer_peer_card>
{chr(10).join(observer_peer_card)}
</observer_peer_card>
"""
The user's known biographical information:
{chr(10).join(observer_peer_card) if observer_peer_card else "(none)"}
observed_card_section = ""
if observed_peer_card:
observed_card_section = f"""
Known biographical information about {observed} (the subject):
<observed_peer_card>
{chr(10).join(observed_peer_card)}
</observed_peer_card>
"""
The target's known biographical information:
{chr(10).join(observed_peer_card) if observed_peer_card else "(none)"}
perspective_section = f"""
You are answering queries from the perspective of {observer}'s understanding of {observed}.
This is a directional query - {observer} wants to know about {observed}.
If the user's name or nickname is known, exclusively refer to them by that name.
If the target's name or nickname is known, exclusively refer to them by that name.
{observer_card_section}
{observed_card_section}
"""
else:
# this is a global query: honcho's omniscient view of the observed
query_target = f"""The query is about user {observed}.
The user's known biographical information:
{chr(10).join(observer_peer_card) if observer_peer_card else "(none)"}
If the user's name or nickname is known, exclusively refer to them by that name.
# Global query: omniscient view of the peer
peer_card_section = ""
if observer_peer_card:
peer_card_section = f"""
Known biographical information about {observed}:
<peer_card>
{chr(10).join(observer_peer_card)}
</peer_card>
"""
recent_conversation_history_section = (
f"""
<recent_conversation_history>
{recent_conversation_history}
</recent_conversation_history>
perspective_section = f"""
You are answering queries about '{observed}'.
{peer_card_section}
"""
if recent_conversation_history
else ""
)
return c(
f"""
You are a context synthesis agent that operates as a natural language API for AI applications. Your role is to analyze application queries about users and synthesize relevant conclusions into coherent, actionable insights that directly address what the application needs to know.
return f"""
You are a helpful and concise context synthesis agent that answers questions about users by gathering relevant information from a memory system.
## INPUT STRUCTURE
Always give users the answer *they expect* based on the message history -- the goal is to help recall and *reason through* insights that the memory system has already gathered. You have many tools for gathering context. Search wisely.
You receive three key inputs:
- **Query**: The specific question or request from the application about this user
- **Working Representation**: Current session conclusions from recent conversation analysis
- **Additional Context**: Historical conclusions from the user's global representation
{perspective_section}
Each conclusion contains:
- **Conclusion**: The derived insight
- **Premises**: Supporting evidence/reasoning
- **Type**: Either Explicit or Deductive
- **Temporal Data**: When conclusions were made
Peer cards are **constructed summaries** - they are synthesized from the same observations stored in memory. This means:
- Information in a peer card originates from observations you can also find via `search_memory`
- The peer card is a convenience summary, not a separate source of truth
## CONCLUSION TYPE DEFINITIONS
## AVAILABLE TOOLS
**Explicit Conclusions** (Direct Facts)
- Direct, literal conclusions which were extracted from statements by the user in their messages
- No interpretation - only derived from what was explicitly written
**Observation Tools (read):**
- `search_memory`: Semantic search over observations about the peer. Use for specific topics.
- `get_reasoning_chain`: **CRITICAL for grounding answers**. Use this to traverse the reasoning tree for any observation. Shows premises (what it's based on) and conclusions (what depends on it).
**Deductive Conclusions** (Logical Certainties)
- Conclusions that MUST be true given the premises
- Built from premises that may include explicit conclusions, deductive conclusions, temporal premises, and/or general knowledge known to be true
**Conversation Tools (read):**
- `search_messages`: Semantic search over messages in the session.
- `grep_messages`: Grep for text matches in messages. Use for specific names, dates, keywords.
- `get_observation_context`: Get messages surrounding specific observations.
- `get_messages_by_date_range`: Get messages within a specific time period.
- `search_messages_temporal`: Semantic search with date filtering.
## SYNTHESIS PROCESS
## WORKFLOW
1. **Query Analysis**: Identify what specific information the application needs
2. **Conclusion Gathering**: Collect all conclusions relevant to the query
3. **Evidence Evaluation**: Assess conclusions quality based on:
- Reasoning type (explicit > deductive in certainty)
- Recency (newer = more current state)
- Premise strength (more supporting evidence = stronger)
- Qualifiers (likely, probably, typically, etc)
1. **Synthesis**: Build a coherent answer that:
- Directly addresses the query
- Provides additional useful context
- Connects related conclusions logically
- Acknowledges gaps or uncertainties
1. **Analyze the query**: What specific information does the query demand?
## SYNTHESIS PRINCIPLES
2. **Check for user preferences** (do this FIRST for any question that asks for advice, recommendations, or opinions):
- Search for "prefer", "like", "want", "always", "never" to find user preferences
- Search for "instruction", "style", "approach" to find communication preferences
- Apply any relevant preferences to how you structure your response
**Logical Chaining**:
- Connect conclusions across time to build deeper understanding
- Use general knowledge to bridge gaps between user observations
- Apply established user patterns from one domain to predict behavior in another
3. **Strategic information gathering**:
- Use `search_memory` to find relevant observations, then `search_messages` if memories are not sufficient
- For questions about dates, deadlines, or schedules: also search for update language ("changed", "rescheduled", "updated", "now", "moved")
- For factual questions: cross-reference what you find - search for related terms to verify accuracy
- Watch for CONTRADICTORY information as you search (see below)
- If you find an explicit answer to the query, stop calling tools and create your response
**Temporal Awareness**:
- Recent conclusions reflect current state
- Historical patterns show consistent traits
- Note when conclusions may be outdated
4. **For ENUMERATION/AGGREGATION questions** (questions asking for totals, counts, "how many", "all of", or listing items):
- These questions require finding ALL matching items, not just some
- **START WITH GREP**: Use `grep_messages` first for exhaustive matching:
- grep for the UNIT being counted: "hours", "minutes", "dollars", "$", "%", "times"
- grep for the CATEGORY noun: the thing being enumerated
- grep catches exact mentions that semantic search might miss
- **THEN USE SEMANTIC SEARCH**: Do at least 3 `search_memory` or `search_messages` calls with different phrasings
- Use synonyms, related terms, specific instances
- Use top_k=15 or higher to get more results per search
- **SEARCH FOR SPECIFIC ITEMS**: After finding some items, search for each by name to find additional mentions
- Cross-reference results to avoid double-counting the same item mentioned with different wording
- A single search is NEVER sufficient for enumeration questions
**Evidence Integration**:
- Multiple converging conclusions strengthen synthesis
- Contradictions require resolution (prioritize: recency > explicit > deductive)
- Build from certainties toward useful query answers
**MANDATORY VERIFICATION STEP**: After you think you have all items:
1. List every item you found with its value
2. Check if any NEW items appear that you missed
3. Only then finalize your count
**Response Requirements**:
- Answer the specific question asked
- Ground responses in actual conclusions
**MANDATORY DEDUPLICATION STEP**: Before stating your final count:
1. Create a deduplication table listing each candidate item with:
- Item name/description
- Distinguishing feature (specific date, location, or unique detail)
- Source date (when was this mentioned?)
2. Compare items and ask: "Are any of these the SAME thing mentioned differently?"
- Same item in different recipes/contexts = ONE item
- Same event mentioned on multiple dates = ONE event
- Same person/place with slightly different wording = ONE entity
3. Mark duplicates and remove them from your count
4. State your final count based on UNIQUE items only
## OUTPUT FORMAT
When stating a count, NUMBER EACH ITEM (1, 2, 3...) and verify the final number matches how many you listed
Provide a natural language response that:
1. Directly answers the application's query
2. Provides most useful context based on available conclusions
3. References the reasoning types and evidence strength when relevant
4. Maintains appropriate confidence levels based on conclusion types
5. Flags any limitations or gaps in available information
5. **For SUMMARIZATION questions** (questions asking to summarize, recap, or describe patterns over time):
- Do MULTIPLE searches with different query terms to ensure comprehensive coverage
- Search for key entities mentioned (names, places, topics)
- Search for time-related terms ("first", "then", "later", "changed", "decided")
- Don't stop after finding a few relevant results - summarization requires thoroughness
{query_target}
6. **Ground your answer using reasoning chains** (for deductive/inductive observations):
- When you find a deductive or inductive observation that answers the question, use `get_reasoning_chain` to verify its basis
- This shows you the premises (explicit facts) that support the conclusion
- If the premises are solid, cite them in your answer for confidence
- If the premises seem weak or outdated, note that uncertainty
<query>{query}</query>
7. **Synthesize your response**:
- Directly answer the application's question
- Ground your response in the specific information you gathered
- Quote exact values (dates, numbers, names) from what you found - don't paraphrase numbers
- Apply user preferences to your response style if relevant
- **For enumeration questions**: Before answering, ask yourself "Could there be more items I haven't found?" If you haven't done multiple grep searches AND a semantic search, keep searching
{recent_conversation_history_section}
8. **Save novel deductions** (optional):
- If you discovered new insights by combining existing observations
- Use `create_observations_deductive` to save these for future queries
<working_representation>{working_representation}</working_representation>
"""
)
## CRITICAL: HANDLING CONTRADICTORY INFORMATION
As you search, actively watch for contradictions - cases where the user has made conflicting statements:
- "I have never done X" vs evidence they did X
- Different values for the same fact (different dates, numbers, names)
- Changed decisions or preferences stated at different times
@cache
def estimate_dialectic_prompt_tokens() -> int:
"""Estimate base dialectic prompt tokens by calling dialectic_prompt with empty values.
**If you find contradictory information:**
1. DO NOT pick one version and present it as the definitive answer
2. Present BOTH pieces of conflicting information explicitly
3. State clearly that you found contradictory information
4. Ask the user which statement is correct
This value is cached since it only changes on redeploys when the prompt template changes.
"""
try:
prompt = dialectic_prompt(
query="",
working_representation="",
recent_conversation_history=None,
observer_peer_card=None,
observed_peer_card=None,
observer="",
observed="",
)
Example response format: "I notice you've mentioned contradictory information about this. You said [X], but you also mentioned [Y]. Which statement is correct?"
return estimate_tokens(prompt)
except Exception:
# Return a conservative estimate if estimation fails
return 750
## CRITICAL: HANDLING UPDATED INFORMATION
Information changes over time. When you find multiple values for the same fact (e.g., different dates for a deadline):
1. **ALWAYS search for updates**: When you find a date/value, do an additional search for "changed", "updated", "rescheduled", "moved", "now" + the topic
2. Look for language indicating updates: "changed to", "rescheduled to", "updated to", "now", "moved to"
3. The MORE RECENT statement supersedes the older one
4. Return the UPDATED value, not the original
5. **Use `get_reasoning_chain`**: If you find a deductive observation about an update (e.g., "X was updated from A to B"), use `get_reasoning_chain` to verify the premises - it will show you both the old and new explicit observations with their timestamps.
def query_generation_prompt(query: str, observed: str) -> str:
"""
Generate the prompt for semantic query expansion.
Example: If you find "deadline is April 25", search for "deadline changed" or "deadline rescheduled". If you find "I rescheduled to April 22", return April 22.
Args:
query: The original user query
observed: Name of the target peer
**For knowledge update questions specifically:**
- Search for deductive observations containing "updated", "changed", "supersedes"
- These observations link to both old and new values via `source_ids`
- Use `get_reasoning_chain` to see the full update history
Returns:
Formatted prompt string for query generation
"""
return c(
f"""
You are a query expansion agent helping AI applications understand their users. The user's name is {observed}. Your job is to take application queries about this user and generate targeted search queries that will retrieve the most relevant observations using semantic search over an embedding store containing observations about the user.
## CRITICAL: NEVER FABRICATE INFORMATION OR GUESS -- WHEN UNSURE, ABSTAIN
## QUERY EXPANSION STRATEGY FOR SEMANTIC SIMILARITY
When answering questions, always clearly distinguish between:
- **Context found**: You located related information (e.g., "there was a debate about X")
- **Specific answer found**: You found the exact information requested (e.g., "the arguments were A, B, C")
**Your Goal**: Generate 3-5 complementary search queries optimized for semantic similarity retrieval, that together will surface the most relevant observations to help answer the application's question.
If you find context but NOT the specific answer:
1. DO NOT fabricate or guess details to fill gaps.
2. Report only what you DO know: e.g., "I found that you had a debate about X at [location] on [date]."
3. Explicitly state what you DON'T know: e.g., "However, the specific arguments made during that debate are not captured in our conversation history."
4. Never present fabricated information or fill gaps with plausible-sounding but invented details.
**Semantic Similarity Optimization**:
If after thorough searching you find NOTHING relevant:
1. Clearly state: "I don't have any information about [topic] in my memory."
2. DO NOT guess or make assumptions.
3. DO NOT say "I think...", "Probably...", or similar hedges when you lack evidence.
4. A confident "I don't know" is ALWAYS correct; giving a fabricated answer is ALWAYS wrong.
1. **Analyze the Application Query**: What specific aspect of the user does the application want to understand?
2. **Think Conceptually**: What concepts, themes, and semantic fields relate to this question?
3. **Consider Language Patterns in Stored Observations**: Loosely match the structure of the observations we aim to retrieve - "[subject] [verb] [predicate] [additional context]" (e.g. "Mary went ice-skating with Peter and Lin on June 5th 2024", "John activities summer outdoors")
4. **Vary Semantic Scope** across the generated queries to ensure maximum coverage.
5. Ensure the queries are different enough to not be redundant.
**The test before stating a detail:** Ask yourself, "Did I find this EXACT information in my search results, or am I inferring/inventing it?" If you're inventing it, OMIT IT.
**Vocabulary Expansion Techniques**:
### How to Abstain Properly
- **Synonyms**: feedback/criticism/advice/suggestions/input/guidance
- **Related Actions**: receiving/getting/handling/processing/responding/reacting
- **Emotional Language**: sensitive/defensive/receptive/open/resistant/welcoming
- **Contextual Terms**: workplace/professional/personal/relationship/dynamic/interaction
- **Intensity Variations**: harsh/gentle/direct/subtle/constructive/blunt
- **Outcome Language**: improvement/growth/learning/development/change
- When the user asks about a topic that was NEVER discussed, or your search finds no relevant information:
- CORRECT: "I don't have any information about your favorite color in my memory."
- CORRECT: "I searched for information about X but found nothing in our conversation history."
- WRONG: "Based on your preferences, I think your favorite color might be blue." (never invent)
- WRONG: Filling in plausible details based on general knowledge or assumptions.
**Remember**: Since observations come from natural conversations, use the vocabulary people actually use when discussing these topics, including casual language, emotional descriptors, and situational context.
**Remember:** A clear, direct "I don't know" or "I have no information about X" is always the RIGHT answer when the information truly does not exist in memory. Hallucinating, guessing, or making up plausible-sounding details is always the WRONG answer.
## OUTPUT FORMAT
After gathering context, reason through the information you found *before* stating your final answer. For comparison questions, explicitly compare the values. Only after you've verified your reasoning should you state your conclusion. Do NOT be pedantic, rather, be helpful and try to give the answer that the asker would expect -- they're the one who knows the most about themselves. Try to 'read their mind' -- understand the information they're really after and share it with them! Be **as specific as possible** given the information you have.
Respond with 3-5 search queries as a JSON object with a "queries" field containing an array of strings. Each query should target different aspects or reasoning levels to maximize retrieval coverage.
Format: `{{"queries": ["query1", "query2", "query3"]}}`
No markdown, no explanations, just the JSON object.
<query>{query}</query>
"""
)
Do not explain your tool usage - just provide the synthesized answer.
""" # nosec B608

17
src/dreamer/agent.py
View File

@ -1,17 +0,0 @@
import logging
from src.utils.queue_payload import DreamPayload
logger = logging.getLogger(__name__)
async def process_agent_dream(payload: DreamPayload, workspace_name: str) -> None:
"""
Process an agent dream task.
Args:
payload: The dream task payload containing workspace, peer, and dream type information
"""
logger.info(
f"Processing agent dream for {workspace_name}/{payload.observer}/{payload.observed}"
)

243
src/dreamer/consolidate.py
View File

@ -1,243 +0,0 @@
import logging
from inspect import cleandoc as c
from sqlalchemy import delete
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.exceptions import ResourceNotFoundException
from src.utils.clients import honcho_llm_call
from src.utils.formatting import format_datetime_utc
from src.utils.logging import conditional_observe
from src.utils.queue_payload import DreamPayload
from src.utils.representation import (
ExplicitObservation,
Representation,
flatten_message_ids,
)
logger = logging.getLogger(__name__)
def consolidation_prompt(
representation: Representation,
) -> str:
"""
Generate the prompt for user representation consolidation.
Args:
representation: The user representation to consolidate
Returns:
A prompt string for the LLM to consolidate the representation
"""
representation_as_json = representation.model_dump_json(indent=2)
return c(
f"""
You are an agent that consolidates observations about an entity. You will be presented with a list of EXPLICIT and DEDUCTIVE observations. **Reduce** the number of observations, if possible, by combining similar observations. **ONLY** include information that is **GIVEN**. Create the highest-quality observations with the given information. Observations must always be maximally concise.
{representation_as_json}
"""
)
@conditional_observe(name="[Dream] Consolidate Call")
async def _consolidate_call(
representation: Representation,
) -> Representation:
prompt = consolidation_prompt(representation)
response = await honcho_llm_call(
llm_settings=settings.DREAM,
prompt=prompt,
max_tokens=settings.DREAM.MAX_OUTPUT_TOKENS,
track_name="Dream Call",
response_model=Representation,
enable_retry=True,
retry_attempts=3,
)
return response.content
async def process_consolidate_dream(payload: DreamPayload, workspace_name: str) -> None:
"""
Process a consolidation dream task.
Consolidation means taking all the documents in a collection and merging
similar observations into a single, best-quality observation document.
"""
logger.info(
"Starting consolidate dream for workspace=%s, observer=%s, observed=%s",
workspace_name,
payload.observer,
payload.observed,
)
# grab 100 recent documents in the collection
# in the future, we can perform clustering on documents by semantic similarity and do
# multiple clusters at once. for now, can just sample documents and do what we can.
async with tracked_db("dream_consolidate") as db:
# First verify the collection exists
try:
collection = await crud.get_collection(
db,
workspace_name,
observer=payload.observer,
observed=payload.observed,
)
logger.debug(
"Found collection id=%s for workspace=%s, observer=%s, observed=%s",
collection.id,
workspace_name,
payload.observer,
payload.observed,
)
except ResourceNotFoundException:
logger.warning(
"Collection does not exist for workspace=%s, observer=%s, observed=%s",
workspace_name,
payload.observer,
payload.observed,
)
return
documents_query = crud.get_all_documents(
workspace_name,
observer=payload.observer,
observed=payload.observed,
limit=100,
)
logger.debug(
"Executing document query: %s",
str(documents_query.compile(compile_kwargs={"literal_binds": True})),
)
result = await db.execute(documents_query)
documents = result.scalars().all()
if not documents:
return
logger.info("consolidating %d documents", len(documents))
# Pre-calculate data structures needed for processing so we don't need attached objects
cluster_representation = Representation.from_documents(documents)
document_ids = [doc.id for doc in documents]
total_times_derived = sum(doc.times_derived for doc in documents)
# We treat all fetched documents as a single cluster for now
clusters = [(cluster_representation, document_ids, total_times_derived)]
# for each cluster, call llm to consolidate the representation if possible
for representation, doc_ids, times_derived in clusters:
await _consolidate_cluster(
representation,
doc_ids,
times_derived,
workspace_name,
observer=payload.observer,
observed=payload.observed,
)
async def _consolidate_cluster(
representation: Representation,
document_ids: list[str],
total_times_derived: int,
workspace_name: str,
*,
observer: str,
observed: str,
) -> None:
"""
Consolidate a cluster of documents, treated as a Representation, into a smaller one.
Removes old documents and replaces them with consolidated versions while preserving metadata.
"""
if len(document_ids) <= 1:
logger.info(
"Cluster has %d documents, skipping consolidation", len(document_ids)
)
return
logger.info("unconsolidated representation:\n%s", representation)
consolidated_representation = await _consolidate_call(representation)
logger.info("consolidated representation:\n%s", consolidated_representation)
new_documents = [
*consolidated_representation.explicit,
*consolidated_representation.deductive,
]
if not new_documents:
return
# Collect all contents for batch embedding
contents: list[str] = []
for obs in new_documents:
if isinstance(obs, ExplicitObservation):
contents.append(obs.content)
else:
contents.append(obs.conclusion)
# Batch embed all contents at once for better performance
embeddings = await embedding_client.simple_batch_embed(contents)
documents_to_create: list[schemas.DocumentCreate] = []
for i, obs in enumerate(new_documents):
if isinstance(obs, ExplicitObservation):
content = obs.content
level = "explicit"
premises = None
else:
content = obs.conclusion
level = "deductive"
premises = obs.premises
# NOTE: other kinds of observations here in the future
metadata = schemas.DocumentMetadata(
message_ids=flatten_message_ids(obs.message_ids),
message_created_at=format_datetime_utc(obs.created_at),
premises=premises,
)
documents_to_create.append(
schemas.DocumentCreate(
content=content,
session_name=obs.session_name,
level=level,
times_derived=total_times_derived,
metadata=metadata,
embedding=embeddings[i],
)
)
async with tracked_db("dream_consolidate_write") as db:
# bulk create documents
await crud.create_documents(
db,
documents_to_create,
workspace_name,
observer=observer,
observed=observed,
)
# delete old documents
await db.execute(
delete(models.Document).where(models.Document.id.in_(document_ids))
)
await db.commit()
logger.info(
"consolidated %d documents into %d new documents",
len(document_ids),
len(new_documents),
)

33
src/dreamer/dream_scheduler.py
View File

@ -153,16 +153,16 @@ class DreamScheduler:
observer=observer,
observed=observed,
)
logger.info(f"Executed dream for {work_unit_key}")
logger.info("Executed dream for %s", work_unit_key)
else:
logger.info(
f"Skipping dream for {work_unit_key} - collection is active"
"Skipping dream for %s - collection is active", work_unit_key
)
except asyncio.CancelledError:
logger.info(f"Dream task cancelled for {work_unit_key}")
logger.debug("Dream task cancelled for %s", work_unit_key)
except Exception as e:
logger.error(f"Error in delayed dream for {work_unit_key}: {str(e)}")
logger.error("Error in delayed dream for %s: %s", work_unit_key, e)
if settings.SENTRY.ENABLED:
sentry_sdk.capture_exception(e)
@ -202,12 +202,33 @@ class DreamScheduler:
# Import here to avoid circular dependency
from src.deriver.enqueue import enqueue_dream
# Find the most recent session for this observer/observed pair
async with tracked_db("dream_session_lookup") as db:
stmt = (
select(models.Document.session_name)
.where(
models.Document.workspace_name == workspace_name,
models.Document.observer == observer,
models.Document.observed == observed,
)
.order_by(models.Document.created_at.desc())
.limit(1)
)
session_name = await db.scalar(stmt)
if not session_name:
logger.warning(
f"No documents found for {workspace_name}/{observer}/{observed}, skipping dream"
)
return
await enqueue_dream(
workspace_name,
observer=observer,
observed=observed,
dream_type=dream_type,
document_count=document_count,
session_name=session_name,
)
async def shutdown(self) -> None:
@ -259,7 +280,7 @@ async def check_and_schedule_dream(
# Calculate documents added since last dream
documents_since_last_dream = current_document_count - last_dream_document_count
logger.info(
logger.debug(
"Dream check",
extra={
"workspace_name": collection.workspace_name,
@ -316,7 +337,7 @@ async def check_and_schedule_dream(
observer=collection.observer,
observed=collection.observed,
)
logger.info(
logger.debug(
"Scheduled dream",
extra={
"workspace_name": collection.workspace_name,

17
src/dreamer/dreamer.py
View File

@ -3,8 +3,8 @@ import logging
import sentry_sdk
from src.config import settings
from src.dreamer.agent import process_agent_dream
from src.dreamer.consolidate import process_consolidate_dream
from src.dependencies import tracked_db
from src.dreamer.orchestrator import run_dream
from src.schemas import DreamType
from src.utils.queue_payload import DreamPayload
@ -31,10 +31,15 @@ DREAM: {payload.dream_type} documents for {workspace_name}/{payload.observer}/{p
try:
match payload.dream_type:
case DreamType.CONSOLIDATE:
await process_consolidate_dream(payload, workspace_name)
case DreamType.AGENT:
await process_agent_dream(payload, workspace_name)
case DreamType.OMNI:
async with tracked_db("dream_orchestrator") as db:
await run_dream(
db=db,
workspace_name=workspace_name,
observer=payload.observer,
observed=payload.observed,
session_name=payload.session_name,
)
except Exception as e:
logger.error(

206
src/dreamer/orchestrator.py Normal file
View File

@ -0,0 +1,206 @@
"""
Dream orchestrator for the specialist-based architecture.
This module coordinates the full dream cycle:
0. [Optional] Surprisal sampling: Pre-filter observations by geometric surprisal
1. Generate probing questions about the peer (or use surprisal-based queries)
2. Run deduction specialist (creates deductive observations, deletes duplicates)
3. Run induction specialist (creates inductive observations from explicit + deductive)
"""
from __future__ import annotations
import logging
import time
import uuid
from typing import Any
from sqlalchemy.ext.asyncio import AsyncSession
from src.config import settings
from src.dreamer.specialists import SPECIALISTS
from src.dreamer.surprisal import SurprisalScore # type: ignore
from src.exceptions import SpecialistExecutionError, SurprisalError
from src.utils.logging import (
accumulate_metric,
log_performance_metrics,
)
logger = logging.getLogger(__name__)
# Predefined probing questions to guide the specialists
# These serve as semantic entry points for searching observations
PROBING_QUESTIONS: list[str] = [
"What information has changed or been updated? Look for dates, deadlines, schedules that moved.",
"What decisions or plans have changed? Look for rescheduled, moved, changed, updated.",
# Temporal and sequential events
"What events happened in sequence? Look for things that happened first, then, after, before.",
"What deadlines, dates, or scheduled events have been mentioned?",
# Identity and background
"What do we know about this person's identity, name, or background?",
# Recent activity
"What has this entity been doing or discussing recently?",
# Preferences and interests
"What are their preferences, likes, or dislikes?",
# Relationships
"Who are the important people in their life (family, friends, colleagues)?",
# Goals and plans
"What goals, plans, or aspirations have they shared? Have any changed?",
]
async def run_dream(
db: AsyncSession,
workspace_name: str,
observer: str,
observed: str,
session_name: str,
) -> None:
"""
Run a full dream cycle with optional surprisal-based sampling.
The dream cycle runs specialists sequentially:
0. [Optional] Surprisal sampling: Pre-filter observations by geometric surprisal
1. Deduction specialist: Creates deductive observations from explicit facts
2. Induction specialist: Creates inductive observations from patterns
Args:
db: Database session
workspace_name: Workspace identifier
observer: Observer peer name
observed: Observed peer name
session_name: Session identifier
"""
run_id = str(uuid.uuid4())[:8]
task_name = f"dream_orchestrator_{run_id}"
start_time = time.perf_counter()
logger.info(
f"[{run_id}] Starting dream cycle for {workspace_name}/{observer}/{observed}"
)
# Phase 0: Surprisal-based sampling (if enabled)
probing_questions = PROBING_QUESTIONS # Default
if settings.DREAM.SURPRISAL.ENABLED:
logger.info(f"[{run_id}] Phase 0: Computing surprisal scores")
try:
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,
)
logger.info(
f"[{run_id}] Surprisal: Found {len(high_surprisal_obs)} high-surprisal observations"
)
accumulate_metric(
task_name, "surprisal_observations", len(high_surprisal_obs), "count"
)
# Hybrid mode: Replace if sufficient high-surprisal observations
if (
len(high_surprisal_obs)
>= settings.DREAM.SURPRISAL.MIN_HIGH_SURPRISAL_FOR_REPLACE
):
probing_questions = _create_queries_from_surprisal(high_surprisal_obs)
logger.info(
f"[{run_id}] ✨ SURPRISAL REPLACE MODE: Using {len(probing_questions)} "
+ "surprisal-based queries instead of standard questions"
)
logger.info(
f"[{run_id}] Targeting observations with surprisal range: "
+ f"{high_surprisal_obs[-1].surprisal:.3f} to {high_surprisal_obs[0].surprisal:.3f}"
)
elif len(high_surprisal_obs) > 0:
# Supplement mode: Add to standard questions
surprisal_queries = _create_queries_from_surprisal(high_surprisal_obs)
probing_questions = surprisal_queries + PROBING_QUESTIONS
logger.info(
f"[{run_id}] ✨ SURPRISAL SUPPLEMENT MODE: Adding {len(surprisal_queries)} "
+ f"surprisal queries to {len(PROBING_QUESTIONS)} standard questions"
)
logger.info(
f"[{run_id}] Targeting observations with surprisal range: "
+ f"{high_surprisal_obs[-1].surprisal:.3f} to {high_surprisal_obs[0].surprisal:.3f}"
)
else:
logger.info(
f"[{run_id}] No high-surprisal observations found using standard probing questions"
)
except SurprisalError as e:
logger.error(f"[{run_id}] Surprisal sampling failed: {e}", exc_info=True)
accumulate_metric(task_name, "surprisal_error", str(e), "blob")
# Fall back to standard probing questions
# Phase 1: Run deduction specialist
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,
session_name=session_name,
probing_questions=probing_questions,
)
logger.info(f"[{run_id}] Deduction completed: {deduction_result[:200]}...")
accumulate_metric(task_name, "deduction_result", deduction_result, "blob")
except SpecialistExecutionError as e:
logger.error(f"[{run_id}] Deduction specialist failed: {e}", exc_info=True)
accumulate_metric(task_name, "deduction_error", str(e), "blob")
# Phase 2: Run induction specialist (after deduction so it can see new deductive obs)
logger.info(f"[{run_id}] Phase 2: Running induction specialist")
induction_specialist = SPECIALISTS["induction"]
try:
induction_result = await induction_specialist.run(
db=db,
workspace_name=workspace_name,
observer=observer,
observed=observed,
session_name=session_name,
probing_questions=probing_questions,
)
logger.info(f"[{run_id}] Induction completed: {induction_result[:200]}...")
accumulate_metric(task_name, "induction_result", induction_result, "blob")
except SpecialistExecutionError as e:
logger.error(f"[{run_id}] Induction specialist failed: {e}", exc_info=True)
accumulate_metric(task_name, "induction_error", str(e), "blob")
# Log final metrics
duration_ms = (time.perf_counter() - start_time) * 1000
accumulate_metric(task_name, "total_duration", duration_ms, "ms")
logger.info(f"[{run_id}] Dream cycle completed in {duration_ms:.0f}ms")
log_performance_metrics("dream_orchestrator", run_id)
def _create_queries_from_surprisal(
high_surprisal_obs: list[SurprisalScore],
) -> list[str]:
"""
Create search queries from high-surprisal observations.
Strategy: Use observation content as semantic search queries.
Truncate if too long (>200 chars).
Args:
high_surprisal_obs: List of SurprisalScore objects
Returns:
List of query strings (max 10)
"""
queries: list[Any] = []
for score in high_surprisal_obs:
content = score.observation.content
if len(content) > 200:
content = content[:200] + "..."
queries.append(content)
return queries[:10] # Limit to 10 queries

475
src/dreamer/specialists.py Normal file
View File

@ -0,0 +1,475 @@
"""
Agentic specialists for the dream cycle.
Each specialist is a fully autonomous agent that:
1. Receives probing questions as entry points
2. Uses tools to search for relevant observations
3. Creates new observations (deductive or inductive)
4. Can delete duplicates (deduction only)
"""
from __future__ import annotations
import logging
import time
import uuid
from abc import ABC, abstractmethod
from collections.abc import Callable
from typing import Any
from sqlalchemy.ext.asyncio import AsyncSession
from src import prometheus
from src.config import settings
from src.utils.agent_tools import (
DEDUCTION_SPECIALIST_TOOLS,
INDUCTION_SPECIALIST_TOOLS,
create_tool_executor,
)
from src.utils.clients import HonchoLLMCallResponse, honcho_llm_call
from src.utils.logging import accumulate_metric, log_performance_metrics
logger = logging.getLogger(__name__)
class BaseSpecialist(ABC):
"""Base class for agentic specialists."""
name: str = "base"
@abstractmethod
def get_tools(self) -> list[dict[str, Any]]:
"""Get the tools available to this specialist."""
...
@abstractmethod
def get_model(self) -> str:
"""Get the model to use for this specialist."""
...
def get_max_tokens(self) -> int:
"""Get max output tokens for this specialist."""
return 16384
def get_max_iterations(self) -> int:
"""Get max tool iterations."""
return 15
@abstractmethod
def build_system_prompt(self, observed: str) -> str:
"""Build the system prompt for this specialist."""
...
@abstractmethod
def build_user_prompt(self, probing_questions: list[str]) -> str:
"""Build the user prompt with probing questions."""
...
async def run(
self,
db: AsyncSession,
workspace_name: str,
observer: str,
observed: str,
session_name: str,
probing_questions: list[str],
) -> str:
"""
Run the specialist agent.
Args:
db: Database session
workspace_name: Workspace identifier
observer: The observing peer
observed: The peer being observed
session_name: Session identifier
probing_questions: Entry point questions to guide exploration
Returns:
Summary of work done
"""
run_id = str(uuid.uuid4())[:8]
task_name = f"dreamer_{self.name}_{run_id}"
start_time = time.perf_counter()
# Build messages
messages: list[dict[str, str]] = [
{"role": "system", "content": self.build_system_prompt(observed)},
{"role": "user", "content": self.build_user_prompt(probing_questions)},
]
# Create tool executor
tool_executor: Callable[
[str, dict[str, Any]], Any
] = await create_tool_executor(
db=db,
workspace_name=workspace_name,
observer=observer,
observed=observed,
session_name=session_name,
include_observation_ids=True,
history_token_limit=settings.DREAM.HISTORY_TOKEN_LIMIT,
)
# Get model with potential override
model = self.get_model()
llm_settings = settings.DREAM.model_copy(update={"MODEL": model})
# Run the agent loop
response: HonchoLLMCallResponse[str] = await honcho_llm_call(
llm_settings=llm_settings,
prompt="", # Ignored since we pass messages
max_tokens=self.get_max_tokens(),
tools=self.get_tools(),
tool_choice=None,
tool_executor=tool_executor,
max_tool_iterations=self.get_max_iterations(),
messages=messages,
track_name=f"Dreamer/{self.name}",
)
# Log metrics
duration_ms = (time.perf_counter() - start_time) * 1000
accumulate_metric(task_name, "total_duration", duration_ms, "ms")
accumulate_metric(
task_name, "tool_calls", len(response.tool_calls_made), "count"
)
accumulate_metric(task_name, "input_tokens", response.input_tokens, "count")
accumulate_metric(task_name, "output_tokens", response.output_tokens, "count")
if prometheus.METRICS_ENABLED:
prometheus.DREAMER_TOKENS_PROCESSED.labels(
specialist_name=self.name,
token_type=prometheus.TokenTypes.INPUT.value,
).inc(response.input_tokens)
prometheus.DREAMER_TOKENS_PROCESSED.labels(
specialist_name=self.name,
token_type=prometheus.TokenTypes.OUTPUT.value,
).inc(response.output_tokens)
logger.info(
f"{self.name}: Completed in {duration_ms:.0f}ms, "
+ f"{len(response.tool_calls_made)} tool calls, "
+ f"{response.input_tokens} in / {response.output_tokens} out"
)
log_performance_metrics(f"dreamer_{self.name}", run_id)
return response.content
class DeductionSpecialist(BaseSpecialist):
"""
Creates deductive observations from explicit observations.
This specialist:
1. Searches for explicit observations using semantic queries
2. Identifies logical implications and connections
3. Creates new deductive observations with premise linkage
4. Deletes duplicate or redundant observations
"""
name: str = "deduction"
def get_tools(self) -> list[dict[str, Any]]:
return DEDUCTION_SPECIALIST_TOOLS
def get_model(self) -> str:
return settings.DREAM.DEDUCTION_MODEL
def get_max_tokens(self) -> int:
return 8192
def get_max_iterations(self) -> int:
return 12
def build_system_prompt(self, observed: str) -> str:
return f"""You are a deductive reasoning specialist for {observed}. Your ONLY job is to create deductive observations by calling tools. Do NOT explain your reasoning - just make tool calls.
## MANDATORY WORKFLOW - YOU MUST FOLLOW THIS PATTERN
For EACH topic, you MUST alternate: search → create → search → create → ...
**CORRECT pattern:**
1. search_memory("topic 1")
2. create_observations([...deductions from topic 1...])
3. search_memory("topic 2")
4. create_observations([...deductions from topic 2...])
5. search_memory("topic 3")
6. create_observations([...deductions from topic 3...])
**WRONG pattern (DO NOT DO THIS):**
1. search_memory("topic 1")
2. search_memory("topic 2")
3. search_memory("topic 3")
4. ... more searches ...
5. create_observations([...]) ← TOO LATE, you'll hit iteration limit!
1. **ALTERNATE SEARCH/CREATE** - After each search, create observations BEFORE your next search.
2. **CREATE OBSERVATIONS** - Your primary goal is to CREATE deductive observations, not just search.
3. **MINIMIZE TEXT OUTPUT** - Do not write explanations. Just call tools.
4. **DELETE OUTDATED INFO** - When you find updated information, DELETE the old observation after creating the update.
## PRIORITY FOCUS AREAS
### 1. KNOWLEDGE UPDATES + DELETION (HIGHEST PRIORITY)
Look for the SAME fact appearing with DIFFERENT values at different times. This is critical!
Examples:
- "meeting is on Tuesday" [old] + "meeting moved to Thursday" [new] → Update + Delete old
- "lives in NYC" [old] + "moved to LA" [new] → Update + Delete old
- "works at Google" [old] + "started job at Meta" [new] → Update + Delete old
**WORKFLOW for knowledge updates:**
1. Create the deductive update observation
2. IMMEDIATELY call `delete_observations` to remove the OUTDATED observation (the old one)
3. Keep the new observation (it's still current)
```json
// Step 1: Create update
{{
"observations": [{{
"content": "[Topic] updated: [old value] → [new value]. Current: [new value]",
"level": "deductive",
"source_ids": ["old_obs_id", "new_obs_id"],
"premises": ["Original: [old fact]", "Update: [new fact]"]
}}]
}}
// Step 2: Delete outdated
{{
"observation_ids": ["old_obs_id"]
}}
```
### 2. CONTRADICTIONS (FLAG FOR CLARIFICATION)
When you find two observations that CANNOT both be true (mutually exclusive statements), create a contradiction observation.
**Update vs Contradiction:**
- UPDATE: Same topic, value changed over time ("meeting on Tuesday" → "meeting on Thursday") - DELETE old
- CONTRADICTION: Logically incompatible statements ("I love coffee" + "I hate coffee") - FLAG for user
```json
{{
"observations": [{{
"content": "Conflicting information about [topic]: [statement A] vs [statement B]",
"level": "contradiction",
"source_ids": ["obs_id_1", "obs_id_2"],
"sources": ["Statement A text", "Statement B text"]
}}]
}}
```
### 3. EVENT ORDERING & TEMPORAL SEQUENCES
Track sequences of events and their order:
- "decided to apply" → "submitted application" → "got interview" → "received offer"
- Create observations noting the sequence: "Applied for job, then interviewed, then received offer"
### 4. INFORMATION EXTRACTION
Create deductions that make implicit information explicit:
- "works as SWE at Google" → "has software engineering skills" + "is employed in tech industry"
- "has 2 kids ages 5 and 8" → "is a parent" + "has school-age children"
## WORKFLOW (REPEAT FOR EACH QUESTION)
1. Call `search_memory` with a relevant query
2. Look at timestamps - are there OLDER and NEWER observations about the same topic?
3. **IMMEDIATELY call `create_observations`** with any deductions you found
4. If you created a knowledge update, call `delete_observations` for the outdated one
## CREATING DEDUCTIVE OBSERVATIONS
```json
{{
"observations": [{{
"content": "The logical conclusion",
"level": "deductive",
"source_ids": ["id1", "id2"],
"premises": ["premise 1 text", "premise 2 text"]
}}]
}}
```
## TOOLS
- `search_memory`: Find observations by semantic query
- `create_observations`: Create new deductive OR contradiction observations (USE THIS!)
- `delete_observations`: Remove outdated observations (USE AFTER KNOWLEDGE UPDATES!)
- `get_recent_observations`: See recent activity
- `get_peer_card`: Retrieve current peer card contents
- `update_peer_card`: Update the peer card with key facts
## PEER CARD UPDATES
The peer card is a concise summary of permanent, stable information about the peer. Update it when you discover important facts that should be easily accessible.
**Peer card format** - Use these prefixes to organize entries:
- Plain facts for biographical info: "Name: Alice", "Works at Google", "Lives in NYC"
- `INSTRUCTION: ...` for standing instructions: "INSTRUCTION: Always call me Al", "INSTRUCTION: Send meeting agendas 24h in advance"
- `PREFERENCE: ...` for preferences: "PREFERENCE: Prefers morning meetings", "PREFERENCE: Likes detailed explanations"
- `TRAIT: ...` for personality traits: "TRAIT: Analytical thinker", "TRAIT: Detail-oriented"
Call `get_peer_card` first to see current contents, then `update_peer_card` with the complete updated list.
REMEMBER:
1. Knowledge updates are your #1 priority. When the same fact has different values at different times, CREATE an update observation AND DELETE the outdated observation.
2. Flag contradictions when statements are logically incompatible (can't both be true).
3. Update the peer card with permanent biographical facts and key insights."""
def build_user_prompt(self, probing_questions: list[str]) -> str:
questions_text = "\n".join(f"- {q}" for q in probing_questions)
return f"""Process these topics by ALTERNATING search and create calls:
{questions_text}
Start now:
1. Search for topic 1
2. Create observations from what you found
3. Search for topic 2
4. Create observations from what you found
... and so on."""
class InductionSpecialist(BaseSpecialist):
"""
Creates inductive observations from explicit and deductive observations.
This specialist:
1. Searches for observations (both explicit and deductive)
2. Identifies patterns and generalizations across multiple observations
3. Creates new inductive observations with source linkage
"""
name: str = "induction"
def get_tools(self) -> list[dict[str, Any]]:
return INDUCTION_SPECIALIST_TOOLS
def get_model(self) -> str:
return settings.DREAM.INDUCTION_MODEL
def get_max_tokens(self) -> int:
return 8192
def get_max_iterations(self) -> int:
return 10
def build_system_prompt(self, observed: str) -> str:
return f"""You are an inductive reasoning specialist for {observed}. Your ONLY job is to create inductive observations by calling tools. Do NOT explain your reasoning - just make tool calls.
## MANDATORY WORKFLOW - YOU MUST FOLLOW THIS PATTERN
For EACH topic, you MUST alternate: search → create → search → create → ...
**CORRECT pattern:**
1. search_memory("topic 1")
2. create_observations([...inductions from topic 1...])
3. search_memory("topic 2")
4. create_observations([...inductions from topic 2...])
5. search_memory("topic 3")
6. create_observations([...inductions from topic 3...])
**WRONG pattern (DO NOT DO THIS):**
1. search_memory("topic 1")
2. search_memory("topic 2")
3. search_memory("topic 3")
4. ... more searches ...
5. create_observations([...]) ← TOO LATE, you'll hit iteration limit!
1. **ALTERNATE SEARCH/CREATE** - After each search, create observations BEFORE your next search.
2. **CREATE OBSERVATIONS** - Your primary goal is to CREATE inductive observations.
3. **MINIMIZE TEXT OUTPUT** - Do not write explanations or summaries. Just call tools.
## PRIORITY FOCUS AREAS
### 1. TEMPORAL & SEQUENTIAL PATTERNS (HIGH PRIORITY)
Look for patterns in HOW things change over time:
- "User tends to reschedule meetings when stressed"
- "User's priorities shift toward family on weekends"
- "User makes major decisions after consulting with [person]"
### 2. EVENT SEQUENCE PATTERNS
Identify recurring sequences of events:
- "When user faces conflict, they: reflect → consult friend → make decision"
- "User's projects follow pattern: enthusiasm → doubt → completion"
### 3. INFORMATION CONSISTENCY PATTERNS
Note patterns in what information stays stable vs changes:
- "User's career goals have remained consistent around [X]"
- "User's living situation changes frequently"
### 4. STANDARD PATTERNS
Also look for:
- **Preferences**: "prefers X", "likes Y" (from multiple mentions)
- **Behaviors**: "tends to X", "usually does Y" (from repeated actions)
- **Personality**: "is generally X" (from multiple indicators)
## WORKFLOW (REPEAT FOR EACH QUESTION)
1. Call `search_memory` with a relevant query
2. Look for PATTERNS across multiple observations (both explicit and deductive levels)
3. Pay special attention to deductive observations about knowledge updates - these reveal change patterns
4. **IMMEDIATELY call `create_observations`** with any patterns you found (need 2+ sources)
5. **ONLY THEN** move to the next question and search again
## CREATING INDUCTIVE OBSERVATIONS
```json
{{
"observations": [{{
"content": "The pattern or generalization",
"level": "inductive",
"source_ids": ["id1", "id2", "id3"],
"sources": ["source 1 text", "source 2 text"],
"pattern_type": "tendency", // preference|behavior|personality|tendency|correlation
"confidence": "medium" // high (5+), medium (3-4), low (2)
}}]
}}
```
REQUIREMENTS:
- Minimum 2 source observations (use source_ids!)
- Confidence based on source count: low=2, medium=3-4, high=5+
- Pattern must generalize, not just restate one fact
## TOOLS
- `search_memory`: Find observations by semantic query
- `create_observations`: Create new inductive observations (USE THIS!)
- `get_recent_observations`: See recent activity
- `get_peer_card`: Retrieve current peer card contents
- `update_peer_card`: Update the peer card with key facts
## PEER CARD UPDATES
The peer card is a concise summary of permanent, stable information about the peer. After identifying high-confidence patterns, update the peer card.
**Peer card format** - Use these prefixes to organize entries:
- Plain facts for biographical info: "Name: Alice", "Works at Google", "Lives in NYC"
- `INSTRUCTION: ...` for standing instructions: "INSTRUCTION: Always call me Al"
- `PREFERENCE: ...` for preferences: "PREFERENCE: Prefers morning meetings"
- `TRAIT: ...` for personality/behavioral traits: "TRAIT: Analytical thinker", "TRAIT: Tends to reschedule when stressed"
Call `get_peer_card` first to see current contents, then `update_peer_card` with the complete updated list.
REMEMBER: Focus on temporal patterns and how things change. Create observations, don't just search. Update the peer card with high-confidence patterns and traits."""
def build_user_prompt(self, probing_questions: list[str]) -> str:
questions_text = "\n".join(f"- {q}" for q in probing_questions)
return f"""Process these topics by ALTERNATING search and create calls:
{questions_text}
Start now:
1. Search for topic 1
2. Create observations from patterns you found (need 2+ sources)
3. Search for topic 2
4. Create observations from patterns you found
... and so on."""
# Singleton instances
SPECIALISTS: dict[str, BaseSpecialist] = {
"deduction": DeductionSpecialist(),
"induction": InductionSpecialist(),
}

467
src/dreamer/surprisal.py Normal file
View File

@ -0,0 +1,467 @@
"""
Surprisal-based observation sampling for dream processing.
Computes geometric surprisal scores for observations using tree-based
data structures, enabling targeted deductive reasoning on anomalous
or novel observations.
"""
from __future__ import annotations
import logging
from dataclasses import dataclass
import numpy as np
from sqlalchemy import func, select
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.dreamer.trees import SurprisalTree, create_tree
logger = logging.getLogger(__name__)
@dataclass
class SurprisalScore:
"""Container for observation with surprisal score."""
observation: models.Document
surprisal: float
embedding: np.ndarray
async def sample_observations_with_surprisal(
db: AsyncSession,
workspace_name: str,
observer: str,
observed: str,
) -> list[SurprisalScore]:
"""
Sample observations and compute surprisal scores.
Workflow:
1. Fetch observations based on SAMPLING_STRATEGY
2. Extract embeddings from DB (already stored)
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
Returns:
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,
)
# Edge case: No observations
if not observations:
logger.warning(
f"No observations found for {workspace_name}/{observer}/{observed}"
)
return []
# Edge case: Too few observations for tree
min_observations = settings.DREAM.SURPRISAL.TREE_K * 2
if len(observations) < min_observations:
logger.warning(
f"Too few observations ({len(observations)} < {min_observations}), skipping surprisal computation"
)
return []
# 2. Extract embeddings
embeddings = _extract_embeddings(observations)
if embeddings.size == 0:
logger.error("Failed to extract embeddings")
return []
# 3. Build tree
tree = _build_tree(embeddings)
# 4. Compute surprisal
scores = _compute_surprisal_scores(observations, embeddings, tree)
# Edge case: Invalid surprisal values
valid_scores = [
s for s in scores if not np.isinf(s.surprisal) and not np.isnan(s.surprisal)
]
if len(valid_scores) < len(scores):
logger.warning(
f"Filtered {len(scores) - len(valid_scores)} invalid surprisal scores"
)
# 5. Normalize surprisal scores to [0, 1] range
normalized_scores = _normalize_scores(valid_scores)
# 6. Rank by normalized surprisal
normalized_scores.sort(key=lambda x: x.surprisal, reverse=True)
# Log top 5 scores BEFORE filtering
top_n = min(5, len(normalized_scores))
percent = settings.DREAM.SURPRISAL.TOP_PERCENT_SURPRISAL * 100
logger.info(f"🎯 Surprisal computation complete. Taking top {percent:.0f}%")
logger.info(f"Top {top_n} observations by normalized surprisal score:")
for i, score in enumerate(normalized_scores[:top_n], 1):
content = score.observation.content
if len(content) > 80:
content = content[:77] + "..."
logger.info(
f" #{i} [surprisal={score.surprisal:.3f}] [level={score.observation.level}] {content}"
)
filtered = _filter_by_percent(normalized_scores)
logger.info(
f"Selected: {len(filtered)}/{len(observations)} observations (top {percent:.0f}%)"
)
# Log summary statistics for filtered results
if filtered:
logger.info(
"📊 Filtered statistics: "
+ f"min={filtered[-1].surprisal:.3f}, "
+ f"max={filtered[0].surprisal:.3f}, "
+ f"mean={sum(s.surprisal for s in filtered) / len(filtered):.3f}"
)
else:
logger.info("No observations exceeded the surprisal threshold")
return filtered
except Exception as e:
logger.error(f"Surprisal sampling failed: {e}", exc_info=True)
# Return empty to allow dream to continue
return []
async def _fetch_observations(
db: AsyncSession,
workspace_name: str,
observer: str,
observed: str,
) -> list[models.Document]:
"""
Fetch observations based on configured sampling strategy.
Args:
db: Database session
workspace_name: Workspace identifier
observer: Observer peer name
observed: Observed peer name
Returns:
List of Document objects
"""
strategy = settings.DREAM.SURPRISAL.SAMPLING_STRATEGY
sample_size = settings.DREAM.SURPRISAL.SAMPLE_SIZE
levels = settings.DREAM.SURPRISAL.INCLUDE_LEVELS
if strategy == "recent":
return await _fetch_recent_observations(
db=db,
workspace_name=workspace_name,
observer=observer,
observed=observed,
limit=sample_size,
levels=levels,
)
elif strategy == "random":
return await _fetch_random_observations(
db=db,
workspace_name=workspace_name,
observer=observer,
observed=observed,
limit=sample_size,
levels=levels,
)
elif strategy == "all":
return await _fetch_all_observations(
db=db,
workspace_name=workspace_name,
observer=observer,
observed=observed,
limit=sample_size,
levels=levels,
)
else:
logger.warning(f"Unknown sampling strategy: {strategy}, using 'recent'")
return await _fetch_recent_observations(
db=db,
workspace_name=workspace_name,
observer=observer,
observed=observed,
limit=sample_size,
levels=levels,
)
async def _fetch_recent_observations(
db: AsyncSession,
workspace_name: str,
observer: str,
observed: str,
limit: int,
levels: list[str],
) -> list[models.Document]:
"""
Fetch most recent observations.
Uses existing get_all_documents() query with level filtering.
Args:
db: Database session
workspace_name: Workspace identifier
observer: Observer peer name
observed: Observed peer name
limit: Maximum number of observations to fetch
levels: Document levels to include
Returns:
List of Document objects ordered by created_at DESC
"""
stmt = get_all_documents(
workspace_name=workspace_name,
observer=observer,
observed=observed,
filters={"level": levels} if levels else None,
limit=limit,
)
result = await db.execute(stmt)
return list(result.scalars().all())
async def _fetch_random_observations(
db: AsyncSession,
workspace_name: str,
observer: str,
observed: str,
limit: int,
levels: list[str],
) -> list[models.Document]:
"""
Fetch random sample of observations.
Uses PostgreSQL's random() function for efficient random sampling.
Args:
db: Database session
workspace_name: Workspace identifier
observer: Observer peer name
observed: Observed peer name
limit: Maximum number of observations to fetch
levels: Document levels to include
Returns:
List of Document objects in random order
"""
stmt = select(models.Document).where(
models.Document.workspace_name == workspace_name,
models.Document.observer == observer,
models.Document.observed == observed,
)
if levels:
stmt = stmt.where(models.Document.level.in_(levels))
# important: limit applied after level filter
stmt = stmt.order_by(func.random()).limit(limit)
result = await db.execute(stmt)
return list(result.scalars().all())
async def _fetch_all_observations(
db: AsyncSession,
workspace_name: str,
observer: str,
observed: str,
limit: int,
levels: list[str],
) -> list[models.Document]:
"""
Fetch all observations up to limit.
Orders by created_at DESC for consistency.
Args:
db: Database session
workspace_name: Workspace identifier
observer: Observer peer name
observed: Observed peer name
limit: Maximum number of observations to fetch
levels: Document levels to include
Returns:
List of Document objects ordered by created_at DESC
"""
stmt = (
select(models.Document)
.where(
models.Document.workspace_name == workspace_name,
models.Document.observer == observer,
models.Document.observed == observed,
)
.order_by(models.Document.created_at.desc())
.limit(limit)
)
if levels:
stmt = stmt.where(models.Document.level.in_(levels))
result = await db.execute(stmt)
return list(result.scalars().all())
def _extract_embeddings(observations: list[models.Document]) -> np.ndarray:
"""
Extract embeddings from observations as numpy array.
Args:
observations: List of Document objects with embeddings
Returns:
np.ndarray of shape (N, 1536) containing embeddings
"""
if not observations:
return np.array([])
embeddings_list = [obs.embedding for obs in observations]
embeddings_array = np.array(embeddings_list, dtype=np.float32)
return embeddings_array
def _build_tree(embeddings: np.ndarray) -> SurprisalTree:
"""
Build tree structure from embeddings.
Args:
embeddings: np.ndarray of shape (N, embedding_dim)
Returns:
SurprisalTree configured per settings
"""
if embeddings.size == 0:
# Return empty tree (will handle gracefully in caller)
return create_tree(settings.DREAM.SURPRISAL.TREE_TYPE)
tree = create_tree(
tree_type=settings.DREAM.SURPRISAL.TREE_TYPE,
k=settings.DREAM.SURPRISAL.TREE_K,
)
tree.batch_insert(embeddings)
return tree
def _compute_surprisal_scores(
observations: list[models.Document],
embeddings: np.ndarray,
tree: SurprisalTree,
) -> list[SurprisalScore]:
"""
Compute surprisal score for each observation.
Args:
observations: List of Document objects
embeddings: np.ndarray of embeddings matching observations
tree: Built SurprisalTree
Returns:
List of SurprisalScore objects (unfiltered, unsorted)
"""
scores: list[SurprisalScore] = []
for obs, embedding in zip(observations, embeddings, strict=False):
surprisal = tree.surprisal(embedding)
scores.append(
SurprisalScore(
observation=obs,
surprisal=surprisal,
embedding=embedding,
)
)
return scores
def _normalize_scores(scores: list[SurprisalScore]) -> list[SurprisalScore]:
"""
Normalize surprisal scores to [0, 1] range using min-max normalization.
Args:
scores: List of SurprisalScore objects with raw surprisal values
Returns:
List of SurprisalScore objects with normalized surprisal values
"""
if not scores:
return []
# Handle edge case: all scores are identical
surprisal_values = [s.surprisal for s in scores]
min_surprisal = min(surprisal_values)
max_surprisal = max(surprisal_values)
if max_surprisal == min_surprisal:
# All scores identical - set all to 0.5 (middle of range)
return [
SurprisalScore(
observation=s.observation, surprisal=0.5, embedding=s.embedding
)
for s in scores
]
# Min-max normalization: (x - min) / (max - min)
normalized: list[SurprisalScore] = []
for score in scores:
normalized_value = (score.surprisal - min_surprisal) / (
max_surprisal - min_surprisal
)
normalized.append(
SurprisalScore(
observation=score.observation,
surprisal=normalized_value,
embedding=score.embedding,
)
)
return normalized
def _filter_by_percent(scores: list[SurprisalScore]) -> list[SurprisalScore]:
"""
Filter observations by top percentage.
Assumes scores are already sorted by surprisal (highest first).
Args:
scores: List of SurprisalScore objects, sorted by surprisal DESC
Returns:
Filtered list of SurprisalScore objects (top N% by surprisal)
"""
if not scores:
return []
# Take top percentage
top_percent = settings.DREAM.SURPRISAL.TOP_PERCENT_SURPRISAL
count = max(1, int(len(scores) * top_percent)) # At least 1 observation
return scores[:count]

64
src/dreamer/trees/__init__.py Normal file
View File

@ -0,0 +1,64 @@
"""
Tree-based structures for computing surprisal from embeddings.
Each tree computes surprisal based on the path to a point in the tree.
"""
from typing import Any
from .base import InternalNode, LeafNode, SurprisalTree, TreeNode
from .covertree import CoverNode, CoverTree
from .graph import GraphSurprisal
from .lsh import LSHSurprisal
from .prototype import PrototypeSurprisal
from .rptree import RPInternalNode, RPTree
from .sklearn_wrapper import SklearnTreeWrapper
def create_tree(tree_type: str, **kwargs: Any) -> SurprisalTree:
"""
Factory function to create different tree types.
Args:
tree_type: Type of tree to create ('rptree', 'kdtree', 'balltree',
'covertree', 'lsh', 'graph', 'prototype')
**kwargs: Additional arguments passed to tree constructor
Returns:
An instance of the specified tree type
Raises:
ValueError: If tree_type is not recognized
"""
if tree_type == "rptree":
return RPTree(**kwargs)
elif tree_type == "kdtree":
return SklearnTreeWrapper(tree_type="kd", **kwargs)
elif tree_type == "balltree":
return SklearnTreeWrapper(tree_type="ball", **kwargs)
elif tree_type == "covertree":
return CoverTree(**kwargs)
elif tree_type == "lsh":
return LSHSurprisal(**kwargs)
elif tree_type == "graph":
return GraphSurprisal(**kwargs)
elif tree_type == "prototype":
return PrototypeSurprisal(**kwargs)
else:
raise ValueError(f"Unknown tree type: {tree_type}")
__all__ = [
"create_tree",
"CoverNode",
"CoverTree",
"GraphSurprisal",
"InternalNode",
"LeafNode",
"LSHSurprisal",
"PrototypeSurprisal",
"RPInternalNode",
"RPTree",
"SklearnTreeWrapper",
"SurprisalTree",
"TreeNode",
]

63
src/dreamer/trees/base.py Normal file
View File

@ -0,0 +1,63 @@
"""
Base classes for tree-based surprisal estimation.
"""
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
import numpy as np
@dataclass
class TreeNode:
"""Base node for tree structures."""
count: int = 0
@dataclass
class LeafNode(TreeNode):
"""Leaf node containing actual points."""
points: list[np.ndarray] = field(default_factory=list)
def __post_init__(self) -> None:
if self.count == 0:
self.count: int = len(self.points)
@dataclass
class InternalNode(TreeNode):
"""Internal node with splitting criterion."""
left: "InternalNode | LeafNode | None" = None
right: "InternalNode | LeafNode | None" = None
class SurprisalTree(ABC):
"""
Abstract base class for tree-based surprisal estimation.
Subclasses implement different spatial indexing strategies.
Not all implementations use a traditional tree structure.
"""
max_leaf_size: int
total_points: int
def __init__(self, max_leaf_size: int = 10) -> None:
self.max_leaf_size = max_leaf_size
self.total_points = 0
@abstractmethod
def insert(self, point: np.ndarray) -> None:
"""Insert a point into the structure."""
@abstractmethod
def surprisal(self, point: np.ndarray) -> float:
"""Compute surprisal for a point."""
def batch_insert(self, points: np.ndarray) -> None:
"""Insert multiple points."""
for point in points:
self.insert(point)

114
src/dreamer/trees/covertree.py Normal file
View File

@ -0,0 +1,114 @@
"""
Cover Tree implementation.
"""
from dataclasses import dataclass, field
import numpy as np
from .base import SurprisalTree, TreeNode
@dataclass
class CoverNode(TreeNode):
"""Node for Cover Tree with point, scale, and children."""
point: np.ndarray | None = None
scale: float = 0.0
children: list["CoverNode"] = field(default_factory=list)
class CoverTree(SurprisalTree):
"""
Cover Tree implementation with surprisal computation.
Organizes points hierarchically by scale.
"""
base: float
root: CoverNode | None
total_points: int
def __init__(self, base: float = 2.0, max_leaf_size: int = 10) -> None:
super().__init__(max_leaf_size)
self.base = base
self.root = None
def insert(self, point: np.ndarray) -> None:
if self.root is None:
self.root = CoverNode(point=point, scale=0.0, count=1)
else:
self._insert_recursive(self.root, point, self.root.scale)
self.total_points += 1
def _insert_recursive(
self, node: CoverNode, point: np.ndarray, scale: float
) -> None:
"""
Insert point into cover tree recursively.
Fixed to ensure proper tree structure and varied paths.
"""
node.count += 1
dist = float(np.linalg.norm(node.point - point))
cover_radius = self.base**scale
if dist <= cover_radius:
for child in node.children:
child_dist = float(np.linalg.norm(child.point - point))
child_radius = self.base ** (scale - 1)
if child_dist <= child_radius:
self._insert_recursive(child, point, scale - 1)
return
new_child = CoverNode(point=point, scale=scale - 1, count=1)
node.children.append(new_child)
else:
new_scale = scale + 1
new_root = CoverNode(
point=node.point, scale=new_scale, count=node.count + 1
)
new_root.children = [node]
new_sibling = CoverNode(point=point, scale=scale, count=1)
new_root.children.append(new_sibling)
self.root = new_root
def surprisal(self, point: np.ndarray) -> float:
"""
Compute surprisal based on path through cover tree.
Uses combination of branch probabilities and distance to final node.
"""
if self.root is None:
return float("inf")
surprisal_value = 0.0
node: CoverNode = self.root
while node.children:
best_child: CoverNode | None = None
best_dist = float("inf")
for child in node.children:
dist = float(np.linalg.norm(child.point - point))
if dist < best_dist:
best_dist = dist
best_child = child
if best_child is None:
break
parent_count = node.count
child_count = best_child.count
p_branch = child_count / parent_count
surprisal_value += -np.log(p_branch + 1e-10)
node = best_child
dist_to_rep = float(np.linalg.norm(node.point - point))
dim = len(point)
distance_surprisal = dim * np.log(dist_to_rep + 0.01)
return float(surprisal_value + distance_surprisal)

124
src/dreamer/trees/graph.py Normal file
View File

@ -0,0 +1,124 @@
"""
Graph-theoretic surprisal using k-NN graph and random walk.
"""
from __future__ import annotations
from typing import Any
import numpy as np
from numpy.typing import NDArray
from sklearn.neighbors import NearestNeighbors
from .base import SurprisalTree
def _knn_indices(
points: NDArray[np.floating[Any]], n_neighbors: int
) -> NDArray[np.intp]:
"""Get k-nearest neighbor indices for each point."""
knn: NearestNeighbors = NearestNeighbors(n_neighbors=n_neighbors, algorithm="auto")
knn.fit(points) # pyright: ignore[reportUnknownMemberType]
_distances, indices = knn.kneighbors(points) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
return indices # pyright: ignore[reportUnknownVariableType]
def _nearest_index(points: NDArray[np.floating[Any]], query: np.ndarray) -> int:
"""Find index of nearest point to query."""
knn: NearestNeighbors = NearestNeighbors(n_neighbors=1)
knn.fit(points) # pyright: ignore[reportUnknownMemberType]
_distances, indices = knn.kneighbors([query]) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
return int(indices[0, 0]) # pyright: ignore[reportUnknownArgumentType]
class GraphSurprisal(SurprisalTree):
"""
Graph-theoretic surprisal using k-NN graph and random walk.
Surprisal based on stationary distribution of random walk.
"""
k: int
max_iter: int
points: list[NDArray[np.floating[Any]]]
stationary_dist: NDArray[np.floating[Any]] | None
graph_built: bool
total_points: int
def __init__(
self, k: int = 5, max_iter: int = 100, max_leaf_size: int = 10
) -> None:
super().__init__(max_leaf_size)
self.k = k
self.max_iter = max_iter
self.points = []
self.stationary_dist = None
self.graph_built = False
def insert(self, point: np.ndarray) -> None:
self.points.append(point)
self.total_points += 1
self.graph_built = False
def batch_insert(self, points: np.ndarray) -> None:
"""More efficient batch insertion."""
self.points.extend(points)
self.total_points += len(points)
self.graph_built = False
def _build_graph_and_compute_stationary(self) -> None:
"""Build k-NN graph and compute stationary distribution."""
if len(self.points) < 2:
return
points_array: NDArray[np.floating[Any]] = np.array(self.points)
k_actual = min(self.k, len(self.points) - 1)
indices = _knn_indices(points_array, k_actual + 1)
n = len(self.points)
transition = np.zeros((n, n))
for i in range(n):
neighbors: NDArray[np.intp] = indices[i, 1:]
for j_idx in neighbors:
j: int = int(j_idx)
transition[i, j] = 1.0
row_sums = transition.sum(axis=1, keepdims=True)
# Add self-loops for isolated nodes to maintain stochasticity
for i in range(n):
if row_sums[i, 0] == 0:
transition[i, i] = 1.0
row_sums[i, 0] = 1.0
transition = transition / row_sums
stationary: NDArray[np.floating[Any]] = np.ones(n) / n
for _ in range(self.max_iter):
new_stationary: NDArray[np.floating[Any]] = transition.T @ stationary
if np.allclose(new_stationary, stationary, atol=1e-6):
break
stationary = new_stationary
self.stationary_dist = stationary / stationary.sum()
self.graph_built = True
def surprisal(self, point: np.ndarray) -> float:
"""
Compute surprisal based on stationary distribution.
Well-connected central facts = high probability = low surprisal
Peripheral isolated facts = low probability = high surprisal
"""
if not self.graph_built:
self._build_graph_and_compute_stationary()
if self.stationary_dist is None or len(self.points) == 0:
return float("inf")
points_array: NDArray[np.floating[Any]] = np.array(self.points)
nearest_idx = _nearest_index(points_array, point)
prob = self.stationary_dist[nearest_idx]
return float(-np.log(prob + 1e-10))

84
src/dreamer/trees/lsh.py Normal file
View File

@ -0,0 +1,84 @@
"""
Locality-Sensitive Hashing based surprisal estimation.
"""
from typing import Any
import numpy as np
from numpy.typing import NDArray
from .base import SurprisalTree
class LSHSurprisal(SurprisalTree):
"""
Locality-Sensitive Hashing based surprisal estimation.
O(1) operations using hash collision frequency as density proxy.
"""
num_tables: int
num_bits: int
tables: list[dict[int, int]]
hash_directions: list[NDArray[np.floating[Any]]]
initialized: bool
total_points: int
def __init__(
self, num_tables: int = 10, num_bits: int = 8, max_leaf_size: int = 10
) -> None:
super().__init__(max_leaf_size)
self.num_tables = num_tables
self.num_bits = num_bits
self.tables = [{} for _ in range(num_tables)]
self.hash_directions = []
self.initialized = False
def _initialize_hash_functions(self, dim: int) -> None:
"""Initialize random projection directions for LSH."""
if not self.initialized:
for _ in range(self.num_tables):
directions: NDArray[np.floating[Any]] = np.random.randn(
self.num_bits, dim
)
directions = directions / np.linalg.norm(
directions, axis=1, keepdims=True
)
self.hash_directions.append(directions)
self.initialized = True
def _hash_vector(self, point: np.ndarray, table_idx: int) -> int:
"""Hash a vector using random projections."""
projections: NDArray[np.floating[Any]] = self.hash_directions[table_idx] @ point
binary: NDArray[np.intp] = (projections > 0).astype(np.intp)
hash_val = int("".join(str(b) for b in binary), 2)
return hash_val
def insert(self, point: np.ndarray) -> None:
if not self.initialized:
self._initialize_hash_functions(len(point))
for i, table in enumerate(self.tables):
bucket = self._hash_vector(point, i)
table[bucket] = table.get(bucket, 0) + 1
self.total_points += 1
def surprisal(self, point: np.ndarray) -> float:
"""
Compute surprisal using hash collision frequency.
High collision = low surprisal (common pattern)
Low collision = high surprisal (rare pattern)
"""
if self.total_points == 0 or not self.initialized:
return float("inf")
counts: list[int] = []
for i, table in enumerate(self.tables):
bucket = self._hash_vector(point, i)
count = table.get(bucket, 0)
counts.append(count)
avg_count = np.mean(counts)
avg_density = float(avg_count) / self.total_points
return float(-np.log(avg_density + 1e-10))

97
src/dreamer/trees/prototype.py Normal file
View File

@ -0,0 +1,97 @@
"""
Prototype-based surprisal using clustering.
"""
from typing import Any
import numpy as np
from numpy.typing import NDArray
from .base import SurprisalTree
class PrototypeSurprisal(SurprisalTree):
"""
Prototype-based surprisal using clustering.
Surprisal proportional to distance from nearest prototype centroid.
"""
n_clusters: int
points: list[NDArray[np.floating[Any]]]
prototypes: NDArray[np.floating[Any]] | None
clusters_built: bool
total_points: int
surprisal_scale: float
def __init__(
self,
n_clusters: int = 10,
max_leaf_size: int = 10,
surprisal_scale: float = 10.0,
) -> None:
"""
Initialize the prototype-based surprisal tree.
Args:
n_clusters: Number of prototype clusters to form.
max_leaf_size: Maximum size for leaf nodes (passed to base class).
surprisal_scale: Multiplier applied to the minimum distance from
prototypes. Default is 10.0 to normalize raw embedding distances
(typically in [0, 1]) to a more interpretable surprisal range.
"""
super().__init__(max_leaf_size)
self.n_clusters = n_clusters
self.points = []
self.prototypes = None
self.clusters_built = False
self.surprisal_scale = surprisal_scale
def insert(self, point: np.ndarray) -> None:
self.points.append(point)
self.total_points += 1
self.clusters_built = False
def batch_insert(self, points: np.ndarray) -> None:
"""More efficient batch insertion."""
self.points.extend(points)
self.total_points += len(points)
self.clusters_built = False
def _build_clusters(self) -> None:
"""Build clusters and identify prototypes."""
if len(self.points) < self.n_clusters:
self.prototypes = np.array(self.points)
self.clusters_built = True
return
from sklearn.cluster import KMeans
points_array: NDArray[np.floating[Any]] = np.array(self.points)
n_clusters_actual = min(self.n_clusters, len(self.points))
kmeans: KMeans = KMeans(
n_clusters=n_clusters_actual, random_state=42, n_init="auto"
)
kmeans.fit(points_array) # pyright: ignore[reportUnknownMemberType]
self.prototypes = kmeans.cluster_centers_ # pyright: ignore[reportUnknownMemberType]
self.clusters_built = True
def surprisal(self, point: np.ndarray) -> float:
"""
Compute surprisal as distance to nearest prototype.
Close to prototype = expected = low surprisal
Far from prototype = surprising = high surprisal
"""
if not self.clusters_built:
self._build_clusters()
if self.prototypes is None or len(self.prototypes) == 0:
return float("inf")
distances: NDArray[np.floating[Any]] = np.linalg.norm(
self.prototypes - point, axis=1
)
min_distance: float = float(np.min(distances))
return min_distance * self.surprisal_scale

157
src/dreamer/trees/rptree.py Normal file
View File

@ -0,0 +1,157 @@
"""
Random Projection Tree implementation.
"""
from dataclasses import dataclass
import numpy as np
from .base import InternalNode, LeafNode, SurprisalTree
@dataclass
class RPInternalNode(InternalNode):
"""Internal node for Random Projection Tree with direction and threshold."""
direction: np.ndarray | None = None
threshold: float = 0.0
class RPTree(SurprisalTree):
"""
Random Projection Tree with surprisal computation.
Uses random projection directions at each split.
"""
root: LeafNode | RPInternalNode | None
total_points: int
def __init__(self, max_leaf_size: int = 10) -> None:
super().__init__(max_leaf_size)
self.root = None
def insert(self, point: np.ndarray) -> None:
if self.root is None:
self.root = LeafNode(points=[point], count=1)
else:
self.root = self._insert(self.root, point)
self.total_points += 1
def _insert(
self, node: LeafNode | RPInternalNode, point: np.ndarray
) -> LeafNode | RPInternalNode:
node.count += 1
if isinstance(node, LeafNode):
node.points.append(point)
if len(node.points) > self.max_leaf_size:
return self._split_leaf(node)
return node
else:
if self._go_left(node, point):
if node.left is not None:
node.left = self._insert_child(node.left, point)
else:
if node.right is not None:
node.right = self._insert_child(node.right, point)
return node
def _insert_child(
self, child: InternalNode | LeafNode, point: np.ndarray
) -> LeafNode | RPInternalNode:
"""Insert into a child node, handling the type narrowing."""
if isinstance(child, LeafNode | RPInternalNode):
return self._insert(child, point)
raise TypeError(f"Unexpected child type: {type(child)}")
def _split_leaf(self, leaf: LeafNode) -> LeafNode | RPInternalNode:
"""
Split a leaf using random projection.
Tries multiple random directions to find a good split.
"""
points = np.array(leaf.points)
if len(points) > 1:
variance = np.var(points, axis=0).sum()
if variance < 1e-10:
return leaf
max_attempts = 5
for _attempt in range(max_attempts):
direction = np.random.randn(points.shape[1])
norm = np.linalg.norm(direction)
if norm < 1e-10:
continue
direction /= norm
projections = points @ direction
if np.std(projections) < 1e-10:
continue
threshold = np.median(projections)
left_mask = projections < threshold
right_mask = ~left_mask
if not left_mask.any() or not right_mask.any():
proj_min, proj_max = projections.min(), projections.max()
if proj_max - proj_min < 1e-10:
continue
threshold = (proj_min + proj_max) / 2
left_mask = projections < threshold
right_mask = ~left_mask
left_points = [p for p, m in zip(leaf.points, left_mask, strict=False) if m]
right_points = [
p for p, m in zip(leaf.points, right_mask, strict=False) if m
]
if left_points and right_points:
return RPInternalNode(
direction=direction,
threshold=float(threshold),
left=LeafNode(points=left_points, count=len(left_points)),
right=LeafNode(points=right_points, count=len(right_points)),
count=leaf.count,
)
return leaf
def _go_left(self, node: RPInternalNode, point: np.ndarray) -> bool:
"""Determine if point should go left. Must match split criterion."""
return bool((point @ node.direction) < node.threshold)
def surprisal(self, point: np.ndarray) -> float:
"""
Compute surprisal as cumulative log-probability along path.
S(x) = -log P(path to x) = Σ -log(n_child / n_parent)
"""
if self.root is None:
return float("inf")
surprisal_value = 0.0
node: LeafNode | RPInternalNode = self.root
while isinstance(node, RPInternalNode):
parent_count = node.count
if self._go_left(node, point) and node.left is not None:
child_count = node.left.count
if isinstance(node.left, LeafNode | RPInternalNode):
node = node.left
else:
break
elif node.right is not None:
child_count = node.right.count
if isinstance(node.right, LeafNode | RPInternalNode):
node = node.right
else:
break
else:
break
p_branch = child_count / parent_count
surprisal_value += -np.log(p_branch + 1e-10)
return surprisal_value

77
src/dreamer/trees/sklearn_wrapper.py Normal file
View File

@ -0,0 +1,77 @@
"""
Wrapper for sklearn's KDTree and BallTree.
"""
from typing import Any
import numpy as np
from numpy.typing import NDArray
from sklearn.neighbors import (
BallTree, # pyright: ignore[reportUnknownVariableType]
KDTree, # pyright: ignore[reportUnknownVariableType]
)
from .base import SurprisalTree
class SklearnTreeWrapper(SurprisalTree):
"""
Wrapper for sklearn's KDTree and BallTree with surprisal computation.
Uses density estimation via k-nearest neighbors.
"""
tree_type: str
k: int
points: list[NDArray[np.floating[Any]]]
tree: KDTree | BallTree | None
total_points: int
def __init__(
self, tree_type: str = "kd", k: int = 5, max_leaf_size: int = 10
) -> None:
super().__init__(max_leaf_size)
self.tree_type = tree_type
self.k = k
self.points = []
self.tree = None
def insert(self, point: np.ndarray) -> None:
self.points.append(point)
self.total_points += 1
self._rebuild_tree()
def batch_insert(self, points: np.ndarray) -> None:
"""More efficient batch insertion."""
self.points.extend(points)
self.total_points += len(points)
self._rebuild_tree()
def _rebuild_tree(self) -> None:
if len(self.points) == 0:
return
points_array: NDArray[np.floating[Any]] = np.array(self.points)
if self.tree_type == "kd":
self.tree = KDTree(points_array)
elif self.tree_type == "ball":
self.tree = BallTree(points_array)
else:
raise ValueError(f"Unknown tree type: {self.tree_type}")
def surprisal(self, point: np.ndarray) -> float:
"""
Compute surprisal using k-NN density estimation.
S(e) ≈ log(V_k(e)) where V_k is the volume of k-ball
"""
if self.tree is None or len(self.points) < self.k: # pyright: ignore[reportUnknownMemberType]
return float("inf")
k_actual = min(self.k, len(self.points))
distances, _indices = self.tree.query([point], k=k_actual) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
avg_distance: float = float(np.mean(distances[0])) # pyright: ignore[reportUnknownArgumentType]
dim = point.shape[0]
surprisal_value: float = dim * np.log(avg_distance + 1e-10)
return surprisal_value

88
src/embedding_client.py
View File

@ -40,6 +40,21 @@ class _EmbeddingClient:
self.max_embedding_tokens: int = min(settings.MAX_EMBEDDING_TOKENS, 2048)
# Gemini batch size is not documented, using conservative estimate
self.max_batch_size: int = 100
elif self.provider == "openrouter":
if api_key is None:
api_key = settings.LLM.OPENAI_COMPATIBLE_API_KEY
if not api_key:
raise ValueError(
"OpenRouter API key (LLM_OPENAI_COMPATIBLE_API_KEY) is required"
)
base_url = (
settings.LLM.OPENAI_COMPATIBLE_BASE_URL
or "https://openrouter.ai/api/v1"
)
self.client = AsyncOpenAI(api_key=api_key, base_url=base_url)
self.model = "openai/text-embedding-3-small"
self.max_embedding_tokens = settings.MAX_EMBEDDING_TOKENS
self.max_batch_size = 2048 # Same as OpenAI
else: # openai
if api_key is None:
api_key = settings.LLM.OPENAI_API_KEY
@ -214,43 +229,64 @@ class _EmbeddingClient:
return batches
async def _process_batch(
self, batch: list[BatchItem]
self, batch: list[BatchItem], max_retries: int = 3
) -> dict[str, dict[int, list[float]]]:
"""
Process a single batch through the embeddings API.
Process a single batch through the embeddings API with retry logic.
Args:
batch: List of BatchItem objects to embed
max_retries: Maximum number of retry attempts (default: 3)
Returns:
Maps text IDs to {chunk_index: embedding_vector} dictionaries
"""
try:
# Organize embeddings by text_id and chunk_index
result: dict[str, dict[int, list[float]]] = defaultdict(dict)
last_exception: Exception | None = None
if isinstance(self.client, genai.Client):
response = await self.client.aio.models.embed_content(
model=self.model,
contents=[item.text for item in batch],
config={"output_dimensionality": 1536},
)
if response.embeddings:
for item, embedding in zip(batch, response.embeddings, strict=True):
if embedding.values:
result[item.text_id][item.chunk_index] = embedding.values
else: # openai
response = await self.client.embeddings.create(
model=self.model, input=[item.text for item in batch]
)
for item, embedding_data in zip(batch, response.data, strict=True):
result[item.text_id][item.chunk_index] = embedding_data.embedding
for attempt in range(max_retries):
try:
# Organize embeddings by text_id and chunk_index
result: dict[str, dict[int, list[float]]] = defaultdict(dict)
return dict(result)
if isinstance(self.client, genai.Client):
response = await self.client.aio.models.embed_content(
model=self.model,
contents=[item.text for item in batch],
config={"output_dimensionality": 1536},
)
if response.embeddings:
for item, embedding in zip(
batch, response.embeddings, strict=True
):
if embedding.values:
result[item.text_id][item.chunk_index] = (
embedding.values
)
else: # openai / openrouter
response = await self.client.embeddings.create(
model=self.model, input=[item.text for item in batch]
)
for item, embedding_data in zip(batch, response.data, strict=True):
result[item.text_id][item.chunk_index] = (
embedding_data.embedding
)
except Exception:
logger.exception("Error processing batch")
raise
return dict(result)
except Exception as e:
last_exception = e
if attempt < max_retries - 1:
# Exponential backoff: 1s, 2s, 4s
wait_time = 2**attempt
logger.warning(
f"Embedding batch failed (attempt {attempt + 1}/{max_retries}), "
+ f"retrying in {wait_time}s: {e}"
)
await asyncio.sleep(wait_time)
else:
logger.exception("Error processing batch after all retries")
raise last_exception or RuntimeError("Batch processing failed")
def _accumulate_embeddings(
self, batch_results: list[dict[str, dict[int, list[float]]]]
@ -344,6 +380,8 @@ class EmbeddingClient:
provider = settings.LLM.EMBEDDING_PROVIDER
if provider == "gemini":
api_key = settings.LLM.GEMINI_API_KEY
elif provider == "openrouter":
api_key = settings.LLM.OPENAI_COMPATIBLE_API_KEY
else:
api_key = settings.LLM.OPENAI_API_KEY

16
src/exceptions.py
View File

@ -109,6 +109,22 @@ class FileProcessingError(HonchoException):
detail = "File processing error"
@final
class SurprisalError(HonchoException):
"""Exception raised when surprisal sampling fails during a dream cycle."""
status_code = 500
detail = "Surprisal sampling failed"
@final
class SpecialistExecutionError(HonchoException):
"""Exception raised when a specialist fails during dream orchestration."""
status_code = 500
detail = "Specialist execution failed"
class LLMError(Exception):
"""Exception raised when an LLM call fails.

2
src/main.py
View File

@ -21,6 +21,7 @@ from src.config import settings
from src.db import engine, request_context
from src.exceptions import HonchoException
from src.routers import (
conclusions,
keys,
messages,
observations,
@ -174,6 +175,7 @@ app.include_router(workspaces.router, prefix="/v2")
app.include_router(peers.router, prefix="/v2")
app.include_router(sessions.router, prefix="/v2")
app.include_router(messages.router, prefix="/v2")
app.include_router(conclusions.router, prefix="/v2")
app.include_router(observations.router, prefix="/v2")
app.include_router(keys.router, prefix="/v2")
app.include_router(webhooks.router, prefix="/v2")

15
src/models.py
View File

@ -371,6 +371,9 @@ class Document(Base):
times_derived: Mapped[int] = mapped_column(
Integer, nullable=False, server_default=text("1")
)
source_ids: Mapped[list[str] | None] = mapped_column(
JSONB, nullable=True, server_default=text("NULL")
)
embedding: MappedColumn[Any] = mapped_column(Vector(1536))
created_at: Mapped[datetime.datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now(), index=True
@ -422,6 +425,12 @@ class Document(Base):
"embedding": "vector_cosine_ops"
}, # Cosine distance operator
),
# GIN index for efficient tree traversal (finding children by source IDs)
Index(
"ix_documents_source_ids_gin",
"source_ids",
postgresql_using="gin",
),
)
@ -458,6 +467,12 @@ class QueueItem(Base):
"message_id",
postgresql_where=text("message_id IS NOT NULL"),
),
Index(
"ux_queue_dream_pending_work_unit_key",
"work_unit_key",
unique=True,
postgresql_where=text("task_type = 'dream' AND processed = false"),
),
Index(
"ix_queue_work_unit_key_processed_id",
"work_unit_key",

55
src/prometheus.py
View File

@ -5,6 +5,7 @@ This module defines all Prometheus metrics for all Honcho processes and exposes
"""
import logging
from enum import Enum
from typing import cast
from prometheus_client import (
@ -72,12 +73,14 @@ MESSAGES_CREATED = NamespacedCounter(
# Incremented in: src/routers/peers.py when successful dialectic calls are made
# Labels:
# - workspace_name: The workspace where the dialectic call was made
# - reasoning_level: The reasoning level used for the call
DIALECTIC_CALLS = NamespacedCounter(
"dialectic_calls_total",
"Total dialectic calls",
[
"namespace",
"workspace_name",
"reasoning_level",
],
)
@ -93,13 +96,19 @@ DERIVER_QUEUE_ITEMS_PROCESSED = NamespacedCounter(
["namespace", "workspace_name", "task_type"],
)
class TokenTypes(Enum):
INPUT = "input"
OUTPUT = "output"
# Tracks the total number of input and output tokens processed by the deriver.
#
# Incremented in: src/deriver/deriver.py after the critical analysis call is made
# Labels:
# - task_type: The type of task that processed the tokens (e.g., "representation", "summary")
# - task_type: The type of task that processed the tokens
# - token_type: The type of tokens ("input" or "output")
# - component: The component of the input (e.g., "peer_card", "working_representation", "prompt", "new_turns", "session_context")
# - component: The component of the input
DERIVER_TOKENS_PROCESSED = NamespacedCounter(
"deriver_tokens_processed_total",
"Total tokens processed by the deriver",
@ -111,17 +120,55 @@ DERIVER_TOKENS_PROCESSED = NamespacedCounter(
],
)
class DeriverTaskTypes(Enum):
INGESTION = "ingestion"
SUMMARY = "summary"
class DeriverComponents(Enum):
PROMPT = "prompt" # used in ingestion and summary
MESSAGES = "messages" # used in ingestion and summary
PREVIOUS_SUMMARY = "previous_summary" # only used for summary
OUTPUT_TOTAL = "output_total"
# Tracks the total number of input and output tokens processed by the dialectic.
#
# Incremented in: src/dialectic/chat.py after the dialectic call is made
# Incremented in: src/dialectic/core.py after the dialectic call is made
# Labels:
# - token_type: The type of tokens ("input" or "output")
# - token_type: The type of tokens
# - component: The component of the input
# - reasoning_level: The reasoning level used for the call
DIALECTIC_TOKENS_PROCESSED = NamespacedCounter(
"dialectic_tokens_processed_total",
"Total tokens processed by the dialectic",
[
"namespace",
"token_type",
"component",
"reasoning_level",
],
)
class DialecticComponents(Enum):
TOTAL = "total"
# Tracks the total number of input and output tokens processed by the dreamer.
#
# Incremented in: src/dreamer/specialists.py after the specialist LLM call is made
# Labels:
# - specialist_name: The name of the specialist ("deduction" or "induction")
# - token_type: The type of tokens ("input" or "output")
DREAMER_TOKENS_PROCESSED = NamespacedCounter(
"dreamer_tokens_processed_total",
"Total tokens processed by the dreamer",
[
"namespace",
"specialist_name",
"token_type",
],
)

150
src/routers/conclusions.py Normal file
View File

@ -0,0 +1,150 @@
import logging
from fastapi import APIRouter, Body, Depends, Path, Query
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import apaginate
from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, schemas
from src.dependencies import db
from src.exceptions import ResourceNotFoundException, ValidationException
from src.security import require_auth
logger = logging.getLogger(__name__)
router = APIRouter(
prefix="/workspaces/{workspace_id}/conclusions",
tags=["conclusions"],
dependencies=[Depends(require_auth(workspace_name="workspace_id"))],
)
@router.post(
"",
response_model=list[schemas.Conclusion],
)
async def create_conclusions(
workspace_id: str = Path(..., description="ID of the workspace"),
body: schemas.ConclusionBatchCreate = Body(
...,
description="Batch of conclusions to create",
),
db: AsyncSession = db,
) -> list[schemas.Conclusion]:
"""
Create one or more conclusions.
Conclusions are theory-of-mind facts derived from interactions between peers.
"""
documents = await crud.create_observations(
db,
observations=body.conclusions,
workspace_name=workspace_id,
)
logger.debug(
"Created %d conclusions in workspace %s",
len(documents),
workspace_id,
)
return [schemas.Conclusion.model_validate(doc) for doc in documents]
@router.post(
"/list",
response_model=Page[schemas.Conclusion],
)
async def list_conclusions(
workspace_id: str = Path(..., description="ID of the workspace"),
options: schemas.ConclusionGet | None = Body(
None,
description="Filtering options for the conclusions list",
),
reverse: bool | None = Query(
False,
description="Whether to reverse the order of results",
),
db: AsyncSession = db,
):
"""
List conclusions using custom filters, ordered by recency unless `reverse` is true.
"""
filters = None
if options and hasattr(options, "filters"):
filters = options.filters
if filters == {}:
filters = None
stmt = crud.get_documents_with_filters(
workspace_name=workspace_id,
filters=filters,
reverse=reverse or False,
)
return await apaginate(db, stmt)
@router.post(
"/query",
response_model=list[schemas.Conclusion],
)
async def query_conclusions(
workspace_id: str = Path(..., description="ID of the workspace"),
body: schemas.ConclusionQuery = Body(
...,
description="Semantic search parameters for conclusions",
),
db: AsyncSession = db,
) -> list[schemas.Conclusion]:
"""
Query conclusions using semantic search.
"""
observer = None
observed = None
if body.filters:
observer = body.filters.get("observer") or body.filters.get("observer_id")
observed = body.filters.get("observed") or body.filters.get("observed_id")
if not observer or not observed:
raise ValidationException(
"observer and observed must be specified for semantic search"
)
documents = await crud.query_documents(
db,
workspace_name=workspace_id,
query=body.query,
observer=observer,
observed=observed,
filters=body.filters,
max_distance=body.distance,
top_k=body.top_k,
)
return [schemas.Conclusion.model_validate(doc) for doc in documents]
@router.delete(
"/{conclusion_id}",
)
async def delete_conclusion(
workspace_id: str = Path(..., description="ID of the workspace"),
conclusion_id: str = Path(..., description="ID of the conclusion to delete"),
db: AsyncSession = db,
):
"""
Delete a specific conclusion (document).
"""
try:
await crud.delete_document_by_id(
db,
workspace_name=workspace_id,
document_id=conclusion_id,
)
logger.debug("Conclusion %s deleted successfully", conclusion_id)
return {"message": "Conclusion deleted successfully"}
except ResourceNotFoundException:
raise
except ValueError as e:
logger.warning(f"Failed to delete conclusion {conclusion_id}: {str(e)}")
raise ResourceNotFoundException("Conclusion not found") from e

6
src/routers/observations.py
View File

@ -22,6 +22,7 @@ router = APIRouter(
@router.post(
"",
response_model=list[schemas.Observation],
deprecated=True,
)
async def create_observations(
workspace_id: str = Path(..., description="ID of the workspace"),
@ -41,7 +42,7 @@ async def create_observations(
"""
documents = await crud.create_observations(
db,
observations=body.observations,
observations=body.conclusions,
workspace_name=workspace_id,
)
@ -56,6 +57,7 @@ async def create_observations(
@router.post(
"/list",
response_model=Page[schemas.Observation],
deprecated=True,
)
async def list_observations(
workspace_id: str = Path(..., description="ID of the workspace"),
@ -94,6 +96,7 @@ async def list_observations(
@router.post(
"/query",
response_model=list[schemas.Observation],
deprecated=True,
)
async def query_observations(
workspace_id: str = Path(..., description="ID of the workspace"),
@ -138,6 +141,7 @@ async def query_observations(
@router.delete(
"/{observation_id}",
deprecated=True,
)
async def delete_observation(
workspace_id: str = Path(..., description="ID of the workspace"),

92
src/routers/peers.py
View File

@ -1,8 +1,8 @@
import json
import logging
from collections.abc import AsyncGenerator, AsyncIterator
from collections.abc import AsyncIterator
from fastapi import APIRouter, Body, Depends, Path, Query
from fastapi.exceptions import HTTPException
from fastapi.responses import StreamingResponse
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import apaginate
@ -11,7 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, prometheus, schemas
from src.config import settings
from src.dependencies import db, tracked_db
from src.dialectic import chat as dialectic_chat
from src.dialectic.chat import agentic_chat, agentic_chat_stream
from src.exceptions import AuthenticationException, ResourceNotFoundException
from src.security import JWTParams, require_auth
from src.utils.search import search
@ -168,64 +168,56 @@ async def chat(
peers=[schemas.PeerCreate(name=peer_id)],
)
if not options.stream:
response = await dialectic_chat(
workspace_name=workspace_id,
session_name=options.session_id,
query=options.query,
stream=options.stream,
observer=peer_id,
# if target is given, that's the observed peer. otherwise, observer==observed
# and it's answered from the omniscient Honcho perspective
observed=options.target if options.target is not None else peer_id,
)
if options.stream:
# Stream the response using Server-Sent Events
async def format_sse_stream(
chunks: AsyncIterator[str],
) -> AsyncIterator[str]:
"""Format chunks as SSE events."""
async for chunk in chunks:
yield f"data: {json.dumps({'delta': {'content': chunk}, 'done': False})}\n\n"
yield f"data: {json.dumps({'done': True})}\n\n"
if prometheus.METRICS_ENABLED:
prometheus.DIALECTIC_CALLS.labels(
workspace_name=workspace_id,
reasoning_level=options.reasoning_level,
).inc()
return schemas.DialecticResponse(content=str(response))
async def parse_stream() -> AsyncGenerator[str, None]:
try:
stream = await dialectic_chat(
workspace_name=workspace_id,
session_name=options.session_id,
query=options.query,
stream=options.stream,
observer=peer_id,
observed=options.target if options.target is not None else peer_id,
)
if prometheus.METRICS_ENABLED:
prometheus.DIALECTIC_CALLS.labels(
return StreamingResponse(
format_sse_stream(
agentic_chat_stream(
workspace_name=workspace_id,
).inc()
if isinstance(stream, AsyncIterator):
async for chunk in stream:
if chunk.content:
# Send each chunk as JSON with nested delta object
stream_chunk = schemas.DialecticStreamChunk(
delta=schemas.DialecticStreamDelta(content=chunk.content)
)
yield f"data: {stream_chunk.model_dump_json()}\n\n"
# Send final done message
final_chunk = schemas.DialecticStreamChunk(
delta=schemas.DialecticStreamDelta(), done=True
session_name=options.session_id,
query=options.query,
observer=peer_id,
observed=options.target if options.target is not None else peer_id,
reasoning_level=options.reasoning_level,
)
yield f"data: {final_chunk.model_dump_json()}\n\n"
else:
raise HTTPException(status_code=500, detail="Invalid stream type")
except Exception as e:
logger.error(f"Error in stream: {str(e)}")
raise HTTPException(status_code=500, detail=str(e)) from e
),
media_type="text/event-stream",
)
return StreamingResponse(
content=parse_stream(), media_type="text/event-stream", status_code=200
response = await agentic_chat(
workspace_name=workspace_id,
session_name=options.session_id,
query=options.query,
observer=peer_id,
# if target is given, that's the observed peer. otherwise, observer==observed
# and it's answered from the omniscient Honcho perspective
observed=options.target if options.target is not None else peer_id,
reasoning_level=options.reasoning_level,
)
if prometheus.METRICS_ENABLED:
prometheus.DIALECTIC_CALLS.labels(
workspace_name=workspace_id,
reasoning_level=options.reasoning_level,
).inc()
return schemas.DialecticResponse(content=str(response))
@router.post(
"/{peer_id}/representation",

47
src/routers/workspaces.py
View File

@ -128,10 +128,41 @@ async def search_workspace(
@router.get(
"/{workspace_id}/deriver/status",
response_model=schemas.DeriverStatus,
"/{workspace_id}/queue/status",
response_model=schemas.QueueStatus,
dependencies=[Depends(require_auth(workspace_name="workspace_id"))],
)
async def get_queue_status(
workspace_id: str = Path(..., description="ID of the workspace"),
observer_id: str | None = Query(
None, description="Optional observer ID to filter by"
),
sender_id: str | None = Query(None, description="Optional sender ID to filter by"),
session_id: str | None = Query(
None, description="Optional session ID to filter by"
),
db: AsyncSession = db,
):
"""Get the processing queue status, optionally scoped to an observer, sender, and/or session."""
try:
return await crud.get_queue_status(
db,
workspace_name=workspace_id,
session_name=session_id,
observer=observer_id,
observed=sender_id,
)
except ValueError as e:
logger.warning(f"Invalid request parameters: {str(e)}")
raise HTTPException(status_code=400, detail=str(e)) from e
@router.get(
"/{workspace_id}/deriver/status",
response_model=schemas.QueueStatus,
dependencies=[Depends(require_auth(workspace_name="workspace_id"))],
deprecated=True,
)
async def get_deriver_status(
workspace_id: str = Path(..., description="ID of the workspace"),
observer_id: str | None = Query(
@ -143,9 +174,9 @@ async def get_deriver_status(
),
db: AsyncSession = db,
):
"""Get the deriver processing status, optionally scoped to an observer, sender, and/or session"""
"""Deprecated: use /queue/status. Provides identical response payload."""
try:
return await crud.get_deriver_status(
return await crud.get_queue_status(
db,
workspace_name=workspace_id,
session_name=session_id,
@ -202,8 +233,14 @@ async def trigger_dream(
observed=observed,
dream_type=dream_type,
document_count=document_count,
session_name=request.session_id,
)
logger.info(
f"Manually triggered dream: {dream_type.value} for {workspace_id}/{observer}/{observed}"
"Manually triggered dream: %s for %s/%s/%s (session: %s)",
dream_type.value,
workspace_id,
observer,
observed,
request.session_id,
)

132
src/schemas.py
View File

@ -6,6 +6,7 @@ from urllib.parse import urlparse
import tiktoken
from pydantic import (
AliasChoices,
BaseModel,
ConfigDict,
Field,
@ -14,7 +15,7 @@ from pydantic import (
model_validator,
)
from src.config import settings
from src.config import ReasoningLevel, settings
from src.utils.representation import Representation
from src.utils.types import DocumentLevel
@ -24,8 +25,7 @@ RESOURCE_NAME_PATTERN = r"^[a-zA-Z0-9_-]+$"
class DreamType(str, Enum):
"""Types of dreams that can be triggered."""
CONSOLIDATE = "consolidate"
AGENT = "agent"
OMNI = "omni"
class DeriverConfiguration(BaseModel):
@ -496,9 +496,26 @@ class DocumentMetadata(BaseModel):
message_created_at: str = Field(
description="The timestamp of the message that this document was derived from. Note that this is not the same as the created_at timestamp of the document. This timestamp is usually only saved with second-level precision."
)
source_ids: list[str] | None = Field(
default=None,
description="Document IDs of source observations for tree traversal -- required for deductive and inductive observations",
)
# Deductive observation fields
premises: list[str] | None = Field(
default=None,
description="The premises of the deduction -- only applicable for deductive observations",
description="Human-readable premise text for display -- only applicable for deductive observations",
)
sources: list[str] | None = Field(
default=None,
description="Human-readable source text for display -- only applicable for inductive observations",
)
pattern_type: str | None = Field(
default=None,
description="Type of pattern identified (preference, behavior, personality, tendency, correlation) -- only applicable for inductive observations",
)
confidence: str | None = Field(
default=None,
description="Confidence level (high, medium, low) -- only applicable for inductive observations",
)
@ -509,7 +526,7 @@ class DocumentCreate(DocumentBase):
)
level: DocumentLevel = Field(
default="explicit",
description="The level of the document (explicit or deductive)",
description="The level of the document (explicit, deductive, inductive, or contradiction)",
)
times_derived: int = Field(
default=1,
@ -518,40 +535,50 @@ class DocumentCreate(DocumentBase):
)
metadata: DocumentMetadata = Field()
embedding: list[float] = Field()
# Tree linkage field
source_ids: list[str] | None = Field(
default=None,
description="Document IDs of source/premise observations -- for deductive and inductive observations",
)
class ObservationGet(BaseModel):
"""Schema for listing observations with optional filters"""
class ConclusionGet(BaseModel):
"""Schema for listing conclusions with optional filters."""
filters: dict[str, Any] | None = None
class Observation(BaseModel):
"""Observation response - external view of a document"""
class Conclusion(BaseModel):
"""Conclusion response - external view of a document."""
id: str
content: str
observer: str = Field(
description="The peer who made the observation",
description="The peer who made the conclusion",
serialization_alias="observer_id",
)
observed: str = Field(
description="The peer being observed", serialization_alias="observed_id"
description="The peer the conclusion is about",
serialization_alias="observed_id",
)
session_name: str = Field(serialization_alias="session_id")
created_at: datetime.datetime
model_config = ConfigDict( # pyright: ignore
from_attributes=True, populate_by_name=True
from_attributes=True,
populate_by_name=True,
)
class ObservationQuery(BaseModel):
"""Query parameters for semantic search of observations"""
class ConclusionQuery(BaseModel):
"""Query parameters for semantic search of conclusions."""
query: str = Field(..., description="Semantic search query")
top_k: int = Field(
default=10, ge=1, le=100, description="Number of results to return"
default=10,
ge=1,
le=100,
description="Number of results to return",
)
distance: float | None = Field(
default=None,
@ -560,24 +587,25 @@ class ObservationQuery(BaseModel):
description="Maximum cosine distance threshold for results",
)
filters: dict[str, Any] | None = Field(
default=None, description="Additional filters to apply"
default=None,
description="Additional filters to apply",
)
class ObservationCreate(BaseModel):
"""Schema for creating a single observation"""
class ConclusionCreate(BaseModel):
"""Schema for creating a single conclusion."""
content: Annotated[str, Field(min_length=1, max_length=65535)]
observer_id: str = Field(..., description="The peer making the observation")
observed_id: str = Field(..., description="The peer being observed")
session_id: str = Field(..., description="The session this observation relates to")
observer_id: str = Field(..., description="The peer making the conclusion")
observed_id: str = Field(..., description="The peer the conclusion is about")
session_id: str = Field(..., description="The session this conclusion relates to")
_token_count: int = PrivateAttr(default=0)
@model_validator(mode="after")
def validate_token_count(self) -> Self:
"""Validate that content doesn't exceed embedding token limit."""
encoding = tiktoken.get_encoding("cl100k_base")
encoding = tiktoken.get_encoding("o200k_base")
tokens = encoding.encode(self.content)
self._token_count = len(tokens)
@ -589,10 +617,35 @@ class ObservationCreate(BaseModel):
return self
class ObservationBatchCreate(BaseModel):
"""Schema for batch observation creation with a max of 100 observations"""
class ConclusionBatchCreate(BaseModel):
"""Schema for batch conclusion creation with a max of 100 conclusions."""
observations: list[ObservationCreate] = Field(..., min_length=1, max_length=100)
conclusions: list[ConclusionCreate] = Field(
...,
min_length=1,
max_length=100,
validation_alias=AliasChoices("conclusions", "observations"),
)
class ObservationGet(ConclusionGet):
"""Deprecated: use ConclusionGet."""
class Observation(Conclusion):
"""Deprecated: use Conclusion."""
class ObservationQuery(ConclusionQuery):
"""Deprecated: use ConclusionQuery."""
class ObservationCreate(ConclusionCreate):
"""Deprecated: use ConclusionCreate."""
class ObservationBatchCreate(ConclusionBatchCreate):
"""Deprecated: use ConclusionBatchCreate."""
class MessageSearchOptions(BaseModel):
@ -620,6 +673,10 @@ class DialecticOptions(BaseModel):
str, Field(min_length=1, max_length=10000, description="Dialectic API Prompt")
]
stream: bool = False
reasoning_level: ReasoningLevel = Field(
default="low",
description="Level of reasoning to apply: minimal, low, medium, high, or extra-high",
)
class DialecticResponse(BaseModel):
@ -689,9 +746,12 @@ class MessageBulkData(BaseModel):
workspace_name: str
class SessionDeriverStatus(BaseModel):
class SessionQueueStatus(BaseModel):
"""Status for a specific session within the processing queue."""
session_id: str | None = Field(
default=None, description="Session ID if filtered by session"
default=None,
description="Session ID if filtered by session",
)
total_work_units: int = Field(description="Total work units")
completed_work_units: int = Field(description="Completed work units")
@ -701,18 +761,29 @@ class SessionDeriverStatus(BaseModel):
pending_work_units: int = Field(description="Work units waiting to be processed")
class DeriverStatus(BaseModel):
class QueueStatus(BaseModel):
"""Aggregated processing queue status."""
total_work_units: int = Field(description="Total work units")
completed_work_units: int = Field(description="Completed work units")
in_progress_work_units: int = Field(
description="Work units currently being processed"
)
pending_work_units: int = Field(description="Work units waiting to be processed")
sessions: dict[str, SessionDeriverStatus] | None = Field(
default=None, description="Per-session status when not filtered by session"
sessions: dict[str, SessionQueueStatus] | None = Field(
default=None,
description="Per-session status when not filtered by session",
)
class SessionDeriverStatus(SessionQueueStatus):
"""Deprecated: use SessionQueueStatus."""
class DeriverStatus(QueueStatus):
"""Deprecated: use QueueStatus."""
# Dream trigger schema
class TriggerDreamRequest(BaseModel):
observer: str = Field(..., description="Observer peer name")
@ -720,6 +791,7 @@ class TriggerDreamRequest(BaseModel):
None, description="Observed peer name (defaults to observer if not specified)"
)
dream_type: DreamType = Field(..., description="Type of dream to trigger")
session_id: str = Field(..., description="Session ID to scope the dream to")
# Webhook endpoint schemas

1622
src/utils/agent_tools.py Normal file
View File

File diff suppressed because it is too large Load Diff

1711
src/utils/clients.py
View File

File diff suppressed because it is too large Load Diff

2
src/utils/config_helpers.py
View File

@ -51,7 +51,7 @@ def get_configuration(
"""
# Start with defaults
config_dict: dict[str, Any] = {
"deriver": {"enabled": True},
"deriver": {"enabled": settings.DERIVER.ENABLED},
"peer_card": {
"use": settings.PEER_CARD.ENABLED,
"create": settings.PEER_CARD.ENABLED,

15
src/utils/filter.py
View File

@ -7,7 +7,7 @@ from sqlalchemy import ColumnElement, Select, and_, case, cast, literal, not_, o
from sqlalchemy.types import Numeric
from ..exceptions import FilterError
from .formatting import parse_datetime_iso
from .formatting import ILIKE_ESCAPE_CHAR, escape_ilike_pattern, parse_datetime_iso
logger = getLogger(__name__)
@ -351,7 +351,8 @@ def _build_comparison_condition(
f"Invalid value for 'in' operator: {op_value}. Expected an iterable (list, tuple, set), got {type(op_value).__name__}"
)
elif operator in ("contains", "icontains"):
return field_accessor.ilike(f"%{op_value}%")
escaped_value = escape_ilike_pattern(str(op_value))
return field_accessor.ilike(f"%{escaped_value}%", escape=ILIKE_ESCAPE_CHAR)
return None
@ -521,11 +522,13 @@ def _build_comparison_conditions(
# For JSONB columns, use JSONB contains
condition = column.contains(op_value)
else:
# For text columns, use ILIKE
condition = column.ilike(f"%{op_value}%")
# For text columns, use ILIKE with escaped pattern
escaped_value = escape_ilike_pattern(str(op_value))
condition = column.ilike(f"%{escaped_value}%", escape=ILIKE_ESCAPE_CHAR)
elif operator == "icontains":
# Case-insensitive contains for text columns
condition = column.ilike(f"%{op_value}%")
# Case-insensitive contains for text columns with escaped pattern
escaped_value = escape_ilike_pattern(str(op_value))
condition = column.ilike(f"%{escaped_value}%", escape=ILIKE_ESCAPE_CHAR)
if condition is not None:
conditions.append(condition)

36
src/utils/formatting.py
View File

@ -2,11 +2,45 @@
Shared formatting utility functions for both dialectic and deriver modules.
This module contains helper functions for processing observations, formatting context,
and handling temporal metadata for the reasoning system.
handling temporal metadata, and string escaping for the reasoning system.
"""
from datetime import datetime, timezone
ILIKE_ESCAPE_CHAR = "\\"
def escape_ilike_pattern(text: str) -> str:
"""
Escape SQL ILIKE/LIKE pattern special characters in user-provided text.
SQL LIKE/ILIKE patterns treat '%' as "match any sequence" and '_' as
"match any single character". Without escaping, a user searching for
"100%" would match "100" followed by anything, not the literal "100%".
This function escapes these wildcards so user input is treated literally.
The escape character itself (backslash) is also escaped.
Args:
text: User-provided search text that may contain %, _, or backslash
Returns:
Escaped text safe for use in ILIKE patterns. Use with escape='\\' parameter.
Example:
>>> escape_ilike_pattern("100%")
'100\\%'
>>> escape_ilike_pattern("file_name")
'file\\_name'
>>> escape_ilike_pattern("path\\to\\file")
'path\\\\to\\\\file'
"""
return (
text.replace(ILIKE_ESCAPE_CHAR, ILIKE_ESCAPE_CHAR + ILIKE_ESCAPE_CHAR)
.replace("%", ILIKE_ESCAPE_CHAR + "%")
.replace("_", ILIKE_ESCAPE_CHAR + "_")
)
def format_datetime_utc(dt: datetime) -> str:
"""

49
src/utils/logging.py
View File

@ -165,6 +165,50 @@ def accumulate_metric(
accumulated_metrics.setdefault(task_name, []).append((label, value, unit))
def log_token_usage_metrics(
task_name: str,
input_tokens: int,
output_tokens: int,
cache_read_input_tokens: int,
cache_creation_input_tokens: int,
) -> None:
"""
Log cache-aware token usage metrics.
Args:
task_name: The task name for metric accumulation
input_tokens: Total input tokens (cached + uncached)
output_tokens: Output tokens generated
cache_read_input_tokens: Tokens read from cache (90% cheaper)
cache_creation_input_tokens: Tokens written to cache (25% more expensive)
Returns:
None
"""
accumulate_metric(task_name, "input_tokens", input_tokens, "tokens")
accumulate_metric(
task_name,
"cache_read_input_tokens",
cache_read_input_tokens,
"tokens",
)
accumulate_metric(
task_name,
"cache_creation_input_tokens",
cache_creation_input_tokens,
"tokens",
)
# Total uncached tokens (what you're paying full price for)
# = total - cache_read (those were cheap) + cache_creation (those cost 1.25x)
uncached_input_tokens = (
input_tokens - cache_read_input_tokens + cache_creation_input_tokens
)
accumulate_metric(
task_name, "uncached_input_tokens", uncached_input_tokens, "tokens"
)
accumulate_metric(task_name, "output_tokens", output_tokens, "tokens")
def log_performance_metrics(
task_slug: str,
task_name: str,
@ -221,9 +265,10 @@ def log_performance_metrics(
content_items: list[RenderableType] = [table]
if blob_metrics:
content_items.append(Text("")) # Empty line separator
for metric, value, _unit in blob_metrics:
content_items.append(Text(f"{metric}:", style="bold cyan"))
content_items.append(
Text.assemble(" ", (f"\n{metric}:", "bold"), " ")
)
content_items.append(Text(str(value)))
panel = Panel(

25
src/utils/peer_card.py
View File

@ -1,25 +0,0 @@
"""
Shared Pydantic models used by both dialectic and deriver modules.
"""
from __future__ import annotations
from pydantic import BaseModel, Field
class PeerCardQuery(BaseModel):
"""
Model for peer card query generation responses.
Contains the new peer card, or None if there are no new key observations.
The notes field is just a place for stupid models to dump useless info.
"""
card: list[str] | None = Field(
default=None,
description="Generated peer card as list of strings. None if no new useful biographical observations.",
)
notes: str | None = Field(
default=None,
description="Optional additional notes from the model; may include non-actionable info.",
)

3
src/utils/queue_payload.py
View File

@ -56,6 +56,7 @@ class DreamPayload(BasePayload):
dream_type: DreamType
observer: str
observed: str
session_name: str
class DeletionPayload(BasePayload):
@ -81,12 +82,14 @@ def create_dream_payload(
*,
observer: str,
observed: str,
session_name: str,
) -> dict[str, Any]:
"""Create a dream payload."""
return DreamPayload(
dream_type=dream_type,
observer=observer,
observed=observed,
session_name=session_name,
).model_dump(mode="json", exclude_none=True)

99
src/utils/reasoning_traces.py Normal file
View File

@ -0,0 +1,99 @@
"""
Utility for logging traces from LLM calls.
This module provides structured JSONL logging of LLM inputs/outputs.
"""
import fcntl
import json
import time
from pathlib import Path
from typing import Any
from pydantic import BaseModel
from src.config import LLMComponentSettings, settings
def get_reasoning_traces_file_path() -> Path | None:
"""Get the traces file path from settings."""
if settings.REASONING_TRACES_FILE:
return Path(settings.REASONING_TRACES_FILE)
return None
def log_reasoning_trace(
task_type: str,
llm_settings: LLMComponentSettings,
prompt: str,
response: Any,
*,
max_tokens: int | None = None,
thinking_budget_tokens: int | None = None,
reasoning_effort: str | None = None,
json_mode: bool = False,
stop_seqs: list[str] | None = None,
messages: list[dict[str, Any]] | None = None,
) -> None:
"""
Log a trace to the configured JSONL file.
Args:
task_type: Type of task (e.g., "minimal_deriver", "dialectic_chat")
llm_settings: LLM settings used for the call
prompt: The full prompt text sent to the LLM (used if messages is None)
response: HonchoLLMCallResponse object with the LLM response
max_tokens: Max output tokens setting
thinking_budget_tokens: Anthropic thinking budget (if used)
reasoning_effort: OpenAI reasoning effort (if used)
json_mode: Whether JSON mode was enabled
stop_seqs: Stop sequences used (if any)
messages: Full conversation history for multi-turn/agentic calls
"""
traces_file = get_reasoning_traces_file_path()
if not traces_file:
return
# Serialize response content - handle Pydantic models
content = response.content
if isinstance(content, BaseModel):
content = content.model_dump()
trace_entry: dict[str, Any] = {
"timestamp": time.time(),
"task_type": task_type,
"provider": llm_settings.PROVIDER,
"model": llm_settings.MODEL,
"settings": {
"max_tokens": max_tokens,
"thinking_budget_tokens": thinking_budget_tokens,
"reasoning_effort": reasoning_effort,
"json_mode": json_mode,
"stop_seqs": stop_seqs,
},
"input": {
"tokens": response.input_tokens,
},
"output": {
"content": content,
"tokens": response.output_tokens,
"finish_reasons": response.finish_reasons,
"thinking_content": response.thinking_content,
},
}
# Use messages for multi-turn/agentic calls, otherwise use prompt
if messages is not None:
trace_entry["input"]["messages"] = messages
else:
trace_entry["input"]["prompt"] = prompt
# Include tool calls if present
if hasattr(response, "tool_calls_made") and response.tool_calls_made:
trace_entry["output"]["tool_calls"] = response.tool_calls_made
# Use file locking to handle concurrent writes from multiple processes
with open(traces_file, "a") as f:
fcntl.flock(f.fileno(), fcntl.LOCK_EX)
f.write(json.dumps(trace_entry) + "\n")
fcntl.flock(f.fileno(), fcntl.LOCK_UN)

399
src/utils/representation.py
View File

@ -2,12 +2,19 @@ from collections.abc import Sequence
from datetime import datetime
from typing import Any
from pydantic import BaseModel, Field
from pydantic import BaseModel, Field, field_validator
from src import models
from src.utils.formatting import parse_datetime_iso
def _strip_microseconds_and_timezone(timestamp: datetime) -> datetime:
"""
Remove microseconds and timezone info from a datetime for stable string formatting.
"""
return timestamp.replace(microsecond=0, tzinfo=None)
def flatten_message_ids(
message_ids: list[int] | list[list[int]] | list[tuple[int, int]],
) -> list[int]:
@ -43,6 +50,7 @@ def flatten_message_ids(
class ObservationMetadata(BaseModel):
id: str = Field(default="", description="Document ID for this observation")
created_at: datetime
message_ids: list[int]
session_name: str
@ -53,13 +61,53 @@ class ExplicitObservationBase(BaseModel):
class DeductiveObservationBase(BaseModel):
source_ids: list[str] = Field(
description="Document IDs of premise observations for tree traversal",
default_factory=list,
)
premises: list[str] = Field(
description="Supporting premises or evidence for this conclusion",
description="Human-readable premise text for display",
default_factory=list,
)
conclusion: str = Field(description="The deductive conclusion")
class InductiveObservationBase(BaseModel):
"""Base model for inductive observations - patterns, generalizations, and personality insights."""
source_ids: list[str] = Field(
description="Document IDs of source observations for tree traversal",
default_factory=list,
)
sources: list[str] = Field(
description="Human-readable source text for display",
default_factory=list,
)
pattern_type: str = Field(
description="Type of pattern: 'preference', 'behavior', 'personality', 'tendency', 'correlation'",
default="pattern",
)
conclusion: str = Field(description="The inductive generalization or pattern")
confidence: str = Field(
description="Confidence level: 'high', 'medium', 'low'",
default="medium",
)
class ContradictionObservationBase(BaseModel):
"""Base model for contradiction observations - when user has made conflicting statements."""
source_ids: list[str] = Field(
description="Document IDs of the contradicting observations",
default_factory=list,
)
sources: list[str] = Field(
description="Human-readable text of the contradicting statements",
default_factory=list,
)
content: str = Field(description="Description of the contradiction")
class PromptRepresentation(BaseModel):
"""
The representation format that is used when getting structured output from an LLM.
@ -69,17 +117,26 @@ class PromptRepresentation(BaseModel):
description="Facts LITERALLY stated by the user - direct quotes or clear paraphrases only, no interpretation or inference. Example: ['The user is 25 years old', 'The user has a dog named Rover']",
default_factory=list,
)
deductive: list[DeductiveObservationBase] = Field(
description="Conclusions that MUST be true given explicit facts and premises - strict logical necessities. Each deduction should have premises and a single conclusion.",
default_factory=list,
)
@field_validator("explicit", mode="before")
@classmethod
def convert_none_to_empty_list(cls, v: Any) -> Any:
"""Convert None to empty list - handles LLMs returning null instead of []."""
if v is None:
return []
return v
class ExplicitObservation(ExplicitObservationBase, ObservationMetadata):
"""Explicit observation with content and metadata."""
def __str__(self) -> str:
return f"[{self.created_at.replace(microsecond=0)}] {self.content}"
return f"[{_strip_microseconds_and_timezone(self.created_at)}] {self.content}"
def str_with_id(self) -> str:
"""Format with ID prefix for use by agents that need to reference observations."""
id_prefix = f"[id:{self.id}] " if self.id else ""
return f"{id_prefix}[{_strip_microseconds_and_timezone(self.created_at)}] {self.content}"
def __hash__(self) -> int:
"""
@ -106,7 +163,13 @@ class DeductiveObservation(DeductiveObservationBase, ObservationMetadata):
def __str__(self) -> str:
premises_text = "\n".join(f" - {premise}" for premise in self.premises)
return f"[{self.created_at.replace(microsecond=0)}] {self.conclusion}\n{premises_text}"
return f"[{_strip_microseconds_and_timezone(self.created_at)}] {self.conclusion}\n{premises_text}"
def str_with_id(self) -> str:
"""Format with ID prefix for use by agents that need to reference observations."""
id_prefix = f"[id:{self.id}] " if self.id else ""
premises_text = "\n".join(f" - {premise}" for premise in self.premises)
return f"{id_prefix}[{_strip_microseconds_and_timezone(self.created_at)}] {self.conclusion}\n{premises_text}"
def str_no_timestamps(self) -> str:
premises_text = "\n".join(f" - {premise}" for premise in self.premises)
@ -132,6 +195,88 @@ class DeductiveObservation(DeductiveObservationBase, ObservationMetadata):
)
class InductiveObservation(InductiveObservationBase, ObservationMetadata):
"""Inductive observation with sources, pattern type, and confidence, plus metadata."""
def __str__(self) -> str:
sources_text = ""
if self.sources:
source_lines = [f" - {source}" for source in self.sources]
sources_text = "\n" + "\n".join(source_lines)
return f"[{_strip_microseconds_and_timezone(self.created_at)}] [{self.confidence}] {self.conclusion}{sources_text}"
def str_with_id(self) -> str:
"""Format with ID prefix for use by agents that need to reference observations."""
id_prefix = f"[id:{self.id}] " if self.id else ""
sources_text = ""
if self.sources:
source_lines = [f" - {source}" for source in self.sources]
sources_text = "\n" + "\n".join(source_lines)
return f"{id_prefix}[{_strip_microseconds_and_timezone(self.created_at)}] [{self.confidence}] {self.conclusion}{sources_text}"
def str_no_timestamps(self) -> str:
sources_text = ""
if self.sources:
source_lines = [f" - {source}" for source in self.sources]
sources_text = "\n" + "\n".join(source_lines)
return f"[{self.confidence}] {self.conclusion}{sources_text}"
def __hash__(self) -> int:
"""Make InductiveObservation hashable for use in sets."""
return hash((self.conclusion, self.created_at, self.session_name))
def __eq__(self, other: object) -> bool:
"""Define equality for InductiveObservation objects."""
if not isinstance(other, InductiveObservation):
return False
return (
self.conclusion == other.conclusion
and self.created_at == other.created_at
and self.session_name == other.session_name
)
class ContradictionObservation(ContradictionObservationBase, ObservationMetadata):
"""Contradiction observation - notes when user has made conflicting statements, plus metadata."""
def __str__(self) -> str:
sources_text = ""
if self.sources:
source_lines = [f" - {source}" for source in self.sources]
sources_text = "\n" + "\n".join(source_lines)
return f"[{_strip_microseconds_and_timezone(self.created_at)}] CONTRADICTION: {self.content}{sources_text}"
def str_with_id(self) -> str:
"""Format with ID prefix for use by agents that need to reference observations."""
id_prefix = f"[id:{self.id}] " if self.id else ""
sources_text = ""
if self.sources:
source_lines = [f" - {source}" for source in self.sources]
sources_text = "\n" + "\n".join(source_lines)
return f"{id_prefix}[{_strip_microseconds_and_timezone(self.created_at)}] CONTRADICTION: {self.content}{sources_text}"
def str_no_timestamps(self) -> str:
sources_text = ""
if self.sources:
source_lines = [f" - {source}" for source in self.sources]
sources_text = "\n" + "\n".join(source_lines)
return f"CONTRADICTION: {self.content}{sources_text}"
def __hash__(self) -> int:
"""Make ContradictionObservation hashable for use in sets."""
return hash((self.content, self.created_at, self.session_name))
def __eq__(self, other: object) -> bool:
"""Define equality for ContradictionObservation objects."""
if not isinstance(other, ContradictionObservation):
return False
return (
self.content == other.content
and self.created_at == other.created_at
and self.session_name == other.session_name
)
class Representation(BaseModel):
"""
A Representation is a traversable and diffable map of observations.
@ -159,12 +304,36 @@ class Representation(BaseModel):
description="Conclusions that MUST be true given explicit facts and premises - strict logical necessities. Each deduction should have premises and a single conclusion.",
default_factory=list,
)
inductive: list[InductiveObservation] = Field(
description="Patterns, generalizations, and personality insights inferred from multiple observations. Higher-level reasoning created by the Dreamer agent.",
default_factory=list,
)
contradiction: list[ContradictionObservation] = Field(
description="Conflicting statements made by the user that need clarification. The dialectic agent should surface these when relevant.",
default_factory=list,
)
def is_empty(self) -> bool:
"""
Check if the representation is empty.
"""
return len(self.explicit) == 0 and len(self.deductive) == 0
return (
len(self.explicit) == 0
and len(self.deductive) == 0
and len(self.inductive) == 0
and len(self.contradiction) == 0
)
def len(self) -> int:
"""
Return the total number of observations in the representation.
"""
return (
len(self.explicit)
+ len(self.deductive)
+ len(self.inductive)
+ len(self.contradiction)
)
def diff_representation(self, other: "Representation") -> "Representation":
"""
@ -174,6 +343,10 @@ class Representation(BaseModel):
diff = Representation()
diff.explicit = [o for o in other.explicit if o not in self.explicit]
diff.deductive = [o for o in other.deductive if o not in self.deductive]
diff.inductive = [o for o in other.inductive if o not in self.inductive]
diff.contradiction = [
o for o in other.contradiction if o not in self.contradiction
]
return diff
def merge_representation(
@ -181,7 +354,7 @@ class Representation(BaseModel):
):
"""
Merge another representation object into this one.
This will automatically deduplicate explicit and deductive observations.
This will automatically deduplicate explicit, deductive, inductive, and contradiction observations.
This *preserves order* of observations so that they retain FIFO order.
NOTE: observations with the *same* timestamp will not have order preserved.
@ -190,13 +363,19 @@ class Representation(BaseModel):
# removing duplicates by going list->set->list
self.explicit = list(set(self.explicit + other.explicit))
self.deductive = list(set(self.deductive + other.deductive))
self.inductive = list(set(self.inductive + other.inductive))
self.contradiction = list(set(self.contradiction + other.contradiction))
# sort by created_at
self.explicit.sort(key=lambda x: x.created_at)
self.deductive.sort(key=lambda x: x.created_at)
self.inductive.sort(key=lambda x: x.created_at)
self.contradiction.sort(key=lambda x: x.created_at)
if max_observations:
self.explicit = self.explicit[-max_observations:]
self.deductive = self.deductive[-max_observations:]
self.inductive = self.inductive[-max_observations:]
self.contradiction = self.contradiction[-max_observations:]
def __str__(self) -> str:
"""
@ -229,6 +408,59 @@ class Representation(BaseModel):
parts.append(f"{i}. {observation}")
parts.append("")
parts.append("INDUCTIVE:\n")
for i, observation in enumerate(self.inductive, 1):
parts.append(f"{i}. {observation}")
parts.append("")
parts.append("CONTRADICTION:\n")
for i, observation in enumerate(self.contradiction, 1):
parts.append(f"{i}. {observation}")
parts.append("")
return "\n".join(parts)
def str_with_ids(self) -> str:
"""
Format representation with observation IDs for agents that need to reference/delete observations.
Returns:
Formatted string with IDs included
Example:
EXPLICIT:
1. [id:abc123] [2025-01-01 12:00:00] The user has a dog named Rover
2. [id:def456] [2025-01-01 12:01:00] The user's dog is 5 years old
DEDUCTIVE:
1. [id:ghi789] [2025-01-01 12:01:00] Rover is 5 years old
- The user has a dog named Rover
- The user's dog is 5 years old
INDUCTIVE:
1. [id:jkl012] [2025-01-01 12:05:00] [high] User tends to be methodical
- id:abc123
- id:def456
"""
parts: list[str] = []
parts.append("EXPLICIT:\n")
for i, observation in enumerate(self.explicit, 1):
parts.append(f"{i}. {observation.str_with_id()}")
parts.append("")
parts.append("DEDUCTIVE:\n")
for i, observation in enumerate(self.deductive, 1):
parts.append(f"{i}. {observation.str_with_id()}")
parts.append("")
parts.append("INDUCTIVE:\n")
for i, observation in enumerate(self.inductive, 1):
parts.append(f"{i}. {observation.str_with_id()}")
parts.append("")
parts.append("CONTRADICTION:\n")
for i, observation in enumerate(self.contradiction, 1):
parts.append(f"{i}. {observation.str_with_id()}")
parts.append("")
return "\n".join(parts)
def str_no_timestamps(self) -> str:
@ -246,6 +478,10 @@ class Representation(BaseModel):
1. Rover is 5 years old
- The user has a dog named Rover
- The user's dog is 5 years old
INDUCTIVE:
1. [high] User tends to be methodical
- id:abc123
- id:def456
"""
parts: list[str] = []
@ -260,13 +496,26 @@ class Representation(BaseModel):
parts.append(f"{i}. {observation.str_no_timestamps()}")
parts.append("")
parts.append("INDUCTIVE:\n")
for i, observation in enumerate(self.inductive, 1):
parts.append(f"{i}. {observation.str_no_timestamps()}")
parts.append("")
parts.append("CONTRADICTION:\n")
for i, observation in enumerate(self.contradiction, 1):
parts.append(f"{i}. {observation.str_no_timestamps()}")
parts.append("")
return "\n".join(parts)
def format_as_markdown(self) -> str:
def format_as_markdown(self, include_ids: bool = False) -> str:
"""
Format a Representation object as markdown.
NOTE: we always strip subsecond precision from the timestamps.
Args:
include_ids: If True, include observation IDs for use with get_reasoning_chain
Returns:
Formatted markdown string
"""
@ -274,21 +523,59 @@ class Representation(BaseModel):
parts: list[str] = []
# Add explicit observations
parts.append("## Explicit Observations\n")
for i, obs in enumerate(self.explicit, 1):
parts.append(f"{i}. {obs}")
parts.append("")
if self.explicit:
parts.append("## Explicit Observations\n")
for obs in self.explicit:
# Don't need IDs for explicit as these are the lowest level of reasoning.
# id_prefix = f"[id:{obs.id}] " if include_ids and obs.id else ""
parts.append(f"{obs}")
parts.append("")
# Add deductive observations
parts.append("## Deductive Observations\n")
for i, obs in enumerate(self.deductive, 1):
parts.append(f"{i}. **Conclusion**: {obs.conclusion}")
if obs.premises:
parts.append(" **Premises**:")
for premise in obs.premises:
parts.append(f" - {premise}")
if self.deductive:
parts.append("## Deductive Observations\n")
for obs in self.deductive:
id_prefix = f"[id:{obs.id}] " if include_ids and obs.id else ""
timestamp = _strip_microseconds_and_timezone(obs.created_at)
parts.append(f"{id_prefix}[{timestamp}] {obs.conclusion}")
if obs.premises:
parts.append(" Premises:")
for premise in obs.premises:
parts.append(f" - {premise}")
parts.append("")
parts.append("")
# Add inductive observations
if self.inductive:
parts.append("## Inductive Observations\n")
for obs in self.inductive:
id_prefix = f"[id:{obs.id}] " if include_ids and obs.id else ""
parts.append(
f"{id_prefix} **Pattern** [{obs.confidence}]: {obs.conclusion}"
)
if obs.pattern_type:
parts.append(f" **Type**: {obs.pattern_type}")
if obs.sources:
parts.append(" **Sources**:")
for source in obs.sources[:5]:
parts.append(f" - {source}")
if len(obs.sources) > 5:
parts.append(f" - ... and {len(obs.sources) - 5} more")
parts.append("")
parts.append("")
# Add contradiction observations
if self.contradiction:
parts.append("## Contradictions\n")
for obs in self.contradiction:
id_prefix = f"[id:{obs.id}] " if include_ids and obs.id else ""
parts.append(f"{id_prefix} **CONTRADICTION**: {obs.content}")
if obs.sources:
parts.append(" **Conflicting statements**:")
for source in obs.sources:
parts.append(f" - {source}")
parts.append("")
parts.append("")
parts.append("")
return "\n".join(parts)
@ -297,6 +584,7 @@ class Representation(BaseModel):
return cls(
explicit=[
ExplicitObservation(
id=doc.id,
created_at=_safe_datetime_from_metadata(
doc.internal_metadata, doc.created_at
),
@ -311,6 +599,7 @@ class Representation(BaseModel):
],
deductive=[
DeductiveObservation(
id=doc.id,
created_at=_safe_datetime_from_metadata(
doc.internal_metadata, doc.created_at
),
@ -319,11 +608,50 @@ class Representation(BaseModel):
doc.internal_metadata.get("message_ids", [])
),
session_name=doc.session_name,
# Support both top-level and metadata locations for backward compatibility
source_ids=doc.source_ids
or doc.internal_metadata.get("premise_ids", []),
premises=doc.internal_metadata.get("premises", []),
)
for doc in documents
if doc.level == "deductive"
],
inductive=[
InductiveObservation(
id=doc.id,
created_at=_safe_datetime_from_metadata(
doc.internal_metadata, doc.created_at
),
conclusion=doc.content,
message_ids=doc.internal_metadata.get("message_ids", []),
session_name=doc.session_name,
# Support both top-level and metadata locations for backward compatibility
source_ids=doc.source_ids
or doc.internal_metadata.get("source_ids", []),
sources=doc.internal_metadata.get("sources", []),
pattern_type=doc.internal_metadata.get("pattern_type", "pattern"),
confidence=doc.internal_metadata.get("confidence", "medium"),
)
for doc in documents
if doc.level == "inductive"
],
contradiction=[
ContradictionObservation(
id=doc.id,
created_at=_safe_datetime_from_metadata(
doc.internal_metadata, doc.created_at
),
content=doc.content,
message_ids=doc.internal_metadata.get("message_ids", []),
session_name=doc.session_name,
# Support both top-level and metadata locations for backward compatibility
source_ids=doc.source_ids
or doc.internal_metadata.get("source_ids", []),
sources=doc.internal_metadata.get("sources", []),
)
for doc in documents
if doc.level == "contradiction"
],
)
@classmethod
@ -334,6 +662,7 @@ class Representation(BaseModel):
session_name: str,
created_at: datetime,
) -> "Representation":
"""Convert PromptRepresentation to Representation."""
return cls(
explicit=[
ExplicitObservation(
@ -344,16 +673,8 @@ class Representation(BaseModel):
)
for e in prompt_representation.explicit
],
deductive=[
DeductiveObservation(
conclusion=d.conclusion,
created_at=created_at,
message_ids=message_ids,
session_name=session_name,
premises=d.premises,
)
for d in prompt_representation.deductive
],
deductive=[],
inductive=[],
)
@ -362,14 +683,16 @@ def _safe_datetime_from_metadata(
) -> datetime:
message_created_at = internal_metadata.get("message_created_at")
if message_created_at is None:
return fallback_datetime.replace(microsecond=0)
return _strip_microseconds_and_timezone(fallback_datetime)
if isinstance(message_created_at, str):
try:
return parse_datetime_iso(message_created_at)
return _strip_microseconds_and_timezone(
parse_datetime_iso(message_created_at)
)
except ValueError:
return fallback_datetime.replace(microsecond=0)
return _strip_microseconds_and_timezone(fallback_datetime)
if isinstance(message_created_at, datetime):
return message_created_at.replace(microsecond=0)
return fallback_datetime.replace(microsecond=0)
return _strip_microseconds_and_timezone(message_created_at)
return _strip_microseconds_and_timezone(fallback_datetime)

13
src/utils/search.py
View File

@ -17,6 +17,7 @@ from src.embedding_client import embedding_client
from src.exceptions import ValidationException
from src.models import session_peers_table
from src.utils.filter import apply_filter
from src.utils.formatting import ILIKE_ESCAPE_CHAR, escape_ilike_pattern
T = TypeVar("T")
@ -123,9 +124,14 @@ async def _fulltext_search(
re.search(r'[~`!@#$%^&*()_+=\[\]{};\':"\\|,.<>/?-]', query)
)
# Escape ILIKE pattern characters to treat user input literally
escaped_query = escape_ilike_pattern(query)
if has_special_chars:
# For queries with special characters, use exact string matching (ILIKE)
search_condition = models.Message.content.ilike(f"%{query}%")
search_condition = models.Message.content.ilike(
f"%{escaped_query}%", escape=ILIKE_ESCAPE_CHAR
)
fulltext_query = stmt.where(search_condition).order_by(
models.Message.created_at.desc()
)
@ -137,7 +143,10 @@ async def _fulltext_search(
# Combine FTS with ILIKE as fallback for better coverage
combined_condition = or_(
fts_condition, models.Message.content.ilike(f"%{query}%")
fts_condition,
models.Message.content.ilike(
f"%{escaped_query}%", escape=ILIKE_ESCAPE_CHAR
),
)
fulltext_query = stmt.where(combined_condition).order_by(

18
src/utils/summarizer.py
View File

@ -19,7 +19,7 @@ from src.models import Message
from src.utils.clients import HonchoLLMCallResponse, honcho_llm_call
from src.utils.formatting import utc_now_iso
from src.utils.logging import accumulate_metric, conditional_observe
from src.utils.tokens import estimate_tokens, track_input_tokens
from src.utils.tokens import estimate_tokens, track_deriver_input_tokens
from .. import crud, models
@ -424,20 +424,20 @@ async def _create_and_save_summary(
else:
prompt_tokens = estimate_long_summary_prompt_tokens()
track_input_tokens(
task_type="summary",
track_deriver_input_tokens(
task_type=prometheus.DeriverTaskTypes.SUMMARY,
components={
"prompt": prompt_tokens,
"messages": messages_tokens,
"previous_summary": previous_summary_tokens,
prometheus.DeriverComponents.PROMPT: prompt_tokens,
prometheus.DeriverComponents.MESSAGES: messages_tokens,
prometheus.DeriverComponents.PREVIOUS_SUMMARY: previous_summary_tokens,
},
)
# Track output tokens
prometheus.DERIVER_TOKENS_PROCESSED.labels(
task_type="summary",
token_type="output", # nosec B106
component="total",
task_type=prometheus.DeriverTaskTypes.SUMMARY.value,
token_type=prometheus.TokenTypes.OUTPUT.value,
component=prometheus.DeriverComponents.OUTPUT_TOTAL.value,
).inc(new_summary["token_count"])
# Save summary to database

13
src/utils/tokens.py
View File

@ -17,17 +17,20 @@ def estimate_tokens(text: str | list[str] | None) -> int:
return len(text) // 4
def track_input_tokens(task_type: str, components: dict[str, int]) -> None:
def track_deriver_input_tokens(
task_type: prometheus.DeriverTaskTypes,
components: dict[prometheus.DeriverComponents, int],
) -> None:
"""
Helper method to track input token components for a given task type.
Args:
task_type: The type of task (e.g., "representation", "peer_card", "summary")
task_type: The type of task
components: Dict mapping component names to token counts
"""
for component, token_count in components.items():
prometheus.DERIVER_TOKENS_PROCESSED.labels(
task_type=task_type,
token_type="input", # nosec B106
component=component,
task_type=task_type.value,
token_type=prometheus.TokenTypes.INPUT.value,
component=component.value,
).inc(token_count)

2
src/utils/types.py
View File

@ -2,4 +2,4 @@ from typing import Literal
SupportedProviders = Literal["anthropic", "openai", "google", "groq", "custom", "vllm"]
TaskType = Literal["webhook", "summary", "representation", "dream", "deletion"]
DocumentLevel = Literal["explicit", "deductive"]
DocumentLevel = Literal["explicit", "deductive", "inductive", "contradiction"]

26
tests/alembic/revisions/test_f1a2b3c4d5e6_add_reasoning_tree_columns.py Normal file
View File

@ -0,0 +1,26 @@
"""Hooks for revision f1a2b3c4d5e6 (reasoning tree columns)."""
from __future__ import annotations
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
@register_before_upgrade("f1a2b3c4d5e6")
def prepare_reasoning_tree_columns(verifier: MigrationVerifier) -> None:
verifier.assert_column_exists("documents", "source_ids", exists=False)
verifier.assert_indexes_not_exist(
[
("documents", "ix_documents_source_ids_gin"),
]
)
@register_after_upgrade("f1a2b3c4d5e6")
def verify_reasoning_tree_columns(verifier: MigrationVerifier) -> None:
verifier.assert_column_exists("documents", "source_ids", nullable=True)
verifier.assert_indexes_exist(
[
("documents", "ix_documents_source_ids_gin"),
]
)

683
tests/bench/beam.py
View File

@ -47,8 +47,7 @@ Optional arguments:
```
--context-length: Context length subset to test (100K, 500K, 1M, 10M) (default: 100K)
--conversation-ids: Comma-separated list of conversation IDs to test (default: all in context length)
--anthropic-api-key: Anthropic API key for response judging (can be set in .env as LLM_ANTHROPIC_API_KEY)
--timeout: Timeout for deriver queue to empty in seconds (default: 10 minutes (600s))
--timeout: Timeout for deriver queue to empty in seconds (default: 10 minutes (600s))
--base-api-port: Base port for Honcho API instances (default: 8000)
--pool-size: Number of Honcho instances in the pool (default: 1)
--batch-size: Number of conversations to run concurrently in each batch (default: 1)
@ -58,14 +57,14 @@ Optional arguments:
```
## Other notes
- Judge is Claude Sonnet 4.5
- Judge uses OpenRouter (configured via LLM_OPENAI_COMPATIBLE_API_KEY and LLM_OPENAI_COMPATIBLE_BASE_URL in tests/bench/.env)
- Default judge model is anthropic/claude-sonnet-4.5 (can be overridden with BEAM_JUDGE_MODEL env var)
- Evaluation follows the paper's nugget-based methodology with 0/0.5/1 scoring
- Event ordering uses Kendall tau-b coefficient
"""
import argparse
import asyncio
import json
import logging
import os
import time
@ -73,53 +72,33 @@ from datetime import datetime
from pathlib import Path
from typing import Any, cast
import tiktoken
from anthropic import AsyncAnthropic
from anthropic.types import MessageParam, ToolParam
from dotenv import load_dotenv
from honcho import AsyncHoncho
from honcho.async_client.session import SessionPeerConfig
from honcho_core.types.workspaces.sessions.message_create_param import (
MessageCreateParam,
)
from scipy.stats import kendalltau # pyright: ignore[reportUnknownVariableType]
from typing_extensions import TypedDict
from openai import AsyncOpenAI
from src.config import settings
from src.utils.metrics_collector import MetricsCollector
load_dotenv()
from .beam_common import (
ConversationResult,
QuestionResult,
calculate_ability_scores,
format_duration,
generate_json_summary,
judge_event_ordering,
judge_nugget_based,
list_conversations,
load_conversation,
print_summary,
)
class QuestionResult(TypedDict):
"""Type definition for question evaluation results."""
question: str
answer: str | None
actual_response: str
memory_ability: str
rubric: list[str]
nugget_scores: list[dict[str, Any]] | None
score: float
passed: bool
reasoning: str
class ConversationResult(TypedDict):
"""Type definition for conversation execution results."""
conversation_id: str
context_length: str
workspace_id: str
total_turns: int
total_messages: int
question_results: list[QuestionResult]
ability_scores: dict[str, float]
overall_score: float
error: str | None
start_time: float
end_time: float
duration_seconds: float
# Load .env from bench directory
bench_dir = Path(__file__).parent
load_dotenv(bench_dir / ".env")
class BEAMRunner:
@ -132,7 +111,6 @@ class BEAMRunner:
data_dir: Path,
base_api_port: int = 8000,
pool_size: int = 1,
anthropic_api_key: str | None = None,
timeout_seconds: int | None = None,
cleanup_workspace: bool = True,
use_get_context: bool = False,
@ -144,7 +122,6 @@ class BEAMRunner:
data_dir: Path to the BEAM data directory
base_api_port: Base port for Honcho API instances (default: 8000)
pool_size: Number of Honcho instances in the pool (default: 1)
anthropic_api_key: Anthropic API key for judging responses
timeout_seconds: Timeout for deriver queue in seconds
cleanup_workspace: If True, delete workspace after executing conversation
use_get_context: If True, use get_context + judge LLM instead of dialectic .chat endpoint
@ -152,7 +129,6 @@ class BEAMRunner:
self.data_dir: Path = data_dir
self.base_api_port: int = base_api_port
self.pool_size: int = pool_size
self.anthropic_api_key: str | None = anthropic_api_key
self.timeout_seconds: int = (
timeout_seconds if timeout_seconds is not None else 600
)
@ -175,15 +151,26 @@ class BEAMRunner:
logging.getLogger("httpx").setLevel(logging.ERROR)
logging.getLogger("httpcore").setLevel(logging.ERROR)
if self.anthropic_api_key:
self.anthropic_client: AsyncAnthropic = AsyncAnthropic(
api_key=self.anthropic_api_key
# Initialize OpenRouter client for judging
openrouter_api_key = os.getenv("LLM_OPENAI_COMPATIBLE_API_KEY")
openrouter_base_url = os.getenv(
"LLM_OPENAI_COMPATIBLE_BASE_URL", "https://openrouter.ai/api/v1"
)
if not openrouter_api_key:
raise ValueError(
"LLM_OPENAI_COMPATIBLE_API_KEY is not set in tests/bench/.env"
)
else:
api_key = os.getenv("LLM_ANTHROPIC_API_KEY")
if not api_key:
raise ValueError("LLM_ANTHROPIC_API_KEY is not set")
self.anthropic_client = AsyncAnthropic(api_key=api_key)
self.openrouter_client: AsyncOpenAI = AsyncOpenAI(
api_key=openrouter_api_key,
base_url=openrouter_base_url,
)
# Model to use for judging (OpenRouter format)
self.judge_model: str = os.getenv(
"BEAM_JUDGE_MODEL", "anthropic/claude-sonnet-4.5"
)
def get_honcho_url_for_index(self, conversation_index: int) -> str:
"""
@ -199,76 +186,6 @@ class BEAMRunner:
port = self.base_api_port + instance_id
return f"http://localhost:{port}"
def _format_duration(self, total_seconds: float) -> str:
"""Format a duration in seconds into a human-readable string."""
minutes = int(total_seconds // 60)
if minutes > 0:
seconds_rounded = int(round(total_seconds - minutes * 60))
if seconds_rounded == 60:
minutes += 1
seconds_rounded = 0
return f"{minutes}m{seconds_rounded:02d}s"
return f"{total_seconds:.2f}s"
def _calculate_tokens(self, text: str) -> int:
"""Calculate tokens for a given text."""
tokenizer = tiktoken.get_encoding("cl100k_base")
try:
return len(
tokenizer.encode(
text,
disallowed_special=(
tokenizer.special_tokens_set - {"<|endoftext|>"}
),
)
)
except Exception:
return len(text) // 4
def load_conversation(
self, context_length: str, conversation_id: str
) -> dict[str, Any]:
"""
Load a BEAM conversation from the data directory.
Args:
context_length: Context length (100K, 500K, 1M, 10M)
conversation_id: Conversation ID
Returns:
Dictionary containing conversation data and probing questions
"""
conv_dir = self.data_dir / context_length / conversation_id
# Load chat data
chat_file = conv_dir / "chat.json"
with open(chat_file) as f:
chat_data = json.load(f)
# Load probing questions
questions_file = conv_dir / "probing_questions" / "probing_questions.json"
with open(questions_file) as f:
questions_data = json.load(f)
return {"chat": chat_data, "questions": questions_data}
def list_conversations(self, context_length: str) -> list[str]:
"""
List all conversation IDs for a given context length.
Args:
context_length: Context length (100K, 500K, 1M, 10M)
Returns:
List of conversation ID strings
"""
context_dir = self.data_dir / context_length
return [
d.name
for d in sorted(context_dir.iterdir())
if d.is_dir() and d.name.isdigit()
]
async def create_honcho_client(
self, workspace_id: str, honcho_url: str
) -> AsyncHoncho:
@ -311,269 +228,69 @@ class BEAMRunner:
return False
await asyncio.sleep(1)
async def judge_nugget_based(
async def trigger_dream_and_wait(
self,
question: str,
rubric: list[str],
actual_response: str,
memory_ability: str,
) -> dict[str, Any]:
honcho_client: AsyncHoncho,
workspace_id: str,
observer: str,
observed: str | None = None,
session_id: str | None = None,
) -> bool:
"""
Use an LLM to judge a response using nugget-based evaluation.
Trigger a dream task and wait for it to complete.
Args:
question: The question asked
rubric: List of nuggets (atomic criteria) to check
actual_response: Actual response from Honcho
memory_ability: The memory ability being tested
honcho_client: Honcho client instance
workspace_id: Workspace identifier
observer: Observer peer name
observed: Observed peer name (defaults to observer)
session_id: Session ID to scope the dream to
Returns:
Judgment result with nugget scores and overall score
True if dream completed successfully, False on timeout
"""
import httpx
observed = observed or observer
honcho_url = self.get_honcho_url_for_index(0)
url = f"{honcho_url}/v2/workspaces/{workspace_id}/trigger_dream"
payload: dict[str, Any] = {
"observer": observer,
"observed": observed,
"dream_type": "omni",
"session_id": session_id or f"{workspace_id}_session",
}
# Trigger the dream via API
try:
# Build the nugget evaluation prompt
nuggets_formatted = "\n".join(
[f"{i + 1}. {nugget}" for i, nugget in enumerate(rubric)]
)
system_prompt = f"""You are an expert judge evaluating AI responses for the {memory_ability} memory ability in the BEAM benchmark.
Your task is to evaluate whether the AI's response satisfies each atomic criterion (nugget) from the rubric.
SCORING INSTRUCTIONS:
For each nugget, assign a score:
- 1.0: The response fully satisfies this criterion
- 0.5: The response partially satisfies this criterion
- 0.0: The response does not satisfy this criterion
Be strict but fair in your evaluation. Focus on whether the response contains the required information or demonstrates the required behavior.
Use the `evaluate_response` tool to submit your evaluation."""
user_prompt = f"""Question: "{question}"
Rubric (atomic criteria to check):
{nuggets_formatted}
Actual Response: "{actual_response}"
Evaluate the response against each nugget in the rubric. Provide a score for each nugget and calculate the overall score as the average of all nugget scores."""
tool_definition: ToolParam = {
"name": "evaluate_response",
"description": "Submit the evaluation results for the response based on the rubric.",
"input_schema": {
"type": "object",
"properties": {
"nugget_scores": {
"type": "array",
"items": {
"type": "object",
"properties": {
"nugget_index": {"type": "integer"},
"score": {"type": "number"},
"reasoning": {"type": "string"},
},
"required": ["nugget_index", "score", "reasoning"],
},
},
"overall_score": {"type": "number"},
"overall_reasoning": {"type": "string"},
},
"required": ["nugget_scores", "overall_score", "overall_reasoning"],
},
}
response = await self.anthropic_client.messages.create(
model="claude-sonnet-4-5",
max_tokens=2000,
temperature=0.0,
system=system_prompt,
messages=[
{
"role": "user",
"content": user_prompt,
}
],
tools=[tool_definition],
tool_choice={"type": "tool", "name": "evaluate_response"},
)
if not response.content:
raise ValueError("Anthropic returned empty response")
# Find the tool use block
tool_use_block = next(
(block for block in response.content if block.type == "tool_use"), None
)
if not tool_use_block:
raise ValueError("No tool use block found in response")
judgment: object = tool_use_block.input
if not isinstance(judgment, dict):
raise ValueError(f"Tool input is not a dictionary: {type(judgment)}")
return cast(dict[str, Any], judgment)
except Exception as e:
self.logger.error(f"Error judging response: {e}")
# Fallback to simple 0 score
return {
"nugget_scores": [
{"nugget_index": i + 1, "score": 0.0, "reasoning": f"Error: {e}"}
for i in range(len(rubric))
],
"overall_score": 0.0,
"overall_reasoning": f"Evaluation failed due to error: {e}",
}
async def judge_event_ordering(
self, question: str, rubric: list[str], actual_response: str
) -> dict[str, Any]:
"""
Judge event ordering questions using Kendall tau-b coefficient.
Args:
question: The question asked
rubric: List of expected events in correct order
actual_response: Actual response from Honcho
Returns:
Judgment with Kendall tau-b score
"""
try:
# First, extract the events mentioned in the response
system_prompt = """You are an expert at extracting ordered lists of events or items from text.
Your task is to extract the ordered list of events/items mentioned in a response.
Use the `extract_ordered_events` tool to submit the extracted list."""
user_prompt = f"""Question: "{question}"
Response: "{actual_response}"
Extract the ordered list of events or items mentioned in the response. Preserve the order as stated in the response."""
tool_definition: ToolParam = {
"name": "extract_ordered_events",
"description": "Submit the ordered list of events extracted from the response.",
"input_schema": {
"type": "object",
"properties": {
"extracted_events": {
"type": "array",
"items": {"type": "string"},
},
},
"required": ["extracted_events"],
},
}
response = await self.anthropic_client.messages.create(
model="claude-sonnet-4-5",
max_tokens=1000,
temperature=0.0,
system=system_prompt,
messages=[
{
"role": "user",
"content": user_prompt,
}
],
tools=[tool_definition],
tool_choice={"type": "tool", "name": "extract_ordered_events"},
)
if not response.content:
raise ValueError("Anthropic returned empty response")
# Find the tool use block
tool_use_block = next(
(block for block in response.content if block.type == "tool_use"), None
)
if not tool_use_block:
raise ValueError("No tool use block found in response")
extracted = tool_use_block.input
if not isinstance(extracted, dict):
raise ValueError(f"Tool input is not a dictionary: {type(extracted)}")
extracted_dict = cast(dict[str, Any], extracted)
raw_events: list[str] = extracted_dict.get("extracted_events", [])
extracted_events = [str(e) for e in raw_events]
# Now compute alignment and Kendall tau-b
# Match extracted events to rubric events using LLM equivalence
alignment = self._align_events(rubric, extracted_events)
# Compute Kendall tau-b
tau: float
if kendalltau is None:
self.logger.warning(
"scipy not installed, cannot compute Kendall tau-b. Install with: uv pip install scipy"
async with httpx.AsyncClient() as client:
response = await client.post(
url,
json=payload,
timeout=30.0,
)
tau = 0.0
elif len(alignment) < 2:
tau = 0.0
else:
# Create rank lists
expected_ranks = list(range(len(alignment)))
actual_ranks = [alignment[i] for i in range(len(alignment))]
result_tuple: Any = kendalltau(expected_ranks, actual_ranks)
# kendalltau returns a tuple, first element is the tau coefficient
tau_value: Any = result_tuple[0]
# Handle the return type properly - convert to float
try:
tau = float(tau_value)
if tau != tau: # Check for NaN
tau = 0.0
except (TypeError, ValueError):
tau = 0.0
return {
"kendall_tau_b": tau,
"extracted_events": extracted_events,
"alignment": alignment,
"overall_score": (tau + 1) / 2, # Normalize to [0, 1]
"overall_reasoning": f"Kendall tau-b coefficient: {tau:.3f}. Extracted {len(extracted_events)} events from response.",
}
if response.status_code != 204:
print(
f"[{workspace_id}] ERROR: Dream trigger failed with status {response.status_code}"
)
print(f"[{workspace_id}] Response body: {response.text}")
return False
except Exception as e:
self.logger.error(f"Error in event ordering evaluation: {e}")
return {
"kendall_tau_b": 0.0,
"extracted_events": [],
"alignment": [],
"overall_score": 0.0,
"overall_reasoning": f"Evaluation failed due to error: {e}",
}
print(f"[{workspace_id}] ERROR: Dream trigger exception: {e}")
return False
def _align_events(
self, expected_events: list[str], extracted_events: list[str]
) -> list[int]:
"""
Align extracted events with expected events using LLM equivalence detection.
print(f"[{workspace_id}] Dream triggered for {observer}/{observed}")
Returns a list of indices mapping extracted events to expected events.
"""
# For each extracted event, find the best match in expected events
alignment: list[int] = []
for extracted in extracted_events:
best_match_idx: int = -1
for i, expected in enumerate(expected_events):
# Use simple string matching for now (can be enhanced with LLM)
if (
expected.lower() in extracted.lower()
or extracted.lower() in expected.lower()
):
best_match_idx = i
break
if best_match_idx >= 0:
alignment.append(best_match_idx)
return alignment
# Wait for dream queue to empty
print(f"[{workspace_id}] Waiting for dream to complete...")
await asyncio.sleep(2) # Give time for dream to be enqueued
success = await self.wait_for_deriver_queue_empty(honcho_client)
if success:
print(f"[{workspace_id}] Dream queue empty")
else:
print(f"[{workspace_id}] Dream queue timeout")
return success
async def _process_single_question(
self,
@ -597,26 +314,48 @@ Extract the ordered list of events or items mentioned in the response. Preserve
print(f" [{ability}] Q{q_idx + 1}: {question[:100]}...")
# Execute question using dialectic
if self.use_get_context:
# For instruction_following, always use get_context + OpenRouter API
# so the LLM can follow user-specified instructions from Honcho context
if self.use_get_context or ability == "instruction_following":
context = await session.get_context(
summary=True,
peer_target="user",
last_user_message=question,
)
context_messages = context.to_anthropic(assistant="assistant")
context_messages = context.to_openai(assistant="assistant")
context_messages.append({"role": "user", "content": question})
response = await self.anthropic_client.messages.create(
model="claude-sonnet-4-5",
max_tokens=2048,
messages=cast(list[MessageParam], context_messages),
# For instruction_following, add a system prompt that tells the LLM
# to follow any stored user preferences/instructions in the context
system_prompt = None
if ability == "instruction_following":
system_prompt = """You are a helpful assistant with memory of the user's preferences and instructions from previous conversations.
The context provided includes observations about the user, which may contain their stated preferences, instructions, or requirements for how you should respond.
IMPORTANT: You MUST follow any instructions or preferences the user has previously stated. For example:
- If the user said "always include X when discussing Y", you must include X when discussing Y
- If the user said "I prefer responses that are Z", format your response accordingly
- If the user gave any standing instructions, follow them
Review the context carefully for any such instructions before responding."""
# Prepare messages: OpenAI format uses system role in messages array
messages: list[dict[str, Any]] = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
messages.extend(cast(list[dict[str, Any]], context_messages))
response = await self.openrouter_client.chat.completions.create(
model=self.judge_model,
max_tokens=settings.DIALECTIC.MAX_OUTPUT_TOKENS,
messages=cast(Any, messages), # type: ignore[arg-type]
)
if not response.content:
if not response.choices or not response.choices[0].message:
actual_response = ""
else:
content_block = response.content[0]
actual_response = getattr(content_block, "text", "")
actual_response = response.choices[0].message.content or ""
else:
actual_response = await user_peer.chat(question)
actual_response = (
@ -625,13 +364,21 @@ Extract the ordered list of events or items mentioned in the response. Preserve
# Judge response based on memory ability
if ability == "event_ordering":
judgment = await self.judge_event_ordering(
question, rubric, actual_response
judgment = await judge_event_ordering(
self.openrouter_client,
self.judge_model,
question,
rubric,
actual_response,
)
nugget_scores = None
else:
judgment = await self.judge_nugget_based(
question, rubric, actual_response, ability
judgment = await judge_nugget_based(
self.openrouter_client,
self.judge_model,
question,
rubric,
actual_response,
)
nugget_scores = judgment.get("nugget_scores")
@ -652,6 +399,15 @@ Extract the ordered list of events or items mentioned in the response. Preserve
status = "PASS" if score >= 0.5 else "FAIL"
print(f" [{ability}] Q{q_idx + 1} Score: {score:.2f} [{status}]")
if score < 0.5 and reasoning:
print(f" Reasoning: {reasoning}")
if rubric:
print(" Rubric:")
for i, rubric_item in enumerate(rubric, 1):
print(f" {i}. {rubric_item}")
if answer:
print(f" Ideal Response: {answer}")
print(f" Our Response: {actual_response}")
return question_result
@ -698,7 +454,9 @@ Extract the ordered list of events or items mentioned in the response. Preserve
try:
# Load conversation data
conv_data = self.load_conversation(context_length, conversation_id)
conv_data = load_conversation(
self.data_dir, context_length, conversation_id
)
chat_data = conv_data["chat"]
questions_data = conv_data["questions"]
@ -812,11 +570,22 @@ Extract the ordered list of events or items mentioned in the response. Preserve
f"\n[{workspace_id}] ERROR: Deriver queue timeout after {self.timeout_seconds}s"
)
print(
f"[{workspace_id}] Failed to complete in {self._format_duration(result['duration_seconds'])}"
f"[{workspace_id}] Failed to complete in {format_duration(result['duration_seconds'])}"
)
return result
print(f"[{workspace_id}] Deriver queue empty. Executing questions...")
print(f"[{workspace_id}] Deriver queue empty. Triggering dream...")
# Single orchestrated dream handles all reasoning types
dream_success = await self.trigger_dream_and_wait(
honcho_client,
workspace_id,
observer="user",
session_id=session_id,
)
if not dream_success:
print(f"[{workspace_id}] Warning: Dream did not complete")
print(f"[{workspace_id}] Dream completed. Executing questions...")
# Execute questions for each memory ability
question_tasks: list[Any] = []
@ -843,15 +612,9 @@ Extract the ordered list of events or items mentioned in the response. Preserve
result["question_results"] = list(results)
# Calculate ability scores
ability_totals: dict[str, list[float]] = {}
for qr in result["question_results"]:
ability = qr["memory_ability"]
if ability not in ability_totals:
ability_totals[ability] = []
ability_totals[ability].append(qr["score"])
for ability, scores in ability_totals.items():
result["ability_scores"][ability] = sum(scores) / len(scores)
result["ability_scores"] = calculate_ability_scores(
result["question_results"]
)
# Calculate overall score
if result["ability_scores"]:
@ -871,7 +634,7 @@ Extract the ordered list of events or items mentioned in the response. Preserve
result["duration_seconds"] = result["end_time"] - result["start_time"]
print(
f"\n[{workspace_id}] Completed in {self._format_duration(result['duration_seconds'])}"
f"\n[{workspace_id}] Completed in {format_duration(result['duration_seconds'])}"
)
print(f"Overall Score: {result['overall_score']:.3f}")
@ -942,102 +705,6 @@ Extract the ordered list of events or items mentioned in the response. Preserve
return all_results, overall_duration
def print_summary(
self, results: list[ConversationResult], total_elapsed_seconds: float
) -> None:
"""Print a summary of all test results."""
print(f"\n{'=' * 80}")
print("BEAM BENCHMARK EXECUTION SUMMARY")
print(f"{'=' * 80}")
total_conversations = len(results)
total_questions = sum(len(r["question_results"]) for r in results)
print(f"Total Conversations: {total_conversations}")
print(f"Total Questions: {total_questions}")
print(f"Total Test Time: {self._format_duration(total_elapsed_seconds)}")
# Calculate average scores by ability
ability_scores: dict[str, list[float]] = {}
for result in results:
for ability, score in result["ability_scores"].items():
if ability not in ability_scores:
ability_scores[ability] = []
ability_scores[ability].append(score)
print("\nAverage Scores by Memory Ability:")
for ability, scores in sorted(ability_scores.items()):
avg_score = sum(scores) / len(scores)
print(f" {ability:30s}: {avg_score:.3f}")
# Overall average
overall_scores = [r["overall_score"] for r in results]
overall_avg = (
sum(overall_scores) / len(overall_scores) if overall_scores else 0.0
)
print(f"\n{'Overall Average Score':30s}: {overall_avg:.3f}")
print(f"{'=' * 80}")
def generate_json_summary(
self,
results: list[ConversationResult],
context_length: str,
total_elapsed_seconds: float,
output_file: Path,
) -> None:
"""Generate a comprehensive JSON summary of test results."""
# Calculate summary statistics
total_conversations = len(results)
total_questions = sum(len(r["question_results"]) for r in results)
# Calculate average scores by ability
ability_scores: dict[str, list[float]] = {}
for result in results:
for ability, score in result["ability_scores"].items():
if ability not in ability_scores:
ability_scores[ability] = []
ability_scores[ability].append(score)
ability_averages = {
ability: sum(scores) / len(scores)
for ability, scores in ability_scores.items()
}
# Overall average
overall_scores = [r["overall_score"] for r in results]
overall_avg = (
sum(overall_scores) / len(overall_scores) if overall_scores else 0.0
)
summary = {
"metadata": {
"context_length": context_length,
"execution_timestamp": datetime.now().isoformat(),
"runner_version": "1.0.0",
"base_api_port": self.base_api_port,
"pool_size": self.pool_size,
"timeout_seconds": self.timeout_seconds,
"deriver_settings": settings.DERIVER.model_dump(),
"dialectic_settings": settings.DIALECTIC.model_dump(),
},
"summary_statistics": {
"total_conversations": total_conversations,
"total_questions": total_questions,
"overall_average_score": overall_avg,
"ability_averages": ability_averages,
},
"timing": {
"total_duration_seconds": total_elapsed_seconds,
},
"detailed_results": results,
}
output_file.parent.mkdir(parents=True, exist_ok=True)
with open(output_file, "w") as f:
json.dump(summary, f, indent=2, default=str)
print(f"\nJSON summary written to: {output_file}")
async def main() -> int:
"""Main entry point for the BEAM test runner."""
@ -1050,7 +717,7 @@ async def main() -> int:
"--context-length",
type=str,
default="100K",
choices=["100K", "500K", "1M", "10M"],
choices=["1K", "100K", "500K", "1M", "10M"],
help="Context length subset to test (default: 100K)",
)
@ -1074,12 +741,6 @@ async def main() -> int:
help="Number of Honcho instances in the pool (default: 1)",
)
parser.add_argument(
"--anthropic-api-key",
type=str,
help="Anthropic API key for response judging (optional)",
)
parser.add_argument(
"--timeout",
type=int,
@ -1125,7 +786,6 @@ async def main() -> int:
data_dir=data_dir,
base_api_port=args.base_api_port,
pool_size=args.pool_size,
anthropic_api_key=args.anthropic_api_key,
timeout_seconds=args.timeout,
cleanup_workspace=args.cleanup_workspace,
use_get_context=args.use_get_context,
@ -1136,14 +796,14 @@ async def main() -> int:
if args.conversation_ids:
conversation_ids = args.conversation_ids.split(",")
else:
conversation_ids = runner.list_conversations(args.context_length)
conversation_ids = list_conversations(data_dir, args.context_length)
# Run conversations
results, total_elapsed = await runner.run_conversations(
args.context_length, conversation_ids, args.batch_size
)
runner.print_summary(results, total_elapsed)
print_summary(results, total_elapsed)
# Generate JSON output
if args.json_output:
@ -1153,8 +813,19 @@ async def main() -> int:
f"tests/bench/eval_results/beam_{args.context_length}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
)
runner.generate_json_summary(
results, args.context_length, total_elapsed, output_file
generate_json_summary(
results,
args.context_length,
total_elapsed,
output_file,
metadata_extra={
"base_api_port": runner.base_api_port,
"pool_size": runner.pool_size,
"timeout_seconds": runner.timeout_seconds,
"deriver_settings": settings.DERIVER.model_dump(),
"dialectic_settings": settings.DIALECTIC.model_dump(),
"dream_settings": settings.DREAM.model_dump(),
},
)
# Export metrics

558
tests/bench/beam_baseline.py Normal file
View File

@ -0,0 +1,558 @@
"""
BEAM Baseline Test Runner (Direct Claude Context)
A script that executes BEAM benchmark tests directly against Claude Sonnet 4.5
by feeding the entire conversation history into the context window.
This serves as a baseline comparison against Honcho's memory framework.
## BEAM Overview
BEAM evaluates long-term memory capabilities across ten distinct memory abilities:
1. Abstention - Determines if models avoid answering without evidence
2. Contradiction Resolution - Detects inconsistencies across distant dialogue turns
3. Event Ordering - Assesses sequence recognition of evolving information
4. Information Extraction - Measures factual recall from lengthy histories
5. Instruction Following - Tests sustained adherence to user constraints
6. Knowledge Update - Evaluates fact revision when new information emerges
7. Multi-Session Reasoning - Probes inference integrating evidence across non-adjacent segments
8. Preference Following - Captures personalized, adaptive responses
9. Summarization - Tests content compression and abstraction
10. Temporal Reasoning - Examines explicit and implicit time-relation understanding
## To use
0. Set up env:
```
uv sync
source .venv/bin/activate
```
1. Run this file with the 100K dataset:
```
python -m tests.bench.beam_baseline --context-length 100K
```
Optional arguments:
```
--context-length: Context length subset to test (100K, 500K, 1M, 10M) (default: 100K)
--conversation-ids: Comma-separated list of conversation IDs to test (default: all in context length)
--batch-size: Number of conversations to run concurrently in each batch (default: 1)
--json-output: Path to write JSON summary results for analytics
```
## Other notes
- Uses OpenRouter API (configured via LLM_OPENAI_COMPATIBLE_API_KEY in tests/bench/.env or env var)
- Default model is anthropic/claude-haiku-4-5
- Evaluation follows the paper's nugget-based methodology with 0/0.5/1 scoring
- Event ordering uses Kendall tau-b coefficient
"""
import argparse
import asyncio
import logging
import os
import time
from datetime import datetime
from pathlib import Path
from typing import Any
from dotenv import load_dotenv
from openai import AsyncOpenAI
from src.config import settings
from .beam_common import (
ConversationResult,
QuestionResult,
calculate_ability_scores,
calculate_tokens,
extract_messages_from_chat_data,
format_duration,
generate_json_summary,
judge_event_ordering,
judge_nugget_based,
list_conversations,
load_conversation,
print_summary,
)
# Load .env from bench directory
bench_dir = Path(__file__).parent
load_dotenv(bench_dir / ".env")
# OpenRouter model format for baseline testing
MODEL_BEING_TESTED = "anthropic/claude-haiku-4.5"
class BEAMBaselineRunner:
"""
Executes BEAM benchmark tests directly against Claude Sonnet 4.5.
"""
def __init__(
self,
data_dir: Path,
):
"""
Initialize the BEAM baseline test runner.
Args:
data_dir: Path to the BEAM data directory
"""
self.data_dir: Path = data_dir
# Configure logging
logging.basicConfig(
level=logging.WARNING, format="%(asctime)s - %(levelname)s - %(message)s"
)
self.logger: logging.Logger = logging.getLogger(__name__)
# Initialize OpenRouter client for model being tested and judging
openrouter_api_key = os.getenv("LLM_OPENAI_COMPATIBLE_API_KEY")
openrouter_base_url = os.getenv(
"LLM_OPENAI_COMPATIBLE_BASE_URL", "https://openrouter.ai/api/v1"
)
if not openrouter_api_key:
raise ValueError(
"LLM_OPENAI_COMPATIBLE_API_KEY is not set in tests/bench/.env"
)
self.openrouter_client: AsyncOpenAI = AsyncOpenAI(
api_key=openrouter_api_key,
base_url=openrouter_base_url,
)
# Model to use for judging (OpenRouter format)
self.judge_model: str = os.getenv(
"BEAM_JUDGE_MODEL", "anthropic/claude-sonnet-4.5"
)
# Model to use for answering questions (OpenRouter format)
self.answer_model: str = MODEL_BEING_TESTED
def _format_conversation_context(
self,
messages: list[dict[str, str]],
max_tokens: int = 140000,
) -> str:
"""
Format conversation messages into a context string, truncating from the
beginning if the total exceeds max_tokens.
Args:
messages: List of messages with 'role' and 'content' keys
max_tokens: Maximum tokens allowed for the context (default: 140000
to account for Claude's tokenizer producing ~30% more tokens
than o200k_base, plus room for system prompt and question)
Returns:
Formatted conversation transcript string
"""
# Calculate tokens for each message and find where to start
# Note: We use o200k_base tokenizer but Claude's tokenizer produces ~30% more
# tokens, so we use a conservative max_tokens limit to compensate
message_tokens: list[int] = []
for msg in messages:
role = msg.get("role", "unknown")
content = msg.get("content", "")
role_label = "User" if role == "user" else "Assistant"
formatted = f"{role_label}: {content}\n"
message_tokens.append(calculate_tokens(formatted))
total_tokens = sum(message_tokens)
# Find the starting index to fit within max_tokens
start_idx = 0
if total_tokens > max_tokens:
running_total = total_tokens
for i, tokens in enumerate(message_tokens):
if running_total <= max_tokens:
break
running_total -= tokens
start_idx = i + 1
truncated_count = start_idx
print(
f" [TRUNCATION] Removed {truncated_count} messages from start "
+ f"({total_tokens:,} -> {running_total:,} tokens)"
)
# Build the context string from start_idx onwards
lines: list[str] = []
lines.append("=== CONVERSATION HISTORY ===\n")
if start_idx > 0:
lines.append(
f"[... {start_idx} earlier messages truncated to fit context window ...]\n"
)
for msg in messages[start_idx:]:
role = msg.get("role", "unknown")
content = msg.get("content", "")
role_label = "User" if role == "user" else "Assistant"
lines.append(f"{role_label}: {content}\n")
lines.append("=== END CONVERSATION HISTORY ===")
return "\n".join(lines)
async def _process_single_question(
self,
conversation_context: str,
ability: str,
q_idx: int,
q_data: dict[str, Any],
semaphore: asyncio.Semaphore,
) -> QuestionResult:
"""Process a single BEAM question."""
async with semaphore:
question = q_data["question"]
rubric = q_data.get("rubric", [])
answer = (
q_data.get("answer")
or q_data.get("ideal_response")
or q_data.get("ideal_answer")
)
print(f" [{ability}] Q{q_idx + 1}: {question[:100]}...")
# Build system prompt with cache control for the conversation context
# The context is identical for all questions, so caching saves tokens/latency
system_prompt = f"""You are a helpful assistant with memory of past conversations.
Below is a history of past conversations. Use this history to answer the user's question accurately.
{conversation_context}"""
# Call model via OpenRouter with full context
try:
response = await self.openrouter_client.chat.completions.create(
model=self.answer_model,
max_tokens=settings.DIALECTIC.MAX_OUTPUT_TOKENS,
messages=[
{
"role": "system",
"content": system_prompt,
},
{
"role": "user",
"content": question,
},
],
)
if not response.choices or not response.choices[0].message.content:
actual_response = ""
else:
actual_response = response.choices[0].message.content
except Exception as e:
self.logger.error(f"Error calling Claude API: {e}")
actual_response = f"Error: {e}"
# Judge response based on memory ability
if ability == "event_ordering":
judgment = await judge_event_ordering(
self.openrouter_client,
self.judge_model,
question,
rubric,
actual_response,
)
nugget_scores = None
else:
judgment = await judge_nugget_based(
self.openrouter_client,
self.judge_model,
question,
rubric,
actual_response,
)
nugget_scores = judgment.get("nugget_scores")
score = judgment.get("overall_score", 0.0)
reasoning = judgment.get("overall_reasoning", "")
question_result: QuestionResult = {
"question": question,
"answer": answer,
"actual_response": actual_response,
"memory_ability": ability,
"rubric": rubric,
"nugget_scores": nugget_scores,
"score": score,
"passed": score >= 0.5,
"reasoning": reasoning,
}
status = "PASS" if score >= 0.5 else "FAIL"
print(f" [{ability}] Q{q_idx + 1} Score: {score:.2f} [{status}]")
if score < 0.5 and reasoning:
print(f" Reasoning: {reasoning}")
if rubric:
print(" Rubric:")
for i, rubric_item in enumerate(rubric, 1):
print(f" {i}. {rubric_item}")
if answer:
print(f" Ideal Response: {answer}")
print(f" Our Response: {actual_response}")
return question_result
async def execute_conversation(
self, context_length: str, conversation_id: str
) -> ConversationResult:
"""
Execute BEAM benchmark for a single conversation using direct Claude context.
Args:
context_length: Context length (100K, 500K, 1M, 10M)
conversation_id: Conversation ID
Returns:
Conversation execution results
"""
start_time = time.time()
print(f"\n{'=' * 80}")
print(
f"Executing BEAM conversation {conversation_id} ({context_length} context) [BASELINE]"
)
print(f"{'=' * 80}")
workspace_id = f"baseline_{context_length}_{conversation_id}"
result: ConversationResult = {
"conversation_id": conversation_id,
"context_length": context_length,
"workspace_id": workspace_id,
"total_turns": 0,
"total_messages": 0,
"question_results": [],
"ability_scores": {},
"overall_score": 0.0,
"error": None,
"start_time": start_time,
"end_time": 0.0,
"duration_seconds": 0.0,
}
try:
# Load conversation data
conv_data = load_conversation(
self.data_dir, context_length, conversation_id
)
chat_data = conv_data["chat"]
questions_data = conv_data["questions"]
# Extract all messages
messages = extract_messages_from_chat_data(chat_data)
result["total_messages"] = len(messages)
result["total_turns"] = len(messages)
# Calculate token count
total_tokens = sum(calculate_tokens(m["content"]) for m in messages)
print(
f"[{workspace_id}] Context: {len(messages)} messages, ~{total_tokens:,} tokens"
)
# Format conversation as context
conversation_context = self._format_conversation_context(messages)
print(f"[{workspace_id}] Executing questions...")
# Execute questions for each memory ability
question_tasks: list[Any] = []
semaphore = asyncio.Semaphore(5)
for ability, questions in questions_data.items():
print(
f"\n[{workspace_id}] Queuing {ability} ({len(questions)} questions)"
)
for q_idx, q_data in enumerate(questions):
question_tasks.append(
self._process_single_question(
conversation_context,
ability,
q_idx,
q_data,
semaphore,
)
)
results = await asyncio.gather(*question_tasks)
result["question_results"] = list(results)
# Calculate ability scores
result["ability_scores"] = calculate_ability_scores(
result["question_results"]
)
# Calculate overall score
if result["ability_scores"]:
result["overall_score"] = sum(result["ability_scores"].values()) / len(
result["ability_scores"]
)
result["end_time"] = time.time()
result["duration_seconds"] = result["end_time"] - result["start_time"]
print(
f"\n[{workspace_id}] Completed in {format_duration(result['duration_seconds'])}"
)
print(f"Overall Score: {result['overall_score']:.3f}")
except Exception as e:
self.logger.error(f"Error executing conversation {conversation_id}: {e}")
result["error"] = str(e)
result["end_time"] = time.time()
result["duration_seconds"] = result["end_time"] - result["start_time"]
return result
async def run_conversations(
self,
context_length: str,
conversation_ids: list[str],
batch_size: int = 1,
) -> tuple[list[ConversationResult], float]:
"""
Run multiple conversations from the BEAM benchmark.
Args:
context_length: Context length (100K, 500K, 1M, 10M)
conversation_ids: List of conversation IDs to run
batch_size: Number of conversations to run concurrently in each batch
Returns:
Tuple of (list of conversation results, total duration)
"""
print(
f"Running {len(conversation_ids)} conversations from {context_length} context length [BASELINE]"
)
overall_start = time.time()
all_results: list[ConversationResult] = []
for i in range(0, len(conversation_ids), batch_size):
batch = conversation_ids[i : i + batch_size]
batch_num = (i // batch_size) + 1
total_batches = (len(conversation_ids) + batch_size - 1) // batch_size
print(f"\n{'=' * 80}")
print(
f"Processing batch {batch_num}/{total_batches} ({len(batch)} conversations)"
)
print(f"{'=' * 80}")
# Run conversations in current batch concurrently
batch_results: list[ConversationResult] = await asyncio.gather(
*[
self.execute_conversation(context_length, conv_id)
for conv_id in batch
]
)
all_results.extend(batch_results)
overall_end = time.time()
overall_duration = overall_end - overall_start
return all_results, overall_duration
async def main() -> int:
"""Main entry point for the BEAM baseline test runner."""
parser = argparse.ArgumentParser(
description="Run BEAM benchmark tests directly against Claude (baseline, no Honcho)",
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument(
"--context-length",
type=str,
default="100K",
choices=["1K", "100K", "500K", "1M", "10M"],
help="Context length subset to test (default: 100K)",
)
parser.add_argument(
"--conversation-ids",
type=str,
help="Comma-separated list of conversation IDs to test (default: all)",
)
parser.add_argument(
"--batch-size",
type=int,
default=1,
help="Number of conversations to run concurrently in each batch (default: 1)",
)
parser.add_argument(
"--json-output",
type=Path,
help="Path to write JSON summary results for analytics (optional)",
)
args = parser.parse_args()
# Setup data directory
data_dir = Path(__file__).parent / "beam_data"
if not data_dir.exists():
print(f"Error: BEAM data directory not found at {data_dir}")
return 1
# Create runner
runner = BEAMBaselineRunner(data_dir=data_dir)
try:
# Determine which conversations to run
if args.conversation_ids:
conversation_ids = args.conversation_ids.split(",")
else:
conversation_ids = list_conversations(data_dir, args.context_length)
# Run conversations
results, total_elapsed = await runner.run_conversations(
args.context_length, conversation_ids, args.batch_size
)
print_summary(results, total_elapsed)
# Generate JSON output
if args.json_output:
output_file = args.json_output
else:
output_file = Path(
f"tests/bench/eval_results/beam_baseline_{args.context_length}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
)
generate_json_summary(
results,
args.context_length,
total_elapsed,
output_file,
metadata_extra={
"runner_type": "baseline_direct_context",
"model": MODEL_BEING_TESTED,
},
)
return 0
except KeyboardInterrupt:
print("\nTest execution interrupted by user")
return 1
except Exception as e:
print(f"Error running tests: {e}")
import traceback
traceback.print_exc()
return 1
if __name__ == "__main__":
exit_code = asyncio.run(main())
exit(exit_code)

643
tests/bench/beam_common.py Normal file
View File

@ -0,0 +1,643 @@
"""
Common utilities for BEAM benchmark test runners.
Shared functionality between the Honcho benchmark and baseline benchmark.
"""
import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Any, cast
import tiktoken
from openai import AsyncOpenAI
from scipy.stats import kendalltau # pyright: ignore[reportUnknownVariableType]
from typing_extensions import TypedDict
logger = logging.getLogger(__name__)
class QuestionResult(TypedDict):
"""Type definition for question evaluation results."""
question: str
answer: str | None
actual_response: str
memory_ability: str
rubric: list[str]
nugget_scores: list[dict[str, Any]] | None
score: float
passed: bool
reasoning: str
class ConversationResult(TypedDict):
"""Type definition for conversation execution results."""
conversation_id: str
context_length: str
workspace_id: str
total_turns: int
total_messages: int
question_results: list[QuestionResult]
ability_scores: dict[str, float]
overall_score: float
error: str | None
start_time: float
end_time: float
duration_seconds: float
def format_duration(total_seconds: float) -> str:
"""Format a duration in seconds into a human-readable string."""
minutes = int(total_seconds // 60)
if minutes > 0:
seconds_rounded = int(round(total_seconds - minutes * 60))
if seconds_rounded == 60:
minutes += 1
seconds_rounded = 0
return f"{minutes}m{seconds_rounded:02d}s"
return f"{total_seconds:.2f}s"
def calculate_tokens(text: str) -> int:
"""Calculate tokens for a given text."""
tokenizer = tiktoken.get_encoding("o200k_base")
try:
return len(
tokenizer.encode(
text,
disallowed_special=(tokenizer.special_tokens_set - {"<|endoftext|>"}),
)
)
except Exception:
return len(text) // 4
def load_conversation(
data_dir: Path, context_length: str, conversation_id: str
) -> dict[str, Any]:
"""
Load a BEAM conversation from the data directory.
Args:
data_dir: Path to the BEAM data directory
context_length: Context length (100K, 500K, 1M, 10M)
conversation_id: Conversation ID
Returns:
Dictionary containing conversation data and probing questions
"""
conv_dir = data_dir / context_length / conversation_id
# Load chat data
chat_file = conv_dir / "chat.json"
with open(chat_file) as f:
chat_data = json.load(f)
# Load probing questions
questions_file = conv_dir / "probing_questions" / "probing_questions.json"
with open(questions_file) as f:
questions_data = json.load(f)
return {"chat": chat_data, "questions": questions_data}
def list_conversations(data_dir: Path, context_length: str) -> list[str]:
"""
List all conversation IDs for a given context length.
Args:
data_dir: Path to the BEAM data directory
context_length: Context length (100K, 500K, 1M, 10M)
Returns:
List of conversation ID strings
"""
context_dir = data_dir / context_length
return [
d.name for d in sorted(context_dir.iterdir()) if d.is_dir() and d.name.isdigit()
]
def extract_messages_from_chat_data(chat_data: list[Any]) -> list[dict[str, str]]:
"""
Extract all messages from BEAM chat data structure.
Handles both standard structure (100K, 500K, 1M) and 10M plan-based structure.
Args:
chat_data: Raw chat data from BEAM JSON
Returns:
List of messages with 'role' and 'content' keys
"""
messages: list[dict[str, str]] = []
for batch in chat_data:
# Check if this is a 10M conversation with plan-based structure
if any(key.startswith("plan-") for key in batch):
# 10M structure: { "plan-1": [...], "plan-2": [...], ... }
for plan_name, plan_batches in batch.items():
if not plan_name.startswith("plan-"):
continue
for plan_batch in plan_batches:
for turn_group in plan_batch.get("turns", []):
for turn in turn_group:
messages.append(
{
"role": turn["role"],
"content": turn["content"],
}
)
else:
# Standard structure for 100K, 500K, 1M
for turn_group in batch.get("turns", []):
for turn in turn_group:
messages.append(
{
"role": turn["role"],
"content": turn["content"],
}
)
return messages
async def judge_nugget_based(
openrouter_client: AsyncOpenAI,
judge_model: str,
question: str,
rubric: list[str],
actual_response: str,
) -> dict[str, Any]:
"""
Use an LLM to judge a response using nugget-based evaluation.
Args:
openrouter_client: OpenAI-compatible client for API calls
judge_model: Model ID to use for judging
question: The question asked
rubric: List of nuggets (atomic criteria) to check
actual_response: Actual response to evaluate
Returns:
Judgment result with nugget scores and overall score
"""
try:
# Build the nugget evaluation prompt
nuggets_formatted = "\n".join(
[f"{i + 1}. {nugget}" for i, nugget in enumerate(rubric)]
)
system_prompt = """You are an expert evaluator tasked with judging whether the LLM's response demonstrates compliance with the specified RUBRIC CRITERIA.
## EVALUATION RUBRIC:
The rubric defines specific requirements, constraints, or expected behaviors that the LLM response should demonstrate.
**IMPORTANT**: Pay careful attention to whether each rubric criterion specifies:
- **Positive requirements** (things the response SHOULD include/do)
- **Negative constraints** (things the response SHOULD NOT include/do, often indicated by "no", "not", "avoid", "absent")
## RESPONSIVENESS REQUIREMENT (anchored to the QUESTION)
A compliant response must be **on-topic with respect to the QUESTION** and attempt to answer it.
- If the response does not address the QUESTION, score **0.0** for all criteria and stop.
- For negative constraints, both must hold: (a) the response is responsive to the QUESTION, and (b) the prohibited element is absent.
## SEMANTIC TOLERANCE RULES:
Judge by meaning, not exact wording.
- Accept **paraphrases** and **synonyms** that preserve intent.
- **Case/punctuation/whitespace** differences must be ignored.
- **Numbers/currencies/dates** may appear in equivalent forms (e.g., "$68,000", "68k", "68,000 USD", or "sixty-eight thousand dollars"). Treat them as equal when numerically equivalent.
- If the rubric expects a number or duration, prefer **normalized comparison** (extract and compare values) over string matching.
## STYLE NEUTRALITY (prevents style contamination):
Ignore tone, politeness, length, and flourish unless the rubric explicitly requires a format/structure (e.g., "itemized list", "no citations", "one sentence").
- Do **not** penalize hedging, voice, or verbosity if content satisfies the rubric.
- Only evaluate format when the rubric **explicitly** mandates it.
## SCORING SCALE:
- **1.0 (Complete Compliance)**: Fully complies with the rubric criterion.
- Positive: required element present, accurate, properly executed (allowing semantic equivalents).
- Negative: prohibited element **absent** AND response is **responsive**.
- **0.5 (Partial Compliance)**: Partially complies.
- Positive: element present but minor inaccuracies/incomplete execution.
- Negative: generally responsive and mostly avoids the prohibited element but with minor/edge violations.
- **0.0 (No Compliance)**: Fails to comply.
- Positive: required element missing or incorrect.
- Negative: prohibited element present **or** response is non-responsive/evasive even if the element is absent.
## EVALUATION INSTRUCTIONS:
1. **Understand the Requirement**: For each rubric criterion, determine if it is asking for something to be present (positive) or absent (negative/constraint).
2. **Parse Compound Statements**: If a rubric criterion contains multiple elements connected by "and" or commas, evaluate whether:
- **All elements** must be present for full compliance (1.0)
- **Some elements** present indicates partial compliance (0.5)
- **No elements** present indicates no compliance (0.0)
3. **Check Compliance**: For each criterion:
- For positive requirements: Look for the presence and quality of the required element
- For negative constraints: Look for the absence of the prohibited element
4. **Assign Score**: Based on compliance with each specific rubric criterion according to the scoring scale above.
5. **Provide Reasoning**: For each criterion, explain whether it was satisfied and justify the score.
Use the `evaluate_response` tool to submit your evaluation with scores and reasoning for each rubric criterion."""
user_prompt = f"""## EVALUATION INPUTS
- QUESTION (what the user asked): {question}
- RUBRIC CRITERIA (what to check):
{nuggets_formatted}
- RESPONSE TO EVALUATE: {actual_response}
Evaluate the response against each rubric criterion. Provide a score and reasoning for each criterion, and calculate the overall score as the average of all criterion scores."""
tool_definition = {
"type": "function",
"function": {
"name": "evaluate_response",
"description": "Submit the evaluation results for the response based on the rubric.",
"parameters": {
"type": "object",
"properties": {
"nugget_scores": {
"type": "array",
"items": {
"type": "object",
"properties": {
"nugget_index": {"type": "integer"},
"score": {"type": "number"},
"reasoning": {"type": "string"},
},
"required": ["nugget_index", "score", "reasoning"],
},
},
"overall_score": {"type": "number"},
"overall_reasoning": {"type": "string"},
},
"required": [
"nugget_scores",
"overall_score",
"overall_reasoning",
],
},
},
}
messages = cast(
list[dict[str, Any]],
[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
)
response = await openrouter_client.chat.completions.create(
model=judge_model,
max_tokens=2000,
temperature=0.0,
messages=cast(Any, messages), # type: ignore[arg-type]
tools=cast(Any, [tool_definition]), # type: ignore[arg-type]
tool_choice={
"type": "function",
"function": {"name": "evaluate_response"},
},
)
if not response.choices or not response.choices[0].message:
raise ValueError("OpenRouter returned empty response")
message = response.choices[0].message
if not message.tool_calls or len(message.tool_calls) == 0:
raise ValueError("No tool calls found in response")
tool_call = message.tool_calls[0]
# Access function tool call attributes (OpenAI format)
if tool_call.function.name != "evaluate_response": # pyright: ignore
raise ValueError(f"Unexpected tool call: {tool_call.function.name}") # pyright: ignore
# Parse the JSON arguments
judgment = json.loads(tool_call.function.arguments) # pyright: ignore
if not isinstance(judgment, dict):
raise ValueError(f"Tool arguments is not a dictionary: {type(judgment)}")
return cast(dict[str, Any], judgment)
except Exception as e:
logger.error(f"Error judging response: {e}")
# Fallback to simple 0 score
return {
"nugget_scores": [
{"nugget_index": i + 1, "score": 0.0, "reasoning": f"Error: {e}"}
for i in range(len(rubric))
],
"overall_score": 0.0,
"overall_reasoning": f"Evaluation failed due to error: {e}",
}
def align_events(expected_events: list[str], extracted_events: list[str]) -> list[int]:
"""
Align extracted events with expected events using string matching.
Returns a list of indices mapping extracted events to expected events.
"""
alignment: list[int] = []
for extracted in extracted_events:
best_match_idx: int = -1
for i, expected in enumerate(expected_events):
# Use simple string matching
if (
expected.lower() in extracted.lower()
or extracted.lower() in expected.lower()
):
best_match_idx = i
break
if best_match_idx >= 0:
alignment.append(best_match_idx)
return alignment
async def judge_event_ordering(
openrouter_client: AsyncOpenAI,
judge_model: str,
question: str,
rubric: list[str],
actual_response: str,
) -> dict[str, Any]:
"""
Judge event ordering questions using Kendall tau-b coefficient.
Args:
openrouter_client: OpenAI-compatible client for API calls
judge_model: Model ID to use for judging
question: The question asked
rubric: List of expected events in correct order
actual_response: Actual response from the system
Returns:
Judgment with Kendall tau-b score
"""
try:
# First, extract the events mentioned in the response
system_prompt = """You are an expert at extracting ordered lists of events or items from text.
Your task is to extract the ordered list of events/items mentioned in a response.
Use the `extract_ordered_events` tool to submit the extracted list."""
user_prompt = f"""Question: "{question}"
Response: "{actual_response}"
Extract the ordered list of events or items mentioned in the response. Preserve the order as stated in the response."""
tool_definition = {
"type": "function",
"function": {
"name": "extract_ordered_events",
"description": "Submit the ordered list of events extracted from the response.",
"parameters": {
"type": "object",
"properties": {
"extracted_events": {
"type": "array",
"items": {"type": "string"},
},
},
"required": ["extracted_events"],
},
},
}
messages = cast(
list[dict[str, Any]],
[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
)
response = await openrouter_client.chat.completions.create(
model=judge_model,
max_tokens=1000,
temperature=0.0,
messages=cast(Any, messages), # type: ignore[arg-type]
tools=cast(Any, [tool_definition]), # type: ignore[arg-type]
tool_choice={
"type": "function",
"function": {"name": "extract_ordered_events"},
},
)
if not response.choices or not response.choices[0].message:
raise ValueError("OpenRouter returned empty response")
message = response.choices[0].message
if not message.tool_calls or len(message.tool_calls) == 0:
raise ValueError("No tool calls found in response")
tool_call = message.tool_calls[0]
# Access function tool call attributes (OpenAI format)
if tool_call.function.name != "extract_ordered_events": # pyright: ignore
raise ValueError(f"Unexpected tool call: {tool_call.function.name}") # pyright: ignore
# Parse the JSON arguments
extracted_dict = json.loads(tool_call.function.arguments) # pyright: ignore
if not isinstance(extracted_dict, dict):
raise ValueError(
f"Tool arguments is not a dictionary: {type(extracted_dict)}"
)
extracted_dict = cast(dict[str, Any], extracted_dict)
raw_events: list[str] = extracted_dict.get("extracted_events", [])
extracted_events = [str(e) for e in raw_events]
# Now compute alignment and Kendall tau-b
alignment = align_events(rubric, extracted_events)
# Compute Kendall tau-b
tau: float
if kendalltau is None:
logger.warning(
"scipy not installed, cannot compute Kendall tau-b. Install with: uv pip install scipy"
)
tau = 0.0
elif len(alignment) < 2:
tau = 0.0
else:
# Create rank lists
expected_ranks = list(range(len(alignment)))
actual_ranks = [alignment[i] for i in range(len(alignment))]
result_tuple: Any = kendalltau(expected_ranks, actual_ranks)
# kendalltau returns a tuple, first element is the tau coefficient
tau_value: Any = result_tuple[0]
# Handle the return type properly - convert to float
try:
tau = float(tau_value)
if tau != tau: # Check for NaN
tau = 0.0
except (TypeError, ValueError):
tau = 0.0
return {
"kendall_tau_b": tau,
"extracted_events": extracted_events,
"alignment": alignment,
"overall_score": (tau + 1) / 2, # Normalize to [0, 1]
"overall_reasoning": f"Kendall tau-b coefficient: {tau:.3f}. Extracted {len(extracted_events)} events from response.",
}
except Exception as e:
logger.error(f"Error in event ordering evaluation: {e}")
return {
"kendall_tau_b": 0.0,
"extracted_events": [],
"alignment": [],
"overall_score": 0.0,
"overall_reasoning": f"Evaluation failed due to error: {e}",
}
def calculate_ability_scores(
question_results: list[QuestionResult],
) -> dict[str, float]:
"""Calculate average scores by memory ability."""
ability_totals: dict[str, list[float]] = {}
for qr in question_results:
ability = qr["memory_ability"]
if ability not in ability_totals:
ability_totals[ability] = []
ability_totals[ability].append(qr["score"])
return {
ability: sum(scores) / len(scores) for ability, scores in ability_totals.items()
}
def print_summary(
results: list[ConversationResult], total_elapsed_seconds: float
) -> None:
"""Print a summary of all test results."""
print(f"\n{'=' * 80}")
print("BEAM BENCHMARK EXECUTION SUMMARY")
print(f"{'=' * 80}")
total_conversations = len(results)
total_questions = sum(len(r["question_results"]) for r in results)
print(f"Total Conversations: {total_conversations}")
print(f"Total Questions: {total_questions}")
print(f"Total Test Time: {format_duration(total_elapsed_seconds)}")
# Calculate average scores by ability
ability_scores: dict[str, list[float]] = {}
for result in results:
for ability, score in result["ability_scores"].items():
if ability not in ability_scores:
ability_scores[ability] = []
ability_scores[ability].append(score)
print("\nAverage Scores by Memory Ability:")
for ability, scores in sorted(ability_scores.items()):
avg_score = sum(scores) / len(scores)
print(f" {ability:30s}: {avg_score:.3f}")
# Overall average
overall_scores = [r["overall_score"] for r in results]
overall_avg = sum(overall_scores) / len(overall_scores) if overall_scores else 0.0
print(f"\n{'Overall Average Score':30s}: {overall_avg:.3f}")
print(f"{'=' * 80}")
def generate_json_summary(
results: list[ConversationResult],
context_length: str,
total_elapsed_seconds: float,
output_file: Path,
metadata_extra: dict[str, Any] | None = None,
) -> None:
"""Generate a comprehensive JSON summary of test results."""
# Calculate summary statistics
total_conversations = len(results)
total_questions = sum(len(r["question_results"]) for r in results)
# Calculate average scores by ability
ability_scores: dict[str, list[float]] = {}
for result in results:
for ability, score in result["ability_scores"].items():
if ability not in ability_scores:
ability_scores[ability] = []
ability_scores[ability].append(score)
ability_averages = {
ability: sum(scores) / len(scores) for ability, scores in ability_scores.items()
}
# Overall average
overall_scores = [r["overall_score"] for r in results]
overall_avg = sum(overall_scores) / len(overall_scores) if overall_scores else 0.0
metadata = {
"context_length": context_length,
"execution_timestamp": datetime.now().isoformat(),
"runner_version": "1.0.0",
}
if metadata_extra:
metadata.update(metadata_extra)
summary = {
"metadata": metadata,
"summary_statistics": {
"total_conversations": total_conversations,
"total_questions": total_questions,
"overall_average_score": overall_avg,
"ability_averages": ability_averages,
},
"timing": {
"total_duration_seconds": total_elapsed_seconds,
},
"detailed_results": results,
}
output_file.parent.mkdir(parents=True, exist_ok=True)
with open(output_file, "w") as f:
json.dump(summary, f, indent=2, default=str)
print(f"\nJSON summary written to: {output_file}")

1447
tests/bench/explicit.py Normal file
View File

File diff suppressed because it is too large Load Diff

40
tests/bench/harness.py
View File

@ -1104,24 +1104,28 @@ Examples:
sys.exit(1)
# Create and run the harness or pool
if args.pool_size > 1:
pool = HonchoHarnessPool(
pool_size=args.pool_size,
base_db_port=args.port,
base_api_port=args.api_port,
base_redis_port=args.redis_port,
project_root=args.project_root,
)
asyncio.run(pool.run())
else:
harness = HonchoHarness(
db_port=args.port,
api_port=args.api_port,
redis_port=args.redis_port,
project_root=args.project_root,
instance_id=0,
)
asyncio.run(harness.run())
try:
if args.pool_size > 1:
pool = HonchoHarnessPool(
pool_size=args.pool_size,
base_db_port=args.port,
base_api_port=args.api_port,
base_redis_port=args.redis_port,
project_root=args.project_root,
)
asyncio.run(pool.run())
else:
harness = HonchoHarness(
db_port=args.port,
api_port=args.api_port,
redis_port=args.redis_port,
project_root=args.project_root,
instance_id=0,
)
asyncio.run(harness.run())
except KeyboardInterrupt:
# Cleanup already handled in run() finally block
pass
if __name__ == "__main__":

62
tests/bench/incorrect_beam_qs.txt Normal file
View File

@ -0,0 +1,62 @@
{
"question": "How long had I been with the person I mentioned meeting at that festival before we started dating?",
"answer": "You said you had been with Stephen for 5 years, and you met him at the Montserrat Film Festival in 2018.",
"difficulty": "medium",
"question_type": "duration_recall",
"conversation_reference": "Turn 1: Early mention of relationship details",
"key_facts_tested": [
"duration: 5 years",
"meeting event: Montserrat Film Festival",
"relationship start timing"
],
"source_chat_ids": [
1
],
"rubric": [
"LLM response should state: 5 years"
]
},
{
"question": "By what date am I aiming to complete all my onboarding modules?",
"answer": "By April 22",
"difficulty": "easy",
"update_type": "goal_date_adjustment",
"tests_retention_of": "updated target completion date for onboarding modules",
"conversation_references": [
"I set a goal to complete all onboarding modules by April 25, aiming for 95% quiz scores.",
"I adjusted my onboarding goal to complete all modules by April 22 to better align with the startup\u2019s revised schedule."
],
"potential_confusion": "LLM might incorrectly recall the original April 25 deadline instead of the updated April 22 date",
"source_chat_ids": {
"original_info": [
108
],
"updated_info": []
},
"rubric": [
"LLM response should state: April 22"
]
},
{
"question": "How many different shoe sizes have I mentioned across my messages?",
"answer": "Two sizes: 11 and 11.5",
"difficulty": "easy",
"reasoning_type": "simple_cross_session_facts",
"sessions_required": 2,
"conversation_references": [
"Choosing between Adidas Ultraboost and Nike React Infinity Run",
"Considering returning Adidas Ultraboost size 11 and reordering size 11.5"
],
"reasoning_steps": [
"Identify the shoe size mentioned in the first user message about deciding between shoes (size not specified here).",
"Find the shoe sizes mentioned in the later user message about returning and reordering Adidas Ultraboost (sizes 11 and 11.5).",
"Count the distinct sizes mentioned across these messages."
],
"source_chat_ids": [
32,
116
],
"rubric": [
"LLM response should state: Two sizes: 11 and 11.5"
]
},

909
tests/bench/locomo.py Normal file
View File

@ -0,0 +1,909 @@
"""
Honcho LoCoMo Benchmark Test Runner
A script that executes LoCoMo benchmark tests against a running Honcho instance.
This script:
1. Loads LoCoMo conversation data from JSON files
2. Creates a workspace for each conversation sample
3. Ingests conversation sessions as messages between two peers
4. Waits for the deriver queue to process everything
5. Triggers a dream for memory consolidation
6. Executes questions and judges responses using an LLM
## LoCoMo Overview
LoCoMo evaluates very long-term conversational memory across five question categories:
1. Single-hop - Direct factual recall from conversations
2. Multi-hop - Reasoning across multiple pieces of information
3. Temporal - Understanding time-based relationships and sequences
4. Commonsense/World knowledge - Applying broader contextual understanding
5. Adversarial - Questions that cannot be answered (filtered out by default)
Reference: https://github.com/snap-research/locomo
Paper: https://arxiv.org/abs/2402.17753
## To use
0. Set up env:
```
uv sync
source .venv/bin/activate
```
NOTE: you may create a .env file in this directory to customize honcho config.
1. Run the test harness:
```
python -m tests.bench.harness
```
2. Run this file with the LoCoMo dataset:
```
python -m tests.bench.locomo --data-file tests/bench/locomo_data/locomo10.json
```
Optional arguments:
```
--anthropic-api-key: Anthropic API key for response judging (can be set in .env as LLM_ANTHROPIC_API_KEY)
--timeout: Timeout for deriver queue to empty in seconds (default: 10 minutes)
--base-api-port: Base port for Honcho API instances (default: 8000)
--pool-size: Number of Honcho instances in the pool (default: 1)
--batch-size: Number of conversations to run concurrently in each batch (default: 1)
--json-output: Path to write JSON summary results for analytics
--cleanup-workspace: Delete workspace after executing each conversation (default: False)
--use-get-context: Use get_context + judge LLM instead of dialectic .chat endpoint (default: False)
--sample-id: Run only the conversation with this sample_id (skips all others)
--test-count: Number of conversations to run (default: all)
--question-count: Number of questions per conversation to run (default: all)
```
"""
import argparse
import asyncio
import logging
import os
import time
from datetime import datetime
from pathlib import Path
from typing import Any, cast
import httpx
from anthropic import AsyncAnthropic
from anthropic.types import MessageParam
from dotenv import load_dotenv
from honcho import AsyncHoncho
from honcho.async_client.session import SessionPeerConfig
from honcho_core.types.workspaces.sessions.message_create_param import (
MessageCreateParam,
)
from openai import AsyncOpenAI
from src.config import settings
from src.utils.metrics_collector import MetricsCollector
from .locomo_common import (
CATEGORY_NAMES,
ConversationResult,
QuestionResult,
calculate_category_scores,
calculate_tokens,
extract_sessions,
filter_questions,
format_duration,
generate_json_summary,
get_evidence_context,
judge_response,
load_locomo_data,
parse_locomo_date,
print_summary,
)
# Load .env from bench directory
bench_dir = Path(__file__).parent
load_dotenv(bench_dir / ".env")
def format_message_with_image(msg: dict[str, Any]) -> tuple[str, dict[str, Any] | None]:
"""
Format a LoCoMo message with optional image caption appended.
Args:
msg: LoCoMo message dict with 'text', optional 'img_url', 'blip_caption', 'query'
Returns:
Tuple of (formatted_content, metadata_dict or None)
"""
text = msg.get("text", "")
blip_caption = msg.get("blip_caption")
img_urls = msg.get("img_url", [])
query = msg.get("query")
# Append caption to content so deriver can see it
content = f"{text}\n\n[Image shared: {blip_caption}]" if blip_caption else text
# Build metadata if image data exists
metadata: dict[str, Any] | None = None
if img_urls or blip_caption or query:
metadata = {}
if img_urls:
metadata["img_urls"] = img_urls
if blip_caption:
metadata["blip_caption"] = blip_caption
if query:
metadata["image_query"] = query
return content, metadata
def determine_question_target(question: str, speaker_a: str, speaker_b: str) -> str:
"""
Determine which speaker a question is asking about based on the question text.
Args:
question: The question text
speaker_a: Name of speaker A (e.g., "Caroline")
speaker_b: Name of speaker B (e.g., "Melanie")
Returns:
The name of the speaker the question is about (speaker_a or speaker_b)
"""
question_lower = question.lower()
speaker_a_lower = speaker_a.lower()
speaker_b_lower = speaker_b.lower()
# Check for possessive forms too (e.g., "Melanie's kids")
a_in_question = (
speaker_a_lower in question_lower or f"{speaker_a_lower}'s" in question_lower
)
b_in_question = (
speaker_b_lower in question_lower or f"{speaker_b_lower}'s" in question_lower
)
if a_in_question and not b_in_question:
return speaker_a
elif b_in_question and not a_in_question:
return speaker_b
else:
# Question mentions both or neither - default to speaker_a
return speaker_a
class LoCoMoRunner:
"""
Executes LoCoMo benchmark tests against a Honcho instance.
"""
def __init__(
self,
base_api_port: int = 8000,
pool_size: int = 1,
anthropic_api_key: str | None = None,
timeout_seconds: int | None = None,
cleanup_workspace: bool = False,
use_get_context: bool = False,
):
"""
Initialize the LoCoMo test runner.
Args:
base_api_port: Base port for Honcho API instances (default: 8000)
pool_size: Number of Honcho instances in the pool (default: 1)
anthropic_api_key: Anthropic API key for judging responses
timeout_seconds: Timeout for deriver queue in seconds
cleanup_workspace: If True, delete workspace after executing conversation
use_get_context: If True, use get_context + judge LLM instead of dialectic .chat endpoint
"""
self.base_api_port: int = base_api_port
self.pool_size: int = pool_size
self.anthropic_api_key: str | None = anthropic_api_key
self.timeout_seconds: int = (
timeout_seconds if timeout_seconds is not None else 600
)
self.cleanup_workspace: bool = cleanup_workspace
self.use_get_context: bool = use_get_context
# Initialize metrics collector
self.metrics_collector: MetricsCollector = MetricsCollector()
self.metrics_collector.start_collection(
f"locomo_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
)
# Configure logging
logging.basicConfig(
level=logging.WARNING, format="%(asctime)s - %(levelname)s - %(message)s"
)
self.logger: logging.Logger = logging.getLogger(__name__)
# Suppress HTTP request logs from the Honcho SDK
logging.getLogger("httpx").setLevel(logging.ERROR)
logging.getLogger("httpcore").setLevel(logging.ERROR)
if self.anthropic_api_key:
self.anthropic_client: AsyncAnthropic = AsyncAnthropic(
api_key=self.anthropic_api_key
)
else:
api_key = os.getenv("LLM_ANTHROPIC_API_KEY")
if not api_key:
raise ValueError("LLM_ANTHROPIC_API_KEY is not set")
self.anthropic_client = AsyncAnthropic(api_key=api_key)
# Initialize OpenAI client for judging responses
openai_api_key = os.getenv("OPENAI_API_KEY")
if not openai_api_key:
raise ValueError("OPENAI_API_KEY is not set")
self.openai_client: AsyncOpenAI = AsyncOpenAI(api_key=openai_api_key)
def get_honcho_url_for_index(self, index: int) -> str:
"""Get the Honcho URL for a given index using round-robin distribution."""
instance_id = index % self.pool_size
port = self.base_api_port + instance_id
return f"http://localhost:{port}"
async def create_honcho_client(
self, workspace_id: str, honcho_url: str
) -> AsyncHoncho:
"""Create a Honcho client for a specific workspace."""
return AsyncHoncho(
environment="local",
workspace_id=workspace_id,
base_url=honcho_url,
)
async def wait_for_deriver_queue_empty(
self, honcho_client: AsyncHoncho, session_id: str | None = None
) -> bool:
"""Wait for the deriver queue to be empty."""
start_time = time.time()
while True:
try:
status = await honcho_client.get_deriver_status(session=session_id)
except Exception:
await asyncio.sleep(1)
elapsed_time = time.time() - start_time
if elapsed_time >= self.timeout_seconds:
return False
continue
if status.pending_work_units == 0 and status.in_progress_work_units == 0:
return True
elapsed_time = time.time() - start_time
if elapsed_time >= self.timeout_seconds:
return False
await asyncio.sleep(1)
async def trigger_dream_and_wait(
self,
honcho_client: AsyncHoncho,
workspace_id: str,
observer: str,
observed: str | None = None,
session_id: str | None = None,
) -> bool:
"""
Trigger a dream task and wait for it to complete.
Args:
honcho_client: Honcho client instance
workspace_id: Workspace identifier
observer: Observer peer name
observed: Observed peer name (defaults to observer)
session_id: Session ID to scope the dream to
Returns:
True if dream completed successfully, False on timeout
"""
observed = observed or observer
honcho_url = self.get_honcho_url_for_index(0)
url = f"{honcho_url}/v2/workspaces/{workspace_id}/trigger_dream"
payload = {
"observer": observer,
"observed": observed,
"dream_type": "omni",
"session_id": session_id or f"{workspace_id}_session",
}
print(f"[{workspace_id}] Triggering dream at {url}")
try:
async with httpx.AsyncClient() as client:
response = await client.post(
url,
json=payload,
timeout=30.0,
)
if response.status_code != 204:
print(
f"[{workspace_id}] ERROR: Dream trigger failed with status {response.status_code}"
)
print(f"[{workspace_id}] Response body: {response.text}")
return False
except Exception as e:
print(f"[{workspace_id}] ERROR: Dream trigger exception: {e}")
return False
print(
f"[{workspace_id}] Dream triggered successfully for {observer}/{observed}"
)
# Wait for dream queue to empty
print(f"[{workspace_id}] Waiting for dream to complete...")
await asyncio.sleep(2)
success = await self.wait_for_deriver_queue_empty(honcho_client)
if success:
print(f"[{workspace_id}] Dream queue empty")
else:
print(f"[{workspace_id}] Dream queue timeout")
return success
async def execute_conversation(
self,
conversation_data: dict[str, Any],
honcho_url: str,
question_count: int | None = None,
) -> ConversationResult:
"""
Execute LoCoMo benchmark for a single conversation.
Args:
conversation_data: Dictionary containing conversation and QA data
honcho_url: URL of the Honcho instance to use
question_count: Optional limit on number of questions to run
Returns:
Conversation execution results
"""
start_time = time.time()
sample_id = conversation_data.get("sample_id", "unknown")
conversation = conversation_data.get("conversation", {})
qa_list = conversation_data.get("qa", [])
speaker_a = conversation.get("speaker_a", "User")
speaker_b = conversation.get("speaker_b", "Assistant")
print(f"\n{'=' * 80}")
print(f"Executing LoCoMo conversation {sample_id}")
print(f"Speakers: {speaker_a} and {speaker_b}")
print(f"{'=' * 80}")
# Create workspace for this conversation
workspace_id = f"locomo_{sample_id}"
honcho_client = await self.create_honcho_client(workspace_id, honcho_url)
result: ConversationResult = {
"sample_id": sample_id,
"speaker_a": speaker_a,
"speaker_b": speaker_b,
"total_sessions": 0,
"total_turns": 0,
"total_tokens": 0,
"question_results": [],
"category_scores": {},
"overall_score": 0.0,
"error": None,
"start_time": start_time,
"end_time": 0.0,
"duration_seconds": 0.0,
}
try:
# Create peers using their actual names as IDs
peer_a = await honcho_client.peer(id=speaker_a)
peer_b = await honcho_client.peer(id=speaker_b)
# Create session for this conversation
session_id = f"{workspace_id}_session"
session = await honcho_client.session(id=session_id)
# Configure peer observation - observe BOTH peers since questions ask about both speakers
await session.add_peers(
[
(
peer_a,
SessionPeerConfig(observe_me=True, observe_others=False),
),
(
peer_b,
SessionPeerConfig(observe_me=True, observe_others=False),
),
]
)
# Extract and ingest all sessions
sessions = extract_sessions(conversation)
result["total_sessions"] = len(sessions)
print(f"[{workspace_id}] Ingesting {len(sessions)} sessions...")
messages: list[MessageCreateParam] = []
total_tokens = 0
for date_str, session_messages in sessions:
session_date = parse_locomo_date(date_str) if date_str else None
for msg in session_messages:
speaker = msg.get("speaker", "")
content, metadata = format_message_with_image(msg)
result["total_turns"] += 1
total_tokens += calculate_tokens(content)
# Map speaker to peer by name
if speaker == speaker_a:
messages.append(
peer_a.message(
content, metadata=metadata, created_at=session_date
)
)
elif speaker == speaker_b:
messages.append(
peer_b.message(
content, metadata=metadata, created_at=session_date
)
)
result["total_tokens"] = total_tokens
# Add messages in batches of 100
for i in range(0, len(messages), 100):
batch = messages[i : i + 100]
await session.add_messages(batch)
print(
f"[{workspace_id}] Ingested {len(messages)} messages (~{total_tokens:,} tokens). Waiting for deriver queue..."
)
# Wait for deriver queue to empty
await asyncio.sleep(1)
queue_empty = await self.wait_for_deriver_queue_empty(honcho_client)
if not queue_empty:
result["error"] = "Deriver queue timeout"
result["end_time"] = time.time()
result["duration_seconds"] = result["end_time"] - result["start_time"]
print(
f"\n[{workspace_id}] ERROR: Deriver queue timeout after {self.timeout_seconds}s"
)
return result
print(
f"[{workspace_id}] Deriver queue empty. Triggering dream consolidation for both peers..."
)
# Trigger dream for memory consolidation for BOTH peers
# Dream for speaker_a
dream_success_a = await self.trigger_dream_and_wait(
honcho_client,
workspace_id,
observer=speaker_a,
session_id=session_id,
)
if not dream_success_a:
print(
f"[{workspace_id}] Warning: Dream for {speaker_a} did not complete, proceeding anyway"
)
else:
print(f"[{workspace_id}] Dream for {speaker_a} completed.")
# Dream for speaker_b
dream_success_b = await self.trigger_dream_and_wait(
honcho_client,
workspace_id,
observer=speaker_b,
session_id=session_id,
)
if not dream_success_b:
print(
f"[{workspace_id}] Warning: Dream for {speaker_b} did not complete, proceeding anyway"
)
else:
print(f"[{workspace_id}] Dream for {speaker_b} completed.")
# Filter questions
filtered_qa = filter_questions(
qa_list,
exclude_adversarial=True,
test_count=question_count,
)
print(f"[{workspace_id}] Executing {len(filtered_qa)} questions...")
# Execute questions
for q_idx, qa in enumerate(filtered_qa):
question = qa.get("question", "")
expected_answer = qa.get("answer", "")
category = qa.get("category", 0)
evidence = qa.get("evidence", [])
category_name = CATEGORY_NAMES.get(category, f"category_{category}")
# Determine which peer the question is about (returns speaker name)
target_speaker = determine_question_target(
question, speaker_a, speaker_b
)
target_peer = peer_a if target_speaker == speaker_a else peer_b
print(
f" Q{q_idx + 1} [{category_name}] (asking {target_speaker}): {question}"
)
try:
if self.use_get_context:
# Use get_context + LLM - target the appropriate peer
context = await session.get_context(
summary=True,
peer_target=target_speaker,
last_user_message=question,
)
context_messages = context.to_anthropic(assistant="assistant")
context_messages.append({"role": "user", "content": question})
response = await self.anthropic_client.messages.create(
model="claude-sonnet-4-5",
max_tokens=1024,
messages=cast(list[MessageParam], context_messages),
)
if not response.content:
raise ValueError("Anthropic returned empty response")
content_block = response.content[0]
actual_response = getattr(content_block, "text", "")
else:
# Use dialectic .chat endpoint on the appropriate peer
actual_response = await target_peer.chat(
question, session=session_id
)
actual_response = (
actual_response if isinstance(actual_response, str) else ""
)
# Get evidence context for the judge
evidence_context = get_evidence_context(conversation, evidence)
# Judge the response
judgment = await judge_response(
self.openai_client,
question,
str(expected_answer),
actual_response,
evidence_context=evidence_context,
)
passed = judgment.get("passed", False)
question_result: QuestionResult = {
"question_id": q_idx,
"question": question,
"expected_answer": str(expected_answer),
"actual_response": actual_response,
"category": category,
"category_name": category_name,
"evidence": evidence,
"judgment": judgment,
"passed": passed,
}
result["question_results"].append(question_result)
status = "PASS" if passed else "FAIL"
print(f" [{status}]")
if not passed:
print(f" Expected: {expected_answer}")
print(f" Got: {actual_response[:200]}...")
except Exception as e:
self.logger.error(f"Error executing question {q_idx}: {e}")
question_result = QuestionResult(
question_id=q_idx,
question=question,
expected_answer=str(expected_answer),
actual_response=f"ERROR: {e}",
category=category,
category_name=category_name,
evidence=evidence,
judgment={"passed": False, "reasoning": str(e)},
passed=False,
)
result["question_results"].append(question_result)
# Calculate category scores
result["category_scores"] = calculate_category_scores(
result["question_results"]
)
# Calculate overall score (pass rate)
if result["question_results"]:
passed_count = sum(
1 for qr in result["question_results"] if qr["passed"]
)
result["overall_score"] = passed_count / len(result["question_results"])
# Cleanup workspace if requested
if self.cleanup_workspace:
try:
await honcho_client.delete_workspace(workspace_id)
print(f"[{workspace_id}] Cleaned up workspace")
except Exception as e:
print(f"Failed to delete workspace: {e}")
result["end_time"] = time.time()
result["duration_seconds"] = result["end_time"] - result["start_time"]
print(
f"\n[{workspace_id}] Completed in {format_duration(result['duration_seconds'])}"
)
print(f"Overall Score: {result['overall_score']:.3f}")
except Exception as e:
self.logger.error(f"Error executing conversation {sample_id}: {e}")
result["error"] = str(e)
result["end_time"] = time.time()
result["duration_seconds"] = result["end_time"] - result["start_time"]
return result
async def run_conversations(
self,
data_file: Path,
batch_size: int = 1,
test_count: int | None = None,
sample_id: str | None = None,
question_count: int | None = None,
) -> tuple[list[ConversationResult], float]:
"""
Run multiple conversations from the LoCoMo benchmark.
Args:
data_file: Path to the LoCoMo JSON file
batch_size: Number of conversations to run concurrently in each batch
test_count: Optional number of conversations to run
sample_id: Optional sample_id to run only that conversation
question_count: Optional limit on questions per conversation
Returns:
Tuple of (list of conversation results, total duration)
"""
conversations = load_locomo_data(data_file)
# Filter by sample_id if specified
if sample_id is not None:
conversations = [
c for c in conversations if c.get("sample_id") == sample_id
]
if not conversations:
print(f"Error: No conversation found with sample_id '{sample_id}'")
return [], 0.0
print(f"Filtering to sample_id '{sample_id}'")
# Limit by test_count
if test_count is not None and test_count > 0:
conversations = conversations[:test_count]
print(f"Limiting to {len(conversations)} conversations")
print(f"Running {len(conversations)} conversations from {data_file}")
if self.pool_size > 1:
print(
f"Distributing conversations across {self.pool_size} Honcho instances"
)
overall_start = time.time()
all_results: list[ConversationResult] = []
for i in range(0, len(conversations), batch_size):
batch = conversations[i : i + batch_size]
batch_num = (i // batch_size) + 1
total_batches = (len(conversations) + batch_size - 1) // batch_size
print(f"\n{'=' * 80}")
print(
f"Processing batch {batch_num}/{total_batches} ({len(batch)} conversations)"
)
print(f"{'=' * 80}")
# Run conversations in current batch concurrently
batch_results: list[ConversationResult] = await asyncio.gather(
*[
self.execute_conversation(
conv,
self.get_honcho_url_for_index(i + idx),
question_count=question_count,
)
for idx, conv in enumerate(batch)
]
)
all_results.extend(batch_results)
overall_end = time.time()
overall_duration = overall_end - overall_start
# Finalize metrics collection
self.metrics_collector.finalize_collection()
return all_results, overall_duration
async def main() -> int:
"""Main entry point for the LoCoMo test runner."""
parser = argparse.ArgumentParser(
description="Run LoCoMo benchmark tests against a Honcho instance",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
%(prog)s --data-file tests/bench/locomo_data/locomo10.json
%(prog)s --data-file locomo10.json --pool-size 4
%(prog)s --data-file locomo10.json --sample-id "sample_0"
%(prog)s --data-file locomo10.json --test-count 5 --question-count 20
""",
)
parser.add_argument(
"--data-file",
type=Path,
required=True,
help="Path to LoCoMo JSON file (required)",
)
parser.add_argument(
"--base-api-port",
type=int,
default=8000,
help="Base port for Honcho API instances (default: 8000)",
)
parser.add_argument(
"--pool-size",
type=int,
default=1,
help="Number of Honcho instances in the pool (default: 1)",
)
parser.add_argument(
"--anthropic-api-key",
type=str,
help="Anthropic API key for response judging (optional)",
)
parser.add_argument(
"--timeout",
type=int,
default=None,
help="Timeout for deriver queue to empty in seconds (default: 10 minutes)",
)
parser.add_argument(
"--batch-size",
type=int,
default=1,
help="Number of conversations to run concurrently in each batch (default: 1)",
)
parser.add_argument(
"--json-output",
type=Path,
help="Path to write JSON summary results for analytics (optional)",
)
parser.add_argument(
"--cleanup-workspace",
action="store_true",
help="Delete workspace after executing each conversation (default: False)",
)
parser.add_argument(
"--use-get-context",
action="store_true",
help="Use get_context + judge LLM instead of dialectic .chat endpoint (default: False)",
)
parser.add_argument(
"--sample-id",
type=str,
help="Run only the conversation with this sample_id (skips all others)",
)
parser.add_argument(
"--test-count",
type=int,
help="Number of conversations to run from the data file (default: all)",
)
parser.add_argument(
"--question-count",
type=int,
help="Number of questions per conversation to run (default: all)",
)
args = parser.parse_args()
# Validate arguments
if not args.data_file.exists():
print(f"Error: Data file {args.data_file} does not exist")
return 1
if args.batch_size <= 0:
print(f"Error: Batch size must be positive, got {args.batch_size}")
return 1
if args.pool_size <= 0:
print(f"Error: Pool size must be positive, got {args.pool_size}")
return 1
# Create test runner
runner = LoCoMoRunner(
base_api_port=args.base_api_port,
pool_size=args.pool_size,
anthropic_api_key=args.anthropic_api_key,
timeout_seconds=args.timeout,
cleanup_workspace=args.cleanup_workspace,
use_get_context=args.use_get_context,
)
try:
# Run conversations
results, total_elapsed = await runner.run_conversations(
args.data_file,
args.batch_size,
args.test_count,
args.sample_id,
args.question_count,
)
print_summary(results, total_elapsed)
# Print metrics summary
runner.metrics_collector.print_summary()
# Generate JSON output
if args.json_output:
output_file = args.json_output
else:
output_file = Path(
f"tests/bench/eval_results/locomo_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
)
generate_json_summary(
results,
total_elapsed,
output_file,
metadata_extra={
"data_file": str(args.data_file),
"base_api_port": runner.base_api_port,
"pool_size": runner.pool_size,
"timeout_seconds": runner.timeout_seconds,
"deriver_settings": settings.DERIVER.model_dump(),
"dialectic_settings": settings.DIALECTIC.model_dump(),
"dream_settings": settings.DREAM.model_dump(),
"summary_settings": settings.SUMMARY.model_dump(),
},
)
# Export metrics to JSON file
metrics_output = Path(
f"tests/bench/perf_metrics/locomo_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
)
runner.metrics_collector.export_to_json(metrics_output)
runner.metrics_collector.cleanup_collection()
# Return exit code based on results
avg_score = (
sum(r["overall_score"] for r in results) / len(results) if results else 0
)
return 0 if avg_score >= 0.5 else 1
except KeyboardInterrupt:
print("\nTest execution interrupted by user")
return 1
except Exception as e:
print(f"Error running tests: {e}")
import traceback
traceback.print_exc()
return 1
if __name__ == "__main__":
exit_code = asyncio.run(main())
exit(exit_code)

534
tests/bench/locomo_baseline.py Normal file
View File

@ -0,0 +1,534 @@
"""
LoCoMo Baseline Test Runner (Direct Claude Context)
A script that executes LoCoMo benchmark tests directly against Claude
by feeding the entire conversation history into the context window.
This serves as a baseline comparison against Honcho's memory framework.
## LoCoMo Overview
LoCoMo evaluates very long-term conversational memory across five question categories:
1. Single-hop - Direct factual recall from conversations
2. Multi-hop - Reasoning across multiple pieces of information
3. Temporal - Understanding time-based relationships and sequences
4. Commonsense/World knowledge - Applying broader contextual understanding
5. Adversarial - Questions that cannot be answered (filtered out by default)
Reference: https://github.com/snap-research/locomo
Paper: https://arxiv.org/abs/2402.17753
## To use
0. Set up env:
```
uv sync
source .venv/bin/activate
```
1. Run this file with the LoCoMo dataset:
```
python -m tests.bench.locomo_baseline --data-file tests/bench/locomo_data/locomo10.json
```
Optional arguments:
```
--batch-size: Number of conversations to run concurrently in each batch (default: 1)
--json-output: Path to write JSON summary results for analytics
--sample-id: Run only the conversation with this sample_id (skips all others)
--test-count: Number of conversations to run (default: all)
--question-count: Number of questions per conversation to run (default: all)
```
## Other notes
- Uses OpenRouter API (configured via LLM_OPENAI_COMPATIBLE_API_KEY in tests/bench/.env or env var)
- Default model is anthropic/claude-haiku-4-5 for baseline comparison
- Evaluation uses LLM judge following the LoCoMo paper methodology
"""
import argparse
import asyncio
import logging
import os
import time
from datetime import datetime
from pathlib import Path
from typing import Any
from dotenv import load_dotenv
from openai import AsyncOpenAI
from src.config import settings
from .locomo_common import (
CATEGORY_NAMES,
ConversationResult,
QuestionResult,
calculate_category_scores,
calculate_tokens,
extract_sessions,
filter_questions,
format_duration,
generate_json_summary,
get_evidence_context,
judge_response,
load_locomo_data,
print_summary,
)
# Load .env from bench directory
bench_dir = Path(__file__).parent
load_dotenv(bench_dir / ".env")
# OpenRouter model format for baseline testing
MODEL_BEING_TESTED = "anthropic/claude-haiku-4.5"
class LoCoMoBaselineRunner:
"""
Executes LoCoMo benchmark tests directly against Claude.
"""
def __init__(self):
"""
Initialize the LoCoMo baseline test runner.
"""
# Configure logging
logging.basicConfig(
level=logging.WARNING, format="%(asctime)s - %(levelname)s - %(message)s"
)
self.logger: logging.Logger = logging.getLogger(__name__)
# Initialize OpenRouter client for model being tested
openrouter_api_key = os.getenv("LLM_OPENAI_COMPATIBLE_API_KEY")
openrouter_base_url = os.getenv(
"LLM_OPENAI_COMPATIBLE_BASE_URL", "https://openrouter.ai/api/v1"
)
if not openrouter_api_key:
raise ValueError(
"LLM_OPENAI_COMPATIBLE_API_KEY is not set in tests/bench/.env or environment"
)
self.openrouter_client: AsyncOpenAI = AsyncOpenAI(
api_key=openrouter_api_key,
base_url=openrouter_base_url,
)
# Initialize OpenAI client for judging responses
openai_api_key = os.getenv("OPENAI_API_KEY")
if not openai_api_key:
raise ValueError("OPENAI_API_KEY is not set")
self.openai_client: AsyncOpenAI = AsyncOpenAI(api_key=openai_api_key)
def _format_conversation_context(
self,
conversation: dict[str, Any],
) -> str:
"""
Format conversation sessions into a context string.
Args:
conversation: The conversation dict containing session data
Returns:
Formatted conversation transcript string
"""
speaker_a = conversation.get("speaker_a", "User")
speaker_b = conversation.get("speaker_b", "Assistant")
lines: list[str] = []
lines.append("=== CONVERSATION HISTORY ===\n")
lines.append(f"This is a conversation between {speaker_a} and {speaker_b}.\n")
sessions = extract_sessions(conversation)
for session_idx, (date_str, messages) in enumerate(sessions, 1):
lines.append(f"\n--- Session {session_idx} ({date_str}) ---\n")
for msg in messages:
speaker = msg.get("speaker", "Unknown")
text = msg.get("text", "")
lines.append(f"{speaker}: {text}\n")
lines.append("\n=== END CONVERSATION HISTORY ===")
return "\n".join(lines)
async def execute_conversation(
self,
conversation_data: dict[str, Any],
question_count: int | None = None,
) -> ConversationResult:
"""
Execute LoCoMo benchmark for a single conversation using direct Claude context.
Args:
conversation_data: Dictionary containing conversation and QA data
question_count: Optional limit on number of questions to run
Returns:
Conversation execution results
"""
start_time = time.time()
sample_id = conversation_data.get("sample_id", "unknown")
conversation = conversation_data.get("conversation", {})
qa_list = conversation_data.get("qa", [])
speaker_a = conversation.get("speaker_a", "User")
speaker_b = conversation.get("speaker_b", "Assistant")
print(f"\n{'=' * 80}")
print(f"Executing LoCoMo conversation {sample_id} [BASELINE]")
print(f"Speakers: {speaker_a} and {speaker_b}")
print(f"{'=' * 80}")
workspace_id = f"baseline_{sample_id}"
result: ConversationResult = {
"sample_id": sample_id,
"speaker_a": speaker_a,
"speaker_b": speaker_b,
"total_sessions": 0,
"total_turns": 0,
"total_tokens": 0,
"question_results": [],
"category_scores": {},
"overall_score": 0.0,
"error": None,
"start_time": start_time,
"end_time": 0.0,
"duration_seconds": 0.0,
}
try:
# Extract sessions and count turns/tokens
sessions = extract_sessions(conversation)
result["total_sessions"] = len(sessions)
total_tokens = 0
for _date_str, messages in sessions:
for msg in messages:
result["total_turns"] += 1
total_tokens += calculate_tokens(msg.get("text", ""))
result["total_tokens"] = total_tokens
print(
f"[{workspace_id}] Context: {len(sessions)} sessions, {result['total_turns']} turns, ~{total_tokens:,} tokens"
)
# Format conversation as context
conversation_context = self._format_conversation_context(conversation)
# Filter questions
filtered_qa = filter_questions(
qa_list,
exclude_adversarial=False,
test_count=question_count,
)
print(f"[{workspace_id}] Executing {len(filtered_qa)} questions...")
# Build system prompt with cache control for the conversation context
system_prompt = f"""You are a helpful assistant with memory of past conversations between {speaker_a} and {speaker_b}.
Below is the history of their past conversations. Use this history to answer the user's question accurately.
{conversation_context}"""
# Execute questions
for q_idx, qa in enumerate(filtered_qa):
question = qa.get("question", "")
expected_answer = qa.get("answer", "")
category = qa.get("category", 0)
evidence = qa.get("evidence", [])
category_name = CATEGORY_NAMES.get(category, f"category_{category}")
print(f" Q{q_idx + 1} [{category_name}]: {question[:80]}...")
try:
# Call model via OpenRouter with full context
response = await self.openrouter_client.chat.completions.create(
model=MODEL_BEING_TESTED,
max_tokens=settings.DIALECTIC.MAX_OUTPUT_TOKENS,
messages=[
{
"role": "system",
"content": system_prompt,
},
{
"role": "user",
"content": question,
},
],
)
if not response.choices or not response.choices[0].message.content:
actual_response = ""
else:
actual_response = response.choices[0].message.content
# Get evidence context for the judge
evidence_context = get_evidence_context(conversation, evidence)
# Judge the response
judgment = await judge_response(
self.openai_client,
question,
str(expected_answer),
actual_response,
evidence_context=evidence_context,
)
passed = judgment.get("passed", False)
question_result: QuestionResult = {
"question_id": q_idx,
"question": question,
"expected_answer": str(expected_answer),
"actual_response": actual_response,
"category": category,
"category_name": category_name,
"evidence": evidence,
"judgment": judgment,
"passed": passed,
}
result["question_results"].append(question_result)
status = "PASS" if passed else "FAIL"
print(f" [{status}]")
if not passed:
print(f" Expected: {expected_answer}")
print(f" Got: {actual_response[:200]}...")
except Exception as e:
self.logger.error(f"Error executing question {q_idx}: {e}")
question_result = QuestionResult(
question_id=q_idx,
question=question,
expected_answer=str(expected_answer),
actual_response=f"ERROR: {e}",
category=category,
category_name=category_name,
evidence=evidence,
judgment={"passed": False, "reasoning": str(e)},
passed=False,
)
result["question_results"].append(question_result)
# Calculate category scores
result["category_scores"] = calculate_category_scores(
result["question_results"]
)
# Calculate overall score (pass rate)
if result["question_results"]:
passed_count = sum(
1 for qr in result["question_results"] if qr["passed"]
)
result["overall_score"] = passed_count / len(result["question_results"])
result["end_time"] = time.time()
result["duration_seconds"] = result["end_time"] - result["start_time"]
print(
f"\n[{workspace_id}] Completed in {format_duration(result['duration_seconds'])}"
)
print(f"Overall Score: {result['overall_score']:.3f}")
except Exception as e:
self.logger.error(f"Error executing conversation {sample_id}: {e}")
result["error"] = str(e)
result["end_time"] = time.time()
result["duration_seconds"] = result["end_time"] - result["start_time"]
return result
async def run_conversations(
self,
data_file: Path,
batch_size: int = 1,
test_count: int | None = None,
sample_id: str | None = None,
question_count: int | None = None,
) -> tuple[list[ConversationResult], float]:
"""
Run multiple conversations from the LoCoMo benchmark.
Args:
data_file: Path to the LoCoMo JSON file
batch_size: Number of conversations to run concurrently in each batch
test_count: Optional number of conversations to run
sample_id: Optional sample_id to run only that conversation
question_count: Optional limit on questions per conversation
Returns:
Tuple of (list of conversation results, total duration)
"""
conversations = load_locomo_data(data_file)
# Filter by sample_id if specified
if sample_id is not None:
conversations = [
c for c in conversations if c.get("sample_id") == sample_id
]
if not conversations:
print(f"Error: No conversation found with sample_id '{sample_id}'")
return [], 0.0
print(f"Filtering to sample_id '{sample_id}'")
# Limit by test_count
if test_count is not None and test_count > 0:
conversations = conversations[:test_count]
print(f"Limiting to {len(conversations)} conversations")
print(f"Running {len(conversations)} conversations from {data_file} [BASELINE]")
overall_start = time.time()
all_results: list[ConversationResult] = []
for i in range(0, len(conversations), batch_size):
batch = conversations[i : i + batch_size]
batch_num = (i // batch_size) + 1
total_batches = (len(conversations) + batch_size - 1) // batch_size
print(f"\n{'=' * 80}")
print(
f"Processing batch {batch_num}/{total_batches} ({len(batch)} conversations)"
)
print(f"{'=' * 80}")
# Run conversations in current batch concurrently
batch_results: list[ConversationResult] = await asyncio.gather(
*[
self.execute_conversation(conv, question_count=question_count)
for conv in batch
]
)
all_results.extend(batch_results)
overall_end = time.time()
overall_duration = overall_end - overall_start
return all_results, overall_duration
async def main() -> int:
"""Main entry point for the LoCoMo baseline test runner."""
parser = argparse.ArgumentParser(
description="Run LoCoMo benchmark tests directly against Claude (baseline, no Honcho)",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
%(prog)s --data-file tests/bench/locomo_data/locomo10.json
%(prog)s --data-file locomo10.json --sample-id "sample_0"
%(prog)s --data-file locomo10.json --test-count 5 --question-count 20
""",
)
parser.add_argument(
"--data-file",
type=Path,
required=True,
help="Path to LoCoMo JSON file (required)",
)
parser.add_argument(
"--batch-size",
type=int,
default=1,
help="Number of conversations to run concurrently in each batch (default: 1)",
)
parser.add_argument(
"--json-output",
type=Path,
help="Path to write JSON summary results for analytics (optional)",
)
parser.add_argument(
"--sample-id",
type=str,
help="Run only the conversation with this sample_id (skips all others)",
)
parser.add_argument(
"--test-count",
type=int,
help="Number of conversations to run from the data file (default: all)",
)
parser.add_argument(
"--question-count",
type=int,
help="Number of questions per conversation to run (default: all)",
)
args = parser.parse_args()
# Validate arguments
if not args.data_file.exists():
print(f"Error: Data file {args.data_file} does not exist")
return 1
if args.batch_size <= 0:
print(f"Error: Batch size must be positive, got {args.batch_size}")
return 1
# Create test runner
runner = LoCoMoBaselineRunner()
try:
# Run conversations
results, total_elapsed = await runner.run_conversations(
args.data_file,
args.batch_size,
args.test_count,
args.sample_id,
args.question_count,
)
print_summary(results, total_elapsed)
# Generate JSON output
if args.json_output:
output_file = args.json_output
else:
output_file = Path(
f"tests/bench/eval_results/locomo_baseline_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
)
generate_json_summary(
results,
total_elapsed,
output_file,
metadata_extra={
"data_file": str(args.data_file),
"runner_type": "baseline_direct_context",
"model": MODEL_BEING_TESTED,
},
)
# Return exit code based on results
avg_score = (
sum(r["overall_score"] for r in results) / len(results) if results else 0
)
return 0 if avg_score >= 0.5 else 1
except KeyboardInterrupt:
print("\nTest execution interrupted by user")
return 1
except Exception as e:
print(f"Error running tests: {e}")
import traceback
traceback.print_exc()
return 1
if __name__ == "__main__":
exit_code = asyncio.run(main())
exit(exit_code)

604
tests/bench/locomo_common.py Normal file
View File

@ -0,0 +1,604 @@
"""
Common utilities for LoCoMo benchmark test runners.
Shared functionality between the Honcho benchmark and baseline benchmark.
LoCoMo evaluates very long-term conversational memory of LLM agents across
five question categories:
1. Single-hop - Direct factual recall from conversations
2. Multi-hop - Reasoning across multiple pieces of information
3. Temporal - Understanding time-based relationships and sequences
4. Commonsense/World knowledge - Applying broader contextual understanding
5. Adversarial - Challenging questions that cannot be answered from the conversation
Reference: https://github.com/snap-research/locomo
Paper: https://arxiv.org/abs/2402.17753
"""
import json
import logging
import re
from datetime import datetime
from pathlib import Path
from typing import Any
import tiktoken
from openai import AsyncOpenAI
from typing_extensions import TypedDict
logger = logging.getLogger(__name__)
# Category ID to name mapping
CATEGORY_NAMES: dict[int, str] = {
1: "single_hop",
2: "multi_hop",
3: "temporal",
4: "commonsense",
5: "adversarial", # Should be filtered out during evaluation
}
class QuestionResult(TypedDict):
"""Type definition for question evaluation results."""
question_id: int
question: str
expected_answer: str
actual_response: str
category: int
category_name: str
evidence: list[str]
judgment: dict[str, Any]
passed: bool
class ConversationResult(TypedDict):
"""Type definition for conversation execution results."""
sample_id: str
speaker_a: str
speaker_b: str
total_sessions: int
total_turns: int
total_tokens: int
question_results: list[QuestionResult]
category_scores: dict[str, dict[str, Any]]
overall_score: float
error: str | None
start_time: float
end_time: float
duration_seconds: float
def format_duration(total_seconds: float) -> str:
"""Format a duration in seconds into a human-readable string."""
minutes = int(total_seconds // 60)
if minutes > 0:
seconds_rounded = int(round(total_seconds - minutes * 60))
if seconds_rounded == 60:
minutes += 1
seconds_rounded = 0
return f"{minutes}m{seconds_rounded:02d}s"
return f"{total_seconds:.2f}s"
def calculate_tokens(text: str) -> int:
"""Calculate tokens for a given text."""
tokenizer = tiktoken.get_encoding("o200k_base")
try:
return len(
tokenizer.encode(
text,
disallowed_special=(tokenizer.special_tokens_set - {"<|endoftext|>"}),
)
)
except Exception:
return len(text) // 4
def load_locomo_data(data_file: Path) -> list[dict[str, Any]]:
"""
Load LoCoMo data from a JSON file.
Args:
data_file: Path to the LoCoMo JSON file
Returns:
List of conversation dictionaries
"""
with open(data_file) as f:
return json.load(f)
def parse_locomo_date(date_str: str) -> datetime:
"""
Parse LoCoMo date format to datetime.
Args:
date_str: Date string in format "H:MM am/pm on D Month, YYYY"
e.g., "1:56 pm on 8 May, 2023"
Returns:
Parsed datetime object
"""
try:
# Handle formats like "1:56 pm on 8 May, 2023"
# Remove "on " and parse
date_str = date_str.replace(" on ", " ")
# Try parsing with different formats
for fmt in ["%I:%M %p %d %B, %Y", "%I:%M %p %d %B %Y"]:
try:
return datetime.strptime(date_str, fmt)
except ValueError:
continue
# If all fail, return a default datetime
logger.warning(f"Could not parse date '{date_str}', using current time")
return datetime.now()
except Exception as e:
logger.warning(f"Error parsing date '{date_str}': {e}")
return datetime.now()
def extract_sessions(
conversation: dict[str, Any],
) -> list[tuple[str, list[dict[str, Any]]]]:
"""
Extract sessions from a LoCoMo conversation.
Args:
conversation: The conversation dict containing session_N and session_N_date_time
Returns:
List of tuples (date_time_str, messages) where messages contain
'speaker', 'text', 'dia_id', and optional 'img_url', 'blip_caption', 'query'
"""
sessions: list[tuple[str, list[dict[str, Any]]]] = []
# Find all session keys
session_keys = sorted(
[k for k in conversation if re.match(r"session_\d+$", k)],
key=lambda x: int(x.split("_")[1]),
)
for session_key in session_keys:
session_num = session_key.split("_")[1]
date_key = f"session_{session_num}_date_time"
date_str = conversation.get(date_key, "")
messages = conversation.get(session_key, [])
sessions.append((date_str, messages))
return sessions
def extract_all_messages(
conversation: dict[str, Any],
) -> list[dict[str, Any]]:
"""
Extract all messages from all sessions in a conversation.
Args:
conversation: The conversation dict
Returns:
List of message dicts with 'speaker', 'text', 'dia_id', and optional image fields
"""
all_messages: list[dict[str, Any]] = []
sessions = extract_sessions(conversation)
for _date_str, messages in sessions:
for msg in messages:
message_dict: dict[str, Any] = {
"speaker": msg.get("speaker", ""),
"text": msg.get("text", ""),
"dia_id": msg.get("dia_id", ""),
}
# Preserve image fields if present
if msg.get("blip_caption"):
message_dict["blip_caption"] = msg["blip_caption"]
if msg.get("img_url"):
message_dict["img_url"] = msg["img_url"]
if msg.get("query"):
message_dict["query"] = msg["query"]
all_messages.append(message_dict)
return all_messages
def get_evidence_context(
conversation: dict[str, Any],
evidence_ids: list[str],
) -> str | None:
"""
Extract evidence messages from a conversation based on dia_id references.
Args:
conversation: The conversation dict containing sessions
evidence_ids: List of dia_id references (e.g., ["D1:3", "D2:8"])
Returns:
Formatted string of evidence messages, or None if no evidence found
"""
if not evidence_ids:
return None
# Build a mapping of dia_id to message
all_messages = extract_all_messages(conversation)
dia_id_to_msg = {msg["dia_id"]: msg for msg in all_messages if msg.get("dia_id")}
# Extract evidence messages
evidence_messages: list[str] = []
for eid in evidence_ids:
if eid in dia_id_to_msg:
msg = dia_id_to_msg[eid]
text = msg["text"]
# Include image caption if present
if msg.get("blip_caption"):
text = f"{text} [Image: {msg['blip_caption']}]"
evidence_messages.append(f"[{eid}] {msg['speaker']}: {text}")
if not evidence_messages:
return None
return "\n".join(evidence_messages)
def filter_questions(
qa_list: list[dict[str, Any]],
exclude_adversarial: bool = False,
question_ids: list[int] | None = None,
test_count: int | None = None,
) -> list[dict[str, Any]]:
"""
Filter questions based on criteria.
Args:
qa_list: List of QA dictionaries
exclude_adversarial: If True, exclude category 5 (adversarial) questions
question_ids: Optional list of specific question indices to include
test_count: Optional limit on number of questions
Returns:
Filtered list of questions
"""
filtered = qa_list
# Filter out adversarial questions (category 5)
if exclude_adversarial:
filtered = [q for q in filtered if q.get("category") != 5]
# Filter by specific question IDs
if question_ids is not None:
filtered = [
q
for i, q in enumerate(filtered)
if i in question_ids or (i + 1) in question_ids
]
# Limit to first N questions
if test_count is not None and test_count > 0:
filtered = filtered[:test_count]
return filtered
def _build_judge_system_prompt(context: str | None) -> str:
"""Build the system prompt for LoCoMo evaluation.
Args:
context: Optional evidence context from the conversation
Returns:
The system prompt for the judge model
"""
return f"""You are evaluating whether a synthesized answer adequately addresses a query about a user based on available conclusions.
## EVIDENCE CONTEXT
{context if context else "No evidence provided."}
## EVALUATION CONTEXT
You will evaluate:
1. **Query**: The specific question asked about the user
2. **Synthesized Answer**: The response generated from available conclusions
3. **Gold Standard Answer**: The expected/correct answer
## EVALUATION CRITERIA
Judge the synthesized answer as SUFFICIENT or INSUFFICIENT based on:
### Content Completeness
- Does the answer address what the query is asking?
- Are all key aspects of the gold answer covered (even if phrased differently)?
- Is critical information missing that would change the answer's usefulness?
### Semantic Accuracy
- Are any factual errors or contradictions present?
## ACCEPTABLE DIFFERENCES
The following differences are ACCEPTABLE and should NOT result in INSUFFICIENT:
- Different phrasing or word choice that still conveys the same or very similar meaning, especially in cases where the question is tentative or open-ended.
- Additional relevant context beyond the gold answer (including evidence supplied above). This includes the case where the synthesized answer is longer and more detailed than the gold answer, potentially even including additional information that is not explicitly stated in the gold answer but is still broadly relevant to the query. Do NOT penalize the synthesized answer for including additional information that is not explicitly stated in the gold answer.
- **The synthesized answer explicitly includes the full gold answer text (even if surrounded by additional or unrelated details). If the gold answer appears within the synthesized answer, you MUST mark the answer as SUFFICIENT.**
- More detailed explanations of reasoning or evidence
- Appropriate confidence qualifiers (e.g., "likely", "probably") when warranted
- Differences in length, with the synthesized answer being longer and even more circuitous or indirect in its addressing of the query, as long as it conveys the same meaning
- Minor format or structure variations
## EVIDENCE-GOLD ANSWER CONSISTENCY CHECK
It is possible for the gold answers to be wrong. Sometimes it may not be fully supported by or follow logically from the evidence messages, instead constituting a guess or assumption. Additionally, the gold answers are generated automatically based on the limited set of evidence messages provided above, whereas if additional context were to be taken into account, the answer might be different. In these cases, we must not penalize the synthesized answer for not being exactly the same as the gold answer.
Before deciding, verify whether the gold answer logically and necessarily follows from the supplied evidence context. If you identify a mismatch or missing logical link **and** the synthesized answer acknowledges this uncertainty or provides a more cautious, evidence-grounded explanation (optionally leveraging additional context beyond the ground truth evidence above), treat the synthesized answer as SUFFICIENT even when it diverges in wording or conclusion from the gold answer. In short:
* If the gold answer over-claims beyond what the evidence shows, do **not** penalize a synthesized answer that appropriately qualifies the claim or offers a plausible alternative consistent with evidence.
* This includes the case where the synthesized answer is ambivalent or uncertain about the answer, as long as it provides sufficient evidence to support not providing a definitive, categorical answer.
* If the synthesized answer clearly explains the gap and gives a better-supported conclusion, mark it SUFFICIENT.
## UNACCEPTABLE DIFFERENCES
The following DO warrant an INSUFFICIENT rating:
- Irreconcilable errors or contradictions with the gold answer **and** the evidence context
- Missing information central to answering the query, such that its absence would change the meaning of the answer
- Does not address the question being asked
## YOUR TASK
First, analyze what the query is asking **and** how well both answers are supported by the evidence context.
Then, provide 2 brief 2-3 sentence arguments for both SUFFICIENT and INSUFFICIENT:
**Arguments for SUFFICIENT:**
- List reasons why the synthesized answer adequately addresses the query
- Note what key information from the gold answer is present or why deviations are justified by the evidence
- Note whether the gold answer is wrong or not necessarily true given the evidence above
**Arguments for INSUFFICIENT:**
- List reasons why the synthesized answer fails to address the question.
Based on weighing these arguments, provide 2-3 sentences to determine if the synthesized answer is sufficient. In your weighing, consider whether the synthesized answer might be a better answer than the gold answer given the evidence above.
Finally, set is_sufficient to true if sufficient or false if insufficient.
Your response MUST be a valid JSON object with EXACTLY these keys:
- arguments_for_sufficient (string)
- arguments_for_insufficient (string)
- final_reasoning (string)
- is_sufficient (boolean)
Return ONLY this JSON object and nothing else."""
def _build_judge_user_prompt(
question: str,
answer: str,
response: str,
) -> str:
"""Build the user prompt for LoCoMo evaluation.
Args:
question: The question asked
answer: Expected answer from the test
response: Actual response from the system under test
Returns:
The user prompt for the judge model
"""
return f"""Query: {question}
Gold Answer: {answer}
Synthesized Answer: {response}"""
async def judge_response(
openai_client: AsyncOpenAI,
question: str,
expected_answer: str,
actual_response: str,
evidence_context: str | None = None,
) -> dict[str, Any]:
"""Use GPT-4o-mini to judge if the actual response matches the expected answer.
This judge is designed to be lenient towards verbose answers that contain
additional context, as long as they address the question. It also considers
whether the gold answer is actually correct given the evidence.
Args:
openai_client: OpenAI client instance
question: The question asked
expected_answer: Expected answer from the test
actual_response: Actual response from the system under test
evidence_context: Optional evidence messages from the conversation
Returns:
Judgment result with pass/fail and reasoning
"""
try:
system_prompt = _build_judge_system_prompt(evidence_context)
user_prompt = _build_judge_user_prompt(
question, expected_answer, actual_response
)
response = await openai_client.chat.completions.create(
model="gpt-4o-mini",
max_tokens=1024,
temperature=0,
n=1,
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
)
if not response.choices:
raise ValueError("OpenAI returned empty response")
eval_response = response.choices[0].message.content
if eval_response is None:
raise ValueError("No text content in response")
# Parse JSON response
try:
# Strip any markdown code blocks if present
json_str = eval_response.strip()
if json_str.startswith("```"):
json_str = json_str.split("```")[1]
if json_str.startswith("json"):
json_str = json_str[4:]
json_str = json_str.strip()
result = json.loads(json_str)
passed = result.get("is_sufficient", False)
reasoning = result.get("final_reasoning", eval_response.strip())
return {
"passed": passed,
"reasoning": reasoning,
"arguments_for_sufficient": result.get("arguments_for_sufficient", ""),
"arguments_for_insufficient": result.get(
"arguments_for_insufficient", ""
),
}
except json.JSONDecodeError:
# Fallback: check for "sufficient" in the response
passed = 'is_sufficient": true' in eval_response.lower() or (
"sufficient" in eval_response.lower()
and "insufficient" not in eval_response.lower()
)
return {
"passed": passed,
"reasoning": eval_response.strip(),
}
except Exception as e:
logger.error(f"Error judging response: {e}")
# Fallback to simple string matching
is_correct = expected_answer.lower() in actual_response.lower()
return {
"passed": is_correct,
"reasoning": f"Fallback string matching due to error: {'Match found' if is_correct else 'No match found'}",
}
def calculate_category_scores(
question_results: list[QuestionResult],
) -> dict[str, dict[str, Any]]:
"""
Calculate scores grouped by question category.
Args:
question_results: List of question results
Returns:
Dictionary mapping category name to statistics
"""
category_stats: dict[str, dict[str, Any]] = {}
for qr in question_results:
cat_name = qr["category_name"]
if cat_name not in category_stats:
category_stats[cat_name] = {
"total": 0,
"passed": 0,
}
category_stats[cat_name]["total"] += 1
if qr["passed"]:
category_stats[cat_name]["passed"] += 1
# Calculate success rates
for cat_name in category_stats:
stats = category_stats[cat_name]
stats["success_rate"] = (
(stats["passed"] / stats["total"]) * 100 if stats["total"] > 0 else 0
)
return category_stats
def print_summary(
results: list[ConversationResult], total_elapsed_seconds: float
) -> None:
"""Print a summary of all test results."""
print(f"\n{'=' * 80}")
print("LOCOMO BENCHMARK EXECUTION SUMMARY")
print(f"{'=' * 80}")
total_conversations = len(results)
total_questions = sum(len(r["question_results"]) for r in results)
print(f"Total Conversations: {total_conversations}")
print(f"Total Questions: {total_questions}")
print(f"Total Test Time: {format_duration(total_elapsed_seconds)}")
# Aggregate category scores across all conversations
category_totals: dict[str, dict[str, Any]] = {}
for result in results:
for cat_name, stats in result["category_scores"].items():
if cat_name not in category_totals:
category_totals[cat_name] = {
"total": 0,
"passed": 0,
}
category_totals[cat_name]["total"] += stats["total"]
category_totals[cat_name]["passed"] += stats["passed"]
print("\nScores by Question Category:")
print(f"{'Category':<20} {'Total':<8} {'Passed':<8} {'Rate':<10}")
print(f"{'-' * 20} {'-' * 8} {'-' * 8} {'-' * 10}")
for cat_name in sorted(category_totals.keys()):
stats = category_totals[cat_name]
rate = (stats["passed"] / stats["total"]) * 100 if stats["total"] > 0 else 0
print(f"{cat_name:<20} {stats['total']:<8} {stats['passed']:<8} {rate:<10.1f}%")
# Overall averages
overall_scores = [r["overall_score"] for r in results]
overall_avg = sum(overall_scores) / len(overall_scores) if overall_scores else 0.0
print(f"\n{'Overall Average Score':<30}: {overall_avg:.3f}")
print(f"{'=' * 80}")
def generate_json_summary(
results: list[ConversationResult],
total_elapsed_seconds: float,
output_file: Path,
metadata_extra: dict[str, Any] | None = None,
) -> None:
"""Generate a comprehensive JSON summary of test results."""
total_conversations = len(results)
total_questions = sum(len(r["question_results"]) for r in results)
# Aggregate category scores
category_totals: dict[str, dict[str, Any]] = {}
for result in results:
for cat_name, stats in result["category_scores"].items():
if cat_name not in category_totals:
category_totals[cat_name] = {
"total": 0,
"passed": 0,
}
category_totals[cat_name]["total"] += stats["total"]
category_totals[cat_name]["passed"] += stats["passed"]
category_averages = {
cat: {
"total": stats["total"],
"passed": stats["passed"],
"success_rate": (stats["passed"] / stats["total"]) * 100
if stats["total"] > 0
else 0,
}
for cat, stats in category_totals.items()
}
# Overall averages
overall_scores = [r["overall_score"] for r in results]
overall_avg = sum(overall_scores) / len(overall_scores) if overall_scores else 0.0
metadata = {
"benchmark": "LoCoMo",
"execution_timestamp": datetime.now().isoformat(),
"runner_version": "1.0.0",
}
if metadata_extra:
metadata.update(metadata_extra)
summary = {
"metadata": metadata,
"summary_statistics": {
"total_conversations": total_conversations,
"total_questions": total_questions,
"overall_average_score": overall_avg,
"category_statistics": category_averages,
},
"timing": {
"total_duration_seconds": total_elapsed_seconds,
},
"detailed_results": results,
}
output_file.parent.mkdir(parents=True, exist_ok=True)
with open(output_file, "w") as f:
json.dump(summary, f, indent=2, default=str)
print(f"\nJSON summary written to: {output_file}")

384
tests/bench/longmem.py
View File

@ -8,7 +8,8 @@ This script:
3. Creates sessions with haystack conversations
4. Adds the answer session if present
5. Waits for the deriver queue to be empty
6. Executes the question and judges the response using an LLM
6. Triggers a dream for memory consolidation
7. Executes the question and judges the response using an LLM
## To use
@ -45,10 +46,11 @@ Optional arguments:
--merge-sessions: Merge all sessions within a question into a single session (default: False)
--cleanup-workspace: Delete workspace after executing each question (default: False)
--use-get-context: Use get_context + judge LLM instead of dialectic .chat endpoint (default: False)
--question-id: Run only the question with this question_id (skips all others)
```
## Other notes
- Judge is Claude Sonnet 4
- Judge is GPT-4o (per LongMemEval paper)
- If processing lots of data, set timeout very high or all will be lost
"""
@ -62,7 +64,7 @@ from datetime import datetime
from pathlib import Path
from typing import Any, cast
import tiktoken
import httpx
from anthropic import AsyncAnthropic
from anthropic.types import MessageParam
from dotenv import load_dotenv
@ -71,11 +73,24 @@ from honcho.async_client.session import SessionPeerConfig
from honcho_core.types.workspaces.sessions.message_create_param import (
MessageCreateParam,
)
from openai import AsyncOpenAI
from typing_extensions import TypedDict
from src.config import settings
from src.utils.metrics_collector import MetricsCollector
from .longmem_common import (
calculate_timing_statistics,
calculate_total_tokens,
calculate_type_statistics,
filter_questions,
format_duration,
judge_response,
load_test_file,
parse_longmemeval_date,
write_json_summary,
)
load_dotenv()
@ -175,6 +190,12 @@ class LongMemEvalRunner:
raise ValueError("LLM_ANTHROPIC_API_KEY is not set")
self.anthropic_client = AsyncAnthropic(api_key=api_key)
# OpenAI client for GPT-4o judge (per LongMemEval paper)
openai_api_key = os.getenv("OPENAI_API_KEY")
if not openai_api_key:
raise ValueError("OPENAI_API_KEY is not set (required for GPT-4o judge)")
self.openai_client: AsyncOpenAI = AsyncOpenAI(api_key=openai_api_key)
def get_honcho_url_for_index(self, question_index: int) -> str:
"""
Get the Honcho URL for a given question index using round-robin distribution.
@ -189,64 +210,8 @@ class LongMemEvalRunner:
port = self.base_api_port + instance_id
return f"http://localhost:{port}"
def _format_duration(self, total_seconds: float) -> str:
"""Format a duration in seconds into a human-readable string.
If the duration is at least one minute, this returns a string in the
form "XmYYs" with zero-padded seconds. Otherwise, it returns the
duration in seconds with two decimal places, e.g., "12.34s".
Args:
total_seconds: The duration in seconds.
Returns:
A formatted duration string.
"""
minutes = int(total_seconds // 60)
if minutes > 0:
seconds_rounded = int(round(total_seconds - minutes * 60))
if seconds_rounded == 60:
minutes += 1
seconds_rounded = 0
return f"{minutes}m{seconds_rounded:02d}s"
return f"{total_seconds:.2f}s"
def _calculate_total_tokens(
self, haystack_sessions: list[list[dict[str, str]]]
) -> int:
"""Calculate total tokens from all messages in all sessions.
Args:
haystack_sessions: List of sessions, each containing messages
Returns:
Total number of tokens across all messages
"""
tokenizer = tiktoken.get_encoding("cl100k_base")
total_tokens = 0
for session_messages in haystack_sessions:
for msg in session_messages:
content = msg.get("content", "")
try:
total_tokens += len(
tokenizer.encode(
content,
disallowed_special=(
tokenizer.special_tokens_set - {"<|endoftext|>"}
),
)
)
except Exception:
total_tokens += len(content) // 4
self.logger.warning(
f"Error tokenizing content. Using rough estimate of {len(content) // 4} tokens"
)
return total_tokens
def _get_latest_tokens_used(self) -> int | None:
"""Get the tokens_used_estimate from the most recent dialectic_chat metric.
def _get_latest_input_tokens_used(self) -> int | None:
"""Get the uncached input tokens from the most recent dialectic_chat metric.
Returns:
Number of tokens used, or None if not found
@ -270,7 +235,7 @@ class LongMemEvalRunner:
if task_name.startswith("dialectic_chat_"):
for metric in data.get("metrics", []):
metric_name = metric.get("name", "")
if metric_name.endswith("tokens_used_estimate"):
if metric_name.endswith("uncached_input_tokens"):
return int(metric.get("value", 0))
except (json.JSONDecodeError, KeyError, ValueError):
continue
@ -280,47 +245,6 @@ class LongMemEvalRunner:
return None
def _parse_date(self, date_str: str) -> datetime:
"""Parse longmemeval date format to datetime.
Args:
date_str: Date string in format "YYYY/MM/DD (Day) HH:MM"
Returns:
Parsed datetime object
Raises:
ValueError: If date format is invalid
"""
try:
# Extract the date and time parts, ignoring the day name in parentheses
# Format: "2023/05/20 (Sat) 02:21"
parts = date_str.split(") ")
if len(parts) != 2:
raise ValueError(f"Invalid date format: {date_str}")
date_part = parts[0].split(" (")[0] # "2023/05/20"
time_part = parts[1] # "02:21"
# Combine and parse
datetime_str = f"{date_part} {time_part}"
return datetime.strptime(datetime_str, "%Y/%m/%d %H:%M")
except (ValueError, IndexError) as e:
raise ValueError(f"Failed to parse date '{date_str}': {e}") from e
def load_test_file(self, test_file: Path) -> list[dict[str, Any]]:
"""
Load longmemeval test definitions from a JSON file.
Args:
test_file: Path to the JSON test file
Returns:
List of test question dictionaries
"""
with open(test_file) as f:
return json.load(f)
async def create_honcho_client(
self, workspace_id: str, honcho_url: str
) -> AsyncHoncho:
@ -362,88 +286,69 @@ class LongMemEvalRunner:
return False
await asyncio.sleep(1)
async def judge_response(
self, question: str, expected_answer: str, actual_response: str
) -> dict[str, Any]:
async def trigger_dream_and_wait(
self,
honcho_client: AsyncHoncho,
workspace_id: str,
observer: str,
observed: str | None = None,
session_id: str | None = None,
) -> bool:
"""
Use an LLM to judge if the actual response matches the expected answer.
Trigger a dream task and wait for it to complete.
Args:
question: The question asked
expected_answer: Expected answer from the test
actual_response: Actual response from Honcho
honcho_client: Honcho client instance
workspace_id: Workspace identifier
observer: Observer peer name
observed: Observed peer name (defaults to observer)
session_id: Session ID to scope the dream to
Returns:
Judgment result with pass/fail and reasoning
True if dream completed successfully, False on timeout
"""
observed = observed or observer
honcho_url = self.get_honcho_url_for_index(0)
url = f"{honcho_url}/v2/workspaces/{workspace_id}/trigger_dream"
payload: dict[str, Any] = {
"observer": observer,
"observed": observed,
"dream_type": "omni",
"session_id": session_id or f"{workspace_id}_session",
}
# Trigger the dream via API
try:
system_prompt = """
You are an expert judge evaluating AI responses to memory questions. Your task is to determine if an actual response contains the correct answer from long-term memory.
CRITICAL JUDGING PRINCIPLES:
1. SEMANTIC UNDERSTANDING: Focus on whether the actual response conveys the same core factual information as expected, even if expressed differently
2. FLEXIBLE INTERPRETATION: Accept responses that are longer, more detailed, or use different phrasing as long as they contain the correct answer
3. MEMORY ACCURACY: The key is whether the AI correctly recalled and stated the factual information from memory
4. PARTIAL CREDIT: If the response shows the AI accessed relevant memories but made minor errors in details, consider partial credit
5. IMPLICIT vs EXPLICIT: Accept responses that clearly imply the correct answer through context
ONLY FAIL when:
- The core factual answer is demonstrably wrong
- The response shows no evidence of accessing the relevant memory
- The AI explicitly states incorrect information that contradicts the expected answer
Always respond with valid JSON: {"passed": boolean, "reasoning": "short (1-3 sentences) explanation of why the response is correct or incorrect"}"""
user_prompt = f"""Question: "{question}"
Expected answer: "{expected_answer}"
Actual response: "{actual_response}"
Evaluate whether the actual response correctly answers the question based on the expected answer. Focus on factual accuracy and evidence that the AI accessed the correct memory."""
response = await self.anthropic_client.messages.create(
model="claude-sonnet-4-5",
max_tokens=300,
temperature=0.0,
system=system_prompt,
messages=[
{
"role": "user",
"content": user_prompt,
}
],
)
if not response.content:
raise ValueError("Anthropic returned empty response")
content_block = response.content[0]
judgment_text = getattr(content_block, "text", None)
if judgment_text is None:
raise ValueError(
f"No text content in response block: {type(content_block)}"
async with httpx.AsyncClient() as client:
response = await client.post(
url,
json=payload,
timeout=30.0,
)
# Extract JSON from the response if it's wrapped in markdown
if "```json" in judgment_text:
json_start = judgment_text.find("```json") + 7
json_end = judgment_text.find("```", json_start)
judgment_text = judgment_text[json_start:json_end].strip()
elif "```" in judgment_text:
json_start = judgment_text.find("```") + 3
json_end = judgment_text.find("```", json_start)
judgment_text = judgment_text[json_start:json_end].strip()
judgment = json.loads(judgment_text)
return judgment
if response.status_code != 204:
print(
f"[{workspace_id}] ERROR: Dream trigger failed with status {response.status_code}"
)
print(f"[{workspace_id}] Response body: {response.text}")
return False
except Exception as e:
self.logger.error(f"Error judging response: {e}")
# Fallback to simple string matching
is_correct = expected_answer.lower() in actual_response.lower()
return {
"passed": is_correct,
"reasoning": f"Fallback string matching due to error: {'Match found' if is_correct else 'No match found'}",
}
print(f"[{workspace_id}] ERROR: Dream trigger exception: {e}")
return False
print(
f"[{workspace_id}] Dream triggered successfully for {observer}/{observed}"
)
# Wait for dream queue to empty
print(f"[{workspace_id}] Waiting for dream to complete...")
await asyncio.sleep(2) # Give time for dream to be enqueued
success = await self.wait_for_deriver_queue_empty(honcho_client)
if success:
print(f"[{workspace_id}] Dream queue empty")
else:
print(f"[{workspace_id}] Dream queue timeout")
return success
async def execute_question(
self, question_data: dict[str, Any], honcho_url: str
@ -517,14 +422,14 @@ Evaluate whether the actual response correctly answers the question based on the
parsed_dates: list[datetime] = []
for date_str in haystack_dates:
try:
parsed_dates.append(self._parse_date(date_str))
parsed_dates.append(parse_longmemeval_date(date_str))
except ValueError as e:
raise ValueError(f"Error parsing date '{date_str}': {e}") from e
haystack_total_messages = sum(len(session) for session in haystack_sessions)
# Calculate total tokens available in the sessions for this question
total_available_tokens = self._calculate_total_tokens(haystack_sessions)
total_available_tokens = calculate_total_tokens(haystack_sessions)
print(
f"[{workspace_id}] processing {len(haystack_sessions)} sessions with {haystack_total_messages} total messages ({total_available_tokens} total tokens)"
@ -533,6 +438,9 @@ Evaluate whether the actual response correctly answers the question based on the
# Determine which peer should be observed based on question type
is_assistant_type = question_type == "single-session-assistant"
# Initialize merged_session_id for potential use in dream trigger
merged_session_id: str | None = None
if self.merge_sessions:
# Create a single merged session for all messages
merged_session_id = f"{workspace_id}_merged"
@ -626,7 +534,6 @@ Evaluate whether the actual response correctly answers the question based on the
)
)
else:
merged_session_id = None
# create separate sessions
# Zip together dates, session IDs, and session content
for session_date, session_id, session_messages in zip(
@ -734,6 +641,36 @@ Evaluate whether the actual response correctly answers the question based on the
results["error"] = "Deriver queue timeout"
return results
# Trigger dream for memory consolidation before questions
print(
f"[{workspace_id}] Deriver queue empty. Triggering dream consolidation..."
)
# Determine session_id for dream
dream_session_id = (
merged_session_id
if self.merge_sessions and merged_session_id
else (
haystack_session_ids[0]
if haystack_session_ids
else f"{workspace_id}_session"
)
)
# Determine observer based on question type
observer_peer = "assistant" if is_assistant_type else "user"
# Single orchestrated dream handles all reasoning types
dream_success = await self.trigger_dream_and_wait(
honcho_client,
workspace_id,
observer=observer_peer,
session_id=dream_session_id,
)
if not dream_success:
print(f"[{workspace_id}] Warning: Dream did not complete")
print(f"[{workspace_id}] Dream completed. Executing question...")
# Execute the question
output_lines.append(f"\nAsking question: {question_with_date}")
@ -796,22 +733,27 @@ Evaluate whether the actual response correctly answers the question based on the
actual_response if isinstance(actual_response, str) else ""
)
tokens_used = self._get_latest_tokens_used()
input_tokens_used = self._get_latest_input_tokens_used()
token_efficiency = None
if tokens_used is not None and total_available_tokens > 0:
efficiency_ratio = tokens_used / total_available_tokens
if input_tokens_used is not None and total_available_tokens > 0:
efficiency_ratio = input_tokens_used / total_available_tokens
token_efficiency = {
"total_available_tokens": total_available_tokens,
"tokens_used": tokens_used,
"tokens_used": input_tokens_used,
"efficiency_ratio": efficiency_ratio,
}
output_lines.append(
f" token efficiency: {efficiency_ratio:.4f} ({tokens_used}/{total_available_tokens} tokens, {efficiency_ratio * 100:.2f}%)"
f" token efficiency: {efficiency_ratio:.4f} ({input_tokens_used}/{total_available_tokens} tokens, {efficiency_ratio * 100:.2f}%)"
)
judgment = await self.judge_response(
question_with_date, expected_answer, actual_response
judgment = await judge_response(
self.openai_client,
question_with_date,
expected_answer,
actual_response,
question_type,
question_id,
)
query_result: QueryResult = {
@ -856,7 +798,7 @@ Evaluate whether the actual response correctly answers the question based on the
results["duration_seconds"] = results["end_time"] - results["start_time"]
output_lines.append(
f"\nQuestion {question_id} completed. Status: {'PASS' if results['passed'] else 'FAIL'} (Duration: {self._format_duration(results['duration_seconds'])})"
f"\nQuestion {question_id} completed. Status: {'PASS' if results['passed'] else 'FAIL'} (Duration: {format_duration(results['duration_seconds'])})"
)
except Exception as e:
@ -870,7 +812,11 @@ Evaluate whether the actual response correctly answers the question based on the
return results
async def run_all_questions(
self, test_file: Path, batch_size: int = 10, test_count: int | None = None
self,
test_file: Path,
batch_size: int = 10,
test_count: int | None = None,
question_id: str | None = None,
) -> tuple[list[TestResult], float]:
"""
Run all questions in a longmemeval test file.
@ -879,18 +825,15 @@ Evaluate whether the actual response correctly answers the question based on the
test_file: Path to the longmemeval JSON file
batch_size: Number of questions to run concurrently in each batch
test_count: Optional number of tests to run (runs first N tests)
question_id: Optional question_id to run (skips all others)
Returns:
Tuple of (list of test results, total duration)
"""
questions = self.load_test_file(test_file)
# Limit to first N questions if test_count is specified
if test_count is not None and test_count > 0:
questions = questions[:test_count]
print(
f"limiting to first {len(questions)} {'question' if len(questions) == 1 else 'questions'} from {test_file}"
)
questions = load_test_file(test_file)
questions = filter_questions(questions, test_file, question_id, test_count)
if not questions:
return [], 0.0
print(
f"found {len(questions)} {'question' if len(questions) == 1 else 'questions'} in {test_file}"
@ -967,7 +910,7 @@ Evaluate whether the actual response correctly answers the question based on the
print(f"Passed: {passed_questions}")
print(f"Failed: {failed_questions}")
print(f"Success Rate: {(passed_questions / total_questions) * 100:.1f}%")
print(f"Total Test Time: {self._format_duration(total_test_time)}")
print(f"Total Test Time: {format_duration(total_test_time)}")
efficiency_ratios: list[float] = []
for result in results:
@ -998,7 +941,7 @@ Evaluate whether the actual response correctly answers the question based on the
question_id = result["question_id"]
question_type = result["question_type"]
status = "PASS" if result.get("passed", False) else "FAIL"
duration = self._format_duration(result["duration_seconds"])
duration = format_duration(result["duration_seconds"])
workspace = result["workspace_id"]
print(
@ -1028,37 +971,10 @@ Evaluate whether the actual response correctly answers the question based on the
failed_questions = total_questions - passed_questions
# Calculate statistics by question type
type_stats: dict[str, dict[str, int | float]] = {}
for result in results:
q_type = result["question_type"]
if q_type not in type_stats:
type_stats[q_type] = {"total": 0, "passed": 0, "failed": 0}
type_stats[q_type]["total"] += 1
if result.get("passed", False):
type_stats[q_type]["passed"] += 1
else:
type_stats[q_type]["failed"] += 1
# Add success rates to type stats
for q_type in type_stats:
stats = type_stats[q_type]
stats["success_rate"] = (
(stats["passed"] / stats["total"]) * 100 if stats["total"] > 0 else 0
)
type_stats = calculate_type_statistics(results)
# Calculate timing statistics
durations = [r["duration_seconds"] for r in results]
timing_stats = {
"total_duration_seconds": total_elapsed_seconds,
"individual_test_durations": {
"min_seconds": min(durations) if durations else 0,
"max_seconds": max(durations) if durations else 0,
"mean_seconds": sum(durations) / len(durations) if durations else 0,
"median_seconds": sorted(durations)[len(durations) // 2]
if durations
else 0,
},
}
timing_stats = calculate_timing_statistics(results, total_elapsed_seconds)
# Calculate token efficiency statistics
efficiency_ratios: list[float] = []
@ -1131,10 +1047,7 @@ Evaluate whether the actual response correctly answers the question based on the
}
if output_file:
output_file.parent.mkdir(parents=True, exist_ok=True)
with open(output_file, "w") as f:
json.dump(summary, f, indent=2, default=str)
print(f"\nJSON summary written to: {output_file}")
write_json_summary(summary, output_file)
async def main() -> int:
@ -1150,6 +1063,7 @@ Examples:
%(prog)s --test-file test.json --pool-size 4 # Use 4 Honcho instances
%(prog)s --test-file test.json --base-api-port 8000 --pool-size 4 # Custom base port with pool
%(prog)s --test-file test.json --test-count 50 # Run only first 50 tests
%(prog)s --test-file test.json --question-id "q123" # Run only question with ID "q123"
""",
)
@ -1224,6 +1138,12 @@ Examples:
help="Number of tests to run from the test file (default: all tests)",
)
parser.add_argument(
"--question-id",
type=str,
help="Run only the question with this question_id (skips all others)",
)
args = parser.parse_args()
# Validate arguments
@ -1257,7 +1177,7 @@ Examples:
try:
# Run all questions
results, total_elapsed = await runner.run_all_questions(
args.test_file, args.batch_size, args.test_count
args.test_file, args.batch_size, args.test_count, args.question_id
)
runner.print_summary(results, total_elapsed_seconds=total_elapsed)

605
tests/bench/longmem_baseline.py Normal file
View File

@ -0,0 +1,605 @@
"""
LongMemEval Baseline Test Runner (Direct Context)
A script that executes longmemeval tests directly against a model
by feeding the entire haystack content into the context window.
## To use
0. Set up env:
```
uv sync
source .venv/bin/activate
```
1. Run this file with a selected test file:
```
python -m tests.bench.longmem_baseline --test-file tests/bench/longmemeval_data/longmemeval_oracle.json
```
Optional arguments:
```
--batch-size: Number of questions to run concurrently in each batch (default: 10)
--json-output: Path to write JSON summary results for analytics
--test-count: Number of tests to run (default: all)
--question-id: Run only the question with this question_id
```
## Other notes
- Uses OpenRouter API (configured via LLM_OPENAI_COMPATIBLE_API_KEY in tests/bench/.env or env var)
- Evaluation uses GPT-4o judge per the LongMemEval paper methodology
"""
import argparse
import asyncio
import logging
import os
import time
from datetime import datetime
from pathlib import Path
from typing import Any
from dotenv import load_dotenv
from openai import AsyncOpenAI
from typing_extensions import TypedDict
from .longmem_common import (
calculate_timing_statistics,
calculate_total_tokens,
calculate_type_statistics,
filter_questions,
format_duration,
judge_response,
load_test_file,
write_json_summary,
)
load_dotenv()
# OpenRouter model format for baseline testing
MODEL_BEING_TESTED = "anthropic/claude-haiku-4.5"
class QueryResult(TypedDict):
"""Type definition for query execution results."""
question: str
expected_answer: str
actual_response: str
judgment: dict[str, Any]
input_tokens: int
output_tokens: int
class TestResult(TypedDict):
"""Type definition for test execution results."""
question_id: str
question_type: str
query_executed: QueryResult | None
passed: bool
error: str | None
start_time: float
end_time: float
duration_seconds: float
total_context_tokens: int
output_lines: list[str]
class LongMemEvalBaselineRunner:
"""
Executes longmemeval tests directly against a model.
"""
def __init__(self):
"""
Initialize the baseline test runner.
"""
# Configure logging
logging.basicConfig(
level=logging.WARNING, format="%(asctime)s - %(levelname)s - %(message)s"
)
self.logger: logging.Logger = logging.getLogger(__name__)
# Initialize OpenRouter client for model being tested
openrouter_api_key = os.getenv("LLM_OPENAI_COMPATIBLE_API_KEY")
openrouter_base_url = os.getenv(
"LLM_OPENAI_COMPATIBLE_BASE_URL", "https://openrouter.ai/api/v1"
)
if not openrouter_api_key:
raise ValueError(
"LLM_OPENAI_COMPATIBLE_API_KEY is not set in tests/bench/.env or environment"
)
self.openrouter_client: AsyncOpenAI = AsyncOpenAI(
api_key=openrouter_api_key,
base_url=openrouter_base_url,
)
# OpenAI client for GPT-4o judge (per LongMemEval paper)
openai_api_key = os.getenv("OPENAI_API_KEY")
if not openai_api_key:
raise ValueError("OPENAI_API_KEY is not set (required for GPT-4o judge)")
self.openai_client: AsyncOpenAI = AsyncOpenAI(api_key=openai_api_key)
def _format_conversation_context(
self,
haystack_sessions: list[list[dict[str, str]]],
haystack_dates: list[str],
_question_type: str,
) -> str:
"""
Format haystack sessions into a conversation transcript for context.
Args:
haystack_sessions: List of sessions, each containing messages
haystack_dates: List of date strings corresponding to sessions
question_type: Type of question (used to determine perspective)
Returns:
Formatted conversation transcript string
"""
lines: list[str] = []
lines.append("=== CONVERSATION HISTORY ===\n")
for session_idx, (session_messages, date_str) in enumerate(
zip(haystack_sessions, haystack_dates, strict=True)
):
lines.append(f"--- Session {session_idx + 1} ({date_str}) ---\n")
for msg in session_messages:
role = msg.get("role", "unknown")
content = msg.get("content", "")
role_label = "User" if role == "user" else "Assistant"
lines.append(f"{role_label}: {content}\n")
lines.append("") # Blank line between sessions
lines.append("=== END CONVERSATION HISTORY ===")
return "\n".join(lines)
async def execute_question(
self, question_data: dict[str, Any], _question_index: int
) -> TestResult:
"""
Execute a single longmemeval question by sending full context to Claude.
Args:
question_data: Dictionary containing question data
question_index: Index of the question (for logging)
Returns:
Test execution results
"""
question_id = question_data["question_id"]
question_type = question_data["question_type"]
question = question_data["question"]
expected_answer = question_data["answer"]
question_date = question_data.get("question_date", "")
question_with_date = (
f"[{question_date}] {question}" if question_date else question
)
output_lines: list[str] = []
output_lines.append(
f"\033[1mExecuting question {question_id} ({question_type})\033[0m"
)
output_lines.append(f"Question: {question_with_date}")
output_lines.append(f"Expected: {expected_answer}")
results: TestResult = {
"question_id": question_id,
"question_type": question_type,
"query_executed": None,
"passed": False,
"error": None,
"start_time": time.time(),
"end_time": 0.0,
"duration_seconds": 0.0,
"total_context_tokens": 0,
"output_lines": output_lines,
}
try:
haystack_dates = question_data.get("haystack_dates", [])
haystack_sessions = question_data.get("haystack_sessions", [])
# Calculate total tokens
total_context_tokens = calculate_total_tokens(haystack_sessions)
results["total_context_tokens"] = total_context_tokens
haystack_total_messages = sum(len(s) for s in haystack_sessions)
output_lines.append(
f"Context: {len(haystack_sessions)} sessions, {haystack_total_messages} messages, ~{total_context_tokens} tokens"
)
# Format conversation history as context
conversation_context = self._format_conversation_context(
haystack_sessions, haystack_dates, question_type
)
# Build system prompt based on question type
if question_type == "single-session-assistant":
perspective = "You are the assistant in these conversations."
else:
perspective = "You are helping a user recall information from their past conversations."
system_prompt = f"""{perspective}
Below is a history of past conversations. Use this history to answer the user's question accurately.
{conversation_context}"""
# Call model via OpenRouter with full context
response = await self.openrouter_client.chat.completions.create(
model=MODEL_BEING_TESTED,
max_tokens=8192,
messages=[
{
"role": "system",
"content": system_prompt,
},
{
"role": "user",
"content": question_with_date,
},
],
)
if not response.choices or not response.choices[0].message.content:
raise ValueError("OpenRouter returned empty response")
actual_response = response.choices[0].message.content
input_tokens = response.usage.prompt_tokens if response.usage else 0
output_tokens = response.usage.completion_tokens if response.usage else 0
output_lines.append(
f" API usage: {input_tokens} input tokens, {output_tokens} output tokens"
)
# Judge the response
judgment = await judge_response(
self.openai_client,
question_with_date,
expected_answer,
actual_response,
question_type,
question_id,
)
query_result: QueryResult = {
"question": question_with_date,
"expected_answer": expected_answer,
"actual_response": actual_response,
"judgment": judgment,
"input_tokens": input_tokens,
"output_tokens": output_tokens,
}
results["query_executed"] = query_result
results["passed"] = judgment["passed"]
output_lines.append(
" judgment: \033[1m\033[32mPASS\033[0m"
if judgment["passed"]
else " judgment: \033[1m\033[31mFAIL\033[0m"
)
if not judgment["passed"]:
output_lines.append(f" got response: \033[3m{actual_response}\033[0m")
output_lines.append(f" expected: {expected_answer}")
output_lines.append(f" reasoning: {judgment['reasoning']}")
except Exception as e:
self.logger.error(f"Error executing question {question_id}: {e}")
results["error"] = str(e)
results["passed"] = False
output_lines.append(f"Error executing question {question_id}: {e}")
results["end_time"] = time.time()
results["duration_seconds"] = results["end_time"] - results["start_time"]
output_lines.append(
f"\nQuestion {question_id} completed. Status: {'PASS' if results['passed'] else 'FAIL'} (Duration: {format_duration(results['duration_seconds'])})"
)
return results
async def run_all_questions(
self,
test_file: Path,
batch_size: int = 10,
test_count: int | None = None,
question_id: str | None = None,
) -> tuple[list[TestResult], float]:
"""
Run all questions in a longmemeval test file.
Args:
test_file: Path to the longmemeval JSON file
batch_size: Number of questions to run concurrently in each batch
test_count: Optional number of tests to run (runs first N tests)
question_id: Optional question_id to run (skips all others)
Returns:
Tuple of (list of test results, total duration)
"""
questions = load_test_file(test_file)
questions = filter_questions(questions, test_file, question_id, test_count)
if not questions:
return [], 0.0
print(
f"found {len(questions)} {'question' if len(questions) == 1 else 'questions'} in {test_file}"
)
overall_start = time.time()
# Process questions in batches
all_results: list[TestResult] = []
for i in range(0, len(questions), batch_size):
batch = questions[i : i + batch_size]
batch_num = (i // batch_size) + 1
total_batches = (len(questions) + batch_size - 1) // batch_size
print(f"\n{'=' * 60}")
print(
f"Processing batch {batch_num}/{total_batches} ({len(batch)} questions)"
)
print(f"{'=' * 60}")
# Run questions in current batch concurrently
batch_results: list[TestResult] = await asyncio.gather(
*[self.execute_question(q, i + idx) for idx, q in enumerate(batch)]
)
# Print detailed per-question outputs for this batch
for result in batch_results:
print(f"\n{'=' * 60}")
print("\n".join(result.get("output_lines", [])))
print(f"{'=' * 60}\n")
all_results.extend(batch_results)
overall_end = time.time()
overall_duration = overall_end - overall_start
return all_results, overall_duration
def print_summary(
self, results: list[TestResult], total_elapsed_seconds: float | None = None
) -> None:
"""Print a summary of all test results."""
print(f"\n{'=' * 80}")
print(
f"LONGMEMEVAL BASELINE TEST SUMMARY (Direct Context with {MODEL_BEING_TESTED})"
)
print(f"{'=' * 80}")
total_questions = len(results)
passed_questions = sum(1 for r in results if r.get("passed", False))
failed_questions = total_questions - passed_questions
total_test_time = (
total_elapsed_seconds
if total_elapsed_seconds is not None
else sum(r["duration_seconds"] for r in results)
)
print(f"Total Questions: {total_questions}")
print(f"Passed: {passed_questions}")
print(f"Failed: {failed_questions}")
print(f"Success Rate: {(passed_questions / total_questions) * 100:.1f}%")
print(f"Total Test Time: {format_duration(total_test_time)}")
# Token usage statistics
total_input_tokens = 0
total_output_tokens = 0
total_context_tokens = 0
for result in results:
total_context_tokens += result.get("total_context_tokens", 0)
query = result.get("query_executed")
if query:
total_input_tokens += query.get("input_tokens", 0)
total_output_tokens += query.get("output_tokens", 0)
print("\nToken Usage:")
print(f" Total Context Tokens (estimated): {total_context_tokens:,}")
print(f" Total Input Tokens (API): {total_input_tokens:,}")
print(f" Total Output Tokens (API): {total_output_tokens:,}")
print("\nDetailed Results:")
print(
f"{'Question ID':<15} {'Type':<25} {'Status':<8} {'Duration':<10} {'Input Tokens':<15}"
)
print(f"{'-' * 15} {'-' * 25} {'-' * 8} {'-' * 10} {'-' * 15}")
for result in results:
question_id = result["question_id"]
question_type = result["question_type"]
status = "PASS" if result.get("passed", False) else "FAIL"
duration = format_duration(result["duration_seconds"])
query = result.get("query_executed")
input_tokens = query.get("input_tokens", 0) if query else 0
print(
f"{question_id:<15} {question_type:<25} {status:<8} {duration:<10} {input_tokens:<15,}"
)
print(f"{'=' * 80}")
def generate_json_summary(
self,
results: list[TestResult],
test_file: Path,
total_elapsed_seconds: float,
output_file: Path | None = None,
) -> None:
"""Generate a comprehensive JSON summary of test results."""
total_questions = len(results)
passed_questions = sum(1 for r in results if r.get("passed", False))
failed_questions = total_questions - passed_questions
# Calculate statistics by question type
type_stats = calculate_type_statistics(results)
# Calculate timing statistics
timing_stats = calculate_timing_statistics(results, total_elapsed_seconds)
# Calculate token usage statistics
total_input_tokens = 0
total_output_tokens = 0
total_context_tokens = 0
for result in results:
total_context_tokens += result.get("total_context_tokens", 0)
query = result.get("query_executed")
if query:
total_input_tokens += query.get("input_tokens", 0)
total_output_tokens += query.get("output_tokens", 0)
token_stats = {
"total_context_tokens_estimated": total_context_tokens,
"total_input_tokens": total_input_tokens,
"total_output_tokens": total_output_tokens,
"mean_input_tokens": total_input_tokens / len(results) if results else 0,
}
summary = {
"metadata": {
"test_file": str(test_file),
"execution_timestamp": datetime.now().isoformat(),
"runner_type": "baseline_direct_context",
"model": MODEL_BEING_TESTED,
},
"summary_statistics": {
"total_questions": total_questions,
"passed": passed_questions,
"failed": failed_questions,
"success_rate_percent": (passed_questions / total_questions) * 100
if total_questions > 0
else 0,
"statistics_by_type": type_stats,
},
"timing": timing_stats,
"token_usage": token_stats,
"detailed_results": [
{
"question_id": result["question_id"],
"question_type": result["question_type"],
"passed": result.get("passed", False),
"duration_seconds": result["duration_seconds"],
"start_time": result["start_time"],
"end_time": result["end_time"],
"total_context_tokens": result.get("total_context_tokens", 0),
"error": result.get("error"),
"query_executed": result.get("query_executed"),
}
for result in results
],
}
if output_file:
write_json_summary(summary, output_file)
async def main() -> int:
"""Main entry point for the baseline test runner."""
parser = argparse.ArgumentParser(
description="Run longmemeval tests directly against a model",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
%(prog)s --test-file tests/bench/longmemeval_data/longmemeval_s.json
%(prog)s --test-file test.json --test-count 50
%(prog)s --test-file test.json --question-id "q123"
""",
)
parser.add_argument(
"--test-file",
type=Path,
required=True,
help="Path to longmemeval JSON file (required)",
)
parser.add_argument(
"--batch-size",
type=int,
default=10,
help="Number of questions to run concurrently in each batch (default: 10)",
)
parser.add_argument(
"--json-output",
type=Path,
help="Path to write JSON summary results for analytics (optional)",
)
parser.add_argument(
"--test-count",
type=int,
help="Number of tests to run from the test file (default: all tests)",
)
parser.add_argument(
"--question-id",
type=str,
help="Run only the question with this question_id (skips all others)",
)
args = parser.parse_args()
# Validate arguments
if not args.test_file.exists():
print(f"Error: Test file {args.test_file} does not exist")
return 1
if args.batch_size <= 0:
print(f"Error: Batch size must be positive, got {args.batch_size}")
return 1
if args.test_count is not None and args.test_count <= 0:
print(f"Error: Test count must be positive, got {args.test_count}")
return 1
# Create test runner
runner = LongMemEvalBaselineRunner()
try:
# Run all questions
results, total_elapsed = await runner.run_all_questions(
args.test_file, args.batch_size, args.test_count, args.question_id
)
runner.print_summary(results, total_elapsed_seconds=total_elapsed)
# Generate JSON output
if args.json_output:
runner.generate_json_summary(
results, args.test_file, total_elapsed, args.json_output
)
else:
default_output = Path(
f"tests/bench/eval_results/baseline_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
)
runner.generate_json_summary(
results, args.test_file, total_elapsed, default_output
)
# Return exit code based on results
all_passed = all(r.get("passed", False) for r in results)
return 0 if all_passed else 1
except KeyboardInterrupt:
print("\nTest execution interrupted by user")
return 1
except Exception as e:
print(f"Error running tests: {e}")
return 1
if __name__ == "__main__":
exit_code = asyncio.run(main())
exit(exit_code)

406
tests/bench/longmem_common.py Normal file
View File

@ -0,0 +1,406 @@
"""
Common utilities for LongMemEval test runners.
"""
import json
import logging
from collections.abc import Sequence
from datetime import datetime
from pathlib import Path
from typing import Any
import tiktoken
from openai import AsyncOpenAI
from typing_extensions import TypedDict
logger = logging.getLogger(__name__)
class BaseQueryResult(TypedDict):
"""Base type definition for query execution results."""
question: str
expected_answer: str
actual_response: str
judgment: dict[str, Any]
class BaseTestResult(TypedDict):
"""Base type definition for test execution results."""
question_id: str
question_type: str
passed: bool
error: str | None
start_time: float
end_time: float
duration_seconds: float
output_lines: list[str]
def format_duration(total_seconds: float) -> str:
"""Format a duration in seconds into a human-readable string.
If the duration is at least one minute, this returns a string in the
form "XmYYs" with zero-padded seconds. Otherwise, it returns the
duration in seconds with two decimal places, e.g., "12.34s".
Args:
total_seconds: The duration in seconds.
Returns:
A formatted duration string.
"""
minutes = int(total_seconds // 60)
if minutes > 0:
seconds_rounded = int(round(total_seconds - minutes * 60))
if seconds_rounded == 60:
minutes += 1
seconds_rounded = 0
return f"{minutes}m{seconds_rounded:02d}s"
return f"{total_seconds:.2f}s"
def calculate_total_tokens(haystack_sessions: list[list[dict[str, str]]]) -> int:
"""Calculate total tokens from all messages in all sessions.
Args:
haystack_sessions: List of sessions, each containing messages
Returns:
Total number of tokens across all messages
"""
tokenizer = tiktoken.get_encoding("o200k_base")
total_tokens = 0
for session_messages in haystack_sessions:
for msg in session_messages:
content = msg.get("content", "")
try:
total_tokens += len(
tokenizer.encode(
content,
disallowed_special=(
tokenizer.special_tokens_set - {"<|endoftext|>"}
),
)
)
except Exception:
total_tokens += len(content) // 4
logger.warning(
f"Error tokenizing content. Using rough estimate of {len(content) // 4} tokens"
)
return total_tokens
def parse_longmemeval_date(date_str: str) -> datetime:
"""Parse longmemeval date format to datetime.
Args:
date_str: Date string in format "YYYY/MM/DD (Day) HH:MM"
Returns:
Parsed datetime object
Raises:
ValueError: If date format is invalid
"""
try:
# Extract the date and time parts, ignoring the day name in parentheses
# Format: "2023/05/20 (Sat) 02:21"
parts = date_str.split(") ")
if len(parts) != 2:
raise ValueError(f"Invalid date format: {date_str}")
date_part = parts[0].split(" (")[0] # "2023/05/20"
time_part = parts[1] # "02:21"
# Combine and parse
datetime_str = f"{date_part} {time_part}"
return datetime.strptime(datetime_str, "%Y/%m/%d %H:%M")
except (ValueError, IndexError) as e:
raise ValueError(f"Failed to parse date '{date_str}': {e}") from e
def load_test_file(test_file: Path) -> list[dict[str, Any]]:
"""Load longmemeval test definitions from a JSON file.
Args:
test_file: Path to the JSON test file
Returns:
List of test question dictionaries
"""
with open(test_file) as f:
return json.load(f)
def _build_judge_prompt(
question_type: str,
question: str,
answer: str,
response: str,
question_id: str,
) -> str:
"""Build the judge prompt matching the official LongMemEval evaluation code.
Based on get_anscheck_prompt() from the official LongMemEval repository.
Args:
question_type: Type of question being evaluated
question: The question asked
answer: Expected answer from the test
response: Actual response from the system under test
question_id: Question ID (used to detect abstention questions)
Returns:
The complete prompt for the judge model
"""
# Check for abstention questions (have '_abs' in question_id)
if "_abs" in question_id:
return (
"I will give you an unanswerable question, an explanation, and a response "
"from a model. Please answer yes if the model correctly identifies the "
"question as unanswerable. The model could say that the information is "
"incomplete, or some other information is given but the asked information "
f"is not.\n\nQuestion: {question}\n\nExplanation: {answer}\n\n"
f"Model Response: {response}\n\nDoes the model correctly identify the "
"question as unanswerable? Answer yes or no only."
)
# Standard prompts by question type
if question_type in (
"single-session-user",
"single-session-assistant",
"multi-session",
):
return (
"I will give you a question, a correct answer, and a response from a model. "
"Please answer yes if the response contains the correct answer. Otherwise, "
"answer no. If the response is equivalent to the correct answer or contains "
"all the intermediate steps to get the correct answer, you should also answer "
"yes. If the response only contains a subset of the information required by "
f"the answer, answer no. \n\nQuestion: {question}\n\nCorrect Answer: {answer}"
f"\n\nModel Response: {response}\n\nIs the model response correct? Answer yes or no only."
)
elif question_type == "temporal-reasoning":
return (
"I will give you a question, a correct answer, and a response from a model. "
"Please answer yes if the response contains the correct answer. Otherwise, "
"answer no. If the response is equivalent to the correct answer or contains "
"all the intermediate steps to get the correct answer, you should also answer "
"yes. If the response only contains a subset of the information required by "
"the answer, answer no. In addition, do not penalize off-by-one errors for "
"the number of days. If the question asks for the number of days/weeks/months, "
"etc., and the model makes off-by-one errors (e.g., predicting 19 days when "
"the answer is 18), the model's response is still correct. \n\n"
f"Question: {question}\n\nCorrect Answer: {answer}\n\nModel Response: {response}"
"\n\nIs the model response correct? Answer yes or no only."
)
elif question_type == "knowledge-update":
return (
"I will give you a question, a correct answer, and a response from a model. "
"Please answer yes if the response contains the correct answer. Otherwise, "
"answer no. If the response contains some previous information along with an "
"updated answer, the response should be considered as correct as long as the "
f"updated answer is the required answer.\n\nQuestion: {question}\n\n"
f"Correct Answer: {answer}\n\nModel Response: {response}\n\n"
"Is the model response correct? Answer yes or no only."
)
elif question_type == "single-session-preference":
return (
"I will give you a question, a rubric for desired personalized response, "
"and a response from a model. Please answer yes if the response satisfies "
"the desired response. Otherwise, answer no. The model does not need to "
"reflect all the points in the rubric. The response is correct as long as "
"it recalls and utilizes the user's personal information correctly.\n\n"
f"Question: {question}\n\nRubric: {answer}\n\nModel Response: {response}"
"\n\nIs the model response correct? Answer yes or no only."
)
else:
# Default case (same as multi-session)
return (
"I will give you a question, a correct answer, and a response from a model. "
"Please answer yes if the response contains the correct answer. Otherwise, "
"answer no. If the response is equivalent to the correct answer or contains "
"all the intermediate steps to get the correct answer, you should also answer "
"yes. If the response only contains a subset of the information required by "
f"the answer, answer no. \n\nQuestion: {question}\n\nCorrect Answer: {answer}"
f"\n\nModel Response: {response}\n\nIs the model response correct? Answer yes or no only."
)
async def judge_response(
openai_client: AsyncOpenAI,
question: str,
expected_answer: str,
actual_response: str,
question_type: str = "default",
question_id: str = "",
) -> dict[str, Any]:
"""Use GPT-4o to judge if the actual response matches the expected answer.
Uses the exact prompt format from the official LongMemEval evaluation code
(evaluate_qa.py) to ensure consistent evaluation.
Args:
openai_client: OpenAI client instance
question: The question asked
expected_answer: Expected answer from the test
actual_response: Actual response from the system under test
question_type: Type of question (temporal-reasoning, knowledge-update,
single-session-preference, single-session-user,
single-session-assistant, multi-session)
question_id: Question ID (used to detect abstention questions with '_abs')
Returns:
Judgment result with pass/fail and reasoning
"""
try:
prompt = _build_judge_prompt(
question_type, question, expected_answer, actual_response, question_id
)
response = await openai_client.chat.completions.create(
model="gpt-4o-2024-08-06",
max_tokens=10,
temperature=0,
n=1,
messages=[{"role": "user", "content": prompt}],
)
if not response.choices:
raise ValueError("OpenAI returned empty response")
eval_response = response.choices[0].message.content
if eval_response is None:
raise ValueError("No text content in response")
# Match official evaluation: check if "yes" appears in lowercased response
passed = "yes" in eval_response.lower()
return {
"passed": passed,
"reasoning": eval_response.strip(),
}
except Exception as e:
logger.error(f"Error judging response: {e}")
# Fallback to simple string matching
is_correct = expected_answer.lower() in actual_response.lower()
return {
"passed": is_correct,
"reasoning": f"Fallback string matching due to error: {'Match found' if is_correct else 'No match found'}",
}
def filter_questions(
questions: list[dict[str, Any]],
test_file: Path,
question_id: str | None = None,
test_count: int | None = None,
) -> list[dict[str, Any]]:
"""Filter questions by question_id and/or test_count.
Args:
questions: List of question dictionaries
test_file: Path to test file (for logging)
question_id: Optional question_id to filter to
test_count: Optional limit on number of questions
Returns:
Filtered list of questions
"""
# Filter by question_id if specified
if question_id is not None:
original_count = len(questions)
questions = [q for q in questions if q.get("question_id") == question_id]
if not questions:
print(
f"Error: No question found with question_id '{question_id}' in {test_file}"
)
return []
print(
f"filtering to question_id '{question_id}' ({len(questions)}/{original_count} {'question' if len(questions) == 1 else 'questions'})"
)
# Limit to first N questions if test_count is specified
if test_count is not None and test_count > 0:
questions = questions[:test_count]
print(
f"limiting to first {len(questions)} {'question' if len(questions) == 1 else 'questions'} from {test_file}"
)
return questions
def calculate_type_statistics(
results: Sequence[Any],
) -> dict[str, dict[str, int | float]]:
"""Calculate pass/fail statistics grouped by question type.
Args:
results: List of test results
Returns:
Dictionary mapping question type to statistics
"""
type_stats: dict[str, dict[str, int | float]] = {}
for result in results:
q_type = result["question_type"]
if q_type not in type_stats:
type_stats[q_type] = {"total": 0, "passed": 0, "failed": 0}
type_stats[q_type]["total"] += 1
if result.get("passed", False):
type_stats[q_type]["passed"] += 1
else:
type_stats[q_type]["failed"] += 1
# Add success rates
for q_type in type_stats:
stats = type_stats[q_type]
stats["success_rate"] = (
(stats["passed"] / stats["total"]) * 100 if stats["total"] > 0 else 0
)
return type_stats
def calculate_timing_statistics(
results: Sequence[Any], total_elapsed_seconds: float
) -> dict[str, Any]:
"""Calculate timing statistics from test results.
Args:
results: List of test results
total_elapsed_seconds: Total elapsed time for all tests
Returns:
Dictionary of timing statistics
"""
durations = [r["duration_seconds"] for r in results]
return {
"total_duration_seconds": total_elapsed_seconds,
"individual_test_durations": {
"min_seconds": min(durations) if durations else 0,
"max_seconds": max(durations) if durations else 0,
"mean_seconds": sum(durations) / len(durations) if durations else 0,
"median_seconds": sorted(durations)[len(durations) // 2]
if durations
else 0,
},
}
def write_json_summary(summary: dict[str, Any], output_file: Path) -> None:
"""Write a JSON summary to a file.
Args:
summary: Summary dictionary to write
output_file: Path to output file
"""
output_file.parent.mkdir(parents=True, exist_ok=True)
with open(output_file, "w") as f:
json.dump(summary, f, indent=2, default=str)
print(f"\nJSON summary written to: {output_file}")

643
tests/bench/obex.py Normal file
View File

@ -0,0 +1,643 @@
"""
Deriver Observation Extraction Benchmark.
Evaluates the deriver's ability to extract observations from messages by comparing
extracted observations against ground truth using embedding similarity.
Test data sources (915 total cases):
- LoCoMo: 543 cases with human-curated observations
- LongMem: 170 cases with Sonnet-generated observations
- BEAM: 202 cases with Sonnet-generated observations
Usage:
uv run python tests/bench/obex.py # Run full evaluation
uv run python tests/bench/obex.py --limit 10 # Test with limited cases
uv run python tests/bench/obex.py --source locomo # Evaluate specific source
uv run python tests/bench/obex.py --threshold 0.80 # Custom similarity threshold
"""
import argparse
import asyncio
import json
import sys
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any
import numpy as np
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
from src.config import settings
from src.deriver.prompts import minimal_deriver_prompt
from src.embedding_client import EmbeddingClient
from src.utils.clients import honcho_llm_call
from src.utils.representation import PromptRepresentation
CANDIDATES_DIR = Path(__file__).parent / "obexeval_data" / "candidates"
EVAL_RESULTS_DIR = Path(__file__).parent / "eval_results"
# =============================================================================
# Data Classes
# =============================================================================
def f1_score(precision: float, recall: float) -> float:
"""Calculate F1 score from precision and recall."""
if precision + recall == 0:
return 0.0
return 2 * (precision * recall) / (precision + recall)
@dataclass
class CaseResult:
"""Result for a single test case."""
case_id: str
source: str
difficulty: str
explicit_precision: float
explicit_recall: float
# deductive_precision: float
# deductive_recall: float
num_extracted: int
num_expected: int
# Actual model output
extracted_explicit: list[str] = field(default_factory=list)
# extracted_deductive: list[str] = field(default_factory=list)
error: str | None = None
@property
def explicit_f1(self) -> float:
return f1_score(self.explicit_precision, self.explicit_recall)
# @property
# def deductive_f1(self) -> float:
# return f1_score(self.deductive_precision, self.deductive_recall)
@dataclass
class SourceStats:
"""Statistics for a source or difficulty grouping."""
count: int
explicit_precision: float
explicit_recall: float
explicit_f1: float
# deductive_precision: float
# deductive_recall: float
# deductive_f1: float
@dataclass
class AggregateResults:
"""Aggregate results across all cases."""
total_cases: int
explicit_precision: float
explicit_recall: float
explicit_f1: float
# deductive_precision: float
# deductive_recall: float
# deductive_f1: float
by_source: dict[str, SourceStats] = field(default_factory=dict)
by_difficulty: dict[str, SourceStats] = field(default_factory=dict)
# =============================================================================
# Data Loading
# =============================================================================
TestCase = dict[str, Any]
def load_test_cases(source: str | None = None) -> list[TestCase]:
"""Load test cases from candidate files."""
all_cases: list[TestCase] = []
for json_file in CANDIDATES_DIR.glob("*_candidates.json"):
with open(json_file) as f:
data: dict[str, Any] = json.load(f)
all_cases.extend(data.get("candidates", []))
if source:
all_cases = [tc for tc in all_cases if tc.get("source") == source]
return all_cases
# =============================================================================
# Deriver Execution
# =============================================================================
def format_messages(messages: list[dict[str, Any]]) -> str:
"""Format messages for the deriver prompt."""
formatted: list[str] = []
for msg in messages:
author: str = msg.get("author", "unknown")
content: str = msg.get("content", "")
timestamp: str = msg.get("timestamp", "")
if timestamp:
try:
dt = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
time_str = dt.strftime("%Y-%m-%d %H:%M")
formatted.append(f"[{time_str}] {author}: {content}")
except ValueError:
formatted.append(f"{author}: {content}")
else:
formatted.append(f"{author}: {content}")
return "\n".join(formatted)
async def run_deriver(
messages: list[dict[str, Any]], target_peer: str
) -> PromptRepresentation:
"""Run the deriver on messages and return extracted observations."""
formatted = format_messages(messages)
prompt = minimal_deriver_prompt(
peer_id=target_peer,
messages=formatted,
)
response = await honcho_llm_call(
llm_settings=settings.DERIVER,
prompt=prompt,
max_tokens=settings.DERIVER.MAX_OUTPUT_TOKENS or 2000,
track_name="Deriver Eval",
response_model=PromptRepresentation,
json_mode=True,
stop_seqs=[" \n", "\n\n\n\n"],
enable_retry=True,
retry_attempts=3,
)
return response.content
# =============================================================================
# Embedding-based Matching
# =============================================================================
class ObservationMatcher:
"""Matches observations using embedding similarity."""
threshold: float
client: EmbeddingClient
def __init__(self, threshold: float = 0.85):
self.threshold = threshold
self.client = EmbeddingClient()
self._cache: dict[str, list[float]] = {}
async def _get_embeddings(self, texts: list[str]) -> list[list[float]]:
"""Get embeddings with caching."""
uncached = [t for t in texts if t not in self._cache]
if uncached:
embeddings = await self.client.simple_batch_embed(uncached)
for text, emb in zip(uncached, embeddings, strict=True):
self._cache[text] = emb
return [self._cache[t] for t in texts]
def _cosine_sim(self, a: list[float], b: list[float]) -> float:
"""Compute cosine similarity."""
a_np, b_np = np.array(a), np.array(b)
return float(np.dot(a_np, b_np) / (np.linalg.norm(a_np) * np.linalg.norm(b_np)))
async def match(
self, extracted: list[str], expected: list[str]
) -> tuple[float, float]:
"""
Match extracted to expected using greedy similarity matching.
Returns: (precision, recall)
"""
if not extracted and not expected:
return 1.0, 1.0
if not extracted:
return 0.0, 0.0
if not expected:
return 0.0, 1.0
# Get embeddings
all_texts = extracted + expected
all_embs = await self._get_embeddings(all_texts)
ext_embs = all_embs[: len(extracted)]
exp_embs = all_embs[len(extracted) :]
# Build similarity matrix
sim_matrix = np.zeros((len(extracted), len(expected)))
for i, e1 in enumerate(ext_embs):
for j, e2 in enumerate(exp_embs):
sim_matrix[i, j] = self._cosine_sim(e1, e2)
# Greedy matching
matches = 0
used_ext: set[int] = set()
used_exp: set[int] = set()
pairs = sorted(
[
(sim_matrix[i, j], i, j)
for i in range(len(extracted))
for j in range(len(expected))
],
reverse=True,
)
for sim, i, j in pairs:
if sim < self.threshold:
break
if i in used_ext or j in used_exp:
continue
matches += 1
used_ext.add(i)
used_exp.add(j)
precision = matches / len(extracted)
recall = matches / len(expected)
return precision, recall
# =============================================================================
# Evaluation
# =============================================================================
async def evaluate_case(case: TestCase, matcher: ObservationMatcher) -> CaseResult:
"""Evaluate a single test case."""
case_id: str = case["id"]
source: str = case.get("source", "unknown")
difficulty: str = case.get("difficulty", "unknown")
messages: list[dict[str, Any]] = case.get("messages", [])
target_peer: str = case.get("target_peer", "user")
expected: dict[str, Any] = case.get("expected_observations", {})
exp_explicit: list[str] = [o["content"] for o in expected.get("explicit", [])]
# exp_deductive: list[str] = [o["conclusion"] for o in expected.get("deductive", [])]
ext_explicit: list[str]
# ext_deductive: list[str]
error: str | None
try:
result = await run_deriver(messages, target_peer)
ext_explicit = [o.content for o in result.explicit]
# ext_deductive = [o.conclusion for o in result.deductive]
error = None
except Exception as e:
ext_explicit = [] # , ext_deductive = [], []
error = str(e)
# Match explicit
exp_p, exp_r = await matcher.match(ext_explicit, exp_explicit)
# Match deductive
# ded_p, ded_r = await matcher.match(ext_deductive, exp_deductive)
return CaseResult(
case_id=case_id,
source=source,
difficulty=difficulty,
explicit_precision=exp_p,
explicit_recall=exp_r,
# deductive_precision=ded_p,
# deductive_recall=ded_r,
num_extracted=len(ext_explicit), # + len(ext_deductive),
num_expected=len(exp_explicit), # + len(exp_deductive),
extracted_explicit=ext_explicit,
# extracted_deductive=ext_deductive,
error=error,
)
def aggregate_results(results: list[CaseResult]) -> AggregateResults:
"""Compute aggregate statistics."""
if not results:
return AggregateResults(0, 0, 0, 0) # , 0, 0, 0)
n = len(results)
exp_p = sum(r.explicit_precision for r in results) / n
exp_r = sum(r.explicit_recall for r in results) / n
# ded_p = sum(r.deductive_precision for r in results) / n
# ded_r = sum(r.deductive_recall for r in results) / n
agg = AggregateResults(
total_cases=n,
explicit_precision=exp_p,
explicit_recall=exp_r,
explicit_f1=f1_score(exp_p, exp_r),
# deductive_precision=ded_p,
# deductive_recall=ded_r,
# deductive_f1=f1_score(ded_p, ded_r),
)
# By source
by_source: dict[str, list[CaseResult]] = {}
for r in results:
by_source.setdefault(r.source, []).append(r)
for source, rs in by_source.items():
src_exp_p = sum(r.explicit_precision for r in rs) / len(rs)
src_exp_r = sum(r.explicit_recall for r in rs) / len(rs)
# src_ded_p = sum(r.deductive_precision for r in rs) / len(rs)
# src_ded_r = sum(r.deductive_recall for r in rs) / len(rs)
agg.by_source[source] = SourceStats(
count=len(rs),
explicit_precision=src_exp_p,
explicit_recall=src_exp_r,
explicit_f1=f1_score(src_exp_p, src_exp_r),
# deductive_precision=src_ded_p,
# deductive_recall=src_ded_r,
# deductive_f1=f1_score(src_ded_p, src_ded_r),
)
# By difficulty
by_diff: dict[str, list[CaseResult]] = {}
for r in results:
by_diff.setdefault(r.difficulty, []).append(r)
for diff, rs in by_diff.items():
diff_exp_p = sum(r.explicit_precision for r in rs) / len(rs)
diff_exp_r = sum(r.explicit_recall for r in rs) / len(rs)
# diff_ded_p = sum(r.deductive_precision for r in rs) / len(rs)
# diff_ded_r = sum(r.deductive_recall for r in rs) / len(rs)
agg.by_difficulty[diff] = SourceStats(
count=len(rs),
explicit_precision=diff_exp_p,
explicit_recall=diff_exp_r,
explicit_f1=f1_score(diff_exp_p, diff_exp_r),
# deductive_precision=diff_ded_p,
# deductive_recall=diff_ded_r,
# deductive_f1=f1_score(diff_ded_p, diff_ded_r),
)
return agg
def print_results(agg: AggregateResults, results: list[CaseResult]) -> None:
"""Print evaluation results."""
print("\n" + "=" * 70)
print("DERIVER OBSERVATION EXTRACTION EVALUATION")
print("=" * 70)
print(f"\nTotal test cases: {agg.total_cases}")
print(f"\n{'EXPLICIT OBSERVATIONS':^35}")
print(f" Precision: {agg.explicit_precision:.3f}")
print(f" Recall: {agg.explicit_recall:.3f}")
print(f" F1: {agg.explicit_f1:.3f}")
# print(f"\n{'DEDUCTIVE OBSERVATIONS':^35}")
# print(f" Precision: {agg.deductive_precision:.3f}")
# print(f" Recall: {agg.deductive_recall:.3f}")
# print(f" F1: {agg.deductive_f1:.3f}")
if agg.by_source:
print(f"\n{'BY SOURCE':^35}")
for source, data in sorted(agg.by_source.items()):
print(
f" {source:12} n={data.count:3} exp_f1={data.explicit_f1:.3f}" # ded_f1={data.deductive_f1:.3f}"
)
if agg.by_difficulty:
print(f"\n{'BY DIFFICULTY':^35}")
for diff, data in sorted(agg.by_difficulty.items()):
print(
f" {diff:12} n={data.count:3} exp_f1={data.explicit_f1:.3f}" # ded_f1={data.deductive_f1:.3f}"
)
# Worst cases (by average of precision and recall)
worst = sorted(
results,
key=lambda r: (
r.explicit_precision + r.explicit_recall
# + r.deductive_precision
# + r.deductive_recall
)
/ 2, # / 4,
)[:5]
print(f"\n{'WORST PERFORMING CASES':^35}")
for r in worst:
avg = (
(
r.explicit_precision + r.explicit_recall
# + r.deductive_precision
# + r.deductive_recall
)
/ 2,
) # / 4,
print(f" {r.case_id[:45]:45} avg={avg:.3f}")
# Errors
errors = [r for r in results if r.error]
if errors:
print(f"\nErrors: {len(errors)} cases failed")
def generate_json_summary(
results: list[CaseResult],
agg: AggregateResults,
source_filter: str | None,
threshold: float,
concurrency: int,
total_elapsed_seconds: float,
output_file: Path,
) -> None:
"""
Generate a comprehensive JSON summary of test results for analytics.
Args:
results: List of case results
agg: Aggregate results
source_filter: Source filter used (if any)
threshold: Similarity threshold used
concurrency: Concurrency level used
total_elapsed_seconds: Total elapsed time for all tests
output_file: Path to write JSON output to
"""
errors = [r for r in results if r.error]
# Calculate per-case average scores
case_scores: list[float] = []
for r in results:
avg = (r.explicit_precision + r.explicit_recall) / 2
# + r.deductive_precision
# + r.deductive_recall
# / 4,
case_scores.append(avg)
output_data: dict[str, Any] = {
"metadata": {
"benchmark": "obex",
"description": "Observation Extraction Benchmark - evaluates deriver's ability to extract observations from messages",
"execution_timestamp": datetime.now().isoformat(),
"runner_version": "1.0.0",
"deriver_settings": settings.DERIVER.model_dump(),
},
"config": {
"source_filter": source_filter,
"threshold": threshold,
"concurrency": concurrency,
"total_cases": len(results),
},
"timing": {
"total_elapsed_seconds": total_elapsed_seconds,
"average_seconds_per_case": total_elapsed_seconds / len(results)
if results
else 0,
},
"summary_statistics": {
"total_cases": agg.total_cases,
"cases_with_errors": len(errors),
"explicit_precision": agg.explicit_precision,
"explicit_recall": agg.explicit_recall,
"explicit_f1": agg.explicit_f1,
# "deductive_precision": agg.deductive_precision,
# "deductive_recall": agg.deductive_recall,
# "deductive_f1": agg.deductive_f1,
"mean_case_score": sum(case_scores) / len(case_scores)
if case_scores
else 0,
"min_case_score": min(case_scores) if case_scores else 0,
"max_case_score": max(case_scores) if case_scores else 0,
},
"statistics_by_source": {k: asdict(v) for k, v in agg.by_source.items()},
"statistics_by_difficulty": {
k: asdict(v) for k, v in agg.by_difficulty.items()
},
"detailed_results": [
{
"case_id": r.case_id,
"source": r.source,
"difficulty": r.difficulty,
"explicit_precision": r.explicit_precision,
"explicit_recall": r.explicit_recall,
"explicit_f1": r.explicit_f1,
# "deductive_precision": r.deductive_precision,
# "deductive_recall": r.deductive_recall,
# "deductive_f1": r.deductive_f1,
"num_extracted": r.num_extracted,
"num_expected": r.num_expected,
"average_f1": r.explicit_f1, # (r.explicit_f1 + r.deductive_f1) / 2,
"model_output": {
"explicit": r.extracted_explicit,
# "deductive": r.extracted_deductive,
},
"error": r.error,
}
for r in results
],
}
# Ensure output directory exists
output_file.parent.mkdir(parents=True, exist_ok=True)
with open(output_file, "w") as f:
json.dump(output_data, f, indent=2)
print(f"\nResults saved to {output_file}")
async def run_evaluation(
source: str | None = None,
limit: int | None = None,
threshold: float = 0.85,
output: str | None = None,
concurrency: int = 10,
) -> None:
"""Run the evaluation pipeline."""
start_time = time.time()
print("Loading test cases...")
cases = load_test_cases(source)
if limit:
cases = cases[:limit]
if not cases:
print("No test cases found!")
return
print(
f"Evaluating {len(cases)} cases (model={settings.DERIVER.MODEL}, threshold={threshold}, concurrency={concurrency})..."
)
matcher = ObservationMatcher(threshold=threshold)
completed = 0
semaphore = asyncio.Semaphore(concurrency)
async def evaluate_with_semaphore(case: TestCase) -> CaseResult:
nonlocal completed
async with semaphore:
result = await evaluate_case(case, matcher)
completed += 1
if completed % 10 == 0 or completed == 1:
print(f" Completed {completed}/{len(cases)}...")
return result
# Run all evaluations concurrently with semaphore limiting
tasks = [evaluate_with_semaphore(case) for case in cases]
results = await asyncio.gather(*tasks)
end_time = time.time()
total_elapsed = end_time - start_time
agg = aggregate_results(list(results))
print_results(agg, list(results))
print(
f"\nTotal time: {total_elapsed:.1f}s ({total_elapsed / len(cases):.2f}s per case)"
)
# Determine output file path
if output:
output_file = Path(output)
else:
# Generate default timestamped filename
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
source_suffix = f"_{source}" if source else ""
output_file = EVAL_RESULTS_DIR / f"obex{source_suffix}_{timestamp}.json"
# Always save results
generate_json_summary(
list(results),
agg,
source,
threshold,
concurrency,
total_elapsed,
output_file,
)
def main() -> None:
parser = argparse.ArgumentParser(
description="Deriver observation extraction benchmark"
)
parser.add_argument(
"--source", type=str, help="Filter by source (locomo, longmem, beam)"
)
parser.add_argument("--limit", type=int, help="Limit number of test cases")
parser.add_argument(
"--threshold",
type=float,
default=0.85,
help="Similarity threshold (default: 0.85)",
)
parser.add_argument("--output", type=str, help="Output path for results JSON")
parser.add_argument(
"--concurrency",
type=int,
default=10,
help="Number of concurrent LLM calls (default: 10)",
)
args = parser.parse_args()
asyncio.run(
run_evaluation(
args.source, args.limit, args.threshold, args.output, args.concurrency
)
)
if __name__ == "__main__":
main()

454
tests/bench/peer_card_bench.py
View File

@ -1,454 +0,0 @@
"""
Peer Card Benchmark
This benchmark exercises the peer card LLM call with varied inputs and compares
outputs across multiple provider/model candidates. Results are graded by an LLM
judge.
Usage example:
python -m tests.bench.peer_card_bench --candidates anthropic:claude-3-7-sonnet-20250219 --candidates openai:gpt-4o-mini-2024-07-18
Environment variables for providers:
- Anthropic: LLM_ANTHROPIC_API_KEY
- OpenAI: LLM_OPENAI_API_KEY or OPENAI_API_KEY
- Google (Gemini): LLM_GEMINI_API_KEY or GEMINI_API_KEY
- Groq: LLM_GROQ_API_KEY or GROQ_API_KEY
"""
import argparse
import asyncio
import json
import os
import time
from collections.abc import Callable, Coroutine
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, cast
from anthropic import AsyncAnthropic
from src.config import settings
from src.deriver.prompts import peer_card_prompt
from src.utils.clients import honcho_llm_call
from src.utils.peer_card import PeerCardQuery
from src.utils.representation import ExplicitObservation, Representation
COLOR_GREEN = "\033[32m"
COLOR_RED = "\033[31m"
COLOR_RESET = "\033[0m"
@dataclass(frozen=True)
class Candidate:
"""Represents a provider/model pair to benchmark."""
provider: str
model: str
@dataclass
class Case:
"""Represents a single benchmark case with expectations for grading.
Attributes:
name: Human-friendly identifier for the case.
old_peer_card: Existing card text to update, or None to create fresh.
new_observations: New input observations that may change the card.
expected_facts: Facts that must be semantically present in the result.
forbidden_facts: Facts that must NOT be present in the result.
"""
name: str
old_peer_card: list[str] | None
new_observations: list[str]
expected_facts: list[str]
forbidden_facts: list[str]
def load_case_file(path: Path) -> Case:
"""Load a single peer-card test case from a JSON file.
The JSON schema must include: name, old_peer_card (nullable), new_observations (list[str]), expected_facts (list[str]).
"""
with path.open() as f:
data = json.load(f)
return Case(
name=str(data["name"]),
old_peer_card=data.get("old_peer_card"),
new_observations=list(data.get("new_observations", [])),
expected_facts=list(data.get("expected_facts", [])),
forbidden_facts=list(data.get("forbidden_facts", [])),
)
def load_cases(tests_dir: Path, test_name: str | None) -> list[Case]:
"""Load all cases from a directory, or a specific case by filename.
Args:
tests_dir: Directory containing JSON case files.
test_name: Optional filename to load a single case (e.g., "create_basic_card.json").
Returns:
List of loaded Case objects.
"""
if test_name:
file_path = tests_dir / test_name
if not file_path.exists():
raise FileNotFoundError(f"Test file {file_path} does not exist")
return [load_case_file(file_path)]
files = sorted(p for p in tests_dir.glob("*.json") if p.is_file())
return [load_case_file(p) for p in files]
def parse_candidates(values: list[str]) -> list[Candidate]:
"""Parse provider:model strings into Candidate objects."""
result: list[Candidate] = []
for v in values:
v = v.strip()
if not v:
continue
if ":" not in v:
raise ValueError(f"Invalid candidate format: {v} (expected provider:model)")
provider, model = v.split(":", 1)
result.append(Candidate(provider=provider.strip(), model=model.strip()))
return result
def deduplicate_preserve_order(items: list[Candidate]) -> list[Candidate]:
"""Return a new list with duplicate provider:model pairs removed, preserving order."""
seen: set[tuple[str, str]] = set()
unique: list[Candidate] = []
for item in items:
key = (item.provider, item.model)
if key in seen:
continue
seen.add(key)
unique.append(item)
return unique
def build_peer_card_caller(
candidate: Candidate,
) -> Callable[[list[str] | None, Representation], Coroutine[Any, Any, PeerCardQuery]]:
"""Create an async callable that invokes the peer card prompt with a specific provider/model."""
resolved_provider = (
"openai" if candidate.provider == "custom" else candidate.provider
)
settings.PEER_CARD.PROVIDER = cast(Any, resolved_provider)
settings.PEER_CARD.MODEL = candidate.model
async def call(
old_peer_card: list[str] | None, new_observations: Representation
) -> PeerCardQuery:
prompt = peer_card_prompt(
old_peer_card=old_peer_card,
new_observations=new_observations.str_no_timestamps(),
)
response = await honcho_llm_call(
llm_settings=settings.PEER_CARD,
prompt=prompt,
max_tokens=settings.PEER_CARD.MAX_OUTPUT_TOKENS,
response_model=PeerCardQuery,
json_mode=True,
reasoning_effort="minimal",
enable_retry=True,
retry_attempts=3,
)
return response.content
return call
def _extract_json_from_text(text: str) -> dict[str, Any]:
"""Extract a JSON object from potentially noisy LLM output.
Strategy in order:
- Try to parse the whole text as JSON
- Try to parse contents of any fenced code blocks (``` or ```json)
- Try the substring from the first '{' to the last '}'
- Scan for balanced-brace substrings and try them in order
Raises ValueError when no valid JSON object can be found.
"""
stripped: str = text.strip()
candidates: list[str] = []
# 1) Fenced code blocks
if "```" in stripped:
idx: int = 0
while True:
start = stripped.find("```", idx)
if start == -1:
break
lang_line_end = stripped.find("\n", start + 3)
if lang_line_end == -1:
break
end = stripped.find("```", lang_line_end + 1)
if end == -1:
break
block = stripped[lang_line_end + 1 : end].strip()
if block:
candidates.append(block)
idx = end + 3
# 2) From first '{' to last '}'
first_brace = stripped.find("{")
last_brace = stripped.rfind("}")
if first_brace != -1 and last_brace != -1 and last_brace > first_brace:
candidates.append(stripped[first_brace : last_brace + 1])
# 3) Balanced-brace scan
depth = 0
start_idx = -1
for i, ch in enumerate(stripped):
if ch == "{":
if depth == 0:
start_idx = i
depth += 1
elif ch == "}":
if depth > 0:
depth -= 1
if depth == 0 and start_idx != -1:
candidates.append(stripped[start_idx : i + 1])
# Try all candidates, prefer ones containing the expected keys
preferred_keys = {"passed", "reasoning"}
fallback_obj: dict[str, Any] | None = None
for cand in candidates:
try:
obj_candidate: object = json.loads(cand)
if isinstance(obj_candidate, dict):
casted_obj: dict[str, Any] = {
str(k): v # pyright: ignore
for k, v in obj_candidate.items() # pyright: ignore
}
if preferred_keys.issubset(set(casted_obj.keys())):
return casted_obj
if fallback_obj is None:
fallback_obj = casted_obj
except Exception:
continue
if fallback_obj is not None:
return fallback_obj
raise ValueError("Could not extract JSON from judge response")
async def judge_response(
anthropic: AsyncAnthropic,
case: Case,
actual_card: list[str],
) -> dict[str, Any]:
"""Use an LLM judge to evaluate whether the card contains the expected facts.
Returns a dict with keys: passed (bool) and reasoning (str).
"""
system_prompt = (
"You are an expert evaluator. Determine if a biographical card satisfies BOTH: "
"(1) it contains all expected facts (semantic match allowed) and "
"(2) it does NOT contain any forbidden facts (semantic match). "
"Allow flexible phrasing and synonyms for matching. A fact is present if its semantic content is clearly stated. "
"Fail if any expected fact is missing or any forbidden fact appears. Always return JSON: "
'{"passed": boolean, "reasoning": string}'
)
expected = "\n".join(f"- {f}" for f in case.expected_facts)
forbidden = "\n".join(f"- {f}" for f in case.forbidden_facts)
card_text = "\n".join(actual_card) if actual_card else "- (none)"
user_prompt = (
f"Case: {case.name}\n\n"
f"Expected facts (must appear, semantic):\n{expected or '- (none)'}\n\n"
f"Forbidden facts (must NOT appear, semantic):\n{forbidden or '- (none)'}\n\n"
f"Biographical card to evaluate:\n{card_text}\n\n"
f"Evaluation criteria: PASS only if all expected facts are present AND all forbidden facts are absent."
)
judgment_text: str | None = None
try:
response = await anthropic.messages.create(
model="claude-sonnet-4-20250514",
max_tokens=1000,
temperature=0.0,
system=system_prompt,
messages=[{"role": "user", "content": user_prompt}],
)
content_block = response.content[0]
judgment_text = getattr(content_block, "text", None)
if not judgment_text:
raise ValueError("Empty judge response")
return _extract_json_from_text(judgment_text)
except Exception as e:
print(judgment_text)
raise ValueError(f"!!!Error judging response for case {case.name}: {e}") from e
async def run_benchmark(candidates: list[Candidate], cases: list[Case]) -> int:
"""Execute cases against candidates and print a concise report.
Returns non-zero when any case fails for any candidate.
"""
anthropic_key = os.getenv("LLM_ANTHROPIC_API_KEY")
if not anthropic_key:
raise ValueError("LLM_ANTHROPIC_API_KEY is required for grading")
anthropic = AsyncAnthropic(api_key=anthropic_key)
any_fail = False
print(f"Running {len(cases)} cases across {len(candidates)} candidates\n")
for candidate in candidates:
print(f"=== Candidate: {candidate.provider}:{candidate.model} ===")
candidate_start_time = time.perf_counter()
try:
caller = build_peer_card_caller(candidate)
except Exception as e:
print(f" SKIP: cannot initialize provider/model ({e})")
any_fail = True
continue
async def run_case(
case: Case,
_caller: Callable[
[list[str] | None, Representation], Coroutine[Any, Any, PeerCardQuery]
] = caller,
) -> tuple[Case, dict[str, Any]]:
card: PeerCardQuery = await _caller(
case.old_peer_card,
Representation(
explicit=[
ExplicitObservation(
content=o,
created_at=datetime.now(timezone.utc),
message_ids=[0],
session_name=case.name,
)
for o in case.new_observations
]
),
)
new_card = card.card
if new_card is None or new_card == []:
new_card = case.old_peer_card or []
judgment = await judge_response(anthropic, case, new_card)
return case, {"card": card, "judgment": judgment}
results = await asyncio.gather(
*(run_case(c) for c in cases), return_exceptions=True
)
passed_count: int = 0
for res in results:
if isinstance(res, BaseException):
print(f" ERROR running case: {res}")
any_fail = True
continue
case, payload = res
judgment = payload["judgment"]
passed = bool(judgment.get("passed"))
status_colored = (
f"{COLOR_GREEN}PASS{COLOR_RESET}"
if passed
else f"{COLOR_RED}FAIL{COLOR_RESET}"
)
if passed:
print(f" {case.name:24} {status_colored}")
passed_count += 1
else:
print(
f" {case.name:24} {status_colored} - {judgment.get('reasoning', '')}"
)
any_fail = True
print(" expected:")
for f in case.expected_facts:
print(f" - {f}")
if case.forbidden_facts:
print(" forbidden (must NOT appear):")
for f in case.forbidden_facts:
print(f" - {f}")
print(" got:")
[print(" " + line) for line in payload["card"].card]
print(" with 'notes' field:")
if payload["card"].notes:
print(" " + payload["card"].notes)
else:
print(" - (none)")
total_count: int = len(cases)
percentage: float = (passed_count / total_count * 100.0) if total_count else 0.0
print(f" Summary: {passed_count}/{total_count} passed ({percentage:.1f}%)")
elapsed = time.perf_counter() - candidate_start_time
print(f" Time: {elapsed:.2f}s\n")
print("Done.")
return 1 if any_fail else 0
def main() -> int:
"""CLI entry point for running the peer card benchmark."""
parser = argparse.ArgumentParser(
description="Benchmark peer card LLM behavior across models"
)
parser.add_argument(
"--candidates",
action="append",
default=None,
help=(
"Provider:model pairs. Repeat or comma-separate. "
"Default: anthropic:claude-3-7-sonnet-20250219"
),
)
parser.add_argument(
"--tests-dir",
type=Path,
default=Path("tests/bench/peer_card_tests"),
help=(
"Directory containing JSON peer-card cases "
"(default: tests/bench/peer_card_tests)"
),
)
parser.add_argument(
"--test",
type=str,
help="Run a specific test file by name (e.g., 'create_basic_card.json')",
)
args = parser.parse_args()
# Use the default candidate only when the flag is not provided at all
candidate_entries: list[str] = (
args.candidates
if args.candidates is not None
else ["anthropic:claude-3-7-sonnet-20250219"]
)
flat: list[str] = []
for entry in candidate_entries:
flat.extend([s.strip() for s in entry.split(",") if s.strip()])
candidates = parse_candidates(flat)
candidates = deduplicate_preserve_order(candidates)
# Load cases from JSON files
if not args.tests_dir.exists():
raise SystemExit(f"Error: Tests directory {args.tests_dir} does not exist")
cases = load_cases(args.tests_dir, args.test)
if not cases:
raise SystemExit(f"Error: No JSON test cases found in {args.tests_dir}")
return asyncio.run(run_benchmark(candidates, cases))
if __name__ == "__main__":
raise SystemExit(main())

2
tests/bench/run_tests.py
View File

@ -552,7 +552,7 @@ Evaluate whether the actual response contains the core correct information from
if session_context.summary:
summary_content = session_context.summary.content
tokenizer = tiktoken.get_encoding("cl100k_base")
tokenizer = tiktoken.get_encoding("o200k_base")
summary_tokens = len(tokenizer.encode(summary_content))
output_lines.append(f" summary: {session_context.summary}")

114
tests/conftest.py
View File

@ -352,6 +352,26 @@ def mock_langfuse():
logging.getLogger().removeHandler(handler)
def _content_to_embedding(content: str) -> list[float]:
"""Generate a deterministic embedding from content hash.
This ensures different content produces different embeddings,
which is critical for deduplication logic to work correctly in tests.
"""
import hashlib
# Hash the content to get a deterministic seed
content_hash = hashlib.sha256(content.encode()).digest()
# Use hash bytes to generate 1536 floats between -1 and 1
embedding: list[float] = []
for i in range(1536):
# Use different bytes from hash (cycling through)
byte_val = content_hash[i % len(content_hash)]
# Normalize to [-1, 1] range
embedding.append((byte_val / 255.0) * 2 - 1)
return embedding
@pytest.fixture(autouse=True)
def mock_openai_embeddings():
"""Mock OpenAI embeddings API calls for testing"""
@ -359,17 +379,20 @@ def mock_openai_embeddings():
patch("src.embedding_client.embedding_client.embed") as mock_embed,
patch("src.embedding_client.embedding_client.batch_embed") as mock_batch_embed,
):
# Mock the embed method to return a fake embedding vector
mock_embed.return_value = [0.1] * 1536
# Mock the embed method to return content-dependent embedding
def embed_side_effect(content: str) -> list[float]:
return _content_to_embedding(content)
# Mock the batch_embed method to return a dict of fake embedding vectors
# Updated to support chunking - each text_id maps to a list of embedding vectors
mock_embed.side_effect = embed_side_effect
# Mock the batch_embed method to return content-dependent embeddings
async def mock_batch_embed_func(
id_resource_dict: dict[str, tuple[str, list[int]]],
) -> dict[str, list[list[float]]]:
return {
text_id: [[0.1] * 1536] for text_id in id_resource_dict
} # Single chunk per text
text_id: [_content_to_embedding(resource[0])]
for text_id, resource in id_resource_dict.items()
}
mock_batch_embed.side_effect = mock_batch_embed_func
@ -381,6 +404,8 @@ def mock_llm_call_functions():
"""Mock LLM functions to avoid needing API keys during tests"""
# Create mock responses for different function types
# Note: critical_analysis_call was removed as the deriver now uses agentic approach
# Note: dialectic_call/dialectic_stream were replaced with agentic_chat
with (
patch(
"src.utils.summarizer.create_short_summary", new_callable=AsyncMock
@ -389,64 +414,30 @@ def mock_llm_call_functions():
"src.utils.summarizer.create_long_summary", new_callable=AsyncMock
) as mock_long_summary,
patch(
"src.deriver.deriver.critical_analysis_call", new_callable=AsyncMock
) as mock_critical_analysis,
patch(
"src.dialectic.chat.dialectic_call", new_callable=AsyncMock
) as mock_dialectic_call,
patch(
"src.dialectic.chat.dialectic_stream", new_callable=AsyncMock
) as mock_dialectic_stream,
"src.routers.peers.agentic_chat", new_callable=AsyncMock
) as mock_agentic_chat,
):
# Import the required models for proper mocking
from src.utils.representation import (
DeductiveObservationBase,
ExplicitObservationBase,
PromptRepresentation,
)
# Mock return values for different function types
mock_short_summary.return_value = "Test short summary content"
mock_long_summary.return_value = "Test long summary content"
# Mock critical_analysis_call to return a proper object with _response attribute
_rep = PromptRepresentation(
explicit=[ExplicitObservationBase(content="Test explicit observation")],
deductive=[
DeductiveObservationBase(
conclusion="Test deductive conclusion",
premises=["Test premise 1", "Test premise 2"],
)
],
)
mock_critical_analysis_result = MagicMock(wraps=_rep)
# Add the _response attribute that contains thinking (used in the actual code)
mock_response = MagicMock()
mock_response.thinking = "Test thinking content"
mock_critical_analysis_result._response = mock_response
mock_critical_analysis.return_value = mock_critical_analysis_result
# Mock dialectic_call to return a string (matching actual return type)
mock_dialectic_call.return_value = "Test dialectic response"
mock_dialectic_stream.return_value = AsyncMock()
# Mock agentic_chat to return a string (matching actual return type)
mock_agentic_chat.return_value = "Test dialectic response"
yield {
"short_summary": mock_short_summary,
"long_summary": mock_long_summary,
"critical_analysis": mock_critical_analysis,
"dialectic_call": mock_dialectic_call,
"dialectic_stream": mock_dialectic_stream,
"agentic_chat": mock_agentic_chat,
}
@pytest.fixture(autouse=True)
def mock_honcho_llm_call():
"""Generic mock for the honcho_llm_call decorator to avoid actual LLM calls during tests"""
from unittest.mock import AsyncMock, MagicMock
from unittest.mock import AsyncMock
from src.utils.representation import (
DeductiveObservationBase,
# DeductiveObservationBase,
ExplicitObservationBase,
PromptRepresentation,
)
@ -469,12 +460,12 @@ def mock_honcho_llm_call():
explicit=[
ExplicitObservationBase(content="Test explicit observation")
],
deductive=[
DeductiveObservationBase(
conclusion="Test deductive conclusion",
premises=["Test premise 1", "Test premise 2"],
),
],
# deductive=[
# DeductiveObservationBase(
# conclusion="Test deductive conclusion",
# premises=["Test premise 1", "Test premise 2"],
# ),
# ],
)
mock_response = MagicMock(wraps=_rep)
# Add the _response attribute that contains thinking (used in the actual code)
@ -564,19 +555,32 @@ def mock_tracked_db(db_session: AsyncSession):
with (
patch("src.dependencies.tracked_db", mock_tracked_db_context),
patch("src.deriver.deriver.tracked_db", mock_tracked_db_context),
patch("src.deriver.queue_manager.tracked_db", mock_tracked_db_context),
patch("src.deriver.consumer.tracked_db", mock_tracked_db_context),
patch("src.deriver.enqueue.tracked_db", mock_tracked_db_context),
patch("src.routers.sessions.tracked_db", mock_tracked_db_context),
patch("src.routers.peers.tracked_db", mock_tracked_db_context),
patch("src.crud.representation.tracked_db", mock_tracked_db_context),
patch("src.dreamer.consolidate.tracked_db", mock_tracked_db_context),
patch("src.dreamer.dreamer.tracked_db", mock_tracked_db_context),
patch("src.dreamer.dream_scheduler.tracked_db", mock_tracked_db_context),
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),
):
yield
@pytest.fixture(autouse=True)
def enable_deriver_for_tests():
"""Enable deriver globally for tests that need queue processing"""
from src.config import settings
original_value = settings.DERIVER.ENABLED
settings.DERIVER.ENABLED = True
yield
settings.DERIVER.ENABLED = original_value
@pytest.fixture(autouse=True)
def mock_crud_collection_operations():
"""Mock CRUD operations that try to commit to database during tests"""

25
tests/deriver/conftest.py
View File

@ -1,8 +1,8 @@
import asyncio
from collections.abc import Awaitable, Callable, Generator, Sequence
from collections.abc import Awaitable, Callable, Sequence
from datetime import datetime, timezone
from typing import Any, Literal, TypeAlias, cast
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, MagicMock
import pytest
from nanoid import generate as generate_nanoid
@ -17,27 +17,6 @@ QueuePayload: TypeAlias = dict[str, Any]
QueuePayloadEntry: TypeAlias = QueuePayload | tuple[QueuePayload, int | None]
@pytest.fixture
def mock_critical_analysis_call() -> Generator[Callable[..., Any], None, None]:
"""Mock the critical analysis call to avoid actual LLM calls"""
async def mock_critical_analysis_call(*_args: Any, **_kwargs: Any) -> MagicMock:
# Create a mock response that matches the expected structure
mock_response = MagicMock()
mock_response.explicit = ["Test explicit observation"]
mock_response.deductive = []
mock_response.thinking = "Test thinking content"
mock_response._response = MagicMock()
mock_response._response.thinking = "Test thinking content"
return mock_response
# Patch the actual function in the deriver module
with patch(
"src.deriver.deriver.critical_analysis_call", mock_critical_analysis_call
):
yield mock_critical_analysis_call
@pytest.fixture
async def sample_session_with_peers(
db_session: AsyncSession,

139
tests/deriver/test_deriver_processing.py
View File

@ -1,14 +1,9 @@
import signal
from collections.abc import Callable, Generator
from datetime import datetime, timedelta, timezone
from typing import Any
from unittest.mock import AsyncMock
import pytest
from sqlalchemy.ext.asyncio import AsyncSession
from src import models
from src.deriver.deriver import process_representation_tasks_batch
from src.utils.representation import Representation
from src.utils.work_unit import construct_work_unit_key
@ -17,18 +12,6 @@ from src.utils.work_unit import construct_work_unit_key
class TestDeriverProcessing:
"""Test suite for deriver processing using the conftest fixtures"""
async def test_mock_critical_analysis_call(
self,
mock_critical_analysis_call: Generator[Callable[..., Any], None, None],
sample_messages: list[models.Message],
):
"""Test that the critical analysis call is properly mocked"""
assert mock_critical_analysis_call is not None
assert len(sample_messages) > 0 # Verify we have messages for testing
# The mock should be in place and return a predefined response
# This ensures no actual LLM calls are made during testing
async def test_work_unit_key_generation(
self,
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
@ -107,74 +90,74 @@ class TestDeriverProcessing:
# Verify the methods were called
assert mock_representation_manager.save_representation.called # type: ignore[attr-defined]
async def test_representation_batch_uses_earliest_cutoff(
self,
db_session: AsyncSession,
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Ensure batching history cutoff uses the earliest payload in the batch."""
captured_cutoffs: list[int] = []
# async def test_representation_batch_uses_earliest_cutoff(
# self,
# db_session: AsyncSession,
# sample_session_with_peers: tuple[models.Session, list[models.Peer]],
# monkeypatch: pytest.MonkeyPatch,
# ) -> None:
# """Ensure batching history cutoff uses the earliest payload in the batch."""
# captured_cutoffs: list[int] = []
async def fake_get_session_context_formatted(*_args: Any, **kwargs: Any) -> str:
captured_cutoffs.append(kwargs["cutoff"])
return "formatted-history"
# async def fake_get_session_context_formatted(*_args: Any, **kwargs: Any) -> str:
# captured_cutoffs.append(kwargs["cutoff"])
# return "formatted-history"
# Mock only the function we need to inspect for the test assertion
monkeypatch.setattr(
"src.deriver.deriver.summarizer.get_session_context_formatted",
fake_get_session_context_formatted,
)
# # Mock only the function we need to inspect for the test assertion
# monkeypatch.setattr(
# "src.deriver.deriver.summarizer.get_session_context_formatted",
# fake_get_session_context_formatted,
# )
# Provide a stub working representation so embedding lookups are skipped.
monkeypatch.setattr(
"src.crud.get_working_representation",
AsyncMock(
return_value=Representation(
explicit=[],
deductive=[],
)
),
)
# # Provide a stub working representation so embedding lookups are skipped.
# monkeypatch.setattr(
# "src.crud.get_working_representation",
# AsyncMock(
# return_value=Representation(
# explicit=[],
# deductive=[],
# )
# ),
# )
# Avoid executing the full reasoning pipeline; we only care about cutoff behavior.
monkeypatch.setattr(
"src.deriver.deriver.CertaintyReasoner.reason",
AsyncMock(return_value=Representation(explicit=[], deductive=[])),
)
# # Avoid executing the full reasoning pipeline; we only care about cutoff behavior.
# monkeypatch.setattr(
# "src.deriver.deriver.CertaintyReasoner.reason",
# AsyncMock(return_value=Representation(explicit=[], deductive=[])),
# )
# Use the real session and workspace from fixtures
session, peers = sample_session_with_peers
alice = peers[0]
# # Use the real session and workspace from fixtures
# session, peers = sample_session_with_peers
# alice = peers[0]
# Create test messages with different IDs in the database
now = datetime.now(timezone.utc)
messages: list[models.Message] = []
for i in range(8):
message = models.Message(
workspace_name=session.workspace_name,
session_name=session.name,
peer_name=alice.name,
content=f"message {i}",
seq_in_session=i + 1,
token_count=10,
created_at=now - timedelta(minutes=7 - i),
)
db_session.add(message)
messages.append(message)
# # Create test messages with different IDs in the database
# now = datetime.now(timezone.utc)
# messages: list[models.Message] = []
# for i in range(8):
# message = models.Message(
# workspace_name=session.workspace_name,
# session_name=session.name,
# peer_name=alice.name,
# content=f"message {i}",
# seq_in_session=i + 1,
# token_count=10,
# created_at=now - timedelta(minutes=7 - i),
# )
# db_session.add(message)
# messages.append(message)
await db_session.commit()
# await db_session.commit()
# Refresh messages to get their IDs
for message in messages:
await db_session.refresh(message)
# # Refresh messages to get their IDs
# for message in messages:
# await db_session.refresh(message)
await process_representation_tasks_batch(
observer=alice.name,
message_level_configuration=None,
observed=alice.name,
messages=messages,
)
# await process_representation_tasks_batch(
# observer=alice.name,
# message_level_configuration=None,
# observed=alice.name,
# messages=messages,
# )
# Verify that the earliest message ID was used as the cutoff
assert captured_cutoffs == [messages[0].id]
# # Verify that the earliest message ID was used as the cutoff
# assert captured_cutoffs == [messages[0].id]

42
tests/deriver/test_queue_processing.py
View File

@ -494,7 +494,7 @@ class TestQueueProcessing:
# Ensure items are only for alice
assert all(qi.payload.get("observed") == alice.name for qi in alice_items)
# Test bob's work unit - starts at message 2 for per-work-unit anchoring
# Test bob's work unit - now includes preceding message for context
bob_work_unit_key = bob_queue_items[0].work_unit_key
bob_aqs = models.ActiveQueueSession(work_unit_key=bob_work_unit_key)
db_session.add(bob_aqs)
@ -507,10 +507,11 @@ class TestQueueProcessing:
aqs_id=bob_aqs.id,
)
# Bob should get 4 messages (2..5)
assert len(bob_messages) == 4
# Bob should get 5 messages (1..5) - includes preceding alice message for context
assert len(bob_messages) == 5
bob_message_ids: set[int] = {m.id for m in bob_messages}
expected_bob_ids = {
messages[0].id, # alice(250) - preceding context
messages[1].id, # bob(400)
messages[2].id, # steve(300)
messages[3].id, # alice(500)
@ -520,7 +521,7 @@ class TestQueueProcessing:
# Ensure items are only for bob
assert all(qi.payload.get("observed") == bob.name for qi in bob_items)
# Test steve's work unit - starts at message 3 for per-work-unit anchoring
# Test steve's work unit - now includes preceding message for context
steve_work_unit_key = steve_queue_items[0].work_unit_key
steve_aqs = models.ActiveQueueSession(work_unit_key=steve_work_unit_key)
db_session.add(steve_aqs)
@ -533,10 +534,11 @@ class TestQueueProcessing:
aqs_id=steve_aqs.id,
)
# Steve should get 5 messages (3..7)
assert len(steve_messages) == 5
# Steve should get 6 messages (2..7) - includes preceding bob message for context
assert len(steve_messages) == 6
steve_message_ids: set[int] = {m.id for m in steve_messages}
expected_steve_ids = {
messages[1].id, # bob(400) - preceding context
messages[2].id, # steve(300)
messages[3].id, # alice(500)
messages[4].id, # bob(500)
@ -634,8 +636,9 @@ class TestQueueProcessing:
# Mock the token limit to 1500 for this test
with patch.object(settings.DERIVER, "REPRESENTATION_BATCH_MAX_TOKENS", 1500):
# Test alice's work unit
# With per-work-unit anchoring, Alice starts at her own first message (message 3)
# Alice's batch: alice(100) + alice(200) = 300 tokens, well under 1500 limit
# With per-work-unit anchoring + preceding context:
# Alice starts at message 3, includes preceding message 2 (steve) for context
# Alice's batch: steve(800) + alice(100) + alice(200) = 1100 tokens, under 1500 limit
if alice_queue_items:
alice_work_unit_key = alice_queue_items[0].work_unit_key
alice_aqs = models.ActiveQueueSession(work_unit_key=alice_work_unit_key)
@ -649,15 +652,16 @@ class TestQueueProcessing:
aqs_id=alice_aqs.id,
)
# Per-work-unit anchoring: Alice starts at message 3 -> [3,4]
assert len(alice_messages2) == 2
# Includes preceding steve message for context -> [2,3,4]
assert len(alice_messages2) == 3
assert [m.id for m in alice_messages2] == [
messages[2].id,
messages[3].id,
messages[1].id, # steve - preceding context
messages[2].id, # alice
messages[3].id, # alice
]
# Test bob's work unit
# With per-work-unit anchoring, Bob starts at his own first message (message 1)
# Bob starts at message 1, no preceding message available
# Bob's batch: bob(800) only, under 1500 limit
if bob_queue_items:
bob_work_unit_key = bob_queue_items[0].work_unit_key
@ -676,8 +680,9 @@ class TestQueueProcessing:
assert bob_messages2[0].id == messages[0].id # bob only
# Test steve's work unit
# With per-work-unit anchoring, Steve starts at his own first message (message 2)
# Steve's batch: steve(800) only, under 1500 limit
# Steve starts at message 2, includes preceding message 1 (bob) for context
# Steve's batch: bob(800) + steve(800) = 1600 tokens, exceeds 1500 limit
# So should only get steve's message
if steve_queue_items:
steve_work_unit_key = steve_queue_items[0].work_unit_key
steve_aqs = models.ActiveQueueSession(work_unit_key=steve_work_unit_key)
@ -691,10 +696,11 @@ class TestQueueProcessing:
aqs_id=steve_aqs.id,
)
# Per-work-unit anchoring: Steve starts at message 2 -> [2]
assert len(steve_messages2) == 1
# Includes preceding bob message for context -> [1,2]
assert len(steve_messages2) == 2
assert [m.id for m in steve_messages2] == [
messages[1].id,
messages[0].id, # bob - preceding context
messages[1].id, # steve
]
@pytest.mark.asyncio

25
tests/deriver/test_representation_crud.py
View File

@ -2,7 +2,6 @@ import datetime
from src.utils.representation import (
DeductiveObservation,
DeductiveObservationBase,
ExplicitObservation,
ExplicitObservationBase,
PromptRepresentation,
@ -73,14 +72,26 @@ def test_representation_formatting_methods():
md = rep.format_as_markdown()
assert "## Explicit Observations" in md
assert "## Deductive Observations" in md
assert "**Conclusion**: owns a pet" in md
assert "owns a pet" in md
assert "Premises:" in md
def test_prompt_representation_conversion():
"""PromptRepresentation.to_representation maps strings to observation objects."""
"""PromptRepresentation.to_representation maps strings to observation objects.
Note: In the current architecture, the Deriver only creates explicit observations.
Deductive and inductive observations are created by the Dreamer agent.
Therefore, from_prompt_representation only converts explicit observations.
"""
pr = PromptRepresentation(
explicit=[ExplicitObservationBase(content="A")],
deductive=[DeductiveObservationBase(conclusion="C", premises=["P1"])],
# Deductive observations in PromptRepresentation are ignored by from_prompt_representation
# because the Deriver only produces explicit observations
# deductive=[
# DeductiveObservationBase(
# conclusion="C", premises=["P1"], source_ids=["id1"]
# )
# ],
)
timestamp = datetime.datetime(2025, 1, 1, 12, 0, 0, tzinfo=datetime.timezone.utc)
rep = Representation.from_prompt_representation(
@ -91,7 +102,7 @@ def test_prompt_representation_conversion():
)
assert isinstance(rep, Representation)
assert [e.content for e in rep.explicit] == ["A"]
assert rep.deductive[0].conclusion == "C"
assert rep.deductive[0].premises == ["P1"]
# Deductive observations from PromptRepresentation are not converted
# (they would be created directly by the Dreamer via the create_observations tool)
assert len(rep.deductive) == 0
assert rep.explicit[0].created_at == timestamp
assert rep.deductive[0].created_at == timestamp

0
tests/dreamer/__init__.py → tests/dialectic/__init__.py
View File

64
tests/dreamer/test_dream_scheduler.py
View File

@ -1,64 +0,0 @@
"""Tests for DreamScheduler singleton pattern."""
import asyncio
import contextlib
from src.dreamer.dream_scheduler import DreamScheduler
def test_dream_scheduler_singleton():
"""Test that DreamScheduler implements proper singleton pattern."""
# Reset singleton state
DreamScheduler.reset_singleton()
# Create first instance
scheduler1 = DreamScheduler()
# Create second instance
scheduler2 = DreamScheduler()
# Both should be the same instance
assert scheduler1 is scheduler2
# Should share the same pending_dreams dict
assert scheduler1.pending_dreams is scheduler2.pending_dreams
async def test_dream_scheduler_initialized_once():
"""Test that DreamScheduler is only initialized once."""
# Reset singleton state
DreamScheduler.reset_singleton()
# Create first instance
scheduler1 = DreamScheduler()
# Create a dummy task to add to pending_dreams
async def dummy_task():
pass
task = asyncio.create_task(dummy_task())
scheduler1.pending_dreams["test_key"] = task
# Create second instance
scheduler2 = DreamScheduler()
# Second instance should have the same data as first
assert "test_key" in scheduler2.pending_dreams
assert scheduler2.pending_dreams["test_key"] is task
# Cleanup
task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await task
def test_dream_scheduler_multiple_instances():
"""Test that creating multiple instances doesn't reset state."""
# Reset singleton state
DreamScheduler.reset_singleton()
instances = [DreamScheduler() for _ in range(5)]
# All instances should be the same
for instance in instances[1:]:
assert instance is instances[0]

56
tests/dreamer/test_queue_manager_singleton.py
View File

@ -1,56 +0,0 @@
"""Tests that QueueManager instances share the same DreamScheduler singleton."""
import asyncio
import contextlib
from src.deriver.queue_manager import QueueManager
from src.dreamer.dream_scheduler import DreamScheduler
def test_queue_manager_shares_dream_scheduler():
"""Test that multiple QueueManager instances share the same DreamScheduler."""
# Reset singleton state
DreamScheduler.reset_singleton()
# Create first QueueManager
manager1 = QueueManager()
# Create second QueueManager
manager2 = QueueManager()
# Both should have the same DreamScheduler instance
assert manager1.dream_scheduler is manager2.dream_scheduler
# Should share the same pending_dreams dict
assert (
manager1.dream_scheduler.pending_dreams
is manager2.dream_scheduler.pending_dreams
)
async def test_queue_manager_preserves_dream_scheduler_state():
"""Test that creating a new QueueManager doesn't reset DreamScheduler state."""
# Reset singleton state
DreamScheduler.reset_singleton()
# Create first QueueManager and modify scheduler state
manager1 = QueueManager()
# Create a dummy task to add to pending_dreams
async def dummy_task():
pass
task = asyncio.create_task(dummy_task())
manager1.dream_scheduler.pending_dreams["test_key"] = task
# Create second QueueManager
manager2 = QueueManager()
# Second manager should see the first manager's state
assert "test_key" in manager2.dream_scheduler.pending_dreams
assert manager2.dream_scheduler.pending_dreams["test_key"] is task
# Cleanup
task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await task

56
tests/integration/test_enqueue.py
View File

@ -89,16 +89,12 @@ class TestEnqueueFunction:
# SESSION MESSAGES
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
async def test_session_with_deriver_disabled(
self,
mock_tracked_db: AsyncMock,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
):
"""Test that deriver disabled sessions skip representation but allows summary"""
mock_tracked_db.return_value.__aenter__.return_value = db_session
test_workspace, test_peer = sample_data
# Create session with deriver disabled
@ -129,15 +125,11 @@ class TestEnqueueFunction:
), f"Expected no queue items, but got {final_count - initial_count}"
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
async def test_session_normal_processing_single_peer(
self,
mock_tracked_db: AsyncMock,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
):
mock_tracked_db.return_value.__aenter__.return_value = db_session
test_workspace, test_peer = sample_data
test_session = await crud.get_or_create_session(
@ -174,16 +166,12 @@ class TestEnqueueFunction:
assert "representation" in task_types
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
async def test_session_with_multiple_peers_none_observe_others(
self,
mock_tracked_db: AsyncMock,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
):
"""Test session processing with multiple peers where some observe others"""
mock_tracked_db.return_value.__aenter__.return_value = db_session
test_workspace, test_peer1 = sample_data
# Create second peer
@ -254,16 +242,12 @@ class TestEnqueueFunction:
assert expected in actual_payloads
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
async def test_session_with_multiple_peers_all_observe_others(
self,
mock_tracked_db: AsyncMock,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
):
"""Test session processing with multiple peers where some observe others"""
mock_tracked_db.return_value.__aenter__.return_value = db_session
test_workspace, test_peer1 = sample_data
# Create second peer
@ -340,16 +324,12 @@ class TestEnqueueFunction:
assert expected in actual_payloads
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
async def test_session_with_multiple_peers_some_observe_others(
self,
mock_tracked_db: AsyncMock,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
):
"""Test session processing with multiple peers where some observe others"""
mock_tracked_db.return_value.__aenter__.return_value = db_session
test_workspace, test_peer1 = sample_data
# Create second peer
@ -438,16 +418,12 @@ class TestEnqueueFunction:
]
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
async def test_session_peer_config_overrides_peer_config(
self,
mock_tracked_db: AsyncMock,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
):
"""Test that session peer config overrides peer config"""
mock_tracked_db.return_value.__aenter__.return_value = db_session
test_workspace, test_peer = sample_data
# Set peer configuration to observe_me=True
@ -480,16 +456,12 @@ class TestEnqueueFunction:
assert final_count - initial_count == 0
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
async def test_multi_sender_scenario(
self,
mock_tracked_db: AsyncMock,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
):
"""Test complex scenario with multiple peers and mixed configurations"""
mock_tracked_db.return_value.__aenter__.return_value = db_session
test_workspace, test_peer1 = sample_data
# Create additional peers
@ -587,16 +559,12 @@ class TestEnqueueFunction:
# RACE CONDITION TESTS - Testing the new logic for peers that have left
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
async def test_sender_left_session_after_message_sent(
self,
mock_tracked_db: AsyncMock,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
):
"""Test that messages from senders who left the session still get processed with default config"""
mock_tracked_db.return_value.__aenter__.return_value = db_session
test_workspace, sender_peer = sample_data
# Create an observer peer
@ -680,16 +648,12 @@ class TestEnqueueFunction:
assert expected in actual_payloads
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
async def test_observer_left_session_no_queue_items_generated(
self,
mock_tracked_db: AsyncMock,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
):
"""Test that peers who left the session don't get representation tasks enqueued"""
mock_tracked_db.return_value.__aenter__.return_value = db_session
test_workspace, sender_peer = sample_data
# Create observer peers - one will leave, one will stay
@ -762,16 +726,12 @@ class TestEnqueueFunction:
assert sender_peer.name in observers
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
async def test_sender_not_in_peer_configuration_uses_defaults(
self,
mock_tracked_db: AsyncMock,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
):
"""Test get_effective_observe_me handles missing sender configuration gracefully"""
mock_tracked_db.return_value.__aenter__.return_value = db_session
test_workspace, existing_peer = sample_data
# Create observer peer
@ -842,16 +802,12 @@ class TestEnqueueFunction:
assert expected in actual_payloads
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
async def test_mixed_active_inactive_peers_complex_scenario(
self,
mock_tracked_db: AsyncMock,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
):
"""Test complex scenario with mix of active/inactive peers and different configurations"""
mock_tracked_db.return_value.__aenter__.return_value = db_session
test_workspace, sender_peer = sample_data
# Create multiple peers with different roles
@ -1142,16 +1098,12 @@ class TestAdvancedEnqueueEdgeCases:
return len(result.scalars().all())
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
async def test_edge_case_all_peers_left_except_sender(
self,
mock_tracked_db: AsyncMock,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
):
"""Test edge case where all observer peers have left the session"""
mock_tracked_db.return_value.__aenter__.return_value = db_session
test_workspace, sender_peer = sample_data
# Create multiple observer peers
@ -1218,16 +1170,12 @@ class TestAdvancedEnqueueEdgeCases:
assert queue_items[0].payload["task_type"] == "representation"
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
async def test_edge_case_sender_and_observer_both_left_different_times(
self,
mock_tracked_db: AsyncMock,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
):
"""Test race condition where both sender and observer left at different times"""
mock_tracked_db.return_value.__aenter__.return_value = db_session
test_workspace, sender_peer = sample_data
observer_peer = models.Peer(
@ -1304,16 +1252,12 @@ class TestAdvancedEnqueueEdgeCases:
assert queue_items[0].payload["task_type"] == "representation"
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
async def test_edge_case_message_from_never_joined_peer(
self,
mock_tracked_db: AsyncMock,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
):
"""Test handling message from peer who was never in the session"""
mock_tracked_db.return_value.__aenter__.return_value = db_session
test_workspace, existing_peer = sample_data
observer_peer = models.Peer(

38
tests/integration/test_representation.py
View File

@ -19,7 +19,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, models
from src.utils.representation import (
DeductiveObservation,
DeductiveObservationBase,
# DeductiveObservationBase,
ExplicitObservation,
ExplicitObservationBase,
PromptRepresentation,
@ -133,8 +133,8 @@ class TestRepresentationWorkflow:
markdown_output = representation.format_as_markdown()
assert "## Explicit Observations" in markdown_output
assert "## Deductive Observations" in markdown_output
assert "**Conclusion**:" in markdown_output
assert "**Premises**:" in markdown_output
assert "User probably has a dog named Rover" in markdown_output
assert "Premises:" in markdown_output
async def test_representation_merging_and_diffing(self):
"""Test representation merge and diff operations"""
@ -313,8 +313,9 @@ class TestDocumentCreationWorkflow:
async def test_representation_from_documents(self):
"""Test converting documents to representation"""
# Create test documents
# Create test documents with IDs (simulating database-assigned IDs)
explicit_doc = models.Document(
id="test_explicit_doc_id",
workspace_name="test_workspace",
observer="test_peer",
observed="test_peer",
@ -329,6 +330,7 @@ class TestDocumentCreationWorkflow:
)
deductive_doc = models.Document(
id="test_deductive_doc_id",
workspace_name="test_workspace",
observer="test_peer",
observed="test_peer",
@ -406,18 +408,17 @@ class TestPromptRepresentationConversion:
"""Test conversion between PromptRepresentation and Representation"""
async def test_prompt_representation_to_representation(self):
"""Test converting PromptRepresentation to Representation"""
"""Test converting PromptRepresentation to Representation.
Note: In the current architecture, the Deriver only creates explicit observations.
Deductive and inductive observations are created by the Dreamer agent.
Therefore, from_prompt_representation only converts explicit observations.
"""
prompt_rep = PromptRepresentation(
explicit=[
ExplicitObservationBase(content="User likes coffee"),
ExplicitObservationBase(content="User works remotely"),
],
deductive=[
DeductiveObservationBase(
conclusion="User probably works from a coffee shop sometimes",
premises=["User likes coffee", "User works remotely"],
)
],
)
timestamp = datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
@ -430,7 +431,9 @@ class TestPromptRepresentationConversion:
)
assert len(representation.explicit) == 2
assert len(representation.deductive) == 1
# Deductive observations from PromptRepresentation are not converted
# (they would be created directly by the Dreamer via the create_observations tool)
assert len(representation.deductive) == 0
# Check explicit observations
assert representation.explicit[0].content == "User likes coffee"
@ -439,17 +442,6 @@ class TestPromptRepresentationConversion:
assert representation.explicit[1].content == "User works remotely"
assert representation.explicit[0].created_at == timestamp
# Check deductive observation
deductive_obs = representation.deductive[0]
assert (
deductive_obs.conclusion
== "User probably works from a coffee shop sometimes"
)
assert deductive_obs.premises == ["User likes coffee", "User works remotely"]
assert deductive_obs.message_ids == [123]
assert deductive_obs.session_name == "test_session"
assert deductive_obs.created_at == timestamp
async def test_empty_prompt_representation_conversion(self):
"""Test converting empty PromptRepresentation"""
empty_prompt_rep = PromptRepresentation()

802
tests/integration/test_token_metrics.py Normal file
View File

@ -0,0 +1,802 @@
"""Integration tests for prometheus token metrics tracking.
These tests verify that DERIVER_TOKENS_PROCESSED and DIALECTIC_TOKENS_PROCESSED
metrics are correctly emitted with accurate token counts when processing messages
and dialectic queries.
The approach uses delta-based verification:
1. Capture counter values before test execution
2. Run the code under test (with mocked LLM)
3. Verify deltas match expected values
"""
from unittest.mock import AsyncMock, patch
import pytest
from nanoid import generate as generate_nanoid
from prometheus_client import REGISTRY
from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, models, schemas
from src.models import Peer, Workspace
from src.schemas import (
ResolvedConfiguration,
ResolvedDeriverConfiguration,
ResolvedDreamConfiguration,
ResolvedPeerCardConfiguration,
ResolvedSummaryConfiguration,
)
from src.utils.clients import HonchoLLMCallResponse
from src.utils.representation import ExplicitObservationBase, PromptRepresentation
from src.utils.summarizer import (
SummaryType,
_create_and_save_summary, # pyright: ignore[reportPrivateUsage]
estimate_short_summary_prompt_tokens,
)
# =============================================================================
# Fixtures
# =============================================================================
class MetricDeltaChecker:
"""Utility class to capture and verify prometheus counter deltas."""
def capture(self, metric_name: str, labels: dict[str, str]) -> float:
"""Capture current value of a counter with specific labels."""
# Counters have _total suffix in prometheus
full_name = (
metric_name if metric_name.endswith("_total") else f"{metric_name}_total"
)
value = REGISTRY.get_sample_value(full_name, labels=labels)
return value or 0.0
def get_delta(
self, metric_name: str, labels: dict[str, str], before: float
) -> float:
"""Get the delta between a before value and current."""
return self.capture(metric_name, labels) - before
def assert_delta(
self,
metric_name: str,
labels: dict[str, str],
before: float,
expected: int | float,
message: str = "",
) -> None:
"""Assert that the delta matches expected value."""
delta = self.get_delta(metric_name, labels, before)
assert (
delta == expected
), f"{message}: expected delta {expected}, got {delta}. Labels: {labels}"
@pytest.fixture
def metric_checker() -> MetricDeltaChecker:
"""Fixture providing a metric delta checker instance."""
return MetricDeltaChecker()
@pytest.fixture
def enable_metrics(monkeypatch: pytest.MonkeyPatch):
"""Enable prometheus metrics with a test namespace."""
monkeypatch.setattr("src.prometheus.METRICS_ENABLED", True)
monkeypatch.setattr("src.config.settings.METRICS.NAMESPACE", "test")
yield
# =============================================================================
# Test Data Helpers
# =============================================================================
async def create_test_session_with_peer(
db_session: AsyncSession,
workspace: Workspace,
peer: Peer,
) -> models.Session:
"""Create a session with a peer configured for observation."""
session = await crud.get_or_create_session(
db_session,
schemas.SessionCreate(
name=str(generate_nanoid()),
peers={peer.name: schemas.SessionPeerConfig(observe_me=True)},
),
workspace.name,
)
await db_session.commit()
return session
async def create_test_messages(
db_session: AsyncSession,
workspace_name: str,
session_name: str,
peer_name: str,
count: int = 1,
content_prefix: str = "Test message",
) -> list[models.Message]:
"""Create test messages in the database."""
messages: list[models.Message] = []
for i in range(count):
message = models.Message(
workspace_name=workspace_name,
session_name=session_name,
peer_name=peer_name,
content=f"{content_prefix} {i}",
public_id=generate_nanoid(),
seq_in_session=i + 1,
token_count=10,
)
db_session.add(message)
messages.append(message)
await db_session.commit()
# Refresh to get IDs
for msg in messages:
await db_session.refresh(msg)
return messages
def create_test_configuration() -> ResolvedConfiguration:
"""Create a test configuration to avoid DB lookups in tests."""
return ResolvedConfiguration(
deriver=ResolvedDeriverConfiguration(enabled=True),
peer_card=ResolvedPeerCardConfiguration(use=False, create=False),
summary=ResolvedSummaryConfiguration(
enabled=True, messages_per_short_summary=20, messages_per_long_summary=60
),
dream=ResolvedDreamConfiguration(enabled=False),
)
def create_mock_deriver_response(
output_tokens: int = 42,
) -> HonchoLLMCallResponse[PromptRepresentation]:
"""Create a mock LLM response for the deriver."""
return HonchoLLMCallResponse(
content=PromptRepresentation(
explicit=[ExplicitObservationBase(content="Test observation from deriver")],
),
input_tokens=100,
output_tokens=output_tokens,
finish_reasons=["end_turn"],
)
def create_mock_dialectic_response(
input_tokens: int = 150, output_tokens: int = 75
) -> HonchoLLMCallResponse[str]:
"""Create a mock LLM response for the dialectic."""
return HonchoLLMCallResponse(
content="This is a test dialectic response about the peer.",
input_tokens=input_tokens,
output_tokens=output_tokens,
cache_read_input_tokens=0,
cache_creation_input_tokens=0,
finish_reasons=["end_turn"],
tool_calls_made=[],
)
# =============================================================================
# Deriver Ingestion Metrics Tests
# =============================================================================
@pytest.mark.asyncio
@pytest.mark.usefixtures("enable_metrics")
class TestDeriverIngestionMetrics:
"""Test token metrics for deriver INGESTION task type."""
async def test_ingestion_tracks_output_tokens(
self,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
metric_checker: MetricDeltaChecker,
):
"""Verify OUTPUT_TOTAL tokens match response.output_tokens from LLM."""
from src.deriver.deriver import process_representation_tasks_batch
workspace, peer = sample_data
session = await create_test_session_with_peer(db_session, workspace, peer)
messages = await create_test_messages(
db_session, workspace.name, session.name, peer.name, count=1
)
expected_output_tokens = 42
mock_response = create_mock_deriver_response(
output_tokens=expected_output_tokens
)
# Capture metrics before
labels = {
"namespace": "test",
"task_type": "ingestion",
"token_type": "output",
"component": "output_total",
}
before = metric_checker.capture("deriver_tokens_processed", labels)
# Mock the LLM call and save_representation (we're testing metrics, not DB writes)
with (
patch(
"src.deriver.deriver.honcho_llm_call",
new=AsyncMock(return_value=mock_response),
),
patch(
"src.crud.representation.RepresentationManager.save_representation",
new=AsyncMock(),
),
):
await process_representation_tasks_batch(
messages=messages,
message_level_configuration=create_test_configuration(),
observer=peer.name,
observed=peer.name,
)
# Verify output tokens metric
metric_checker.assert_delta(
"deriver_tokens_processed",
labels,
before,
expected_output_tokens,
"Ingestion output tokens",
)
async def test_ingestion_tracks_prompt_input_tokens(
self,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
metric_checker: MetricDeltaChecker,
):
"""Verify PROMPT component is tracked for ingestion input."""
from src.deriver.deriver import process_representation_tasks_batch
from src.deriver.prompts import estimate_minimal_deriver_prompt_tokens
workspace, peer = sample_data
session = await create_test_session_with_peer(db_session, workspace, peer)
messages = await create_test_messages(
db_session, workspace.name, session.name, peer.name, count=1
)
mock_response = create_mock_deriver_response()
# Get expected prompt tokens
expected_prompt_tokens = estimate_minimal_deriver_prompt_tokens()
labels = {
"namespace": "test",
"task_type": "ingestion",
"token_type": "input",
"component": "prompt",
}
before = metric_checker.capture("deriver_tokens_processed", labels)
with (
patch(
"src.deriver.deriver.honcho_llm_call",
new=AsyncMock(return_value=mock_response),
),
patch(
"src.crud.representation.RepresentationManager.save_representation",
new=AsyncMock(),
),
):
await process_representation_tasks_batch(
messages=messages,
message_level_configuration=create_test_configuration(),
observer=peer.name,
observed=peer.name,
)
metric_checker.assert_delta(
"deriver_tokens_processed",
labels,
before,
expected_prompt_tokens,
"Ingestion prompt input tokens",
)
async def test_ingestion_tracks_messages_input_tokens(
self,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
metric_checker: MetricDeltaChecker,
):
"""Verify MESSAGES component is tracked for ingestion input."""
from src.deriver.deriver import process_representation_tasks_batch
workspace, peer = sample_data
session = await create_test_session_with_peer(db_session, workspace, peer)
messages = await create_test_messages(
db_session,
workspace.name,
session.name,
peer.name,
count=1,
content_prefix="Hello this is a test message",
)
mock_response = create_mock_deriver_response()
labels = {
"namespace": "test",
"task_type": "ingestion",
"token_type": "input",
"component": "messages",
}
before = metric_checker.capture("deriver_tokens_processed", labels)
with (
patch(
"src.deriver.deriver.honcho_llm_call",
new=AsyncMock(return_value=mock_response),
),
patch(
"src.crud.representation.RepresentationManager.save_representation",
new=AsyncMock(),
),
):
await process_representation_tasks_batch(
messages=messages,
message_level_configuration=create_test_configuration(),
observer=peer.name,
observed=peer.name,
)
# Verify messages tokens were tracked (should be > 0)
delta = metric_checker.get_delta("deriver_tokens_processed", labels, before)
assert delta > 0, f"Expected messages input tokens > 0, got {delta}"
# =============================================================================
# Deriver Summary Metrics Tests
# =============================================================================
@pytest.mark.asyncio
@pytest.mark.usefixtures("enable_metrics")
class TestDeriverSummaryMetrics:
"""Test token metrics for deriver SUMMARY task type."""
async def test_summary_tracks_output_tokens(
self,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
metric_checker: MetricDeltaChecker,
):
"""Verify OUTPUT_TOTAL tokens are tracked for summary."""
workspace, peer = sample_data
session = await create_test_session_with_peer(db_session, workspace, peer)
# Create messages for summary
messages = await create_test_messages(
db_session, workspace.name, session.name, peer.name, count=5
)
last_message = messages[-1]
# Mock _create_summary to return a summary with known token count
expected_output_tokens = 25
mock_summary = {
"content": "This is a test summary.",
"token_count": expected_output_tokens,
"message_id": last_message.id,
"message_public_id": last_message.public_id,
}
labels = {
"namespace": "test",
"task_type": "summary",
"token_type": "output",
"component": "output_total",
}
before = metric_checker.capture("deriver_tokens_processed", labels)
with (
patch(
"src.utils.summarizer._create_summary",
new=AsyncMock(return_value=(mock_summary, False)), # is_fallback=False
),
patch(
"src.utils.summarizer._save_summary",
new=AsyncMock(),
),
):
await _create_and_save_summary(
db=db_session,
workspace_name=workspace.name,
session_name=session.name,
message_id=last_message.id,
message_seq_in_session=last_message.seq_in_session,
summary_type=SummaryType.SHORT,
message_public_id=last_message.public_id,
configuration=create_test_configuration(),
)
# Verify output tokens match the summary token_count
metric_checker.assert_delta(
"deriver_tokens_processed",
labels,
before,
expected_output_tokens,
"Summary output tokens",
)
async def test_summary_tracks_prompt_input_tokens(
self,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
metric_checker: MetricDeltaChecker,
):
"""Verify PROMPT component is tracked for summary input."""
workspace, peer = sample_data
session = await create_test_session_with_peer(db_session, workspace, peer)
messages = await create_test_messages(
db_session, workspace.name, session.name, peer.name, count=5
)
last_message = messages[-1]
expected_prompt_tokens = estimate_short_summary_prompt_tokens()
mock_summary = {
"content": "Test summary.",
"token_count": 10,
"message_id": last_message.id,
"message_public_id": last_message.public_id,
}
labels = {
"namespace": "test",
"task_type": "summary",
"token_type": "input",
"component": "prompt",
}
before = metric_checker.capture("deriver_tokens_processed", labels)
with (
patch(
"src.utils.summarizer._create_summary",
new=AsyncMock(return_value=(mock_summary, False)),
),
patch(
"src.utils.summarizer._save_summary",
new=AsyncMock(),
),
):
await _create_and_save_summary(
db=db_session,
workspace_name=workspace.name,
session_name=session.name,
message_id=last_message.id,
message_seq_in_session=last_message.seq_in_session,
summary_type=SummaryType.SHORT,
message_public_id=last_message.public_id,
configuration=create_test_configuration(),
)
metric_checker.assert_delta(
"deriver_tokens_processed",
labels,
before,
expected_prompt_tokens,
"Summary prompt input tokens",
)
async def test_summary_tracks_messages_input_tokens(
self,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
metric_checker: MetricDeltaChecker,
):
"""Verify MESSAGES component is tracked for summary input."""
workspace, peer = sample_data
session = await create_test_session_with_peer(db_session, workspace, peer)
messages = await create_test_messages(
db_session, workspace.name, session.name, peer.name, count=5
)
last_message = messages[-1]
# Get the actual messages that would be included in the summary
# (to match what _create_and_save_summary computes via get_messages_by_seq_range)
actual_messages = await crud.get_messages_by_seq_range(
db_session,
workspace.name,
session.name,
start_seq=1,
end_seq=last_message.seq_in_session,
)
expected_messages_tokens = sum(m.token_count for m in actual_messages)
mock_summary = {
"content": "Test summary.",
"token_count": 10,
"message_id": last_message.id,
"message_public_id": last_message.public_id,
}
labels = {
"namespace": "test",
"task_type": "summary",
"token_type": "input",
"component": "messages",
}
before = metric_checker.capture("deriver_tokens_processed", labels)
with (
patch(
"src.utils.summarizer._create_summary",
new=AsyncMock(return_value=(mock_summary, False)),
),
patch(
"src.utils.summarizer._save_summary",
new=AsyncMock(),
),
):
await _create_and_save_summary(
db=db_session,
workspace_name=workspace.name,
session_name=session.name,
message_id=last_message.id,
message_seq_in_session=last_message.seq_in_session,
summary_type=SummaryType.SHORT,
message_public_id=last_message.public_id,
configuration=create_test_configuration(),
)
# Verify messages tokens match what summarizer actually computed
delta = metric_checker.get_delta("deriver_tokens_processed", labels, before)
assert (
delta == expected_messages_tokens
), f"Expected messages input tokens {expected_messages_tokens}, got {delta}"
assert delta > 0, "Expected at least some message tokens to be tracked"
async def test_summary_fallback_does_not_track(
self,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
metric_checker: MetricDeltaChecker,
):
"""Verify metrics are NOT emitted when _create_summary returns is_fallback=True."""
workspace, peer = sample_data
session = await create_test_session_with_peer(db_session, workspace, peer)
messages = await create_test_messages(
db_session, workspace.name, session.name, peer.name, count=5
)
last_message = messages[-1]
mock_summary = {
"content": "Fallback summary.",
"token_count": 10,
"message_id": last_message.id,
"message_public_id": last_message.public_id,
}
# Capture all relevant metric labels before
output_labels = {
"namespace": "test",
"task_type": "summary",
"token_type": "output",
"component": "output_total",
}
prompt_labels = {
"namespace": "test",
"task_type": "summary",
"token_type": "input",
"component": "prompt",
}
before_output = metric_checker.capture(
"deriver_tokens_processed", output_labels
)
before_prompt = metric_checker.capture(
"deriver_tokens_processed", prompt_labels
)
with patch(
"src.utils.summarizer._create_summary",
new=AsyncMock(return_value=(mock_summary, True)), # is_fallback=True
):
await _create_and_save_summary(
db=db_session,
workspace_name=workspace.name,
session_name=session.name,
message_id=last_message.id,
message_seq_in_session=last_message.seq_in_session,
summary_type=SummaryType.SHORT,
message_public_id=last_message.public_id,
configuration=create_test_configuration(),
)
# Verify NO change in metrics when fallback
output_delta = metric_checker.get_delta(
"deriver_tokens_processed", output_labels, before_output
)
prompt_delta = metric_checker.get_delta(
"deriver_tokens_processed", prompt_labels, before_prompt
)
assert (
output_delta == 0
), f"Expected no output token change on fallback, got {output_delta}"
assert (
prompt_delta == 0
), f"Expected no prompt token change on fallback, got {prompt_delta}"
# =============================================================================
# Dialectic Token Metrics Tests
# =============================================================================
@pytest.mark.asyncio
@pytest.mark.usefixtures("enable_metrics")
class TestDialecticTokenMetrics:
"""Test token metrics for dialectic calls."""
async def test_dialectic_tracks_input_tokens(
self,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
metric_checker: MetricDeltaChecker,
):
"""Verify INPUT tokens are tracked from LLM response."""
from src.dialectic.core import DialecticAgent
workspace, peer = sample_data
session = await create_test_session_with_peer(db_session, workspace, peer)
expected_input_tokens = 150
mock_response = create_mock_dialectic_response(
input_tokens=expected_input_tokens, output_tokens=75
)
labels = {
"namespace": "test",
"token_type": "input",
"component": "total",
"reasoning_level": "low",
}
before = metric_checker.capture("dialectic_tokens_processed", labels)
agent = DialecticAgent(
db=db_session,
workspace_name=workspace.name,
session_name=session.name,
observer=peer.name,
observed=peer.name,
)
with patch(
"src.dialectic.core.honcho_llm_call",
new=AsyncMock(return_value=mock_response),
):
await agent.answer("What do you know about this user?")
metric_checker.assert_delta(
"dialectic_tokens_processed",
labels,
before,
expected_input_tokens,
"Dialectic input tokens",
)
async def test_dialectic_tracks_output_tokens(
self,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
metric_checker: MetricDeltaChecker,
):
"""Verify OUTPUT tokens are tracked from LLM response."""
from src.dialectic.core import DialecticAgent
workspace, peer = sample_data
session = await create_test_session_with_peer(db_session, workspace, peer)
expected_output_tokens = 75
mock_response = create_mock_dialectic_response(
input_tokens=150, output_tokens=expected_output_tokens
)
labels = {
"namespace": "test",
"token_type": "output",
"component": "total",
"reasoning_level": "low",
}
before = metric_checker.capture("dialectic_tokens_processed", labels)
agent = DialecticAgent(
db=db_session,
workspace_name=workspace.name,
session_name=session.name,
observer=peer.name,
observed=peer.name,
)
with patch(
"src.dialectic.core.honcho_llm_call",
new=AsyncMock(return_value=mock_response),
):
await agent.answer("What do you know about this user?")
metric_checker.assert_delta(
"dialectic_tokens_processed",
labels,
before,
expected_output_tokens,
"Dialectic output tokens",
)
async def test_dialectic_metrics_disabled_no_emission(
self,
db_session: AsyncSession,
sample_data: tuple[Workspace, Peer],
metric_checker: MetricDeltaChecker,
monkeypatch: pytest.MonkeyPatch,
):
"""Verify metrics are NOT emitted when METRICS_ENABLED=False."""
from src.dialectic.core import DialecticAgent
# Explicitly disable metrics
monkeypatch.setattr("src.prometheus.METRICS_ENABLED", False)
monkeypatch.setattr("src.config.settings.METRICS.NAMESPACE", "test")
workspace, peer = sample_data
session = await create_test_session_with_peer(db_session, workspace, peer)
mock_response = create_mock_dialectic_response(
input_tokens=200, output_tokens=100
)
# Capture before values
input_labels = {
"namespace": "test",
"token_type": "input",
"component": "total",
}
output_labels = {
"namespace": "test",
"token_type": "output",
"component": "total",
}
before_input = metric_checker.capture(
"dialectic_tokens_processed", input_labels
)
before_output = metric_checker.capture(
"dialectic_tokens_processed", output_labels
)
agent = DialecticAgent(
db=db_session,
workspace_name=workspace.name,
session_name=session.name,
observer=peer.name,
observed=peer.name,
)
with patch(
"src.dialectic.core.honcho_llm_call",
new=AsyncMock(return_value=mock_response),
):
await agent.answer("What do you know about this user?")
# Verify NO change in metrics
input_delta = metric_checker.get_delta(
"dialectic_tokens_processed", input_labels, before_input
)
output_delta = metric_checker.get_delta(
"dialectic_tokens_processed", output_labels, before_output
)
assert (
input_delta == 0
), f"Expected no input token change when disabled, got {input_delta}"
assert (
output_delta == 0
), f"Expected no output token change when disabled, got {output_delta}"

90
tests/test_llm_mock.py
View File

@ -1,90 +0,0 @@
from datetime import datetime, timezone
from typing import cast
from unittest.mock import MagicMock
import pytest
from src.models import Message
from src.utils.representation import (
DeductiveObservation,
ExplicitObservation,
PromptRepresentation,
Representation,
)
@pytest.mark.asyncio
async def test_generic_honcho_llm_call_mock():
"""Test that the generic honcho_llm_call mock is working for existing decorated functions"""
# Import a function that we know is decorated with honcho_llm_call
from src.deriver.deriver import critical_analysis_call
# Call the decorated function - this should use our mock
result: PromptRepresentation = await critical_analysis_call(
peer_id="test_peer_id",
peer_card=["test_peer_card"],
message_created_at=datetime(2023, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
working_representation=Representation(
explicit=[
ExplicitObservation(
content="test explicit observation",
created_at=datetime(2023, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
message_ids=[1],
session_name="test_session",
)
],
deductive=[
DeductiveObservation(
conclusion="test deductive conclusion",
premises=["test premise 1", "test premise 2"],
created_at=datetime(2023, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
message_ids=[1],
session_name="test_session",
)
],
),
history="test history",
new_turns=["test new turn"],
)
# Verify that we get a mock result, not an actual LLM call
assert result is not None
# The result should have the attributes we expect from our mock
assert hasattr(result, "explicit")
assert hasattr(result, "deductive")
assert hasattr(result, "_response")
@pytest.mark.asyncio
async def test_summarizer_decorated_functions_with_mock():
"""Test that summarizer decorated functions work with our mock"""
# Import functions that we know are decorated with honcho_llm_call
from src.utils.summarizer import create_long_summary, create_short_summary
# Create mock messages for testing
mock_message = MagicMock(spec=Message)
mock_message.content = "Test message content"
mock_message.peer_name = "test_peer"
mock_messages = cast(list[Message], [mock_message])
# Call the decorated functions - these should use our mock
short_result = await create_short_summary(
messages=mock_messages, input_tokens=100, previous_summary="Previous summary"
)
long_result = await create_long_summary(
messages=mock_messages, previous_summary="Previous summary"
)
# Verify that we get mock results, not actual LLM calls
assert short_result is not None
assert long_result is not None
# For functions with return_call_response=True, we should get a string or object with content
# The existing mock returns a string, so we check if it's a string
assert isinstance(short_result, str | object)
assert isinstance(long_result, str | object)
# If it's not a string, check for content attribute
if not isinstance(short_result, str):
assert hasattr(short_result, "content")
if not isinstance(long_result, str):
assert hasattr(long_result, "content")

Some files were not shown because too many files have changed in this diff Show More