eval-protocol 0.0.3__py3-none-any.whl

eval_protocol/mcp/clients.py

@@ -0,0 +1,195 @@
+import asyncio
+import json
+import logging
+from contextlib import AsyncExitStack
+from typing import Any, Dict, List, Optional
+
+import aiohttp  # Still needed for type hints if we expose the session, but primary interaction changes
+import mcp.types  # Reverted to mcp.types; Explicit import for clarity
+from mcp import types as mcp_types  # Reverted to mcp.types; Explicit import for clarity
+from mcp.client.session import DEFAULT_CLIENT_INFO, ClientSession
+from mcp.client.streamable_http import streamablehttp_client
+from omegaconf import DictConfig
+
+logger = logging.getLogger(__name__)
+
+
+class IntermediaryMCPClient:
+    """
+    Client for interacting with the RewardKitIntermediaryServer using mcp.client components.
+    """
+
+    def __init__(self, intermediary_server_url: str):
+        if not intermediary_server_url:
+            raise ValueError("intermediary_server_url must be provided.")
+        self.server_url = intermediary_server_url.rstrip("/")  # Should be like http://localhost:8001/mcp
+
+        self._exit_stack: Optional[AsyncExitStack] = None
+        self._mcp_session: Optional[ClientSession] = None
+
+    async def connect(self):
+        """Establishes connection and MCP session."""
+        if self._mcp_session and not self._mcp_session.is_closed:
+            logger.debug("Already connected.")
+            return
+
+        self._exit_stack = AsyncExitStack()
+        try:
+            logger.debug(f"Attempting to connect to Intermediary MCP server at {self.server_url}")
+            read_stream, write_stream, http_session_info = await self._exit_stack.enter_async_context(
+                streamablehttp_client(self.server_url)
+            )
+            # http_session_info might contain the underlying aiohttp session if needed, and mcp_session_id
+            # logger.debug(f"Streamable HTTP transport established. HTTP session info: {http_session_info}")
+
+            self._mcp_session = await self._exit_stack.enter_async_context(
+                ClientSession(read_stream, write_stream, client_info=DEFAULT_CLIENT_INFO)
+            )
+            await self._mcp_session.initialize()
+            logger.info(f"IntermediaryMCPClient connected and MCP session initialized with {self.server_url}")
+        except Exception as e:
+            if self._exit_stack:  # pragma: no cover
+                await self._exit_stack.aclose()
+                self._exit_stack = None
+            self._mcp_session = None
+            logger.error(
+                f"Failed to connect or initialize MCP session with {self.server_url}: {e}",
+                exc_info=True,
+            )
+            raise
+
+    async def close(self):
+        """Closes the MCP session and underlying transport."""
+        if self._exit_stack:
+            logger.debug(f"Closing IntermediaryMCPClient connection to {self.server_url}")
+            await self._exit_stack.aclose()
+            self._exit_stack = None
+            self._mcp_session = None
+            logger.info(f"IntermediaryMCPClient connection to {self.server_url} closed.")
+
+    async def _ensure_connected(self):
+        # ClientSession doesn't have a public is_closed.
+        # We rely on _mcp_session being None or connect() re-establishing.
+        # The AsyncExitStack handles actual closure of resources.
+        if not self._mcp_session:
+            logger.debug("Session not established, attempting to connect...")
+            await self.connect()
+
+        # After attempting to connect, if _mcp_session is still None, it means connection failed.
+        if not self._mcp_session:
+            raise RuntimeError("Failed to establish or re-establish MCP session.")
+
+    async def _call_intermediary_tool(self, tool_name: str, tool_args_payload: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Helper to make a raw tool call to the intermediary server and parse the result.
+        The tool_args_payload is the "arguments" field for the intermediary's tool.
+        """
+        await self._ensure_connected()
+        if not self._mcp_session:  # For type checker
+            raise RuntimeError("MCP session not available after ensure_connected.")
+
+        logger.debug(f"Calling intermediary tool '{tool_name}' with payload: {tool_args_payload}")
+
+        mcp_response: mcp_types.CallToolResult = await self._mcp_session.call_tool(tool_name, tool_args_payload)
+
+        logger.debug(f"Raw MCP response from intermediary for '{tool_name}': {mcp_response}")
+
+        if mcp_response.isError or not mcp_response.content or not hasattr(mcp_response.content[0], "text"):
+            error_message = f"Tool call '{tool_name}' to intermediary failed."
+            if mcp_response.isError and mcp_response.content and hasattr(mcp_response.content[0], "text"):
+                error_message += f" Details: {mcp_response.content[0].text}"
+            elif mcp_response.isError:
+                error_message += " No detailed error message in content."
+            logger.error(error_message)
+            try:
+                if mcp_response.content and hasattr(mcp_response.content[0], "text"):
+                    parsed_error = json.loads(mcp_response.content[0].text)
+                    if isinstance(parsed_error, dict) and "error" in parsed_error:
+                        raise RuntimeError(f"{error_message} Nested error: {parsed_error['error']}")
+            except (json.JSONDecodeError, TypeError):
+                pass
+            raise RuntimeError(error_message)
+
+        try:
+            parsed_result = json.loads(mcp_response.content[0].text)
+            logger.debug(f"Parsed JSON result from intermediary for '{tool_name}': {parsed_result}")
+            return parsed_result
+        except json.JSONDecodeError as e:
+            logger.error(
+                f"Failed to parse JSON from intermediary's tool '{tool_name}' response content: {mcp_response.content[0].text}. Error: {e}"
+            )
+            raise RuntimeError(f"Failed to parse JSON response from intermediary tool '{tool_name}'.")
+
+    async def initialize_session(self, backend_requests: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """
+        Initializes a session with the IntermediaryServer, requesting backend instances.
+        """
+        payload_for_intermediary_tool = {"args": {"backends": backend_requests}}
+        return await self._call_intermediary_tool(
+            tool_name="initialize_session",
+            tool_args_payload=payload_for_intermediary_tool,
+        )
+
+    async def call_backend_tool(
+        self,
+        rk_session_id: str,
+        instance_id: str,
+        backend_name_ref: str,
+        tool_name: str,
+        tool_args: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        """
+        Calls a tool on a specific backend instance managed by the IntermediaryServer.
+        """
+        payload_for_intermediary_tool = {
+            "args": {
+                "rk_session_id": rk_session_id,
+                "instance_id": instance_id,
+                "backend_name_ref": backend_name_ref,
+                "tool_name": tool_name,
+                "tool_args": tool_args,
+            }
+        }
+        return await self._call_intermediary_tool(
+            tool_name="call_backend_tool",
+            tool_args_payload=payload_for_intermediary_tool,
+        )
+
+    async def list_backend_tools(
+        self, rk_session_id: str, instance_id: str, backend_name_ref: str
+    ) -> mcp_types.ListToolsResult:
+        """
+        Lists tools available on a specific backend instance via the IntermediaryServer.
+        """
+        payload_for_intermediary_tool = {
+            "args": {
+                "rk_session_id": rk_session_id,
+                "instance_id": instance_id,
+                "backend_name_ref": backend_name_ref,
+            }
+        }
+        # _call_intermediary_tool returns a Dict[str, Any] which is the parsed JSON
+        # from the intermediary's response. This dict should be the model_dump of ListToolsResult.
+        raw_result_dict = await self._call_intermediary_tool(
+            tool_name="list_backend_tools",
+            tool_args_payload=payload_for_intermediary_tool,
+        )
+        # Parse the dictionary back into the Pydantic model
+        return mcp_types.ListToolsResult(**raw_result_dict)
+
+    async def cleanup_session(self, rk_session_id: str) -> Dict[str, Any]:
+        """
+        Cleans up a session on the IntermediaryServer.
+        """
+        payload_for_intermediary_tool = {"args": {"rk_session_id": rk_session_id}}
+        result = await self._call_intermediary_tool(
+            tool_name="cleanup_session", tool_args_payload=payload_for_intermediary_tool
+        )
+        return result
+
+    async def __aenter__(self):
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()

eval_protocol/mcp/execution/__init__.py

@@ -0,0 +1,23 @@
+"""
+MCP Execution Framework
+
+This module handles policy execution, tool calling, and rollout coordination.
+"""
+
+from .base_policy import LLMBasePolicy
+from .policy import AnthropicPolicy, OpenAIPolicy, FireworksPolicy
+from .manager import ExecutionManager
+
+# FireworksPolicy is conditionally imported by policy.py
+_FIREWORKS_AVAILABLE = FireworksPolicy is not None
+
+__all__ = [
+    "LLMBasePolicy",
+    "AnthropicPolicy",
+    "OpenAIPolicy", 
+    "ExecutionManager",
+]
+
+# Only export FireworksPolicy if it's available
+if _FIREWORKS_AVAILABLE:
+    __all__.append("FireworksPolicy")

eval_protocol/mcp/execution/base_policy.py

@@ -0,0 +1,227 @@
+"""
+Base Policy for LLM Policies
+
+This module contains the LLMBasePolicy abstract base class that provides
+common functionality for all LLM-based policies (Fireworks, OpenAI, Anthropic, etc.)
+"""
+
+import asyncio
+import json
+import logging
+import os
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from ...playback_policy import PlaybackPolicyBase
+from ..types import LLMUsageStats, MCPToolCall
+
+logger = logging.getLogger(__name__)
+
+
+class LLMBasePolicy(PlaybackPolicyBase, ABC):
+    """
+    Base class for LLM policies that work with MCP environments via tool calling.
+
+    This abstraction enables shared code between FireworksPolicy and OpenAIPolicy.
+    Maintains conversation history per environment for proper OpenAI-style trajectories.
+    """
+
+    def __init__(
+        self,
+        model_id: str,
+        temperature: float = 0.2,
+        max_tokens: int = 4096,
+        max_tools_per_turn: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Initialize base policy with automatic record/playback detection.
+
+        Args:
+            model_id: Model identifier
+            temperature: Sampling temperature (0.0 to 2.0)
+            max_tokens: Maximum tokens to generate per request
+            max_tools_per_turn: Maximum number of tool calls per turn (None = unlimited, 1 = single tool)
+        """
+        # Initialize playback functionality (parent class handles EP_PLAYBACK_FILE automatically)
+        super().__init__(**kwargs)
+
+        # Store policy configuration
+        self.model_id = model_id
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.max_tools_per_turn = max_tools_per_turn
+
+        # Initialize conversation state tracking for proper OpenAI trajectories
+        self.initialized = False
+
+    @abstractmethod
+    async def _make_llm_call(self, messages: List[Dict], tools: List[Dict]) -> Dict:
+        """
+        Make an LLM API call. Subclasses must implement this.
+
+        Args:
+            messages: Conversation messages
+            tools: Available tools in OpenAI format
+
+        Returns:
+            LLM response with choices[0].message containing content and tool_calls
+        """
+        pass
+
+    @abstractmethod
+    def _convert_mcp_tools_to_llm_format(self, mcp_tools: List[Dict]) -> List[Dict]:
+        """
+        Convert MCP tool schemas to LLM-specific format.
+
+        Args:
+            mcp_tools: List of MCP tool definitions
+
+        Returns:
+            List of LLM-compatible tool definitions
+        """
+        pass
+
+    def add_tool_response(
+        self,
+        env_index: int,
+        tool_call: MCPToolCall,
+        tool_response: Union[str, List[Dict[str, Any]]],
+        conversation_history: List[Dict[str, Any]],
+        reward: float = 0.0,
+        terminated: bool = False,
+        info: Optional[Dict[str, Any]] = None,
+    ):
+        """Add tool call and response to conversation history with control plane metadata."""
+        # Use the preserved tool_call_id directly
+        if tool_call.tool_call_id is None:
+            raise ValueError("Tool call ID is required for tool response recording")
+
+        tool_message = {
+            "role": "tool",
+            "tool_call_id": tool_call.tool_call_id,
+            "content": tool_response,
+        }
+
+        # Add control plane metadata if provided
+        if reward != 0.0 or terminated or info:
+
+            tool_message["metadata"] = {
+                "reward": reward,
+                "terminated": terminated,
+                "info": info or {},
+            }
+
+        conversation_history.append(tool_message)
+
+    def log_conversation_state_for_playback(
+        self, env_index: int, step: int, conversation_history: List[Dict[str, Any]]
+    ):
+        """
+        Log the current conversation state in the format required for playback.
+
+        Expected format: {"env_index": 0, "step": 0, "messages": [{..}, {..}]}
+
+        Args:
+            env_index: Environment index
+            step: Current step number
+        """
+        # Use EP_PLAYBACK_FILE environment variable for recording
+        playback_file = os.environ.get("EP_PLAYBACK_FILE")
+        if not playback_file:
+            return  # No recording file specified
+
+        playback_entry = {
+            "env_index": env_index,
+            "step": step,
+            "messages": conversation_history.copy(),
+        }
+
+        # TODO: because we're using threads now, the ordering will be wrong.
+
+        with open(playback_file, "a") as f:
+            f.write(json.dumps(playback_entry) + "\n")
+
+    async def _generate_live_tool_calls(
+        self,
+        tool_schemas: List[Dict],
+        env_index: int,
+        conversation_history: List[Dict[str, Any]],
+    ) -> Tuple[List[MCPToolCall], LLMUsageStats]:
+        """
+        Generate tool calls using conversation history for proper OpenAI trajectories.
+
+        Args:
+            tool_schemas: Available MCP tools for this environment
+            env_index: Environment index
+            user_prompt: Current user prompt with observation
+
+        Returns:
+            List of MCPToolCall objects
+        """
+        # Convert MCP tools to LLM format
+        llm_tools = self._convert_mcp_tools_to_llm_format(tool_schemas)
+
+        logger.debug(
+            f"Environment {env_index} - Converted {len(tool_schemas)} MCP tools to {len(llm_tools)} LLM tools"
+        )
+        logger.debug(f"Environment {env_index} - Conversation length: {len(conversation_history)} messages")
+
+        try:
+            # Make API call with conversation history
+            response = await self._make_llm_call(conversation_history, llm_tools)
+        except Exception as e:
+            logger.error(f"LLM API call failed for env {env_index}: {e}")
+            raise e
+
+        # ADD ASSISTANT MESSAGE TO ACTUAL CONVERSATION HISTORY
+        # This is crucial for proper tool call ID management in add_tool_response
+        assistant_message_for_history = {
+            "role": "assistant",
+            "content": response["choices"][0]["message"]["content"],
+        }
+        usage_stats = LLMUsageStats(
+            prompt_tokens=response["usage"]["prompt_tokens"],
+            completion_tokens=response["usage"]["completion_tokens"],
+            total_tokens=response["usage"]["total_tokens"],
+        )
+
+        # Extract tool call from response
+        message = response["choices"][0]["message"]
+        logger.debug(f"Environment {env_index} - Response message: {message}")
+
+        # Add ALL tool calls if present with the actual API response IDs
+        if message.get("tool_calls"):
+            assistant_message_for_history["tool_calls"] = message["tool_calls"]
+
+        # Add to actual conversation history
+        conversation_history.append(assistant_message_for_history)
+
+        if message.get("tool_calls") and len(message["tool_calls"]) > 0:
+            tool_calls = message["tool_calls"]
+
+            # Handle multiple tool calls - create MCPToolCall for each
+            mcp_tool_calls = []
+            for tool_call in tool_calls:
+                mcp_tool_call = MCPToolCall(
+                    tool_name=tool_call["function"]["name"],
+                    arguments=json.loads(tool_call["function"]["arguments"]),
+                    tool_call_id=tool_call["id"],
+                )
+                mcp_tool_calls.append(mcp_tool_call)
+
+            if self.max_tools_per_turn:
+                mcp_tool_calls = mcp_tool_calls[: self.max_tools_per_turn]
+
+            return mcp_tool_calls, usage_stats
+        else:
+            # No tool calls in response - this is normal when episode ends or LLM provides only text
+            logger.info(f"No tool calls in response for env {env_index}, message content: {message.get('content')}")
+            return [
+                MCPToolCall(
+                    tool_name="_no_tool_call",
+                    arguments={
+                        "reason": "no_tool_call_generated",
+                    },
+                )
+            ], usage_stats

eval_protocol/mcp/execution/fireworks_policy.py

@@ -0,0 +1,209 @@
+"""
+Fireworks AI Policy Implementation
+
+This module contains the FireworksPolicy class that integrates with Fireworks AI's LLM API
+for tool calling and conversation management in MCP environments.
+"""
+
+import asyncio
+import json
+import logging
+import os
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, List, Optional
+
+from .base_policy import LLMBasePolicy
+from ..types import MCPToolCall
+
+logger = logging.getLogger(__name__)
+
+
+class FireworksPolicy(LLMBasePolicy):
+    """
+    Fireworks AI policy implementation that works with ANY MCP environment via tool calling.
+
+    NO environment-specific logic - everything comes from MCP tools and dataset prompts.
+    Supports both live mode (using Fireworks LLM) and playback mode (replaying recorded trajectories).
+    """
+
+    from fireworks import DeploymentTypeLiteral
+
+    def __init__(
+        self,
+        model_id: str,
+        temperature: float = 0.2,
+        deployment_type: DeploymentTypeLiteral = "serverless",
+        max_tokens: int = 4096,
+        max_tools_per_turn: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Initialize Fireworks policy.
+
+        Args:
+            model_id: Fireworks model identifier (e.g., "accounts/fireworks/models/qwen3-235b-a22b")
+            temperature: Sampling temperature (0.0 to 2.0)
+            deployment_type: "serverless", "on-demand", "auto", or "on-demand-lora"
+            max_tokens: Maximum tokens to generate per request
+            max_tools_per_turn: Maximum number of tool calls per turn (None = unlimited, 1 = single tool)
+        """
+        super().__init__(model_id, temperature, max_tokens, max_tools_per_turn, **kwargs)
+
+        self.deployment_type = deployment_type
+
+        # Only initialize Fireworks LLM in live mode (not in playback mode)
+        if not self._is_playback:
+            # Import Fireworks Build SDK - optional at module level
+            try:
+                from fireworks import LLM
+            except ImportError:
+                raise ImportError(
+                    "The 'fireworks-ai' package is required for FireworksPolicy. "
+                    "Please install it with 'pip install fireworks-ai'"
+                )
+
+            # Verify authentication
+            from ...auth import get_fireworks_api_key
+
+            api_key = get_fireworks_api_key()
+            if not api_key:
+                raise ValueError(
+                    "FIREWORKS_API_KEY environment variable or ~/.fireworks/auth.ini file is required "
+                    "to use FireworksPolicy. See the reward-kit documentation for setup instructions."
+                )
+
+            # Set the API key for the Fireworks SDK
+            os.environ["FIREWORKS_API_KEY"] = api_key
+
+            # Initialize the LLM instance using Build SDK
+            try:
+                self.llm = LLM(
+                    model=self.model_id,
+                    deployment_type=self.deployment_type,
+                    temperature=self.temperature,
+                )
+                logger.info(f"✅ Initialized Fireworks LLM: {self.model_id} ({self.deployment_type})")
+            except Exception as e:
+                raise RuntimeError(f"Failed to initialize Fireworks LLM '{self.model_id}': {e}")
+            # Create dedicated executor for non-blocking LLM calls
+            self.llm_executor = ThreadPoolExecutor(
+                max_workers=16,  # Allow up to 16 concurrent LLM API calls
+                thread_name_prefix="fireworks-api",
+            )
+        else:
+            # In playback mode, skip expensive LLM initialization
+            self.llm = None
+            logger.info(f"🎬 Playback mode: Skipping Fireworks LLM initialization for performance")
+
+    def __del__(self):
+        """Clean up executor on garbage collection."""
+        if hasattr(self, "llm_executor"):
+            self.llm_executor.shutdown(wait=False)
+
+    def _clean_messages_for_api(self, messages: List[Dict]) -> List[Dict]:
+        """
+        Clean messages by removing metadata fields that Fireworks API doesn't accept.
+
+        Args:
+            messages: Conversation messages with potential metadata
+
+        Returns:
+            Clean messages without metadata fields
+        """
+        clean_messages = []
+        for msg in messages:
+            clean_msg = msg.copy()
+            # Remove metadata field if present
+            if "metadata" in clean_msg:
+                del clean_msg["metadata"]
+            clean_messages.append(clean_msg)
+        return clean_messages
+
+    async def _make_llm_call(self, messages: List[Dict], tools: List[Dict]) -> Dict:
+        """
+        Make a Fireworks API call.
+
+        Args:
+            messages: Conversation messages (may contain metadata)
+            tools: Available tools in OpenAI format
+
+        Returns:
+            API response in OpenAI format
+        """
+        llm = self.llm
+        if llm is None:
+            raise RuntimeError("Fireworks LLM not initialized")
+
+        # Clean messages by removing metadata before sending to API
+        clean_messages = self._clean_messages_for_api(messages)
+
+        current_request = {
+            "messages": clean_messages,
+            "tools": tools,
+            "temperature": self.temperature,
+            "max_tokens": self.max_tokens,
+        }
+
+        loop = asyncio.get_event_loop()
+        response = await loop.run_in_executor(
+            self.llm_executor, lambda: llm.chat.completions.create(**current_request)
+        )
+
+        # Convert Fireworks response to standard format
+        return {
+            "choices": [
+                {
+                    "message": {
+                        "content": response.choices[0].message.content,
+                        "tool_calls": (
+                            [
+                                {
+                                    "id": tc.id,
+                                    "type": tc.type,
+                                    "function": {
+                                        "name": tc.function.name,
+                                        "arguments": tc.function.arguments,
+                                    },
+                                }
+                                for tc in (response.choices[0].message.tool_calls or [])
+                            ]
+                            if response.choices[0].message.tool_calls
+                            else []
+                        ),
+                    }
+                }
+            ],
+            "usage": {
+                "prompt_tokens": response.usage.prompt_tokens,
+                "completion_tokens": response.usage.completion_tokens,
+                "total_tokens": response.usage.total_tokens,
+            },
+        }
+
+    def _convert_mcp_tools_to_llm_format(self, mcp_tools: List[Dict]) -> List[Dict]:
+        """
+        Convert MCP tool schemas to OpenAI function calling format for Fireworks.
+
+        Args:
+            mcp_tools: List of MCP tool definitions
+
+        Returns:
+            List of OpenAI-compatible tool definitions
+        """
+        openai_tools = []
+
+        for mcp_tool in mcp_tools:
+            openai_tool = {
+                "type": "function",
+                "function": {
+                    "name": mcp_tool["name"],
+                    "description": mcp_tool.get("description", f"Execute {mcp_tool['name']} action"),
+                    "parameters": mcp_tool.get(
+                        "input_schema",
+                        {"type": "object", "properties": {}, "required": []},
+                    ),
+                },
+            }
+            openai_tools.append(openai_tool)
+
+        return openai_tools