Update app.py

app.py

@@ -1,56 +1,75 @@
-import asyncio
-import json
-from concurrent.futures import ThreadPoolExecutor
-from typing import List
-
-import g4f
-
-from fastapi import FastAPI, HTTPException, Header
+from fastapi import FastAPI, Request, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
+from typing import List
+import json
+import uuid
+import time
+import asyncio
+import g4f
 
-# =====================================
+# =========================================
 # CONFIG
-# =====================================
+# =========================================
 
 API_KEY = "sk-your-secret-key"
 
-executor = ThreadPoolExecutor(max_workers=5)
-
-# =====================================
-# APP
-# =====================================
+# =========================================
+# FASTAPI
+# =========================================
 
 app = FastAPI(
-    title="AI Gateway"
+    title="Universal OpenAI API",
+    version="1.0.0"
 )
 
+# =========================================
+# CORS
+# =========================================
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
+    allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# =====================================
+# =========================================
 # MODELS
-# =====================================
+# =========================================
+
+AVAILABLE_MODELS = [
+    "gpt-4o-mini",
+    "gpt-4",
+    "claude-3-haiku",
+    "llama-3.1-70b",
+    "mixtral-8x7b"
+]
+
+# =========================================
+# Pydantic
+# =========================================
 
 class Message(BaseModel):
     role: str
     content: str
 
 class ChatRequest(BaseModel):
-    model: str = "gpt-4o-mini"
+    model: str
     messages: List[Message]
     stream: bool = False
+    temperature: float = 0.7
+    max_tokens: int = 4096
 
-# =====================================
+# =========================================
 # AUTH
-# =====================================
+# =========================================
 
-def verify_api_key(auth: str):
+def verify_api_key(req: Request):
+
+    auth = req.headers.get("Authorization")
 
     if not auth:
         raise HTTPException(
@@ -61,7 +80,7 @@ def verify_api_key(auth: str):
     if not auth.startswith("Bearer "):
         raise HTTPException(
             status_code=401,
-            detail="Invalid Authorization"
+            detail="Invalid Authorization Format"
         )
 
     token = auth.replace("Bearer ", "")
@@ -72,117 +91,171 @@ def verify_api_key(auth: str):
             detail="Invalid API Key"
         )
 
-# =====================================
-# GENERATE
-# =====================================
+# =========================================
+# ROOT
+# =========================================
 
-def generate(model, messages):
+@app.get("/")
+async def root():
 
-    response = g4f.ChatCompletion.create(
-        model=model,
-        messages=messages
-    )
+    return {
+        "status": "online"
+    }
+
+# =========================================
+# MODELS ENDPOINT
+# =========================================
 
-    return response
+@app.get("/v1/models")
+async def models():
 
-# =====================================
-# STREAM
-# =====================================
+    return {
+        "object": "list",
+        "data": [
+            {
+                "id": model,
+                "object": "model",
+                "created": int(time.time()),
+                "owned_by": "9router"
+            }
+            for model in AVAILABLE_MODELS
+        ]
+    }
 
-async def stream_generate(model, messages):
+# =========================================
+# CHAT COMPLETIONS
+# =========================================
 
-    loop = asyncio.get_event_loop()
+@app.post("/v1/chat/completions")
+async def chat_completions(
+    req: Request,
+    body: ChatRequest
+):
 
-    response = await loop.run_in_executor(
-        executor,
-        lambda: g4f.ChatCompletion.create(
-            model=model,
-            messages=messages,
-            stream=True
-        )
-    )
+    verify_api_key(req)
 
-    for chunk in response:
+    messages = [
+        {
+            "role": m.role,
+            "content": m.content
+        }
+        for m in body.messages
+    ]
 
-        if chunk:
+    provider = g4f.Provider.Blackbox
 
-            payload = {
-                "choices": [
-                    {
-                        "delta": {
-                            "content": chunk
-                        }
-                    }
-                ]
-            }
+    # =====================================
+    # STREAMING
+    # =====================================
 
-            yield f"data: {json.dumps(payload)}\\n\\n"
+    if body.stream:
 
-    yield "data: [DONE]\\n\\n"
+        async def generate_stream():
 
-# =====================================
-# HOME
-# =====================================
+            try:
 
-@app.get("/")
-async def home():
+                response = g4f.ChatCompletion.create(
+                    model=body.model,
+                    provider=provider,
+                    messages=messages,
+                    stream=True
+                )
 
-    return {
-        "status": "online"
-    }
+                for chunk in response:
 
-# =====================================
-# CHAT
-# =====================================
+                    if chunk:
 
-@app.post("/v1/chat/completions")
-async def chat(
-    req: ChatRequest,
-    authorization: str = Header(None)
-):
+                        payload = {
+                            "id": f"chatcmpl-{uuid.uuid4().hex}",
+                            "object": "chat.completion.chunk",
+                            "created": int(time.time()),
+                            "model": body.model,
+                            "choices": [
+                                {
+                                    "index": 0,
+                                    "delta": {
+                                        "content": chunk
+                                    },
+                                    "finish_reason": None
+                                }
+                            ]
+                        }
 
-    verify_api_key(authorization)
+                        yield f"data: {json.dumps(payload)}\n\n"
 
-    messages = [
-        m.model_dump()
-        for m in req.messages
-    ]
+                        await asyncio.sleep(0)
 
-    # =========================
-    # STREAM
-    # =========================
+                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"
+
+            except Exception as e:
+
+                error_payload = {
+                    "error": {
+                        "message": str(e),
+                        "type": "server_error"
+                    }
+                }
 
-    if req.stream:
+                yield f"data: {json.dumps(error_payload)}\n\n"
 
         return StreamingResponse(
-            stream_generate(
-                req.model,
-                messages
-            ),
+            generate_stream(),
             media_type="text/event-stream"
         )
 
-    # =========================
-    # NORMAL
-    # =========================
+    # =====================================
+    # NORMAL RESPONSE
+    # =====================================
 
-    loop = asyncio.get_event_loop()
+    try:
 
-    response = await loop.run_in_executor(
-        executor,
-        lambda: generate(
-            req.model,
-            messages
+        response = g4f.ChatCompletion.create(
+            model=body.model,
+            provider=provider,
+            messages=messages
         )
-    )
 
-    return {
-        "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": [
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": response
+                    },
+                    "finish_reason": "stop"
                 }
+            ],
+            "usage": {
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "total_tokens": 0
             }
-        ]
-    }
+        })
+
+    except Exception as e:
+
+        raise HTTPException(
+            status_code=500,
+            detail=str(e)
+        )