factory pattern works

Author: Dylan Huang

eval_protocol/pytest/default_pydantic_ai_rollout_processor.py

@@ -4,19 +4,13 @@
 from collections.abc import Callable
 import logging
 import time
-import types
-from pydantic_ai.models import Model
 from pydantic_ai.usage import UsageLimits
 from typing_extensions import override
 from eval_protocol.models import EvaluationRow, Message
-from openai.types import CompletionUsage
 from eval_protocol.pytest.rollout_processor import RolloutProcessor
 from eval_protocol.pytest.types import RolloutProcessorConfig
 from openai.types.chat import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageParam
 from openai.types.chat.chat_completion import Choice as ChatCompletionChoice
-from openai.types.chat.chat_completion_assistant_message_param import (
-    ChatCompletionAssistantMessageParam,
-)
 from pydantic import TypeAdapter
 from pydantic_ai import Agent
 from pydantic_ai._utils import generate_tool_call_id
@@ -27,8 +21,6 @@
     ToolReturnPart,
     UserPromptPart,
 )
-from pydantic_ai.models.anthropic import AnthropicModel
-from pydantic_ai.models.google import GoogleModel
 from pydantic_ai.models.openai import OpenAIModel
 from pydantic_ai.providers.openai import OpenAIProvider
 
@@ -39,64 +31,27 @@ class PydanticAgentRolloutProcessor(RolloutProcessor):
     """Rollout processor for Pydantic AI agents. Mainly converts
     EvaluationRow.messages to and from Pydantic AI ModelMessage format."""
 
-    def __init__(self, setup_agent: Callable[..., Agent] | Agent, usage_limits: UsageLimits | None = None):
+    def __init__(
+        self, agent_factory: Callable[[RolloutProcessorConfig], Agent], usage_limits: UsageLimits | None = None
+    ):
         # dummy model used for its helper functions for processing messages
-        self.util: OpenAIModel = OpenAIModel("dummy-model", provider=OpenAIProvider(api_key="dummy"))
+        self._util: OpenAIModel = OpenAIModel("dummy-model", provider=OpenAIProvider(api_key="dummy"))
+        self._setup_agent = agent_factory
 
     @override
     def __call__(self, rows: list[EvaluationRow], config: RolloutProcessorConfig) -> list[asyncio.Task[EvaluationRow]]:
         """Create agent rollout tasks and return them for external handling."""
 
         semaphore = config.semaphore
 
-        # validate that the "agent" field is present with a valid Pydantic AI Agent instance in the completion_params dict
-        if "agent" not in config.kwargs:
-            raise ValueError("kwargs must contain an 'agent' field with a valid Pydantic AI Agent instance")
-        if not isinstance(config.kwargs["agent"], Agent) and not isinstance(
-            config.kwargs["agent"], types.FunctionType
-        ):
-            raise ValueError(
-                "kwargs['agent'] must be a valid Pydantic AI Agent instance or a function that returns an Agent"
-            )
-
-        if isinstance(config.kwargs["agent"], types.FunctionType):
-            setup_agent = config.kwargs["agent"]
-            if not isinstance(config.completion_params["model"], dict):
-                raise ValueError(
-                    "completion_params['model'] must be a dict mapping agent argument names to model config dicts (with 'model' and 'provider' keys)"
-                )
-            kwargs: dict[str, Model] = {}
-            for k, v in config.completion_params["model"].items():  # pyright: ignore[reportUnknownVariableType]
-                if v["model"] and v["model"].startswith("anthropic:"):  # pyright: ignore[reportUnknownMemberType]
-                    kwargs[k] = AnthropicModel(
-                        v["model"].removeprefix("anthropic:"),  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
-                    )
-                elif v["model"] and v["model"].startswith("google:"):  # pyright: ignore[reportUnknownMemberType]
-                    kwargs[k] = GoogleModel(
-                        v["model"].removeprefix("google:"),  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
-                    )
-                else:
-                    kwargs[k] = OpenAIModel(
-                        v["model"],  # pyright: ignore[reportUnknownArgumentType]
-                        provider=v["provider"],  # pyright: ignore[reportUnknownArgumentType]
-                    )
-            agent_instance: Agent = setup_agent(**kwargs)  # pyright: ignore[reportAny]
-            model = None
-        else:
-            agent_instance = config.kwargs["agent"]  # pyright: ignore[reportAssignmentType]
-            model = OpenAIModel(
-                config.completion_params["model"],  # pyright: ignore[reportAny]
-                provider=config.completion_params["provider"],  # pyright: ignore[reportAny]
-            )
+        agent = self._setup_agent(config)
 
         async def process_row(row: EvaluationRow) -> EvaluationRow:
             """Process a single row with agent rollout."""
             start_time = time.perf_counter()
 
             model_messages = [self.convert_ep_message_to_pyd_message(m, row) for m in row.messages]
