Enable model variables & start working on RAG task

eveai_chat/__init__.py

@@ -7,6 +7,9 @@ from common.extensions import db, socketio, jwt, kms_client, cors, session
 from config.logging_config import LOGGING
 from eveai_chat.socket_handlers import chat_handler
 from common.utils.cors_utils import create_cors_after_request
+from common.utils.celery_utils import make_celery, init_celery
+
+
 
 
 def create_app(config_file=None):
@@ -20,6 +23,9 @@ def create_app(config_file=None):
     logging.config.dictConfig(LOGGING)
     register_extensions(app)
 
+    app.celery = make_celery(app.name, app.config)
+    init_celery(app.celery, app)
+
     # Register Blueprints
     register_blueprints(app)
 
@@ -49,13 +55,6 @@ def register_extensions(app):
     cors.init_app(app, resources={r"/chat/*": {"origins": "*"}})
     app.after_request(create_cors_after_request('/chat'))
 
-    # Session setup
-    # redis_config = app.config['SESSION_REDIS']
-    # redis_client = Redis(host=redis_config['host'],
-    #                      port=redis_config['port'],
-    #                      db=redis_config['db'],
-    #                      password=redis_config['password']
-    #                      )
     session.init_app(app)
 
 

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)