mantisdk 0.1.0__py3-none-any.whl

mantisdk/algorithm/gepa/lib/adapters/mcp_adapter/mcp_client.py

@@ -0,0 +1,364 @@
+"""
+Unified MCP Client - Supports stdio, SSE, and StreamableHTTP transports.
+
+This utility provides a single abstraction for connecting to MCP servers
+using different transport mechanisms.
+"""
+
+import asyncio
+import json
+import logging
+from abc import ABC, abstractmethod
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class BaseMCPClient(ABC):
+    """Abstract base class for MCP clients."""
+
+    def __init__(self):
+        self.request_id = 0
+
+    @abstractmethod
+    async def start(self):
+        """Start the MCP connection."""
+        pass
+
+    @abstractmethod
+    async def send_request(self, method: str, params: dict | None = None) -> dict:
+        """Send JSON-RPC request and get response."""
+        pass
+
+    @abstractmethod
+    async def close(self):
+        """Close the connection."""
+        pass
+
+    async def initialize(self) -> dict:
+        """Initialize MCP session (common across all transports)."""
+        result = await self.send_request(
+            "initialize",
+            {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {},
+                "clientInfo": {"name": "gepa-mcp-adapter", "version": "1.0"},
+            },
+        )
+        await self._send_initialized_notification()
+        return result
+
+    @abstractmethod
+    async def _send_initialized_notification(self):
+        """Send initialized notification (transport-specific)."""
+        pass
+
+    async def list_tools(self) -> list[dict]:
+        """List available tools."""
+        result = await self.send_request("tools/list")
+        return result.get("tools", [])
+
+    async def call_tool(self, name: str, arguments: dict) -> dict:
+        """Call a tool."""
+        return await self.send_request("tools/call", {"name": name, "arguments": arguments})
+
+
+class StdioMCPClient(BaseMCPClient):
+    """MCP client using stdio transport (subprocess-based)."""
+
+    def __init__(self, command: str, args: list[str]):
+        super().__init__()
+        self.command = command
+        self.args = args
+        self.process = None
+
+    async def start(self):
+        """Start the MCP server process."""
+        logger.info(f"Starting stdio MCP server: {self.command} {' '.join(self.args)}")
+        self.process = await asyncio.create_subprocess_exec(
+            self.command,
+            *self.args,
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.DEVNULL,
+        )
+
+    async def send_request(self, method: str, params: dict | None = None) -> dict:
+        """Send JSON-RPC request via stdio."""
+        if not self.process or not self.process.stdin or not self.process.stdout:
+            raise RuntimeError("Process not started or streams not available")
+
+        self.request_id += 1
+        request = {"jsonrpc": "2.0", "method": method, "id": self.request_id}
+
+        if params is not None:
+            request["params"] = params
+
+        # Send request
+        request_str = json.dumps(request) + "\n"
+        self.process.stdin.write(request_str.encode())
+        await self.process.stdin.drain()
+
+        # Read response
+        response_line = await self.process.stdout.readline()
+        response = json.loads(response_line.decode())
+
+        if "error" in response:
+            raise Exception(f"MCP error: {response['error']}")
+
+        return response.get("result", {})
+
+    async def _send_initialized_notification(self):
+        """Send initialized notification via stdio."""
+        if not self.process or not self.process.stdin:
+            raise RuntimeError("Process not started or stdin not available")
+
+        notification = {"jsonrpc": "2.0", "method": "notifications/initialized"}
+        notification_str = json.dumps(notification) + "\n"
+        self.process.stdin.write(notification_str.encode())
+        await self.process.stdin.drain()
+
+    async def close(self):
+        """Close the subprocess."""
+        if self.process and self.process.stdin:
+            self.process.stdin.close()
+            await self.process.wait()
+            logger.info("Stdio MCP connection closed")
+
+
+class SSEMCPClient(BaseMCPClient):
+    """MCP client using Server-Sent Events transport."""
+
+    def __init__(self, url: str, headers: dict[str, str] | None = None, timeout: float = 30):
+        super().__init__()
+        self.url = url
+        self.headers = headers or {}
+        self.timeout = timeout
+        self.read_stream = None
+        self.write_stream = None
+        self._sse_context = None
+
+    async def start(self):
+        """Start the SSE connection."""
+        from mcp.client.sse import sse_client  # type: ignore[import-untyped]
+
+        logger.info(f"Connecting to SSE MCP server at {self.url}")
+
+        self._sse_context = sse_client(
+            url=self.url,
+            headers=self.headers,
+            timeout=self.timeout,
+            sse_read_timeout=300,
+        )
+
+        streams = await self._sse_context.__aenter__()
+        self.read_stream, self.write_stream = streams
+        logger.info("SSE connection established")
+
+    async def send_request(self, method: str, params: dict | None = None) -> dict:
+        """Send JSON-RPC request via SSE."""
+        from mcp.shared.message import SessionMessage  # type: ignore[import-untyped]
+        from mcp.types import JSONRPCMessage, JSONRPCRequest  # type: ignore[import-untyped]
+
+        if not self.read_stream or not self.write_stream:
+            raise RuntimeError("SSE streams not initialized")
+
+        self.request_id += 1
+        request_dict = {
+            "jsonrpc": "2.0",
+            "method": method,
+            "id": self.request_id,
+        }
+
+        if params is not None:
+            request_dict["params"] = params
+
+        logger.debug(f"Sending SSE request: {method} (id={self.request_id})")
+
+        request = JSONRPCRequest(**request_dict)
+        session_message = SessionMessage(message=JSONRPCMessage(request))
+        await self.write_stream.send(session_message)
+
+        # Read response
+        response_message = await self.read_stream.receive()
+
+        if hasattr(response_message.message.root, "error"):
+            error = response_message.message.root.error
+            raise Exception(f"MCP error: {error}")
+
+        if hasattr(response_message.message.root, "result"):
+            return response_message.message.root.result
+
+        raise Exception(f"Unexpected response format: {response_message}")
+
+    async def _send_initialized_notification(self):
+        """Send initialized notification via SSE."""
+        from mcp.shared.message import SessionMessage  # type: ignore[import-untyped]
+        from mcp.types import JSONRPCMessage, JSONRPCNotification  # type: ignore[import-untyped]
+
+        if not self.write_stream:
+            raise RuntimeError("SSE write stream not initialized")
+
+        notification = JSONRPCNotification(
+            jsonrpc="2.0",
+            method="notifications/initialized",
+        )
+
+        session_message = SessionMessage(message=JSONRPCMessage(notification))
+        await self.write_stream.send(session_message)
+
+    async def close(self):
+        """Close the SSE connection."""
+        if self._sse_context:
+            try:
+                await self._sse_context.__aexit__(None, None, None)
+                logger.info("SSE connection closed")
+            except Exception as e:
+                logger.warning(f"Error closing SSE connection: {e}")
+
+
+class StreamableHTTPMCPClient(BaseMCPClient):
+    """MCP client using StreamableHTTP transport (production-grade)."""
+
+    def __init__(
+        self,
+        url: str,
+        headers: dict[str, str] | None = None,
+        timeout: float = 30,
+        sse_read_timeout: float = 300,
+    ):
+        super().__init__()
+        self.url = url
+        self.headers = headers or {}
+        self.timeout = timeout
+        self.sse_read_timeout = sse_read_timeout
+        self.read_stream = None
+        self.write_stream = None
+        self._transport_context = None
+
+    async def start(self):
+        """Start the StreamableHTTP connection."""
+        from mcp.client.streamable_http import streamable_http_client  # type: ignore[import-untyped]
+
+        logger.info(f"Connecting to StreamableHTTP MCP server at {self.url}")
+
+        self._transport_context = streamable_http_client(
+            url=self.url,
+            headers=self.headers,
+            timeout=self.timeout,
+            sse_read_timeout=self.sse_read_timeout,
+        )
+
+        streams = await self._transport_context.__aenter__()
+        self.read_stream, self.write_stream = streams
+        logger.info("StreamableHTTP connection established")
+
+    async def send_request(self, method: str, params: dict | None = None) -> dict:
+        """Send JSON-RPC request via StreamableHTTP."""
+        from mcp.shared.message import SessionMessage  # type: ignore[import-untyped]
+        from mcp.types import JSONRPCMessage, JSONRPCRequest  # type: ignore[import-untyped]
+
+        if not self.read_stream or not self.write_stream:
+            raise RuntimeError("StreamableHTTP streams not initialized")
+
+        self.request_id += 1
+        request_dict = {
+            "jsonrpc": "2.0",
+            "method": method,
+            "id": self.request_id,
+        }
+
+        if params is not None:
+            request_dict["params"] = params
+
+        logger.debug(f"Sending StreamableHTTP request: {method} (id={self.request_id})")
+
+        request = JSONRPCRequest(**request_dict)
+        session_message = SessionMessage(message=JSONRPCMessage(request))
+        await self.write_stream.send(session_message)
+
+        # Read response
+        response_message = await self.read_stream.receive()
+
+        if hasattr(response_message.message.root, "error"):
+            error = response_message.message.root.error
+            raise Exception(f"MCP error: {error}")
+
+        if hasattr(response_message.message.root, "result"):
+            return response_message.message.root.result
+
+        raise Exception(f"Unexpected response format: {response_message}")
+
+    async def _send_initialized_notification(self):
+        """Send initialized notification via StreamableHTTP."""
+        from mcp.shared.message import SessionMessage  # type: ignore[import-untyped]
+        from mcp.types import JSONRPCMessage, JSONRPCNotification  # type: ignore[import-untyped]
+
+        if not self.write_stream:
+            raise RuntimeError("StreamableHTTP write stream not initialized")
+
+        notification = JSONRPCNotification(
+            jsonrpc="2.0",
+            method="notifications/initialized",
+        )
+
+        session_message = SessionMessage(message=JSONRPCMessage(notification))
+        await self.write_stream.send(session_message)
+
+    async def close(self):
+        """Close the StreamableHTTP connection."""
+        if self._transport_context:
+            try:
+                await self._transport_context.__aexit__(None, None, None)
+                logger.info("StreamableHTTP connection closed")
+            except Exception as e:
+                logger.warning(f"Error closing StreamableHTTP connection: {e}")
+
+
+def create_mcp_client(
+    server_params: Any = None,
+    remote_url: str | None = None,
+    remote_transport: str = "sse",
+    remote_headers: dict[str, str] | None = None,
+    remote_timeout: float = 30,
+    sse_read_timeout: float = 300,
+) -> BaseMCPClient:
+    """
+    Factory function to create the appropriate MCP client.
+
+    Args:
+        server_params: StdioServerParameters for local server
+        remote_url: URL for remote server
+        remote_transport: "sse" or "streamable_http"
+        remote_headers: HTTP headers for remote connections
+        remote_timeout: Timeout for HTTP operations
+        sse_read_timeout: Timeout for SSE streaming
+
+    Returns:
+        BaseMCPClient instance (Stdio, SSE, or StreamableHTTP)
+
+    Raises:
+        ValueError: If configuration is invalid
+    """
+    if server_params and remote_url:
+        raise ValueError("Provide either server_params (local) or remote_url (remote), not both")
+    if not server_params and not remote_url:
+        raise ValueError("Must provide either server_params (local) or remote_url (remote)")
+
+    if server_params:
+        return StdioMCPClient(command=server_params.command, args=server_params.args)
+    elif remote_url:  # Type guard ensures remote_url is not None
+        if remote_transport == "sse":
+            return SSEMCPClient(url=remote_url, headers=remote_headers, timeout=remote_timeout)
+        elif remote_transport == "streamable_http":
+            return StreamableHTTPMCPClient(
+                url=remote_url,
+                headers=remote_headers,
+                timeout=remote_timeout,
+                sse_read_timeout=sse_read_timeout,
+            )
+        else:
+            raise ValueError(f"Unknown remote transport: {remote_transport}. Must be 'sse' or 'streamable_http'")
+    else:
+        # This should never happen due to earlier checks
+        raise ValueError("Must provide either server_params (local) or remote_url (remote)")

