refactor test_pydantic_multi_agent to work with factory setup

Author: Dylan Huang

tests/pytest/test_pydantic_multi_agent.py

@@ -1,12 +1,13 @@
 """
 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.
+To test your Pydantic AI multi-agent application, you can pass a factory that
+sets up the agenet based on the completion_params. The function should accept a
+RolloutProcessorConfig. In completion_params, you can provide mappings of model
+to agent based on key.
 """
 
+from pydantic_ai.models.openai import OpenAIModel
 import pytest
 
 from eval_protocol.models import EvaluationRow, Message
@@ -18,6 +19,8 @@
 from pydantic_ai.models import Model
 from pydantic_ai.usage import UsageLimits
 
+from eval_protocol.pytest.types import RolloutProcessorConfig
+
 
 def setup_agent(joke_generation_model: Model, joke_selection_model: Model) -> Agent:
     """
@@ -45,22 +48,31 @@ async def joke_factory(ctx: RunContext[None], count: int) -> list[str]:  # pyrig
     return joke_selection_agent
 
 
+def agent_factory(config: RolloutProcessorConfig) -> Agent:
+    joke_generation_model = OpenAIModel(
+        config.completion_params["model"]["joke_generation_model"], provider="fireworks"
+    )
+    joke_selection_model = OpenAIModel(config.completion_params["model"]["joke_selection_model"], provider="fireworks")
+    return setup_agent(
+        joke_generation_model,
+        joke_selection_model,
+    )
+
+
 @pytest.mark.asyncio
 @evaluation_test(
     input_messages=[[[Message(role="user", content="Tell me a joke.")]]],
     completion_params=[
         # multi-agent
         {
-            "joke_generation_model": {
-                "model": "fireworks_ai/accounts/fireworks/models/kimi-k2-instruct",
-            },
-            "joke_selection_model": {
-                "model": "fireworks_ai/accounts/fireworks/models/deepseek-v3p1",
-            },
+            "model": {
+                "joke_generation_model": "accounts/fireworks/models/kimi-k2-instruct",
+                "joke_selection_model": "accounts/fireworks/models/deepseek-v3p1",
+            }
         },
     ],
     rollout_processor=PydanticAgentRolloutProcessor(
-        setup_agent, UsageLimits(request_limit=5, total_tokens_limit=1000)
+        agent_factory, UsageLimits(request_limit=5, total_tokens_limit=1000)
     ),
     mode="pointwise",
 )