Improving chat functionality significantly throughout the application.

This commit is contained in:
Josako
2024-06-12 11:07:18 +02:00
parent 27b6de8734
commit be311c440b
22 changed files with 604 additions and 127 deletions

View File

@@ -1,9 +1,13 @@
import uuid
from flask_jwt_extended import create_access_token, get_jwt_identity, verify_jwt_in_request, decode_token
from flask_socketio import emit, disconnect
from flask import current_app, request
from flask import current_app, request, session
from sqlalchemy.exc import SQLAlchemyError
from common.extensions import socketio, kms_client
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
@@ -26,6 +30,10 @@ def handle_connect():
token = create_access_token(identity={"tenant_id": tenant_id, "api_key": api_key})
current_app.logger.debug(f'SocketIO: Connection handling created token: {token} for tenant {tenant_id}')
# Create a unique session ID
if 'session_id' not in session:
session['session_id'] = str(uuid.uuid4())
# Communicate connection to client
emit('connect', {'status': 'Connected', 'tenant_id': tenant_id})
emit('authenticated', {'token': token}) # Emit custom event with the token
@@ -71,14 +79,15 @@ def handle_message(data):
task = current_celery.send_task('ask_question', queue='llm_interactions', args=[
current_tenant_id,
data['message'],
data['language'],
session['session_id'],
])
current_app.logger.debug(f'SocketIO: Message offloading for tenant {current_tenant_id}, '
f'Question: {task.id}')
response = {
'tenantId': data['tenantId'],
'message': 'Processing question ...',
'message': f'Processing question ... Session ID = {session["session_id"]}',
'taskId': task.id,
'algorithm': 'alg1'
}
current_app.logger.debug(f"SocketIO: Message handling sent bot response: {response}")
emit('bot_response', response, broadcast=True)
@@ -99,16 +108,39 @@ def check_task_status(data):
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':
elif task_result.state == 'SUCCESS':
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
})
result = task_result.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)
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)})
emit('task_status', {'status': task_result.state, 'message': str(task_result.info)})
@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
def validate_api_key(tenant_id, api_key):