eveAI/eveai_chat/socket_handlers/chat_handler.py

import uuid
from functools import wraps

from flask_jwt_extended import create_access_token, get_jwt_identity, verify_jwt_in_request, decode_token
from flask_socketio import emit, disconnect, join_room, leave_room
from flask import current_app, request, session
from sqlalchemy.exc import SQLAlchemyError
from datetime import datetime, timedelta
from prometheus_client import Counter, Histogram
from time import time

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

# 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'])


# Decorator to measure SocketIO events
def track_socketio_event(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        event_type = func.__name__
        socketio_message_counter.labels(event_type=event_type).inc()
        start_time = time()
        result = func(*args, **kwargs)
        latency = time() - start_time
        socketio_message_latency.labels(event_type=event_type).observe(latency)
        return result
    return wrapper


@socketio.on('connect')
@track_socketio_event
def handle_connect():
    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}')

        if not validate_api_key(tenant_id, api_key):
            raise Exception("Invalid tenant_id - api_key combination")

        # Create JWT token
        token = create_access_token(identity={"tenant_id": tenant_id, "api_key": api_key})

        # Create a unique room for this client
        room = f"{tenant_id}_{request.sid}"
        join_room(room)

        # Create a unique session ID
        if 'session_id' not in session:
            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
    except Exception as e:
        current_app.logger.error(f'SocketIO: Connection failed: {e}')
        # communicate connection problem to client
        emit('connect_error', {'status': 'Connection Failed'})
        disconnect()


@socketio.on('disconnect')
@track_socketio_event
def handle_disconnect():
    room = session.get('room')
    if room:
        leave_room(room)


@socketio.on('heartbeat')
def handle_heartbeat():
    last_activity = session.get('last_activity')
    if datetime.now() - last_activity > current_app.config.get('SOCKETIO_MAX_IDLE_TIME'):
        disconnect()


@socketio.on('user_message')
def handle_message(data):
    try:
        session['last_activity'] = datetime.now()
        current_tenant_id = validate_incoming_data(data)
        room = session.get('room')

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


@socketio.on('check_task_status')
def check_task_status(data):
    current_app.logger.debug(f'SocketIO: Checking Task Status ... {data}')
    task_id = data.get('task_id')
    room = session.get('room')
    if not task_id:
        emit('task_status', {'status': 'error', 'message': 'Missing task ID'}, room=room)
        return

    task_result = current_celery.AsyncResult(task_id)
    if task_result.state == 'PENDING':
        emit('task_status', {'status': 'pending', 'taskId': task_id}, room=room)
    elif task_result.state == 'SUCCESS':
        result = task_result.result
        current_app.logger.debug(f'SocketIO: Task {task_id} returned: {result}')
        response = {
            'status': 'success',
            'taskId': task_id,
            'answer': result['answer'],
            'citations': result['citations'],
            'algorithm': result['algorithm'],
            'interaction_id': result['interaction_id'],
        }
        emit('task_status', response, room=room)
    else:
        current_app.logger.error(f'SocketIO: Task {task_id} has failed. Error: {task_result.info}')
        emit('task_status', {'status': task_result.state, 'message': str(task_result.info)}, room=room)


@socketio.on('feedback')
def handle_feedback(data):
    try:
        current_tenant_id = validate_incoming_data(data)

        interaction_id = data.get('interactionId')
        feedback = data.get('feedback')  # 'up' or 'down'

        Database(current_tenant_id).switch_schema()

        interaction = Interaction.query.get_or_404(interaction_id)
        interaction.appreciation = 0 if feedback == 'down' else 100
        try:
            db.session.commit()
            emit('feedback_received', {'status': 'success', 'interaction_id': interaction_id})
        except SQLAlchemyError as e:
            current_app.logger.error(f'SocketIO: Feedback handling failed: {e}')
            db.session.rollback()
            emit('feedback_received', {'status': 'Could not register feedback', 'interaction_id': interaction_id})
            raise e
    except Exception as e:
        current_app.logger.error(f'SocketIO: Feedback handling failed: {e}')
        disconnect()


def validate_api_key(tenant_id, api_key):
    tenant = Tenant.query.get_or_404(tenant_id)
    decrypted_api_key = simple_encryption.decrypt_api_key(tenant.encrypted_chat_api_key)

    return decrypted_api_key == api_key


def validate_incoming_data(data):
    token = data.get('token')
    if not token:
        raise Exception("Missing token")

    decoded_token = decode_token(token)
    if not decoded_token:
        raise Exception("Invalid token")

    token_sub = decoded_token.get('sub')

    if not token_sub:
        raise Exception("Missing token subject")

    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")

    return current_tenant_id