Update app.py

app.py

@@ -1,39 +1,316 @@
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import StreamingResponse, JSONResponse
-from pydantic import BaseModel
-from typing import List, Optional
+"""
+Universal AI Gateway v5.0
+━━━━━━━━━━━━━━━━━━━━━━━━━
+- OpenAI-compatible API
+- g4f multi-provider routing
+- Duck.ai integration (GPT-4o-mini, Claude, Llama, Mixtral)
+- Auto-fallback between providers
+- Streaming + non-streaming
+"""
+
+from __future__ import annotations
 
 import asyncio
 import json
-import uuid
-import time
 import logging
+import time
+import uuid
+from typing import AsyncIterator, List, Optional, Dict, Any
+
+import aiohttp
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, StreamingResponse
+from pydantic import BaseModel
 
 import g4f
-from g4f.client import Client
+from g4f.client import AsyncClient
+from g4f import Provider
 
-# =====================================================
+# ─────────────────────────────────────────────────────────────────────────────
 # LOGGING
-# =====================================================
+# ─────────────────────────────────────────────────────────────────────────────
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+)
+logger = logging.getLogger("ai-gateway")
+
+# ─────────────────────────────────────────────────────────────────────────────
+# CONFIG
+# ─────────────────────────────────────────────────────────────────────────────
+
+API_KEY = "sk-your-secret-key"
+
+# ─────────────────────────────────────────────────────────────────────────────
+# DUCK.AI INTEGRATION
+# Direct integration with https://duck.ai (DuckDuckGo AI Chat)
+# Supported models: gpt-4o-mini, claude-3-haiku, llama-3.3-70b, mixtral-8x7b, o4-mini
+# ─────────────────────────────────────────────────────────────────────────────
+
+DUCK_MODEL_MAP: Dict[str, str] = {
+    # Public names  →  Duck.ai internal model IDs
+    "gpt-4o-mini":          "gpt-4o-mini",
+    "claude-3-haiku":       "claude-3-haiku-20240307",
+    "claude-haiku":         "claude-3-haiku-20240307",
+    "llama-3.3-70b":        "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
+    "llama-3.1-70b":        "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
+    "llama":                "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
+    "mixtral-8x7b":         "mistralai/Mixtral-8x7B-Instruct-v0.1",
+    "mixtral":              "mistralai/Mixtral-8x7B-Instruct-v0.1",
+    "o4-mini":              "o4-mini",
+    "o3-mini":              "o4-mini",
+}
+
+DUCK_MODELS_LIST = list(dict.fromkeys(DUCK_MODEL_MAP.keys()))
+
+DUCK_VQD_URL   = "https://duckduckgo.com/duckchat/v1/status"
+DUCK_CHAT_URL  = "https://duckduckgo.com/duckchat/v1/chat"
+
+DUCK_HEADERS_BASE = {
+    "User-Agent": (
+        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+        "AppleWebKit/537.36 (KHTML, like Gecko) "
+        "Chrome/124.0.0.0 Safari/537.36"
+    ),
+    "Accept":           "text/event-stream",
+    "Accept-Language":  "en-US,en;q=0.9",
+    "Referer":          "https://duckduckgo.com/",
+    "Origin":           "https://duckduckgo.com",
+}
+
+
+async def duck_get_vqd(model_id: str, session: aiohttp.ClientSession) -> str:
+    """Fetch the x-vqd-4 token required by Duck.ai."""
+    headers = {**DUCK_HEADERS_BASE, "x-vqd-accept": "1"}
+    async with session.get(DUCK_VQD_URL, headers=headers) as resp:
+        if resp.status != 200:
+            raise RuntimeError(f"Duck.ai VQD fetch failed: HTTP {resp.status}")
+        vqd = resp.headers.get("x-vqd-4", "")
+        if not vqd:
+            raise RuntimeError("Duck.ai did not return x-vqd-4 token")
+        return vqd
+
+
+async def duck_chat_stream(
+    messages: List[Dict[str, str]],
+    model: str,
+) -> AsyncIterator[str]:
+    """
+    Async generator that yields text chunks from Duck.ai.
+    Handles VQD token refresh automatically.
+    """
+    duck_model = DUCK_MODEL_MAP.get(model, "gpt-4o-mini")
+
+    async with aiohttp.ClientSession() as session:
+        vqd = await duck_get_vqd(duck_model, session)
+
+        payload = {
+            "model":    duck_model,
+            "messages": messages,
+        }
+
+        headers = {
+            **DUCK_HEADERS_BASE,
+            "Content-Type": "application/json",
+            "x-vqd-4":      vqd,
+        }
+
+        async with session.post(
+            DUCK_CHAT_URL,
+            json=payload,
+            headers=headers,
+        ) as resp:
+            if resp.status != 200:
+                body = await resp.text()
+                raise RuntimeError(
+                    f"Duck.ai chat failed: HTTP {resp.status} — {body[:200]}"
+                )
 
