TODO: refactor rolloutprocessor to not use __call__

Author: Dylan Huang

eval_protocol/pytest/default_pydantic_ai_rollout_processor.py

@@ -1,9 +1,8 @@
 import asyncio
 import logging
 import types
-from typing import List
+from typing import List, Literal
 
-from attr import dataclass
 from openai.types.chat.chat_completion_assistant_message_param import ChatCompletionAssistantMessageParam
 
 from eval_protocol.models import EvaluationRow, Message
@@ -18,14 +17,15 @@
 from pydantic_ai.messages import ModelMessage
 from pydantic_ai._utils import generate_tool_call_id
 from pydantic_ai import Agent
+from pydantic_ai.usage import UsageLimits
 from pydantic_ai.messages import (
     ModelRequest,
     SystemPromptPart,
     ToolReturnPart,
     UserPromptPart,
 )
 from pydantic_ai.providers.openai import OpenAIProvider
-from typing_extensions import TypedDict
+from typing_extensions import Callable
 
 logger = logging.getLogger(__name__)
 
@@ -34,9 +34,33 @@ class PydanticAgentRolloutProcessor(RolloutProcessor):
     """Rollout processor for Pydantic AI agents. Mainly converts
     EvaluationRow.messages to and from Pydantic AI ModelMessage format."""
 
-    def __init__(self):
+    def __init__(self, setup_agent: Callable[..., Agent], usage_limits: UsageLimits = None):
         # dummy model used for its helper functions for processing messages
         self.util = OpenAIModel("dummy-model", provider=OpenAIProvider(api_key="dummy"))
+        self.setup_agent = setup_agent
+        self.usage_limits = usage_limits
+
+    def _map_litellm_to_pydantic_ai(
+        self, model_name: str
+    ) -> Literal[
+        "openai",
+        "deepseek",
+        "azure",
+        "openrouter",
+        "grok",
+        "fireworks",
+        "together",
+    ]:
+        mapping = {
+            "fireworks_ai": "fireworks",
+            "together_ai": "together",
+            "xai": "grok",
+            "azure_ai": "azure",
+        }
+        provider = model_name.split("/")[0]
+        if provider in mapping:
+            provider = mapping[provider]
+        return provider  # type: ignore
 
     def __call__(self, rows: List[EvaluationRow], config: RolloutProcessorConfig) -> List[asyncio.Task[EvaluationRow]]:
         """Create agent rollout tasks and return them for external handling."""
@@ -60,20 +84,28 @@ def __call__(self, rows: List[EvaluationRow], config: RolloutProcessorConfig) ->
                 raise ValueError(
                     "completion_params['model'] must be a dict mapping agent argument names to model config dicts (with 'model' and 'provider' keys)"
                 )
-            kwargs = {}
-            for k, v in config.completion_params["model"].items():
-                if v["model"] and v["model"].startswith("anthropic:"):
-                    kwargs[k] = AnthropicModel(
-                        v["model"].removeprefix("anthropic:"),
+            kwargs: dict[str, OpenAIModel | GoogleModel | AnthropicModel] = {}
+            for agent, model_config in config.completion_params["model"].items():
+                if "model" not in model_config:
+                    raise ValueError(f"model_config for agent {agent} must contain a 'model' key")
+                model_name = model_config["model"]
+                if model_name.startswith("anthropic/"):
+                    kwargs[agent] = AnthropicModel(
+                        model_name.removeprefix("anthropic/"),
+                    )
+                elif model_name.startswith("google/"):
+                    kwargs[agent] = GoogleModel(
+                        model_name.removeprefix("google/"),
                     )
-                elif v["model"] and v["model"].startswith("google:"):
-                    kwargs[k] = GoogleModel(
-                        v["model"].removeprefix("google:"),
+                elif model_name.startswith("gemini/"):
+                    kwargs[agent] = GoogleModel(
+                        model_name.removeprefix("gemini/"),
                     )
                 else:
-                    kwargs[k] = OpenAIModel(
-                        v["model"],
-                        provider=v["provider"],
+                    provider = self._map_litellm_to_pydantic_ai(model_name)
+                    kwargs[agent] = OpenAIModel(
+                        model_name.removeprefix(f"{provider}/"),
+                        provider=provider,
                     )
             agent = setup_agent(**kwargs)
             model = None
@@ -144,5 +176,4 @@ def convert_ep_message_to_pyd_message(self, message: Message, row: EvaluationRow
                     )
                 ]
             )
-        else:
-            raise ValueError(f"Unknown role: {message.role}")
+        raise ValueError(f"Unknown role: {message.role}")

eval_protocol/pytest/evaluation_test.py

@@ -544,11 +544,6 @@ def _log_eval_error(status: Status, rows: Optional[List[EvaluationRow]] | None,
                             row.input_metadata.row_id = generate_id(seed=0, index=index)
 
                     completion_params = kwargs["completion_params"]
-                    if completion_params and ("model" not in completion_params or not completion_params["model"]):
-                        raise ValueError(
-                            "No model provided. Please provide a model in the completion parameters object."
-                        )
-
                     # Create eval metadata with test function info and current commit hash
                     eval_metadata = EvalMetadata(
                         name=test_func.__name__,

tests/pytest/test_pydantic_multi_agent.py

@@ -47,12 +47,8 @@ async def joke_factory(ctx: RunContext[None], count: int) -> list[str]:
 
 @pytest.mark.asyncio
 @evaluation_test(
-    input_messages=[Message(role="user", content="Tell me a joke.")],
+    input_messages=[[Message(role="user", content="Tell me a joke.")]],
     completion_params=[
-        # single agent
-        {
-            "model": "fireworks_ai/accounts/fireworks/models/kimi-k2-instruct",
-        },
         # multi-agent
         {
             "joke_generation_model": {
@@ -62,21 +58,10 @@ async def joke_factory(ctx: RunContext[None], count: int) -> list[str]:
                 "model": "fireworks_ai/accounts/fireworks/models/deepseek-v3p1",
             },
         },
-        {
-            "joke_generation_model": {
-                "model": "fireworks_ai/accounts/fireworks/models/kimi-k2-instruct",
-            },
-            "joke_selection_model": {
-                "model": "fireworks_ai/accounts/fireworks/models/kimi-k2-instruct",
-            },
-        },
     ],
-    rollout_processor=PydanticAgentRolloutProcessor(),
-    rollout_processor_kwargs={
-        "agent": setup_agent,
-        # PydanticAgentRolloutProcessor will pass usage_limits into the "run" call
-        "usage_limits": UsageLimits(request_limit=5, total_tokens_limit=1000),
-    },
+    rollout_processor=PydanticAgentRolloutProcessor.__init__(
+        setup_agent, UsageLimits(request_limit=5, total_tokens_limit=1000)
+    ),
     mode="pointwise",
 )
 async def test_pydantic_multi_agent(row: EvaluationRow) -> EvaluationRow: