- Created a new eveai_chat plugin to support the new dynamic possibilities of the Specialists. Currently only supports standard Rag retrievers (i.e. no extra arguments).

2024-11-26 13:35:29 +01:00
parent 7702a6dfcc
commit 07d89d204f
42 changed files with 1771 additions and 989 deletions
--- a/eveai_chat/socket_handlers/chat_handler.py
+++ b/eveai_chat/socket_handlers/chat_handler.py
@@ -8,18 +8,77 @@ from sqlalchemy.exc import SQLAlchemyError
 from datetime import datetime, timedelta
 from prometheus_client import Counter, Histogram
 from time import time
+import re

 from common.extensions import socketio, db, simple_encryption
 from common.models.user import Tenant
 from common.models.interaction import Interaction
 from common.utils.celery_utils import current_celery
 from common.utils.database import Database
+from common.utils.token_validation import TokenValidator
+from common.utils.eveai_exceptions import EveAISocketInputException

 # Define custom metrics
 socketio_message_counter = Counter('socketio_message_count', 'Count of SocketIO messages', ['event_type'])
 socketio_message_latency = Histogram('socketio_message_latency_seconds', 'Latency of SocketIO message processing', ['event_type'])


+class RoomManager:
+    def __init__(self):
+        self.active_rooms = {}  # Store active room metadata
+
+    def validate_room_format(self, room_id: str) -> bool:
+        """Validate room ID format: tenant_id_sessionid_timestamp"""
+        pattern = r'^\d+_[a-zA-Z0-9]+_\d+$'
+        return bool(re.match(pattern, room_id))
+
+    def is_room_active(self, room_id: str) -> bool:
+        return room_id in self.active_rooms
+
+    def validate_room_ownership(self, room_id: str, tenant_id: int, token: str) -> bool:
+        if not self.is_room_active(room_id):
+            return False
+
+        room_data = self.active_rooms[room_id]
+        return (room_data['tenant_id'] == tenant_id and
+                room_data['token'] == token)
+
+    def create_room(self, tenant_id: int, token: str) -> str:
+        """Create new room with metadata"""
+        timestamp = int(datetime.now().timestamp())
+        room_id = f"{tenant_id}_{request.sid}_{timestamp}"
+
+        self.active_rooms[room_id] = {
+            'tenant_id': tenant_id,
+            'token': token,
+            'created_at': datetime.now(),
+            'last_activity': datetime.now()
+        }
+
+        return room_id
+
+    def update_room_activity(self, room_id: str):
+        """Update room's last activity timestamp"""
+        if room_id in self.active_rooms:
+            self.active_rooms[room_id]['last_activity'] = datetime.now()
+
+    def cleanup_inactive_rooms(self, max_age_hours: int = 1):
+        """Remove inactive rooms"""
+        now = datetime.now()
+        cutoff = now - timedelta(hours=max_age_hours)
+
+        inactive_rooms = [
+            room_id for room_id, data in self.active_rooms.items()
+            if data['last_activity'] < cutoff
+        ]
+
+        for room_id in inactive_rooms:
+            del self.active_rooms[room_id]
+
+
+room_manager = RoomManager()
+
+
 # Decorator to measure SocketIO events
 def track_socketio_event(func):
    @wraps(func)
@@ -37,42 +96,89 @@ def track_socketio_event(func):
@socketio.on('connect')
@track_socketio_event
 def handle_connect():
+    """Handle incoming socket connections with enhanced security"""
    try:
-        tenant_id = request.args.get('tenantId')
-        if not tenant_id:
-            raise Exception("Missing Tenant ID")
-        api_key = request.args.get('apiKey')
-        if not api_key:
-            raise Exception("Missing API Key")
-        current_app.logger.info(f'SocketIO: Connection handling found Tenant {tenant_id} with API Key {api_key}')
+        current_app.logger.debug('Handle Connection')
+        token = request.args.get('token')
+        if not token:
+            raise ValueError("Missing token")

-        if not validate_api_key(tenant_id, api_key):
-            raise Exception("Invalid tenant_id - api_key combination")
+        current_app.logger.debug(f"Token received: {token}")

-        # Create JWT token
-        token = create_access_token(identity={"tenant_id": tenant_id, "api_key": api_key})
+        if not token:
+            raise ValueError("Missing token")

-        # Create a unique room for this client
-        room = f"{tenant_id}_{request.sid}"
+        current_app.logger.info(f"Trying to connect with: {token}")
+
+        validator = TokenValidator()
+        validation_result = validator.validate_token(token)
+
+        if not validation_result.is_valid:
+            current_app.logger.error(f"Socket connection failed: {validation_result.error_message}")
+            emit('connect_error', {'error': validation_result.error_message})
+            disconnect()
+            return
+
+        # Create room and setup session
+        room = room_manager.create_room(validation_result.tenant_id, token)
        join_room(room)

-        # Create a unique session ID
-        if 'session_id' not in session:
-            session['session_id'] = str(uuid.uuid4())
-
+        session['session_id'] = str(uuid.uuid4())
        session['last_activity'] = datetime.now()
        session['room'] = room

-        # Communicate connection to client
-        emit('connect', {'status': 'Connected', 'tenant_id': tenant_id, 'room': room})
-        emit('authenticated', {'token': token, 'room': room})  # Emit custom event with the token
+        # Emit success events
+        emit('connect', {
+            'status': 'Connected',
+            'tenant_id': validation_result.tenant_id,
+            'room': room
+        })
+        emit('authenticated', {'token': token, 'room': room})
+        current_app.logger.info(f"Socket connection succeeded: {token} / {room}")
+
    except Exception as e:
-        current_app.logger.error(f'SocketIO: Connection failed: {e}')
-        # communicate connection problem to client
+        current_app.logger.error(f"Socket connection failed: {str(e)}")
        emit('connect_error', {'status': 'Connection Failed'})
        disconnect()


+@socketio.on('rejoin_room')
+def handle_room_rejoin(data):
+    try:
+        token = data.get('token')
+        tenant_id = data.get('tenant_id')
+        previous_room = data.get('previousRoom')
+
+        validator = TokenValidator()
+        validation_result = validator.validate_token(token, require_session=True)
+        if not validation_result.is_valid:
+            emit('room_rejoin_result', {'success': False, 'error': validation_result.error_message})
+            return
+
+        if not all([token, tenant_id, previous_room]):
+            raise ValueError("Missing required rejoin data")
+
+        # Validate room ownership
+        if not room_manager.validate_room_ownership(previous_room, tenant_id, token):
+            raise ValueError("Invalid room ownership")
+
+        # Rejoin room
+        join_room(previous_room)
+        session['room'] = previous_room
+        room_manager.update_room_activity(previous_room)
+
+        emit('room_rejoin_result', {
+            'success': True,
+            'room': previous_room
+        })
+    except Exception as e:
+        current_app.logger.error(f'Room rejoin failed: {e}')
+        emit('room_rejoin_result', {
+            'success': False,
+            'error': str(e)
+        })
+
+
@socketio.on('disconnect')
@track_socketio_event
 def handle_disconnect():
@@ -90,37 +196,71 @@ def handle_heartbeat():

@socketio.on('user_message')
 def handle_message(data):
+    current_app.logger.debug(f"SocketIO: Received message: {data}")
    try:
+        validator = TokenValidator()
+        validation_result = validator.validate_token(data.get('token'))
+
+        if not validation_result.is_valid:
+            emit('error', {'message': validation_result.error_message})
+            return
+
+        current_app.logger.debug(f"SocketIO: token validated: {validation_result}")
+
+        room = session.get('room')
+        current_app.logger.debug(f"SocketIO: Room in session: {room}, Room in arguments: {data.get('room')}")
+
+        current_app.logger.debug(f"SocketIO: Room: {room}")
+        if not room or not room_manager.is_room_active(room):
+            raise Exception("Invalid or inactive room")
+        current_app.logger.debug(f"SocketIO: Room active: {room}")
+
+        if not room_manager.validate_room_ownership(room, data['tenant_id'], data['token']):
+            raise Exception("Room ownership validation failed")
+        current_app.logger.debug(f"SocketIO: Room ownership validated: {room}")
+
+        room_manager.update_room_activity(room)
+        current_app.logger.debug(f"SocketIO: Room activity updated: {room}")
        session['last_activity'] = datetime.now()
        current_tenant_id = validate_incoming_data(data)
-        room = session.get('room')
+
+        current_app.logger.debug(f"SocketIO: Incoming data validated: {current_tenant_id}")

        # Offload actual processing of question
        task = current_celery.send_task('execute_specialist',
                                        queue='llm_interactions',
                                        args=[
                                            current_tenant_id,
-                                            data['specialistId'],
+                                            data['specialist_id'],
                                            data['arguments'],
                                            session['session_id'],
                                            data['timezone'],
                                            room
                                        ])
        response = {
-            'tenantId': data['tenantId'],
+            'tenantId': current_tenant_id,
            'message': f'Processing question ... Session ID = {session["session_id"]}',
            'taskId': task.id,
+            'room': room,
        }
-        current_app.logger.debug(f"Sent message with {data}, response {response}")
+        current_app.logger.debug(f"SocketIO: Sent response {response}")
        emit('bot_response', response, room=room)
    except Exception as e:
        current_app.logger.error(f'SocketIO: Message handling failed: {str(e)}')
-        disconnect()
+        emit('error', {'message': 'Failed to process message'}, room=room)


@socketio.on('check_task_status')
 def check_task_status(data):
    current_app.logger.debug(f'SocketIO: Checking Task Status ... {data}')
+
+    validator = TokenValidator()
+    validation_result = validator.validate_token(data.get('token'))
+
+    if not validation_result.is_valid:
+        emit('feedback_received', {'status': 'error', 'error': validation_result.error_message})
+        return
+
    task_id = data.get('task_id')
    room = session.get('room')
    if not task_id:
@@ -145,6 +285,7 @@ def check_task_status(data):
                'insufficient_info': specialist_result.get('insufficient_info', False)
            },
            'interaction_id': result['interaction_id'],
+            'room': room
        }
        emit('task_status', response, room=room)
    else:
@@ -153,7 +294,15 @@ def check_task_status(data):

@socketio.on('feedback')
 def handle_feedback(data):
+    current_app.logger.debug(f'SocketIO: Received feedback: {data}')
    try:
+        validator = TokenValidator()
+        validation_result = validator.validate_token(data.get('token'))
+
+        if not validation_result.is_valid:
+            emit('feedback_received', {'status': 'error', 'error': validation_result.error_message})
+            return
+
        current_tenant_id = validate_incoming_data(data)

        interaction_id = data.get('interactionId')
@@ -163,9 +312,15 @@ def handle_feedback(data):

        interaction = Interaction.query.get_or_404(interaction_id)
        interaction.appreciation = 0 if feedback == 'down' else 100
+
+        room = session.get('room')
+        if not room:
+            emit('feedback_received', {'status': 'error', 'message': 'No active room'})
+            return
+
        try:
            db.session.commit()
-            emit('feedback_received', {'status': 'success', 'interaction_id': interaction_id})
+            emit('feedback_received', {'status': 'success', 'interaction_id': interaction_id, 'room': room}, room=room)
        except SQLAlchemyError as e:
            current_app.logger.error(f'SocketIO: Feedback handling failed: {e}')
            db.session.rollback()
@@ -184,25 +339,20 @@ def validate_api_key(tenant_id, api_key):


 def validate_incoming_data(data):
+    current_app.logger.debug(f'SocketIO: Validating incoming data: {data}')
    token = data.get('token')
    if not token:
-        raise Exception("Missing token")
+        raise EveAISocketInputException("SocketIO: Missing token in input")

    decoded_token = decode_token(token)
    if not decoded_token:
-        raise Exception("Invalid token")
+        raise EveAISocketInputException("SocketIO: Invalid token in input")

-    token_sub = decoded_token.get('sub')
+    current_app.logger.debug(f'SocketIO: Decoded token: {decoded_token}')

-    if not token_sub:
-        raise Exception("Missing token subject")
+    current_tenant_id = decoded_token.get('sub')

-    current_tenant_id = token_sub.get('tenant_id')
    if not current_tenant_id:
-        raise Exception("Missing tenant_id")
-
-    current_api_key = token_sub.get('api_key')
-    if not current_api_key:
-        raise Exception("Missing api_key")
+        raise EveAISocketInputException("SocketIO: Missing tenant_id (sub) in input")

    return current_tenant_id