gpt-oss example e2e

Author: xzrderek

eval_protocol/pytest/default_mcp_gym_rollout_processor.py

@@ -1,18 +1,17 @@
 import asyncio
+import atexit
 import os
+import signal
+import socket
 import subprocess
 import time
-import socket
 from pathlib import Path
 from typing import List, Optional
 
 import eval_protocol as ep
 from eval_protocol.models import EvaluationRow, Message
 from eval_protocol.pytest.types import RolloutProcessorConfig
 
-import atexit
-import signal
-
 
 class MCPServerManager:
     """Manages MCP server lifecycle for testing."""
@@ -182,50 +181,52 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         return False  # Don't suppress exceptions
 
 
-
-async def default_mcp_gym_rollout_processor(rows: List[EvaluationRow], config: RolloutProcessorConfig) -> List[EvaluationRow]:
+async def default_mcp_gym_rollout_processor(
+    rows: List[EvaluationRow], config: RolloutProcessorConfig
+) -> List[EvaluationRow]:
     """
     Rollout processor for tau bench environments.
-    
+
     This processor starts an MCP server, creates tau bench environments, and runs rollouts
     using the eval_protocol framework, following the pattern from test_tau2_e2e.py.
-    
+
     Args:
         rows: List of EvaluationRow objects containing messages and dataset info in input_metadata
         config: RolloutProcessorConfig with model and other parameters
-    
+
     Returns:
         List of EvaluationRow objects with completed conversations
     """
     server = MCPServerManager(config.server_script_path, port=9700)
-    
+
     try:
         server.start()
-        
+
         policy = ep.LiteLLMPolicy(
             model_id=config.model,
-            temperature=config.input_params.get('temperature', 0.0),
-            max_tokens=config.input_params.get('max_tokens', 4096),
+            temperature=config.input_params.get("temperature", 0.0),
+            max_tokens=config.input_params.get("max_tokens", 4096),
+            reasoning_effort=config.input_params.get("reasoning_effort", None),
         )
-        
+
         # Create MCP environments directly from evaluation_rows
         envs = ep.make(
-            'http://localhost:9700/mcp/',
+            "http://localhost:9700/mcp/",
             evaluation_rows=rows,
             model_id=policy.model_id,
         )
-        
+
         # Run rollout with environments and policy
         evaluation_rows = await ep.rollout(
-            envs, 
-            policy=policy, 
+            envs,
+            policy=policy,
             evaluation_rows=rows,
-            steps=config.steps, 
-            max_concurrent_rollouts=config.max_concurrent_rollouts
+            steps=config.steps,
+            max_concurrent_rollouts=config.max_concurrent_rollouts,
         )
-        
+
         return evaluation_rows
-        
+
     finally:
         # Always clean up the server
         server.stop()

tests/pytest/test_tau_bench_airline.py

@@ -10,10 +10,9 @@
 from pathlib import Path
 from typing import Any, Dict, List
 
-from eval_protocol.models import EvaluateResult, EvaluationRow, Message, InputMetadata, CompletionParams
+from eval_protocol.models import CompletionParams, EvaluateResult, EvaluationRow, InputMetadata, Message
 from eval_protocol.pytest import evaluation_test
 from eval_protocol.pytest.default_mcp_gym_rollout_processor import default_mcp_gym_rollout_processor
-
 from vendor.tau2.data_model.message import (
     AssistantMessage,
     SystemMessage,
@@ -28,20 +27,21 @@
 from vendor.tau2.evaluator.evaluator_nl_assertions import NLAssertionsEvaluator
 from vendor.tau2.registry import registry
 
+
 def tau_bench_airline_to_evaluation_row(data: List[Dict[str, Any]]) -> List[EvaluationRow]:
     """
     Convert entries from airline dataset to EvaluationRow objects.
     """
     rows = []
     test_dir = Path(__file__).parent.parent.parent / "examples" / "tau2_mcp" / "tests"
-    
+
     # Load system prompt from file so we can change it in one place
     domain = data[0]["environment_context"]["domain"]
     prompt_file = test_dir / f"system_prompts/{domain}_agent_system_prompt.md"
-    
+
     with open(prompt_file, "r") as f:
         system_prompt = f.read().strip()
-    
+
     for row in data:
         eval_row = EvaluationRow(
             messages=[Message(role="system", content=system_prompt)],
@@ -52,19 +52,20 @@ def tau_bench_airline_to_evaluation_row(data: List[Dict[str, Any]]) -> List[Eval
                     "user_simulation": row["user_simulation"],
                     "evaluation_criteria": row["evaluation_criteria"],
                     "user_prompt_template": row["user_prompt_template"],
-                }
+                },
             ),
         )
-        
+
         rows.append(eval_row)
-    
+
     return rows
 
