update logic

Author: zihaolin96

eval_protocol/pytest/default_klavis_sandbox_rollout_processor.py

@@ -30,7 +30,6 @@ def __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:
@@ -63,27 +62,46 @@ 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()
+            agent: Agent | None = None
+            temp_config_path: str | None = None
+            sandbox: CreateSandboxResponse | None = None
 
             try:
+                # Step 0: Create a sandbox for this row
+                sandbox = self._init_sandbox()
+                logger.info(f"Sandbox created: {sandbox}")
+
                 # Step 1: Initialize data in the sandbox
+                init_data: Dict[str, Any] | None = None
                 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")
-
+                else:
+                    # Allow datasets to provide initialization payload directly
+                    init_data = (
+                        (row.input_metadata.session_data or {}).get("initialize_data")
+                        if row.input_metadata is not None
+                        else None
+                    )
+                
+                if init_data:
+                    logger.info(f"Initializing {self.server_name} sandbox {sandbox.sandbox_id}")
+                    initialize_method = getattr(
+                        self.klavis_client.sandbox, f"initialize_{sandbox.server_name.value}_sandbox"
+                    )
+                    init_response = initialize_method(sandbox_id=sandbox.sandbox_id, **init_data)
+                    logger.info(f"Initialization response: {init_response}")
+                    
                 # 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)
+                temp_config_path = self.create_mcp_config(
+                    server_url=sandbox.server_url, server_key=sandbox.server_name.value
+                )
+                logger.info(f"MCP config created: {temp_config_path}")
 
                 # 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")
@@ -106,16 +124,16 @@ async def process_row(row: EvaluationRow) -> EvaluationRow:
                 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)
+                dump_method = getattr(self.klavis_client.sandbox, f"dump_{sandbox.server_name.value}_sandbox")
+                dump_response = dump_method(sandbox_id=sandbox.sandbox_id)
                 sandbox_data = dump_response.data
+                logger.info(f"Sandbox data: {sandbox_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["sandbox_id"] = sandbox.sandbox_id
                 row.execution_metadata.extra["server_name"] = self.server_name
 
             except Exception as e:
@@ -133,15 +151,14 @@ async def process_row(row: EvaluationRow) -> EvaluationRow:
                     os.unlink(temp_config_path)
 
                 # Release sandbox
-                if self.sandbox.sandbox_id:
+                if sandbox and 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
+                            server_name=sandbox.server_name, sandbox_id=sandbox.sandbox_id
                         )
-                        logger.info(f"Sandbox {self.sandbox.sandbox_id} released successfully")
+                        logger.info(f"Sandbox {sandbox.sandbox_id} released successfully")
                     except Exception as e:
-                        logger.error(f"Error releasing sandbox {self.sandbox.sandbox_id}: {str(e)}", exc_info=True)
+                        logger.error(f"Error releasing sandbox {sandbox.sandbox_id}: {str(e)}", exc_info=True)
 
                 row.execution_metadata.rollout_duration_seconds = time.perf_counter() - start_time
 

tests/pytest/datasets/klavis_gmail_sandbox_test.jsonl

@@ -1,2 +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 [email protected] with subject 'Project Update' about project progress"}
+{"initialize_data": {"messages": [{"subject": "Project Update", "to": "[email protected]", "body": "The project is progressing well. We should have the final deliverables by next week.", "cc": "", "bcc": "", "from": "[email protected]", "reply_to": "", "labels": ["INBOX"]}, {"subject": "Spam Newsletter", "to": "zihao@klavisai.com", "body": "Check out our amazing deals! Click here now!", "cc": "", "bcc": "", "from": "[email protected]", "reply_to": "", "labels": ["INBOX"]}], "drafts": []}, "messages": "Please delete the email with subject \"Spam Newsletter\" from my inbox.", "ground_truth": {"messages": [{"subject": "Project Update", "to": "zihao@klavisai.com", "body": "The project is progressing well. We should have the final deliverables by next week.", "cc": "", "bcc": "", "from": "[email protected]", "reply_to": "", "labels": ["INBOX"]}], "drafts": []}}
+{"initialize_data": {"messages": [], "drafts": []}, "messages": "Please directly send an email to [email protected] with subject \"Meeting Tomorrow\" and body \"Hi Zihao, just confirming our meeting tomorrow at 2pm. Best regards.\"", "ground_truth": {"messages": [{"subject": "Meeting Tomorrow", "to": "zihao@klavisai.com", "body": "Hi Zihao, just confirming our meeting tomorrow at 2pm. Best regards.", "cc": "", "bcc": "", "from": "", "reply_to": "", "labels": ["SENT"]}], "drafts": []}}

