76 lines
2.5 KiB
Python
76 lines
2.5 KiB
Python
from datetime import datetime as dt, timezone as tz
|
|
from flask import request, redirect, url_for, render_template, Blueprint, session, current_app, jsonify
|
|
from flask_security import hash_password, roles_required, roles_accepted
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
from flask_jwt_extended import create_access_token, jwt_required, get_jwt_identity
|
|
from flask_socketio import emit, join_room, leave_room
|
|
import ast
|
|
|
|
|
|
from common.models.user import User, Tenant
|
|
from common.models.interaction import ChatSession, Interaction, InteractionEmbedding
|
|
from common.models.document import Embedding
|
|
from common.extensions import db, socketio, kms_client
|
|
from common.utils.database import Database
|
|
|
|
chat_bp = Blueprint('chat_bp', __name__, url_prefix='/chat')
|
|
|
|
|
|
@chat_bp.route('/register_client', methods=['POST'])
|
|
def register_client():
|
|
tenant_id = request.json.get('tenant_id')
|
|
api_key = request.json.get('api_key')
|
|
|
|
# Validate tenant_id and api_key here (e.g., check against the database)
|
|
if validate_tenant(tenant_id, api_key):
|
|
access_token = create_access_token(identity={'tenant_id': tenant_id, 'api_key': api_key})
|
|
return jsonify({'token': access_token}), 200
|
|
else:
|
|
return jsonify({'message': 'Invalid credentials'}), 401
|
|
|
|
|
|
@socketio.on('connect', namespace='/chat')
|
|
@jwt_required()
|
|
def handle_connect():
|
|
current_tenant = get_jwt_identity()
|
|
print(f'Tenant {current_tenant["tenant_id"]} connected')
|
|
|
|
|
|
@socketio.on('message', namespace='/chat')
|
|
@jwt_required()
|
|
def handle_message(data):
|
|
current_tenant = get_jwt_identity()
|
|
print(f'Tenant {current_tenant["tenant_id"]} sent a message: {data}')
|
|
# Store interaction in the database
|
|
emit('response', {'data': 'Message received'}, broadcast=True)
|
|
|
|
|
|
def validate_tenant(tenant_id, api_key):
|
|
tenant = Tenant.query.get_or_404(tenant_id)
|
|
encrypted_api_key = ast.literal_eval(tenant.encrypted_chat_api_key)
|
|
|
|
decrypted_api_key = kms_client.decrypt_api_key(encrypted_api_key)
|
|
|
|
return decrypted_api_key == api_key
|
|
|
|
|
|
|
|
# @chat_bp.route('/', methods=['GET', 'POST'])
|
|
# def chat():
|
|
# return render_template('chat.html')
|
|
#
|
|
#
|
|
# @chat.record_once
|
|
# def on_register(state):
|
|
# # TODO: write initialisation code when the blueprint is registered (only once)
|
|
# # socketio.init_app(state.app)
|
|
# pass
|
|
#
|
|
#
|
|
# @socketio.on('message', namespace='/chat')
|
|
# def handle_message(message):
|
|
# # TODO: write message handling code to actually realise chat
|
|
# # print('Received message:', message)
|
|
# # socketio.emit('response', {'data': message}, namespace='/chat')
|
|
# pass
|