Fix tests

Author: StreetLamb

src/mcp_agent/workflows/llm/augmented_llm_azure.py

@@ -365,7 +365,7 @@ async def generate_structured(
 
     @classmethod
     def convert_message_to_message_param(
-        cls, message: ResponseMessage, **kwargs
+        cls, message: ResponseMessage
     ) -> AssistantMessage:
         """Convert a response object to an input parameter object to allow LLM calls to be chained."""
         assistant_message = AssistantMessage(

tests/workflows/llm/test_augmented_llm_anthropic.py

@@ -210,31 +210,56 @@ async def test_generate_str(self, mock_llm, default_usage):
     @pytest.mark.asyncio
     async def test_generate_structured(self, mock_llm, default_usage):
         """
-        Tests structured output generation using Instructor.
+        Tests structured output generation using native Anthropic API.
         """
+        from unittest.mock import patch
+        import json
 
         # Define a simple response model
         class TestResponseModel(BaseModel):
             name: str
             value: int
 
-        # Mock the generate_str method to return a string that will be parsed by the instructor mock
-        mock_llm.generate_str = AsyncMock(return_value="name: Test, value: 42")
-
-        # Patch executor.execute to return the expected TestResponseModel instance
-        mock_llm.executor.execute = AsyncMock(
-            return_value=TestResponseModel(name="Test", value=42)
+        # Create a mock Message with tool_use block containing the structured data
+        tool_use_block = ToolUseBlock(
+            type="tool_use",
+            id="tool_123",
+            name="return_structured_output",
+            input={"name": "Test", "value": 42},
         )
 
-        # Call the method
-        result = await AnthropicAugmentedLLM.generate_structured(
-            mock_llm, "Test query", TestResponseModel
+        mock_message = Message(
+            type="message",
+            id="msg_123",
+            role="assistant",
+            content=[tool_use_block],
+            model="claude-3-7-sonnet-latest",
+            stop_reason="tool_use",
+            usage=default_usage,
         )
 
-        # Assertions
-        assert isinstance(result, TestResponseModel)
-        assert result.name == "Test"
-        assert result.value == 42
+        # Mock the AsyncAnthropic client and streaming
+        with patch(
+            "mcp_agent.workflows.llm.augmented_llm_anthropic.AsyncAnthropic"
+        ) as MockAsyncAnthropic:
+            mock_client = MockAsyncAnthropic.return_value
+            mock_stream = AsyncMock()
+            mock_stream.get_final_message = AsyncMock(return_value=mock_message)
+            mock_stream.__aenter__ = AsyncMock(return_value=mock_stream)
+            mock_stream.__aexit__ = AsyncMock(return_value=None)
+            mock_client.messages.stream = MagicMock(return_value=mock_stream)
+            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+            mock_client.__aexit__ = AsyncMock(return_value=None)
+
+            # Call the method
+            result = await AnthropicAugmentedLLM.generate_structured(
+                mock_llm, "Test query", TestResponseModel
+            )
+
+            # Assertions
+            assert isinstance(result, TestResponseModel)
+            assert result.name == "Test"
+            assert result.value == 42
 
     # Test 4: With History
     @pytest.mark.asyncio
@@ -779,6 +804,8 @@ async def test_generate_structured_with_mixed_message_types(self, mock_llm):
         """
         Tests generate_structured() method with mixed message types.
         """
+        from unittest.mock import patch
+        import json
 
         # Define a simple response model
         class TestResponseModel(BaseModel):
@@ -795,19 +822,51 @@ class TestResponseModel(BaseModel):
             ),
         ]
 
-        mock_llm.generate_str = AsyncMock(return_value="name: MixedTypes, value: 123")
-        # Patch executor.execute to return the expected TestResponseModel instance
-        mock_llm.executor.execute = AsyncMock(
-            return_value=TestResponseModel(name="MixedTypes", value=123)
+        # Create a mock Message with tool_use block containing the structured data
+        tool_use_block = ToolUseBlock(
+            type="tool_use",
+            id="tool_456",
+            name="return_structured_output",
+            input={"name": "MixedTypes", "value": 123},
         )
 
-        # Call generate_structured with mixed message types
-        result = await mock_llm.generate_structured(messages, TestResponseModel)
+        mock_message = Message(
+            type="message",
+            id="msg_456",
+            role="assistant",
+            content=[tool_use_block],
+            model="claude-3-7-sonnet-latest",
+            stop_reason="tool_use",
+            usage=Usage(
+                cache_creation_input_tokens=0,
+                cache_read_input_tokens=0,
+                input_tokens=100,
+                output_tokens=50,
+                server_tool_use=None,
+                service_tier=None,
+            ),
+        )
 
-        # Assertions
-        assert isinstance(result, TestResponseModel)
-        assert result.name == "MixedTypes"
-        assert result.value == 123
+        # Mock the AsyncAnthropic client and streaming
+        with patch(
+            "mcp_agent.workflows.llm.augmented_llm_anthropic.AsyncAnthropic"
+        ) as MockAsyncAnthropic:
+            mock_client = MockAsyncAnthropic.return_value
+            mock_stream = AsyncMock()
+            mock_stream.get_final_message = AsyncMock(return_value=mock_message)
+            mock_stream.__aenter__ = AsyncMock(return_value=mock_stream)
+            mock_stream.__aexit__ = AsyncMock(return_value=None)
+            mock_client.messages.stream = MagicMock(return_value=mock_stream)
+            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+            mock_client.__aexit__ = AsyncMock(return_value=None)
+
+            # Call generate_structured with mixed message types
+            result = await mock_llm.generate_structured(messages, TestResponseModel)
+
+            # Assertions
+            assert isinstance(result, TestResponseModel)
+            assert result.name == "MixedTypes"
+            assert result.value == 123
 
     # Test 25: System Prompt Not None in API Call
     @pytest.mark.asyncio

tests/workflows/llm/test_augmented_llm_google.py

@@ -183,29 +183,23 @@ class TestResponseModel(BaseModel):
             name: str
             value: int
 
-        # Mock the generate_str method
-        mock_llm.generate_str = AsyncMock(return_value="name: Test, value: 42")
-
-        # Mock instructor from_genai
-        with patch("instructor.from_genai") as mock_instructor:
-            mock_client = MagicMock()
-            mock_client.chat.completions.create.return_value = TestResponseModel(
-                name="Test", value=42
-            )
-            mock_instructor.return_value = mock_client
+        # Create a proper GenerateContentResponse with JSON content
+        import json
+        json_content = json.dumps({"name": "Test", "value": 42})
+        response = self.create_text_response(json_content)
 
-            # Patch executor.execute to be an async mock returning the expected value
-            mock_llm.executor.execute = AsyncMock(
-                return_value=TestResponseModel(name="Test", value=42)
-            )
+        # Patch executor.execute to return the GenerateContentResponse with JSON
+        mock_llm.executor.execute = AsyncMock(
+            return_value=response
+        )
 
-            # Call the method
-            result = await mock_llm.generate_structured("Test query", TestResponseModel)
+        # Call the method
+        result = await mock_llm.generate_structured("Test query", TestResponseModel)
 
-            # Assertions
-            assert isinstance(result, TestResponseModel)
-            assert result.name == "Test"
-            assert result.value == 42
+        # Assertions
+        assert isinstance(result, TestResponseModel)
+        assert result.name == "Test"
+        assert result.value == 42
 
     # Test 4: With History
     @pytest.mark.asyncio
@@ -773,26 +767,20 @@ class TestResponseModel(BaseModel):
             ),
         ]
 
-        # Mock the generate_str method
-        mock_llm.generate_str = AsyncMock(return_value="name: MixedTypes, value: 123")
+        # Create a proper GenerateContentResponse with JSON content
+        import json
+        json_content = json.dumps({"name": "MixedTypes", "value": 123})
+        response = self.create_text_response(json_content)
 
-        # Patch instructor.from_genai to return the expected model
-        with patch("instructor.from_genai") as mock_instructor:
-            mock_client = MagicMock()
-            mock_client.chat.completions.create.return_value = TestResponseModel(
-                name="MixedTypes", value=123
-            )
-            mock_instructor.return_value = mock_client
-
-            # Patch executor.execute to be an async mock returning the expected value
-            mock_llm.executor.execute = AsyncMock(
-                return_value=TestResponseModel(name="MixedTypes", value=123)
-            )
+        # Patch executor.execute to return the GenerateContentResponse with JSON
+        mock_llm.executor.execute = AsyncMock(
+            return_value=response
+        )
 
-            result = await mock_llm.generate_structured(messages, TestResponseModel)
-            assert isinstance(result, TestResponseModel)
-            assert result.name == "MixedTypes"
-            assert result.value == 123
+        result = await mock_llm.generate_structured(messages, TestResponseModel)
+        assert isinstance(result, TestResponseModel)
+        assert result.name == "MixedTypes"
+        assert result.value == 123
 
     @pytest.mark.asyncio
     async def test_parallel_tool_calls(self, mock_llm: GoogleAugmentedLLM):

tests/workflows/llm/test_augmented_llm_openai.py

@@ -174,38 +174,33 @@ async def test_generate_str(self, mock_llm, default_usage):
     @pytest.mark.asyncio
     async def test_generate_structured(self, mock_llm, default_usage):
         """
-        Tests structured output generation using Instructor.
+        Tests structured output generation using native OpenAI API.
         """
+        import json
 
         # Define a simple response model
         class TestResponseModel(BaseModel):
             name: str
             value: int
 
-        # Set up mocks for the two-stage process
-        # First for the text generation
-        mock_llm.generate_str = AsyncMock(return_value="name: Test, value: 42")
-
-        # Then for Instructor's structured data extraction
-        with patch("instructor.from_openai") as mock_instructor:
-            mock_client = MagicMock()
-            mock_client.chat.completions.create.return_value = TestResponseModel(
-                name="Test", value=42
-            )
-            mock_instructor.return_value = mock_client
+        # Create a proper ChatCompletion response with JSON content
+        json_content = json.dumps({"name": "Test", "value": 42})
+        completion_response = self.create_text_response(
+            json_content, usage=default_usage
+        )
 
-            # Patch executor.execute to be an async mock returning the expected value
-            mock_llm.executor.execute = AsyncMock(
-                return_value=TestResponseModel(name="Test", value=42)
-            )
+        # Patch executor.execute to return the ChatCompletion with JSON
+        mock_llm.executor.execute = AsyncMock(
+            return_value=completion_response
+        )
 
-            # Call the method
-            result = await mock_llm.generate_structured("Test query", TestResponseModel)
+        # Call the method
+        result = await mock_llm.generate_structured("Test query", TestResponseModel)
 
-            # Assertions
-            assert isinstance(result, TestResponseModel)
-            assert result.name == "Test"
-            assert result.value == 42
+        # Assertions
+        assert isinstance(result, TestResponseModel)
+        assert result.name == "Test"
+        assert result.value == 42
 
     # Test 4: With History
     @pytest.mark.asyncio
@@ -612,6 +607,7 @@ async def test_generate_structured_with_mixed_message_types(self, mock_llm):
         """
         Tests generate_structured() method with mixed message types.
         """
+        import json
 
         # Define a simple response model
         class TestResponseModel(BaseModel):
@@ -628,10 +624,19 @@ class TestResponseModel(BaseModel):
             ),
         ]
 
-        mock_llm.generate_str = AsyncMock(return_value="name: MixedTypes, value: 123")
-        # Patch executor.execute to return the expected TestResponseModel instance
+        # Create a proper ChatCompletion response with JSON content
+        json_content = json.dumps({"name": "MixedTypes", "value": 123})
+        completion_response = self.create_text_response(
+            json_content, usage=CompletionUsage(
+                completion_tokens=100,
+                prompt_tokens=150,
+                total_tokens=250
+            )
+        )
+
+        # Patch executor.execute to return the ChatCompletion with JSON
         mock_llm.executor.execute = AsyncMock(
-            return_value=TestResponseModel(name="MixedTypes", value=123)
+            return_value=completion_response
         )
 
         # Call generate_structured with mixed message types