deepset-mcp 0.0.2__py3-none-any.whl

deepset_mcp/benchmark/runner/cli_tests.py

@@ -0,0 +1,226 @@
+import typer
+
+from deepset_mcp.benchmark.runner.config_loader import (
+    find_all_test_case_paths,
+    load_test_case_by_name,
+    load_test_case_from_path,
+)
+from deepset_mcp.benchmark.runner.models import TestCaseConfig
+from deepset_mcp.benchmark.runner.setup_actions import setup_all, setup_test_case
+from deepset_mcp.benchmark.runner.teardown_actions import teardown_all, teardown_test_case
+
+tests_app = typer.Typer(help="Commands for setting up and tearing down test-cases.")
+
+
+@tests_app.command("list")
+def list_cases(
+    test_dir: str | None = typer.Option(
+        None,
+        help="Directory where all test-case YAMLs live (`benchmark/tasks/*.yml`).",
+    ),
+) -> None:
+    """List all available test cases stored under `test_dir`."""
+    paths = find_all_test_case_paths(test_dir)
+    if not paths:
+        typer.secho(f"No test-case files found in {test_dir}", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    for p in paths:
+        typer.echo(f" • {p.stem}")
+
+
+@tests_app.command("setup")
+def create_case(
+    test_name: str = typer.Argument(..., help="Test-case name (without .yml)."),
+    workspace_name: str = typer.Option(
+        "default", "--workspace", "-w", help="Workspace in which to create pipelines and indexes."
+    ),
+    api_key: str | None = typer.Option(
+        None,
+        "--api-key",
+        "-k",
+        help="Explicit DP_API_KEY to use (overrides environment).",
+    ),
+    test_dir: str | None = typer.Option(
+        None,
+        help="Directory where test-case YAMLs are stored.",
+    ),
+) -> None:
+    """Load a single test-case by name and create its pipeline + index (if any) on deepset."""
+    try:
+        test_cfg = load_test_case_by_name(name=test_name, task_dir=test_dir)
+    except FileNotFoundError:
+        typer.secho(f"Test-case '{test_name}' not found under {test_dir}.", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+    except Exception as e:
+        typer.secho(f"Failed to load test-case '{test_name}': {e}", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    typer.secho(f"→ Creating resources for '{test_name}' in '{workspace_name}'…", fg=typer.colors.GREEN)
+    try:
+        setup_test_case(test_cfg=test_cfg, workspace_name=workspace_name, api_key=api_key)
+    except Exception as e:
+        typer.secho(f"✘ Failed to set up '{test_name}': {e}", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    typer.secho(f"✔ '{test_name}' ready.", fg=typer.colors.GREEN)
+
+
+@tests_app.command("setup-all")
+def create_all(
+    workspace_name: str = typer.Option(
+        "default", "--workspace", "-w", help="Workspace in which to create pipelines and indexes."
+    ),
+    api_key: str | None = typer.Option(
+        None,
+        "--api-key",
+        "-k",
+        help="Explicit DP_API_KEY to use (overrides environment).",
+    ),
+    concurrency: int = typer.Option(
+        5,
+        "--concurrency",
+        "-c",
+        help="Maximum number of test-cases to set up in parallel.",
+    ),
+    test_dir: str | None = typer.Option(
+        None,
+        help="Directory where test-case YAMLs are stored.",
+    ),
+) -> None:
+    """Load every test-case under `task_dir` and create pipelines + indexes in `workspace_name` in parallel."""
+    paths = find_all_test_case_paths(test_dir)
+    if not paths:
+        typer.secho(f"No test-case files found in {test_dir}", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    # 1) Load all configs
+    test_cfgs: list[TestCaseConfig] = []
+    for p in paths:
+        try:
+            cfg = load_test_case_from_path(path=p)
+            test_cfgs.append(cfg)
+        except Exception as e:
+            typer.secho(f"Skipping '{p.stem}' (load error: {e})", fg=typer.colors.YELLOW)
+
+    if not test_cfgs:
+        typer.secho("No valid test-case configs to create.", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    typer.secho(
+        f"→ Creating {len(test_cfgs)} test-cases in '{workspace_name}' (concurrency={concurrency})…",
+        fg=typer.colors.GREEN,
+    )
+    try:
+        setup_all(
+            test_cfgs=test_cfgs,
+            workspace_name=workspace_name,
+            api_key=api_key,
+            concurrency=concurrency,
+        )
+    except Exception as e:
+        typer.secho(f"✘ Some test-cases failed during creation: {e}", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    typer.secho("✔ All test-cases attempted.", fg=typer.colors.GREEN)
+
+
+@tests_app.command("teardown")
+def delete_case(
+    test_name: str = typer.Argument(..., help="Test-case name (without .yml)."),
+    workspace_name: str = typer.Option(
+        "default", "--workspace", "-w", help="Workspace from which to delete pipelines and indexes."
+    ),
+    api_key: str | None = typer.Option(
+        None,
+        "--api-key",
+        "-k",
+        help="Explicit DP_API_KEY to use (overrides environment).",
+    ),
+    test_dir: str | None = typer.Option(
+        None,
+        help="Directory where test-case YAMLs are stored.",
+    ),
+) -> None:
+    """Teardown a single test-case by name and delete its pipeline + index (if any) from deepset."""
+    try:
+        test_cfg = load_test_case_by_name(name=test_name, task_dir=test_dir)
+    except FileNotFoundError:
+        typer.secho(f"Test-case '{test_name}' not found under {test_dir}.", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+    except Exception as e:
+        typer.secho(f"Failed to load test-case '{test_name}': {e}", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    typer.secho(f"→ Deleting resources for '{test_name}' from '{workspace_name}'…", fg=typer.colors.GREEN)
+    try:
+        teardown_test_case(test_cfg=test_cfg, workspace_name=workspace_name, api_key=api_key)
+    except Exception as e:
+        typer.secho(f"✘ Failed to teardown '{test_name}': {e}", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    typer.secho(f"✔ '{test_name}' resources deleted.", fg=typer.colors.GREEN)
+
+
+@tests_app.command("teardown-all")
+def delete_all(
+    workspace_name: str = typer.Option(
+        "default", "--workspace", "-w", help="Workspace from which to delete pipelines and indexes."
+    ),
+    api_key: str | None = typer.Option(
+        None,
+        "--api-key",
+        "-k",
+        help="Explicit DP_API_KEY to use (overrides environment).",
+    ),
+    concurrency: int = typer.Option(
+        5,
+        "--concurrency",
+        "-c",
+        help="Maximum number of test-cases to teardown in parallel.",
+    ),
+    test_dir: str | None = typer.Option(
+        None,
+        help="Directory where test-case YAMLs are stored.",
+    ),
+) -> None:
+    """Teardown every test-case under `task_dir` and delete pipelines and indexes from deepset."""
+    paths = find_all_test_case_paths(test_dir)
+    if not paths:
+        typer.secho(f"No test-case files found in {test_dir}", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    # 1) Load all configs
+    test_cfgs: list[TestCaseConfig] = []
+    for p in paths:
+        try:
+            cfg = load_test_case_from_path(path=p)
+            test_cfgs.append(cfg)
+        except Exception as e:
+            typer.secho(f"Skipping '{p.stem}' (load error: {e})", fg=typer.colors.YELLOW)
+
+    if not test_cfgs:
+        typer.secho("No valid test-case configs to delete.", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    typer.secho(
+        f"→ Deleting {len(test_cfgs)} test-cases from '{workspace_name}' (concurrency={concurrency})…",
+        fg=typer.colors.GREEN,
+    )
+    try:
+        teardown_all(
+            test_cfgs=test_cfgs,
+            workspace_name=workspace_name,
+            api_key=api_key,
+            concurrency=concurrency,
+        )
+    except Exception as e:
+        typer.secho(f"✘ Some test-cases failed during deletion: {e}", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    typer.secho("✔ All test-cases teardown attempted.", fg=typer.colors.GREEN)
+
+
+def create_tests_app() -> typer.Typer:
+    """Create the tests CLI app."""
+    return tests_app

deepset_mcp/benchmark/runner/cli_utils.py

@@ -0,0 +1,61 @@
+import os
+from pathlib import Path
+
+import typer
+from pydantic import ValidationError
+
+from deepset_mcp.benchmark.runner.config import BenchmarkConfig
+from deepset_mcp.benchmark.runner.models import AgentConfig
+
+
+def override_deepset_env_vars(api_key: str | None, workspace: str | None) -> None:
+    """Overrides deepset-specific environment variables."""
+    if api_key is not None:
+        os.environ["DEEPSET_API_KEY"] = api_key
+
+    if workspace is not None:
+        os.environ["DEEPSET_WORKSPACE"] = workspace
+
+
+def validate_and_setup_configs(
+    agent_config: str, test_case_base_dir: str | None, output_dir: str | None
+) -> tuple[AgentConfig, BenchmarkConfig]:
+    """Validate and setup configurations."""
+    # Validate agent config path
+    agent_config_path = Path(agent_config)
+    if not agent_config_path.exists():
+        typer.secho(f"Agent config file not found: {agent_config}", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    test_case_base_path = None
+    if test_case_base_dir is not None:
+        test_case_base_path = Path(test_case_base_dir)
+        if not test_case_base_path.exists():
+            typer.secho(f"Test case base directory not found: {test_case_base_dir}", fg=typer.colors.RED)
+            raise typer.Exit(code=1)
+
+    benchmark_kwargs = {}
+    if test_case_base_path is not None:
+        benchmark_kwargs["test_case_base_dir"] = test_case_base_path
+
+    if output_dir is not None:
+        benchmark_kwargs["output_dir"] = Path(output_dir)
+
+    # Load and validate configurations
+    try:
+        benchmark_config = BenchmarkConfig(**benchmark_kwargs)  # type: ignore
+    except ValidationError as e:
+        typer.secho("Configuration error:", fg=typer.colors.RED)
+        for error in e.errors():
+            field = " -> ".join(str(loc) for loc in error["loc"])
+            typer.secho(f"  {field}: {error['msg']}", fg=typer.colors.RED)
+        typer.secho("\nPlease ensure all required environment variables are set", fg=typer.colors.YELLOW)
+        raise typer.Exit(code=1)
+
+    try:
+        agent_cfg = AgentConfig.from_file(agent_config_path)
+    except Exception as e:
+        typer.secho(f"Invalid agent config: {e}", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    return agent_cfg, benchmark_config

deepset_mcp/benchmark/runner/config.py

@@ -0,0 +1,73 @@
+import os
+from pathlib import Path
+from typing import Any
+
+from pydantic import Field, field_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class BenchmarkConfig(BaseSettings):
+    """Core configuration for the benchmark runner."""
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_file_encoding="utf-8",
+        case_sensitive=False,
+        extra="ignore",  # Ignore extra env vars
+    )
+
+    deepset_workspace: str = Field(default_factory=lambda: os.environ.get("DEEPSET_WORKSPACE", ""))
+    deepset_api_key: str = Field(default_factory=lambda: os.environ.get("DEEPSET_API_KEY", ""))
+
+    # Optional fields with defaults
+    output_dir: Path = Field(default_factory=Path.cwd)
+    test_case_base_dir: Path = Field(default_factory=lambda: Path(__file__).parent.parent / "tasks")
+
+    # Store all other available env vars
+    additional_env_vars: dict[str, str] = Field(default_factory=dict)
+
+    @field_validator("deepset_workspace", "deepset_api_key")
+    @classmethod
+    def validate_not_empty(cls, v: str) -> str:
+        """Validate that the env var is not empty."""
+        if not v or not v.strip():
+            raise ValueError("DEEPSET_WORKSPACE or DEEPSET_API_KEY is empty.")
+        return v
+
+    def model_post_init(self, __context: Any) -> None:
+        """After initialization, collect all available env vars."""
+        # Collect all env vars that might be useful (excluding system ones)
+        ignore_prefixes = ("PATH", "HOME", "USER", "SHELL", "TERM", "PWD", "LC_")
+
+        for key, value in os.environ.items():
+            # Skip system variables and already captured ones
+            if (
+                not any(key.startswith(prefix) for prefix in ignore_prefixes)
+                and key not in ("DEEPSET_WORKSPACE", "DEEPSET_API_KEY")
+                and value
+            ):  # Only include non-empty values
+                self.additional_env_vars[key] = value
+
+    def check_required_env_vars(self, required_vars: list[str]) -> tuple[bool, list[str]]:
+        """
+        Check if all required environment variables are available.
+
+        Returns:
+            Tuple of (all_available, missing_vars)
+        """
+        available = self.get_all_env_vars()
+        missing = [var for var in required_vars if var not in available]
+        return len(missing) == 0, missing
+
+    def get_env_var(self, key: str) -> str:
+        """Get a specific environment variable."""
+        all_vars = self.get_all_env_vars()
+        return all_vars[key]
+
+    def get_all_env_vars(self) -> dict[str, str]:
+        """Get all available environment variables."""
+        return {
+            "DEEPSET_WORKSPACE": self.deepset_workspace,
+            "DEEPSET_API_KEY": self.deepset_api_key,
+            **self.additional_env_vars,
+        }

deepset_mcp/benchmark/runner/config_loader.py

@@ -0,0 +1,64 @@
+import glob
+from pathlib import Path
+
+from pydantic import ValidationError
+
+from deepset_mcp.benchmark.runner.models import TestCaseConfig
+
+
+def _default_task_dir() -> Path:
+    """Return the path to the `benchmark/tasks` directory, resolved relative to this file."""
+    return Path(__file__).parent.parent / "tasks"
+
+
+def find_all_test_case_paths(task_dir: str | Path | None = None) -> list[Path]:
+    """
+    Return a list of all `.yml` or `.yaml` files under `task_dir`.
+
+    If `task_dir` is None, we resolve to `benchmark/tasks` (relative to this file).
+    """
+    if task_dir is None:
+        base = _default_task_dir()
+    else:
+        base = Path(task_dir)
+
+    pattern1 = base / "*.yml"
+    pattern2 = base / "*.yaml"
+    return [Path(p) for p in glob.glob(str(pattern1))] + [Path(p) for p in glob.glob(str(pattern2))]
+
+
+def load_test_case_from_path(path: Path) -> TestCaseConfig:
+    """
+    Read a single test-case YAML at `path` using TestCaseConfig.from_file().
+
+    Raises RuntimeError if validation or loading fails.
+    """
+    try:
+        return TestCaseConfig.from_file(path)
+    except (ValidationError, FileNotFoundError) as e:
+        raise RuntimeError(f"Failed to load {path}: {e}") from e
+
+
+def load_test_case_by_name(name: str, task_dir: str | Path | None = None) -> TestCaseConfig:
+    """
+    Given a test‐case “name” (without extension), locate the corresponding `.yml` or `.yaml`under `task_dir`.
+
+    If `task_dir` is None, defaults to `benchmark/tasks` relative to this file.
+    Returns a loaded TestCaseConfig or raises FileNotFoundError if not found.
+    """
+    if task_dir is None:
+        base = _default_task_dir()
+    else:
+        base = Path(task_dir)
+
+    candidates: list[Path] = []
+    for ext in (".yml", ".yaml"):
+        p = base / f"{name}{ext}"
+        if p.exists():
+            candidates.append(p)
+
+    if not candidates:
+        raise FileNotFoundError(f"No test-case named '{name}' under {base}")
+
+    # If multiple matches exist, pick the first
+    return load_test_case_from_path(candidates[0])

deepset_mcp/benchmark/runner/interactive.py

@@ -0,0 +1,140 @@
+from collections.abc import Callable
+from typing import Any
+
+from haystack.tools import Tool, Toolset
+from rich.console import Console
+from rich.panel import Panel
+from rich.prompt import Prompt
+
+TOOL_CONFIRMATION_CHAR = "y"
+TOOL_REJECTED_CHAR = "n"
+TOOL_AUTO_CONFIRM_CHAR = "a"
+
+
+class InteractiveToolsetWrapper:
+    """Simple wrapper that adds interactive confirmation to any toolset."""
+
+    def __init__(self, toolset: Toolset):
+        """
+        Initialize the wrapper.
+
+        Args:
+            toolset: The base toolset to wrap
+        """
+        self.base_toolset = toolset
+        self.auto_confirm_tools: set[str] = set()
+        self.console = Console()
+
+        # Create wrapped toolset
+        self._wrapped_toolset = self._create_wrapped_toolset()
+
+    def _create_wrapped_toolset(self) -> Toolset:
+        """Create a new toolset with wrapped tools."""
+        wrapped_tools = []
+
+        for tool in self.base_toolset.tools:
+            wrapped_tool = Tool(
+                name=tool.name,
+                description=tool.description,
+                parameters=tool.parameters,
+                function=self._wrap_function(tool),
+            )
+            wrapped_tools.append(wrapped_tool)
+
+        return Toolset(tools=wrapped_tools)
+
+    def _wrap_function(self, tool: Tool) -> Callable[..., Any]:
+        """Wrap a tool function with confirmation logic."""
+        original_function = tool.function
+        tool_name = tool.name
+
+        def wrapped_function(**kwargs: Any) -> Any:
+            # Check if auto-confirmed
+            if tool_name in self.auto_confirm_tools:
+                self.console.print(f"[green]✓ Auto-executing '{tool_name}'[/green]")
+                return original_function(**kwargs)
+
+            # Ask for confirmation
+            action = self._ask_confirmation(tool_name, kwargs)
+
+            if action == "reject":
+                # Get feedback message
+                feedback = Prompt.ask("Feedback message (optional)", default="")
+                return {
+                    "status": "rejected",
+                    "tool": tool_name,
+                    "feedback": feedback or "Tool execution rejected by user",
+                }
+            elif action == "confirm_auto":
+                # Add to auto-confirm and execute
+                self.auto_confirm_tools.add(tool_name)
+                self.console.print(f"[green]✓ '{tool_name}' added to auto-confirm list[/green]")
+
+            # Execute tool (for both "confirm" and "confirm_auto")
+            return original_function(**kwargs)
+
+        return wrapped_function
+
+    def _ask_confirmation(self, tool_name: str, params: dict[str, Any]) -> str:
+        """Ask user for confirmation with Rich formatting."""
+        # Build tool call display
+        lines = [f"[bold yellow]Tool:[/bold yellow] {tool_name}"]
+
+        if params:
+            lines.append("\n[bold yellow]Arguments:[/bold yellow]")
+            for key, value in params.items():
+                lines.append(f"\n[cyan]{key}:[/cyan]")
+                # Format the value with proper indentation
+                value_str = str(value)
+                if "\n" in value_str:
+                    # Multi-line value - indent each line
+                    for line in value_str.split("\n"):
+                        lines.append(f"  {line}")
+                else:
+                    lines.append(f"  {value_str}")
+
+        # Create panel with tool information
+        panel = Panel("\n".join(lines), title="🔧 Tool Execution Request", border_style="blue")
+        self.console.print(panel)
+
+        # Show options
+        self.console.print("\n[bold]Options:[/bold]")
+        self.console.print(f"  [green]{TOOL_CONFIRMATION_CHAR}[/green]  - Confirm execution")
+        self.console.print(f"  [yellow]{TOOL_AUTO_CONFIRM_CHAR}[/yellow]  - Confirm and auto-confirm this tool")
+        self.console.print(f"  [red]{TOOL_REJECTED_CHAR}[/red]  - Reject execution")
+
+        # Get user choice
+        while True:
+            choice = Prompt.ask(
+                "\nYour choice", choices=[TOOL_CONFIRMATION_CHAR, TOOL_AUTO_CONFIRM_CHAR, TOOL_REJECTED_CHAR]
+            )
+
+            if choice == TOOL_CONFIRMATION_CHAR:
+                return "confirm"
+            elif choice == TOOL_AUTO_CONFIRM_CHAR:
+                return "confirm_auto"
+            elif choice == TOOL_REJECTED_CHAR:
+                return "reject"
+
+    @property
+    def toolset(self) -> Toolset:
+        """Get the wrapped toolset."""
+        return self._wrapped_toolset
+
+    def close(self) -> None:
+        """Close the underlying toolset if it has a close method."""
+        if hasattr(self.base_toolset, "close"):
+            self.base_toolset.close()
+
+
+def wrap_toolset_interactive(toolset: Toolset) -> InteractiveToolsetWrapper:
+    """
+    Wrap any toolset with interactive confirmation.
+
+    Args:
+        toolset: The toolset to wrap
+
+    Returns:
+        InteractiveToolsetWrapper instance
+    """
+    return InteractiveToolsetWrapper(toolset)