CoolFace
Apppublic

WebashalarForML/scratch_chat

sourceHugging Facemitupdated 1y agoView on Hugging Face
0likes
websocket_example.py185 linesDownload Raw Back to examples
1"""
2WebSocket integration example for the multi-language chat agent.
3
4This example demonstrates how to set up and use the WebSocket communication layer
5with the chat agent services.
6"""
7
8import os
9import redis
10from flask import Flask
11from flask_socketio import SocketIO
12
13# Import WebSocket components
14from chat_agent.websocket import initialize_websocket_handlers
15from chat_agent.services.chat_agent import create_chat_agent
16from chat_agent.services.session_manager import create_session_manager
17from chat_agent.services.language_context import create_language_context_manager
18from chat_agent.services.chat_history import create_chat_history_manager
19from chat_agent.services.groq_client import create_groq_client
20
21
22def create_app_with_websockets():
23    """
24    Create a Flask app with WebSocket support configured.
25    
26    This example shows how to integrate all the services and set up WebSocket handlers.
27    """
28    # Create Flask app
29    app = Flask(__name__)
30    app.config['SECRET_KEY'] = 'your-secret-key-here'
31    app.config['TESTING'] = True
32    
33    # Create SocketIO instance
34    socketio = SocketIO(app, cors_allowed_origins="*")
35    
36    # Create Redis client (in production, use proper Redis configuration)
37    redis_client = redis.Redis(host='localhost', port=6379, db=0, decode_responses=False)
38    
39    # Create service instances (these would normally be created with proper configuration)
40    try:
41        # Create Groq client (requires API key)
42        groq_api_key = os.getenv('GROQ_API_KEY', 'your-groq-api-key')
43        groq_client = create_groq_client(groq_api_key)
44        
45        # Create language context manager
46        language_context_manager = create_language_context_manager(redis_client)
47        
48        # Create session manager
49        session_manager = create_session_manager(redis_client)
50        
51        # Create chat history manager
52        chat_history_manager = create_chat_history_manager(redis_client)
53        
54        # Create chat agent
55        chat_agent = create_chat_agent(
56            groq_client, language_context_manager, 
57            session_manager, chat_history_manager
58        )
59        
60        # Initialize WebSocket handlers
61        initialize_websocket_handlers(
62            socketio, chat_agent, session_manager, redis_client
63        )
64        
65        print("✓ WebSocket handlers initialized successfully")
66        
67    except Exception as e:
68        print(f"❌ Failed to initialize services: {e}")
69        print("Note: This example requires proper service configuration")
70    
71    return app, socketio
72
73
74def websocket_client_example():
75    """
76    Example of how a client would interact with the WebSocket API.
77    
78    This shows the expected message formats and event flow.
79    """
80    print("\n=== WebSocket Client Example ===")
81    
82    # Connection authentication
83    auth_data = {
84        'session_id': 'example-session-123',
85        'user_id': 'example-user-456'
86    }
87    print(f"1. Connect with auth: {auth_data}")
88    
89    # Send a chat message
90    message_data = {
91        'content': 'Hello! Can you help me with Python programming?',
92        'session_id': 'example-session-123'
93    }
94    print(f"2. Send message: {message_data}")
95    
96    # Expected response events:
97    print("3. Expected response events:")
98    print("   - message_received: Acknowledgment")
99    print("   - processing_status: Processing started")
100    print("   - response_start: Response generation started")
101    print("   - response_chunk: Streaming response chunks")
102    print("   - response_complete: Response finished")
103    
104    # Switch programming language
105    language_switch_data = {
106        'language': 'javascript',
107        'session_id': 'example-session-123'
108    }
109    print(f"4. Switch language: {language_switch_data}")
110    print("   - Expected: language_switched event")
111    
112    # Typing indicators
113    print("5. Typing indicators:")
114    print("   - Send: typing_start event")
115    print("   - Send: typing_stop event")
116    print("   - Receive: user_typing / user_typing_stop events")
117    
118    # Health check
119    ping_data = {'timestamp': '2024-01-01T12:00:00Z'}
120    print(f"6. Health check: ping {ping_data}")
121    print("   - Expected: pong event with timestamps")
122    
123    # Get session info
124    print("7. Get session info: get_session_info event")
125    print("   - Expected: session_info event with session details")
126
127
128def main():
129    """Main example function."""
130    print("WebSocket Integration Example")
131    print("=" * 40)
132    
133    # Show how to create app with WebSocket support
134    try:
135        app, socketio = create_app_with_websockets()
136        print("✓ Flask app with WebSocket support created")
137    except Exception as e:
138        print(f"❌ Failed to create app: {e}")
139    
140    # Show client interaction examples
141    websocket_client_example()
142    
143    print("\n=== WebSocket Events Summary ===")
144    print("Server Events (sent by server):")
145    print("  - connection_status: Connection established/status")
146    print("  - message_received: Message acknowledgment")
147    print("  - processing_status: Processing state updates")
148    print("  - response_start: Response generation started")
149    print("  - response_chunk: Streaming response content")
150    print("  - response_complete: Response generation finished")
151    print("  - language_switched: Language context changed")
152    print("  - user_typing: User is typing indicator")
153    print("  - user_typing_stop: User stopped typing")
154    print("  - pong: Health check response")
155    print("  - session_info: Session information")
156    print("  - error: Error messages")
157    
158    print("\nClient Events (sent by client):")
159    print("  - connect: Establish connection (with auth)")
160    print("  - message: Send chat message")
161    print("  - language_switch: Change programming language")
162    print("  - typing_start: Start typing indicator")
163    print("  - typing_stop: Stop typing indicator")
164    print("  - ping: Health check request")
165    print("  - get_session_info: Request session information")
166    print("  - disconnect: Close connection")
167    
168    print("\n=== Security Features ===")
169    print("✓ Message validation and sanitization")
170    print("✓ Rate limiting (30 messages per minute)")
171    print("✓ XSS protection with HTML escaping")
172    print("✓ Malicious content detection")
173    print("✓ Session-based authentication")
174    print("✓ Connection timeout management")
175    
176    print("\n=== Performance Features ===")
177    print("✓ Redis-based connection management")
178    print("✓ In-memory connection caching")
179    print("✓ Streaming response support")
180    print("✓ Connection pooling and cleanup")
181    print("✓ Typing indicators for better UX")
182
183
184if __name__ == '__main__':
185    main()