22import logging
33import os
44
5- from eval_protocol .models import EvaluateResult , EvaluationRow , Message
6- from eval_protocol .pytest import KlavisSandboxRolloutProcessor , evaluation_test
5+ import pytest
76from openai import AsyncOpenAI
87from pydantic import BaseModel
98
9+ from eval_protocol .models import EvaluateResult , EvaluationRow
10+ from eval_protocol .pytest import KlavisSandboxRolloutProcessor
11+ from eval_protocol .pytest .evaluation_test import evaluation_test
12+ from tests .klavis_helpers import MODEL , klavis_gmail_sandbox_dataset_adapter
13+
1014logger = logging .getLogger (__name__ )
1115
1216
1317class ResponseFormat (BaseModel ):
1418 score : float
15-
16-
17- def klavis_gmail_sandbox_dataset_adapter (rows : list [dict ]) -> list [EvaluationRow ]:
18- """Dataset adapter for sandbox JSONL rows.
19-
20- Supports the new schema:
21- - initialize_data: dict (passed to Klavis sandbox initializer)
22- - messages: str (task instruction)
23- - ground_truth: dict (expected final sandbox state)
24-
25- """
26- adapted : list [EvaluationRow ] = []
27- system_prompt = (
28- "You are a helpful assistant with access to Gmail. "
29- "You can send emails, draft emails, and manage messages, etc."
30- )
31-
32- for r in rows :
33- if isinstance (r .get ("messages" ), str ) and "initialize_data" in r :
34- init_data = r .get ("initialize_data" ) or {}
35- task = r .get ("messages" ) or ""
36- ground_truth = r .get ("ground_truth" )
37-
38- row = EvaluationRow (
39- messages = [
40- Message (role = "system" , content = system_prompt ),
41- Message (role = "user" , content = task ),
42- ],
43- ground_truth = ground_truth ,
44- )
45- row .input_metadata .session_data = {
46- "initialize_data" : init_data ,
47- "task" : task ,
48- }
49- adapted .append (row )
50- else :
51- adapted .append (EvaluationRow (** r ))
52-
53- return adapted
5419
5520
21+ @pytest .mark .skipif (
22+ not (os .environ .get ("KLAVIS_API_KEY" ) and os .environ .get ("FIREWORKS_API_KEY" )),
23+ reason = "Live Klavis Gmail test requires KLAVIS_API_KEY and FIREWORKS_API_KEY" ,
24+ )
5625@evaluation_test (
5726 input_dataset = ["tests/pytest/datasets/klavis_gmail_sandbox_test.jsonl" ],
5827 rollout_processor = KlavisSandboxRolloutProcessor (
5928 server_name = "gmail" ,
6029 ),
61- completion_params = [{"model" : "fireworks_ai/accounts/fireworks/models/kimi-k2p5 " }],
30+ completion_params = [{"model" : f "fireworks_ai/{ MODEL } " }],
6231 mode = "pointwise" ,
6332 dataset_adapter = klavis_gmail_sandbox_dataset_adapter ,
6433)
6534async def test_pytest_gmail_sandbox (row : EvaluationRow ) -> EvaluationRow :
6635 """
6736 Evaluate Gmail sandbox results by comparing with ground truth using LLM judge.
68-
37+
6938 The sandbox data is exported after agent execution and compared with expected output.
7039 Sandbox data is available in row.execution_metadata.extra["sandbox_data"].
7140 """
41+ assert row .rollout_status .is_finished (), row .rollout_status
7242 ground_truth = row .ground_truth
7343 sandbox_data = row .execution_metadata .extra .get ("sandbox_data" , {}) if row .execution_metadata .extra else {}
7444 final_message = row .messages [- 1 ].content if row .messages else ""
@@ -83,11 +53,10 @@ async def test_pytest_gmail_sandbox(row: EvaluationRow) -> EvaluationRow:
8353 async with AsyncOpenAI (
8454 api_key = os .environ ["FIREWORKS_API_KEY" ], base_url = "https://api.fireworks.ai/inference/v1"
8555 ) as client :
86-
8756 evaluation_prompt = f"""You are evaluating an AI agent's performance on a Gmail sandbox task.
8857
8958Task:
90- { task or (row .messages [- 1 ].content if row .messages else ' N/A' )}
59+ { task or (row .messages [- 1 ].content if row .messages else " N/A" )}
9160
9261Initial Gmail Sandbox State (initialize_data):
9362{ json .dumps (initialize_data , indent = 2 , default = str )}
@@ -110,7 +79,7 @@ async def test_pytest_gmail_sandbox(row: EvaluationRow) -> EvaluationRow:
11079
11180 try :
11281 response = await client .chat .completions .create (
113- model = "accounts/fireworks/models/kimi-k2p5" ,
82+ model = MODEL ,
11483 messages = [
11584 {
11685 "role" : "system" ,
@@ -134,9 +103,6 @@ async def test_pytest_gmail_sandbox(row: EvaluationRow) -> EvaluationRow:
134103 row .evaluation_result = EvaluateResult (score = score )
135104 except Exception as e :
136105 logger .error (f"Error during LLM evaluation: { str (e )} " , exc_info = True )
137- row .evaluation_result = EvaluateResult (
138- score = 0.0 ,
139- reason = f"Evaluation error: { str (e )} " ,
140- )
106+ raise
141107
142108 return row
0 commit comments