Update main.py

main.py

@@ -1,11 +1,10 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel
 import requests
 import json
 
 app = FastAPI()
 
-# Model and API Key validation
 valid_api_keys = ['PARTH-SADARIA-NI-API-CHAWI', 'HEET-NI-CHEESEWADI-API-KEY']
 model_aliases = {
     'Llama-3.1-Nemotron-70B-Instruct': 'nvidia/Llama-3.1-Nemotron-70B-Instruct',
@@ -22,8 +21,8 @@ class Payload(BaseModel):
     messages: list
 
 @app.post("/api/v1/chat/completions")
-async def get_completion(payload: Payload):
-    api_key = payload.headers.get("Authorization")
+async def get_completion(payload: Payload, request: Request):
+    api_key = request.headers.get("Authorization")
     
     # API Key validation
     if api_key not in valid_api_keys:
@@ -31,6 +30,9 @@ async def get_completion(payload: Payload):
     
     # Model alias resolution
     user_model = payload.model
+    if user_model not in model_aliases:
+        raise HTTPException(status_code=400, detail="Invalid model name")
+    
     full_model_name = model_aliases.get(user_model, user_model)
     
     is_deepinfra_model = full_model_name in model_aliases.values()
@@ -38,14 +40,16 @@ async def get_completion(payload: Payload):
     # Determine the URL to send the request to
     url = "https://api.deepinfra.com/v1/openai/chat/completions" if is_deepinfra_model else "https://gpt.tiptopuni.com/api/openai/v1/chat/completions"
     
-    # Send the API request to the appropriate service
-    response = requests.post(url, json={**payload.dict(), "model": full_model_name})
-    
-    if response.status_code != 200:
-        raise HTTPException(status_code=response.status_code, detail="Failed to fetch data")
+    try:
+        response = requests.post(url, json={**payload.dict(), "model": full_model_name})
+        response.raise_for_status()  # Raises HTTPError if the response status code is 4xx/5xx
+    except requests.exceptions.RequestException as e:
+        raise HTTPException(status_code=500, detail=f"Request failed: {e}")
     
-    # Return the response from the model API
-    return response.json()
+    try:
+        return response.json()
+    except ValueError:
+        raise HTTPException(status_code=500, detail="Non-JSON response received")
 
 # Run the server with Uvicorn using the 'main' module
 if __name__ == "__main__":