Improvements on the chat UI, prepare for supporting multiple models, and adding feedback to interactions.

eveai_chat/socket_handlers/chat_handler.py

@@ -9,6 +9,7 @@ from common.extensions import socketio, kms_client, db
 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
 
 
 @socketio.on('connect')
@@ -35,7 +36,9 @@ def handle_connect():
             session['session_id'] = str(uuid.uuid4())
 
         # Communicate connection to client
+        current_app.logger.debug(f'SocketIO: Connection handling sending status to client for tenant {tenant_id}')
         emit('connect', {'status': 'Connected', 'tenant_id': tenant_id})
+        current_app.logger.debug(f'SocketIO: Connection handling sending authentication token to client')
         emit('authenticated', {'token': token})  # Emit custom event with the token
         current_app.logger.debug(f'SocketIO: Connection handling sent token to client for tenant {tenant_id}')
     except Exception as e:
@@ -55,25 +58,8 @@ def handle_message(data):
     try:
         current_app.logger.debug(f"SocketIO: Message handling received message from tenant {data['tenantId']}: "
                                  f"{data['message']} with token {data['token']}")
-        token = data.get('token')
-        if not token:
-            raise Exception("Missing token")
 
-        # decoded_token = decode_token(token.split(" ")[1]) # remove "Bearer "
-        decoded_token = decode_token(token)
-        if not decoded_token:
-            raise Exception("Invalid token")
-        current_app.logger.debug(f"SocketIO: Message handling decoded token: {decoded_token}")
-
-        token_sub = 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")
+        current_tenant_id = validate_incoming_data(data)
 
         # Offload actual processing of question
         task = current_celery.send_task('ask_question', queue='llm_interactions', args=[
@@ -128,19 +114,30 @@ def check_task_status(data):
 
 @socketio.on('feedback')
 def handle_feedback(data):
-    interaction_id = data.get('interaction_id')
-    feedback = data.get('feedback')  # 'up' or 'down'
-    # Store feedback in the database associated with the interaction_id
-    interaction = Interaction.query.get_or_404(interaction_id)
-    interaction.feedback = 0 if feedback == 'down' else 1
     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
+        current_app.logger.debug(f'SocketIO: Feedback handling received feedback with data: {data}')
+
+        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)
+        current_app.logger.debug(f'Processing feedback for interaction: {interaction}')
+        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.debug(f'SocketIO: Feedback handling failed: {e}')
+        disconnect()
 
 
 def validate_api_key(tenant_id, api_key):
@@ -148,3 +145,29 @@ def validate_api_key(tenant_id, api_key):
     decrypted_api_key = kms_client.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")
+    tenant_id = token_sub.get('tenant_id')
+
+    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