- Adding a Tenant Type

- Allow filtering on Tenant Types & searching for parts of Tenant names
- Implement health checks
- Start Prometheus monitoring (needs to be finalized)
- Refine audio_processor and srt_processor to reduce duplicate code and support for larger files
- Introduce repopack to reason in LLMs about the code

eveai_chat/__init__.py

@@ -3,7 +3,7 @@ import logging.config
 from flask import Flask, jsonify
 import os
 
-from common.extensions import db, socketio, jwt, cors, session, simple_encryption
+from common.extensions import db, socketio, jwt, cors, session, simple_encryption, metrics
 from config.logging_config import LOGGING
 from eveai_chat.socket_handlers import chat_handler
 from common.utils.cors_utils import create_cors_after_request
@@ -32,17 +32,6 @@ def create_app(config_file=None):
     app.celery = make_celery(app.name, app.config)
     init_celery(app.celery, app)
 
-    # Register Blueprints
-    # register_blueprints(app)
-
-    @app.route('/ping')
-    def ping():
-        return 'pong'
-
-    @app.route('/health', methods=['GET'])
-    def health():
-        return jsonify({'status': 'ok'}), 200
-
     app.logger.info("EveAI Chat Server Started Successfully")
     app.logger.info("-------------------------------------------------------------------------------------------------")
     return app
@@ -61,8 +50,8 @@ def register_extensions(app):
                       ping_interval=app.config.get('SOCKETIO_PING_INTERVAL'),
                       )
     jwt.init_app(app)
-    # kms_client.init_app(app)
     simple_encryption.init_app(app)
+    metrics.init_app(app)
 
     # Cors setup
     cors.init_app(app, resources={r"/chat/*": {"origins": "*"}})
@@ -71,6 +60,7 @@ def register_extensions(app):
     session.init_app(app)
 
 
+
 def register_blueprints(app):
-    from .views.chat_views import chat_bp
-    app.register_blueprint(chat_bp)
+    from views.healthz_views import healthz_bp
+    app.register_blueprint(healthz_bp)

eveai_chat/socket_handlers/chat_handler.py

@@ -1,10 +1,13 @@
 import uuid
+from functools import wraps
 
 from flask_jwt_extended import create_access_token, get_jwt_identity, verify_jwt_in_request, decode_token
 from flask_socketio import emit, disconnect, join_room, leave_room
 from flask import current_app, request, session
 from sqlalchemy.exc import SQLAlchemyError
 from datetime import datetime, timedelta
+from prometheus_client import Counter, Histogram
+from time import time
 
 from common.extensions import socketio, db, simple_encryption
 from common.models.user import Tenant
@@ -12,8 +15,27 @@ from common.models.interaction import Interaction
 from common.utils.celery_utils import current_celery
 from common.utils.database import Database
 
+# Define custom metrics
+socketio_message_counter = Counter('socketio_message_count', 'Count of SocketIO messages', ['event_type'])
+socketio_message_latency = Histogram('socketio_message_latency_seconds', 'Latency of SocketIO message processing', ['event_type'])
+
+
+# Decorator to measure SocketIO events
+def track_socketio_event(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        event_type = func.__name__
+        socketio_message_counter.labels(event_type=event_type).inc()
+        start_time = time()
+        result = func(*args, **kwargs)
+        latency = time() - start_time
+        socketio_message_latency.labels(event_type=event_type).observe(latency)
+        return result
+    return wrapper
+
 
 @socketio.on('connect')
+@track_socketio_event
 def handle_connect():
     try:
         current_app.logger.debug(f'SocketIO: Connection handling started using {request.args}')
@@ -58,6 +80,7 @@ def handle_connect():
 
 
 @socketio.on('disconnect')
+@track_socketio_event
 def handle_disconnect():
     room = session.get('room')
     if room:

eveai_chat/views/chat_views.py

@@ -1,77 +0,0 @@
-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})
-        current_app.logger.debug(f'Tenant Registration: Tenant {tenant_id} registered successfully')
-        return jsonify({'token': access_token}), 200
-    else:
-        current_app.logger.debug(f'Tenant Registration: Invalid tenant_id ({tenant_id}) or api_key ({api_key})')
-        return jsonify({'message': 'Invalid credentials'}), 401
-
-
-@socketio.on('connect', namespace='/chat')
-@jwt_required()
-def handle_connect():
-    current_tenant = get_jwt_identity()
-    current_app.logger.debug(f'Tenant {current_tenant["tenant_id"]} connected')
-
-
-@socketio.on('message', namespace='/chat')
-@jwt_required()
-def handle_message(data):
-    current_tenant = get_jwt_identity()
-    current_app.logger.debug(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

eveai_chat/views/healthz_views.py

@@ -0,0 +1,70 @@
+from flask import Blueprint, current_app, request
+from flask_healthz import HealthError
+from sqlalchemy.exc import SQLAlchemyError
+from celery.exceptions import TimeoutError as CeleryTimeoutError
+from common.extensions import db, metrics, minio_client
+from common.utils.celery_utils import current_celery
+from eveai_chat.socket_handlers.chat_handler import socketio_message_counter, socketio_message_latency
+
+healthz_bp = Blueprint('healthz', __name__, url_prefix='/_healthz')
+
+
+def liveness():
+    try:
+        # Basic check to see if the app is running
+        return True
+    except Exception:
+        raise HealthError("Liveness check failed")
+
+
+def readiness():
+    checks = {
+        "database": check_database(),
+        "celery": check_celery(),
+        # Add more checks as needed
+    }
+
+    if not all(checks.values()):
+        raise HealthError("Readiness check failed")
+
+
+def check_database():
+    try:
+        # Perform a simple database query
+        db.session.execute("SELECT 1")
+        return True
+    except SQLAlchemyError:
+        current_app.logger.error("Database check failed", exc_info=True)
+        return False
+
+
+def check_celery():
+    try:
+        # Send a simple task to Celery
+        result = current_celery.send_task('tasks.ping', queue='llm_interactions')
+        response = result.get(timeout=10)  # Wait for up to 10 seconds for a response
+        return response == 'pong'
+    except CeleryTimeoutError:
+        current_app.logger.error("Celery check timed out", exc_info=True)
+        return False
+    except Exception as e:
+        current_app.logger.error(f"Celery check failed: {str(e)}", exc_info=True)
+        return False
+
+
+@healthz_bp.route('/metrics')
+@metrics.do_not_track()
+def prometheus_metrics():
+    return metrics.generate_latest()
+
+
+def init_healtz(app):
+    app.config.update(
+        HEALTHZ={
+            "live": "healthz_views.liveness",
+            "ready": "healthz_views.readiness",
+        }
+    )
+    # Register SocketIO metrics with Prometheus
+    metrics.register(socketio_message_counter)
+    metrics.register(socketio_message_latency)