Update app.py

app.py

@@ -16,21 +16,21 @@ from pydantic import BaseModel
 
 API_KEY = "sk-your-secret-key"
 
-PROVIDERS = [
-    g4f.Provider.DuckDuckGo,
-    g4f.Provider.Blackbox,
-]
-
 executor = ThreadPoolExecutor(max_workers=5)
 
 # =====================================
-# APP
+# FASTAPI
 # =====================================
 
 app = FastAPI(
-    title="HuggingFace AI Gateway"
+    title="AI Gateway",
+    version="2.0.0"
 )
 
+# =====================================
+# CORS
+# =====================================
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -52,15 +52,21 @@ class ChatRequest(BaseModel):
     stream: bool = False
 
 # =====================================
-# AUTH
+# API KEY
 # =====================================
 
 def verify_api_key(auth: str):
 
+    if not auth:
+        raise HTTPException(
+            status_code=401,
+            detail="Missing Authorization Header"
+        )
+
     if not auth.startswith("Bearer "):
         raise HTTPException(
             status_code=401,
-            detail="Invalid Authorization"
+            detail="Invalid Authorization Format"
         )
 
     token = auth.replace("Bearer ", "")
@@ -72,106 +78,92 @@ def verify_api_key(auth: str):
         )
 
 # =====================================
-# GENERATE
+# NORMAL RESPONSE
 # =====================================
 
-def generate(model, messages):
-
-    last_error = None
+def generate_response(model, messages):
 
-    for provider in PROVIDERS:
-
-        try:
-
-            response = g4f.ChatCompletion.create(
-                model=model,
-                provider=provider,
-                messages=messages
-            )
-
-            return response
+    try:
 
-        except Exception as e:
+        response = g4f.ChatCompletion.create(
+            model=model,
+            messages=messages
+        )
 
-            print(f"{provider.__name__}: {e}")
+        return response
 
-            last_error = e
+    except Exception as e:
 
-    raise Exception(str(last_error))
+        raise Exception(str(e))
 
 # =====================================
-# STREAM
+# STREAMING RESPONSE
 # =====================================
 
-async def stream_generate(model, messages):
+async def stream_response(model, messages):
 
     loop = asyncio.get_event_loop()
 
-    for provider in PROVIDERS:
-
-        try:
+    try:
 
-            response = await loop.run_in_executor(
-                executor,
-                lambda: g4f.ChatCompletion.create(
-                    model=model,
-                    provider=provider,
-                    messages=messages,
-                    stream=True
-                )
+        response = await loop.run_in_executor(
+            executor,
+            lambda: g4f.ChatCompletion.create(
+                model=model,
+                messages=messages,
+                stream=True
             )
+        )
 
-            for chunk in response:
+        for chunk in response:
 
-                if chunk:
+            if chunk:
 
-                    payload = {
-                        "choices": [
-                            {
-                                "delta": {
-                                    "content": chunk
-                                }
+                payload = {
+                    "choices": [
+                        {
+                            "delta": {
+                                "content": chunk
                             }
-                        ]
-                    }
-
-                    yield f"data: {json.dumps(payload)}\\n\\n"
-
-            yield "data: [DONE]\\n\\n"
+                        }
+                    ]
+                }
 
-            return
+                yield f"data: {json.dumps(payload)}\\n\\n"
 
-        except Exception as e:
+        yield "data: [DONE]\\n\\n"
 
-            print(f"Streaming Error: {e}")
+    except Exception as e:
 
-            continue
+        error_payload = {
+            "error": str(e)
+        }
 
-    yield f"data: {json.dumps({'error':'All providers failed'})}\\n\\n"
+        yield f"data: {json.dumps(error_payload)}\\n\\n"
 
 # =====================================
-# API
+# HOME
 # =====================================
 
 @app.get("/")
 async def home():
 
     return {
-        "status": "online"
+        "status": "online",
+        "service": "AI Gateway"
     }
 
+# =====================================
+# CHAT COMPLETIONS
+# =====================================
+
 @app.post("/v1/chat/completions")
 async def chat(
     req: ChatRequest,
     authorization: str = Header(None)
 ):
 
-    if not authorization:
-        raise HTTPException(
-            status_code=401,
-            detail="Missing Authorization"
-        )
-
+    # تحقق API Key
     verify_api_key(authorization)
 
     messages = [
@@ -180,13 +172,13 @@ async def chat(
     ]
 
     # =================================
-    # STREAM
+    # STREAMING
     # =================================
 
     if req.stream:
 
         return StreamingResponse(
-            stream_generate(
+            stream_response(
                 req.model,
                 messages
             ),
@@ -194,7 +186,7 @@ async def chat(
         )
 
     # =================================
-    # NORMAL
+    # NORMAL RESPONSE
     # =================================
 
     loop = asyncio.get_event_loop()
@@ -203,7 +195,7 @@ async def chat(
 
         response = await loop.run_in_executor(
             executor,
-            lambda: generate(
+            lambda: generate_response(
                 req.model,
                 messages
             )