Improving chat functionality significantly throughout the application.
This commit is contained in:
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user