Merge pull request #348 from plastic-labs/ben/trim-test-trial-tools

feat: give bench tools access to remote honchos
This commit is contained in:
doria 2026-01-26 14:18:39 -05:00 committed by GitHub
commit 5371463f39
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 1346 additions and 1602 deletions

8
src/cache/client.py vendored
View File

@ -3,7 +3,9 @@ from __future__ import annotations
import asyncio
import logging
from typing import cast
from urllib.parse import urlparse, urlunparse
import redis.asyncio as aioredis
import sentry_sdk
from cashews import cache
from cashews.picklers import PicklerType
@ -126,9 +128,11 @@ async def is_deriver_flush_enabled() -> bool:
if not is_cache_enabled():
return False
try:
import redis.asyncio as aioredis
# Strip query parameters - redis-py doesn't support custom params like ?suppress=true
parsed = urlparse(settings.CACHE.URL)
clean_url = urlunparse(parsed._replace(query=""))
redis_client = aioredis.from_url(settings.CACHE.URL) # pyright: ignore[reportUnknownMemberType]
redis_client = aioredis.from_url(clean_url) # pyright: ignore[reportUnknownMemberType]
try:
result = await redis_client.get(DERIVER_FLUSH_KEY)
return result == b"1"

View File

@ -66,7 +66,6 @@ Optional arguments:
import argparse
import asyncio
import os
import time
from datetime import datetime
from pathlib import Path
from typing import Any, cast
@ -82,7 +81,6 @@ from .beam_common import (
ConversationResult,
QuestionResult,
calculate_ability_scores,
format_duration,
generate_json_summary,
judge_event_ordering,
judge_nugget_based,
@ -91,11 +89,11 @@ from .beam_common import (
print_summary,
)
from .runner_common import (
ReasoningLevel,
RunnerMixin,
BaseRunner,
ItemContext,
RunnerConfig,
add_common_arguments,
create_openai_client,
export_metrics,
validate_common_arguments,
)
@ -104,48 +102,33 @@ bench_dir = Path(__file__).parent
load_dotenv(bench_dir / ".env")
class BEAMRunner(RunnerMixin):
class BEAMRunner(BaseRunner[ConversationResult]):
"""
Executes BEAM benchmark tests against a Honcho instance.
"""
def __init__(
self,
config: RunnerConfig,
data_dir: Path,
base_api_port: int = 8000,
pool_size: int = 1,
timeout_seconds: int | None = None,
cleanup_workspace: bool = True,
use_get_context: bool = False,
redis_url: str = "redis://localhost:6379/0",
reasoning_level: ReasoningLevel | None = None,
context_length: str,
conversation_ids: list[str] | None = None,
):
"""
Initialize the BEAM test runner.
Args:
config: Common runner configuration
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)
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
redis_url: Redis URL for flush mode signaling (default: redis://localhost:6379/0)
reasoning_level: Reasoning level for dialectic chat (default: None)
context_length: Context length (100K, 500K, 1M, 10M)
conversation_ids: Optional list of specific conversation IDs to run
"""
self.data_dir: Path = data_dir
self.base_api_port: int = base_api_port
self.pool_size: int = pool_size
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
self.redis_url: str = redis_url
self.reasoning_level: ReasoningLevel | None = reasoning_level
self.context_length: str = context_length
self.conversation_ids: list[str] | None = conversation_ids
# Initialize common components (metrics, logging)
self._init_common("beam")
# Initialize base class
super().__init__(config)
# Initialize OpenRouter client for judging
openrouter_base_url = os.getenv(
@ -161,10 +144,122 @@ class BEAMRunner(RunnerMixin):
"BEAM_JUDGE_MODEL", "anthropic/claude-sonnet-4.5"
)
def get_metrics_prefix(self) -> str:
return "beam"
def load_items(self) -> list[Any]:
"""Load conversation IDs to process."""
if self.conversation_ids:
return self.conversation_ids
return list_conversations(self.data_dir, self.context_length)
def get_workspace_id(self, item: Any) -> str:
"""Return workspace ID for a conversation."""
return f"beam_{self.context_length}_{item}"
def get_session_id(self, item: Any, workspace_id: str) -> str:
"""Return session ID for a conversation."""
return f"{workspace_id}_session"
async def setup_peers(self, ctx: ItemContext, item: Any) -> None:
"""Create user and assistant peers."""
ctx.peers["user"] = await ctx.honcho_client.aio.peer(id="user")
ctx.peers["assistant"] = await ctx.honcho_client.aio.peer(id="assistant")
async def setup_session(self, ctx: ItemContext, item: Any) -> None:
"""Create and configure session - observe the user peer."""
user_peer = ctx.peers["user"]
assistant_peer = ctx.peers["assistant"]
ctx.session = await ctx.honcho_client.aio.session(
id=ctx.session_id, configuration=self._get_session_configuration()
)
await ctx.session.aio.add_peers(
[
(user_peer, SessionPeerConfig(observe_me=True, observe_others=False)),
(
assistant_peer,
SessionPeerConfig(observe_me=False, observe_others=False),
),
]
)
async def ingest_messages(self, ctx: ItemContext, item: Any) -> int:
"""Ingest conversation turns into the session."""
conversation_id = item
user_peer = ctx.peers["user"]
assistant_peer = ctx.peers["assistant"]
# Load conversation data
conv_data = load_conversation(
self.data_dir, self.context_length, conversation_id
)
chat_data = conv_data["chat"]
# Store questions data for later use
ctx.peers["_questions_data"] = conv_data["questions"]
messages: list[MessageCreateParams] = []
# Handle different data structures for 10M vs other sizes
for batch in chat_data:
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.extend(
self._process_turn(turn, user_peer, assistant_peer)
)
else:
# Standard structure for 100K, 500K, 1M
for turn_group in batch.get("turns", []):
for turn in turn_group:
messages.extend(
self._process_turn(turn, user_peer, assistant_peer)
)
# Add messages in batches of 100
for i in range(0, len(messages), 100):
batch = messages[i : i + 100]
await ctx.session.aio.add_messages(batch)
return len(messages)
def _process_turn(
self, turn: dict[str, Any], user_peer: Any, assistant_peer: Any
) -> list[MessageCreateParams]:
"""Process a single turn, handling long message splitting."""
role = turn["role"]
content = turn["content"]
messages: list[MessageCreateParams] = []
if len(content) > 25000:
chunks = [content[i : i + 25000] for i in range(0, len(content), 25000)]
for chunk in chunks:
if role == "user":
messages.append(user_peer.message(chunk))
elif role == "assistant":
messages.append(assistant_peer.message(chunk))
else:
if role == "user":
messages.append(user_peer.message(content))
elif role == "assistant":
messages.append(assistant_peer.message(content))
return messages
def get_dream_observers(self, item: Any) -> list[str]:
"""Return the observer - always user for BEAM."""
return ["user"]
async def _process_single_question(
self,
session: Any,
user_peer: Any,
ctx: ItemContext,
ability: str,
q_idx: int,
q_data: dict[str, Any],
@ -184,9 +279,8 @@ class BEAMRunner(RunnerMixin):
# Execute question using dialectic
# 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.aio.context(
if self.config.use_get_context or ability == "instruction_following":
context = await ctx.session.aio.context(
summary=True,
peer_target="user",
search_query=question,
@ -194,8 +288,7 @@ class BEAMRunner(RunnerMixin):
context_messages = context.to_openai(assistant="assistant")
context_messages.append({"role": "user", "content": question})
# For instruction_following, add a system prompt that tells the LLM
# to follow any stored user preferences/instructions in the context
# For instruction_following, add a system prompt
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.
@ -209,7 +302,6 @@ IMPORTANT: You MUST follow any instructions or preferences the user has previous
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})
@ -218,7 +310,7 @@ Review the context carefully for any such instructions before responding."""
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]
messages=cast(Any, messages),
)
if not response.choices or not response.choices[0].message:
@ -226,9 +318,9 @@ Review the context carefully for any such instructions before responding."""
else:
actual_response = response.choices[0].message.content or ""
else:
actual_response = await user_peer.aio.chat(
actual_response = await ctx.peers["user"].aio.chat(
question,
reasoning_level=self.reasoning_level,
reasoning_level=self.config.reasoning_level,
)
actual_response = (
actual_response if isinstance(actual_response, str) else ""
@ -283,35 +375,16 @@ Review the context carefully for any such instructions before responding."""
return question_result
async def execute_conversation(
self, context_length: str, conversation_id: str, honcho_url: str
async def execute_questions(
self, ctx: ItemContext, item: Any
) -> ConversationResult:
"""
Execute BEAM benchmark for a single conversation.
Args:
context_length: Context length (100K, 500K, 1M, 10M)
conversation_id: Conversation ID
honcho_url: URL of the Honcho instance to use
Returns:
Conversation execution results
"""
start_time = time.time()
print(f"\n{'=' * 80}")
print(
f"Executing BEAM conversation {conversation_id} ({context_length} context)"
)
print(f"{'=' * 80}")
# Create workspace for this conversation
workspace_id = f"beam_{context_length}_{conversation_id}"
honcho_client = self.create_honcho_client(workspace_id, honcho_url)
"""Execute all questions for the conversation."""
conversation_id = item
workspace_id = ctx.workspace_id
result: ConversationResult = {
"conversation_id": conversation_id,
"context_length": context_length,
"context_length": self.context_length,
"workspace_id": workspace_id,
"total_turns": 0,
"total_messages": 0,
@ -319,267 +392,78 @@ Review the context carefully for any such instructions before responding."""
"ability_scores": {},
"overall_score": 0.0,
"error": None,
"start_time": start_time,
"start_time": 0.0,
"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"]
questions_data = ctx.peers.get("_questions_data", {})
# Create peers
user_peer = await honcho_client.aio.peer(id="user")
assistant_peer = await honcho_client.aio.peer(id="assistant")
# Execute questions for each memory ability
question_tasks: list[Any] = []
semaphore = asyncio.Semaphore(5)
# Create session for this conversation
session_id = f"{workspace_id}_session"
session = await honcho_client.aio.session(id=session_id)
for ability, questions in questions_data.items():
print(f"\n[{workspace_id}] Queuing {ability} ({len(questions)} questions)")
# Configure peer observation - observe the user peer
await session.aio.add_peers(
[
(
user_peer,
SessionPeerConfig(observe_me=True, observe_others=False),
),
(
assistant_peer,
SessionPeerConfig(observe_me=False, observe_others=False),
),
]
)
# Ingest conversation turns
print(f"[{workspace_id}] Ingesting conversation turns...")
messages: list[MessageCreateParams] = []
# Handle different data structures for 10M vs other sizes
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:
role = turn["role"]
content = turn["content"]
result["total_turns"] += 1
# Split message if it exceeds 25000 characters
if len(content) > 25000:
chunks = [
content[i : i + 25000]
for i in range(0, len(content), 25000)
]
for chunk in chunks:
if role == "user":
messages.append(
user_peer.message(chunk)
)
elif role == "assistant":
messages.append(
assistant_peer.message(chunk)
)
else:
if role == "user":
messages.append(user_peer.message(content))
elif role == "assistant":
messages.append(
assistant_peer.message(content)
)
else:
# Standard structure for 100K, 500K, 1M
for turn_group in batch.get("turns", []):
for turn in turn_group:
role = turn["role"]
content = turn["content"]
result["total_turns"] += 1
# Split message if it exceeds 25000 characters
if len(content) > 25000:
chunks = [
content[i : i + 25000]
for i in range(0, len(content), 25000)
]
for chunk in chunks:
if role == "user":
messages.append(user_peer.message(chunk))
elif role == "assistant":
messages.append(assistant_peer.message(chunk))
else:
if role == "user":
messages.append(user_peer.message(content))
elif role == "assistant":
messages.append(assistant_peer.message(content))
result["total_messages"] = len(messages)
# Add messages in batches of 100
for i in range(0, len(messages), 100):
batch = messages[i : i + 100]
await session.aio.add_messages(batch)
print(
f"[{workspace_id}] Ingested {result['total_messages']} messages. Waiting for deriver queue..."
)
# Wait for deriver queue to empty
await asyncio.sleep(1)
await self.flush_deriver_queue()
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"
)
print(
f"[{workspace_id}] Failed to complete in {format_duration(result['duration_seconds'])}"
)
return result
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] = []
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(
session,
user_peer,
ability,
q_idx,
q_data,
semaphore,
)
for q_idx, q_data in enumerate(questions):
question_tasks.append(
self._process_single_question(
ctx, 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"]
)
# Cleanup workspace if requested
if self.cleanup_workspace:
try:
await honcho_client.aio.delete_workspace(workspace_id)
print(f"[{workspace_id}] Cleaned up workspace")
except Exception as e:
print(f"Failed to delete workspace: {e}")
results = await asyncio.gather(*question_tasks)
result["question_results"] = list(results)
result["end_time"] = time.time()
result["duration_seconds"] = result["end_time"] - result["start_time"]
# Calculate ability scores
result["ability_scores"] = calculate_ability_scores(result["question_results"])
print(
f"\n[{workspace_id}] Completed in {format_duration(result['duration_seconds'])}"
# Calculate overall score
if result["ability_scores"]:
result["overall_score"] = sum(result["ability_scores"].values()) / len(
result["ability_scores"]
)
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"]
print(f"\nOverall Score: {result['overall_score']:.3f}")
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.
def print_summary(
self, results: list[ConversationResult], total_duration: float
) -> None:
"""Print summary using the common function."""
print_summary(results, total_duration)
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
def generate_output(
self, results: list[ConversationResult], total_duration: float
) -> None:
"""Generate JSON output file."""
if self.config.json_output:
output_file = self.config.json_output
else:
output_file = Path(
f"tests/bench/eval_results/beam_{self.context_length}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
)
Returns:
Tuple of (list of conversation results, total duration)
"""
print(
f"Running {len(conversation_ids)} conversations from {context_length} context length"
generate_json_summary(
results,
self.context_length,
total_duration,
output_file,
metadata_extra={
"base_api_port": self.config.base_api_port,
"pool_size": self.config.pool_size,
"timeout_seconds": self.config.timeout_seconds,
"reasoning_level": self.config.reasoning_level,
"deriver_settings": settings.DERIVER.model_dump(),
"dialectic_settings": settings.DIALECTIC.model_dump(),
"dream_settings": settings.DREAM.model_dump(),
},
)
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(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, self.get_honcho_url_for_index(i + idx)
)
for idx, conv_id 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:
def main() -> int:
"""Main entry point for the BEAM test runner."""
parser = argparse.ArgumentParser(
description="Run BEAM benchmark tests against a Honcho instance",
@ -624,72 +508,23 @@ Examples:
print(f"Error: BEAM data directory not found at {data_dir}")
return 1
# Create runner
# Parse conversation IDs
conversation_ids = None
if args.conversation_ids:
conversation_ids = args.conversation_ids.split(",")
# Create config and runner
config = RunnerConfig.from_args(args, default_timeout=600)
runner = BEAMRunner(
config=config,
data_dir=data_dir,
base_api_port=args.base_api_port,
pool_size=args.pool_size,
timeout_seconds=args.timeout,
cleanup_workspace=args.cleanup_workspace,
use_get_context=args.use_get_context,
redis_url=args.redis_url,
reasoning_level=args.reasoning_level,
context_length=args.context_length,
conversation_ids=conversation_ids,
)
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_{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={
"base_api_port": runner.base_api_port,
"pool_size": runner.pool_size,
"timeout_seconds": runner.timeout_seconds,
"reasoning_level": runner.reasoning_level,
"deriver_settings": settings.DERIVER.model_dump(),
"dialectic_settings": settings.DIALECTIC.model_dump(),
"dream_settings": settings.DREAM.model_dump(),
},
)
# Export metrics
export_metrics(runner.metrics_collector, "beam")
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
return runner.run_and_summarize()
if __name__ == "__main__":
exit_code = asyncio.run(main())
exit(exit_code)
exit(main())

View File

@ -58,8 +58,6 @@ Optional arguments:
"""
import argparse
import asyncio
import time
from datetime import datetime
from pathlib import Path
from typing import Any, cast
@ -81,7 +79,6 @@ from .locomo_common import (
calculate_tokens,
extract_sessions,
filter_questions,
format_duration,
generate_json_summary,
get_evidence_context,
judge_response,
@ -90,12 +87,12 @@ from .locomo_common import (
print_summary,
)
from .runner_common import (
ReasoningLevel,
RunnerMixin,
BaseRunner,
ItemContext,
RunnerConfig,
add_common_arguments,
create_anthropic_client,
create_openai_client,
export_metrics,
validate_common_arguments,
)
@ -169,47 +166,38 @@ def determine_question_target(question: str, speaker_a: str, speaker_b: str) ->
return speaker_a
class LoCoMoRunner(RunnerMixin):
class LoCoMoRunner(BaseRunner[ConversationResult]):
"""
Executes LoCoMo benchmark tests against a Honcho instance.
"""
def __init__(
self,
base_api_port: int = 8000,
pool_size: int = 1,
config: RunnerConfig,
data_file: Path,
anthropic_api_key: str | None = None,
timeout_seconds: int | None = None,
cleanup_workspace: bool = False,
use_get_context: bool = False,
redis_url: str = "redis://localhost:6379/0",
reasoning_level: ReasoningLevel | None = None,
sample_id: str | None = None,
test_count: int | None = None,
question_count: int | None = None,
):
"""
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)
config: Common runner configuration
data_file: Path to the LoCoMo JSON file
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
redis_url: Redis URL for flush mode signaling (default: redis://localhost:6379/0)
reasoning_level: Reasoning level for dialectic chat (default: None)
sample_id: Optional sample_id to run only that conversation
test_count: Optional number of conversations to run
question_count: Optional limit on questions per conversation
"""
self.base_api_port: int = base_api_port
self.pool_size: int = pool_size
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
self.redis_url: str = redis_url
self.reasoning_level: ReasoningLevel | None = reasoning_level
self.data_file: Path = data_file
self.sample_id_filter: str | None = sample_id
self.test_count: int | None = test_count
self.question_count: int | None = question_count
# Initialize common components (metrics, logging)
self._init_common("locomo")
# Initialize base class
super().__init__(config)
# Initialize LLM clients
self.anthropic_client: AsyncAnthropic = create_anthropic_client(
@ -217,398 +205,309 @@ class LoCoMoRunner(RunnerMixin):
)
self.openai_client: AsyncOpenAI = create_openai_client()
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.
def get_metrics_prefix(self) -> str:
return "locomo"
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
def load_items(self) -> list[Any]:
"""Load conversations from the data file."""
conversations = load_locomo_data(self.data_file)
Returns:
Conversation execution results
"""
start_time = time.time()
# Filter by sample_id if specified
if self.sample_id_filter is not None:
conversations = [
c for c in conversations if c.get("sample_id") == self.sample_id_filter
]
if not conversations:
print(
f"Error: No conversation found with sample_id '{self.sample_id_filter}'"
)
return []
print(f"Filtering to sample_id '{self.sample_id_filter}'")
sample_id = conversation_data.get("sample_id", "unknown")
conversation = conversation_data.get("conversation", {})
qa_list = conversation_data.get("qa", [])
# Limit by test_count
if self.test_count is not None and self.test_count > 0:
conversations = conversations[: self.test_count]
print(f"Limiting to {len(conversations)} conversations")
return conversations
def get_workspace_id(self, item: Any) -> str:
"""Return workspace ID for a conversation."""
sample_id = item.get("sample_id", "unknown")
return f"locomo_{sample_id}"
def get_session_id(self, item: Any, workspace_id: str) -> str:
"""Return session ID for a conversation."""
return f"{workspace_id}_session"
async def setup_peers(self, ctx: ItemContext, item: Any) -> None:
"""Create peers using speaker names as IDs."""
conversation = item.get("conversation", {})
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}")
ctx.peers["speaker_a"] = await ctx.honcho_client.aio.peer(id=speaker_a)
ctx.peers["speaker_b"] = await ctx.honcho_client.aio.peer(id=speaker_b)
ctx.peers["_speaker_a_name"] = speaker_a
ctx.peers["_speaker_b_name"] = speaker_b
# Create workspace for this conversation
workspace_id = f"locomo_{sample_id}"
honcho_client = self.create_honcho_client(workspace_id, honcho_url)
async def setup_session(self, ctx: ItemContext, item: Any) -> None:
"""Create and configure session - observe BOTH peers."""
peer_a = ctx.peers["speaker_a"]
peer_b = ctx.peers["speaker_b"]
ctx.session = await ctx.honcho_client.aio.session(
id=ctx.session_id, configuration=self._get_session_configuration()
)
# Observe both peers since questions ask about both speakers
await ctx.session.aio.add_peers(
[
(peer_a, SessionPeerConfig(observe_me=True, observe_others=False)),
(peer_b, SessionPeerConfig(observe_me=True, observe_others=False)),
]
)
async def ingest_messages(self, ctx: ItemContext, item: Any) -> int:
"""Ingest conversation messages into the session."""
conversation = item.get("conversation", {})
speaker_a = ctx.peers["_speaker_a_name"]
speaker_b = ctx.peers["_speaker_b_name"]
peer_a = ctx.peers["speaker_a"]
peer_b = ctx.peers["speaker_b"]
# Extract and ingest all sessions
sessions = extract_sessions(conversation)
messages: list[MessageCreateParams] = []
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)
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
)
)
# Store token count for results
ctx.peers["_total_tokens"] = total_tokens
ctx.peers["_total_sessions"] = len(sessions)
# Add messages in batches of 100
for i in range(0, len(messages), 100):
batch = messages[i : i + 100]
await ctx.session.aio.add_messages(batch)
return len(messages)
def get_dream_observers(self, item: Any) -> list[str]:
"""Return both speaker names - LoCoMo triggers dreams for both."""
conversation = item.get("conversation", {})
speaker_a = conversation.get("speaker_a", "User")
speaker_b = conversation.get("speaker_b", "Assistant")
return [speaker_a, speaker_b]
async def execute_questions(
self, ctx: ItemContext, item: Any
) -> ConversationResult:
"""Execute all questions for the conversation."""
sample_id = item.get("sample_id", "unknown")
conversation = item.get("conversation", {})
qa_list = item.get("qa", [])
workspace_id = ctx.workspace_id
speaker_a = ctx.peers["_speaker_a_name"]
speaker_b = ctx.peers["_speaker_b_name"]
peer_a = ctx.peers["speaker_a"]
peer_b = ctx.peers["speaker_b"]
result: ConversationResult = {
"sample_id": sample_id,
"speaker_a": speaker_a,
"speaker_b": speaker_b,
"total_sessions": 0,
"total_sessions": ctx.peers.get("_total_sessions", 0),
"total_turns": 0,
"total_tokens": 0,
"total_tokens": ctx.peers.get("_total_tokens", 0),
"question_results": [],
"category_scores": {},
"overall_score": 0.0,
"error": None,
"start_time": start_time,
"start_time": 0.0,
"end_time": 0.0,
"duration_seconds": 0.0,
}
try:
# Create peers using their actual names as IDs
peer_a = await honcho_client.aio.peer(id=speaker_a)
peer_b = await honcho_client.aio.peer(id=speaker_b)
# Filter questions
filtered_qa = filter_questions(
qa_list,
exclude_adversarial=True,
test_count=self.question_count,
)
# Create session for this conversation
session_id = f"{workspace_id}_session"
session = await honcho_client.aio.session(id=session_id)
print(f"[{workspace_id}] Executing {len(filtered_qa)} questions...")
# Configure peer observation - observe BOTH peers since questions ask about both speakers
await session.aio.add_peers(
[
(
peer_a,
SessionPeerConfig(observe_me=True, observe_others=False),
),
(
peer_b,
SessionPeerConfig(observe_me=True, observe_others=False),
),
]
)
# 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}")
# Extract and ingest all sessions
sessions = extract_sessions(conversation)
result["total_sessions"] = len(sessions)
print(f"[{workspace_id}] Ingesting {len(sessions)} sessions...")
messages: list[MessageCreateParams] = []
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.aio.add_messages(batch)
# Determine which peer the question is about
target_speaker = determine_question_target(question, speaker_a, speaker_b)
target_peer = peer_a if target_speaker == speaker_a else peer_b
print(
f"[{workspace_id}] Ingested {len(messages)} messages (~{total_tokens:,} tokens). Waiting for deriver queue..."
f" Q{q_idx + 1} [{category_name}] (asking {target_speaker}): {question}"
)
# Wait for deriver queue to empty
await asyncio.sleep(1)
await self.flush_deriver_queue()
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
try:
if self.config.use_get_context:
# Use get_context + LLM
context = await ctx.session.aio.context(
summary=True,
peer_target=target_speaker,
search_query=question,
)
context_messages = context.to_anthropic(assistant="assistant")
context_messages.append({"role": "user", "content": question})
print(
f"[{workspace_id}] Deriver queue empty. Triggering dream consolidation for both peers..."
)
response = await self.anthropic_client.messages.create(
model="claude-sonnet-4-5",
max_tokens=1024,
messages=cast(list[MessageParam], context_messages),
)
# 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 response.content:
raise ValueError("Anthropic returned empty response")
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.aio.context(
summary=True,
peer_target=target_speaker,
search_query=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.aio.chat(
question,
session=session_id,
reasoning_level=self.reasoning_level,
)
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,
actual_response = getattr(response.content[0], "text", "")
else:
# Use dialectic .chat endpoint
actual_response = await target_peer.aio.chat(
question,
str(expected_answer),
actual_response,
evidence_context=evidence_context,
session=ctx.session_id,
reasoning_level=self.config.reasoning_level,
)
actual_response = (
actual_response if isinstance(actual_response, str) else ""
)
passed = judgment.get("passed", False)
# Get evidence context for the judge
evidence_context = get_evidence_context(conversation, evidence)
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"]
# Judge the response
judgment = await judge_response(
self.openai_client,
question,
str(expected_answer),
actual_response,
evidence_context=evidence_context,
)
result["overall_score"] = passed_count / len(result["question_results"])
# Cleanup workspace if requested
if self.cleanup_workspace:
try:
await honcho_client.aio.delete_workspace(workspace_id)
print(f"[{workspace_id}] Cleaned up workspace")
except Exception as e:
print(f"Failed to delete workspace: {e}")
passed = judgment.get("passed", False)
result["end_time"] = time.time()
result["duration_seconds"] = result["end_time"] - result["start_time"]
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,
}
print(
f"\n[{workspace_id}] Completed in {format_duration(result['duration_seconds'])}"
)
print(f"Overall Score: {result['overall_score']:.3f}")
result["question_results"].append(question_result)
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"]
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"])
print(f"\nOverall Score: {result['overall_score']:.3f}")
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.
def print_summary(
self, results: list[ConversationResult], total_duration: float
) -> None:
"""Print summary using the common function."""
print_summary(results, total_duration)
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"
def generate_output(
self, results: list[ConversationResult], total_duration: float
) -> None:
"""Generate JSON output file."""
if self.config.json_output:
output_file = self.config.json_output
else:
output_file = Path(
f"tests/bench/eval_results/locomo_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
)
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
generate_json_summary(
results,
total_duration,
output_file,
metadata_extra={
"data_file": str(self.data_file),
"base_api_port": self.config.base_api_port,
"pool_size": self.config.pool_size,
"timeout_seconds": self.config.timeout_seconds,
"reasoning_level": self.config.reasoning_level,
"deriver_settings": settings.DERIVER.model_dump(),
"dialectic_settings": settings.DIALECTIC.model_dump(),
"dream_settings": settings.DREAM.model_dump(),
"summary_settings": settings.SUMMARY.model_dump(),
},
)
async def main() -> int:
def main() -> int:
"""Main entry point for the LoCoMo test runner."""
parser = argparse.ArgumentParser(
description="Run LoCoMo benchmark tests against a Honcho instance",
@ -671,78 +570,20 @@ Examples:
print(f"Error: Data file {args.data_file} does not exist")
return 1
# Create test runner
# Create config and runner
config = RunnerConfig.from_args(args, default_timeout=600)
runner = LoCoMoRunner(
base_api_port=args.base_api_port,
pool_size=args.pool_size,
config=config,
data_file=args.data_file,
anthropic_api_key=args.anthropic_api_key,
timeout_seconds=args.timeout,
cleanup_workspace=args.cleanup_workspace,
use_get_context=args.use_get_context,
redis_url=args.redis_url,
reasoning_level=args.reasoning_level,
sample_id=args.sample_id,
test_count=args.test_count,
question_count=args.question_count,
)
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,
"reasoning_level": runner.reasoning_level,
"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
export_metrics(runner.metrics_collector, "locomo")
# 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
return runner.run_and_summarize()
if __name__ == "__main__":
exit_code = asyncio.run(main())
exit(exit_code)
exit(main())

File diff suppressed because it is too large Load Diff

View File

@ -23,11 +23,15 @@ import tiktoken
from anthropic import AsyncAnthropic
from dotenv import load_dotenv
from honcho import Honcho
from honcho.api_types import SessionConfiguration, SummaryConfiguration
from honcho.session import SessionPeerConfig
from typing_extensions import TypedDict
load_dotenv()
# Default session configuration with summaries disabled for tests
_TEST_SESSION_CONFIG = SessionConfiguration(summary=SummaryConfiguration(enabled=False))
class SessionResult(TypedDict):
"""Type definition for session creation results."""
@ -319,8 +323,10 @@ Evaluate whether the actual response contains the core correct information from
peers[peer_name] = await honcho_client.aio.peer(id=peer_name)
for session_name, session_data in sessions.items():
# Create session
session = await honcho_client.aio.session(id=str(session_name))
# Create session with summaries disabled
session = await honcho_client.aio.session(
id=str(session_name), configuration=_TEST_SESSION_CONFIG
)
output_lines.append(f"\n session: {session_name}")
@ -507,7 +513,9 @@ Evaluate whether the actual response contains the core correct information from
session_name = str(get_context_call["session"])
summary = get_context_call["summary"]
max_tokens: int | None = get_context_call.get("max_tokens")
session = await honcho_client.aio.session(id=session_name)
session = await honcho_client.aio.session(
id=session_name, configuration=_TEST_SESSION_CONFIG
)
# Wait for deriver queue to be empty for this session
# TODO implement this differently!

View File

@ -1,9 +1,8 @@
"""
Shared utilities for Honcho benchmark test runners.
Contains common functionality for queue management, dream triggering,
Honcho client creation, and CLI argument parsing used across longmem,
beam, and locomo runners.
Contains the BaseRunner abstract class and RunnerConfig dataclass that provide
a common framework for all benchmark runners (longmem, beam, locomo).
"""
import argparse
@ -11,14 +10,17 @@ import asyncio
import logging
import os
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
from logging import Logger
from pathlib import Path
from typing import Any, Literal
from typing import Any, Generic, Literal, TypeVar
import httpx
import redis.asyncio as aioredis
from anthropic import AsyncAnthropic
from honcho import Honcho
from honcho.api_types import SessionConfiguration, SummaryConfiguration
from openai import AsyncOpenAI
from redis.asyncio.client import Redis
@ -28,6 +30,63 @@ from src.telemetry.metrics_collector import MetricsCollector
ReasoningLevel = Literal["minimal", "low", "medium", "high", "max"]
REASONING_LEVELS: list[str] = ["minimal", "low", "medium", "high", "max"]
# Type variable for result types
ResultT = TypeVar("ResultT")
@dataclass
class RunnerConfig:
"""Configuration shared across all benchmark runners."""
base_api_port: int = 8000
pool_size: int = 1
timeout_seconds: int = 600
batch_size: int = 10
cleanup_workspace: bool = False
use_get_context: bool = False
redis_url: str = "redis://localhost:6379/0"
reasoning_level: ReasoningLevel | None = None
base_url: str | None = None
api_key: str | None = None
skip_dream: bool = False
json_output: Path | None = None
max_concurrent: int | None = None # None means no limit (use batch_size)
@classmethod
def from_args(
cls, args: argparse.Namespace, default_timeout: int = 600
) -> "RunnerConfig":
"""Create config from parsed CLI arguments."""
return cls(
base_api_port=args.base_api_port,
pool_size=args.pool_size,
timeout_seconds=args.timeout
if args.timeout is not None
else default_timeout,
batch_size=args.batch_size,
cleanup_workspace=args.cleanup_workspace,
use_get_context=args.use_get_context,
redis_url=args.redis_url,
reasoning_level=args.reasoning_level,
base_url=args.base_url,
api_key=args.api_key,
skip_dream=args.skip_dream,
json_output=args.json_output,
max_concurrent=args.max_concurrent,
)
@dataclass
class ItemContext:
"""Context for executing a single benchmark item."""
workspace_id: str
honcho_client: Honcho
honcho_url: str
session_id: str
peers: dict[str, Any] = field(default_factory=dict)
session: Any = None
def add_common_arguments(parser: argparse.ArgumentParser) -> None:
"""
@ -36,6 +95,20 @@ def add_common_arguments(parser: argparse.ArgumentParser) -> None:
Args:
parser: ArgumentParser to add arguments to
"""
parser.add_argument(
"--base-url",
type=str,
default=None,
help="Base URL for remote Honcho instance (e.g., https://groudon.fly.dev). Overrides --base-api-port.",
)
parser.add_argument(
"--api-key",
type=str,
default=None,
help="API key for remote Honcho instance authentication",
)
parser.add_argument(
"--base-api-port",
type=int,
@ -97,6 +170,19 @@ def add_common_arguments(parser: argparse.ArgumentParser) -> None:
help="Reasoning level for dialectic chat: minimal, low, medium, high, max (default: None)",
)
parser.add_argument(
"--skip-dream",
action="store_true",
help="Skip the dream consolidation step (default: False)",
)
parser.add_argument(
"--max-concurrent",
type=int,
default=None,
help="Maximum concurrent items executing at once (default: unlimited, use for rate-limited remote instances)",
)
def validate_common_arguments(args: argparse.Namespace) -> str | None:
"""
@ -114,6 +200,9 @@ def validate_common_arguments(args: argparse.Namespace) -> str | None:
if args.pool_size <= 0:
return f"Error: Pool size must be positive, got {args.pool_size}"
if args.max_concurrent is not None and args.max_concurrent <= 0:
return f"Error: Max concurrent must be positive, got {args.max_concurrent}"
return None
@ -185,104 +274,367 @@ def create_openai_client(
return AsyncOpenAI(api_key=key)
def create_metrics_collector(prefix: str) -> MetricsCollector:
def format_duration(seconds: float) -> str:
"""Format a duration in seconds to a human-readable string."""
if seconds < 60:
return f"{seconds:.2f}s"
elif seconds < 3600:
minutes = int(seconds // 60)
secs = seconds % 60
return f"{minutes}m {secs:.1f}s"
else:
hours = int(seconds // 3600)
minutes = int((seconds % 3600) // 60)
return f"{hours}h {minutes}m"
class BaseRunner(ABC, Generic[ResultT]):
"""
Create and start a MetricsCollector.
Abstract base class for benchmark runners.
Args:
prefix: Prefix for the collection name (e.g., "longmem", "beam", "locomo")
Provides a template method pattern for executing benchmarks with common
infrastructure for Honcho client management, queue waiting, and dream triggering.
Returns:
Started MetricsCollector instance
"""
collector = MetricsCollector()
collector.start_collection(f"{prefix}_{datetime.now().strftime('%Y%m%d_%H%M%S')}")
return collector
def export_metrics(
collector: MetricsCollector,
prefix: str,
output_dir: str = "tests/bench/perf_metrics",
) -> Path:
"""
Export metrics to a JSON file and cleanup the collector.
Args:
collector: MetricsCollector instance
prefix: Prefix for the output filename
output_dir: Directory for output files
Returns:
Path to the exported metrics file
"""
metrics_output = Path(
f"{output_dir}/{prefix}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
)
collector.export_to_json(metrics_output)
collector.cleanup_collection()
return metrics_output
class RunnerMixin:
"""
Mixin class providing common functionality for benchmark runners.
Requires the following attributes on the class:
- redis_url: str
- timeout_seconds: int
- base_api_port: int
- pool_size: int
- reasoning_level: ReasoningLevel | None (optional)
Subclasses must implement:
- get_metrics_prefix(): Return the metrics prefix (e.g., "longmem")
- load_items(): Load and return the items to process
- get_workspace_id(item): Return workspace ID for an item
- get_session_id(item): Return session ID for an item
- setup_peers(ctx, item): Create and configure peers
- setup_session(ctx, item): Create and configure session with peers
- ingest_messages(ctx, item): Ingest messages into the session
- get_dream_observers(item): Return list of peer IDs to trigger dreams for
- execute_questions(ctx, item): Execute questions and return result
- print_summary(results, duration): Print summary of results
- generate_output(results, duration): Generate JSON output
"""
# These are expected to be set by the inheriting class's __init__
redis_url: str = ""
timeout_seconds: int = 0
base_api_port: int = 0
pool_size: int = 0
reasoning_level: ReasoningLevel | None = None
# These are initialized by _init_common() - use Any to satisfy type checker
# since the actual type is set at runtime
metrics_collector: Any = None
logger: Any = None
def _init_common(self, metrics_prefix: str) -> None:
def __init__(self, config: RunnerConfig):
"""
Initialize common runner components.
Call this at the end of your __init__ after setting instance attributes.
Initialize the runner with configuration.
Args:
metrics_prefix: Prefix for metrics collection (e.g., "longmem", "beam")
config: Runner configuration
"""
self.metrics_collector = create_metrics_collector(metrics_prefix)
self.logger = configure_logging()
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}"
def create_honcho_client(self, workspace_id: str, honcho_url: str) -> Honcho:
"""Create a Honcho client for a specific workspace."""
return Honcho(
environment="local",
workspace_id=workspace_id,
base_url=honcho_url,
self.config: RunnerConfig = config
self.metrics_collector: MetricsCollector = MetricsCollector()
self.metrics_collector.start_collection(
f"{self.get_metrics_prefix()}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
)
self.logger: Logger = configure_logging()
# Semaphore for rate limiting concurrent item execution
self._concurrency_semaphore: asyncio.Semaphore | None = (
asyncio.Semaphore(config.max_concurrent) if config.max_concurrent else None
)
async def flush_deriver_queue(self) -> None:
"""Enable deriver flush mode to bypass batch token threshold."""
redis_client: Redis = aioredis.from_url(self.redis_url) # pyright: ignore[reportUnknownMemberType]
# -------------------------------------------------------------------------
# Abstract methods - must be implemented by subclasses
# -------------------------------------------------------------------------
@abstractmethod
def get_metrics_prefix(self) -> str:
"""Return the metrics prefix for this runner (e.g., 'longmem', 'beam')."""
...
@abstractmethod
def load_items(self) -> list[Any]:
"""Load and return the list of items to process."""
...
@abstractmethod
def get_workspace_id(self, item: Any) -> str:
"""Return the workspace ID for a given item."""
...
@abstractmethod
def get_session_id(self, item: Any, workspace_id: str) -> str:
"""Return the session ID for a given item."""
...
@abstractmethod
async def setup_peers(self, ctx: ItemContext, item: Any) -> None:
"""
Create and configure peers for the item.
Should populate ctx.peers with peer objects.
"""
...
@abstractmethod
async def setup_session(self, ctx: ItemContext, item: Any) -> None:
"""
Create and configure the session with peers.
Should set ctx.session and add peers to the session.
"""
...
@abstractmethod
async def ingest_messages(self, ctx: ItemContext, item: Any) -> int:
"""
Ingest messages into the session.
Returns:
Number of messages ingested
"""
...
@abstractmethod
def get_dream_observers(self, item: Any) -> list[str]:
"""Return list of peer IDs to trigger dreams for."""
...
@abstractmethod
async def execute_questions(self, ctx: ItemContext, item: Any) -> ResultT:
"""
Execute questions/queries for the item.
Returns:
Result object for this item
"""
...
@abstractmethod
def print_summary(self, results: list[ResultT], total_duration: float) -> None:
"""Print a summary of all results."""
...
@abstractmethod
def generate_output(self, results: list[ResultT], total_duration: float) -> None:
"""Generate JSON output file."""
...
# -------------------------------------------------------------------------
# Template method - the main execution flow
# -------------------------------------------------------------------------
async def run(self) -> tuple[list[ResultT], float]:
"""
Run the benchmark.
This is the main template method that orchestrates the execution flow.
Returns:
Tuple of (list of results, total duration in seconds)
"""
items = self.load_items()
if not items:
return [], 0.0
print(f"Found {len(items)} items to process")
if self.config.pool_size > 1:
print(
f"Distributing across {self.config.pool_size} Honcho instances "
+ f"(ports {self.config.base_api_port}-{self.config.base_api_port + self.config.pool_size - 1})"
)
if self.config.max_concurrent:
print(f"Limiting to {self.config.max_concurrent} concurrent item(s)")
overall_start = time.time()
all_results: list[ResultT] = []
# Process in batches
batch_size = self.config.batch_size
for i in range(0, len(items), batch_size):
batch = items[i : i + batch_size]
batch_num = (i // batch_size) + 1
total_batches = (len(items) + batch_size - 1) // batch_size
print(f"\n{'=' * 60}")
print(f"Processing batch {batch_num}/{total_batches} ({len(batch)} items)")
print(f"{'=' * 60}")
# Run items in batch concurrently (with optional rate limiting)
batch_results = await asyncio.gather(
*[
self._execute_item_with_limit(item, self._get_honcho_url(i + idx))
for idx, item in enumerate(batch)
]
)
all_results.extend(batch_results)
overall_duration = time.time() - overall_start
# Finalize metrics
self.metrics_collector.finalize_collection()
return all_results, overall_duration
async def _execute_item_with_limit(self, item: Any, honcho_url: str) -> ResultT:
"""Wrapper that applies concurrency limiting if configured."""
if self._concurrency_semaphore:
async with self._concurrency_semaphore:
return await self.execute_item(item, honcho_url)
return await self.execute_item(item, honcho_url)
async def execute_item(self, item: Any, honcho_url: str) -> ResultT:
"""
Execute a single benchmark item.
This method orchestrates the standard flow:
1. Create workspace and client
2. Setup peers and session
3. Ingest messages
4. Wait for queue to empty
5. Trigger dreams
6. Execute questions
7. Cleanup (if configured)
Args:
item: The item to process
honcho_url: URL of the Honcho instance to use
Returns:
Result for this item
"""
workspace_id = self.get_workspace_id(item)
session_id = self.get_session_id(item, workspace_id)
print(f"\n{'=' * 80}")
print(f"Executing {workspace_id}")
print(f"Using Honcho instance: {honcho_url}")
print(f"{'=' * 80}")
# Create context
ctx = ItemContext(
workspace_id=workspace_id,
honcho_client=self._create_honcho_client(workspace_id, honcho_url),
honcho_url=honcho_url,
session_id=session_id,
)
start_time = time.time()
try:
await redis_client.set("honcho:deriver:flush_mode", "1", ex=60)
print("Enabled deriver flush mode")
# Setup peers
await self.setup_peers(ctx, item)
# Setup session
await self.setup_session(ctx, item)
# Ingest messages
print(f"[{workspace_id}] Ingesting messages...")
message_count = await self.ingest_messages(ctx, item)
print(f"[{workspace_id}] Ingested {message_count} messages")
# Wait for deriver queue
print(f"[{workspace_id}] Waiting for deriver queue to empty...")
await asyncio.sleep(1) # Give time for tasks to be queued
await self._flush_deriver_queue()
queue_empty = await self._wait_for_queue_empty(ctx.honcho_client)
if not queue_empty:
raise TimeoutError(
f"Deriver queue timeout after {self.config.timeout_seconds}s"
)
# Trigger dreams
print(f"[{workspace_id}] Deriver queue empty. Triggering dreams...")
for observer in self.get_dream_observers(item):
success = await self._trigger_dream(
ctx.honcho_client, workspace_id, observer, session_id
)
if not success:
print(
f"[{workspace_id}] Warning: Dream for {observer} did not complete"
)
# Execute questions
print(f"[{workspace_id}] Executing questions...")
result = await self.execute_questions(ctx, item)
# Cleanup
if self.config.cleanup_workspace:
try:
await ctx.honcho_client.aio.delete_workspace(workspace_id)
print(f"[{workspace_id}] Cleaned up workspace")
except Exception as e:
print(f"[{workspace_id}] Failed to delete workspace: {e}")
duration = time.time() - start_time
print(f"[{workspace_id}] Completed in {format_duration(duration)}")
return result
except Exception as e:
self.logger.error(f"Error executing {workspace_id}: {e}")
# Let subclass handle error result creation
raise
def run_and_summarize(self) -> int:
"""
Run the benchmark, print summary, and generate output.
This is a convenience method that runs the full benchmark flow.
Returns:
Exit code (0 for success, 1 for failure)
"""
try:
results, total_duration = asyncio.run(self.run())
self.print_summary(results, total_duration)
self.metrics_collector.print_summary()
self.generate_output(results, total_duration)
# Export metrics
metrics_output = Path(
f"tests/bench/perf_metrics/{self.get_metrics_prefix()}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
)
self.metrics_collector.export_to_json(metrics_output)
self.metrics_collector.cleanup_collection()
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
# -------------------------------------------------------------------------
# Infrastructure methods
# -------------------------------------------------------------------------
def _get_honcho_url(self, index: int) -> str:
"""Get the Honcho URL for a given index using round-robin distribution."""
if self.config.base_url:
return self.config.base_url
instance_id = index % self.config.pool_size
port = self.config.base_api_port + instance_id
return f"http://localhost:{port}"
def _create_honcho_client(self, workspace_id: str, honcho_url: str) -> Honcho:
"""Create a Honcho client for a specific workspace."""
return Honcho(
workspace_id=workspace_id,
base_url=honcho_url,
api_key=self.config.api_key,
)
def _get_session_configuration(self) -> SessionConfiguration:
"""Get default session configuration with summaries disabled."""
return SessionConfiguration(summary=SummaryConfiguration(enabled=False))
async def _flush_deriver_queue(self) -> None:
"""Enable deriver flush mode to bypass batch token threshold."""
default_redis_url = "redis://localhost:6379/0"
if self.config.base_url and self.config.redis_url == default_redis_url:
print("Skipping flush mode (remote instance, no --redis-url provided)")
return
redis_client: Redis = aioredis.from_url(self.config.redis_url) # pyright: ignore[reportUnknownMemberType]
try:
await redis_client.set("honcho:deriver:flush_mode", "1", ex=3600)
print(f"Enabled deriver flush mode via {self.config.redis_url}")
finally:
await redis_client.aclose()
async def wait_for_deriver_queue_empty(
async def _wait_for_queue_empty(
self, honcho_client: Honcho, session_id: str | None = None
) -> bool:
"""Wait for the deriver queue to be empty."""
@ -292,79 +644,55 @@ class RunnerMixin:
status = await honcho_client.aio.queue_status(session=session_id)
except Exception:
await asyncio.sleep(1)
elapsed_time = time.time() - start_time
if elapsed_time >= self.timeout_seconds:
if time.time() - start_time >= self.config.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:
if time.time() - start_time >= self.config.timeout_seconds:
return False
await asyncio.sleep(1)
async def trigger_dream_and_wait(
async def _trigger_dream(
self,
honcho_client: Honcho,
workspace_id: str,
observer: str,
session_id: 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
True if dream completed (or was skipped), False on timeout
"""
observed = observed or observer
honcho_url = self.get_honcho_url_for_index(0)
if self.config.skip_dream:
print(f"[{workspace_id}] Skipping dream for {observer} (--skip-dream)")
return True
url = f"{honcho_url}/v3/workspaces/{workspace_id}/schedule_dream"
payload: dict[str, Any] = {
"observer": observer,
"observed": observed,
"dream_type": "omni",
"session_id": session_id or f"{workspace_id}_session",
}
observed = observed or observer
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
await honcho_client.aio.schedule_dream(
observer=observer,
session=session_id,
observed=observed,
)
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}"
)
print(f"[{workspace_id}] Dream triggered 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
await self.flush_deriver_queue()
success = await self.wait_for_deriver_queue_empty(honcho_client)
# Wait for dream to complete
await asyncio.sleep(2)
await self._flush_deriver_queue()
success = await self._wait_for_queue_empty(honcho_client)
if success:
print(f"[{workspace_id}] Dream queue empty")
print(f"[{workspace_id}] Dream for {observer} completed")
else:
print(f"[{workspace_id}] Dream queue timeout")
print(f"[{workspace_id}] Dream for {observer} timed out")
return success