init

Author: zihaolin96

eval_protocol/pytest/__init__.py

@@ -1,5 +1,6 @@
 from .default_agent_rollout_processor import AgentRolloutProcessor
 from .default_dataset_adapter import default_dataset_adapter
+from .default_klavis_sandbox_rollout_processor import KlavisSandboxRolloutProcessor
 from .default_mcp_gym_rollout_processor import MCPGymRolloutProcessor
 from .default_no_op_rollout_processor import NoOpRolloutProcessor
 from .default_single_turn_rollout_process import SingleTurnRolloutProcessor
@@ -31,6 +32,7 @@
 
 __all__ = [
     "AgentRolloutProcessor",
+    "KlavisSandboxRolloutProcessor",
     "MCPGymRolloutProcessor",
     "RolloutProcessor",
     "SingleTurnRolloutProcessor",

eval_protocol/pytest/default_klavis_sandbox_rollout_processor.py

@@ -0,0 +1,157 @@
+import asyncio
+import json
+import logging
+import os
+import tempfile
+import time
+from typing import Any, Callable, Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+from eval_protocol.models import EvaluationRow
+from eval_protocol.pytest.rollout_processor import RolloutProcessor
+from eval_protocol.pytest.types import RolloutProcessorConfig
+
+from eval_protocol.pytest.default_agent_rollout_processor import Agent
+from klavis import Klavis
+from klavis.types import CreateSandboxResponse, SandboxMcpServer
+from openai.types import CompletionUsage
+
+logger = logging.getLogger(__name__)
+
+
+class KlavisSandboxRolloutProcessor(RolloutProcessor):
+    def __init__(
+        self,
+        server_name: str,
+        initialize_data_factory: Optional[Callable[[EvaluationRow], Dict[str, Any]]] = None,
+    ):
+        super().__init__()
+        self.server_name = server_name
+        self.initialize_data_factory = initialize_data_factory
+        self.klavis_client = Klavis(api_key=os.environ.get("KLAVIS_API_KEY"))
+        self.sandbox = self._init_sandbox()
+        
+    def _init_sandbox(self) -> CreateSandboxResponse:
+        try:
+            server_name_enum = SandboxMcpServer(self.server_name)
+            return self.klavis_client.sandbox.create_sandbox(server_name=server_name_enum)
+        except Exception as e:
+            logger.error(f"Error creating sandbox: {str(e)}", exc_info=True)
+            raise
+    
+    @staticmethod
+    def create_mcp_config(server_url: str, server_key: str = "main", auth_token: str | None = None) -> str:
+        """Create a temporary MCP config file and return its path."""
+        config = {
+            "mcpServers": {
+                server_key: {
+                    "url": server_url,
+                    "transport": "streamable_http",
+                    **({"authorization": f"Bearer {auth_token}"} if auth_token else {})
+                }
+            }
+        }
+        
+        # Create a temp file that persists for the session
+        fd, path = tempfile.mkstemp(suffix=".json", prefix="mcp_config_")
+        with os.fdopen(fd, 'w') as f:
+            json.dump(config, f)
+        return path
+
+    def __call__(
+        self, rows: List[EvaluationRow], config: RolloutProcessorConfig
+    ) -> List[asyncio.Task[EvaluationRow]]:
+        """Process evaluation rows with Klavis sandbox lifecycle management"""
+        if not self.sandbox:
+            raise RuntimeError("Sandbox not initialized")
+
+        semaphore = config.semaphore
+
+        async def process_row(row: EvaluationRow) -> EvaluationRow:
+            """Process a single row with complete sandbox lifecycle"""
+            
+            start_time = time.perf_counter()
+
+            try:
+                # Step 1: Initialize data in the sandbox
+                if self.initialize_data_factory:
+                    logger.info(f"Initializing {self.server_name} sandbox {self.sandbox.sandbox_id}")
+                    init_data = self.initialize_data_factory(row)
+                    initialize_method = getattr(self.klavis_client.sandbox, f"initialize_{self.sandbox.server_name}_sandbox")
+                    initialize_method(sandbox_id=self.sandbox.sandbox_id, **init_data)
+                    logger.info(f"Sandbox initialized successfully")
+
+                # Step 2: Create temporary MCP config with sandbox URL
+                temp_config_path = self.create_mcp_config(server_url=self.sandbox.server_url, server_key=self.sandbox.server_name)
+
+                # Step 3: Run agent with sandbox MCP server
+                logger.info(f"Running agent for row {row.execution_metadata.rollout_id} with {self.server_name} sandbox")
+                agent = Agent(
+                    model=row.input_metadata.completion_params["model"],
+                    row=row,
+                    config_path=temp_config_path,
+                    logger=config.logger,
+                )
+                await agent.setup()
+                await agent.call_agent()
+
+                # Update usage metadata
+                row.execution_metadata.usage = CompletionUsage(
+                    prompt_tokens=agent.usage.get("prompt_tokens", 0),
+                    completion_tokens=agent.usage.get("completion_tokens", 0),
+                    total_tokens=agent.usage.get("total_tokens", 0),
+                )
+                row = agent.evaluation_row
+                logger.info(f"Agent execution completed for row {row.execution_metadata.rollout_id}")
+
+                # Step 4: Export sandbox data
+                logger.info(f"Exporting {self.server_name} sandbox data")
+                dump_method = getattr(self.klavis_client.sandbox, f"dump_{self.sandbox.server_name}_sandbox")
+                dump_response = dump_method(sandbox_id=self.sandbox.sandbox_id)
+                sandbox_data = dump_response.data
+
+                # Store sandbox data in row metadata for evaluation
+                if not row.execution_metadata.extra:
+                    row.execution_metadata.extra = {}
+                row.execution_metadata.extra["sandbox_data"] = sandbox_data
+                row.execution_metadata.extra["sandbox_id"] = self.sandbox.sandbox_id
+                row.execution_metadata.extra["server_name"] = self.server_name
+
+            except Exception as e:
+                logger.error(f"Error processing row {row.execution_metadata.rollout_id}: {str(e)}", exc_info=True)
+                if not row.execution_metadata.extra:
+                    row.execution_metadata.extra = {}
+                row.execution_metadata.extra["error"] = str(e)
+                raise
+
+            finally:
+                # Cleanup agent MCP client and temp config
+                if agent and agent.mcp_client:
+                    await agent.mcp_client.cleanup()
+                if temp_config_path and os.path.exists(temp_config_path):
+                    os.unlink(temp_config_path)
+                
+                # Release sandbox
+                if self.sandbox.sandbox_id:
+                    try:
+                        logger.info(f"Releasing {self.server_name} sandbox {self.sandbox.sandbox_id}")
+                        self.klavis_client.sandbox.delete_sandbox(
+                            server_name=self.sandbox.server_name, sandbox_id=self.sandbox.sandbox_id
+                        )
+                        logger.info(f"Sandbox {self.sandbox.sandbox_id} released successfully")
+                    except Exception as e:
+                        logger.error(f"Error releasing sandbox {self.sandbox.sandbox_id}: {str(e)}", exc_info=True)
+
+                row.execution_metadata.rollout_duration_seconds = time.perf_counter() - start_time
+
+            return row
+
+        async def _sem_wrapper(r: EvaluationRow) -> EvaluationRow:
+            async with semaphore:
+                result = await process_row(r)
+                return result
+
+        # Create and return tasks
+        tasks = [asyncio.create_task(_sem_wrapper(row)) for row in rows]
+        return tasks

pyproject.toml

@@ -134,6 +134,9 @@ openenv = [
 dspy = [
     "dspy>=3.0.0",
 ]
+klavis = [
+    "klavis>=2.18.0",
+]
 
 # Optional deps for LangGraph example/tests
 langgraph = [

tests/pytest/datasets/klavis_gmail_sandbox_test.jsonl

@@ -0,0 +1,2 @@
+{"messages": [{"role": "system", "content": "You are a helpful assistant with access to Gmail. You can send emails, draft emails, and manage messages."}, {"role": "user", "content": "Send an email to john@example.com with subject 'Meeting Tomorrow' and body 'Hi John, Just confirming our meeting tomorrow at 2pm. Best regards.'"}], "ground_truth": "One email sent to john@example.com with subject 'Meeting Tomorrow' containing meeting confirmation"}
+{"messages": [{"role": "system", "content": "You are a helpful assistant with access to Gmail. You can send emails, draft emails, and manage messages."}, {"role": "user", "content": "Draft an email to sarah@company.com with subject 'Project Update' and body 'Hi Sarah, The project is progressing well. I will send you the detailed report by Friday.'"}], "ground_truth": "One draft email created for sarah@company.com with subject 'Project Update' about project progress"}

tests/pytest/test_pytest_klavis_sandbox.py

@@ -0,0 +1,104 @@
+import json
+import logging
+import os
+
+from eval_protocol.models import EvaluateResult, EvaluationRow
+from eval_protocol.pytest import KlavisSandboxRolloutProcessor, evaluation_test
+from openai import AsyncOpenAI
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class ResponseFormat(BaseModel):
+    score: float
+    reasoning: str
+
+
+@evaluation_test(
+    input_dataset=["tests/pytest/datasets/klavis_gmail_sandbox_test.jsonl"],
+    rollout_processor=KlavisSandboxRolloutProcessor(
+        server_name="gmail",
+        # Optional: provide custom initialization data factory
+        # initialize_data_factory=lambda row: {"messages": [], "drafts": []},
+    ),
+    completion_params=[{"model": "fireworks_ai/accounts/fireworks/models/deepseek-v3p2"}],
+    mode="pointwise",
+)
+async def test_pytest_gmail_sandbox(row: EvaluationRow) -> EvaluationRow:
+    """
+    Evaluate Gmail sandbox results by comparing with ground truth using LLM judge.
+    
+    The sandbox data is exported after agent execution and compared with expected output.
+    Sandbox data is available in row.execution_metadata.metadata["sandbox_data"].
+    """
+    ground_truth = row.ground_truth
+    sandbox_data = row.execution_metadata.extra.get("sandbox_data", {}) if row.execution_metadata.extra else {}
+    final_message = row.messages[-1].content if row.messages else ""
+
+    logger.info(f"Evaluating row {row.execution_metadata.rollout_id}")
+    logger.info(f"Final message: {final_message}")
+    logger.info(f"Sandbox data: {json.dumps(sandbox_data, indent=2, default=str)}")
+    logger.info(f"Ground truth: {ground_truth}")
+
+    async with AsyncOpenAI(
+        api_key=os.environ["FIREWORKS_API_KEY"], base_url="https://api.fireworks.ai/inference/v1"
+    ) as client:
+        # Use LLM to judge if the sandbox data matches the ground truth
+        evaluation_prompt = f"""You are evaluating an AI agent's performance on a Gmail task.
+
+Task: {row.messages[0].content if row.messages else 'N/A'}
+
+Ground Truth: {ground_truth}
+
+Agent's Final Response: {final_message}
+
+Gmail Sandbox State After Execution:
+{json.dumps(sandbox_data, indent=2, default=str)}
+
+Evaluate whether the agent successfully completed the task by checking:
+1. Did the agent understand and attempt the task?
+2. Does the sandbox data reflect the expected outcome described in the ground truth?
+3. Are there any emails sent/drafted that match the task requirements?
+
+Return:
+- score: 1.0 if task completed successfully, 0.5 if partially completed, 0.0 if failed
+- reasoning: Explain your evaluation in 1-2 sentences
+"""
+
+        try:
+            response = await client.chat.completions.create(
+                model="accounts/fireworks/models/deepseek-v3p2",
+                messages=[
+                    {
+                        "role": "system",
+                        "content": "You are a precise evaluator of AI agent performance. Analyze the task, execution, and results carefully.",
+                    },
+                    {"role": "user", "content": evaluation_prompt},
+                ],
+                response_format={
+                    "type": "json_schema",
+                    "json_schema": {"name": "ResponseFormat", "schema": ResponseFormat.model_json_schema()},
+                },
+                temperature=0.0,
+            )
+
+            response_text = response.choices[0].message.content
+            logger.info(f"LLM judge response: {response_text}")
+
+            parsed = json.loads(response_text or "{}")
+            score = parsed.get("score", 0.0)
+            reasoning = parsed.get("reasoning", "No reasoning provided")
+
+            row.evaluation_result = EvaluateResult(
+                score=score,
+                reason=reasoning,
+            )
+        except Exception as e:
+            logger.error(f"Error during LLM evaluation: {str(e)}", exc_info=True)
+            row.evaluation_result = EvaluateResult(
+                score=0.0,
+                reason=f"Evaluation error: {str(e)}",
+            )
+
+    return row