deepset-mcp 0.0.2__py3-none-any.whl

deepset_mcp/benchmark/runner/agent_benchmark_runner.py

@@ -0,0 +1,561 @@
+import asyncio
+import json
+import logging
+from collections.abc import Callable
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from haystack.dataclasses.chat_message import ChatMessage
+from haystack.dataclasses.streaming_chunk import StreamingChunk
+
+from deepset_mcp.api.client import AsyncDeepsetClient
+from deepset_mcp.benchmark.runner.agent_loader import load_agent
+from deepset_mcp.benchmark.runner.config import BenchmarkConfig
+from deepset_mcp.benchmark.runner.config_loader import (
+    find_all_test_case_paths,
+    load_test_case_by_name,
+)
+from deepset_mcp.benchmark.runner.models import AgentConfig, TestCaseConfig
+from deepset_mcp.benchmark.runner.streaming import StreamingCallbackManager
+from deepset_mcp.benchmark.runner.teardown_actions import teardown_test_case_async
+from deepset_mcp.benchmark.runner.tracing import enable_tracing
+
+logger = logging.getLogger(__name__)
+
+
+class AgentBenchmarkRunner:
+    """Main class for running agent benchmarks against test cases."""
+
+    def __init__(
+        self,
+        agent_config: AgentConfig,
+        benchmark_config: BenchmarkConfig,
+        streaming: bool = True,
+    ):
+        """
+        Initialize the benchmark runner.
+
+        Args:
+            agent_config: Configuration for the agent to test.
+            benchmark_config: Benchmark configuration.
+            streaming: Whether to enable streaming output during agent execution.
+        """
+        self.agent_config = agent_config
+        self.benchmark_config = benchmark_config
+        self.streaming = streaming
+
+        # Create a single timestamp for this benchmark run
+        self.run_timestamp = datetime.now()
+
+        try:
+            secret_key = self.benchmark_config.get_env_var("LANGFUSE_SECRET_KEY")
+            public_key = self.benchmark_config.get_env_var("LANGFUSE_PUBLIC_KEY")
+            logger.info("Langfuse environment variables detected. Enabling tracing with Langfuse.")
+            enable_tracing(secret_key=secret_key, public_key=public_key, name="deepset-mcp")
+        except KeyError:
+            pass
+
+        agent, commit_hash = load_agent(config=agent_config, benchmark_config=benchmark_config)
+
+        self.agent = agent
+        self.commit_hash = commit_hash
+
+        # Create the run ID once for all test cases
+        self.run_id = (
+            f"{self.agent_config.display_name}-{self.commit_hash}_{self.run_timestamp.strftime('%Y%m%d_%H%M%S')}"
+        )
+
+    def _create_streaming_callback(self, test_case_name: str) -> Callable[[StreamingChunk], Any]:
+        """
+        Create a streaming callback function for a specific test case.
+
+        Args:
+            test_case_name: Name of the test case for logging context
+
+        Returns:
+            Callback function for streaming
+        """
+        callback = StreamingCallbackManager()
+
+        async def streaming_callback(chunk: StreamingChunk) -> Any:
+            return await callback(chunk)
+
+        return streaming_callback
+
+    async def run_single_test(self, test_case_name: str) -> dict[str, Any]:
+        """
+        Run the agent against a single test case.
+
+        Args:
+            test_case_name: Name of the test case to run
+
+        Returns:
+            Dictionary containing run results and metadata
+        """
+        logger.info(f"Running test case: {test_case_name}")
+
+        try:
+            # Load test case configuration
+            test_config = load_test_case_by_name(
+                name=test_case_name,
+                task_dir=str(self.benchmark_config.test_case_base_dir)
+                if self.benchmark_config.test_case_base_dir
+                else None,
+            )
+
+            index_yaml_config = test_config.get_index_yaml_text()
+            index_name = test_config.index_name
+            if index_yaml_config and index_name:
+                async with AsyncDeepsetClient(api_key=self.benchmark_config.deepset_api_key) as client:
+                    await client.indexes(workspace=self.benchmark_config.deepset_workspace).create(
+                        name=index_name, yaml_config=index_yaml_config
+                    )
+
+            pre_agent_validation = None
+            query_yaml_config = test_config.get_query_yaml_text()
+            query_name = test_config.query_name
+            if query_yaml_config and query_name:
+                async with AsyncDeepsetClient(api_key=self.benchmark_config.deepset_api_key) as client:
+                    await client.pipelines(workspace=self.benchmark_config.deepset_workspace).create(
+                        name=query_name, yaml_config=query_yaml_config
+                    )
+                    pre_agent_validation = await client.pipelines(
+                        workspace=self.benchmark_config.deepset_workspace
+                    ).validate(yaml_config=query_yaml_config)
+
+            # Prepare streaming callback if streaming is enabled
+            streaming_callback = None
+            if self.streaming:
+                streaming_callback = self._create_streaming_callback(test_case_name)
+                print(f"\n🤖 [{test_case_name}] Agent starting...\n")
+
+            agent_output = await self.agent.run_async(
+                messages=[ChatMessage.from_user(test_config.prompt)], streaming_callback=streaming_callback
+            )
+
+            if self.streaming:
+                print(f"\n\n✅ [{test_case_name}] Agent completed.\n")
+
+            post_agent_validation = None
+            if query_name:
+                async with AsyncDeepsetClient(api_key=self.benchmark_config.deepset_api_key) as client:
+                    pipeline_resource = client.pipelines(workspace=self.benchmark_config.deepset_workspace)
+                    updated_pipeline = await pipeline_resource.get(pipeline_name=query_name)
+                    assert updated_pipeline.yaml_config, "Pipeline YAML config not found"
+                    post_agent_validation = await pipeline_resource.validate(yaml_config=updated_pipeline.yaml_config)
+
+            # Process the results
+            processed_data = self._format_results(
+                agent_output=agent_output,
+                test_config=test_config,
+                is_pre_agent_valid=pre_agent_validation.valid if pre_agent_validation else None,
+                is_post_agent_valid=post_agent_validation.valid if post_agent_validation else None,
+                post_yaml=updated_pipeline.yaml_config if post_agent_validation else None,
+            )
+
+            # Save results to filesystem
+            test_dir = self._save_run_results(
+                processed_data=processed_data,
+                test_case_name=test_case_name,
+                output_base_dir=self.benchmark_config.output_dir,
+            )
+
+            logger.info(f"Test case {test_case_name} completed. Results saved to: {test_dir}")
+
+            return {
+                "status": "success",
+                "test_case": test_case_name,
+                "output_dir": str(test_dir),
+                "processed_data": processed_data,
+            }
+
+        except Exception as e:
+            logger.error(f"Error running test case {test_case_name}: {e}")
+            return {"status": "error", "test_case": test_case_name, "error": str(e)}
+
+    async def run_single_test_with_cleanup(self, test_case_name: str) -> dict[str, Any]:
+        """
+        Run a single test case with automatic cleanup of created resources.
+
+        Args:
+            test_case_name: Name of the test case to run
+
+        Returns:
+            Dictionary containing run results and metadata
+        """
+        result = await self.run_single_test(test_case_name)
+
+        # Perform cleanup regardless of test result
+        try:
+            # Load test config for cleanup
+            test_config = load_test_case_by_name(
+                name=test_case_name,
+                task_dir=self.benchmark_config.test_case_base_dir if self.benchmark_config.test_case_base_dir else None,
+            )
+
+            # Cleanup resources
+            await teardown_test_case_async(
+                test_cfg=test_config,
+                workspace_name=self.benchmark_config.deepset_workspace,
+                api_key=self.benchmark_config.deepset_api_key,
+            )
+
+            logger.info(f"Cleanup completed for test case: {test_case_name}")
+            result["cleanup_status"] = "success"
+
+        except Exception as e:
+            logger.error(f"Error during cleanup for {test_case_name}: {e}")
+            result["cleanup_status"] = "error"
+            result["cleanup_error"] = str(e)
+
+        return result
+
+    def run_all_tests(self, test_case_path: Path) -> tuple[list[dict[str, Any]], dict[str, Any]]:
+        """
+        Run the agent against all available test cases.
+
+        Args:
+            test_case_path: Directory containing test case files
+
+        Returns:
+            Tuple of (test results list, summary statistics dict)
+        """
+        # Find all test case files
+        test_paths = find_all_test_case_paths(test_case_path)
+
+        if not test_paths:
+            logger.warning(f"No test cases found in {test_case_path}")
+            empty_summary = {
+                "total_prompt_tokens": 0,
+                "total_completion_tokens": 0,
+                "tests_completed": 0,
+                "tests_failed": 0,
+                "avg_tool_calls": 0.0,
+                "pass_rate_percent": 0.0,
+                "fail_rate_percent": 0.0,
+            }
+            return [], empty_summary
+
+        logger.info(f"Found {len(test_paths)} test cases to run")
+
+        # Run tests sequentially with cleanup
+        results = []
+        for test_path in test_paths:
+            test_name = test_path.stem
+            result = asyncio.run(self.run_single_test_with_cleanup(test_name))
+            results.append(result)
+
+        # Create run summary CSV and get summary data
+        summary_data = self._create_run_summary_csv(results)
+
+        return results, summary_data
+
+    async def run_all_tests_async(
+        self,
+        test_case_path: Path,
+        concurrency: int = 1,  # Keep concurrency low to avoid resource conflicts
+    ) -> tuple[list[dict[str, Any]], dict[str, Any]]:
+        """
+        Run all test cases asynchronously with controlled concurrency.
+
+        Args:
+            test_case_path: Directory containing test case files
+            concurrency: Number of concurrent test runs (default: 1 for safety)
+
+        Returns:
+            Tuple of (test results list, summary statistics dict)
+        """
+        # Find all test case files
+        test_paths = find_all_test_case_paths(test_case_path)
+
+        if not test_paths:
+            logger.warning(f"No test cases found in {test_case_path}")
+            empty_summary = {
+                "total_prompt_tokens": 0,
+                "total_completion_tokens": 0,
+                "tests_completed": 0,
+                "tests_failed": 0,
+                "avg_tool_calls": 0.0,
+                "pass_rate_percent": 0.0,
+                "fail_rate_percent": 0.0,
+            }
+            return [], empty_summary
+
+        logger.info(f"Found {len(test_paths)} test cases to run with concurrency={concurrency}")
+
+        # Create semaphore for concurrency control
+        semaphore = asyncio.Semaphore(concurrency)
+
+        async def run_with_semaphore(test_name: str) -> dict[str, Any]:
+            async with semaphore:
+                return await self.run_single_test_with_cleanup(test_name)
+
+        # Create tasks for all test cases
+        tasks = [asyncio.create_task(run_with_semaphore(test_path.stem)) for test_path in test_paths]
+
+        # Wait for all tasks to complete
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        # Handle any exceptions
+        processed_results: list[dict[str, Any]] = []
+        for i, result in enumerate(results):
+            if isinstance(result, Exception):
+                test_name = test_paths[i].stem
+                logger.error(f"Exception in test {test_name}: {result}")
+                processed_results.append({"status": "error", "test_case": test_name, "error": str(result)})
+            else:
+                processed_results.append(result)  # type: ignore
+
+        # Create run summary CSV and get summary data
+        summary_data = self._create_run_summary_csv(processed_results)
+
+        return processed_results, summary_data
+
+    def _format_results(
+        self,
+        agent_output: dict[str, Any],
+        test_config: TestCaseConfig,
+        is_pre_agent_valid: bool | None = None,
+        is_post_agent_valid: bool | None = None,
+        post_yaml: str | None = None,
+    ) -> dict[str, Any]:
+        """Format the agent output and metadata for saving to file."""
+        return {
+            "metadata": {
+                "commit_hash": self.commit_hash,
+                "agent_display_name": self.agent_config.display_name,
+                "test_case_name": test_config.name,
+                "timestamp": self.run_timestamp.isoformat(),
+                "run_id": self.run_id,
+            },
+            "validation": {
+                "pre_validation": "PASS"
+                if is_pre_agent_valid is True
+                else ("FAIL" if is_pre_agent_valid is False else None),
+                "post_validation": "PASS"
+                if is_post_agent_valid is True
+                else ("FAIL" if is_post_agent_valid is False else None),
+            },
+            "messages": {
+                "serialized": [message.to_dict() for message in agent_output["messages"]],
+                "stats": self._extract_assistant_message_stats(agent_output["messages"]),
+            },
+            "pipeline_yaml": post_yaml,
+        }
+
+    def _create_run_summary_csv(self, results: list[dict[str, Any]]) -> dict[str, Any]:
+        """
+        Create a summary CSV file for the entire benchmark run.
+
+        Args:
+            results: List of test results from the benchmark run
+
+        Returns:
+            Dictionary containing the summary statistics
+        """
+        # Initialize counters
+        total_prompt_tokens = 0
+        total_completion_tokens = 0
+        tests_completed = 0
+        tests_failed = 0
+        total_tool_calls = 0
+        tests_with_validation = 0
+        validation_passes = 0
+
+        for result in results:
+            if result["status"] == "success":
+                tests_completed += 1
+                processed_data = result["processed_data"]
+
+                # Sum token counts
+                stats = processed_data["messages"]["stats"]
+                total_prompt_tokens += stats["total_prompt_tokens"]
+                total_completion_tokens += stats["total_completion_tokens"]
+                total_tool_calls += stats["total_tool_calls"]
+
+                # Check validation results (exclude cases where pre or post validation is None)
+                validation = processed_data["validation"]
+                pre_val = validation["pre_validation"]
+                post_val = validation["post_validation"]
+
+                # Only count validation if both pre and post validation exist
+                if pre_val is not None and post_val is not None:
+                    tests_with_validation += 1
+
+                    # Expected pattern: pre_validation should FAIL, post_validation should PASS
+                    # This indicates the agent successfully fixed the broken pipeline
+                    if pre_val == "FAIL" and post_val == "PASS":
+                        validation_passes += 1
+            else:
+                tests_failed += 1
+
+        # Calculate averages and rates
+        avg_tool_calls = total_tool_calls / tests_completed if tests_completed > 0 else 0
+        pass_rate = (validation_passes / tests_with_validation * 100) if tests_with_validation > 0 else 0
+        fail_rate = 100 - pass_rate if tests_with_validation > 0 else 0
+
+        # Create summary dict
+        summary_data = {
+            "total_prompt_tokens": total_prompt_tokens,
+            "total_completion_tokens": total_completion_tokens,
+            "tests_completed": tests_completed,
+            "tests_failed": tests_failed,
+            "avg_tool_calls": round(avg_tool_calls, 2),
+            "pass_rate_percent": round(pass_rate, 2),
+            "fail_rate_percent": round(fail_rate, 2),
+        }
+
+        # Create CSV content
+        csv_data = [
+            "total_prompt_tokens,total_completion_tokens,tests_completed,tests_failed,avg_tool_calls,pass_rate_percent,fail_rate_percent",
+            f"{total_prompt_tokens},{total_completion_tokens},{tests_completed},{tests_failed},{avg_tool_calls:.2f},{pass_rate:.2f},{fail_rate:.2f}",
+        ]
+
+        # Save to main run directory
+        run_dir = self.benchmark_config.output_dir / "agent_runs" / self.run_id
+        run_dir.mkdir(exist_ok=True, parents=True)
+        summary_file = run_dir / "run_summary.csv"
+
+        with open(summary_file, "w", encoding="utf-8") as f:
+            f.write("\n".join(csv_data))
+
+        logger.info(f"Run summary saved to: {summary_file}")
+
+        return summary_data
+
+    @staticmethod
+    def _extract_assistant_message_stats(messages: list[ChatMessage]) -> dict[str, str | int]:
+        """
+        Extract statistics from ChatMessage objects with role=assistant.
+
+        Args:
+            messages: List of ChatMessage objects
+
+        Returns:
+            Dict containing aggregated statistics and model info
+        """
+        total_tool_calls = 0
+        total_prompt_tokens = 0
+        total_completion_tokens = 0
+        model = None
+
+        for message in messages:
+            # Only process assistant messages
+            if not message.is_from("assistant"):
+                continue
+
+            # Count tool calls
+            tool_calls = message.tool_calls
+            total_tool_calls += len(tool_calls)
+
+            # Extract token counts and model from meta
+            meta = message.meta
+            if "usage" in meta:
+                usage = meta["usage"]
+                prompt_tokens = usage.get("prompt_tokens")
+                total_prompt_tokens += prompt_tokens if prompt_tokens is not None else 0
+                completion_tokens = usage.get("completion_tokens")
+                total_completion_tokens += completion_tokens if completion_tokens is not None else 0
+
+            # Extract model (should be consistent across messages)
+            if "model" in meta and model is None:
+                model = meta["model"]
+
+        return {
+            "total_tool_calls": total_tool_calls,
+            "total_prompt_tokens": total_prompt_tokens,
+            "total_completion_tokens": total_completion_tokens,
+            "model": model or "unknown",
+        }
+
+    @staticmethod
+    def _save_run_results(processed_data: dict[str, Any], test_case_name: str, output_base_dir: Path) -> Path:
+        """
+        Save the processed run results to the filesystem.
+
+        Args:
+            processed_data: Output from process_pipeline_result
+            test_case_name: Name of the test case
+            output_base_dir: Base directory for saving results
+
+        Returns:
+            Path to the created test case directory
+        """
+        metadata = processed_data["metadata"]
+        run_dir = output_base_dir / "agent_runs" / metadata["run_id"]
+        test_case_dir: Path = run_dir / test_case_name
+        test_case_dir.mkdir(exist_ok=True, parents=True)
+
+        # Save messages.json
+        messages_file = test_case_dir / "messages.json"
+        with open(messages_file, "w", encoding="utf-8") as f:
+            json.dump(processed_data["messages"]["serialized"], f, indent=2, ensure_ascii=False)
+
+        # Save test_results.csv
+        csv_file = test_case_dir / "test_results.csv"
+        pre_validation = processed_data["validation"]["pre_validation"] or "N/A"
+        post_validation = processed_data["validation"]["post_validation"] or "N/A"
+        csv_data = [
+            "commit,test_case,agent,prompt_tokens,completion_tokens,tool_calls,model,pre_validation,post_validation",
+            f"{metadata['commit_hash']},{test_case_name},{metadata['agent_display_name']},"
+            f"{processed_data['messages']['stats']['total_prompt_tokens']},"
+            f"{processed_data['messages']['stats']['total_completion_tokens']},"
+            f"{processed_data['messages']['stats']['total_tool_calls']},"
+            f"{processed_data['messages']['stats']['model']},"
+            f"{pre_validation},"
+            f"{post_validation}",
+        ]
+
+        with open(csv_file, "w", encoding="utf-8") as f:
+            f.write("\n".join(csv_data))
+
+        # Save post_run_pipeline.yml
+        if processed_data["pipeline_yaml"]:
+            pipeline_file = test_case_dir / "post_run_pipeline.yml"
+            with open(pipeline_file, "w", encoding="utf-8") as f:
+                f.write(processed_data["pipeline_yaml"])
+
+        return test_case_dir
+
+
+def run_agent_benchmark(
+    agent_config: AgentConfig,
+    benchmark_config: BenchmarkConfig,
+    test_case_name: str | None = None,
+    concurrency: int = 1,
+    streaming: bool = False,
+) -> tuple[list[dict[str, Any]], dict[str, Any]]:
+    """
+    Convenience function to run agent benchmarks.
+
+    Args:
+        agent_config_path: Path to agent configuration file
+        benchmark_config: Benchmark configuration.
+        test_case_name: Specific test case to run (if None, runs all)
+        concurrency: Number of concurrent test runs
+        streaming: If True, run in streaming mode
+
+    Returns:
+        List of test results
+    """
+    # Create runner
+    runner = AgentBenchmarkRunner(
+        agent_config=agent_config,
+        benchmark_config=benchmark_config,
+        streaming=streaming,
+    )
+
+    if test_case_name:
+        # Run single test case
+        result = asyncio.run(runner.run_single_test_with_cleanup(test_case_name))
+        results = [result]
+        # Create run summary CSV for single test case
+        summary_data = runner._create_run_summary_csv(results)
+        return results, summary_data
+    else:
+        # Run all test cases
+        if concurrency == 1:
+            return runner.run_all_tests(benchmark_config.test_case_base_dir)
+        else:
+            return asyncio.run(runner.run_all_tests_async(benchmark_config.test_case_base_dir, concurrency))

