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