Update app.py

app.py

@@ -1,34 +1,26 @@
 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 asyncio
 import json
-import time
 import uuid
+import time
 
 import g4f
-import g4f.models
 
-# =========================================================
+# ==================================================
 # CONFIG
-# =========================================================
+# ==================================================
 
-API_KEY = "sk-your-secret-key"
+API_KEY = "sk-123456"
 
-# =========================================================
+# ==================================================
 # APP
-# =========================================================
-
-app = FastAPI(
-    title="DuckAI Gateway",
-    version="3.0.0"
-)
+# ==================================================
 
-# =========================================================
-# CORS
-# =========================================================
+app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
@@ -38,20 +30,33 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# =========================================================
+# ==================================================
+# MODELS
+# ==================================================
+
+MODELS = [
+    "gpt-4o-mini",
+    "gpt-4",
+    "claude-3-haiku",
+    "llama-3.1-70b",
+    "mixtral-8x7b"
+]
+
+# ==================================================
 # PROVIDERS
-# =========================================================
+# ==================================================
 
-PROVIDERS = [
-    g4f.Provider.DuckDuckGo,
-    g4f.Provider.Blackbox,
-    g4f.Provider.Free2GPT,
-    g4f.Provider.Bing
-]
+PROVIDERS = []
 
-# =========================================================
-# MODELS
-# =========================================================
+if hasattr(g4f.Provider, "DuckDuckGo"):
+    PROVIDERS.append(g4f.Provider.DuckDuckGo)
+
+if hasattr(g4f.Provider, "Blackbox"):
+    PROVIDERS.append(g4f.Provider.Blackbox)
+
+# ==================================================
+# REQUEST MODELS
+# ==================================================
 
 class Message(BaseModel):
     role: str
@@ -61,128 +66,62 @@ 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):
+def verify_key(req: Request):
 
     auth = req.headers.get("Authorization")
 
     if not auth:
-        raise HTTPException(
-            status_code=401,
-            detail="Missing Authorization Header"
-        )
+        raise HTTPException(status_code=401, detail="Missing API Key")
 
     if not auth.startswith("Bearer "):
-        raise HTTPException(
-            status_code=401,
-            detail="Invalid Authorization Format"
-        )
+        raise HTTPException(status_code=401, detail="Invalid Auth")
 
     token = auth.replace("Bearer ", "")
 
     if token != API_KEY:
-        raise HTTPException(
-            status_code=403,
-            detail="Invalid API Key"
-        )
+        raise HTTPException(status_code=403, detail="Invalid API Key")
 
-# =========================================================
+# ==================================================
 # ROOT
-# =========================================================
+# ==================================================
 
 @app.get("/")
 async def root():
 
     return {
-        "status": "online",
-        "service": "DuckAI Gateway"
+        "status": "online"
     }
 
-# =========================================================
+# ==================================================
 # MODELS
-# =========================================================
+# ==================================================
 
 @app.get("/v1/models")
-async def models():
-
-    data = []
-
-    try:
-
-        all_models = list(g4f.models._all_models)
+async def get_models():
 