deepset_mcp/benchmark/runner/agent_loader.py

@@ -0,0 +1,110 @@
+import importlib
+import json
+import os
+import subprocess
+from collections.abc import Callable
+from typing import cast
+
+from haystack.components.agents.agent import Agent
+
+from deepset_mcp.benchmark.runner.config import BenchmarkConfig
+from deepset_mcp.benchmark.runner.models import AgentConfig
+
+
+def load_agent(
+    config: AgentConfig,
+    benchmark_config: BenchmarkConfig,
+    interactive: bool = False,
+) -> tuple[Agent, str | None]:
+    """
+    Load an agent based on the configuration.
+
+    This function:
+    - Loads the agent from either qualified name or JSON file
+    - Checks required environment variables (for qualified name approach)
+    - Collects metadata (timestamp, git commit hash)
+
+    Args:
+        config: AgentConfig instance specifying how to load the agent
+        benchmark_config: BenchmarkConfig instance specifying the benchmark configuration.
+        interactive: Whether to load the agent in interactive mode.
+
+    Returns:
+        LoadedAgent containing the agent instance and metadata
+
+    Raises:
+        ImportError: If qualified function cannot be imported
+        AttributeError: If function doesn't exist in module
+        ValueError: If function is not callable or doesn't return proper tuple
+        FileNotFoundError: If JSON file cannot be found
+        EnvironmentError: If required environment variables are not set
+        json.JSONDecodeError: If JSON file is invalid
+    """
+    # Get git commit hash
+    git_commit_hash = None
+    try:
+        result = subprocess.run(
+            ["git", "rev-parse", "--short", "HEAD"], capture_output=True, text=True, check=True, cwd=os.getcwd()
+        )
+        git_commit_hash = result.stdout.strip()
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        # Git not available or not in a git repo
+        pass
+
+    # Load the agent
+    if config.agent_factory_function:
+        agent_func = _import_factory_from_qualified_name(config.agent_factory_function)
+        if interactive:
+            agent = agent_func(
+                benchmark_config,
+                interactive=True,
+            )
+        else:
+            agent = agent_func(benchmark_config)
+    elif config.agent_json:
+        if interactive:
+            raise ValueError("Interactive mode is not supported for JSON-based agents.")
+        agent = _load_from_json(config.agent_json)
+    else:
+        # This should never happen due to validation, but just in case
+        raise ValueError("No agent source specified")
+
+    is_complete, missing = benchmark_config.check_required_env_vars(config.required_env_vars)
+
+    if not is_complete:
+        raise OSError(f"Required environment variables not set. Missing: {', '.join(missing)}.")
+
+    return agent, git_commit_hash
+
+
+def _import_factory_from_qualified_name(qualified_name: str) -> Callable[..., Agent]:
+    """Load agent from qualified function name."""
+    try:
+        module_path, function_name = qualified_name.rsplit(".", 1)
+    except ValueError as e:
+        raise ValueError(
+            f"Invalid qualified name format: '{qualified_name}'. Expected 'module.path.function_name'"
+        ) from e
+
+    try:
+        module = importlib.import_module(module_path)
+    except ImportError as e:
+        raise ImportError(f"Could not import module '{module_path}': {e}") from e
+
+    try:
+        get_agent_func = getattr(module, function_name)
+    except AttributeError as e:
+        raise AttributeError(f"Function '{function_name}' not found in module '{module_path}'") from e
+
+    if not callable(get_agent_func):
+        raise ValueError(f"'{qualified_name}' is not callable")
+
+    return cast(Callable[..., Agent], get_agent_func)
+
+
+def _load_from_json(json_path: str) -> Agent:
+    """Load agent from JSON file."""
+    with open(json_path, encoding="utf-8") as f:
+        agent_dict = json.load(f)
+
+    return Agent.from_dict(agent_dict)