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:
commit
5371463f39
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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!
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue