Update main.py (#6)

- Update main.py (b1c818e9b0ac26da5e2145259aaf22caefcd32b4)


Co-authored-by: True Saiyan <[email protected]>

main.py

@@ -50,15 +50,13 @@ from serpapi import GoogleSearch
 
 from fastapi import FastAPI, UploadFile, File, Response
 from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
-from fastapi import Depends, FastAPI, HTTPException, status
+from fastapi import Depends HTTPException, status
 from fastapi.openapi.utils import get_openapi
-from fastapi.responses import StreamingResponse
-from fastapi import HTTPException
-from fastapi import FastAPI, Request, Header
+from fastapi.responses import StreamingResponse, JSONResponse, FileResponse
+from fastapi import Request, Header
 from fastapi import Body, Query
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
-from fastapi.responses import FileResponse
 
 from pymongo import MongoClient
 
@@ -124,7 +122,15 @@ contact_support = """
 We are aware that AI is currently offline. This seems to be caused by the API
 We are investigating and doing our best to get things back online as soon as possible. Thank you for your patience
 
-Contact Support @xtdevs
+~ Contact Support @xtdevs
+"""
+
+internal_error = """
+There has been an Internal error. We are aware of this error and notice that it can be
+caused by your search terms being to explict, too confusing, or it can be caused by the API.
+Please modify your search terms and/or try again later thank you for your understanding.
+
+~ @xtdevs Team
 """
 
 def get_all_api_keys():
@@ -548,7 +554,7 @@ def dalle_3xl(
     except requests.exceptions.RequestException:
         raise HTTPException(
             status_code=500,
-            detail=CustomErrorResponseModel(detail=[{}])
+            detail=internal_error
         )
 
     try:
@@ -556,7 +562,7 @@ def dalle_3xl(
     except Exception:
         raise HTTPException(
             status_code=500,
-            detail=CustomErrorResponseModel(detail=[{}])
+            detail=contact_support
         )
 
     if encoded_string:
@@ -578,14 +584,14 @@ def open_dalle(
     except requests.exceptions.RequestException:
         raise HTTPException(
             status_code=500,
-            detail=CustomErrorResponseModel(detail=[{}])
+            detail=internal_error
         )
     try:
         encoded_string = base64.b64encode(response.content).decode("utf-8")
     except Exception:
         raise HTTPException(
             status_code=500,
-            detail=CustomErrorResponseModel(detail=[{}])
+            detail=contact_support
         )
     if encoded_string:
         return SuccessResponse(status="True", randydev={"data": encoded_string})
@@ -614,7 +620,7 @@ def Picsart_Pro(
         logger.error(f"Error in Picsart request: {e}")
         raise HTTPException(
             status_code=500,
-            detail=CustomErrorResponseModel(detail=[{}])
+            detail=internal_error
         )
 
     try:
@@ -626,7 +632,7 @@ def Picsart_Pro(
         logger.error(f"Error in Picsart image retrieval: {e}")
         raise HTTPException(
             status_code=500,
-            detail=CustomErrorResponseModel(detail=[{}])
+            detail=contact_support
         )
 
     if encoded_string:
@@ -648,14 +654,14 @@ def Anime_Styled(
     except requests.exceptions.RequestException:
         raise HTTPException(
             status_code=500,
-            detail=CustomErrorResponseModel(detail=[{}])
+            detail=internal_error
         )
     try:
         encoded_string = base64.b64encode(response.content).decode("utf-8")
     except Exception:
         raise HTTPException(
             status_code=500,
-            detail=CustomErrorResponseModel(detail=[{}])
+            detail=contact_support
         )
     if encoded_string:
         return SuccessResponse(status="True", randydev={"data": encoded_string})
@@ -672,14 +678,14 @@ def image_unsplash(item: GetImageUnsplash):
     except requests.exceptions.RequestException:
         raise HTTPException(
             status_code=500,
-            detail=CustomErrorResponseModel(detail=[{}])
+            detail=internal_error
         )
     try:
         encoded_string = base64.b64encode(response.content).decode("utf-8")
     except Exception:
         raise HTTPException(
             status_code=500,
-            detail=CustomErrorResponseModel(detail=[{}])
+            detail=contact_support
         )
     if encoded_string:
         return SuccessResponse(status="True", randydev={"data": encoded_string})
@@ -788,7 +794,7 @@ def v1beta3_google_ai(
             message = results.get("output")
         return SuccessResponse(status="True", randydev={"message": message})
     except:
-        return SuccessResponse(status="False", randydev={"message": contact_support})
+        return SuccessResponse(status="False", randydev={"message": internal_error})
 
 @app.post("/ryuzaki/gemini-ai-pro", response_model=SuccessResponse, responses={422: {"model": ErrorStatus}})
 def gemini_pro(item: GeminiPro):
@@ -862,7 +868,7 @@ def v1beta2_google_ai(
             message = results.get("content")
         return SuccessResponse(status="True", randydev={"message": message})
     except:
-        return SuccessResponse(status="False", randydev={"message": contact_support})
+        return SuccessResponse(status="False", randydev={"message": internal_error})
 
 @app.post("/ryuzaki/new-monitor", response_model=SuccessResponse, responses={422: {"model": ErrorStatus}})
 def new_monitor(
@@ -1369,6 +1375,17 @@ description = """
 
 •Developed by [@xtdevs](https://t.me/xtdevs)
 """
+
+def custom_exception_handler(request: Request, exc: HTTPException) -> JSONResponse:
+    error_detail = [{"error": str(exc.detail)}]
+    custom_error_model = CustomErrorResponseModel(detail=error_detail)
+
+    return JSONResponse(
+        status_code=exc.status_code,
+        content=custom_error_model.dict(),
+        headers=exc.headers,
+    )
+
 def custom_openapi():
     if app.openapi_schema:
         return app.openapi_schema
@@ -1386,3 +1403,5 @@ def custom_openapi():
     return app.openapi_schema
 
 app.openapi = custom_openapi
+
+app.add_exception_handler(HTTPException, custom_exception_handler)