Add MCP based AI Chat

Author: robin-rpr

app.py

@@ -1,23 +1,29 @@
 from flask import Flask, render_template, request, redirect, url_for, abort, current_app as app
 from utils.text import time_ago, format_text, clean_text, alphanumeric_text, clean_text, sanitize_text, date_text
 from routes.asn import asn_blueprint, asn_api_v1_blueprint
-from utils.socket import sock, broadcaster
+from routes.ai import ai_blueprint, ai_api_v1_blueprint
+from utils.socket import sock, subscriber, chat
+from utils.database import get_db, close_db
+from utils.mcp import list_tools, run_tool
 from utils.cache import cache, caching
 from utils.limiter import limiter
 from flask_compress import Compress
 from flask_talisman import Talisman
 from utils.seo import get_sitemap
 from datetime import timedelta
 from logging import StreamHandler
+from sqlalchemy import text
 from config import Config
 import logging
+import gevent
 import atexit
-import sass # type: ignore
+import json
+import sass
 import sys
 import re
-import gevent
+
 from gevent import monkey
-monkey.patch_all()
+monkey.patch_all(ssl=False)
 
 def compile_scss():
     scss_file = 'static/styles/main.scss'
@@ -162,23 +168,26 @@ def sitemap_index_xml():
 
 
     """
-    Subscriptions
+    WebSocket
     """
 
     # Initialize socket
-    broadcaster.init_app(app)
+    subscriber.init_app(app)
+    chat.init_app(app)
 
     # Register socket shutdown
-    atexit.register(lambda: broadcaster.shutdown(wait=False))
+    atexit.register(lambda: subscriber.shutdown(wait=False))
+    atexit.register(lambda: chat.shutdown(wait=False))
 
     # Start socket
-    broadcaster.start()
+    subscriber.start()
+    chat.start()
 
     @sock.route('/ws/v1/subscribe/<resource>')
     def ws_v1_subscribe(ws, resource):
         try:
             app.logger.info(f"WebSocket connected for {resource}")
-            broadcaster.register(ws, resource=resource)
+            subscriber.register(ws, resource=resource)
 
             while True:
                 try:
@@ -197,14 +206,41 @@ def ws_v1_subscribe(ws, resource):
             app.logger.error(f"WebSocket connection error for {resource}: {str(e)}", exc_info=True)
         finally:
             app.logger.info(f"WebSocket disconnected for {resource}")
-            broadcaster.unregister(ws, resource)
+            subscriber.unregister(ws, resource)
+
+    @sock.route('/ws/v1/chat')
+    def ws_v1_chat(ws):
+        try:
+            app.logger.info(f"WebSocket connected for chat")
+            chat.register(ws)
+
+            while True:
+                # Receive message from client
+                message = ws.receive()
+                if not message:
+                    break
+                    
+                try:
+                    data = json.loads(message)
+                    chat.handle(ws, data)
+                        
+                except json.JSONDecodeError:
+                    ws.send(json.dumps({'type': 'error', 'message': 'Invalid JSON format'}))
+                except Exception as e:
+                    app.logger.error("WebSocket chat error: %s", str(e), exc_info=True)
+                    ws.send(json.dumps({'type': 'error', 'message': str(e)}))
+                    
+        except Exception as e:
+            app.logger.error("WebSocket connection error: %s", str(e), exc_info=True)
+        finally:
+            app.logger.info(f"WebSocket disconnected for chat")
+            chat.unregister(ws)
 
 
     """
-    Basic
+    Base
     """
 
-
     @app.route('/')
     @caching(timeout=86400) # 24 hours
     def index():
@@ -221,6 +257,8 @@ def index():
     # Register Blueprints
     app.register_blueprint(asn_blueprint, url_prefix='/as')
     app.register_blueprint(asn_api_v1_blueprint, url_prefix='/api/v1/as')
+    app.register_blueprint(ai_blueprint, url_prefix='/ai')
+    app.register_blueprint(ai_api_v1_blueprint, url_prefix='/api/v1/ai')
 
     return app
 

config.py

@@ -19,6 +19,9 @@ class Config():
     # Postmark
     POSTMARK_API_KEY = os.getenv('POSTMARK_API_KEY')
 
+    # OpenRouter
+    OPENROUTER_API_KEY = os.getenv('OPENROUTER_API_KEY')
+    
     # Environment
     ENVIRONMENT = os.getenv('ENVIRONMENT', 'development').lower()
     LOG_LEVEL = os.getenv('LOG_LEVEL', 'INFO' if ENVIRONMENT == 'production' else 'DEBUG').upper()