-        for model_name in all_models:
-
-            data.append({
-                "id": str(model_name),
+    return {
+        "object": "list",
+        "data": [
+            {
+                "id": model,
                 "object": "model",
                 "created": int(time.time()),
                 "owned_by": "g4f"
-            })
-
-    except Exception:
-
-        fallback = [
-            "gpt-4o-mini",
-            "gpt-4",
-            "claude-3-haiku",
-            "llama-3.1-70b",
-            "mixtral-8x7b"
+            }
+            for model in MODELS
         ]
-
-        for model_name in fallback:
-
-            data.append({
-                "id": model_name,
-                "object": "model",
-                "created": int(time.time()),
-                "owned_by": "g4f"
-            })
-
-    return {
-        "object": "list",
-        "data": data
     }
 
-# =========================================================
-# NORMAL RESPONSE
-# =========================================================
+# ==================================================
+# STREAM
+# ==================================================
 
-def generate_response(model, messages):
-
-    last_error = None
-
-    for provider in PROVIDERS:
-
-        try:
-
-            response = g4f.ChatCompletion.create(
-                model=model,
-                provider=provider,
-                messages=messages
-            )
-
-            return response
-
-        except Exception as e:
-
-            last_error = str(e)
-
-            continue
-
-    raise Exception(last_error or "All providers failed")
-
-# =========================================================
-# STREAM RESPONSE
-# =========================================================
-
-async def generate_stream(model, messages):
+async def stream_response(model, messages):
 
     last_error = None
 
@@ -208,11 +147,9 @@ async def generate_stream(model, messages):
                         "model": model,
                         "choices": [
                             {
-                                "index": 0,
                                 "delta": {
                                     "content": chunk
-                                },
-                                "finish_reason": None
+                                }
                             }
                         ]
                     }
@@ -221,22 +158,6 @@ async def generate_stream(model, messages):
 
                     await asyncio.sleep(0)
 
-            done_payload = {
-                "id": f"chatcmpl-{uuid.uuid4().hex}",
-                "object": "chat.completion.chunk",
-                "created": int(time.time()),
-                "model": model,
-                "choices": [
-                    {
-                        "index": 0,
-                        "delta": {},
-                        "finish_reason": "stop"
-                    }
-                ]
-            }
-
-            yield f"data: {json.dumps(done_payload)}\n\n"
-
             yield "data: [DONE]\n\n"
 
             return
@@ -247,26 +168,22 @@ async def generate_stream(model, messages):
 
             continue
 
-    error_payload = {
+    error = {
         "error": {
-            "message": last_error or "All providers failed",
-            "type": "server_error"
+            "message": last_error or "All providers failed"
         }
     }
 
-    yield f"data: {json.dumps(error_payload)}\n\n"
+    yield f"data: {json.dumps(error)}\n\n"
 
-# =========================================================
-# CHAT COMPLETIONS
-# =========================================================
+# ==================================================
+# CHAT
+# ==================================================
 
 @app.post("/v1/chat/completions")
-async def chat_completions(
-    req: Request,
-    body: ChatRequest
-):
+async def chat(req: Request, body: ChatRequest):
 
-    verify_api_key(req)
+    verify_key(req)
 
     messages = [
         {
@@ -276,56 +193,55 @@ async def chat_completions(
         for m in body.messages
     ]
 
-    # =====================================================
+    # ==========================================
     # STREAM
-    # =====================================================
+    # ==========================================
 
     if body.stream:
 
         return StreamingResponse(
-            generate_stream(
-                body.model,
-                messages
-            ),
+            stream_response(body.model, messages),
             media_type="text/event-stream"
         )
 
-    # =====================================================
+    # ==========================================
     # NORMAL
-    # =====================================================
+    # ==========================================
 
-    try:
+    last_error = None
 
-        response = generate_response(
-            body.model,
-            messages
-        )
+    for provider in PROVIDERS:
 
-        return JSONResponse({
-            "id": f"chatcmpl-{uuid.uuid4().hex}",
-            "object": "chat.completion",
-            "created": int(time.time()),
-            "model": body.model,
-            "choices": [
-                {
-                    "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": response
-                    },
-                    "finish_reason": "stop"
-                }
-            ],
-            "usage": {
-                "prompt_tokens": 0,
-                "completion_tokens": 0,
-                "total_tokens": 0
-            }
-        })
+        try:
 
-    except Exception as e:
+            response = g4f.ChatCompletion.create(
+                model=body.model,
+                provider=provider,
+                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
+                        }
+                    }
+                ]
+            })
+
+        except Exception as e:
+
+            last_error = str(e)
+
+            continue
 
-        raise HTTPException(
-            status_code=500,
-            detail=str(e)
-        )
+    raise HTTPException(
+        status_code=500,
+        detail=last_error or "All providers failed"
+    )