-logging.basicConfig(level=logging.INFO)
+            async for raw_line in resp.content:
+                line = raw_line.decode("utf-8", errors="replace").strip()
+                if not line or not line.startswith("data:"):
+                    continue
+                data_str = line[len("data:"):].strip()
+                if data_str == "[DONE]":
+                    break
+                try:
+                    data = json.loads(data_str)
+                    chunk = data.get("message", "")
+                    if chunk:
+                        yield chunk
+                except json.JSONDecodeError:
+                    continue
+
+
+async def duck_chat_complete(
+    messages: List[Dict[str, str]],
+    model: str,
+) -> str:
+    """Collect full Duck.ai response (non-streaming)."""
+    parts: List[str] = []
+    async for chunk in duck_chat_stream(messages, model):
+        parts.append(chunk)
+    return "".join(parts)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# G4F PROVIDER ROUTING
+# Maps model names → preferred g4f provider with fallback chain
+# ─────────────────────────────────────────────────────────────────────────────
+
+G4F_PROVIDER_MAP: Dict[str, Any] = {
+    # OpenAI GPT family
+    "gpt-4":                Provider.OpenaiChat,
+    "gpt-4o":               Provider.OpenaiChat,
+    "gpt-4o-mini":          Provider.OpenaiChat,
+    "gpt-4.1":              Provider.OpenaiChat,
+    "gpt-4.1-mini":         Provider.OpenaiChat,
+    "gpt-5":                Provider.OpenaiChat,
+    "auto":                 Provider.OpenaiChat,
+
+    # Anthropic Claude
+    "claude-3-haiku":       Provider.Anthropic,
+    "claude-3-sonnet":      Provider.Anthropic,
+    "claude-3-opus":        Provider.Anthropic,
+    "claude-3-5-sonnet":    Provider.Anthropic,
+    "claude-3-7-sonnet":    Provider.Anthropic,
+    "claude-sonnet-4":      Provider.Anthropic,
+
+    # Google Gemini
+    "gemini-pro":           Provider.GeminiPro,
+    "gemini-1.5-pro":       Provider.GeminiPro,
+    "gemini-2.5-flash":     Provider.Gemini,
+    "gemini-2.5-pro":       Provider.Gemini,
+
+    # Meta Llama
+    "llama-3.1-70b":        Provider.Cerebras,
+    "llama-3.3-70b":        Provider.Cerebras,
+    "llama-3.1-8b":         Provider.Cerebras,
+
+    # Mistral / Mixtral
+    "mixtral-8x7b":         Provider.HuggingFace,
+    "mistral-7b":           Provider.HuggingFace,
+
+    # DeepSeek
+    "deepseek-chat":        Provider.DeepSeek,
+    "deepseek-r1":          Provider.DeepSeek,
+
+    # Qwen
+    "qwen-2.5":             Provider.Qwen,
+    "qwen-3":               Provider.Qwen,
+    "qwen-2-72b":           Provider.Qwen,
+
+    # Grok
+    "grok-3":               Provider.Grok,
+    "grok-4":               Provider.Grok,
+
+    # Copilot
+    "copilot":              Provider.Copilot,
+    "o1":                   Provider.Copilot,
+    "o3-mini":              Provider.Copilot,
+    "o4-mini":              Provider.Copilot,
+
+    # Blackbox (multi-model)
+    "blackbox":             Provider.BlackboxPro,
+    "openai/gpt-5":         Provider.BlackboxPro,
+    "x-ai/grok-4":          Provider.BlackboxPro,
+
+    # Perplexity
+    "perplexity":           Provider.Perplexity,
+    "perplexity-turbo":     Provider.Perplexity,
+
+    # Pi
+    "pi":                   Provider.Pi,
+}
+
+
+def get_g4f_provider(model: str) -> Optional[Any]:
+    """Return the best g4f provider for a model name."""
+    # Exact match
+    if model in G4F_PROVIDER_MAP:
+        return G4F_PROVIDER_MAP[model]
+    # Prefix match (e.g. "gpt-4o-mini-2024" → gpt-4o-mini)
+    for key, provider in G4F_PROVIDER_MAP.items():
+        if model.startswith(key) or key.startswith(model):
+            return provider
+    return None
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# MODELS REGISTRY
+# ─────────────────────────────────────────────────────────────────────────────
+
+def build_models_list() -> List[Dict]:
+    """Build the /v1/models response combining Duck.ai + g4f providers."""
+    now = int(time.time())
+    seen: set = set()
+    models: List[Dict] = []
+
+    # Duck.ai models
+    for m in DUCK_MODELS_LIST:
+        if m not in seen:
+            seen.add(m)
+            models.append({
+                "id":         m,
+                "object":     "model",
+                "created":    now,
+                "owned_by":   "duck-ai",
+                "provider":   "duck.ai",
+            })
 
