Fix structured output tests

ag2ai · Dec 5, 2024 · c40da67 · c40da67
1 parent 5e15988
commit c40da67
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/test/agentchat/test_structured_output.py b/test/agentchat/test_structured_output.py
@@ -37,6 +37,9 @@ class ResponseModel(BaseModel):
         reasoning: str
         difficulty: float
 
+    for config in config_list:
+        config["response_format"] = MathReasoning
+
     llm_config = {"config_list": config_list, "cache_seed": 43}
 
     user_proxy = autogen.UserProxyAgent(
@@ -84,13 +87,12 @@ def format(self) -> str:
 @pytest.fixture
 def mock_assistant():
     """Set up a mocked AssistantAgent with a predefined response format."""
-    config_list = [{"model": "gpt-4o", "api_key": MOCK_OPEN_AI_API_KEY}]
+    config_list = [{"model": "gpt-4o", "api_key": MOCK_OPEN_AI_API_KEY, "response_format": MathReasoning}]
     llm_config = {"config_list": config_list, "cache_seed": 43}
 
     assistant = autogen.AssistantAgent(
         name="Assistant",
         llm_config=llm_config,
-        response_format=MathReasoning,
     )
 
     oai_client_mock = MagicMock()