Update src/main_MM.py

src/main_MM.py

@@ -9,24 +9,19 @@ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLI
 '''
 
 import os
-from flask import Flask,request,jsonify
+from flask import Flask, request, jsonify
 from flask_swagger_ui import get_swaggerui_blueprint
-
-from routing.injectionRouter import injection_router
-from routing.detoxifyRouter import detoxify_router
-from routing.embedingRouter import embed_router
-from routing.privacyRouter import privacy_router
-from routing.topicRouter import topic_router
-from routing.sentimentRouter import sentiment_router
-from routing.invisibletextRouter import invisibletext_router
-from routing.gibberishRouter import gibberish_router
-from routing.bancodeRouter import bancode_router
-# from routing.safety_router import img_router
-from config.logger import CustomLogger,request_id_var
+from config.logger import CustomLogger, request_id_var
 from waitress import serve
-from werkzeug.exceptions import HTTPException,UnsupportedMediaType,BadRequest
+from werkzeug.exceptions import HTTPException, UnsupportedMediaType
 from mapper.mapper import *
 
+# ----------------------
+# Create necessary log directories to avoid runtime warnings
+os.makedirs("/responsible-ai/logs", exist_ok=True)
+os.makedirs("/aicloud/logs", exist_ok=True)
+# ----------------------
+
 SWAGGER_URL = '/rai/v1/raimoderationmodels/docs'  
 API_URL = '/static/swagger.json'
 
@@ -35,43 +30,87 @@ swaggerui_blueprint = get_swaggerui_blueprint(
     API_URL
 )
 
-log=CustomLogger()
+log = CustomLogger()
 app = Flask(__name__)
 app.register_blueprint(swaggerui_blueprint)
-app.register_blueprint(injection_router,url_prefix='/rai/v1/raimoderationmodels')
-app.register_blueprint(detoxify_router,url_prefix='/rai/v1/raimoderationmodels')
-app.register_blueprint(embed_router,url_prefix='/rai/v1/raimoderationmodels')
-app.register_blueprint(privacy_router,url_prefix='/rai/v1/raimoderationmodels')
-app.register_blueprint(topic_router,url_prefix='/rai/v1/raimoderationmodels')
-app.register_blueprint(sentiment_router,url_prefix='/rai/v1/raimoderationmodels')
-app.register_blueprint(invisibletext_router,url_prefix='/rai/v1/raimoderationmodels')
-app.register_blueprint(gibberish_router,url_prefix='/rai/v1/raimoderationmodels')
-app.register_blueprint(bancode_router,url_prefix='/rai/v1/raimoderationmodels')
-# app.register_blueprint(img_router,url_prefix='/rai/v1/raimoderationmodels')
-
 
+# ----------------------
+# Function to load all models before starting server
+def load_models():
+    from routing.injectionRouter import init_injection_model
+    from routing.detoxifyRouter import init_detoxify_model
+    from routing.embedingRouter import init_embed_model
+    from routing.privacyRouter import init_privacy_model
+    from routing.topicRouter import init_topic_model
+
+    init_injection_model()
+    init_detoxify_model()
+    init_embed_model()
+    init_privacy_model()
+    init_topic_model()
+# ----------------------
+
+# ----------------------
+# Function to register all routers
+def register_routers():
+    from routing.injectionRouter import injection_router
+    from routing.detoxifyRouter import detoxify_router
+    from routing.embedingRouter import embed_router
+    from routing.privacyRouter import privacy_router
+    from routing.topicRouter import topic_router
+    from routing.sentimentRouter import sentiment_router
+    from routing.invisibletextRouter import invisibletext_router
+    from routing.gibberishRouter import gibberish_router
+    from routing.bancodeRouter import bancode_router
+    # from routing.safety_router import img_router
+
+    routers = [
+        injection_router,
+        detoxify_router,
+        embed_router,
+        privacy_router,
+        topic_router,
+        sentiment_router,
+        invisibletext_router,
+        gibberish_router,
+        bancode_router
+        # img_router
+    ]
+
+    for r in routers:
+        app.register_blueprint(r, url_prefix='/rai/v1/raimoderationmodels')
+# ----------------------
+
+# ----------------------
+# Error handling
 def handle_http_exception(exc):
-    """Handles HTTP exceptions, returning a JSON response."""
     response = jsonify({"error": exc.description})
     response.status_code = exc.code
     return response
 
 @app.errorhandler(HTTPException)
 def handle_all_http_exceptions(exc):
-    """Global exception handler for HTTP errors."""
     return handle_http_exception(exc)
 
 def handle_unsupported_mediatype(exc):
-    """Handles unsupported media type exceptions."""
-    return jsonify({"error": "Unsupported media type"}), 415  # 415 Unsupported Media Type
+    return jsonify({"error": "Unsupported media type"}), 415
 
 @app.errorhandler(UnsupportedMediaType)
 def handle_all_unsupported_mediatype_exceptions(exc):
-    """Global exception handler for unsupported media types."""
     return handle_unsupported_mediatype(exc)
-
+# ----------------------
 
 if __name__ == "__main__":
-   serve(app, host='0.0.0.0', port=7860, threads=int(os.getenv('THREADS',1)),connection_limit=int(os.getenv('CONNECTION_LIMIT',500)), channel_timeout=int(os.getenv('CHANNEL_TIMEOUT',120)))
-   #app.run()
+    load_models()       # Load models before registering routes
+    register_routers()  # Register routers after models are loaded
+
+    serve(
+        app,
+        host='0.0.0.0',
+        port=7860,
+        threads=int(os.getenv('THREADS', 1)),
+        connection_limit=int(os.getenv('CONNECTION_LIMIT', 500)),
+        channel_timeout=int(os.getenv('CHANNEL_TIMEOUT', 120))
+    )
+