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

749 lines
26 KiB
Python
Raw Blame History

"""
Honcho LoCoMo Benchmark Test Runner
A script that executes LoCoMo benchmark tests against a running Honcho instance.
This script:
1. Loads LoCoMo conversation data from JSON files
2. Creates a workspace for each conversation sample
3. Ingests conversation sessions as messages between two peers
4. Waits for the deriver queue to process everything
5. Triggers a dream for memory consolidation
6. Executes questions and judges responses using an LLM
## LoCoMo Overview
LoCoMo evaluates very long-term conversational memory across five question categories:
1. Single-hop - Direct factual recall from conversations
2. Multi-hop - Reasoning across multiple pieces of information
3. Temporal - Understanding time-based relationships and sequences
4. Commonsense/World knowledge - Applying broader contextual understanding
5. Adversarial - Questions that cannot be answered (filtered out by default)
Reference: https://github.com/snap-research/locomo
Paper: https://arxiv.org/abs/2402.17753
## To use
0. Set up env:
```
uv sync
source .venv/bin/activate
```
NOTE: you may create a .env file in this directory to customize honcho config.
1. Run the test harness:
```
python -m tests.bench.harness
```
2. Run this file with the LoCoMo dataset:
```
python -m tests.bench.locomo --data-file tests/bench/locomo_data/locomo10.json
```
Optional arguments:
```
--anthropic-api-key: Anthropic API key for response judging (can be set in .env as LLM_ANTHROPIC_API_KEY)
--timeout: Timeout for deriver queue to empty in seconds (default: 10 minutes)
--base-api-port: Base port for Honcho API instances (default: 8000)
--pool-size: Number of Honcho instances in the pool (default: 1)
--batch-size: Number of conversations to run concurrently in each batch (default: 1)
--json-output: Path to write JSON summary results for analytics
--cleanup-workspace: Delete workspace after executing each conversation (default: False)
--use-get-context: Use get_context + judge LLM instead of dialectic .chat endpoint (default: False)
--sample-id: Run only the conversation with this sample_id (skips all others)
--test-count: Number of conversations to run (default: all)
--question-count: Number of questions per conversation to run (default: all)
```
"""
import argparse
import asyncio
import time
from datetime import datetime
from pathlib import Path
from typing import Any, cast
from anthropic import AsyncAnthropic
from anthropic.types import MessageParam
from dotenv import load_dotenv
from honcho.api_types import MessageCreateParams
from honcho.session import SessionPeerConfig
from openai import AsyncOpenAI
from src.config import settings
from .locomo_common import (
CATEGORY_NAMES,
ConversationResult,
QuestionResult,
calculate_category_scores,
calculate_tokens,
extract_sessions,
filter_questions,
format_duration,
generate_json_summary,
get_evidence_context,
judge_response,
load_locomo_data,
parse_locomo_date,
print_summary,
)
from .runner_common import (
ReasoningLevel,
RunnerMixin,
add_common_arguments,
create_anthropic_client,
create_openai_client,
export_metrics,
validate_common_arguments,
)
# Load .env from bench directory
bench_dir = Path(__file__).parent
load_dotenv(bench_dir / ".env")
def format_message_with_image(msg: dict[str, Any]) -> tuple[str, dict[str, Any] | None]:
"""
Format a LoCoMo message with optional image caption appended.
Args:
msg: LoCoMo message dict with 'text', optional 'img_url', 'blip_caption', 'query'
Returns:
Tuple of (formatted_content, metadata_dict or None)
"""
text = msg.get("text", "")
blip_caption = msg.get("blip_caption")
img_urls = msg.get("img_url", [])
query = msg.get("query")
# Append caption to content so deriver can see it
content = f"{text}\n\n[Image shared: {blip_caption}]" if blip_caption else text
# Build metadata if image data exists
metadata: dict[str, Any] | None = None
if img_urls or blip_caption or query:
metadata = {}
if img_urls:
metadata["img_urls"] = img_urls
if blip_caption:
metadata["blip_caption"] = blip_caption
if query:
metadata["image_query"] = query
return content, metadata
def determine_question_target(question: str, speaker_a: str, speaker_b: str) -> str:
"""
Determine which speaker a question is asking about based on the question text.
Args:
question: The question text
speaker_a: Name of speaker A (e.g., "Caroline")
speaker_b: Name of speaker B (e.g., "Melanie")
Returns:
The name of the speaker the question is about (speaker_a or speaker_b)
"""
question_lower = question.lower()
speaker_a_lower = speaker_a.lower()
speaker_b_lower = speaker_b.lower()
# Check for possessive forms too (e.g., "Melanie's kids")
a_in_question = (
speaker_a_lower in question_lower or f"{speaker_a_lower}'s" in question_lower
)
b_in_question = (
speaker_b_lower in question_lower or f"{speaker_b_lower}'s" in question_lower
)
if a_in_question and not b_in_question:
return speaker_a
elif b_in_question and not a_in_question:
return speaker_b
else:
# Question mentions both or neither - default to speaker_a
return speaker_a
class LoCoMoRunner(RunnerMixin):
"""
Executes LoCoMo benchmark tests against a Honcho instance.
"""
def __init__(
self,
base_api_port: int = 8000,
pool_size: int = 1,
anthropic_api_key: str | None = None,
timeout_seconds: int | None = None,
cleanup_workspace: bool = False,
use_get_context: bool = False,
redis_url: str = "redis://localhost:6379/0",
reasoning_level: ReasoningLevel | 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)
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)
"""
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
# Initialize common components (metrics, logging)
self._init_common("locomo")
# Initialize LLM clients
self.anthropic_client: AsyncAnthropic = create_anthropic_client(
anthropic_api_key
)
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.
Args:
conversation_data: Dictionary containing conversation and QA data
honcho_url: URL of the Honcho instance to use
question_count: Optional limit on number of questions to run
Returns:
Conversation execution results
"""
start_time = time.time()
sample_id = conversation_data.get("sample_id", "unknown")
conversation = conversation_data.get("conversation", {})
qa_list = conversation_data.get("qa", [])
speaker_a = conversation.get("speaker_a", "User")
speaker_b = conversation.get("speaker_b", "Assistant")
print(f"\n{'=' * 80}")
print(f"Executing LoCoMo conversation {sample_id}")
print(f"Speakers: {speaker_a} and {speaker_b}")
print(f"{'=' * 80}")
# Create workspace for this conversation
workspace_id = f"locomo_{sample_id}"
honcho_client = self.create_honcho_client(workspace_id, honcho_url)
result: ConversationResult = {
"sample_id": sample_id,
"speaker_a": speaker_a,
"speaker_b": speaker_b,
"total_sessions": 0,
"total_turns": 0,
"total_tokens": 0,
"question_results": [],
"category_scores": {},
"overall_score": 0.0,
"error": None,
"start_time": start_time,
"end_time": 0.0,
"duration_seconds": 0.0,
}
try:
# Create peers using their actual names as IDs
peer_a = await honcho_client.aio.peer(id=speaker_a)
peer_b = await honcho_client.aio.peer(id=speaker_b)
# Create session for this conversation
session_id = f"{workspace_id}_session"
session = await honcho_client.aio.session(id=session_id)
# 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),
),
]
)
# 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)
print(
f"[{workspace_id}] Ingested {len(messages)} messages (~{total_tokens:,} tokens). Waiting for deriver queue..."
)
# Wait for deriver queue to empty
await asyncio.sleep(1)
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
print(
f"[{workspace_id}] Deriver queue empty. Triggering dream consolidation for both peers..."
)
# Trigger dream for memory consolidation for BOTH peers
# Dream for speaker_a
dream_success_a = await self.trigger_dream_and_wait(
honcho_client,
workspace_id,
observer=speaker_a,
session_id=session_id,
)
if not dream_success_a:
print(
f"[{workspace_id}] Warning: Dream for {speaker_a} did not complete, proceeding anyway"
)
else:
print(f"[{workspace_id}] Dream for {speaker_a} completed.")
# Dream for speaker_b
dream_success_b = await self.trigger_dream_and_wait(
honcho_client,
workspace_id,
observer=speaker_b,
session_id=session_id,
)
if not dream_success_b:
print(
f"[{workspace_id}] Warning: Dream for {speaker_b} did not complete, proceeding anyway"
)
else:
print(f"[{workspace_id}] Dream for {speaker_b} completed.")
# Filter questions
filtered_qa = filter_questions(
qa_list,
exclude_adversarial=True,
test_count=question_count,
)
print(f"[{workspace_id}] Executing {len(filtered_qa)} questions...")
# Execute questions
for q_idx, qa in enumerate(filtered_qa):
question = qa.get("question", "")
expected_answer = qa.get("answer", "")
category = qa.get("category", 0)
evidence = qa.get("evidence", [])
category_name = CATEGORY_NAMES.get(category, f"category_{category}")
# Determine which peer the question is about (returns speaker name)
target_speaker = determine_question_target(
question, speaker_a, speaker_b
)
target_peer = peer_a if target_speaker == speaker_a else peer_b
print(
f" Q{q_idx + 1} [{category_name}] (asking {target_speaker}): {question}"
)
try:
if self.use_get_context:
# Use get_context + LLM - target the appropriate peer
context = await session.aio.context(
summary=True,
peer_target=target_speaker,
last_user_message=question,
)
context_messages = context.to_anthropic(assistant="assistant")
context_messages.append({"role": "user", "content": question})
response = await self.anthropic_client.messages.create(
model="claude-sonnet-4-5",
max_tokens=1024,
messages=cast(list[MessageParam], context_messages),
)
if not response.content:
raise ValueError("Anthropic returned empty response")
content_block = response.content[0]
actual_response = getattr(content_block, "text", "")
else:
# Use dialectic .chat endpoint on the appropriate peer
actual_response = await target_peer.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,
question,
str(expected_answer),
actual_response,
evidence_context=evidence_context,
)
passed = judgment.get("passed", False)
question_result: QuestionResult = {
"question_id": q_idx,
"question": question,
"expected_answer": str(expected_answer),
"actual_response": actual_response,
"category": category,
"category_name": category_name,
"evidence": evidence,
"judgment": judgment,
"passed": passed,
}
result["question_results"].append(question_result)
status = "PASS" if passed else "FAIL"
print(f" [{status}]")
if not passed:
print(f" Expected: {expected_answer}")
print(f" Got: {actual_response[:200]}...")
except Exception as e:
self.logger.error(f"Error executing question {q_idx}: {e}")
question_result = QuestionResult(
question_id=q_idx,
question=question,
expected_answer=str(expected_answer),
actual_response=f"ERROR: {e}",
category=category,
category_name=category_name,
evidence=evidence,
judgment={"passed": False, "reasoning": str(e)},
passed=False,
)
result["question_results"].append(question_result)
# Calculate category scores
result["category_scores"] = calculate_category_scores(
result["question_results"]
)
# Calculate overall score (pass rate)
if result["question_results"]:
passed_count = sum(
1 for qr in result["question_results"] if qr["passed"]
)
result["overall_score"] = passed_count / len(result["question_results"])
# Cleanup workspace if requested
if self.cleanup_workspace:
try:
await honcho_client.aio.delete_workspace(workspace_id)
print(f"[{workspace_id}] Cleaned up workspace")
except Exception as e:
print(f"Failed to delete workspace: {e}")
result["end_time"] = time.time()
result["duration_seconds"] = result["end_time"] - result["start_time"]
print(
f"\n[{workspace_id}] Completed in {format_duration(result['duration_seconds'])}"
)
print(f"Overall Score: {result['overall_score']:.3f}")
except Exception as e:
self.logger.error(f"Error executing conversation {sample_id}: {e}")
result["error"] = str(e)
result["end_time"] = time.time()
result["duration_seconds"] = result["end_time"] - result["start_time"]
return result
async def run_conversations(
self,
data_file: Path,
batch_size: int = 1,
test_count: int | None = None,
sample_id: str | None = None,
question_count: int | None = None,
) -> tuple[list[ConversationResult], float]:
"""
Run multiple conversations from the LoCoMo benchmark.
Args:
data_file: Path to the LoCoMo JSON file
batch_size: Number of conversations to run concurrently in each batch
test_count: Optional number of conversations to run
sample_id: Optional sample_id to run only that conversation
question_count: Optional limit on questions per conversation
Returns:
Tuple of (list of conversation results, total duration)
"""
conversations = load_locomo_data(data_file)
# Filter by sample_id if specified
if sample_id is not None:
conversations = [
c for c in conversations if c.get("sample_id") == sample_id
]
if not conversations:
print(f"Error: No conversation found with sample_id '{sample_id}'")
return [], 0.0
print(f"Filtering to sample_id '{sample_id}'")
# Limit by test_count
if test_count is not None and test_count > 0:
conversations = conversations[:test_count]
print(f"Limiting to {len(conversations)} conversations")
print(f"Running {len(conversations)} conversations from {data_file}")
if self.pool_size > 1:
print(
f"Distributing conversations across {self.pool_size} Honcho instances"
)
overall_start = time.time()
all_results: list[ConversationResult] = []
for i in range(0, len(conversations), batch_size):
batch = conversations[i : i + batch_size]
batch_num = (i // batch_size) + 1
total_batches = (len(conversations) + batch_size - 1) // batch_size
print(f"\n{'=' * 80}")
print(
f"Processing batch {batch_num}/{total_batches} ({len(batch)} conversations)"
)
print(f"{'=' * 80}")
# Run conversations in current batch concurrently
batch_results: list[ConversationResult] = await asyncio.gather(
*[
self.execute_conversation(
conv,
self.get_honcho_url_for_index(i + idx),
question_count=question_count,
)
for idx, conv in enumerate(batch)
]
)
all_results.extend(batch_results)
overall_end = time.time()
overall_duration = overall_end - overall_start
# Finalize metrics collection
self.metrics_collector.finalize_collection()
return all_results, overall_duration
async def main() -> int:
"""Main entry point for the LoCoMo test runner."""
parser = argparse.ArgumentParser(
description="Run LoCoMo benchmark tests against a Honcho instance",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
%(prog)s --data-file tests/bench/locomo_data/locomo10.json
%(prog)s --data-file locomo10.json --pool-size 4
%(prog)s --data-file locomo10.json --sample-id "sample_0"
%(prog)s --data-file locomo10.json --test-count 5 --question-count 20
%(prog)s --data-file locomo10.json --reasoning-level high
""",
)
parser.add_argument(
"--data-file",
type=Path,
required=True,
help="Path to LoCoMo JSON file (required)",
)
# Add common arguments shared across all runners
add_common_arguments(parser)
# LoCoMo-specific arguments
parser.add_argument(
"--anthropic-api-key",
type=str,
help="Anthropic API key for response judging (optional)",
)
parser.add_argument(
"--sample-id",
type=str,
help="Run only the conversation with this sample_id (skips all others)",
)
parser.add_argument(
"--test-count",
type=int,
help="Number of conversations to run from the data file (default: all)",
)
parser.add_argument(
"--question-count",
type=int,
help="Number of questions per conversation to run (default: all)",
)
args = parser.parse_args()
# Validate common arguments
error = validate_common_arguments(args)
if error:
print(error)
return 1
# Validate locomo-specific arguments
if not args.data_file.exists():
print(f"Error: Data file {args.data_file} does not exist")
return 1
# Create test runner
runner = LoCoMoRunner(
base_api_port=args.base_api_port,
pool_size=args.pool_size,
anthropic_api_key=args.anthropic_api_key,
timeout_seconds=args.timeout,
cleanup_workspace=args.cleanup_workspace,
use_get_context=args.use_get_context,
redis_url=args.redis_url,
reasoning_level=args.reasoning_level,
)
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
if __name__ == "__main__":
exit_code = asyncio.run(main())
exit(exit_code)
Reference in New Issue View Git Blame Copy Permalink