Update database/connections.py

database/connections.py

@@ -1,24 +1,19 @@
-import streamlit as st
 from pymongo import MongoClient
 from core.logger import logger
 from core.settings import cnf
 
+import streamlit as st
+
+@st.cache_resource
 def get_mongo_client():
-    """Initialize MongoDB connection with enhanced error handling"""
-    try:
-        uri = cnf.MONGO_URI
-        dbname = cnf.MONGO_DB
-        if not uri:
-            raise RuntimeError("MONGO_URI is not set")
-        client = MongoClient(uri)
-        db = client[dbname]
-        db.command('ping')
-        logger.info("MongoDB connection established successfully")
-        return db
-    except Exception as e:
-        print(f"Database connection failed: {e}")
-        return None
 
+    uri = cnf.MONGO_URI
+    dbname = cnf.MONGO_DB
+    client = MongoClient(uri)
+    db = client[dbname]
+    db.command("ping")
+    logger.info("MongoDB connection established successfully")
+    return db
 
 def get_results_collection():
     db = get_mongo_client()
@@ -26,6 +21,13 @@ def get_results_collection():
         return None
     return db["img_gen"]
 
+def get_image_collection():
+    db = get_mongo_client()
+    if db is None:
+        return None
+    return db["image_analyses"]
+
+
 def get_video_collection():
     db = get_mongo_client()
     if db is None:
@@ -36,4 +38,4 @@ def get_script_collection():
     db = get_mongo_client()
     if db is None:
         return None
-    return db["script_generator"]
+    return db["script_generation"]