Enable model variables & start working on RAG task

eveai_chat/socket_handlers/chat_handler.py

@@ -4,6 +4,7 @@ from flask import current_app, request
 
 from common.extensions import socketio, kms_client
 from common.models.user import Tenant
+from common.utils.celery_utils import current_celery
 
 
 @socketio.on('connect')
@@ -66,13 +67,17 @@ def handle_message(data):
         if not current_api_key:
             raise Exception("Missing api_key")
 
-        # Store interaction in the database
-
+        # Offload actual processing of question
+        task = current_celery.send_task('ask_question', queue='llm_interactions', args=[
+            current_tenant_id,
+            data['message'],
+        ])
+        current_app.logger.debug(f'SocketIO: Message offloading for tenant {current_tenant_id}, '
+                                 f'Question: {task.id}')
         response = {
             'tenantId': data['tenantId'],
-            'message': f'This is a bot response. Responding to message {data['message']} '
-                       f'from tenant {current_tenant_id}',
-            'messageId': 'bot-message-id',
+            'message': 'Processing question ...',
+            'taskId': task.id,
             'algorithm': 'alg1'
         }
         current_app.logger.debug(f"SocketIO: Message handling sent bot response: {response}")
@@ -82,6 +87,30 @@ def handle_message(data):
         disconnect()
 
 
+@socketio.on('check_task_status')
+def check_task_status(data):
+    task_id = data.get('task_id')
+    current_app.logger.debug(f'SocketIO: Check task status for task_id: {task_id}')
+    if not task_id:
+        emit('task_status', {'status': 'error', 'message': 'Missing task ID'})
+        return
+
+    task_result = current_celery.AsyncResult(task_id)
+    if task_result.state == 'PENDING':
+        current_app.logger.debug(f'SocketIO: Task {task_id} is pending')
+        emit('task_status', {'status': 'pending', 'taskId': task_id})
+    elif task_result.state != 'FAILURE':
+        current_app.logger.debug(f'SocketIO: Task {task_id} has finished. Status: {task_result.state}, '
+                                 f'Result: {task_result.result}')
+        emit('task_status', {
+            'status': task_result.state,
+            'result': task_result.result
+        })
+    else:
+        current_app.logger.error(f'SocketIO: Task {task_id} has failed. Error: {task_result.info}')
+        emit('task_status', {'status': 'failure', 'message': str(task_result.info)})
+
+
 def validate_api_key(tenant_id, api_key):
     tenant = Tenant.query.get_or_404(tenant_id)
     decrypted_api_key = kms_client.decrypt_api_key(tenant.encrypted_chat_api_key)