Test fixes and test updates, some failing

tests/integration/test_agent_manager_integration.py

@@ -28,11 +28,11 @@ class TestAgentManagerIntegration:
         
         # Mock prompt client
         prompt_client = AsyncMock()
-        prompt_client.agent_react.return_value = {
+        prompt_client.agent_react.return_value = """Thought: I need to search for information about machine learning
-            "thought": "I need to search for information about machine learning",
+Action: knowledge_query
-            "action": "knowledge_query",
+Args: {
-            "arguments": {"question": "What is machine learning?"}
+    "question": "What is machine learning?"
-        }
+}"""
         
         # Mock graph RAG client
         graph_rag_client = AsyncMock()
@@ -147,10 +147,8 @@ class TestAgentManagerIntegration:
     async def test_agent_manager_final_answer(self, agent_manager, mock_flow_context):
         """Test agent manager returning final answer"""
         # Arrange
-        mock_flow_context("prompt-request").agent_react.return_value = {
+        mock_flow_context("prompt-request").agent_react.return_value = """Thought: I have enough information to answer the question
-            "thought": "I have enough information to answer the question",
+Final Answer: Machine learning is a field of AI that enables computers to learn from data."""
-            "final-answer": "Machine learning is a field of AI that enables computers to learn from data."
-        }
         
         question = "What is machine learning?"
         history = []
@@ -195,10 +193,8 @@ class TestAgentManagerIntegration:
     async def test_agent_manager_react_with_final_answer(self, agent_manager, mock_flow_context):
         """Test ReAct cycle ending with final answer"""
         # Arrange
-        mock_flow_context("prompt-request").agent_react.return_value = {
+        mock_flow_context("prompt-request").agent_react.return_value = """Thought: I can provide a direct answer
-            "thought": "I can provide a direct answer",
+Final Answer: Machine learning is a branch of artificial intelligence."""
-            "final-answer": "Machine learning is a branch of artificial intelligence."
-        }
         
         question = "What is machine learning?"
         history = []
@@ -258,11 +254,11 @@ class TestAgentManagerIntegration:
 
         for tool_name, expected_service in tool_scenarios:
             # Arrange
-            mock_flow_context("prompt-request").agent_react.return_value = {
+            mock_flow_context("prompt-request").agent_react.return_value = f"""Thought: I need to use {tool_name}
-                "thought": f"I need to use {tool_name}",
+Action: {tool_name}
-                "action": tool_name,
+Args: {{
-                "arguments": {"question": "test question"}
+    "question": "test question"
-            }
+}}"""
             
             think_callback = AsyncMock()
             observe_callback = AsyncMock()
@@ -288,11 +284,11 @@ class TestAgentManagerIntegration:
     async def test_agent_manager_unknown_tool_error(self, agent_manager, mock_flow_context):
         """Test agent manager error handling for unknown tool"""
         # Arrange
-        mock_flow_context("prompt-request").agent_react.return_value = {
+        mock_flow_context("prompt-request").agent_react.return_value = """Thought: I need to use an unknown tool
-            "thought": "I need to use an unknown tool",
+Action: unknown_tool
-            "action": "unknown_tool",
+Args: {
-            "arguments": {"param": "value"}
+    "param": "value"
-        }
+}"""
         
         think_callback = AsyncMock()
         observe_callback = AsyncMock()
@@ -325,11 +321,11 @@ class TestAgentManagerIntegration:
         question = "Find information about AI and summarize it"
         
         # Mock multi-step reasoning
-        mock_flow_context("prompt-request").agent_react.return_value = {
+        mock_flow_context("prompt-request").agent_react.return_value = """Thought: I need to search for AI information first
-            "thought": "I need to search for AI information first",
+Action: knowledge_query
-            "action": "knowledge_query",
+Args: {
-            "arguments": {"question": "What is artificial intelligence?"}
+    "question": "What is artificial intelligence?"
-        }
+}"""
 
         # Act
         action = await agent_manager.reason(question, [], mock_flow_context)
@@ -373,11 +369,12 @@ class TestAgentManagerIntegration:
 
         for test_case in test_cases:
             # Arrange
-            mock_flow_context("prompt-request").agent_react.return_value = {
+            # Format arguments as JSON
-                "thought": f"Using {test_case['action']}",
+            import json
-                "action": test_case['action'],
+            args_json = json.dumps(test_case['arguments'], indent=4)
-                "arguments": test_case['arguments']
+            mock_flow_context("prompt-request").agent_react.return_value = f"""Thought: Using {test_case['action']}
-            }
+Action: {test_case['action']}
+Args: {args_json}"""
             
             think_callback = AsyncMock()
             observe_callback = AsyncMock()
@@ -465,6 +462,167 @@ class TestAgentManagerIntegration:
             # Reset mocks
             mock_flow_context("graph-rag-request").reset_mock()
 
+    @pytest.mark.asyncio
+    async def test_agent_manager_malformed_response_handling(self, agent_manager, mock_flow_context):
+        """Test agent manager handling of malformed text responses"""
+        # Test various malformed responses
+        malformed_responses = [
+            # Missing required fields
+            "Thought: I need to do something",
+            # Invalid JSON in Args
+            """Thought: I need to search
+Action: knowledge_query
+Args: {invalid json}""",
+            # Empty response
+            "",
+            # Only whitespace
+            "   \n\t  ",
+            # Missing Args for action
+            """Thought: I need to search
+Action: knowledge_query""",
+            # Incomplete JSON
+            """Thought: I need to search
+Action: knowledge_query
+Args: {
+    "question": "test"
+""",
+        ]
+        
+        for response in malformed_responses:
+            mock_flow_context("prompt-request").agent_react.return_value = response
+            
+            # Act & Assert
+            with pytest.raises(RuntimeError) as exc_info:
+                await agent_manager.reason("test question", [], mock_flow_context)
+            
+            assert "Failed to parse agent response" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    async def test_agent_manager_text_parsing_edge_cases(self, agent_manager, mock_flow_context):
+        """Test edge cases in text parsing"""
+        # Test response with markdown code blocks
+        mock_flow_context("prompt-request").agent_react.return_value = """```
+Thought: I need to search for information
+Action: knowledge_query
+Args: {
+    "question": "What is AI?"
+}
+```"""
+        
+        action = await agent_manager.reason("test", [], mock_flow_context)
+        assert isinstance(action, Action)
+        assert action.thought == "I need to search for information"
+        assert action.name == "knowledge_query"
+        
+        # Test response with extra whitespace
+        mock_flow_context("prompt-request").agent_react.return_value = """
+
+Thought:    I need to think about this    
+Action:   knowledge_query   
+Args:    {
+    "question": "test"
+}
+
+"""
+        
+        action = await agent_manager.reason("test", [], mock_flow_context)
+        assert isinstance(action, Action)
+        assert action.thought == "I need to think about this"
+        assert action.name == "knowledge_query"
+
+    @pytest.mark.asyncio
+    async def test_agent_manager_multiline_content(self, agent_manager, mock_flow_context):
+        """Test handling of multi-line thoughts and final answers"""
+        # Multi-line thought
+        mock_flow_context("prompt-request").agent_react.return_value = """Thought: I need to consider multiple factors:
+1. The user's question is complex
+2. I should search for comprehensive information
+3. This requires using the knowledge query tool
+Action: knowledge_query
+Args: {
+    "question": "complex query"
+}"""
+        
+        action = await agent_manager.reason("test", [], mock_flow_context)
+        assert isinstance(action, Action)
+        assert "multiple factors" in action.thought
+        assert "knowledge query tool" in action.thought
+        
+        # Multi-line final answer
+        mock_flow_context("prompt-request").agent_react.return_value = """Thought: I have gathered enough information
+Final Answer: Here is a comprehensive answer:
+1. First point about the topic
+2. Second point with details
+3. Final conclusion
+
+This covers all aspects of the question."""
+        
+        action = await agent_manager.reason("test", [], mock_flow_context)
+        assert isinstance(action, Final)
+        assert "First point" in action.final
+        assert "Final conclusion" in action.final
+        assert "all aspects" in action.final
+
+    @pytest.mark.asyncio
+    async def test_agent_manager_json_args_special_characters(self, agent_manager, mock_flow_context):
+        """Test JSON arguments with special characters and edge cases"""
+        # Test with special characters in JSON
+        mock_flow_context("prompt-request").agent_react.return_value = """Thought: Processing special characters
+Action: knowledge_query
+Args: {
+    "question": "What about \"quotes\" and 'apostrophes'?",
+    "context": "Line 1\\nLine 2\\tTabbed",
+    "special": "Symbols: @#$%^&*()_+-=[]{}|;':,.<>?"
+}"""
+        
+        action = await agent_manager.reason("test", [], mock_flow_context)
+        assert isinstance(action, Action)
+        assert action.arguments["question"] == 'What about "quotes" and \'apostrophes\'?'
+        assert "\\n" in action.arguments["context"]
+        assert "@#$%^&*" in action.arguments["special"]
+        
+        # Test with nested JSON
+        mock_flow_context("prompt-request").agent_react.return_value = """Thought: Complex arguments
+Action: web_search
+Args: {
+    "query": "test",
+    "options": {
+        "limit": 10,
+        "filters": ["recent", "relevant"],
+        "metadata": {
+            "source": "user",
+            "timestamp": "2024-01-01"
+        }
+    }
+}"""
+        
+        action = await agent_manager.reason("test", [], mock_flow_context)
+        assert isinstance(action, Action)
+        assert action.arguments["options"]["limit"] == 10
+        assert "recent" in action.arguments["options"]["filters"]
+        assert action.arguments["options"]["metadata"]["source"] == "user"
+
+    @pytest.mark.asyncio
+    async def test_agent_manager_final_answer_json_format(self, agent_manager, mock_flow_context):
+        """Test final answers that contain JSON-like content"""
+        # Final answer with JSON content
+        mock_flow_context("prompt-request").agent_react.return_value = """Thought: I can provide the data in JSON format
+Final Answer: {
+    "result": "success",
+    "data": {
+        "name": "Machine Learning",
+        "type": "AI Technology",
+        "applications": ["NLP", "Computer Vision", "Robotics"]
+    },
+    "confidence": 0.95
+}"""
+        
+        action = await agent_manager.reason("test", [], mock_flow_context)
+        assert isinstance(action, Final)
+        # The final answer should preserve the JSON structure as a string
+        assert '"result": "success"' in action.final
+        assert '"applications":' in action.final
+
     @pytest.mark.asyncio
     @pytest.mark.slow
     async def test_agent_manager_performance_with_large_history(self, agent_manager, mock_flow_context):