-logger = logging.getLogger(__name__)
+    # g4f provider models
+    for model_name, provider in G4F_PROVIDER_MAP.items():
+        if model_name not in seen:
+            seen.add(model_name)
+            pname = getattr(provider, "__name__", str(provider))
+            models.append({
+                "id":         model_name,
+                "object":     "model",
+                "created":    now,
+                "owned_by":   "g4f",
+                "provider":   pname,
+            })
 
-# =====================================================
-# APP
-# =====================================================
+    # Extra g4f working providers
+    for pname in dir(Provider):
+        if pname.startswith("_") or pname[0].islower():
+            continue
+        try:
+            p = getattr(Provider, pname)
+            if not (hasattr(p, "working") and p.working):
+                continue
+            pmodels = getattr(p, "models", None)
+            if not pmodels or callable(pmodels):
+                continue
+            for m in pmodels:
+                m = str(m)
+                if m and m not in seen:
+                    seen.add(m)
+                    models.append({
+                        "id":         m,
+                        "object":     "model",
+                        "created":    now,
+                        "owned_by":   "g4f",
+                        "provider":   pname,
+                    })
+        except Exception:
+            continue
+
+    return models
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# FASTAPI APP
+# ─────────────────────────────────────────────────────────────────────────────
 
 app = FastAPI(
-    title="DuckAI Gateway",
-    version="6.0.0"
+    title="Universal AI Gateway",
+    version="5.0.0",
+    description="OpenAI-compatible gateway with Duck.ai + g4f multi-provider support",
 )
 
-# =====================================================
-# CORS
-# =====================================================
-
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -42,259 +319,237 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# =====================================================
-# MODELS
-# =====================================================
+# ─────────────────────────────────────────────────────────────────────────────
+# SCHEMAS
+# ─────────────────────────────────────────────────────────────────────────────
 
 class Message(BaseModel):
     role: str
     content: str
 
+
 class ChatRequest(BaseModel):
-    model: str
+    model: str = "gpt-4o-mini"
     messages: List[Message]
-    stream: Optional[bool] = False
+    stream: bool = False
     temperature: Optional[float] = 0.7
     max_tokens: Optional[int] = 4096
 
-# =====================================================
-# CLIENT
-# =====================================================
-
-client = Client()
-
-# =====================================================
-# ROOT
-# =====================================================
-
-@app.get("/")
-async def root():
 
