TODO: make test_pydantic_multi_agent work

Author: Dylan Huang

eval_protocol/pytest/default_pydantic_ai_rollout_processor.py

@@ -1,6 +1,6 @@
-import os
 import asyncio
 import logging
+import types
 from typing import List
 
 from openai.types.chat.chat_completion_assistant_message_param import ChatCompletionAssistantMessageParam
@@ -23,7 +23,6 @@
     UserPromptPart,
 )
 from pydantic_ai.providers.openai import OpenAIProvider
-from pydantic_ai.providers.fireworks import FireworksProvider
 
 logger = logging.getLogger(__name__)
 
@@ -45,20 +44,40 @@ def __call__(self, rows: List[EvaluationRow], config: RolloutProcessorConfig) ->
         # 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):
-            raise ValueError("kwargs['agent'] must be a valid Pydantic AI Agent instance")
-
-        agent: Agent = config.kwargs["agent"]
+        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"
+            )
 
-        model = OpenAIModel(
-            config.completion_params["model"],
-            provider=config.completion_params["provider"],
-        )
+        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 = {}
+            for model_name, model_config in config.completion_params["model"].items():
+                kwargs[model_name] = OpenAIModel(
+                    model_config["model"],
+                    provider=model_config["provider"],
+                )
+            agent = setup_agent(**kwargs)
+            model = None
+        else:
+            agent = config.kwargs["agent"]
+            model = OpenAIModel(
+                config.completion_params["model"],
+                provider=config.completion_params["provider"],
+            )
 
         async def process_row(row: EvaluationRow) -> EvaluationRow:
             """Process a single row with agent rollout."""
             model_messages = [self.convert_ep_message_to_pyd_message(m, row) for m in row.messages]
-            response = await agent.run(message_history=model_messages, model=model)
+            response = await agent.run(
+                message_history=model_messages, model=model, usage_limits=config.kwargs.get("usage_limits")
+            )
             row.messages = await self.convert_pyd_message_to_ep_message(response.all_messages())
             return row
 

tests/pytest/test_pydantic_multi_agent.py

@@ -0,0 +1,71 @@
+"""
+Copied and modified for eval-protocol from https://ai.pydantic.dev/multi-agent-applications/#agent-delegation
+
+To test your Pydantic AI multi-agent application, you can pass a function that
+sets up the agents and their tools. The function should accept parameters that
+map a model to each agent. In completion_params, you can provide mappings of
+model to agent based on key.
+"""
+
+import pytest
+
+from eval_protocol.models import EvaluationRow, Message
+from eval_protocol.pytest import evaluation_test
+from pydantic_ai import Agent
+
+from eval_protocol.pytest.default_pydantic_ai_rollout_processor import PydanticAgentRolloutProcessor
+from pydantic_ai import RunContext
+from pydantic_ai.models import Model
+from pydantic_ai.usage import UsageLimits
+
+
+def setup_agent(joke_generation_model: Model, joke_selection_model: Model) -> Agent:
+    """
+    This is an extra step that most applications will probably need to do to
+    parameterize the model that their agents use. But we believe that this is a
+    necessary step for multi-agent applications if developers want to solve the
+    model selection problem.
+    """
+    joke_selection_agent = Agent(
+        model=joke_selection_model,
+        system_prompt=(
+            "Use the `joke_factory` to generate some jokes, then choose the best. You must return just a single joke."
+        ),
+    )
+    joke_generation_agent = Agent(joke_generation_model, output_type=list[str])
+
+    @joke_selection_agent.tool
+    async def joke_factory(ctx: RunContext[None], count: int) -> list[str]:
+        r = await joke_generation_agent.run(
+            f"Please generate {count} jokes.",
+            usage=ctx.usage,
+        )
+        return r.output
+
+    return joke_selection_agent
+
+
+@pytest.mark.asyncio
+@evaluation_test(
+    input_messages=[Message(role="user", content="Hello, how are you?")],
+    completion_params=[
+        {
+            "model": {
+                "joke_generation_model": {"model": "accounts/fireworks/models/gpt-oss-120b", "provider": "fireworks"},
+                "joke_selection_model": {"model": "accounts/fireworks/models/gpt-oss-20b", "provider": "fireworks"},
+            }
+        },
+    ],
+    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=30),
+    },
+    mode="pointwise",
+)
+async def test_pydantic_multi_agent(row: EvaluationRow) -> EvaluationRow:
+    """
+    Super simple hello world test for Pydantic AI.
+    """
+    return row