407 lines
15 KiB
Python
407 lines
15 KiB
Python
"""
|
|
Common utilities for LongMemEval test runners.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
from collections.abc import Sequence
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import tiktoken
|
|
from openai import AsyncOpenAI
|
|
from typing_extensions import TypedDict
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class BaseQueryResult(TypedDict):
|
|
"""Base type definition for query execution results."""
|
|
|
|
question: str
|
|
expected_answer: str
|
|
actual_response: str
|
|
judgment: dict[str, Any]
|
|
|
|
|
|
class BaseTestResult(TypedDict):
|
|
"""Base type definition for test execution results."""
|
|
|
|
question_id: str
|
|
question_type: str
|
|
passed: bool
|
|
error: str | None
|
|
start_time: float
|
|
end_time: float
|
|
duration_seconds: float
|
|
output_lines: list[str]
|
|
|
|
|
|
def format_duration(total_seconds: float) -> str:
|
|
"""Format a duration in seconds into a human-readable string.
|
|
|
|
If the duration is at least one minute, this returns a string in the
|
|
form "XmYYs" with zero-padded seconds. Otherwise, it returns the
|
|
duration in seconds with two decimal places, e.g., "12.34s".
|
|
|
|
Args:
|
|
total_seconds: The duration in seconds.
|
|
|
|
Returns:
|
|
A formatted duration 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_total_tokens(haystack_sessions: list[list[dict[str, str]]]) -> int:
|
|
"""Calculate total tokens from all messages in all sessions.
|
|
|
|
Args:
|
|
haystack_sessions: List of sessions, each containing messages
|
|
|
|
Returns:
|
|
Total number of tokens across all messages
|
|
"""
|
|
tokenizer = tiktoken.get_encoding("o200k_base")
|
|
total_tokens = 0
|
|
|
|
for session_messages in haystack_sessions:
|
|
for msg in session_messages:
|
|
content = msg.get("content", "")
|
|
try:
|
|
total_tokens += len(
|
|
tokenizer.encode(
|
|
content,
|
|
disallowed_special=(
|
|
tokenizer.special_tokens_set - {"<|endoftext|>"}
|
|
),
|
|
)
|
|
)
|
|
except Exception:
|
|
total_tokens += len(content) // 4
|
|
logger.warning(
|
|
f"Error tokenizing content. Using rough estimate of {len(content) // 4} tokens"
|
|
)
|
|
|
|
return total_tokens
|
|
|
|
|
|
def parse_longmemeval_date(date_str: str) -> datetime:
|
|
"""Parse longmemeval date format to datetime.
|
|
|
|
Args:
|
|
date_str: Date string in format "YYYY/MM/DD (Day) HH:MM"
|
|
|
|
Returns:
|
|
Parsed datetime object
|
|
|
|
Raises:
|
|
ValueError: If date format is invalid
|
|
"""
|
|
try:
|
|
# Extract the date and time parts, ignoring the day name in parentheses
|
|
# Format: "2023/05/20 (Sat) 02:21"
|
|
parts = date_str.split(") ")
|
|
if len(parts) != 2:
|
|
raise ValueError(f"Invalid date format: {date_str}")
|
|
|
|
date_part = parts[0].split(" (")[0] # "2023/05/20"
|
|
time_part = parts[1] # "02:21"
|
|
|
|
# Combine and parse
|
|
datetime_str = f"{date_part} {time_part}"
|
|
return datetime.strptime(datetime_str, "%Y/%m/%d %H:%M")
|
|
except (ValueError, IndexError) as e:
|
|
raise ValueError(f"Failed to parse date '{date_str}': {e}") from e
|
|
|
|
|
|
def load_test_file(test_file: Path) -> list[dict[str, Any]]:
|
|
"""Load longmemeval test definitions from a JSON file.
|
|
|
|
Args:
|
|
test_file: Path to the JSON test file
|
|
|
|
Returns:
|
|
List of test question dictionaries
|
|
"""
|
|
with open(test_file) as f:
|
|
return json.load(f)
|
|
|
|
|
|
def _build_judge_prompt(
|
|
question_type: str,
|
|
question: str,
|
|
answer: str,
|
|
response: str,
|
|
question_id: str,
|
|
) -> str:
|
|
"""Build the judge prompt matching the official LongMemEval evaluation code.
|
|
|
|
Based on get_anscheck_prompt() from the official LongMemEval repository.
|
|
|
|
Args:
|
|
question_type: Type of question being evaluated
|
|
question: The question asked
|
|
answer: Expected answer from the test
|
|
response: Actual response from the system under test
|
|
question_id: Question ID (used to detect abstention questions)
|
|
|
|
Returns:
|
|
The complete prompt for the judge model
|
|
"""
|
|
# Check for abstention questions (have '_abs' in question_id)
|
|
if "_abs" in question_id:
|
|
return (
|
|
"I will give you an unanswerable question, an explanation, and a response "
|
|
"from a model. Please answer yes if the model correctly identifies the "
|
|
"question as unanswerable. The model could say that the information is "
|
|
"incomplete, or some other information is given but the asked information "
|
|
f"is not.\n\nQuestion: {question}\n\nExplanation: {answer}\n\n"
|
|
f"Model Response: {response}\n\nDoes the model correctly identify the "
|
|
"question as unanswerable? Answer yes or no only."
|
|
)
|
|
|
|
# Standard prompts by question type
|
|
if question_type in (
|
|
"single-session-user",
|
|
"single-session-assistant",
|
|
"multi-session",
|
|
):
|
|
return (
|
|
"I will give you a question, a correct answer, and a response from a model. "
|
|
"Please answer yes if the response contains the correct answer. Otherwise, "
|
|
"answer no. If the response is equivalent to the correct answer or contains "
|
|
"all the intermediate steps to get the correct answer, you should also answer "
|
|
"yes. If the response only contains a subset of the information required by "
|
|
f"the answer, answer no. \n\nQuestion: {question}\n\nCorrect Answer: {answer}"
|
|
f"\n\nModel Response: {response}\n\nIs the model response correct? Answer yes or no only."
|
|
)
|
|
elif question_type == "temporal-reasoning":
|
|
return (
|
|
"I will give you a question, a correct answer, and a response from a model. "
|
|
"Please answer yes if the response contains the correct answer. Otherwise, "
|
|
"answer no. If the response is equivalent to the correct answer or contains "
|
|
"all the intermediate steps to get the correct answer, you should also answer "
|
|
"yes. If the response only contains a subset of the information required by "
|
|
"the answer, answer no. In addition, do not penalize off-by-one errors for "
|
|
"the number of days. If the question asks for the number of days/weeks/months, "
|
|
"etc., and the model makes off-by-one errors (e.g., predicting 19 days when "
|
|
"the answer is 18), the model's response is still correct. \n\n"
|
|
f"Question: {question}\n\nCorrect Answer: {answer}\n\nModel Response: {response}"
|
|
"\n\nIs the model response correct? Answer yes or no only."
|
|
)
|
|
elif question_type == "knowledge-update":
|
|
return (
|
|
"I will give you a question, a correct answer, and a response from a model. "
|
|
"Please answer yes if the response contains the correct answer. Otherwise, "
|
|
"answer no. If the response contains some previous information along with an "
|
|
"updated answer, the response should be considered as correct as long as the "
|
|
f"updated answer is the required answer.\n\nQuestion: {question}\n\n"
|
|
f"Correct Answer: {answer}\n\nModel Response: {response}\n\n"
|
|
"Is the model response correct? Answer yes or no only."
|
|
)
|
|
elif question_type == "single-session-preference":
|
|
return (
|
|
"I will give you a question, a rubric for desired personalized response, "
|
|
"and a response from a model. Please answer yes if the response satisfies "
|
|
"the desired response. Otherwise, answer no. The model does not need to "
|
|
"reflect all the points in the rubric. The response is correct as long as "
|
|
"it recalls and utilizes the user's personal information correctly.\n\n"
|
|
f"Question: {question}\n\nRubric: {answer}\n\nModel Response: {response}"
|
|
"\n\nIs the model response correct? Answer yes or no only."
|
|
)
|
|
else:
|
|
# Default case (same as multi-session)
|
|
return (
|
|
"I will give you a question, a correct answer, and a response from a model. "
|
|
"Please answer yes if the response contains the correct answer. Otherwise, "
|
|
"answer no. If the response is equivalent to the correct answer or contains "
|
|
"all the intermediate steps to get the correct answer, you should also answer "
|
|
"yes. If the response only contains a subset of the information required by "
|
|
f"the answer, answer no. \n\nQuestion: {question}\n\nCorrect Answer: {answer}"
|
|
f"\n\nModel Response: {response}\n\nIs the model response correct? Answer yes or no only."
|
|
)
|
|
|
|
|
|
async def judge_response(
|
|
openai_client: AsyncOpenAI,
|
|
question: str,
|
|
expected_answer: str,
|
|
actual_response: str,
|
|
question_type: str = "default",
|
|
question_id: str = "",
|
|
) -> dict[str, Any]:
|
|
"""Use GPT-4o to judge if the actual response matches the expected answer.
|
|
|
|
Uses the exact prompt format from the official LongMemEval evaluation code
|
|
(evaluate_qa.py) to ensure consistent evaluation.
|
|
|
|
Args:
|
|
openai_client: OpenAI client instance
|
|
question: The question asked
|
|
expected_answer: Expected answer from the test
|
|
actual_response: Actual response from the system under test
|
|
question_type: Type of question (temporal-reasoning, knowledge-update,
|
|
single-session-preference, single-session-user,
|
|
single-session-assistant, multi-session)
|
|
question_id: Question ID (used to detect abstention questions with '_abs')
|
|
|
|
Returns:
|
|
Judgment result with pass/fail and reasoning
|
|
"""
|
|
try:
|
|
prompt = _build_judge_prompt(
|
|
question_type, question, expected_answer, actual_response, question_id
|
|
)
|
|
|
|
response = await openai_client.chat.completions.create(
|
|
model="gpt-4o-2024-08-06",
|
|
max_tokens=10,
|
|
temperature=0,
|
|
n=1,
|
|
messages=[{"role": "user", "content": prompt}],
|
|
)
|
|
|
|
if not response.choices:
|
|
raise ValueError("OpenAI returned empty response")
|
|
|
|
eval_response = response.choices[0].message.content
|
|
if eval_response is None:
|
|
raise ValueError("No text content in response")
|
|
|
|
# Match official evaluation: check if "yes" appears in lowercased response
|
|
passed = "yes" in eval_response.lower()
|
|
|
|
return {
|
|
"passed": passed,
|
|
"reasoning": eval_response.strip(),
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error judging response: {e}")
|
|
# Fallback to simple string matching
|
|
is_correct = expected_answer.lower() in actual_response.lower()
|
|
return {
|
|
"passed": is_correct,
|
|
"reasoning": f"Fallback string matching due to error: {'Match found' if is_correct else 'No match found'}",
|
|
}
|
|
|
|
|
|
def filter_questions(
|
|
questions: list[dict[str, Any]],
|
|
test_file: Path,
|
|
question_id: str | None = None,
|
|
test_count: int | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
"""Filter questions by question_id and/or test_count.
|
|
|
|
Args:
|
|
questions: List of question dictionaries
|
|
test_file: Path to test file (for logging)
|
|
question_id: Optional question_id to filter to
|
|
test_count: Optional limit on number of questions
|
|
|
|
Returns:
|
|
Filtered list of questions
|
|
"""
|
|
# Filter by question_id if specified
|
|
if question_id is not None:
|
|
original_count = len(questions)
|
|
questions = [q for q in questions if q.get("question_id") == question_id]
|
|
if not questions:
|
|
print(
|
|
f"Error: No question found with question_id '{question_id}' in {test_file}"
|
|
)
|
|
return []
|
|
print(
|
|
f"filtering to question_id '{question_id}' ({len(questions)}/{original_count} {'question' if len(questions) == 1 else 'questions'})"
|
|
)
|
|
|
|
# Limit to first N questions if test_count is specified
|
|
if test_count is not None and test_count > 0:
|
|
questions = questions[:test_count]
|
|
print(
|
|
f"limiting to first {len(questions)} {'question' if len(questions) == 1 else 'questions'} from {test_file}"
|
|
)
|
|
|
|
return questions
|
|
|
|
|
|
def calculate_type_statistics(
|
|
results: Sequence[Any],
|
|
) -> dict[str, dict[str, int | float]]:
|
|
"""Calculate pass/fail statistics grouped by question type.
|
|
|
|
Args:
|
|
results: List of test results
|
|
|
|
Returns:
|
|
Dictionary mapping question type to statistics
|
|
"""
|
|
type_stats: dict[str, dict[str, int | float]] = {}
|
|
for result in results:
|
|
q_type = result["question_type"]
|
|
if q_type not in type_stats:
|
|
type_stats[q_type] = {"total": 0, "passed": 0, "failed": 0}
|
|
type_stats[q_type]["total"] += 1
|
|
if result.get("passed", False):
|
|
type_stats[q_type]["passed"] += 1
|
|
else:
|
|
type_stats[q_type]["failed"] += 1
|
|
|
|
# Add success rates
|
|
for q_type in type_stats:
|
|
stats = type_stats[q_type]
|
|
stats["success_rate"] = (
|
|
(stats["passed"] / stats["total"]) * 100 if stats["total"] > 0 else 0
|
|
)
|
|
|
|
return type_stats
|
|
|
|
|
|
def calculate_timing_statistics(
|
|
results: Sequence[Any], total_elapsed_seconds: float
|
|
) -> dict[str, Any]:
|
|
"""Calculate timing statistics from test results.
|
|
|
|
Args:
|
|
results: List of test results
|
|
total_elapsed_seconds: Total elapsed time for all tests
|
|
|
|
Returns:
|
|
Dictionary of timing statistics
|
|
"""
|
|
durations = [r["duration_seconds"] for r in results]
|
|
return {
|
|
"total_duration_seconds": total_elapsed_seconds,
|
|
"individual_test_durations": {
|
|
"min_seconds": min(durations) if durations else 0,
|
|
"max_seconds": max(durations) if durations else 0,
|
|
"mean_seconds": sum(durations) / len(durations) if durations else 0,
|
|
"median_seconds": sorted(durations)[len(durations) // 2]
|
|
if durations
|
|
else 0,
|
|
},
|
|
}
|
|
|
|
|
|
def write_json_summary(summary: dict[str, Any], output_file: Path) -> None:
|
|
"""Write a JSON summary to a file.
|
|
|
|
Args:
|
|
summary: Summary dictionary to write
|
|
output_file: Path to output file
|
|
"""
|
|
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}")
|