<BACK> - Added check connection function.

Author: eric114610

backend/app/db.py

@@ -12,7 +12,7 @@
     load_dotenv(dotenv_path=os.path.join(os.path.dirname(__file__), '..', '.env'))
 
 client = MongoClient(os.getenv("MONGO_URI"))
-db = client.smartreply
+db = client.epic
 
 def get_user_db(UserEmail: str):
     # Create collections if they don't exist
@@ -107,3 +107,10 @@ def get_records_to_process(UserEmail: str):
 def set_records_to_process(UserEmail: str, count: int):
     db[UserEmail][SETTINGS].update_one({"Email": UserEmail}, {"$set": {"Records_to_process": count}})
     return count
+
+def get_collections():
+    try:
+        db.list_collection_names()
+        return True
+    except Exception as e:
+        return False

backend/app/gemini_client.py

@@ -48,3 +48,11 @@ def generate_mail_class(input: list[ThreadInput], userPrompt:str, processCount:
 
     response = model.generate_content(prompt)
     return response.text
+
+def get_gemini_connection():
+    try:
+        genai.list_models()
+        return True
+    except Exception as e:
+        return False
+    

backend/app/routes.py

@@ -1,8 +1,8 @@
 from fastapi import APIRouter
 from pydantic import BaseModel
 import time
-from backend.app.gemini_client import generate_mail_class
-from backend.app.db import save_thread, query_thread, get_stats, get_user_db, get_prompt, save_prompt, reset_stats, delete_records_by_oldest, delete_records_by_category, get_records_to_process, set_records_to_process
+from backend.app.gemini_client import *
+from backend.app.db import *
 from backend.app.schemas import RecordInput, ThreadInput, UserData, PromptInput, DeleteInput, RecordsToProcessInput
 
 router = APIRouter()
@@ -153,3 +153,10 @@ async def apply_records_to_process(input: RecordsToProcessInput):
     set_records_to_process(input.UserEmail, input.RecordsToProcess)
     print(f"Records to process set to {input.RecordsToProcess} for {input.UserEmail}.")
     return {"message": "Records to process setting applied successfully."}
+
+@router.get("/get_connection")
+async def get_connection():
+    if (get_collections() and get_gemini_connection()):
+        return {"status": True}
+    else:
+        return {"status": False}