Update app.py

app.py

@@ -3,22 +3,21 @@ from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from typing import List
-import asyncio
+import g4f
 import json
 import uuid
 import time
+import asyncio
 
-import g4f
-
-# ==================================================
-# CONFIG
-# ==================================================
+# ======================================================
+# API KEY
+# ======================================================
 
 API_KEY = "sk-123456"
 
-# ==================================================
+# ======================================================
 # APP
-# ==================================================
+# ======================================================
 
 app = FastAPI()
 
@@ -30,33 +29,20 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ==================================================
+# ======================================================
 # MODELS
-# ==================================================
+# ======================================================
 
 MODELS = [
     "gpt-4o-mini",
     "gpt-4",
     "claude-3-haiku",
     "llama-3.1-70b",
-    "mixtral-8x7b"
 ]
 
-# ==================================================
-# PROVIDERS
-# ==================================================
-
-PROVIDERS = []
-
-if hasattr(g4f.Provider, "DuckDuckGo"):
-    PROVIDERS.append(g4f.Provider.DuckDuckGo)
-
-if hasattr(g4f.Provider, "Blackbox"):
-    PROVIDERS.append(g4f.Provider.Blackbox)
-
-# ==================================================
-# REQUEST MODELS
-# ==================================================
+# ======================================================
+# REQUEST
+# ======================================================
 
 class Message(BaseModel):
     role: str
@@ -67,42 +53,43 @@ class ChatRequest(BaseModel):
     messages: List[Message]
     stream: bool = False
 
-# ==================================================
+# ======================================================
 # AUTH
-# ==================================================
+# ======================================================
 
-def verify_key(req: Request):
+def verify_api_key(req: Request):
 
     auth = req.headers.get("Authorization")
 
     if not auth:
-        raise HTTPException(status_code=401, detail="Missing API Key")
+        raise HTTPException(401, "Missing API Key")
 
     if not auth.startswith("Bearer "):
-        raise HTTPException(status_code=401, detail="Invalid Auth")
+        raise HTTPException(401, "Invalid Authorization")
 
     token = auth.replace("Bearer ", "")
 
     if token != API_KEY:
-        raise HTTPException(status_code=403, detail="Invalid API Key")
+        raise HTTPException(403, "Invalid API Key")
 
-# ==================================================
+# ======================================================
 # ROOT
-# ==================================================
+# ======================================================
 
 @app.get("/")
 async def root():
 
     return {
-        "status": "online"
+        "status": "online",
+        "provider": "Duck.ai"
     }
 
-# ==================================================
+# ======================================================
 # MODELS
-# ==================================================
+# ======================================================
 
 @app.get("/v1/models")
-async def get_models():
+async def models():
 
     return {
         "object": "list",
@@ -111,79 +98,69 @@ async def get_models():
                 "id": model,
                 "object": "model",
                 "created": int(time.time()),
-                "owned_by": "g4f"
+                "owned_by": "duck.ai"
             }
             for model in MODELS
         ]
     }
 
-# ==================================================
+# ======================================================
 # STREAM
-# ==================================================
-
-async def stream_response(model, messages):
+# ======================================================
 
-    last_error = None
+async def stream_generator(model, messages):
 
-    for provider in PROVIDERS:
+    try:
 
-        try:
-
-            response = g4f.ChatCompletion.create(
-                model=model,
-                provider=provider,
-                messages=messages,
-                stream=True
-            )
+        response = g4f.ChatCompletion.create(
+            model=model,
+            provider=g4f.Provider.DuckDuckGo,
+            messages=messages,
+            stream=True
+        )
 
-            for chunk in response:
+        for chunk in response:
 
-                if chunk:
+            if chunk:
 
-                    payload = {
-                        "id": f"chatcmpl-{uuid.uuid4().hex}",
-                        "object": "chat.completion.chunk",
-                        "created": int(time.time()),
-                        "model": model,
-                        "choices": [
-                            {
-                                "delta": {
-                                    "content": chunk
-                                }
+                payload = {
+                    "id": f"chatcmpl-{uuid.uuid4().hex}",
+                    "object": "chat.completion.chunk",
+                    "created": int(time.time()),
+                    "model": model,
+                    "choices": [
+                        {
+                            "delta": {
+                                "content": chunk
                             }
-                        ]
-                    }
-
-                    yield f"data: {json.dumps(payload)}\n\n"
-
-                    await asyncio.sleep(0)
-
-            yield "data: [DONE]\n\n"
+                        }
+                    ]
+                }
 
-            return
+                yield f"data: {json.dumps(payload)}\n\n"
 
-        except Exception as e:
+                await asyncio.sleep(0)
 
-            last_error = str(e)
+        yield "data: [DONE]\n\n"
 
-            continue
+    except Exception as e:
 
-    error = {
-        "error": {
-            "message": last_error or "All providers failed"
+        error = {
+            "error": {
+                "message": str(e)
+            }
         }
-    }
 
-    yield f"data: {json.dumps(error)}\n\n"
+        yield f"data: {json.dumps(error)}\n\n"
 
-# ==================================================
+# ======================================================
 # CHAT
-# ==================================================
+# ======================================================
 
 @app.post("/v1/chat/completions")
 async def chat(req: Request, body: ChatRequest):
 
-    verify_key(req)
+    verify_api_key(req)
 
     messages = [
         {
@@ -193,55 +170,50 @@ async def chat(req: Request, body: ChatRequest):
         for m in body.messages
     ]
 
-    # ==========================================
+    # ==================================================
     # STREAM
-    # ==========================================
+    # ==================================================
 
     if body.stream:
 
         return StreamingResponse(
-            stream_response(body.model, messages),
+            stream_generator(
+                body.model,
+                messages
+            ),
             media_type="text/event-stream"
         )
 
-    # ==========================================
+    # ==================================================
     # NORMAL
-    # ==========================================
-
-    last_error = None
+    # ==================================================
 
-    for provider in PROVIDERS:
+    try:
 
-        try:
-
-            response = g4f.ChatCompletion.create(
-                model=body.model,
-                provider=provider,
-                messages=messages
-            )
+        response = g4f.ChatCompletion.create(
+            model=body.model,
+            provider=g4f.Provider.DuckDuckGo,
+            messages=messages
+        )
 
-            return JSONResponse({
-                "id": f"chatcmpl-{uuid.uuid4().hex}",
-                "object": "chat.completion",
-                "created": int(time.time()),
-                "model": body.model,
-                "choices": [
-                    {
-                        "message": {
-                            "role": "assistant",
-                            "content": response
-                        }
+        return JSONResponse({
+            "id": f"chatcmpl-{uuid.uuid4().hex}",
+            "object": "chat.completion",
+            "created": int(time.time()),
+            "model": body.model,
+            "choices": [
+                {
+                    "message": {
+                        "role": "assistant",
+                        "content": response
                     }
-                ]
-            })
-
-        except Exception as e:
-
-            last_error = str(e)
+                }
+            ]
+        })
 
-            continue
+    except Exception as e:
 
-    raise HTTPException(
-        status_code=500,
-        detail=last_error or "All providers failed"
-    )
+        raise HTTPException(
+            status_code=500,
+            detail=str(e)
+        )