mock_server.py

#!/usr/bin/env python3
"""
Simple mock HTTP server for testing ChatBotChatBot response collection.
"""

import json
import asyncio
from datetime import datetime
from typing import Dict, Any
from http.server import HTTPServer, BaseHTTPRequestHandler
import urllib.parse


class MockTargetAgent:
    """Mock target agent that provides predictable responses."""
    
    def __init__(self):
        self.session_memory = []
        self.session_id = None
    
    def handle_request(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        """Handle a mock request and return a response."""
        message = payload.get('message', '')
        session_id = payload.get('session_id')
        
        if session_id:
            self.session_id = session_id
        
        # Store message in session memory
        self.session_memory.append(message)
        
        # Generate mock response based on input
        response = self._generate_mock_response(message)
        
        return {
            'response': response,
            'session_id': self.session_id or f"mock-session-{len(self.session_memory)}",
            'timestamp': datetime.now().isoformat(),
            'model': 'mock-agent-v1'
        }
    
    def _generate_mock_response(self, message: str) -> str:
        """Generate a mock response based on the input message."""
        message_lower = message.lower()
        
        # Math responses
        if '2 + 2' in message_lower or '2+2' in message_lower:
            return "The answer is 4."
        elif '15 * 3' in message_lower or '15*3' in message_lower:
            return "15 multiplied by 3 equals 45."
        elif '10 divided by 2' in message_lower or '10/2' in message_lower:
            return "10 divided by 2 is 5."
        elif 'square root of 144' in message_lower:
            return "The square root of 144 is 12."
        elif '7 * 8' in message_lower or '7*8' in message_lower:
            return "7 times 8 equals 56."
        
        # Customer service responses
        elif 'return' in message_lower:
            return "I'd be happy to help with your return request. Can you provide your order number and the reason for the return?"
        elif 'refund policy' in message_lower:
            return "Our return policy allows returns within 30 days of purchase with original receipt. Refunds are processed within 3-5 business days."
        elif 'order' in message_lower and ('arrived' in message_lower or 'delivery' in message_lower):
            return "I can help you track your order. Please provide your order number and I'll check the delivery status for you."
        elif 'cancel' in message_lower and 'order' in message_lower:
            return "I can help you cancel your order if it hasn't shipped yet. Please provide your order number."
        elif 'damaged product' in message_lower:
            return "I'm sorry to hear about the damaged product. I can arrange a replacement for you right away."
        
        # General knowledge responses
        elif 'capital of france' in message_lower:
            return "The capital of France is Paris."
        elif 'leap year' in message_lower and 'days' in message_lower:
            return "A leap year has 366 days."
        elif 'red and blue' in message_lower and ('mix' in message_lower or 'color' in message_lower):
            return "When you mix red and blue, you get purple."
        elif 'largest planet' in message_lower:
            return "Jupiter is the largest planet in our solar system."
        elif 'romeo and juliet' in message_lower:
            return "Romeo and Juliet was written by William Shakespeare."
        
        # Greetings
        elif 'hello' in message_lower or 'hi' in message_lower:
            return "Hello! I'm a helpful assistant. How can I help you today?"
        
        # Default response
        else:
            return f"Thank you for your message about '{message[:50]}...'. I understand you're asking about this topic. How can I best assist you?"


class MockServerHandler(BaseHTTPRequestHandler):
    """HTTP request handler for the mock server."""
    
    def __init__(self, *args, **kwargs):
        self.mock_agent = MockTargetAgent()
        super().__init__(*args, **kwargs)
    
    def do_POST(self):
        """Handle POST requests."""
        if self.path == '/chat':
            try:
                # Read request body
                content_length = int(self.headers['Content-Length'])
                post_data = self.rfile.read(content_length)
                
                # Parse JSON payload
                payload = json.loads(post_data.decode('utf-8'))
                
                # Get response from mock agent
                response_data = self.mock_agent.handle_request(payload)
                
                # Send response
                self.send_response(200)
                self.send_header('Content-Type', 'application/json')
                self.send_header('Access-Control-Allow-Origin', '*')
                self.end_headers()
                
                response_json = json.dumps(response_data, indent=2)
                self.wfile.write(response_json.encode('utf-8'))
                
                print(f"Handled request: {payload.get('message', '')[:50]}...")
                
            except Exception as e:
                # Send error response
                self.send_response(500)
                self.send_header('Content-Type', 'application/json')
                self.end_headers()
                
                error_response = {'error': str(e)}
                self.wfile.write(json.dumps(error_response).encode('utf-8'))
                print(f"Error handling request: {e}")
        else:
            self.send_response(404)
            self.end_headers()
    
    def do_GET(self):
        """Handle GET requests."""
        if self.path == '/health':
            self.send_response(200)
            self.send_header('Content-Type', 'application/json')
            self.end_headers()
            
            health_response = {'status': 'healthy', 'timestamp': datetime.now().isoformat()}
            self.wfile.write(json.dumps(health_response).encode('utf-8'))
        else:
            self.send_response(404)
            self.end_headers()
    
    def log_message(self, format, *args):
        """Override to reduce log noise."""
        pass


def run_mock_server(port=8001):
    """Run the mock server."""
    server_address = ('', port)
    
    # Create a new MockTargetAgent for each request handler
    def handler_with_agent(*args, **kwargs):
        handler = MockServerHandler(*args, **kwargs)
        handler.mock_agent = MockTargetAgent()
        return handler
    
    httpd = HTTPServer(server_address, MockServerHandler)
    print(f"Mock target agent server running on http://localhost:{port}")
    print("Endpoints:")
    print(f"  POST http://localhost:{port}/chat - Send messages")
    print(f"  GET  http://localhost:{port}/health - Health check")
    print("\nPress Ctrl+C to stop the server")
    
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print("\nShutting down mock server...")
        httpd.shutdown()


if __name__ == "__main__":
    import sys
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8001
    run_mock_server(port)