Refactor mock completion function in mini agents tests for improved clarity and functionality

Author: MervinPraison

src/praisonai/tests/unit/agent/test_mini_agents_sequential.py

@@ -13,29 +13,88 @@
 from unittest.mock import patch
 
 def mock_completion(*args, **kwargs):
-    """Mock completion function to avoid calling actual OpenAI API"""
-    class MockResponse:
-        def __init__(self, content):
-            self.content = content
-            self.choices = [type('obj', (object,), {'message': type('obj', (object,), {'content': content})})]
+    """Mock litellm.completion function to avoid calling actual API"""
+    stream = kwargs.get('stream', False)
+    
+    # Determine response content based on messages
+    messages = kwargs.get('messages', [])
+    response_content = "mock response"
+    
+    for message in messages:
+        content = message.get('content', '') if isinstance(message, dict) else str(message)
+        if 'Generate the number 42' in content:
+            response_content = "42"
+            break
+        elif 'multiply it by 2' in content:
+            response_content = "84"
+            break
+    
+    if stream:
+        # Return streaming response (iterator)
+        class MockDelta:
+            def __init__(self, content):
+                self.content = content
+        
+        class MockStreamChoice:
+            def __init__(self, content):
+                self.delta = MockDelta(content)
+        
+        class MockStreamChunk:
+            def __init__(self, content):
+                self.choices = [MockStreamChoice(content)]
+        
+        # Return a list that can be iterated (simulating streaming chunks)
+        return [MockStreamChunk(response_content)]
+    else:
+        # Return non-streaming response
+        class MockMessage:
+            def __init__(self, content):
+                self.content = content
+            
+            def get(self, key, default=None):
+                if key == "tool_calls":
+                    return None  # No tool calls in our simple test
+                return getattr(self, key, default)
+            
+            def __getitem__(self, key):
+                if hasattr(self, key):
+                    return getattr(self, key)
+                raise KeyError(key)
+        
+        class MockChoice:
+            def __init__(self, content):
+                self.message = MockMessage(content)
+            
+            def __getitem__(self, key):
+                if key == "message":
+                    return self.message
+                if hasattr(self, key):
+                    return getattr(self, key)
+                raise KeyError(key)
+        
+        class MockResponse:
+            def __init__(self, content):
+                self.choices = [MockChoice(content)]
 
-    if 'messages' in kwargs:
-        if any('Generate the number 42' in str(m.get('content', '')) for m in kwargs.get('messages', [])):
-            return MockResponse("42")
-        elif any('multiply it by 2' in str(m.get('content', '')) for m in kwargs.get('messages', [])):
-            return MockResponse("84")
-    return MockResponse("mock response")
+            def __getitem__(self, key):
+                # Support dictionary-style access
+                if key == "choices":
+                    return self.choices
+                if hasattr(self, key):
+                    return getattr(self, key)
+                raise KeyError(key)
+        
+        return MockResponse(response_content)
 
-@patch('praisonai.inc.models.PraisonAIModel.chat', side_effect=mock_completion)
-@patch('praisonai.inc.models.PraisonAIModel.stream_chat', side_effect=mock_completion)
-def test_mini_agents_sequential_data_passing(mock_stream, mock_chat):
+@patch('litellm.completion', side_effect=mock_completion)
+def test_mini_agents_sequential_data_passing(mock_litellm):
     """Test that output from previous task is passed to next task in Mini Agents"""
 
     print("Testing Mini Agents Sequential Data Passing...")
 
     # Create two agents for sequential processing
-    agent1 = Agent(instructions="Generate the number 42 as your output. Only return the number 42, nothing else.", model_name="gpt-3.5-turbo")
-    agent2 = Agent(instructions="Take the input number and multiply it by 2. Only return the result number, nothing else.", model_name="gpt-3.5-turbo")
+    agent1 = Agent(instructions="Generate the number 42 as your output. Only return the number 42, nothing else.", llm={'model': 'gpt-3.5-turbo'})
+    agent2 = Agent(instructions="Take the input number and multiply it by 2. Only return the result number, nothing else.", llm={'model': 'gpt-3.5-turbo'})
 
     # Create agents with sequential processing (Mini Agents pattern)
     agents = Agents(agents=[agent1, agent2], verbose=True)