@@ -31,12 +34,13 @@ def validate():
             'FLASK_HOST': Config.FLASK_HOST,
             'KAFKA_FQDN': Config.KAFKA_FQDN,
             'KAFKA_JMX_FQDN': Config.KAFKA_JMX_FQDN,
-            'POSTMARK_API_KEY': Config.POSTMARK_API_KEY,
             'POSTGRES_HOST': Config.POSTGRES_HOST,
             'POSTGRES_PORT': Config.POSTGRES_PORT,
             'POSTGRES_DB': Config.POSTGRES_DB,
             'POSTGRES_USER': Config.POSTGRES_USER,
             'POSTGRES_PASSWORD': Config.POSTGRES_PASSWORD,
+            'POSTMARK_API_KEY': Config.POSTMARK_API_KEY,
+            'OPENROUTER_API_KEY': Config.OPENROUTER_API_KEY,
         }
         for key, val in required.items():
             if not val:

manage.py

@@ -12,15 +12,21 @@ def run(workers="1", host="localhost", port=8080, reload=False):
     """
     Runs the application.
     """
-    # Validate the configuration
+    # Validate.
     Config.validate()
 
     if reload:
-        # Get the app instance
         from app import create_app
         app = create_app()
-        # Use the built-in Flask development server which works with flask-sock
-        app.run(host=host, port=port, debug=True, extra_files=['templates/**/*.html', 'static/**/*.scss'])
+        app.run(
+            host=host,
+            port=port,
+            debug=True,
+            extra_files=[
+                'templates/**/*.html',
+                'static/**/*.scss'
+            ]
+        )
     else:
         subprocess.run([
             "gunicorn", 

requirements.txt

@@ -20,4 +20,5 @@ confluent-kafka==2.5.3
 rocksdb-py==0.0.5
 scapy==2.6.1
 SQLAlchemy==2.0.36
-psycopg2-binary==2.9.10
+psycopg2-binary==2.9.10
+openai==2.1.0

routes/ai.py

@@ -0,0 +1,32 @@
+from flask import Blueprint, render_template, jsonify, current_app as app
+from utils.ai import client
+
+# Create Blueprints
+ai_blueprint = Blueprint('ai', __name__)
+ai_api_v1_blueprint = Blueprint('ai_api_v1', __name__)
+
+@ai_blueprint.route('/')
+def page():
+    return render_template('pages/ai.html')
+
+@ai_api_v1_blueprint.route('/models', methods=['GET'])
+def api_v1_ai_models():
+    try:
+        models_response = client.models.list()
+        models = []
+        
+        for model in models_response.data:
+            if hasattr(model, 'id'):
+                models.append({
+                    'id': model.id,
+                    'name': getattr(model, 'name', model.id),
+                    'description': getattr(model, 'description', ''),
+                    'context_length': getattr(model, 'context_length', None)
+                })
+        
+        return jsonify({'models': models})
+    except Exception as e:
+        app.logger.error("Error fetching models: %s", str(e), exc_info=True)
+        return jsonify({'error': str(e)}), 500
+
+

routes/asn.py

@@ -347,12 +347,6 @@ def api_v1_asn_advertised_prefixes(asn):
     finally:
         close_db(db)
 
-@asn_blueprint.route("/<int:asn>")
-#@caching(timeout=86400) # 24 hours
+@asn_blueprint.route('/<int:asn>')
 def asn(asn):
-    return render_template('pages/asn.html', asn=asn)
-
-# Register blueprints
-def init_app(app):
-    app.register_blueprint(asn_blueprint)
-    app.register_blueprint(asn_api_v1_blueprint)
+    return render_template('pages/asn.html', asn=asn)

templates/base.html

@@ -59,7 +59,6 @@
         <!-- Analytics -->
         <script defer data-domain="bgp-data.net" src="https://plausible.io/js/script.outbound-links.tagged-events.js"></script>
         <script>window.plausible = window.plausible || function() { (window.plausible.q = window.plausible.q || []).push(arguments) }</script>
-        <script type="text/javascript" src="https://a.usbrowserspeed.com/cs?pid=ddae2e0bce828a30a7b24f94f87290780f71120eaf9f11353f234c3bd86512d3&puid=68dbaa52e5531cd2ffc2c510"></script>
 
         <!-- Additional headers -->
         {% block headers %}{% endblock %}