randebu/src/backend/app/services/ai_agent/mock_client.py

"""Mock client for testing the ConversationalAgent without hitting real APIs."""

from typing import List, Dict, Any, Optional


class MockMiniMaxClient:
    """Mock client that returns predefined responses for testing.
    
    Usage:
        mock = MockMiniMaxClient()
        mock.add_response({\"choices\": [...]})  # Add responses in order
        mock.add_response({\"choices\": [...]})  # Second call gets this
        
        agent = ConversationalAgent(client=mock)
        result = agent.chat(\"hello\")
    """
    
    def __init__(self, responses: List[Dict[str, Any]] = None):
        self.responses = responses or []
        self.call_count = 0
        self.calls: List[Dict[str, Any]] = []  # Record all calls for assertions
    
    def add_response(self, response: Dict[str, Any]):
        """Add a response to be returned on the next call."""
        self.responses.append(response)
    
    def chat(
        self,
        messages: List[Dict[str, str]],
        system_prompt: str,
        tools: Optional[List[Dict[str, Any]]] = None,
        temperature: float = 0.7,
        max_tokens: int = 2000,
        thinking_budget: int = 1500,
    ) -> Dict[str, Any]:
        """Return the next predefined response."""
        # Record the call for debugging/tests
        self.calls.append({
            "messages": messages,
            "system_prompt": system_prompt[:100] if system_prompt else None,
            "tool_calls_count": len(tools) if tools else 0,
        })
        
        if self.call_count < len(self.responses):
            response = self.responses[self.call_count]
            self.call_count += 1
            return response
        
        # Default response if no more predefined responses
        return {
            "choices": [{
                "message": {
                    "content": "Mock response - no more responses configured",
                    "role": "assistant",
                }
            }]
        }
    
    def reset(self):
        """Reset call count and calls list."""
        self.call_count = 0
        self.calls = []
    
    def verify_call(self, call_index: int, expected_messages: int = None, expected_tool_count: int = None):
        """Verify a specific call was made correctly."""
        if call_index >= len(self.calls):
            raise AssertionError(f"Call {call_index} was not made. Total calls: {len(self.calls)}")
        
        call = self.calls[call_index]
        if expected_messages is not None:
            actual = len(call["messages"])
            if actual != expected_messages:
                raise AssertionError(
                    f"Call {call_index}: expected {expected_messages} messages, got {actual}"
                )
        if expected_tool_count is not None:
            actual = call["tool_calls_count"]
            if actual != expected_tool_count:
                raise AssertionError(
                    f"Call {call_index}: expected {expected_tool_count} tools, got {actual}"
                )


class MockMiniMaxClientWithToolCall(MockMiniMaxClient):
    """Mock client that generates tool call responses based on message content."""
    
    def __init__(self, tool_handlers: Dict[str, Dict[str, Any]] = None):
        """tool_handlers: dict mapping tool names to their responses."""
        super().__init__()
        self.tool_handlers = tool_handlers or {}
    
    def chat(
        self,
        messages: List[Dict[str, str]],
        system_prompt: str,
        tools: Optional[List[Dict[str, Any]]] = None,
        temperature: float = 0.7,
        max_tokens: int = 2000,
        thinking_budget: int = 1500,
    ) -> Dict[str, Any]:
        """Check if last message contains a tool call request and return appropriate response."""
        self.calls.append({
            "messages": messages,
            "system_prompt": system_prompt[:100] if system_prompt else None,
            "tool_calls_count": len(tools) if tools else 0,
        })
        
        # Get the last message
        if not messages:
            return {"choices": [{"message": {"content": "No messages", "role": "assistant"}}]}
        
        last_msg = messages[-1]
        
        # If it's a tool result, look for the tool that was called
        if last_msg.get("role") == "user" and "[TOOL RESULT]" in last_msg.get("content", ""):
            # Extract tool name from the content
            content = last_msg["content"]
            # Format: [TOOL RESULT] tool_name: result
            for tool_name in self.tool_handlers:
                if f"[TOOL RESULT] {tool_name}:" in content:
                    return self.tool_handlers[tool_name]
        
        # Check if we should generate a tool call
        last_user_msg = ""
        for msg in reversed(messages):
            if msg.get("role") == "user" and "[TOOL RESULT]" not in msg.get("content", ""):
                last_user_msg = msg.get("content", "")
                break
        
        # Check each tool's trigger
        for tool_name, handler in self.tool_handlers.items():
            trigger = handler.get("trigger", "")
            if trigger and trigger.lower() in last_user_msg.lower():
                return {
                    "choices": [{
                        "message": {
                            "content": handler.get("content", ""),
                            "role": "assistant",
                            "tool_calls": [{
                                "id": f"call_{tool_name}",
                                "type": "function",
                                "function": {
                                    "name": tool_name,
                                    "arguments": handler.get("arguments", "{}")
                                }
                            }]
                        }
                    }]
                }
        
        # Default: return first available response or empty
        if self.responses:
            response = self.responses[0]
            self.call_count += 1
            return response
        
        return {
            "choices": [{
                "message": {
                    "content": "How can I help you with your trading bot?",
                    "role": "assistant",
                }
            }]
        }