mantisdk/algorithm/gepa/lib/adapters/terminal_bench_adapter/README.md

@@ -0,0 +1,9 @@
+### Terminal-bench adapter
+
+This adapter is used to optimize the system prompt/terminal-use instruction for the default Terminus agent through custom a `GEPAAdapter` implementation.
+
+To run this example, you need to install `pip install terminal-bench` and run the following command:
+
+```bash
+python src/gepa/examples/terminal-bench/train_terminus.py --model_name=gpt-5-mini
+```

mantisdk/algorithm/gepa/lib/adapters/terminal_bench_adapter/terminal_bench_adapter.py

@@ -0,0 +1,217 @@
+import json
+import os
+import subprocess
+from datetime import datetime
+from pathlib import Path
+
+from pydantic import BaseModel
+from terminal_bench.agents.terminus_1 import CommandBatchResponse
+
+from mantisdk.algorithm.gepa.lib import EvaluationBatch, GEPAAdapter
+
+
+class TerminalBenchTask(BaseModel):
+    task_id: str
+    model_name: str
+
+
+def run_agent_tb(
+    task_ids: str | list[str],
+    run_id: str,
+    model_name: str,
+    instruction_prompt: str,
+    dataset_name: str = "terminal-bench-core",
+    dataset_version: str = "head",
+    agent_import_path: str = "train_terminus:TerminusWrapper",
+    n_concurrent: int = 6,
+    prompt_template_path: str = "prompt-templates/instruction_prompt.txt",
+):
+    """Run the replay agent for multiple task IDs using tb run command."""
+
+    env = os.environ.copy()
+    # write instruction prompt to file
+    with open(prompt_template_path, "w") as f:
+        f.write(instruction_prompt)
+
+    cmd = [
+        "tb",
+        "run",
+        "--dataset-name",
+        dataset_name,
+        "--dataset-version",
+        dataset_version,
+        "--agent-import-path",
+        agent_import_path,
+        "--model-name",
+        model_name,
+        "--run-id",
+        run_id,
+        "--n-concurrent",
+        str(n_concurrent),
+        "--output-path",
+        str(Path(os.getcwd()) / "runs"),
+    ]
+    if isinstance(task_ids, list):
+        for task_id in task_ids:
+            cmd.extend(["--task-id", task_id])
+    else:
+        cmd.extend(["--task-id", task_ids])
+
+    print(f"Running command: {' '.join(cmd)}")
+
+    try:
+        result = subprocess.run(cmd, env=env, cwd=Path(prompt_template_path).parent.parent, check=True)
+        print(f"Command completed successfully with return code: {result.returncode}")
+        return result.returncode
+    except subprocess.CalledProcessError as e:
+        print(f"Command failed with return code: {e.returncode}")
+        return e.returncode
+    except Exception as e:
+        print(f"Error running command: {e}")
+        return 1
+
+
+def get_results(task_id: str, run_id: str) -> tuple[int, list]:
+    def _read_episode_response(episode_dir: Path) -> CommandBatchResponse | None:
+        """Helper method to read and parse response.json from an episode directory."""
+        response_file = episode_dir / "response.json"
+        if response_file.exists():
+            try:
+                response_content = response_file.read_text()
+                return CommandBatchResponse.model_validate_json(response_content)
+            except Exception:
+                pass
+        return None
+
+    def _get_logging_dir(task_id: str, run_id: str):
+        logging_dir_base = Path("runs") / run_id / task_id
+        for dir in logging_dir_base.iterdir():
+            if dir.is_dir() and dir.name.startswith(task_id):
+                return dir
+        raise ValueError(f"No logging directory found for task {task_id} and run {run_id}")
+
+    logging_dir = _get_logging_dir(task_id, run_id)
+    result_json = logging_dir / "results.json"
+    with open(result_json) as f:
+        result = json.load(f)
+    if result.get("parser_results", None):
+        score = sum(x == "passed" for x in result["parser_results"].values())
+    else:
+        score = 0
+
+    if result.get("is_resolved", None):
+        success = True
+    else:
+        success = False
+
+    failed_reason = result.get("failure_mode", "unknown")
+
+    trajectory_path = logging_dir / "agent-logs"
+    episode_dirs = []
+    for dir in trajectory_path.iterdir():
+        if dir.is_dir() and dir.name.startswith("episode-"):
+            episode_dirs.append(dir)
+
+    if episode_dirs:
+        # Sort by episode number to get the last one
+        episode_dirs.sort(key=lambda x: int(x.name.split("-")[1]))
+        last_episode_dir = episode_dirs[-1]
+
+    last_episode_dir_trajectory = last_episode_dir / "debug.json"
+    with open(last_episode_dir_trajectory) as f:
+        trajectory = json.load(f)
+
+        if "input" in trajectory and isinstance(trajectory["input"], list):
+            messages = trajectory["input"]
+
+        # Add the last assistant response using helper method
+        parsed_response = _read_episode_response(last_episode_dir)
+
+        if parsed_response:
+            assistant_message = {
+                "role": "assistant",
+                "content": parsed_response.model_dump_json(),
+            }
+            messages.append(assistant_message)
+
+    return success, score, failed_reason, messages
+
+
+class TerminusAdapter(GEPAAdapter):
+    def __init__(
+        self,
+        n_concurrent: int = 6,
+        instruction_prompt_path: str = "prompt-templates/instruction_prompt.txt",
+    ):
+        self.n_concurrent = n_concurrent
+        self.instruction_prompt_path = instruction_prompt_path
+
+    def evaluate(
+        self,
+        batch: list[TerminalBenchTask],
+        candidate: dict[str, str],
+        capture_traces: bool = False,
+    ) -> EvaluationBatch:
+        outputs = []
+        scores = []
+        trajectories = []
+        example_run_id = "temp_gepa_run" + "_" + datetime.now().strftime("%Y%m%d%H%M%S")
+        example_model_name = batch[0].model_name
+
+        run_agent_tb(
+            [task.task_id for task in batch],
+            example_run_id,
+            example_model_name,
+            instruction_prompt=candidate["instruction_prompt"],
+            n_concurrent=self.n_concurrent,
+            prompt_template_path=self.instruction_prompt_path,
+        )
+
+        for example in batch:
+            try:
+                success, score, failed_reason, messages = get_results(example.task_id, example_run_id)
+            except Exception as e:
+                print(f"Error running example {example.task_id} {example_run_id}: {e}")
+                success = False
+                score = 0
+                failed_reason = str(e)
+                messages = []
+
+            outputs.append(
+                f"Terminal Bench outputs are omitted. Please see runs/{example_run_id}/{example.task_id}/ for detailed logging."
+            )
+            scores.append(score)
+            trajectories.append(
+                {
+                    "messages": messages,
+                    "instruction_prompt": candidate["instruction_prompt"],
+                    "failed_reason": failed_reason,
+                    "success": success,
+                }
+            )
+        return EvaluationBatch(
+            outputs=outputs,
+            scores=scores,
+            trajectories=trajectories,
+        )
+
+    def make_reflective_dataset(
+        self,
+        candidate: dict[str, str],
+        eval_batch: EvaluationBatch,
+        components_to_update: list[str],
+    ):
+        reflective_dataset = {"instruction_prompt": []}
+        for _score, trajectory in zip(eval_batch.scores, eval_batch.trajectories, strict=False):
+            if trajectory["success"]:
+                feedback = "Successfully solved the task!"
+            else:
+                feedback = f"Failed to solve the task. Reason: {trajectory['failed_reason']}"
+            reflective_dataset["instruction_prompt"].append(
+                {
+                    "Message History": trajectory["messages"],
+                    "Instruction Prompt": candidate["instruction_prompt"],
+                    "Feedback": feedback,
+                }
+            )
+        return reflective_dataset