Update app.py

app.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, Request, HTTPException
-from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
 from typing import List
 import g4f
@@ -9,17 +9,24 @@ import uuid
 import time
 import asyncio
 
-# ======================================================
-# API KEY
-# ======================================================
+# =====================================================
+# CONFIG
+# =====================================================
 
-API_KEY = "sk-123456"
+API_KEY = "sk-your-secret-key"
 
-# ======================================================
-# APP
-# ======================================================
+# =====================================================
+# FASTAPI
+# =====================================================
 
-app = FastAPI()
+app = FastAPI(
+    title="Universal AI Gateway",
+    version="2.1.0"
+)
+
+# =====================================================
+# CORS
+# =====================================================
 
 app.add_middleware(
     CORSMiddleware,
@@ -29,20 +36,21 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ======================================================
+# =====================================================
 # MODELS
-# ======================================================
+# =====================================================
 
 MODELS = [
     "gpt-4o-mini",
     "gpt-4",
     "claude-3-haiku",
     "llama-3.1-70b",
+    "mixtral-8x7b"
 ]
 
-# ======================================================
-# REQUEST
-# ======================================================
+# =====================================================
+# REQUEST MODELS
+# =====================================================
 
 class Message(BaseModel):
     role: str
@@ -52,29 +60,39 @@ class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
     stream: bool = False
+    temperature: float = 0.7
+    max_tokens: int = 4096
 
-# ======================================================
+# =====================================================
 # AUTH
-# ======================================================
+# =====================================================
 
 def verify_api_key(req: Request):
 
     auth = req.headers.get("Authorization")
 
     if not auth:
-        raise HTTPException(401, "Missing API Key")
+        return True
 
     if not auth.startswith("Bearer "):
-        raise HTTPException(401, "Invalid Authorization")
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid Authorization Format"
+        )
 
     token = auth.replace("Bearer ", "")
 
     if token != API_KEY:
-        raise HTTPException(403, "Invalid API Key")
+        raise HTTPException(
+            status_code=403,
+            detail="Invalid API Key"
+        )
+
+    return True
 
-# ======================================================
+# =====================================================
 # ROOT
-# ======================================================
+# =====================================================
 
 @app.get("/")
 async def root():
@@ -84,9 +102,9 @@ async def root():
         "provider": "Duck.ai"
     }
 
-# ======================================================
-# MODELS
-# ======================================================
+# =====================================================
+# MODELS ENDPOINT
+# =====================================================
 
 @app.get("/v1/models")
 async def models():
@@ -104,15 +122,16 @@ async def models():
         ]
     }
 
-# ======================================================
-# STREAM
-# ======================================================
+# =====================================================
+# STREAM GENERATOR
+# =====================================================
 
-async def stream_generator(model, messages):
+async def generate_stream(model, messages):
 
     try:
 
-        response = g4f.ChatCompletion.create(
+        response = await asyncio.to_thread(
+            g4f.ChatCompletion.create,
             model=model,
             provider=g4f.Provider.DuckDuckGo,
             messages=messages,
@@ -130,9 +149,11 @@ async def stream_generator(model, messages):
                     "model": model,
                     "choices": [
                         {
+                            "index": 0,
                             "delta": {
-                                "content": chunk
-                            }
+                                "content": str(chunk)
+                            },
+                            "finish_reason": None
                         }
                     ]
                 }
@@ -145,20 +166,24 @@ async def stream_generator(model, messages):
 
     except Exception as e:
 
-        error = {
+        error_payload = {
             "error": {
-                "message": str(e)
+                "message": str(e),
+                "type": "server_error"
             }
         }
 
-        yield f"data: {json.dumps(error)}\n\n"
+        yield f"data: {json.dumps(error_payload)}\n\n"
 
-# ======================================================
-# CHAT
-# ======================================================
+# =====================================================
+# CHAT COMPLETIONS
+# =====================================================
 
 @app.post("/v1/chat/completions")
-async def chat(req: Request, body: ChatRequest):
+async def chat_completions(
+    req: Request,
+    body: ChatRequest
+):
 
     verify_api_key(req)
 
@@ -170,27 +195,32 @@ async def chat(req: Request, body: ChatRequest):
         for m in body.messages
     ]
 
-    # ==================================================
+    # =================================================
     # STREAM
-    # ==================================================
+    # =================================================
 
     if body.stream:
 
         return StreamingResponse(
-            stream_generator(
+            generate_stream(
                 body.model,
                 messages
             ),
-            media_type="text/event-stream"
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "X-Accel-Buffering": "no"
+            }
         )
 
-    # ==================================================
-    # NORMAL
-    # ==================================================
+    # =================================================
+    # NORMAL RESPONSE
+    # =================================================
 
     try:
 
-        response = g4f.ChatCompletion.create(
+        response = await asyncio.to_thread(
+            g4f.ChatCompletion.create,
             model=body.model,
             provider=g4f.Provider.DuckDuckGo,
             messages=messages
@@ -203,10 +233,12 @@ async def chat(req: Request, body: ChatRequest):
             "model": body.model,
             "choices": [
                 {
+                    "index": 0,
                     "message": {
                         "role": "assistant",
-                        "content": response
-                    }
+                        "content": str(response)
+                    },
+                    "finish_reason": "stop"
                 }
             ]
         })
@@ -216,4 +248,18 @@ async def chat(req: Request, body: ChatRequest):
         raise HTTPException(
             status_code=500,
             detail=str(e)
-        )
+        )
+
+# =====================================================
+# RUN
+# =====================================================
+
+if __name__ == "__main__":
+
+    import uvicorn
+
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=7860
+    )