API key working, CORS working, SocketIO working (but no JWT), Chat client v1, Session implemented (server side)

eveai_chat/views/chat_views.py

@@ -1,34 +1,75 @@
-# from . import user_bp
-import uuid
 from datetime import datetime as dt, timezone as tz
-from flask import request, redirect, url_for, flash, render_template, Blueprint, session, current_app
+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
+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('/', methods=['GET', 'POST'])
-def chat():
-    return render_template('chat.html')
+@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
 
 
-@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('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')
-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
+@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