Update main.py

main.py

@@ -32,20 +32,6 @@ if not secret_api_endpoint or not secret_api_endpoint_2 or not secret_api_endpoi
 alternate_models = {"gpt-4o-mini", "claude-3-haiku", "llama-3.1-70b", "mixtral-8x7b"}
 
 available_model_ids = []
-def check_server_status():
-    server_down = True
-    
-    def decorator(func):
-        async def wrapper(*args, **kwargs):
-            if server_down:
-                raise HTTPException(
-                    status_code=503, 
-                    detail="Server is currently unavailable. All services are temporarily down."
-                )
-            return await func(*args, **kwargs)
-        return wrapper
-    return decorator
-    
 class Payload(BaseModel):
     model: str
     messages: list
@@ -157,30 +143,41 @@ async def fetch_models():
     return await get_models()
 
 
+from fastapi import HTTPException
+from fastapi.responses import JSONResponse
+
+server_status = False  # Set to False to simulate the server being down
+
 @app.post("/chat/completions")
 @app.post("/v1/chat/completions")
-@check_server_status()
-async def get_completion(payload: Payload,request: Request):
+async def get_completion(payload: Payload, request: Request):
+    # Check server status
+    if not server_status:
+        return JSONResponse(
+            status_code=503,
+            content={"message": "Server is down. Please try again later."}
+        )
+    
     model_to_use = payload.model if payload.model else "gpt-4o-mini"
+    
     # Validate model availability
     if model_to_use not in available_model_ids:
         raise HTTPException(
-            status_code=400, 
+            status_code=400,
             detail=f"Model '{model_to_use}' is not available. Check /models for the available model list."
         )
-
+    
     # Proceed with the request handling
     payload_dict = payload.dict()
     payload_dict["model"] = model_to_use
-
+    
     # Select the appropriate endpoint
     endpoint = secret_api_endpoint_2 if model_to_use in alternate_models else secret_api_endpoint
     current_time = (datetime.datetime.utcnow() + datetime.timedelta(hours=5, minutes=30)).strftime("%Y-%m-%d %I:%M:%S %p")
     aaip = request.client.host
     print(f"Time: {current_time}, {aaip}")
     print(payload_dict)
-
-
+    
     async def stream_generator(payload_dict):
         scraper = cloudscraper.create_scraper()  # Create a CloudScraper session
         try:
@@ -198,12 +195,11 @@ async def get_completion(payload: Payload,request: Request):
                 raise HTTPException(status_code=404, detail="The requested resource was not found.")
             elif response.status_code >= 500:
                 raise HTTPException(status_code=500, detail="Server error. Try again later.")
-
+            
             # Stream response lines to the client
             for line in response.iter_lines():
                 if line:
                     yield line.decode('utf-8') + "\n"
-
         except requests.exceptions.RequestException as req_err:
             # Handle request-specific errors
             print(response.text)
@@ -212,7 +208,7 @@ async def get_completion(payload: Payload,request: Request):
             # Handle unexpected errors
             print(response.text)
             raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")
-
+    
     return StreamingResponse(stream_generator(payload_dict), media_type="application/json")
 # Remove the duplicated endpoint and combine the functionality
 @app.get("/images/generations") #pollinations.ai thanks to them :)