From 98d1841032e1eb1c420349fa55a2538257551c60 Mon Sep 17 00:00:00 2001 From: Sehlani042 <257166922+Sehlani042@users.noreply.github.com> Date: Thu, 2 Jul 2026 08:42:29 +0800 Subject: [PATCH] Add optional agent trajectory export --- README.md | 9 ++- benchmark/runner.py | 43 ++++++++++-- benchmark/trajectory_writer.py | 53 +++++++++++++++ tests/test_runner_trajectory_export.py | 90 ++++++++++++++++++++++++++ tests/test_trajectory_writer.py | 50 ++++++++++++++ 5 files changed, 239 insertions(+), 6 deletions(-) create mode 100644 benchmark/trajectory_writer.py create mode 100644 tests/test_runner_trajectory_export.py create mode 100644 tests/test_trajectory_writer.py diff --git a/README.md b/README.md index 168d00c..006936d 100755 --- a/README.md +++ b/README.md @@ -127,6 +127,12 @@ source .env python run_benchmark.py --models gpt-oss-20b \ --tasks-file tasks/mcpbench_tasks_multi_3server_runner_format.json +## optional: save per-task agent trajectories for debugging or analysis +source .env +python run_benchmark.py --models gpt-oss-20b \ +--tasks-file tasks/mcpbench_tasks_single_runner_format.json \ +--trajectory-output-dir trajectories + ``` ### Optional: Add other model providers @@ -280,7 +286,8 @@ mcp-bench/ │ ├── evaluator.py # LLM-as-judge evaluation metrics │ ├── runner.py # Benchmark orchestrator │ ├── results_aggregator.py # Results aggregation and statistics -│ └── results_formatter.py # Results formatting and display +│ ├── results_formatter.py # Results formatting and display +│ └── trajectory_writer.py # Per-task trajectory text export ├── config/ # Configuration management │ ├── __init__.py │ ├── benchmark_config.yaml # Benchmark configuration diff --git a/benchmark/runner.py b/benchmark/runner.py index cb97637..5367a33 100644 --- a/benchmark/runner.py +++ b/benchmark/runner.py @@ -36,6 +36,7 @@ from benchmark.evaluator import TaskEvaluator from benchmark.results_aggregator import ResultsAggregator from benchmark.results_formatter import ResultsFormatter, execution_results_to_text +from benchmark.trajectory_writer import write_task_trajectory from utils.local_server_config import LocalServerConfigLoader import config.config_loader as config_loader @@ -176,6 +177,7 @@ def __init__( filter_problematic_tools: Optional[bool] = None, concurrent_summarization: Optional[bool] = None, use_fuzzy_descriptions: Optional[bool] = None, + trajectory_output_dir: Optional[str] = None, # Dependency injection parameters local_config_loader: Optional[LocalServerConfigLoader] = None, aggregator: Optional[ResultsAggregator] = None, @@ -197,6 +199,7 @@ def __init__( self.filter_problematic_tools = filter_problematic_tools if filter_problematic_tools is not None else config_loader.is_problematic_tools_filter_enabled() self.concurrent_summarization = concurrent_summarization if concurrent_summarization is not None else config_loader.is_concurrent_summarization_enabled() self.use_fuzzy_descriptions = use_fuzzy_descriptions if use_fuzzy_descriptions is not None else config_loader.use_fuzzy_descriptions() + self.trajectory_output_dir = trajectory_output_dir self.enable_concrete_description_ref = config_loader.is_concrete_description_ref_enabled() self.commands_config = None @@ -1026,10 +1029,27 @@ async def _evaluate_task_result(self, task_execution_info: Dict[str, Any], execu if evaluation: logger.info("\n TASK METRICS:") self.formatter.format_single_task_report(task_id, evaluation, []) - + + trajectory_file = None + if self.trajectory_output_dir: + trajectory_text = evaluator.llm_judge._create_execution_summary( + result.get('execution_results', []), + result.get('total_rounds', 0), + accumulated_info + ) + trajectory_file = write_task_trajectory( + output_dir=self.trajectory_output_dir, + task_id=task_id, + model_name=model_name, + server_name=server_name, + task_description=task_description, + trajectory_text=trajectory_text + ) + logger.info(f"Trajectory saved to {trajectory_file}") + logger.info("="*80 + "\n") - - return { + + final_result = { 'task_id': task_id, 'server_name': server_name, 'model_name': model_name, @@ -1047,6 +1067,9 @@ async def _evaluate_task_result(self, task_execution_info: Dict[str, Any], execu 'total_prompt_tokens': result.get('total_prompt_tokens', 0), 'total_tokens': result.get('total_tokens', 0) } + if trajectory_file: + final_result['trajectory_file'] = trajectory_file + return final_result def parse_arguments(): @@ -1081,6 +1104,12 @@ def parse_arguments(): metavar='FILE', help='Output file for results (default: auto-generated timestamp name)' ) + + parser.add_argument( + '--trajectory-output-dir', + metavar='DIR', + help='Directory for per-task agent trajectory text files (default: disabled)' + ) parser.add_argument( '--tasks-file', @@ -1194,7 +1223,8 @@ def _create_runner_and_get_models(args, tasks_file, enable_distraction): enable_judge_stability=not args.disable_judge_stability, filter_problematic_tools=not args.disable_filter_problematic_tools, concurrent_summarization=not args.disable_concurrent_summarization, - use_fuzzy_descriptions=not args.disable_fuzzy + use_fuzzy_descriptions=not args.disable_fuzzy, + trajectory_output_dir=args.trajectory_output_dir ) available_models = list(runner.model_configs.keys()) @@ -1265,6 +1295,9 @@ def _print_configuration(selected_models, available_models, runner, args): if args.output: print(f" Output file: {args.output}") + if args.trajectory_output_dir: + print(f" Trajectory output dir: {args.trajectory_output_dir}") + async def main(): """Main entry point for multi-model benchmark runner""" # Step 1: Parse and validate arguments @@ -1336,4 +1369,4 @@ async def main(): sys.exit(1) if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file + asyncio.run(main()) diff --git a/benchmark/trajectory_writer.py b/benchmark/trajectory_writer.py new file mode 100644 index 0000000..db3bea9 --- /dev/null +++ b/benchmark/trajectory_writer.py @@ -0,0 +1,53 @@ +""" +Utilities for saving per-task MCP-Bench agent trajectories. +""" + +import re +from pathlib import Path +from typing import Union + + +PathLike = Union[str, Path] + + +def sanitize_path_component(value: str) -> str: + """Return a filesystem-safe path component.""" + sanitized = re.sub(r"[^A-Za-z0-9_.-]+", "_", str(value)).strip("._") + return sanitized or "unknown" + + +def build_trajectory_path(output_dir: PathLike, model_name: str, task_id: str) -> Path: + """Build the destination path for a task trajectory file.""" + return ( + Path(output_dir) + / sanitize_path_component(model_name) + / f"{sanitize_path_component(task_id)}.txt" + ) + + +def write_task_trajectory( + output_dir: PathLike, + task_id: str, + model_name: str, + server_name: str, + task_description: str, + trajectory_text: str, +) -> str: + """Write a per-task trajectory text file and return its path.""" + trajectory_path = build_trajectory_path(output_dir, model_name, task_id) + trajectory_path.parent.mkdir(parents=True, exist_ok=True) + + content = "\n".join([ + f"Task ID: {task_id}", + f"Model: {model_name}", + f"Server: {server_name}", + "", + "Task Description:", + task_description or "", + "", + "--- TRAJECTORY ---", + trajectory_text or "", + "", + ]) + trajectory_path.write_text(content, encoding="utf-8") + return str(trajectory_path) diff --git a/tests/test_runner_trajectory_export.py b/tests/test_runner_trajectory_export.py new file mode 100644 index 0000000..96f9c0f --- /dev/null +++ b/tests/test_runner_trajectory_export.py @@ -0,0 +1,90 @@ +import json +import tempfile +import unittest +from pathlib import Path + +from benchmark.runner import BenchmarkRunner + + +class StubJudgeProvider: + async def get_completion(self, system_prompt, user_prompt, max_tokens): + return json.dumps({ + "task_fulfillment_reasoning": "ok", + "grounding_reasoning": "ok", + "tool_appropriateness_reasoning": "ok", + "parameter_accuracy_reasoning": "ok", + "dependency_awareness_reasoning": "ok", + "parallelism_efficiency_reasoning": "ok", + "task_fulfillment": 8, + "grounding": 6, + "tool_appropriateness": 10, + "parameter_accuracy": 4, + "dependency_awareness": 7, + "parallelism_and_efficiency": 3, + }) + + def clean_and_parse_json(self, raw_json): + return json.loads(raw_json) + + +class RunnerTrajectoryExportTest(unittest.IsolatedAsyncioTestCase): + async def test_evaluate_task_result_writes_trajectory_when_output_dir_is_set(self): + with tempfile.TemporaryDirectory() as tmp_dir: + runner = BenchmarkRunner( + judge_provider=StubJudgeProvider(), + trajectory_output_dir=tmp_dir, + ) + + result = await runner._evaluate_task_result( + task_execution_info={ + "task_id": "openapi_explorer_000", + "task_description": "Explore the API.", + "concrete_task_description": "Explore the API.", + "server_name": "openapi_explorer", + }, + execution_result={ + "result": { + "execution_results": [ + { + "tool": "openapi_search", + "success": True, + "parameters": {"query": "pets"}, + } + ], + "solution": "Done", + "total_rounds": 2, + "available_tools": { + "openapi_search": { + "server": "openapi_explorer", + "description": "Search OpenAPI specs", + } + }, + "planning_json_compliance": 1.0, + "accumulated_information": "Round 1 called openapi_search.", + }, + "task_execution_start_time": 0, + "execution_time": 1.0, + "agent_execution_time": 0.5, + }, + model_name="gpt-oss-20b", + server_name="openapi_explorer", + ) + + trajectory_path = Path(result["trajectory_file"]) + self.assertTrue(trajectory_path.exists()) + self.assertEqual( + trajectory_path, + Path(tmp_dir) / "gpt-oss-20b" / "openapi_explorer_000.txt", + ) + + content = trajectory_path.read_text(encoding="utf-8") + self.assertIn("Task ID: openapi_explorer_000", content) + self.assertIn("Model: gpt-oss-20b", content) + self.assertIn("Server: openapi_explorer", content) + self.assertIn("Total rounds: 2", content) + self.assertIn("Tools executed: 1", content) + self.assertIn("Round 1 called openapi_search.", content) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_trajectory_writer.py b/tests/test_trajectory_writer.py new file mode 100644 index 0000000..cb8e5cd --- /dev/null +++ b/tests/test_trajectory_writer.py @@ -0,0 +1,50 @@ +import tempfile +import unittest +from pathlib import Path + +from benchmark.trajectory_writer import build_trajectory_path, write_task_trajectory + + +class TrajectoryWriterTest(unittest.TestCase): + def test_build_trajectory_path_sanitizes_model_and_task_id(self): + path = build_trajectory_path( + "/tmp/trajectories", + model_name="openai/gpt-oss:20b", + task_id="openapi explorer/000", + ) + + self.assertEqual( + path, + Path("/tmp/trajectories/openai_gpt-oss_20b/openapi_explorer_000.txt"), + ) + + def test_write_task_trajectory_creates_parent_dir_and_content(self): + with tempfile.TemporaryDirectory() as tmp_dir: + path = write_task_trajectory( + output_dir=tmp_dir, + task_id="openapi_explorer_000", + model_name="gpt-oss-20b", + server_name="openapi_explorer", + task_description="Explore the API.", + trajectory_text="Total rounds: 2; Tools executed: 1", + ) + + written = Path(path) + self.assertTrue(written.exists()) + self.assertEqual( + written, + Path(tmp_dir) / "gpt-oss-20b" / "openapi_explorer_000.txt", + ) + + content = written.read_text(encoding="utf-8") + self.assertIn("Task ID: openapi_explorer_000", content) + self.assertIn("Model: gpt-oss-20b", content) + self.assertIn("Server: openapi_explorer", content) + self.assertIn("Task Description:", content) + self.assertIn("Explore the API.", content) + self.assertIn("--- TRAJECTORY ---", content) + self.assertIn("Total rounds: 2; Tools executed: 1", content) + + +if __name__ == "__main__": + unittest.main()