deepset-mcp 0.0.2rc1__py3-none-any.whl

deepset_mcp/benchmark/runner/models.py

@@ -0,0 +1,203 @@
+from pathlib import Path
+from typing import Any, Self
+
+import yaml
+from pydantic import BaseModel, Field, PrivateAttr, ValidationError, model_validator
+
+
+class TestCaseConfig(BaseModel):
+    """
+    Pydantic v2 model for a single “benchmark/tasks/<test>.yml” file.
+
+    Provides a `from_file()` constructor that:
+    - Reads the YAML from disk.
+    - Resolves `query_yaml`, `index_yaml`, `expected_query` as paths relative to the YAML file’s directory.
+    - Validates that at least one of query/index is present, that paired fields (query_name/index_name) exist,
+      and that each referenced file actually exists on disk.
+    """
+
+    name: str = Field(
+        ...,
+        description="Unique identifier for this test case (snake_case, no spaces).",
+        pattern=r"^[a-z0-9_]+$",
+    )
+    objective: str = Field(..., description="A short description of what this test is about.")
+    prompt: str = Field(..., description="The prompt text to send to the Agent.")
+    query_yaml: str | None = Field(None, description="Relative or absolute path to a query pipeline YAML.")
+    query_name: str | None = Field(
+        None, description="Name to assign to the ‘query’ pipeline if `query_yaml` is present."
+    )
+    index_yaml: str | None = Field(None, description="Relative or absolute path to an indexing pipeline YAML.")
+    index_name: str | None = Field(None, description="Name to assign to the Index if `index_yaml` is present.")
+    expected_query: str | None = Field(
+        None, description="(Optional) Relative or absolute path to a 'gold' query pipeline YAML."
+    )
+    expected_index: str | None = Field(None, description="(Optional) Relative or absolute path to a 'gold' index YAML.")
+    tags: list[str] = Field(default_factory=list, description="Tags (e.g. [api-outputs, debug]).")
+    judge_prompt: str | None = Field(
+        None,
+        description="(Optional) Prompt to use for a judge LLM to verify correctness.",
+    )
+
+    # These PrivateAttrs will hold the raw YAML‐as‐text after reading from disk:
+    _query_yaml_text: str | None = PrivateAttr(default=None)
+    _index_yaml_text: str | None = PrivateAttr(default=None)
+    _expected_query_text: str | None = PrivateAttr(default=None)
+    _expected_index_text: str | None = PrivateAttr(default=None)
+
+    @model_validator(mode="before")
+    def _check_at_least_one(cls, values: dict[str, str]) -> dict[str, str]:
+        """Before any field‐level validation, ensure at least one of `query_yaml` or `index_yaml` is provided."""
+        if not values.get("query_yaml") and not values.get("index_yaml"):
+            raise ValueError("At least one of `query_yaml` or `index_yaml` must be provided.")
+        return values
+
+    @model_validator(mode="after")
+    def _load_yaml_files(self) -> Self:
+        """
+        Hook to load YAML contents from disk.
+
+        After all standard field validation has passed, this hook:
+
+        1) If `query_yaml` is set:
+           - Ensures `query_name` is also set.
+           - Reads the file from disk into `_query_yaml_text`.
+        2) If `index_yaml` is set:
+           - Ensures `index_name` is also set.
+           - Reads the file from disk into `_index_yaml_text`.
+        3) If `expected_query` is set:
+           - Reads that file into `_expected_query_text`.
+        Any missing paired field or missing file will raise an error.
+        """
+        # 1) Query pipeline YAML → text
+        if self.query_yaml:
+            if not self.query_name:
+                raise ValueError("`query_name` must be provided if `query_yaml` is set.")
+            path = Path(self.query_yaml)
+            if not path.is_file():
+                raise FileNotFoundError(f"query_yaml file not found: {self.query_yaml}")
+            self._query_yaml_text = path.read_text(encoding="utf-8")
+
+        # 2) Index YAML → text
+        if self.index_yaml:
+            if not self.index_name:
+                raise ValueError("`index_name` must be provided if `index_yaml` is set.")
+            path = Path(self.index_yaml)
+            if not path.is_file():
+                raise FileNotFoundError(f"index_yaml file not found: {self.index_yaml}")
+            self._index_yaml_text = path.read_text(encoding="utf-8")
+
+        # 3) Expected “gold” pipeline YAML → text
+        if self.expected_query:
+            path = Path(self.expected_query)
+            if not path.is_file():
+                raise FileNotFoundError(f"expected_query file not found: {self.expected_query}")
+            self._expected_query_text = path.read_text(encoding="utf-8")
+
+        if self.expected_index:
+            path = Path(self.expected_index)
+            if not path.is_file():
+                raise FileNotFoundError(f"expected_index file not found: {self.expected_index}")
+            self._expected_index_text = path.read_text(encoding="utf-8")
+
+        return self
+
+    def get_query_yaml_text(self) -> str | None:
+        """Return the raw text of the query‐pipeline YAML (or None if not set)."""
+        return self._query_yaml_text
+
+    def get_index_yaml_text(self) -> str | None:
+        """Return the raw text of the index YAML (or None if not set)."""
+        return self._index_yaml_text
+
+    def get_expected_query_text(self) -> str | None:
+        """Return the raw text of the expected “gold” pipeline YAML (or None)."""
+        return self._expected_query_text
+
+    def get_expected_index_text(self) -> str | None:
+        """Return the raw text of the expected 'gold' index YAML (or None)."""
+        return self._expected_index_text
+
+    @classmethod
+    def from_file(cls, cfg_path: Path) -> Self:
+        """
+        Read a test-case YAML from `cfg_path`, then initialize and return a TestCaseConfig instance.
+
+        Raises:
+          - FileNotFoundError if cfg_path doesn’t exist.
+          - ValidationError if any field is invalid or any referenced file is missing.
+        """
+        if not cfg_path.is_file():
+            raise FileNotFoundError(f"Test-case config not found: {cfg_path}")
+
+        raw = yaml.safe_load(cfg_path.read_text(encoding="utf-8"))
+        if not isinstance(raw, dict):
+            raise ValidationError(f"Invalid test-case YAML at {cfg_path}; expected a mapping.")
+
+        base_dir = cfg_path.parent
+
+        # For each of query_yaml, index_yaml, expected_query: if present and relative, make it absolute.
+        for field_name in ("query_yaml", "index_yaml", "expected_query"):
+            p = raw.get(field_name)
+            if p:
+                # Only rewrite if it’s not already absolute
+                candidate = Path(p)
+                if not candidate.is_absolute():
+                    raw[field_name] = str((base_dir / candidate).resolve())
+
+        return cls(**raw)
+
+
+class AgentConfig(BaseModel):
+    """Agent configuration with flexible loading patterns."""
+
+    agent_json: str | None = Field(None, description="Relative or absolute path to an agent JSON file.")
+
+    agent_factory_function: str | None = Field(None, description="Qualified name of Agent factory function.")
+
+    display_name: str = Field(..., description="Display name for the agent.")
+
+    interactive: bool = Field(False, description="Whether to run the agent in interactive mode.")
+
+    required_env_vars: list[str] = Field(
+        default_factory=list, description="Required environment variables to run the agent."
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_mutually_exclusive(cls, values: dict[str, Any]) -> dict[str, Any]:
+        """Ensure exactly one loading method is specified."""
+        methods = [values.get("agent_json"), values.get("agent_factory_function")]
+
+        if sum(bool(m) for m in methods) != 1:
+            raise ValueError("Exactly one of agent_json or agent_factory_function must be provided")
+        return values
+
+    @model_validator(mode="after")
+    def validate_files_exist(self) -> Self:
+        """Validate that referenced files exist."""
+        if self.agent_json:
+            json_path = Path(self.agent_json)
+            if not json_path.is_file():
+                raise FileNotFoundError(f"Agent JSON file not found: {self.agent_json}")
+        return self
+
+    @classmethod
+    def from_file(cls, cfg_path: Path) -> Self:
+        """Read agent config from YAML file."""
+        if not cfg_path.is_file():
+            raise FileNotFoundError(f"Agent config not found: {cfg_path}")
+
+        raw = yaml.safe_load(cfg_path.read_text(encoding="utf-8"))
+        if not isinstance(raw, dict):
+            raise ValueError(f"Invalid agent config YAML at {cfg_path}; expected a mapping.")
+
+        base_dir = cfg_path.parent
+
+        # Resolve relative paths for agent_json
+        if "agent_json" in raw and raw["agent_json"]:
+            json_path = Path(raw["agent_json"])
+            if not json_path.is_absolute():
+                raw["agent_json"] = str((base_dir / json_path).resolve())
+
+        return cls(**raw)

deepset_mcp/benchmark/runner/repl.py

@@ -0,0 +1,67 @@
+import asyncio
+from typing import Any
+
+import typer
+from haystack.dataclasses.chat_message import ChatMessage
+from haystack.dataclasses.streaming_chunk import StreamingChunk
+
+from deepset_mcp.benchmark.runner.agent_loader import load_agent
+from deepset_mcp.benchmark.runner.config import BenchmarkConfig
+from deepset_mcp.benchmark.runner.models import AgentConfig
+from deepset_mcp.benchmark.runner.streaming import StreamingCallbackManager
+
+
+async def run_repl_session_async(agent_config: AgentConfig, benchmark_config: BenchmarkConfig) -> None:
+    """Starts an interactive REPL session with the specified agent."""
+    agent, _ = load_agent(
+        config=agent_config,
+        benchmark_config=benchmark_config,
+        interactive=agent_config.interactive,
+    )
+    history: list[ChatMessage] = []
+
+    typer.secho(f"Starting interactive session with '{agent_config.display_name}'.", fg=typer.colors.CYAN)
+    typer.secho("Type 'exit' or 'quit' to end the session.", fg=typer.colors.CYAN)
+
+    while True:
+        try:
+            user_input = typer.prompt("\n👤 You")
+            if user_input.lower() in ["exit", "quit"]:
+                typer.secho("Ending session. Goodbye!", fg=typer.colors.MAGENTA)
+                break
+
+            # Add user message to history
+            history.append(ChatMessage.from_user(user_input))
+
+            # Setup streaming
+            streaming_callback_manager = StreamingCallbackManager()
+
+            async def streaming_callback(
+                chunk: StreamingChunk,
+                manager: StreamingCallbackManager = streaming_callback_manager,
+            ) -> Any:
+                return await manager(chunk)
+
+            # Run the agent
+            typer.secho("\n🤖 Agent\n\n", fg=typer.colors.BLUE, nl=False)
+            agent_output = await agent.run_async(messages=history, streaming_callback=streaming_callback)
+
+            # The streaming callback handles printing the final text output.
+            # We replace our local history with the full history from the agent
+            # to preserve the tool calls and results.
+            if agent_output and "messages" in agent_output:
+                history = agent_output["messages"]
+
+        except (KeyboardInterrupt, EOFError):
+            typer.secho("\nEnding session. Goodbye!", fg=typer.colors.MAGENTA)
+            break
+        except Exception as e:
+            typer.secho(f"\nAn error occurred: {e}", fg=typer.colors.RED)
+            # Remove the user message that caused the error to prevent it from being re-processed
+            if history and history[-1].is_from("user"):
+                history.pop()
+
+
+def run_repl_session(agent_config: AgentConfig, benchmark_config: BenchmarkConfig) -> None:
+    """Synchronous wrapper for the REPL session."""
+    asyncio.run(run_repl_session_async(agent_config, benchmark_config))

deepset_mcp/benchmark/runner/setup_actions.py

@@ -0,0 +1,238 @@
+import asyncio
+import os
+from pathlib import Path
+from typing import Any
+
+from deepset_mcp.api.client import AsyncDeepsetClient
+from deepset_mcp.benchmark.runner.models import TestCaseConfig
+
+
+def _get_api_key(explicit_key: str | None) -> str:
+    """
+    Return whichever API key to use: explicit_key takes precedence, otherwise read DP_API_KEY from the environment.
+
+    If still missing, raise ValueError.
+    """
+    if explicit_key:
+        return explicit_key
+    env_key = os.getenv("DP_API_KEY")
+    if not env_key:
+        raise ValueError("No API key provided: pass --api-key or set DP_API_KEY in env.")
+    return env_key
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# 1) LOW-LEVEL: “setup_pipeline” and “setup_index” using AsyncDeepsetClient as a context manager
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+async def setup_pipeline_async(
+    *,
+    yaml_path: str | None = None,
+    yaml_content: str | None = None,
+    pipeline_name: str,
+    workspace_name: str,
+    api_key: str | None = None,
+) -> None:
+    """
+    Create a new pipeline in the given workspace. Exactly one of (yaml_path, yaml_content) must be provided.
+
+    Uses DP_API_KEY or explicit api_key.
+    """
+    if (yaml_path and yaml_content) or (not yaml_path and not yaml_content):
+        raise ValueError("Exactly one of `yaml_path` or `yaml_content` must be specified.")
+
+    if yaml_path is not None:
+        yaml_str = Path(yaml_path).read_text(encoding="utf-8")
+    else:
+        yaml_str = yaml_content  # type: ignore
+
+    key_to_use = _get_api_key(api_key)
+    async with AsyncDeepsetClient(api_key=key_to_use) as client:
+        await client.pipelines(workspace=workspace_name).create(
+            name=pipeline_name,
+            yaml_config=yaml_str,
+        )
+    return None
+
+
+async def setup_index_async(
+    *,
+    yaml_path: str | None = None,
+    yaml_content: str | None = None,
+    index_name: str,
+    workspace_name: str,
+    api_key: str | None = None,
+    description: str | None = None,
+) -> None:
+    """
+    Create a new index in the given workspace. Exactly one of (yaml_path, yaml_content) must be provided.
+
+    Uses DP_API_KEY or explicit api_key.
+    """
+    if (yaml_path and yaml_content) or (not yaml_path and not yaml_content):
+        raise ValueError("Exactly one of `yaml_path` or `yaml_content` must be specified.")
+
+    if yaml_path is not None:
+        yaml_str = Path(yaml_path).read_text(encoding="utf-8")
+    else:
+        yaml_str = yaml_content  # type: ignore
+
+    key_to_use = _get_api_key(api_key)
+    async with AsyncDeepsetClient(api_key=key_to_use) as client:
+        await client.indexes(workspace=workspace_name).create(
+            name=index_name,
+            yaml_config=yaml_str,
+            description=description or "",
+        )
+    return None
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# 2) MID-LEVEL: setup a full test-case (pipeline + index if present)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+async def setup_test_case_async(
+    *,
+    test_cfg: TestCaseConfig,
+    workspace_name: str,
+    api_key: str | None = None,
+) -> None:
+    """
+    Given a TestCaseConfig, create the index and the pipeline in the specified workspace.
+
+    Uses DP_API_KEY or explicit api_key.
+    """
+    # 1) If there’s an index to create:
+    if test_cfg.index_yaml:
+        assert test_cfg.index_name is not None  # already validated by Pydantic model; added to satisfy mypy
+        await setup_index_async(
+            yaml_content=test_cfg.get_index_yaml_text(),
+            index_name=test_cfg.index_name,
+            workspace_name=workspace_name,
+            api_key=api_key,
+            description=f"Index for test {test_cfg.name}",
+        )
+
+    # 2) If there’s a “query pipeline” to create:
+    if test_cfg.query_yaml:
+        assert test_cfg.query_name is not None  # already validated by Pydantic model; added to satisfy mypy
+        await setup_pipeline_async(
+            yaml_content=test_cfg.get_query_yaml_text(),
+            pipeline_name=test_cfg.query_name,
+            workspace_name=workspace_name,
+            api_key=api_key,
+        )
+
+    return None
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# 3) HIGH-LEVEL: parallel “setup all” with configurable concurrency
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+async def setup_all_async(
+    *,
+    test_cfgs: list[TestCaseConfig],
+    workspace_name: str,
+    api_key: str | None = None,
+    concurrency: int = 5,
+) -> None:
+    """
+    Given a list of TestCaseConfig, create all indexes and pipelines in parallel.
+
+    Uses DP_API_KEY or explicit api_key.
+    """
+    semaphore = asyncio.Semaphore(concurrency)
+    tasks: list[asyncio.Task[Any]] = []
+
+    async def sem_task(cfg: TestCaseConfig) -> str:
+        async with semaphore:
+            await setup_test_case_async(test_cfg=cfg, workspace_name=workspace_name, api_key=api_key)
+            return cfg.name
+
+    for cfg in test_cfgs:
+        tasks.append(asyncio.create_task(sem_task(cfg)))
+
+    done, _ = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED)
+    errors: list[Exception] = []
+    for t in done:
+        if t.exception():
+            errors.append(t.exception())  # type: ignore
+
+    if errors:
+        raise RuntimeError(f"Errors during setup: {errors}")
+
+    return None
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# 4) SYNC WRAPPERS for all of the above (now accept api_key)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+def setup_pipeline(
+    *,
+    yaml_path: str | None = None,
+    yaml_content: str | None = None,
+    pipeline_name: str,
+    workspace_name: str,
+    api_key: str | None = None,
+) -> None:
+    """Synchronous wrapper for setup_pipeline_async. Blocks until the pipeline is created."""
+    return asyncio.run(
+        setup_pipeline_async(
+            yaml_path=yaml_path,
+            yaml_content=yaml_content,
+            pipeline_name=pipeline_name,
+            workspace_name=workspace_name,
+            api_key=api_key,
+        )
+    )
+
+
+def setup_index(
+    *,
+    yaml_path: str | None = None,
+    yaml_content: str | None = None,
+    index_name: str,
+    workspace_name: str,
+    api_key: str | None = None,
+    description: str | None = None,
+) -> None:
+    """Synchronous wrapper for setup_index_async. Blocks until the index is created."""
+    return asyncio.run(
+        setup_index_async(
+            yaml_path=yaml_path,
+            yaml_content=yaml_content,
+            index_name=index_name,
+            workspace_name=workspace_name,
+            api_key=api_key,
+            description=description,
+        )
+    )
+
+
+def setup_test_case(
+    *,
+    test_cfg: TestCaseConfig,
+    workspace_name: str,
+    api_key: str | None = None,
+) -> None:
+    """Synchronous wrapper: blocks until both pipeline and index (if any) are created."""
+    return asyncio.run(setup_test_case_async(test_cfg=test_cfg, workspace_name=workspace_name, api_key=api_key))
+
+
+def setup_all(
+    *,
+    test_cfgs: list[TestCaseConfig],
+    workspace_name: str,
+    api_key: str | None = None,
+    concurrency: int = 5,
+) -> None:
+    """Synchronous wrapper for setup_all_async. Blocks until all test-cases are created."""
+    return asyncio.run(
+        setup_all_async(test_cfgs=test_cfgs, workspace_name=workspace_name, api_key=api_key, concurrency=concurrency)
+    )