examples/websocket_example.py

"""
WebSocket integration example for the multi-language chat agent.

This example demonstrates how to set up and use the WebSocket communication layer
with the chat agent services.
"""

import os
import redis
from flask import Flask
from flask_socketio import SocketIO

# Import WebSocket components
from chat_agent.websocket import initialize_websocket_handlers
from chat_agent.services.chat_agent import create_chat_agent
from chat_agent.services.session_manager import create_session_manager
from chat_agent.services.language_context import create_language_context_manager
from chat_agent.services.chat_history import create_chat_history_manager
from chat_agent.services.groq_client import create_groq_client


def create_app_with_websockets():
    """
    Create a Flask app with WebSocket support configured.
    
    This example shows how to integrate all the services and set up WebSocket handlers.
    """
    # Create Flask app
    app = Flask(__name__)
    app.config['SECRET_KEY'] = 'your-secret-key-here'
    app.config['TESTING'] = True
    
    # Create SocketIO instance
    socketio = SocketIO(app, cors_allowed_origins="*")
    
    # Create Redis client (in production, use proper Redis configuration)
    redis_client = redis.Redis(host='localhost', port=6379, db=0, decode_responses=False)
    
    # Create service instances (these would normally be created with proper configuration)
    try:
        # Create Groq client (requires API key)
        groq_api_key = os.getenv('GROQ_API_KEY', 'your-groq-api-key')
        groq_client = create_groq_client(groq_api_key)
        
        # Create language context manager
        language_context_manager = create_language_context_manager(redis_client)
        
        # Create session manager
        session_manager = create_session_manager(redis_client)
        
        # Create chat history manager
        chat_history_manager = create_chat_history_manager(redis_client)
        
        # Create chat agent
        chat_agent = create_chat_agent(
            groq_client, language_context_manager, 
            session_manager, chat_history_manager
        )
        
        # Initialize WebSocket handlers
        initialize_websocket_handlers(
            socketio, chat_agent, session_manager, redis_client
        )
        
        print("✓ WebSocket handlers initialized successfully")
        
    except Exception as e:
        print(f"❌ Failed to initialize services: {e}")
        print("Note: This example requires proper service configuration")
    
    return app, socketio


def websocket_client_example():
    """
    Example of how a client would interact with the WebSocket API.
    
    This shows the expected message formats and event flow.
    """
    print("\n=== WebSocket Client Example ===")
    
    # Connection authentication
    auth_data = {
        'session_id': 'example-session-123',
        'user_id': 'example-user-456'
    }
    print(f"1. Connect with auth: {auth_data}")
    
    # Send a chat message
    message_data = {
        'content': 'Hello! Can you help me with Python programming?',
        'session_id': 'example-session-123'
    }
    print(f"2. Send message: {message_data}")
    
    # Expected response events:
    print("3. Expected response events:")
    print("   - message_received: Acknowledgment")
    print("   - processing_status: Processing started")
    print("   - response_start: Response generation started")
    print("   - response_chunk: Streaming response chunks")
    print("   - response_complete: Response finished")
    
    # Switch programming language
    language_switch_data = {
        'language': 'javascript',
        'session_id': 'example-session-123'
    }
    print(f"4. Switch language: {language_switch_data}")
    print("   - Expected: language_switched event")
    
    # Typing indicators
    print("5. Typing indicators:")
    print("   - Send: typing_start event")
    print("   - Send: typing_stop event")
    print("   - Receive: user_typing / user_typing_stop events")
    
    # Health check
    ping_data = {'timestamp': '2024-01-01T12:00:00Z'}
    print(f"6. Health check: ping {ping_data}")
    print("   - Expected: pong event with timestamps")
    
    # Get session info
    print("7. Get session info: get_session_info event")
    print("   - Expected: session_info event with session details")


def main():
    """Main example function."""
    print("WebSocket Integration Example")
    print("=" * 40)
    
    # Show how to create app with WebSocket support
    try:
        app, socketio = create_app_with_websockets()
        print("✓ Flask app with WebSocket support created")
    except Exception as e:
        print(f"❌ Failed to create app: {e}")
    
    # Show client interaction examples
    websocket_client_example()
    
    print("\n=== WebSocket Events Summary ===")
    print("Server Events (sent by server):")
    print("  - connection_status: Connection established/status")
    print("  - message_received: Message acknowledgment")
    print("  - processing_status: Processing state updates")
    print("  - response_start: Response generation started")
    print("  - response_chunk: Streaming response content")
    print("  - response_complete: Response generation finished")
    print("  - language_switched: Language context changed")
    print("  - user_typing: User is typing indicator")
    print("  - user_typing_stop: User stopped typing")
    print("  - pong: Health check response")
    print("  - session_info: Session information")
    print("  - error: Error messages")
    
    print("\nClient Events (sent by client):")
    print("  - connect: Establish connection (with auth)")
    print("  - message: Send chat message")
    print("  - language_switch: Change programming language")
    print("  - typing_start: Start typing indicator")
    print("  - typing_stop: Stop typing indicator")
    print("  - ping: Health check request")
    print("  - get_session_info: Request session information")
    print("  - disconnect: Close connection")
    
    print("\n=== Security Features ===")
    print("✓ Message validation and sanitization")
    print("✓ Rate limiting (30 messages per minute)")
    print("✓ XSS protection with HTML escaping")
    print("✓ Malicious content detection")
    print("✓ Session-based authentication")
    print("✓ Connection timeout management")
    
    print("\n=== Performance Features ===")
    print("✓ Redis-based connection management")
    print("✓ In-memory connection caching")
    print("✓ Streaming response support")
    print("✓ Connection pooling and cleanup")
    print("✓ Typing indicators for better UX")


if __name__ == '__main__':
    main()