+
 @evaluation_test(
     input_dataset=["tests/pytest/data/airline_dataset.jsonl"],
     dataset_adapter=tau_bench_airline_to_evaluation_row,
     model=["fireworks_ai/accounts/fireworks/models/gpt-oss-120b"],
-    rollout_input_params=[{"temperature": 0.8, "max_tokens": 4096}],
+    rollout_input_params=[{"temperature": 0.8, "max_tokens": 4096, "reasoning_effort": "high"}],
     rollout_processor=default_mcp_gym_rollout_processor,
     threshold_of_success=0.4,
     num_runs=1,
@@ -75,22 +76,22 @@ def tau_bench_airline_to_evaluation_row(data: List[Dict[str, Any]]) -> List[Eval
 def test_tau_bench_airline_evaluation(row: EvaluationRow) -> EvaluationRow:
     """
     Test tau bench airline evaluation using the pytest framework.
-    
+
     This test now uses the tau_bench_airline_reward function which automatically
     extracts evaluation criteria from dataset entries. No wrapper needed!
-    
+
     Args:
         row: EvaluationRow object from tau bench airline dataset after rollout
-    
+
     Returns:
         EvaluationRow with tau2 evaluation results
     """
     messages = row.messages
-    
+
     # Get evaluation criteria and user_simulation from input_metadata.dataset_info
     dataset_info = row.input_metadata.dataset_info if row.input_metadata else {}
     evaluation_criteria = dataset_info.get("evaluation_criteria", {})
-    
+
     nl_assertions = evaluation_criteria.get("nl_assertions", [])
     communicate_info = evaluation_criteria.get("communicate_info", [])
     actions = evaluation_criteria.get("actions", [])
@@ -130,7 +131,7 @@ def test_tau_bench_airline_evaluation(row: EvaluationRow) -> EvaluationRow:
         actions=actions,
         reward_basis=[
             RewardType.DB,
-            RewardType.ACTION,
+            RewardType.COMMUNICATE,
         ],
     )
 
@@ -226,4 +227,4 @@ def test_tau_bench_airline_evaluation(row: EvaluationRow) -> EvaluationRow:
         reason=reason,
         metrics={},
     )
-    return row
+    return row

vendor/tau2/evaluator/evaluator_nl_assertions.py

@@ -1,4 +1,7 @@
 import json
+from typing import List
+
+from pydantic import BaseModel
 
 from vendor.tau2.config import DEFAULT_LLM_NL_ASSERTIONS, DEFAULT_LLM_NL_ASSERTIONS_ARGS
 from vendor.tau2.data_model.message import Message, SystemMessage, UserMessage
@@ -7,6 +10,20 @@
 from vendor.tau2.utils.llm_utils import generate
 
 
+class NLAssertionResult(BaseModel):
+    """Individual NL assertion evaluation result."""
+
+    expectedOutcome: str
+    reasoning: str
+    metExpectation: bool
+
+
+class NLAssertionsResponse(BaseModel):
+    """Complete NL assertions evaluation response."""
+
+    results: List[NLAssertionResult]
+
+
 class NLAssertionsEvaluator:
     """
     Judge that evaluates whether a trajectory adheres to all the natural-language assertions.
@@ -37,9 +54,7 @@ def calculate_reward(
                 reward_breakdown={RewardType.NL_ASSERTION: 1.0},
             )
 
-        nl_assertions_checks = cls.evaluate_nl_assertions(
-            full_trajectory, nl_assertions
-        )
+        nl_assertions_checks = cls.evaluate_nl_assertions(full_trajectory, nl_assertions)
 
         # Calculate reward: 1 if all expectations are met, 0 otherwise
         all_expectations_met = all(result.met for result in nl_assertions_checks)
@@ -70,9 +85,7 @@ def evaluate_nl_assertions(
             - metExpectation: Boolean indicating if the assertion was met
             - reasoning: Explanation for the evaluation
         """
-        trajectory_str = "\n".join(
-            [f"{message.role}: {message.content}" for message in trajectory]
-        )
+        trajectory_str = "\n".join([f"{message.role}: {message.content}" for message in trajectory])
         # System prompt similar to the TypeScript implementation
         system_prompt = """
         TASK
@@ -86,7 +99,7 @@ def evaluate_nl_assertions(
         - `reasoning`: a short explanation for your classification
         - `metExpectation`: `true` if the agent satisfies the expected outcomes, `false` otherwise
         - `expectedOutcome`: repeat the expectation from the input that you are grading
-        
+
         Example response structure:
         {
             "results": [
@@ -102,7 +115,7 @@ def evaluate_nl_assertions(
         user_prompt = f"""
         conversation:
         {trajectory_str}
-        
+
         expectedOutcomes:
         {nl_assertions}
         """
@@ -115,8 +128,12 @@ def evaluate_nl_assertions(
         assistant_message = generate(
             model=DEFAULT_LLM_NL_ASSERTIONS,
             messages=messages,
-            **DEFAULT_LLM_NL_ASSERTIONS_ARGS,
-        )
+            temperature=0.0,
+            response_format={
+                "type": "json_schema",
+                "json_schema": {"name": "NLAssertionsResponse", "schema": NLAssertionsResponse.model_json_schema()},
+            },
+        )  # Adding constrained generation to ensure the response is a valid JSON object
         result_data = json.loads(assistant_message.content)
         return [
             NLAssertionCheck(