Create app.py

app.py

@@ -0,0 +1,188 @@
+import asyncio
+import json
+from concurrent.futures import ThreadPoolExecutor
+from typing import List
+
+import g4f
+
+from fastapi import FastAPI, HTTPException, Header
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+
+# =====================================
+# CONFIG
+# =====================================
+
+API_KEY = "sk-your-secret-key"
+
+executor = ThreadPoolExecutor(max_workers=5)
+
+# =====================================
+# APP
+# =====================================
+
+app = FastAPI(
+    title="AI Gateway"
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# =====================================
+# MODELS
+# =====================================
+
+class Message(BaseModel):
+    role: str
+    content: str
+
+class ChatRequest(BaseModel):
+    model: str = "gpt-4o-mini"
+    messages: List[Message]
+    stream: bool = False
+
+# =====================================
+# AUTH
+# =====================================
+
+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"
+        )
+
+    token = auth.replace("Bearer ", "")
+
+    if token != API_KEY:
+        raise HTTPException(
+            status_code=403,
+            detail="Invalid API Key"
+        )
+
+# =====================================
+# GENERATE
+# =====================================
+
+def generate(model, messages):
+
+    response = g4f.ChatCompletion.create(
+        model=model,
+        messages=messages
+    )
+
+    return response
+
+# =====================================
+# STREAM
+# =====================================
+
+async def stream_generate(model, messages):
+
+    loop = asyncio.get_event_loop()
+
+    response = await loop.run_in_executor(
+        executor,
+        lambda: g4f.ChatCompletion.create(
+            model=model,
+            messages=messages,
+            stream=True
+        )
+    )
+
+    for chunk in response:
+
+        if chunk:
+
+            payload = {
+                "choices": [
+                    {
+                        "delta": {
+                            "content": chunk
+                        }
+                    }
+                ]
+            }
+
+            yield f"data: {json.dumps(payload)}\\n\\n"
+
+    yield "data: [DONE]\\n\\n"
+
+# =====================================
+# HOME
+# =====================================
+
+@app.get("/")
+async def home():
+
+    return {
+        "status": "online"
+    }
+
+# =====================================
+# CHAT
+# =====================================
+
+@app.post("/v1/chat/completions")
+async def chat(
+    req: ChatRequest,
+    authorization: str = Header(None)
+):
+
+    verify_api_key(authorization)
+
+    messages = [
+        m.model_dump()
+        for m in req.messages
+    ]
+
+    # =========================
+    # STREAM
+    # =========================
+
+    if req.stream:
+
+        return StreamingResponse(
+            stream_generate(
+                req.model,
+                messages
+            ),
+            media_type="text/event-stream"
+        )
+
+    # =========================
+    # NORMAL
+    # =========================
+
+    loop = asyncio.get_event_loop()
+
+    response = await loop.run_in_executor(
+        executor,
+        lambda: generate(
+            req.model,
+            messages
+        )
+    )
+
+    return {
+        "choices": [
+            {
+                "message": {
+                    "role": "assistant",
+                    "content": response
+                }
+            }
+        ]
+    }