Update main.py

main.py

@@ -4,18 +4,16 @@ from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import StreamingResponse, HTMLResponse
 from pydantic import BaseModel
 import httpx
+from functools import lru_cache
 
-# Load environment variables from .env file
 load_dotenv()
 
 app = FastAPI()
 
-# Get API keys and secret endpoint from environment variables
 api_keys_str = os.getenv('API_KEYS')
 valid_api_keys = api_keys_str.split(',') if api_keys_str else []
 secret_api_endpoint = os.getenv('SECRET_API_ENDPOINT')
 
-# Check if the endpoint is set in the environment
 if not secret_api_endpoint:
     raise HTTPException(status_code=500, detail="API endpoint is not configured in environment variables.")
 
@@ -50,8 +48,6 @@ async def root():
     """
     return HTMLResponse(content=html_content)
 
-# Remove cache from get_models
-@app.get("/v1/models")
 async def get_models():
     async with httpx.AsyncClient() as client:
         try:
@@ -61,33 +57,36 @@ async def get_models():
         except httpx.RequestError as e:
             raise HTTPException(status_code=500, detail=f"Request failed: {e}")
 
+@app.get("/v1/models")
+async def fetch_models():
+    return await get_models()
+
 @app.post("/v1/chat/completions")
 async def get_completion(payload: Payload, request: Request):
     api_key = request.headers.get("Authorization")
     
-    # Validate API key
     if api_key not in valid_api_keys:
         raise HTTPException(status_code=403, detail="Forbidden: Invalid API key. Join dsc.gg/chadgang and DM @mr_leaderyt on Discord for a free API key :)")
     
-    # Prepare the payload for streaming
     payload_dict = {**payload.dict(), "stream": True}
 
-    # Define an asynchronous generator to stream the response line by line
-    async def stream_generator():
+    async def stream_generator(payload_dict):
         async with httpx.AsyncClient() as client:
             try:
-                async with client.stream("POST", secret_api_endpoint, json=payload_dict, timeout=10) as response:
+                async with client.stream("POST", f"{secret_api_endpoint}/v1/chat/completions", json=payload_dict, timeout=10) as response:
                     response.raise_for_status()
                     async for line in response.aiter_lines():
                         if line:
-                            yield f"{line}\n"  # Add a newline to distinguish each line
-            except httpx.RequestError as e:
-                raise HTTPException(status_code=500, detail=f"Streaming failed: {e}")
-
-    # Return the streaming response
-    return StreamingResponse(stream_generator(), media_type="application/json")
+                            yield f"{line}\n"
+            except httpx.HTTPStatusError as status_err:
+                raise HTTPException(status_code=status_err.response.status_code, detail=f"HTTP error: {status_err}")
+            except httpx.RequestError as req_err:
+                raise HTTPException(status_code=500, detail=f"Streaming failed: {req_err}")
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")
+    
+    return StreamingResponse(stream_generator(payload_dict), media_type="application/json")
 
-# Log the API endpoints
 @app.on_event("startup")
 async def startup_event():
     print("API endpoints:")
@@ -95,7 +94,6 @@ async def startup_event():
     print("GET /models")
     print("POST /v1/chat/completions")
 
-# Run the server with Uvicorn using the 'main' module
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)