Update app.py

app.py

@@ -3,11 +3,13 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
 from typing import List
-import g4f
+import asyncio
 import json
-import uuid
 import time
-import asyncio
+import uuid
+
+import g4f
+import g4f.models
 
 # =====================================================
 # CONFIG
@@ -21,7 +23,7 @@ API_KEY = "sk-your-secret-key"
 
 app = FastAPI(
     title="Universal AI Gateway",
-    version="2.1.0"
+    version="3.0.0"
 )
 
 # =====================================================
@@ -40,18 +42,6 @@ app.add_middleware(
 # MODELS
 # =====================================================
 
-MODELS = [
-    "gpt-4o-mini",
-    "gpt-4",
-    "claude-3-haiku",
-    "llama-3.1-70b",
-    "mixtral-8x7b"
-]
-
-# =====================================================
-# REQUEST MODELS
-# =====================================================
-
 class Message(BaseModel):
     role: str
     content: str
@@ -72,7 +62,10 @@ def verify_api_key(req: Request):
     auth = req.headers.get("Authorization")
 
     if not auth:
-        return True
+        raise HTTPException(
+            status_code=401,
+            detail="Missing Authorization Header"
+        )
 
     if not auth.startswith("Bearer "):
         raise HTTPException(
@@ -88,7 +81,24 @@ def verify_api_key(req: Request):
             detail="Invalid API Key"
         )
 
-    return True
+# =====================================================
+# PROVIDERS
+# =====================================================
+
+def get_providers():
+
+    providers = []
+
+    if hasattr(g4f.Provider, "DuckDuckGo"):
+        providers.append(g4f.Provider.DuckDuckGo)
+
+    if hasattr(g4f.Provider, "Blackbox"):
+        providers.append(g4f.Provider.Blackbox)
+
+    if hasattr(g4f.Provider, "Free2GPT"):
+        providers.append(g4f.Provider.Free2GPT)
+
+    return providers
 
 # =====================================================
 # ROOT
@@ -99,81 +109,55 @@ async def root():
 
     return {
         "status": "online",
-        "provider": "Duck.ai"
+        "service": "Universal AI Gateway",
+        "version": "3.0.0"
     }
 
 # =====================================================
-# MODELS ENDPOINT
+# MODELS
 # =====================================================
 
 @app.get("/v1/models")
 async def models():
 
-    return {
-        "object": "list",
-        "data": [
-            {
-                "id": model,
-                "object": "model",
-                "created": int(time.time()),
-                "owned_by": "duck.ai"
-            }
-            for model in MODELS
-        ]
-    }
-
-# =====================================================
-# STREAM GENERATOR
-# =====================================================
-
-async def generate_stream(model, messages):
+    data = []
 
     try:
 
-        response = await asyncio.to_thread(
-            g4f.ChatCompletion.create,
-            model=model,
-            provider=g4f.Provider.DuckDuckGo,
-            messages=messages,
-            stream=True
-        )
+        all_models = list(g4f.models._all_models)
 
-        for chunk in response:
-
-            if chunk:
-
-                payload = {
-                    "id": f"chatcmpl-{uuid.uuid4().hex}",
-                    "object": "chat.completion.chunk",
-                    "created": int(time.time()),
-                    "model": model,
-                    "choices": [
-                        {
-                            "index": 0,
-                            "delta": {
-                                "content": str(chunk)
-                            },
-                            "finish_reason": None
-                        }
-                    ]
-                }
+        for model_name in all_models:
 
-                yield f"data: {json.dumps(payload)}\n\n"
+            data.append({
+                "id": str(model_name),
+                "object": "model",
+                "created": int(time.time()),
+                "owned_by": "g4f"
+            })
 
-                await asyncio.sleep(0)
+    except Exception:
 
-        yield "data: [DONE]\n\n"
+        fallback = [
+            "gpt-4o-mini",
+            "gpt-4",
+            "claude-3-haiku",
+            "llama-3.1-70b",
+            "mixtral-8x7b"
+        ]
 
-    except Exception as e:
+        for model_name in fallback:
 
-        error_payload = {
-            "error": {
-                "message": str(e),
-                "type": "server_error"
-            }
-        }
+            data.append({
+                "id": model_name,
+                "object": "model",
+                "created": int(time.time()),
+                "owned_by": "g4f"
+            })
 
-        yield f"data: {json.dumps(error_payload)}\n\n"
+    return {
+        "object": "list",
+        "data": data
+    }
 
 # =====================================================
 # CHAT COMPLETIONS
@@ -196,19 +180,94 @@ async def chat_completions(
     ]
 
     # =================================================
-    # STREAM
+    # STREAMING
     # =================================================
 
     if body.stream:
 
+        async def generate_stream():
+
+            providers = get_providers()
+
+            last_error = None
+
+            for provider in providers:
+
+                try:
+
+                    response = await asyncio.to_thread(
+                        g4f.ChatCompletion.create,
+                        model=body.model,
+                        provider=provider,
+                        messages=messages,
+                        stream=True
+                    )
+
+                    for chunk in response:
+
+                        if chunk:
+
+                            payload = {
+                                "id": f"chatcmpl-{uuid.uuid4().hex}",
+                                "object": "chat.completion.chunk",
+                                "created": int(time.time()),
+                                "model": body.model,
+                                "choices": [
+                                    {
+                                        "index": 0,
+                                        "delta": {
+                                            "content": str(chunk)
+                                        },
+                                        "finish_reason": None
+                                    }
+                                ]
+                            }
+
+                            yield f"data: {json.dumps(payload)}\n\n"
+
+                            await asyncio.sleep(0)
+
+                    done_payload = {
+                        "id": f"chatcmpl-{uuid.uuid4().hex}",
+                        "object": "chat.completion.chunk",
+                        "created": int(time.time()),
+                        "model": body.model,
+                        "choices": [
+                            {
+                                "index": 0,
+                                "delta": {},
+                                "finish_reason": "stop"
+                            }
+                        ]
+                    }
+
+                    yield f"data: {json.dumps(done_payload)}\n\n"
+
+                    yield "data: [DONE]\n\n"
+
+                    return
+
+                except Exception as e:
+
+                    last_error = str(e)
+
+                    continue
+
+            error_payload = {
+                "error": {
+                    "message": last_error or "All providers failed",
+                    "type": "server_error"
+                }
+            }
+
+            yield f"data: {json.dumps(error_payload)}\n\n"
+
         return StreamingResponse(
-            generate_stream(
-                body.model,
-                messages
-            ),
+            generate_stream(),
             media_type="text/event-stream",
             headers={
                 "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
                 "X-Accel-Buffering": "no"
             }
         )
@@ -217,38 +276,53 @@ async def chat_completions(
     # NORMAL RESPONSE
     # =================================================
 
-    try:
+    providers = get_providers()
 
-        response = await asyncio.to_thread(
-            g4f.ChatCompletion.create,
-            model=body.model,
-            provider=g4f.Provider.DuckDuckGo,
-            messages=messages
-        )
+    last_error = None
+
+    for provider in providers:
+
+        try:
 
-        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": str(response)
-                    },
-                    "finish_reason": "stop"
+            response = await asyncio.to_thread(
+                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": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": str(response)
+                        },
+                        "finish_reason": "stop"
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": 0,
+                    "completion_tokens": 0,
+                    "total_tokens": 0
                 }
-            ]
-        })
+            })
 
-    except Exception as e:
+        except Exception as e:
 
-        raise HTTPException(
-            status_code=500,
-            detail=str(e)
-        )
+            last_error = str(e)
+
+            continue
+
+    raise HTTPException(
+        status_code=500,
+        detail=last_error or "All providers failed"
+    )
 
 # =====================================================
 # RUN