fix test_pydantic_agent.py

Author: Dylan Huang

eval_protocol/pytest/default_pydantic_ai_rollout_processor.py

@@ -32,7 +32,9 @@ class PydanticAgentRolloutProcessor(RolloutProcessor):
     EvaluationRow.messages to and from Pydantic AI ModelMessage format."""
 
     def __init__(
-        self, agent_factory: Callable[[RolloutProcessorConfig], Agent], usage_limits: UsageLimits | None = None
+        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"))

tests/pytest/test_pydantic_agent.py

@@ -1,21 +1,26 @@
+from pydantic_ai.agent import Agent
+from pydantic_ai.models.openai import OpenAIModel
 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 eval_protocol.pytest.types import RolloutProcessorConfig
 
-agent = Agent()
+
+def agent_factory(config: RolloutProcessorConfig) -> Agent:
+    model = OpenAIModel(config.completion_params["model"], provider="fireworks")
+    return Agent(model=model)
 
 
 @pytest.mark.asyncio
 @evaluation_test(
     input_messages=[[[Message(role="user", content="Hello, how are you?")]]],
     completion_params=[
-        {"model": "fireworks_ai/accounts/fireworks/models/gpt-oss-120b"},
+        {"model": "accounts/fireworks/models/gpt-oss-120b"},
     ],
-    rollout_processor=PydanticAgentRolloutProcessor(agent),
+    rollout_processor=PydanticAgentRolloutProcessor(agent_factory),
     mode="pointwise",
 )
 async def test_pydantic_agent(row: EvaluationRow) -> EvaluationRow: