From ffff6afb1a853274aefda2e8e1c43d1a799f230f Mon Sep 17 00:00:00 2001 From: Benjamin McCormick Date: Fri, 23 Jan 2026 15:47:51 -0500 Subject: [PATCH 1/3] chore: refactor test runners to be cleaner, deduplicate code, add params for testing remote servers --- tests/bench/beam.py | 569 +++++++------------ tests/bench/locomo.py | 725 ++++++++++-------------- tests/bench/longmem.py | 1018 +++++++++++++--------------------- tests/bench/runner_common.py | 551 +++++++++++++----- 4 files changed, 1278 insertions(+), 1585 deletions(-) diff --git a/tests/bench/beam.py b/tests/bench/beam.py index 3a9f3290..0567802e 100644 --- a/tests/bench/beam.py +++ b/tests/bench/beam.py @@ -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,120 @@ 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) + + 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 +277,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 +286,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 +300,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 +308,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 +316,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 +373,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 +390,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 +506,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()) diff --git a/tests/bench/locomo.py b/tests/bench/locomo.py index 4fe3fc2b..f91a1285 100644 --- a/tests/bench/locomo.py +++ b/tests/bench/locomo.py @@ -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,307 @@ 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) + + # 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 +568,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()) diff --git a/tests/bench/longmem.py b/tests/bench/longmem.py index 0ed1b7f3..7271e189 100644 --- a/tests/bench/longmem.py +++ b/tests/bench/longmem.py @@ -55,7 +55,6 @@ Optional arguments: """ import argparse -import asyncio import json import time from datetime import datetime @@ -77,19 +76,19 @@ from .longmem_common import ( calculate_total_tokens, calculate_type_statistics, filter_questions, - format_duration, judge_response, load_test_file, parse_longmemeval_date, write_json_summary, ) from .runner_common import ( - ReasoningLevel, - RunnerMixin, + BaseRunner, + ItemContext, + RunnerConfig, add_common_arguments, create_anthropic_client, create_openai_client, - export_metrics, + format_duration, validate_common_arguments, ) @@ -126,53 +125,40 @@ class TestResult(TypedDict): start_time: float end_time: float duration_seconds: float - output_lines: list[str] -class LongMemEvalRunner(RunnerMixin): +class LongMemEvalRunner(BaseRunner[TestResult]): """ Executes longmemeval JSON tests against a Honcho instance. """ def __init__( self, - base_api_port: int = 8000, - pool_size: int = 1, + config: RunnerConfig, + test_file: Path, anthropic_api_key: str | None = None, - timeout_seconds: int | None = None, merge_sessions: bool = False, - cleanup_workspace: bool = False, - use_get_context: bool = False, - redis_url: str = "redis://localhost:6379/0", - reasoning_level: ReasoningLevel | None = None, + test_count: int | None = None, + question_id: str | None = None, ): """ Initialize the 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 + test_file: Path to the longmemeval JSON file anthropic_api_key: Anthropic API key for judging responses - timeout_seconds: Timeout for deriver queue in seconds merge_sessions: If True, merge all sessions within a question into one session - cleanup_workspace: If True, delete workspace after executing question (default: False) - 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) + test_count: Optional number of tests to run (runs first N tests) + question_id: Optional question_id to run (skips all others) """ - 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 10000 - ) + self.test_file: Path = test_file self.merge_sessions: bool = merge_sessions - 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.test_count: int | None = test_count + self.question_id_filter: str | None = question_id - # Initialize common components (metrics, logging) - self._init_common("longmem") + # Initialize base class (sets up metrics collector and logger) + super().__init__(config) # Initialize LLM clients self.anthropic_client: AsyncAnthropic = create_anthropic_client( @@ -180,138 +166,143 @@ class LongMemEvalRunner(RunnerMixin): ) self.openai_client: AsyncOpenAI = create_openai_client() - def _get_latest_input_tokens_used(self) -> int | None: - """Get the uncached input tokens from the most recent dialectic_chat metric. + def get_metrics_prefix(self) -> str: + return "longmem" - Returns: - Number of tokens used, or None if not found - """ - metrics_file = Path(settings.LOCAL_METRICS_FILE) - if not metrics_file.exists(): - return None + def load_items(self) -> list[Any]: + """Load questions from the test file.""" + questions = load_test_file(self.test_file) + questions = filter_questions( + questions, self.test_file, self.question_id_filter, self.test_count + ) + return questions - # Read the file and find the most recent dialectic_chat metric - try: - with open(metrics_file) as f: - lines = f.readlines() + def get_workspace_id(self, item: Any) -> str: + """Return workspace ID for a question.""" + return f"{item['question_id']}_{item['question_type']}" - # Search backwards through the file for the most recent dialectic_chat - for line in reversed(lines): - if not line.strip(): - continue - try: - data = json.loads(line) - task_name = data.get("task_name", "") - if task_name.startswith("dialectic_chat_"): - for metric in data.get("metrics", []): - metric_name = metric.get("name", "") - if metric_name.endswith("uncached_input_tokens"): - return int(metric.get("value", 0)) - except (json.JSONDecodeError, KeyError, ValueError): - continue - - except Exception as e: - self.logger.warning(f"Error reading metrics file: {e}") - - return None - - async def execute_question( - self, question_data: dict[str, Any], honcho_url: str - ) -> TestResult: - """ - Execute a single longmemeval question. - - Args: - question_data: Dictionary containing question data - honcho_url: URL of the Honcho instance to use - - Returns: - Test execution results - """ - question_id = question_data["question_id"] - question_type = question_data["question_type"] - question = question_data["question"] - expected_answer = question_data["answer"] - question_date = question_data.get("question_date", "") - - question_with_date = ( - f"[{question_date}] {question}" if question_date else question + def get_session_id(self, item: Any, workspace_id: str) -> str: + """Return session ID for a question.""" + if self.merge_sessions: + return f"{workspace_id}_merged" + # For non-merged, we use the first haystack session ID + haystack_session_ids = item.get("haystack_session_ids", []) + return ( + haystack_session_ids[0] + if haystack_session_ids + else f"{workspace_id}_session" ) - output_lines: list[str] = [] - output_lines.append( - f"\033[1mExecuting question {question_id} ({question_type})\033[0m" - ) - output_lines.append(f"Question: {question_with_date}") - output_lines.append(f"Expected: {expected_answer}") - output_lines.append(f"Using Honcho instance: {honcho_url}") + 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") - # Create workspace for this question - workspace_id = f"{question_id}_{question_type}" - honcho_client = self.create_honcho_client(workspace_id, honcho_url) + async def setup_session(self, ctx: ItemContext, item: Any) -> None: + """Create and configure session with appropriate observation settings.""" + is_assistant_type = item["question_type"] == "single-session-assistant" + user_peer = ctx.peers["user"] + assistant_peer = ctx.peers["assistant"] - results: TestResult = { - "question_id": question_id, - "question_type": question_type, - "workspace_id": workspace_id, - "sessions_created": [], - "query_executed": None, - "passed": False, - "error": None, - "start_time": time.time(), - "end_time": 0.0, - "duration_seconds": 0.0, - "output_lines": output_lines, - } + if self.merge_sessions: + # Create a single merged session + ctx.session = await ctx.honcho_client.aio.session(id=ctx.session_id) - try: - user_peer = await honcho_client.aio.peer(id="user") - assistant_peer = await honcho_client.aio.peer(id="assistant") - - # Process haystack sessions - haystack_dates = question_data.get("haystack_dates", []) - haystack_sessions = question_data.get("haystack_sessions", []) - haystack_session_ids = question_data.get("haystack_session_ids", []) - - # Validate alignment of dates, session IDs, and sessions - if len(haystack_dates) != len(haystack_sessions): - raise ValueError( - f"Misaligned data: {len(haystack_dates)} dates but {len(haystack_sessions)} sessions" + if is_assistant_type: + await ctx.session.aio.add_peers( + [ + ( + user_peer, + SessionPeerConfig(observe_me=False, observe_others=False), + ), + ( + assistant_peer, + SessionPeerConfig(observe_me=True, observe_others=False), + ), + ] ) - if len(haystack_session_ids) != len(haystack_sessions): - raise ValueError( - f"Misaligned data: {len(haystack_session_ids)} session IDs but {len(haystack_sessions)} sessions" + else: + await ctx.session.aio.add_peers( + [ + ( + user_peer, + SessionPeerConfig(observe_me=True, observe_others=False), + ), + ( + assistant_peer, + SessionPeerConfig(observe_me=False, observe_others=False), + ), + ] ) + else: + # Sessions are created during ingestion for non-merged mode + ctx.session = None - # Parse all dates upfront to catch parsing errors early - parsed_dates: list[datetime] = [] - for date_str in haystack_dates: - try: - parsed_dates.append(parse_longmemeval_date(date_str)) - except ValueError as e: - raise ValueError(f"Error parsing date '{date_str}': {e}") from e + async def ingest_messages(self, ctx: ItemContext, item: Any) -> int: + """Ingest haystack messages into session(s).""" + is_assistant_type = item["question_type"] == "single-session-assistant" + user_peer = ctx.peers["user"] + assistant_peer = ctx.peers["assistant"] - haystack_total_messages = sum(len(session) for session in haystack_sessions) + haystack_dates = item.get("haystack_dates", []) + haystack_sessions = item.get("haystack_sessions", []) + haystack_session_ids = item.get("haystack_session_ids", []) - # Calculate total tokens available in the sessions for this question - total_available_tokens = calculate_total_tokens(haystack_sessions) + # Parse dates + parsed_dates = [parse_longmemeval_date(d) for d in haystack_dates] - print( - f"[{workspace_id}] processing {len(haystack_sessions)} sessions with {haystack_total_messages} total messages ({total_available_tokens} total tokens)" - ) + total_messages = 0 - # Determine which peer should be observed based on question type - is_assistant_type = question_type == "single-session-assistant" + if self.merge_sessions: + # Collect all messages from all sessions + all_messages: list[MessageCreateParams] = [] + for session_date, session_messages in zip( + parsed_dates, haystack_sessions, strict=True + ): + for msg in session_messages: + role = msg["role"] + content = msg["content"] - # Initialize merged_session_id for potential use in dream trigger - merged_session_id: str | None = None + # Split long messages + if len(content) > 25000: + chunks = [ + content[i : i + 25000] + for i in range(0, len(content), 25000) + ] + for chunk in chunks: + if role == "user": + all_messages.append( + user_peer.message(chunk, created_at=session_date) + ) + elif role == "assistant": + all_messages.append( + assistant_peer.message( + chunk, created_at=session_date + ) + ) + else: + if role == "user": + all_messages.append( + user_peer.message(content, created_at=session_date) + ) + elif role == "assistant": + all_messages.append( + assistant_peer.message(content, created_at=session_date) + ) - if self.merge_sessions: - # Create a single merged session for all messages - merged_session_id = f"{workspace_id}_merged" - session = await honcho_client.aio.session(id=merged_session_id) + # Add messages in batches + for i in range(0, len(all_messages), 100): + batch = all_messages[i : i + 100] + await ctx.session.aio.add_messages(batch) + + total_messages = len(all_messages) + else: + # Create separate sessions for each haystack session + for session_date, session_id, session_messages in zip( + parsed_dates, haystack_session_ids, haystack_sessions, strict=True + ): + session = await ctx.honcho_client.aio.session(id=session_id) - # Configure peer observation based on question type if is_assistant_type: await session.aio.add_peers( [ @@ -347,426 +338,218 @@ class LongMemEvalRunner(RunnerMixin): ] ) - # Collect all messages from all sessions in chronological order - all_messages: list[MessageCreateParams] = [] - for session_date, session_messages in zip( - parsed_dates, haystack_sessions, strict=True - ): - for msg in session_messages: - role = msg["role"] - content = msg["content"] + honcho_messages: list[MessageCreateParams] = [] + for msg in session_messages: + role = msg["role"] + content = msg["content"] - # 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": - all_messages.append( - user_peer.message( - chunk, created_at=session_date - ) - ) - elif role == "assistant": - all_messages.append( - assistant_peer.message( - chunk, created_at=session_date - ) - ) - else: + if len(content) > 25000: + chunks = [ + content[i : i + 25000] + for i in range(0, len(content), 25000) + ] + for chunk in chunks: if role == "user": - all_messages.append( - user_peer.message(content, created_at=session_date) + honcho_messages.append( + user_peer.message(chunk, created_at=session_date) ) elif role == "assistant": - all_messages.append( + honcho_messages.append( assistant_peer.message( - content, created_at=session_date + chunk, created_at=session_date ) ) + else: + if role == "user": + honcho_messages.append( + user_peer.message(content, created_at=session_date) + ) + elif role == "assistant": + honcho_messages.append( + assistant_peer.message(content, created_at=session_date) + ) - # Add messages in batches of 100 (max supported by add_messages) - if all_messages: - for i in range(0, len(all_messages), 100): - batch = all_messages[i : i + 100] - await session.aio.add_messages(batch) + for i in range(0, len(honcho_messages), 100): + batch = honcho_messages[i : i + 100] + await session.aio.add_messages(batch) - results["sessions_created"].append( - SessionResult( - name=merged_session_id, message_count=len(all_messages) - ) + total_messages += len(honcho_messages) + + return total_messages + + def get_dream_observers(self, item: Any) -> list[str]: + """Return the observer based on question type.""" + is_assistant_type = item["question_type"] == "single-session-assistant" + return ["assistant"] if is_assistant_type else ["user"] + + def _get_latest_input_tokens_used(self) -> int | None: + """Get the uncached input tokens from the most recent dialectic_chat metric.""" + metrics_file = Path(settings.LOCAL_METRICS_FILE) + if not metrics_file.exists(): + return None + + try: + with open(metrics_file) as f: + lines = f.readlines() + + for line in reversed(lines): + if not line.strip(): + continue + try: + data = json.loads(line) + task_name = data.get("task_name", "") + if task_name.startswith("dialectic_chat_"): + for metric in data.get("metrics", []): + metric_name = metric.get("name", "") + if metric_name.endswith("uncached_input_tokens"): + return int(metric.get("value", 0)) + except (json.JSONDecodeError, KeyError, ValueError): + continue + except Exception as e: + self.logger.warning(f"Error reading metrics file: {e}") + + return None + + async def execute_questions(self, ctx: ItemContext, item: Any) -> TestResult: + """Execute the question and judge the response.""" + start_time = time.time() + workspace_id = ctx.workspace_id + + question_id = item["question_id"] + question_type = item["question_type"] + question = item["question"] + expected_answer = item["answer"] + question_date = item.get("question_date", "") + + question_with_date = ( + f"[{question_date}] {question}" if question_date else question + ) + is_assistant_type = question_type == "single-session-assistant" + + # Calculate total tokens for efficiency metrics + haystack_sessions = item.get("haystack_sessions", []) + total_available_tokens = calculate_total_tokens(haystack_sessions) + + result: TestResult = { + "question_id": question_id, + "question_type": question_type, + "workspace_id": workspace_id, + "sessions_created": [], # Populated during ingestion tracking + "query_executed": None, + "passed": False, + "error": None, + "start_time": start_time, + "end_time": 0.0, + "duration_seconds": 0.0, + } + + try: + print(f" Asking: {question_with_date}") + + if self.config.use_get_context: + # Use get_context instead of dialectic .chat endpoint + if not self.merge_sessions or ctx.session is None: + raise ValueError("Merged session required for get_context mode") + + peer_id = "assistant" if is_assistant_type else "user" + context = await ctx.session.aio.context( + summary=True, + peer_target=peer_id, + search_query=question, ) + + context_messages = context.to_anthropic(assistant="assistant") + context_messages.append({"role": "user", "content": question_with_date}) + + 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") + + actual_response = getattr(response.content[0], "text", "") else: - # create separate sessions - # Zip together dates, session IDs, and session content - for session_date, session_id, session_messages in zip( - parsed_dates, haystack_session_ids, haystack_sessions, strict=True - ): - session = await honcho_client.aio.session(id=session_id) - - # Configure peer observation based on question type - if is_assistant_type: - # For assistant questions, observe the assistant peer - await session.aio.add_peers( - [ - ( - user_peer, - SessionPeerConfig( - observe_me=False, observe_others=False - ), - ), - ( - assistant_peer, - SessionPeerConfig( - observe_me=True, observe_others=False - ), - ), - ] - ) - else: - # For user questions, observe the user peer (default behavior) - await session.aio.add_peers( - [ - ( - user_peer, - SessionPeerConfig( - observe_me=True, observe_others=False - ), - ), - ( - assistant_peer, - SessionPeerConfig( - observe_me=False, observe_others=False - ), - ), - ] - ) - - honcho_messages: list[MessageCreateParams] = [] - for msg in session_messages: - role = msg["role"] - content = msg["content"] - - # 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: - # Use the session date as the timestamp for all messages in this session - if role == "user": - honcho_messages.append( - user_peer.message( - chunk, created_at=session_date - ) - ) - elif role == "assistant": - honcho_messages.append( - assistant_peer.message( - chunk, created_at=session_date - ) - ) - else: - # Use the session date as the timestamp for all messages in this session - if role == "user": - honcho_messages.append( - user_peer.message(content, created_at=session_date) - ) - elif role == "assistant": - honcho_messages.append( - assistant_peer.message( - content, created_at=session_date - ) - ) - - if honcho_messages: - for i in range(0, len(honcho_messages), 100): - batch = honcho_messages[i : i + 100] - await session.aio.add_messages(batch) - - results["sessions_created"].append( - SessionResult( - name=session_id, message_count=len(honcho_messages) - ) - ) - - print( - f"[{workspace_id}] fired all messages.\nwaiting for deriver queue to be empty... will time out in {self.timeout_seconds} seconds" - ) - await asyncio.sleep( - 1 - ) # Give time for at least some tasks to be queued, so deriver queue size check doesn't immediately return 0 - - # Enable flush mode to bypass batch token threshold - await self.flush_deriver_queue() - - queue_empty = await self.wait_for_deriver_queue_empty(honcho_client) - if not queue_empty: - output_lines.append("Deriver queue never emptied!!!") - results["error"] = "Deriver queue timeout" - return results - - # Trigger dream for memory consolidation before questions - print( - f"[{workspace_id}] Deriver queue empty. Triggering dream consolidation..." - ) - - # Determine session_id for dream - dream_session_id = ( - merged_session_id - if self.merge_sessions and merged_session_id - else ( - haystack_session_ids[0] - if haystack_session_ids - else f"{workspace_id}_session" + # Use dialectic .chat endpoint + peer = ( + ctx.peers["assistant"] if is_assistant_type else ctx.peers["user"] + ) + actual_response = await peer.aio.chat( + question_with_date, + reasoning_level=self.config.reasoning_level, ) - ) - - # Determine observer based on question type - observer_peer = "assistant" if is_assistant_type else "user" - - # Single orchestrated dream handles all reasoning types - dream_success = await self.trigger_dream_and_wait( - honcho_client, - workspace_id, - observer=observer_peer, - session_id=dream_session_id, - ) - if not dream_success: - print(f"[{workspace_id}] Warning: Dream did not complete") - print(f"[{workspace_id}] Dream completed. Executing question...") - - # Execute the question - output_lines.append(f"\nAsking question: {question_with_date}") - - try: - if self.use_get_context: - # Use get_context instead of dialectic .chat endpoint - # Get the session to retrieve context from - if not self.merge_sessions or merged_session_id is None: - raise ValueError( - "Merged session ID is required when using get_context. Set --merge-sessions to True." - ) - session = await honcho_client.aio.session(id=merged_session_id) - - # Get context for the appropriate peer - peer_id = "assistant" if is_assistant_type else "user" - context = await session.aio.context( - summary=True, - peer_target=peer_id, - search_query=question, - ) - - # Format context using to_anthropic method - context_messages = context.to_anthropic(assistant="assistant") - - # Add the question as the final user message - context_messages.append( - {"role": "user", "content": question_with_date} - ) - - # Call Anthropic API to generate response - 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 the appropriate peer based on question type - if is_assistant_type: - # For assistant questions, use the assistant peer - actual_response = await assistant_peer.aio.chat( - question_with_date, - reasoning_level=self.reasoning_level, - ) - else: - # For user questions, use the user peer (default behavior) - actual_response = await user_peer.aio.chat( - question_with_date, - reasoning_level=self.reasoning_level, - ) - - # Clean up 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}") - actual_response = ( actual_response if isinstance(actual_response, str) else "" ) - input_tokens_used = self._get_latest_input_tokens_used() - - token_efficiency = None - if input_tokens_used is not None and total_available_tokens > 0: - efficiency_ratio = input_tokens_used / total_available_tokens - token_efficiency = { - "total_available_tokens": total_available_tokens, - "tokens_used": input_tokens_used, - "efficiency_ratio": efficiency_ratio, - } - output_lines.append( - f" token efficiency: {efficiency_ratio:.4f} ({input_tokens_used}/{total_available_tokens} tokens, {efficiency_ratio * 100:.2f}%)" - ) - - judgment = await judge_response( - self.openai_client, - question_with_date, - expected_answer, - actual_response, - question_type, - question_id, - ) - - query_result: QueryResult = { - "question": question_with_date, - "expected_answer": expected_answer, - "actual_response": actual_response, - "judgment": judgment, - "token_efficiency": token_efficiency, + # Get token efficiency + input_tokens_used = self._get_latest_input_tokens_used() + token_efficiency = None + if input_tokens_used is not None and total_available_tokens > 0: + efficiency_ratio = input_tokens_used / total_available_tokens + token_efficiency = { + "total_available_tokens": total_available_tokens, + "tokens_used": input_tokens_used, + "efficiency_ratio": efficiency_ratio, } - - results["query_executed"] = query_result - results["passed"] = judgment["passed"] - - output_lines.append( - " judgment: \033[1m\033[32mPASS\033[0m" - if judgment["passed"] - else " judgment: \033[1m\033[31mFAIL\033[0m" + print( + f" Token efficiency: {efficiency_ratio:.4f} ({input_tokens_used}/{total_available_tokens})" ) - if not judgment["passed"]: - output_lines.append( - f" got response: \033[3m{actual_response}\033[0m" - ) - output_lines.append(f" expected: {expected_answer}") - output_lines.append(f" reasoning: {judgment['reasoning']}") - except Exception as e: - self.logger.error(f"Error executing question: {e}") - query_result = QueryResult( - question=question_with_date, - expected_answer=expected_answer, - actual_response=f"ERROR: {e}", - judgment={ - "passed": False, - "reasoning": f"Question execution failed: {e}", - }, - token_efficiency=None, - ) - results["query_executed"] = query_result - results["passed"] = False - - results["end_time"] = time.time() - results["duration_seconds"] = results["end_time"] - results["start_time"] - - output_lines.append( - f"\nQuestion {question_id} completed. Status: {'PASS' if results['passed'] else 'FAIL'} (Duration: {format_duration(results['duration_seconds'])})" + # Judge the response + judgment = await judge_response( + self.openai_client, + question_with_date, + expected_answer, + actual_response, + question_type, + question_id, ) + result["query_executed"] = QueryResult( + question=question_with_date, + expected_answer=expected_answer, + actual_response=actual_response, + judgment=judgment, + token_efficiency=token_efficiency, + ) + result["passed"] = judgment["passed"] + + status = ( + "\033[1m\033[32mPASS\033[0m" + if judgment["passed"] + else "\033[1m\033[31mFAIL\033[0m" + ) + print(f" Judgment: {status}") + if not judgment["passed"]: + print(f" Got: \033[3m{actual_response}\033[0m") + print(f" Expected: {expected_answer}") + print(f" Reasoning: {judgment['reasoning']}") + except Exception as e: - self.logger.error(f"Error executing question {question_id}: {e}") - results["error"] = str(e) - results["passed"] = False - results["end_time"] = time.time() - results["duration_seconds"] = results["end_time"] - results["start_time"] - output_lines.append(f"Error executing question {question_id}: {e}") - - return results - - async def run_all_questions( - self, - test_file: Path, - batch_size: int = 10, - test_count: int | None = None, - question_id: str | None = None, - ) -> tuple[list[TestResult], float]: - """ - Run all questions in a longmemeval test file. - - Args: - test_file: Path to the longmemeval JSON file - batch_size: Number of questions to run concurrently in each batch - test_count: Optional number of tests to run (runs first N tests) - question_id: Optional question_id to run (skips all others) - - Returns: - Tuple of (list of test results, total duration) - """ - questions = load_test_file(test_file) - questions = filter_questions(questions, test_file, question_id, test_count) - if not questions: - return [], 0.0 - - print( - f"found {len(questions)} {'question' if len(questions) == 1 else 'questions'} in {test_file}" - ) - if self.pool_size > 1: - print( - f"distributing questions across {self.pool_size} Honcho instances (ports {self.base_api_port}-{self.base_api_port + self.pool_size - 1})" + self.logger.error(f"Error executing question: {e}") + result["query_executed"] = QueryResult( + question=question_with_date, + expected_answer=expected_answer, + actual_response=f"ERROR: {e}", + judgment={ + "passed": False, + "reasoning": f"Question execution failed: {e}", + }, + token_efficiency=None, ) + result["passed"] = False + result["error"] = str(e) - overall_start = time.time() + result["end_time"] = time.time() + result["duration_seconds"] = result["end_time"] - result["start_time"] - # Process questions in batches - all_results: list[TestResult] = [] + return result - for i in range(0, len(questions), batch_size): - batch = questions[i : i + batch_size] - batch_num = (i // batch_size) + 1 - total_batches = (len(questions) + batch_size - 1) // batch_size - - print(f"\n{'=' * 60}") - print( - f"Processing batch {batch_num}/{total_batches} ({len(batch)} questions)" - ) - print(f"{'=' * 60}") - - # Run questions in current batch concurrently, distributing via round-robin - batch_results: list[TestResult] = await asyncio.gather( - *[ - self.execute_question(q, self.get_honcho_url_for_index(i + idx)) - for idx, q in enumerate(batch) - ] - ) - - # Print detailed per-question outputs for this batch - for result in batch_results: - print(f"\n{'=' * 60}") - print("\n".join(result.get("output_lines", []))) - print(f"{'=' * 60}\n") - - all_results.extend(batch_results) - - overall_end = time.time() - overall_duration = overall_end - overall_start - - # Finalize metrics collection - self.metrics_collector.finalize_collection() - - return all_results, overall_duration - - def print_summary( - self, results: list[TestResult], total_elapsed_seconds: float | None = None - ) -> None: - """ - Print a summary of all test results. - - Args: - results: List of test results - total_elapsed_seconds: Total elapsed time - """ + def print_summary(self, results: list[TestResult], total_duration: float) -> None: + """Print a summary of all test results.""" print(f"\n{'=' * 80}") print("LONGMEMEVAL TEST EXECUTION SUMMARY") print(f"{'=' * 80}") @@ -774,19 +557,19 @@ class LongMemEvalRunner(RunnerMixin): total_questions = len(results) passed_questions = sum(1 for r in results if r.get("passed", False)) failed_questions = total_questions - passed_questions - total_test_time = ( - total_elapsed_seconds - if total_elapsed_seconds is not None - else sum(r["duration_seconds"] for r in results) - ) print(f"Total Questions: {total_questions}") print(f"Passed: {passed_questions}") print(f"Failed: {failed_questions}") - print(f"Success Rate: {(passed_questions / total_questions) * 100:.1f}%") - print(f"Total Test Time: {format_duration(total_test_time)}") + print( + f"Success Rate: {(passed_questions / total_questions) * 100:.1f}%" + if total_questions > 0 + else "N/A" + ) + print(f"Total Test Time: {format_duration(total_duration)}") - efficiency_ratios: list[float] = [] + # Token efficiency stats + efficiency_ratios: list[Any] = [] for result in results: query = result.get("query_executed") if query: @@ -796,64 +579,39 @@ class LongMemEvalRunner(RunnerMixin): if efficiency_ratios: avg_efficiency = sum(efficiency_ratios) / len(efficiency_ratios) - min_efficiency = min(efficiency_ratios) - max_efficiency = max(efficiency_ratios) print("\nToken Efficiency:") print( f" Average: {avg_efficiency:.4f} ({avg_efficiency * 100:.2f}% of available tokens used)" ) - print(f" Min: {min_efficiency:.4f} ({min_efficiency * 100:.2f}%)") - print(f" Max: {max_efficiency:.4f} ({max_efficiency * 100:.2f}%)") + print(f" Min: {min(efficiency_ratios):.4f}") + print(f" Max: {max(efficiency_ratios):.4f}") print("\nDetailed Results:") - print( - f"{'Question ID':<15} {'Type':<20} {'Status':<8} {'Duration':<10} {'Workspace ID':<30}" - ) - print(f"{'-' * 15} {'-' * 20} {'-' * 8} {'-' * 10} {'-' * 30}") + print(f"{'Question ID':<15} {'Type':<20} {'Status':<8} {'Duration':<10}") + print(f"{'-' * 15} {'-' * 20} {'-' * 8} {'-' * 10}") for result in results: question_id = result["question_id"] question_type = result["question_type"] status = "PASS" if result.get("passed", False) else "FAIL" duration = format_duration(result["duration_seconds"]) - workspace = result["workspace_id"] - - print( - f"{question_id:<15} {question_type:<20} {status:<8} {duration:<10} {workspace:<30}" - ) + print(f"{question_id:<15} {question_type:<20} {status:<8} {duration:<10}") print(f"{'=' * 80}") - def generate_json_summary( - self, - results: list[TestResult], - test_file: Path, - total_elapsed_seconds: float, - output_file: Path | None = None, - ) -> None: - """ - Generate a comprehensive JSON summary of test results for analytics. - - Args: - results: List of test results - test_file: Path to the test file that was executed - total_elapsed_seconds: Total elapsed time for all tests - output_file: Optional path to write JSON output to - """ + def generate_output(self, results: list[TestResult], total_duration: float) -> None: + """Generate JSON output file.""" total_questions = len(results) passed_questions = sum(1 for r in results if r.get("passed", False)) - failed_questions = total_questions - passed_questions - # Calculate statistics by question type + # Calculate statistics type_stats = calculate_type_statistics(results) + timing_stats = calculate_timing_statistics(results, total_duration) - # Calculate timing statistics - timing_stats = calculate_timing_statistics(results, total_elapsed_seconds) - - # Calculate token efficiency statistics - efficiency_ratios: list[float] = [] - total_available_tokens_list: list[int] = [] - tokens_used_list: list[int] = [] + # Token efficiency stats + efficiency_ratios: list[Any] = [] + total_available_tokens_list: list[Any] = [] + tokens_used_list: list[Any] = [] for result in results: query = result.get("query_executed") if query: @@ -879,16 +637,15 @@ class LongMemEvalRunner(RunnerMixin): "total_questions_with_metrics": len(efficiency_ratios), } - # Create the full summary summary = { "metadata": { - "test_file": str(test_file), + "test_file": str(self.test_file), "execution_timestamp": datetime.now().isoformat(), - "runner_version": "1.0.0", - "base_api_port": self.base_api_port, - "pool_size": self.pool_size, - "timeout_seconds": self.timeout_seconds, - "reasoning_level": self.reasoning_level, + "runner_version": "2.0.0", + "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(), @@ -897,7 +654,7 @@ class LongMemEvalRunner(RunnerMixin): "summary_statistics": { "total_questions": total_questions, "passed": passed_questions, - "failed": failed_questions, + "failed": total_questions - passed_questions, "success_rate_percent": (passed_questions / total_questions) * 100 if total_questions > 0 else 0, @@ -921,14 +678,19 @@ class LongMemEvalRunner(RunnerMixin): ], } - if output_file: - write_json_summary(summary, output_file) + # Determine output file + if self.config.json_output: + output_file = self.config.json_output + else: + output_file = Path( + f"tests/bench/eval_results/longmemeval_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json" + ) + + write_json_summary(summary, output_file) -async def main() -> int: - """ - Main entry point for the longmemeval test runner. - """ +def main() -> int: + """Main entry point for the longmemeval test runner.""" parser = argparse.ArgumentParser( description="Run longmemeval tests against a Honcho instance", formatter_class=argparse.RawDescriptionHelpFormatter, @@ -995,58 +757,20 @@ Examples: print(f"Error: Test count must be positive, got {args.test_count}") return 1 - # Create test runner + # Create config and runner + config = RunnerConfig.from_args(args, default_timeout=10000) + runner = LongMemEvalRunner( - base_api_port=args.base_api_port, - pool_size=args.pool_size, + config=config, + test_file=args.test_file, anthropic_api_key=args.anthropic_api_key, - timeout_seconds=args.timeout, merge_sessions=args.merge_sessions, - cleanup_workspace=args.cleanup_workspace, - use_get_context=args.use_get_context, - redis_url=args.redis_url, - reasoning_level=args.reasoning_level, + test_count=args.test_count, + question_id=args.question_id, ) - try: - # Run all questions - results, total_elapsed = await runner.run_all_questions( - args.test_file, args.batch_size, args.test_count, args.question_id - ) - runner.print_summary(results, total_elapsed_seconds=total_elapsed) - - # Print metrics summary - runner.metrics_collector.print_summary() - - # Generate JSON output if requested - if args.json_output: - runner.generate_json_summary( - results, args.test_file, total_elapsed, args.json_output - ) - else: - # Always generate a default JSON output file with timestamp - default_output = Path( - f"tests/bench/eval_results/longmemeval_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json" - ) - runner.generate_json_summary( - results, args.test_file, total_elapsed, default_output - ) - - # Export metrics to JSON file - export_metrics(runner.metrics_collector, "longmem") - - # Return exit code based on results - all_passed = all(r.get("passed", False) for r in results) - return 0 if all_passed else 1 - - except KeyboardInterrupt: - print("\nTest execution interrupted by user") - return 1 - except Exception as e: - print(f"Error running tests: {e}") - return 1 + return runner.run_and_summarize() if __name__ == "__main__": - exit_code = asyncio.run(main()) - exit(exit_code) + exit(main()) diff --git a/tests/bench/runner_common.py b/tests/bench/runner_common.py index b1072025..a12b4834 100644 --- a/tests/bench/runner_common.py +++ b/tests/bench/runner_common.py @@ -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,11 +10,13 @@ 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 @@ -28,6 +29,61 @@ 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 + + @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, + ) + + +@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 +92,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 +167,12 @@ 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)", + ) + def validate_common_arguments(args: argparse.Namespace) -> str | None: """ @@ -185,104 +261,349 @@ 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() + 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() - 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}" + # ------------------------------------------------------------------------- + # Abstract methods - must be implemented by subclasses + # ------------------------------------------------------------------------- - def create_honcho_client(self, workspace_id: str, honcho_url: str) -> Honcho: - """Create a Honcho client for a specific workspace.""" - return Honcho( - environment="local", + @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})" + ) + + 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 + batch_results = await asyncio.gather( + *[ + self.execute_item(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(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, - base_url=honcho_url, + honcho_client=self._create_honcho_client(workspace_id, honcho_url), + honcho_url=honcho_url, + session_id=session_id, ) - async def flush_deriver_queue(self) -> None: + start_time = time.time() + + try: + # 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, + ) + + 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] + if self.config.base_url: + print("Skipping flush mode (remote instance)") + 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=60) print("Enabled deriver flush mode") 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 +613,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 From 258fd3736bd03f851941fa9707b7fd1c23b89974 Mon Sep 17 00:00:00 2001 From: Benjamin McCormick Date: Fri, 23 Jan 2026 16:13:00 -0500 Subject: [PATCH 2/3] fix: turn off summaries on test runs directly --- tests/bench/beam.py | 4 +++- tests/bench/locomo.py | 4 +++- tests/bench/longmem.py | 8 ++++++-- tests/bench/run_tests.py | 14 +++++++++++--- tests/bench/runner_common.py | 34 ++++++++++++++++++++++++++++++++-- 5 files changed, 55 insertions(+), 9 deletions(-) diff --git a/tests/bench/beam.py b/tests/bench/beam.py index 0567802e..b61c0e81 100644 --- a/tests/bench/beam.py +++ b/tests/bench/beam.py @@ -171,7 +171,9 @@ class BEAMRunner(BaseRunner[ConversationResult]): user_peer = ctx.peers["user"] assistant_peer = ctx.peers["assistant"] - ctx.session = await ctx.honcho_client.aio.session(id=ctx.session_id) + ctx.session = await ctx.honcho_client.aio.session( + id=ctx.session_id, configuration=self._get_session_configuration() + ) await ctx.session.aio.add_peers( [ diff --git a/tests/bench/locomo.py b/tests/bench/locomo.py index f91a1285..f9400088 100644 --- a/tests/bench/locomo.py +++ b/tests/bench/locomo.py @@ -256,7 +256,9 @@ class LoCoMoRunner(BaseRunner[ConversationResult]): peer_a = ctx.peers["speaker_a"] peer_b = ctx.peers["speaker_b"] - ctx.session = await ctx.honcho_client.aio.session(id=ctx.session_id) + 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( diff --git a/tests/bench/longmem.py b/tests/bench/longmem.py index 7271e189..09a417c8 100644 --- a/tests/bench/longmem.py +++ b/tests/bench/longmem.py @@ -206,7 +206,9 @@ class LongMemEvalRunner(BaseRunner[TestResult]): if self.merge_sessions: # Create a single merged session - ctx.session = await ctx.honcho_client.aio.session(id=ctx.session_id) + ctx.session = await ctx.honcho_client.aio.session( + id=ctx.session_id, configuration=self._get_session_configuration() + ) if is_assistant_type: await ctx.session.aio.add_peers( @@ -301,7 +303,9 @@ class LongMemEvalRunner(BaseRunner[TestResult]): for session_date, session_id, session_messages in zip( parsed_dates, haystack_session_ids, haystack_sessions, strict=True ): - session = await ctx.honcho_client.aio.session(id=session_id) + session = await ctx.honcho_client.aio.session( + id=session_id, configuration=self._get_session_configuration() + ) if is_assistant_type: await session.aio.add_peers( diff --git a/tests/bench/run_tests.py b/tests/bench/run_tests.py index 13ed113b..fcf141de 100644 --- a/tests/bench/run_tests.py +++ b/tests/bench/run_tests.py @@ -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! diff --git a/tests/bench/runner_common.py b/tests/bench/runner_common.py index a12b4834..dea8e544 100644 --- a/tests/bench/runner_common.py +++ b/tests/bench/runner_common.py @@ -20,6 +20,7 @@ from typing import Any, Generic, Literal, TypeVar 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 @@ -49,6 +50,7 @@ class RunnerConfig: 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( @@ -70,6 +72,7 @@ class RunnerConfig: api_key=args.api_key, skip_dream=args.skip_dream, json_output=args.json_output, + max_concurrent=args.max_concurrent, ) @@ -173,6 +176,13 @@ def add_common_arguments(parser: argparse.ArgumentParser) -> None: 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: """ @@ -190,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 @@ -309,6 +322,10 @@ class BaseRunner(ABC, Generic[ResultT]): 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 + ) # ------------------------------------------------------------------------- # Abstract methods - must be implemented by subclasses @@ -410,6 +427,8 @@ class BaseRunner(ABC, Generic[ResultT]): 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] = [] @@ -425,10 +444,10 @@ class BaseRunner(ABC, Generic[ResultT]): print(f"Processing batch {batch_num}/{total_batches} ({len(batch)} items)") print(f"{'=' * 60}") - # Run items in batch concurrently + # Run items in batch concurrently (with optional rate limiting) batch_results = await asyncio.gather( *[ - self.execute_item(item, self._get_honcho_url(i + idx)) + self._execute_item_with_limit(item, self._get_honcho_url(i + idx)) for idx, item in enumerate(batch) ] ) @@ -442,6 +461,13 @@ class BaseRunner(ABC, Generic[ResultT]): 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. @@ -591,6 +617,10 @@ class BaseRunner(ABC, Generic[ResultT]): 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.""" if self.config.base_url: From 6c3f671636fe6e3ec30518c6e5ad8bb9584b120b Mon Sep 17 00:00:00 2001 From: Benjamin McCormick Date: Fri, 23 Jan 2026 17:57:58 -0500 Subject: [PATCH 3/3] fix: clean Redis URL by removing query parameters in cache client and update flush mode logging in benchmark runner --- src/cache/client.py | 8 ++++++-- tests/bench/runner_common.py | 9 +++++---- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/cache/client.py b/src/cache/client.py index 0381dd4a..0b8480ca 100644 --- a/src/cache/client.py +++ b/src/cache/client.py @@ -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" diff --git a/tests/bench/runner_common.py b/tests/bench/runner_common.py index dea8e544..0cfcc79b 100644 --- a/tests/bench/runner_common.py +++ b/tests/bench/runner_common.py @@ -623,13 +623,14 @@ class BaseRunner(ABC, Generic[ResultT]): async def _flush_deriver_queue(self) -> None: """Enable deriver flush mode to bypass batch token threshold.""" - if self.config.base_url: - print("Skipping flush mode (remote instance)") + 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=60) - print("Enabled deriver flush mode") + 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()