save

Author: Dylan Huang

eval_protocol/pytest/default_pydantic_ai_rollout_processor.py

@@ -13,7 +13,7 @@
 from openai.types.chat import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageParam
 from openai.types.chat.chat_completion import Choice as ChatCompletionChoice
 from pydantic import TypeAdapter
-from pydantic_ai import Agent
+from pydantic_ai import Agent, ModelSettings
 from pydantic_ai._utils import generate_tool_call_id
 from pydantic_ai.messages import ModelMessage
 from pydantic_ai.messages import (
@@ -46,7 +46,6 @@ def __call__(self, rows: list[EvaluationRow], config: RolloutProcessorConfig) ->
         """Create agent rollout tasks and return them for external handling."""
 
         semaphore = config.semaphore
-
         agent = self._setup_agent(config)
 
         async def process_row(row: EvaluationRow) -> EvaluationRow:
@@ -70,7 +69,10 @@ async def process_row(row: EvaluationRow) -> EvaluationRow:
             row.tools = tools
 
             model_messages = [self.convert_ep_message_to_pyd_message(m, row) for m in row.messages]
-            response = await agent.run(message_history=model_messages, usage_limits=config.kwargs.get("usage_limits"))
+            settings = self.construct_model_settings(agent, row)
+            response = await agent.run(
+                message_history=model_messages, usage_limits=config.kwargs.get("usage_limits"), model_settings=settings
+            )
             row.messages = await self.convert_pyd_message_to_ep_message(response.all_messages())
 
             # TODO: pydantic ai accumulates usage info across all models in multi-agent setup, so this simple tracking doesn't work for cost. to discuss with @dphuang2 when he's back.
@@ -98,6 +100,24 @@ async def convert_pyd_message_to_ep_message(self, messages: list[ModelMessage])
         oai_messages: list[ChatCompletionMessageParam] = await self._util._map_messages(messages)
         return [Message(**m) for m in oai_messages]  # pyright: ignore[reportArgumentType]
 
+    def construct_model_settings(self, agent: Agent, row: EvaluationRow) -> ModelSettings:
+        model = agent.model
+        if model and not isinstance(model, str) and model.settings:
+            # We must copy model settings to avoid concurrency issues by modifying the same object in-place
+            settings = model.settings.copy()
+        if settings is None:
+            settings = ModelSettings()
+        settings["extra_body"] = settings.get("extra_body", {})
+        extra_body = settings["extra_body"]
+        if isinstance(extra_body, dict):
+            extra_body["metadata"] = settings.get("metadata", {})
+            extra_body["metadata"]["row_id"] = row.input_metadata.row_id
+            extra_body["metadata"]["invocation_id"] = row.execution_metadata.invocation_id
+            extra_body["metadata"]["rollout_id"] = row.execution_metadata.rollout_id
+            extra_body["metadata"]["run_id"] = row.execution_metadata.run_id
+            extra_body["metadata"]["experiment_id"] = row.execution_metadata.experiment_id
+        return settings
+
     def convert_ep_message_to_pyd_message(self, message: Message, row: EvaluationRow) -> ModelMessage:
         if message.role == "assistant":
             type_adapter = TypeAdapter(ChatCompletionMessage)

tests/chinook/pydantic/test_pydantic_complex_queries.py

@@ -1,7 +1,7 @@
 import os
 from pydantic import BaseModel
 from pydantic_ai import Agent
-from pydantic_ai.models.openai import OpenAIChatModel
+from pydantic_ai.models.openai import OpenAIChatModel, OpenAIChatModelSettings
 import pytest
 
 from eval_protocol.models import EvaluateResult, EvaluationRow
@@ -14,15 +14,18 @@
 LLM_JUDGE_PROMPT = (
     "Your job is to compare the response to the expected answer.\n"
     "The response will be a narrative report of the query results.\n"
-    "If the response contains the same or well summarized information as the expected answer, return 1.0.\n"
-    "If the response does not contain the same information or is missing information, return 0.0."
+    "Return a score between 0.0 and 1.0, where 1.0 means the response contains all or well summarized information as the expected answer, "
+    "0.0 means the response does not contain the same information or is missing all key information, "
+    "and values in between represent partial credit for responses that are partially correct or contain some but not all of the expected information."
 )
 
 
 def agent_factory(config: RolloutProcessorConfig) -> Agent:
     model_name = config.completion_params["model"]
-    provider = config.completion_params["provider"]
-    model = OpenAIChatModel(model_name, provider=provider)
+    provider = config.completion_params.get("provider")
+    reasoning = config.completion_params.get("reasoning")
+    settings = OpenAIChatModelSettings(openai_reasoning_effort=reasoning)
+    model = OpenAIChatModel(model_name, provider=provider or "openai", settings=settings)
     return setup_agent(model)
 
 
@@ -38,8 +41,19 @@ def agent_factory(config: RolloutProcessorConfig) -> Agent:
             "model": "accounts/fireworks/models/kimi-k2-instruct",
             "provider": "fireworks",
         },
+        {
+            "model": "accounts/fireworks/models/deepseek-v3p1",
+            "provider": "fireworks",
+        },
+        {
+            "model": "accounts/fireworks/models/kimi-k2-instruct-0905",
+            "provider": "fireworks",
+        },
+        {"model": "gpt-5"},
+        {"model": "gpt-5", "reasoning": "high"},
     ],
     rollout_processor=PydanticAgentRolloutProcessor(agent_factory),
+    num_runs=2,
 )
 async def test_pydantic_complex_queries(row: EvaluationRow) -> EvaluationRow:
     """
@@ -58,8 +72,7 @@ async def test_pydantic_complex_queries(row: EvaluationRow) -> EvaluationRow:
         )
     else:
         model = OpenAIChatModel(
-            "accounts/fireworks/models/kimi-k2-instruct",
-            provider="fireworks",
+            "gpt-5",
         )
 
         class Response(BaseModel):

tests/chinook/pydantic/test_pydantic_complex_queries_responses.py

@@ -19,7 +19,7 @@
 
 def agent_factory(config: RolloutProcessorConfig) -> Agent:
     model_name = config.completion_params["model"]
-    reasoning = config.completion_params["reasoning"]
+    reasoning = config.completion_params.get("reasoning")
     settings = OpenAIResponsesModelSettings(
         openai_reasoning_effort=reasoning,
     )
@@ -37,7 +37,6 @@ def agent_factory(config: RolloutProcessorConfig) -> Agent:
     completion_params=[
         {
             "model": "gpt-5",
-            "reasoning": "high",
         },
     ],
     rollout_processor=PydanticAgentRolloutProcessor(agent_factory),