-    return {
-        "status": "online",
-        "service": "DuckAI Gateway"
+# ─────────────────────────────────────────────────────────────────────────────
+# AUTH
+# ─────────────────────────────────────────────────────────────────────────────
+
+def verify_api_key(req: Request) -> bool:
+    auth = req.headers.get("Authorization", "")
+    if not auth:
+        return True  # Allow unauthenticated for testing
+    if not auth.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Invalid Authorization format")
+    token = auth.removeprefix("Bearer ").strip()
+    if token != API_KEY:
+        raise HTTPException(status_code=403, detail="Invalid API key")
+    return True
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# SSE HELPERS
+# ─────────────────────────────────────────────────────────────────────────────
+
+def sse_chunk(chunk_id: str, model: str, content: str) -> str:
+    payload = {
+        "id":      chunk_id,
+        "object":  "chat.completion.chunk",
+        "created": int(time.time()),
+        "model":   model,
+        "choices": [{
+            "index":         0,
+            "delta":         {"content": content},
+            "finish_reason": None,
+        }],
     }
+    return f"data: {json.dumps(payload)}\n\n"
+
+
+def sse_done(chunk_id: str, model: str) -> str:
+    payload = {
+        "id":      chunk_id,
+        "object":  "chat.completion.chunk",
+        "created": int(time.time()),
+        "model":   model,
+        "choices": [{
+            "index":         0,
+            "delta":         {},
+            "finish_reason": "stop",
+        }],
+    }
+    return f"data: {json.dumps(payload)}\n\ndata: [DONE]\n\n"
 
-# =====================================================
-# MODELS
-# =====================================================
-
-@app.get("/v1/models")
-async def models():
-
-    output = []
-
-    try:
-
-        all_models = list(g4f.models._all_models)
-
-        for model in all_models:
-
-            output.append({
-                "id": str(model),
-                "object": "model",
-                "created": int(time.time()),
-                "owned_by": "g4f"
-            })
-
-    except Exception:
-
-        fallback = [
-            "gpt-4o-mini",
-            "gpt-4o",
-            "gpt-4",
-            "claude-3-haiku",
-            "gemini-pro"
-        ]
-
-        for model in fallback:
-
-            output.append({
-                "id": model,
-                "object": "model",
-                "created": int(time.time()),
-                "owned_by": "g4f"
-            })
 
+def build_response(chunk_id: str, model: str, content: str) -> Dict:
     return {
-        "object": "list",
-        "data": output
+        "id":      chunk_id,
+        "object":  "chat.completion",
+        "created": int(time.time()),
+        "model":   model,
+        "choices": [{
+            "index":         0,
+            "message":       {"role": "assistant", "content": content},
+            "finish_reason": "stop",
+        }],
+        "usage": {
+            "prompt_tokens":     0,
+            "completion_tokens": 0,
+            "total_tokens":      0,
+        },
     }
 
-# =====================================================
-# CHAT
-# =====================================================
 
-@app.post("/v1/chat/completions")
-async def chat(body: ChatRequest):
+# ─────────────────────────────────────────────────────────────────────────────
+# CORE ROUTING LOGIC
+# Priority: Duck.ai → specific g4f provider → g4f auto-routing → error
+# ─────────────────────────────────────────────────────────────────────────────
 
-    messages = [
-        {
-            "role": m.role,
-            "content": m.content
-        }
-        for m in body.messages
-    ]
+def is_duck_model(model: str) -> bool:
+    return model in DUCK_MODEL_MAP
 
-    # =================================================
-    # STREAM
-    # =================================================
 
-    if body.stream:
+async def route_stream(
+    messages: List[Dict],
+    model: str,
+    chunk_id: str,
+) -> AsyncIterator[str]:
+    """Unified streaming router."""
 
-        async def event_stream():
+    # ── Duck.ai ──────────────────────────────────────────────────────────────
+    if is_duck_model(model):
+        logger.info(f"[STREAM] Duck.ai → model={model}")
+        try:
+            async for chunk in duck_chat_stream(messages, model):
+                yield sse_chunk(chunk_id, model, chunk)
+            yield sse_done(chunk_id, model)
+            return
+        except Exception as e:
+            logger.warning(f"[STREAM] Duck.ai failed ({e}), trying g4f fallback")
+
+    # ── g4f provider ─────────────────────────────────────────────────────────
+    provider = get_g4f_provider(model)
+    provider_name = getattr(provider, "__name__", "auto") if provider else "auto"
+    logger.info(f"[STREAM] g4f provider={provider_name} model={model}")
 
+    try:
+        client = AsyncClient(provider=provider)
+        response = await client.chat.completions.create(
+            model=model,
+            messages=messages,
+            stream=True,
+        )
+        async for chunk in response:
             try:
+                content = chunk.choices[0].delta.content or ""
+                if content:
+                    yield sse_chunk(chunk_id, model, content)
+            except Exception:
+                continue
+        yield sse_done(chunk_id, model)
 
-                response = await asyncio.wait_for(
-                    asyncio.to_thread(
-                        client.chat.completions.create,
-                        model=body.model,
-                        messages=messages,
-                        stream=True
-                    ),
-                    timeout=60
-                )
-
-                chunk_id = f"chatcmpl-{uuid.uuid4().hex}"
-
-                for chunk in response:
-
-                    try:
-
-                        content = ""
-
-                        if (
-                            hasattr(chunk, "choices")
-                            and chunk.choices
-                            and chunk.choices[0].delta
-                        ):
-                            content = chunk.choices[0].delta.content
+    except Exception as e:
+        logger.error(f"[STREAM] g4f error: {e}")
+        error_payload = {"error": {"message": str(e), "type": "server_error"}}
+        yield f"data: {json.dumps(error_payload)}\n\n"
 
-                        if content:
 
-                            payload = {
-                                "id": chunk_id,
-                                "object": "chat.completion.chunk",
-                                "created": int(time.time()),
-                                "model": body.model,
-                                "choices": [
-                                    {
-                                        "index": 0,
-                                        "delta": {
-                                            "content": content
-                                        },
-                                        "finish_reason": None
-                                    }
-                                ]
-                            }
+async def route_complete(
+    messages: List[Dict],
+    model: str,
+) -> str:
+    """Unified non-streaming router."""
 
-                            yield f"data: {json.dumps(payload)}\n\n"
+    # ── Duck.ai ──────────────────────────────────────────────────────────────
+    if is_duck_model(model):
+        logger.info(f"[COMPLETE] Duck.ai → model={model}")
+        try:
+            return await duck_chat_complete(messages, model)
+        except Exception as e:
+            logger.warning(f"[COMPLETE] Duck.ai failed ({e}), trying g4f fallback")
+
+    # ── g4f provider ─────────────────────────────────────────────────────────
+    provider = get_g4f_provider(model)
+    provider_name = getattr(provider, "__name__", "auto") if provider else "auto"
+    logger.info(f"[COMPLETE] g4f provider={provider_name} model={model}")
+
+    client = AsyncClient(provider=provider)
+    response = await client.chat.completions.create(
+        model=model,
+        messages=messages,
+        stream=False,
+    )
+    try:
+        return response.choices[0].message.content or ""
+    except Exception:
+        return str(response)
 
-                            await asyncio.sleep(0)
 
-                    except Exception as e:
-                        logger.error(e)
+# ─────────────────────────────────────────────────────────────────────────────
+# ROUTES
+# ─────────────────────────────────────────────────────────────────────────────
 
-                done_payload = {
-                    "id": chunk_id,
-                    "object": "chat.completion.chunk",
-                    "created": int(time.time()),
-                    "model": body.model,
-                    "choices": [
-                        {
-                            "index": 0,
-                            "delta": {},
-                            "finish_reason": "stop"
-                        }
-                    ]
-                }
+@app.get("/")
+async def root():
+    return {
+        "status":  "online",
+        "service": "Universal AI Gateway",
+        "version": "5.0.0",
+        "docs":    "/docs",
+        "models":  "/v1/models",
+    }
 
-                yield f"data: {json.dumps(done_payload)}\n\n"
 
-                yield "data: [DONE]\n\n"
+@app.get("/v1/models")
+async def get_models():
+    return {
+        "object": "list",
+        "data":   build_models_list(),
+    }
 
-            except Exception as e:
 
-                logger.error(e)
+@app.post("/v1/chat/completions")
+async def chat_completions(req: Request, body: ChatRequest):
+    verify_api_key(req)
 
-                payload = {
-                    "error": {
-                        "message": str(e)
-                    }
-                }
+    messages = [{"role": m.role, "content": m.content} for m in body.messages]
+    chunk_id = f"chatcmpl-{uuid.uuid4().hex}"
 
-                yield f"data: {json.dumps(payload)}\n\n"
+    logger.info(
+        f"Request: model={body.model!r} stream={body.stream} "
+        f"messages={len(messages)}"
+    )
 
+    # ── Streaming ─────────────────────────────────────────────────────────
+    if body.stream:
         return StreamingResponse(
-            event_stream(),
+            route_stream(messages, body.model, chunk_id),
             media_type="text/event-stream",
             headers={
-                "Cache-Control": "no-cache",
-                "Connection": "keep-alive"
-            }
+                "Cache-Control":    "no-cache",
+                "Connection":       "keep-alive",
+                "X-Accel-Buffering": "no",
+            },
         )
 
-    # =================================================
-    # NORMAL
-    # =================================================
-
+    # ── Non-streaming ─────────────────────────────────────────────────────
     try:
-
-        response = await asyncio.wait_for(
-            asyncio.to_thread(
-                client.chat.completions.create,
-                model=body.model,
-                messages=messages
-            ),
-            timeout=60
-        )
-
-        text = ""
-
-        try:
-            text = response.choices[0].message.content
-        except:
-            text = str(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": text
-                    },
-                    "finish_reason": "stop"
-                }
-            ]
-        })
-
+        content = await route_complete(messages, body.model)
+        return JSONResponse(build_response(chunk_id, body.model, content))
     except Exception as e:
+        logger.error(f"Chat error: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
 
-        logger.error(e)
 
-        return JSONResponse(
-            status_code=500,
-            content={
-                "error": str(e)
-            }
-        )
-
-# =====================================================
-# RUN
-# =====================================================
+# ─────────────────────────────────────────────────────────────────────────────
+# ENTRY POINT
+# ─────────────────────────────────────────────────────────────────────────────
 
 if __name__ == "__main__":
-
     import uvicorn
-
-    uvicorn.run(
-        app,
-        host="0.0.0.0",
-        port=7860
-    )
+    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")