tests/pytest/test_pytest_klavis_sandbox.py

@@ -2,7 +2,7 @@
 import logging
 import os
 
-from eval_protocol.models import EvaluateResult, EvaluationRow
+from eval_protocol.models import EvaluateResult, EvaluationRow, Message
 from eval_protocol.pytest import KlavisSandboxRolloutProcessor, evaluation_test
 from openai import AsyncOpenAI
 from pydantic import BaseModel
@@ -12,29 +12,68 @@
 
 class ResponseFormat(BaseModel):
     score: float
-    reasoning: str
+    
+
+def klavis_gmail_sandbox_dataset_adapter(rows: list[dict]) -> list[EvaluationRow]:
+    """Dataset adapter for sandbox JSONL rows.
+
+    Supports the new schema:
+      - initialize_data: dict (passed to Klavis sandbox initializer)
+      - messages: str (task instruction)
+      - ground_truth: dict (expected final sandbox state)
+
+    """
+    adapted: list[EvaluationRow] = []
+    system_prompt = (
+        "You are a helpful assistant with access to Gmail. "
+        "You can send emails, draft emails, and manage messages, etc."
+    )
+
+    for r in rows:
+        if isinstance(r.get("messages"), str) and "initialize_data" in r:
+            init_data = r.get("initialize_data") or {}
+            task = r.get("messages") or ""
+            ground_truth = r.get("ground_truth")
+
+            row = EvaluationRow(
+                messages=[
+                    Message(role="system", content=system_prompt),
+                    Message(role="user", content=task),
+                ],
+                ground_truth=ground_truth,
+            )
+            row.input_metadata.session_data = {
+                "initialize_data": init_data,
+                "task": task,
+            }
+            adapted.append(row)
+        else:
+            adapted.append(EvaluationRow(**r))
+
+    return adapted
 
 
 @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"}],
+    completion_params=[{"model": "fireworks_ai/accounts/fireworks/models/kimi-k2-thinking"}],
     mode="pointwise",
+    dataset_adapter=klavis_gmail_sandbox_dataset_adapter,
 )
 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"].
+    Sandbox data is available in row.execution_metadata.extra["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 ""
+    initialize_data = (row.input_metadata.session_data or {}).get("initialize_data", {})
+    task = (row.input_metadata.session_data or {}).get("task", "")
 
     logger.info(f"Evaluating row {row.execution_metadata.rollout_id}")
     logger.info(f"Final message: {final_message}")
@@ -44,31 +83,34 @@ async def test_pytest_gmail_sandbox(row: EvaluationRow) -> EvaluationRow:
     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'}
+        evaluation_prompt = f"""You are evaluating an AI agent's performance on a Gmail sandbox task.
+
+Task:
+{task or (row.messages[-1].content if row.messages else 'N/A')}
 
-Ground Truth: {ground_truth}
+Initial Gmail Sandbox State (initialize_data):
+{json.dumps(initialize_data, indent=2, default=str)}
 
-Agent's Final Response: {final_message}
+Expected Final Gmail Sandbox State (ground_truth):
+{json.dumps(ground_truth, indent=2, default=str)}
 
 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?
+1. Does the final sandbox state match the expected ground_truth state?
+2. If there are small formatting differences, judge semantically
+3. Use the initial state only as context; the key is whether the correct changes happened.
 
 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",
+                model="accounts/fireworks/models/kimi-k2-thinking",
                 messages=[
                     {
                         "role": "system",
@@ -88,12 +130,8 @@ async def test_pytest_gmail_sandbox(row: EvaluationRow) -> EvaluationRow:
 
             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,
-            )
+            row.evaluation_result = EvaluateResult(score=score)
         except Exception as e:
             logger.error(f"Error during LLM evaluation: {str(e)}", exc_info=True)
             row.evaluation_result = EvaluateResult(