-            response = await agent_instance.run(
-                message_history=model_messages, model=model, usage_limits=config.kwargs.get("usage_limits")
-            )
+            response = await agent.run(message_history=model_messages, usage_limits=config.kwargs.get("usage_limits"))
             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.
@@ -121,15 +76,15 @@ async def _sem_wrapper(r: EvaluationRow) -> EvaluationRow:
         return tasks
 
     async def convert_pyd_message_to_ep_message(self, messages: list[ModelMessage]) -> list[Message]:
-        oai_messages: list[ChatCompletionMessageParam] = await self.util._map_messages(messages)
+        oai_messages: list[ChatCompletionMessageParam] = await self._util._map_messages(messages)
         return [Message(**m) for m in oai_messages]  # pyright: ignore[reportArgumentType]
 
     def convert_ep_message_to_pyd_message(self, message: Message, row: EvaluationRow) -> ModelMessage:
         if message.role == "assistant":
             type_adapter = TypeAdapter(ChatCompletionMessage)
             oai_message = type_adapter.validate_python(message)
             # Fix: Provide required finish_reason and index, and ensure created is int (timestamp)
-            return self.util._process_response(
+            return self._util._process_response(
                 ChatCompletion(
                     choices=[ChatCompletionChoice(message=oai_message, finish_reason="stop", index=0)],
                     object="chat.completion",

tests/chinook/pydantic/test_pydantic_chinook.py

@@ -6,6 +6,7 @@
 from eval_protocol.pytest import evaluation_test
 
 from eval_protocol.pytest.default_pydantic_ai_rollout_processor import PydanticAgentRolloutProcessor
+from eval_protocol.pytest.types import RolloutProcessorConfig
 from tests.chinook.pydantic.agent import setup_agent
 import os
 from pydantic_ai.models.openai import OpenAIModel
@@ -20,20 +21,23 @@
 )
 
 
+def agent_factory(config: RolloutProcessorConfig) -> Agent:
+    model_name = config.completion_params["model"]
+    provider = config.completion_params["provider"]
+    model = OpenAIModel(model_name, provider=provider)
+    return setup_agent(model)
+
+
 @pytest.mark.asyncio
 @evaluation_test(
     input_messages=[[[Message(role="user", content="What is the total number of tracks in the database?")]]],
     completion_params=[
         {
-            "model": {
-                "orchestrator_agent_model": {
-                    "model": "accounts/fireworks/models/kimi-k2-instruct",
-                    "provider": "fireworks",
-                }
-            }
+            "model": "accounts/fireworks/models/kimi-k2-instruct",
+            "provider": "fireworks",
         },
     ],
-    rollout_processor=PydanticAgentRolloutProcessor(setup_agent),
+    rollout_processor=PydanticAgentRolloutProcessor(agent_factory),
     mode="pointwise",
 )
 async def test_simple_query(row: EvaluationRow) -> EvaluationRow:
@@ -91,15 +95,11 @@ class Response(BaseModel):
     input_rows=[collect_dataset()],
     completion_params=[
         {
-            "model": {
-                "orchestrator_agent_model": {
-                    "model": "accounts/fireworks/models/kimi-k2-instruct",
-                    "provider": "fireworks",
-                }
-            }
+            "model": "accounts/fireworks/models/kimi-k2-instruct",
+            "provider": "fireworks",
         },
     ],
-    rollout_processor=PydanticAgentRolloutProcessor(setup_agent),
+    rollout_processor=PydanticAgentRolloutProcessor(agent_factory),
     mode="pointwise",
 )
 async def test_complex_queries(row: EvaluationRow) -> EvaluationRow: