Update main.py

main.py

@@ -1,9 +1,9 @@
 import os
 import json
 import time
-import logging
 import asyncio
-import g4f
+import logging
+import inspect
 
 from fastapi import FastAPI, HTTPException, Header
 from fastapi.responses import StreamingResponse
@@ -11,449 +11,445 @@ from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from typing import List, Optional
 
+import g4f
+import g4f.Provider as Provider
+import litellm
 
-# ================= Logging =================
-
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s %(levelname)s %(message)s"
-)
+# ----------------------------
+# Logging
+# ----------------------------
 
+logging.basicConfig(level=logging.INFO)
 logger=logging.getLogger(__name__)
 
-
-# ================= App =================
+# ----------------------------
+# App
+# ----------------------------
 
 app=FastAPI(
-    title="G4F Smart Gateway",
+    title="AI Gateway",
     version="2.0"
 )
 
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
-    allow_credentials=False,
     allow_methods=["*"],
     allow_headers=["*"]
 )
 
 API_KEY=os.getenv(
     "API_KEY",
-    "your_fallback_secret"
+    "your_secret"
 )
 
+DEFAULT_MODEL=os.getenv(
+    "DEFAULT_MODEL",
+    "groq/llama-3.3-70b-versatile"
+)
 
-# ================= Models =================
+# ----------------------------
+# Models
+# ----------------------------
 
-class ChatMessage(BaseModel):
+class Message(BaseModel):
     role:str
     content:str
 
 
 class ChatRequest(BaseModel):
-    model:str="gpt-4o-mini"
-    messages:List[ChatMessage]
+
+    model:str=DEFAULT_MODEL
+
+    messages:List[Message]
+
     stream:bool=False
+
     provider:Optional[str]=None
-    temperature:Optional[float]=0.7
-    max_tokens:Optional[int]=2048
 
 
-# ================= Auth =================
+# ----------------------------
+# Auth
+# ----------------------------
 
-def verify_key(auth:str):
+def verify(auth):
 
     if not auth:
-        raise HTTPException(
-            status_code=401,
-            detail="Missing Authorization"
-        )
-
-    parts=auth.split()
 
-    if len(parts)!=2:
         raise HTTPException(
             status_code=401,
-            detail="Malformed Authorization"
+            detail="Missing token"
         )
 
-    if parts[0]!="Bearer":
-        raise HTTPException(
-            status_code=401,
-            detail="Invalid Authorization"
-        )
+    if auth != f"Bearer {API_KEY}":
 
-    if parts[1]!=API_KEY:
         raise HTTPException(
             status_code=401,
-            detail="Invalid API Key"
+            detail="Unauthorized"
         )
 
 
-# ================= Provider Selection =================
+# ----------------------------
+# g4f provider discovery
+# ----------------------------
 
-def choose_provider(model:str):
+SKIP={
 
-    model=model.lower()
+"BaseProvider",
+"RetryProvider",
+"AsyncProvider"
 
-    try:
+}
 
-        if "qwen" in model:
 
-            return getattr(
-                g4f.Provider,
-                "DeepInfra",
-                None
-            )
+def collect_models(cls):
 
-        elif "perplexity" in model:
+    result=[]
 
-            return getattr(
-                g4f.Provider,
-                "PerplexityLabs",
-                None
-            )
+    for attr in [
 
-        elif "llama" in model:
+        "default_model",
+        "models",
+        "model"
 
-            return getattr(
-                g4f.Provider,
-                "DeepInfra",
-                None
-            )
+    ]:
 
-        elif "claude" in model:
+        v=getattr(
+            cls,
+            attr,
+            None
+        )
 
-            return getattr(
-                g4f.Provider,
-                "OpenaiChat",
-                None
-            )
+        if not v:
+            continue
 
-        elif "gemini" in model:
+        if isinstance(v,str):
 
-            return getattr(
-                g4f.Provider,
-                "OpenaiChat",
-                None
-            )
+            result.append(v)
 
-        elif "gpt" in model:
+        elif isinstance(
+            v,
+            (list,tuple)
+        ):
 
-            return getattr(
-                g4f.Provider,
-                "OpenaiChat",
-                None
+            result.extend(
+                [str(x) for x in v]
             )
 
-    except:
-        pass
-
-    return None
+    return list(
+        set(result)
+    )
 
 
-# ================= Health =================
+# ----------------------------
+# health
+# ----------------------------
 
 @app.get("/")
+
 async def root():
 
     return {
-        "status":"online"
+
+        "status":"online",
+
+        "default":DEFAULT_MODEL
+
     }
 
 
-# ================= Models =================
+# ----------------------------
+# models
+# ----------------------------
 
 @app.get("/v1/models")
+
 async def models(
-    authorization:str=Header(None)
+authorization:str=Header(None)
 ):
 
-    verify_key(
-        authorization
-    )
+    verify(authorization)
+
+    data=[]
 
     try:
 
-        found=[]
+        # LiteLLM models
 
-        if hasattr(
-            g4f.models,
-            "ModelUtils"
-        ):
+        ll_models=[
 
-            found=list(
-                g4f.models.ModelUtils.convert.keys()
-            )
+            "groq/llama-3.3-70b-versatile",
 
-        found=sorted(
-            list(set(found))
-        )
+            "groq/llama-3.1-8b-instant",
+
+            "openrouter/qwen/qwen-2.5-72b-instruct",
+
+            "huggingface/Qwen/Qwen2.5-72B-Instruct",
 
-        if not found:
+            "openrouter/deepseek/deepseek-chat",
 
-            found=[
+            "openai/gpt-4o",
 
-                "gpt-4o-mini",
-                "gpt-4",
-                "gpt-3.5-turbo",
-                "qwen-2.5-72b",
-                "llama-3-70b",
-                "perplexity"
+            "openai/gpt-4o-mini"
+
+        ]
+
+        for m in ll_models:
+
+            data.append({
+
+                "id":m,
+
+                "object":"model",
+
+                "owned_by":"litellm"
+
+            })
+
+        # g4f dynamic providers
+
+        for name in dir(Provider):
+
+            if name.startswith("_"):
+                continue
+
+            if name in SKIP:
+                continue
+
+            cls=getattr(
+                Provider,
+                name
+            )
 
-            ]
+            if not inspect.isclass(cls):
+                continue
 
-        return {
+            models=collect_models(
+                cls
+            )
 
-            "object":"list",
+            for m in models:
 
-            "data":[
+                data.append({
 
-                {
                     "id":m,
+
                     "object":"model",
-                    "owned_by":"g4f"
-                }
 
-                for m in found
+                    "owned_by":name
 
-            ]
-        }
+                })
 
     except Exception as e:
 
-        logger.exception(e)
-
-        return {
+        logger.error(e)
 
-            "object":"list",
+    return {
 
-            "data":[
+        "object":"list",
 
-                {
-                    "id":"gpt-4o-mini",
-                    "object":"model"
-                }
+        "data":data
 
-            ]
-        }
+    }
 
 
-# ================= Chat =================
+# ----------------------------
+# Chat
+# ----------------------------
 
 @app.post("/v1/chat/completions")
+
 async def chat(
-    body:ChatRequest,
-    authorization:str=Header(None)
+
+body:ChatRequest,
+
+authorization:str=Header(None)
+
 ):
 
-    verify_key(
+    verify(
         authorization
     )
 
-    try:
+    messages=[
 
-        messages=[
+        {
 
-            {
-                "role":m.role,
-                "content":m.content
-            }
+            "role":m.role,
 
-            for m in body.messages
-        ]
+            "content":m.content
 
+        }
 
-        provider=None
+        for m in body.messages
+    ]
 
-        if body.provider:
 
-            provider=getattr(
-                g4f.Provider,
-                body.provider,
-                None
-            )
+    # =====================
+    # Streaming
+    # =====================
 
-        else:
+    if body.stream:
 
-            provider=choose_provider(
-                body.model
-            )
+        async def generate():
 
+            try:
 
-        logger.info(
-            f"Model={body.model} Provider={provider}"
-        )
+                # LiteLLM first
+
+                response=litellm.completion(
 
+                    model=body.model,
 
-        # ========= STREAM =========
+                    messages=messages,
 
-        if body.stream:
+                    stream=True
 
-            def generate():
+                )
 
-                try:
+                for chunk in response:
 
-                    response=g4f.ChatCompletion.create(
+                    content=""
 
-                        model=body.model,
-                        messages=messages,
-                        provider=provider,
-                        stream=True
+                    try:
 
-                    )
+                        content=chunk.choices[0].delta.content
+                    except:
+                        pass
 
-                    for chunk in response:
+                    if content:
 
                         payload={
 
                             "id":"chatcmpl",
 
-                            "object":
-                            "chat.completion.chunk",
+                            "object":"chat.completion.chunk",
 
-                            "created":
-                            int(time.time()),
+                            "created":int(time.time()),
 
-                            "model":
-                            body.model,
+                            "model":body.model,
 
-                            "choices":[
-                                {
-                                    "delta":{
-                                        "content":
-                                        str(chunk)
-                                    },
-                                    "index":0
-                                }
-                            ]
-                        }
+                            "choices":[{
 
-                        yield (
-                            f"data: {json.dumps(payload)}\n\n"
-                        )
+                                "delta":{
 
-                    yield "data: [DONE]\n\n"
+                                    "content":content
 
-                except Exception as e:
+                                },
 
-                    logger.exception(e)
+                                "index":0
 
-                    yield (
-                        f"data:{json.dumps({'error':str(e)})}\n\n"
-                    )
+                            }]
+                        }
 
+                        yield f"data:{json.dumps(payload)}\n\n"
 
-            return StreamingResponse(
-                generate(),
-                media_type="text/event-stream"
-            )
+            except:
 
+                logger.info(
+                    "Fallback g4f"
+                )
 
-        # ========= NORMAL =========
+                response=g4f.ChatCompletion.create(
 
-        response=await asyncio.to_thread(
+                    model=body.model,
 
-            g4f.ChatCompletion.create,
+                    messages=messages,
 
-            model=body.model,
-            messages=messages,
-            provider=provider
+                    stream=True
+                )
 
-        )
+                for chunk in response:
 
+                    payload={
 
-        return {
+                        "choices":[{
 
-            "id":"chatcmpl",
+                            "delta":{
 
-            "object":"chat.completion",
+                                "content":str(chunk)
 
-            "created":
-            int(time.time()),
+                            }
+                        }]
+                    }
 
-            "model":
-            body.model,
+                    yield f"data:{json.dumps(payload)}\n\n"
 
-            "choices":[
+            yield "data:[DONE]\n\n"
 
-                {
+        return StreamingResponse(
 
-                    "index":0,
+            generate(),
 
-                    "message":{
+            media_type="text/event-stream"
+        )
 
-                        "role":
-                        "assistant",
 
-                        "content":
-                        str(response)
 
-                    },
+    # =====================
+    # Normal
+    # =====================
 
-                    "finish_reason":
-                    "stop"
+    try:
 
-                }
+        response=await asyncio.to_thread(
 
-            ]
-        }
+            litellm.completion,
 
-    except Exception as e:
+            model=body.model,
 
-        logger.exception(e)
+            messages=messages
 
-        raise HTTPException(
-            status_code=500,
-            detail=str(e)
         )
 
+        content=response.choices[0].message.content
 
-# ================= Test =================
 
-@app.get("/test/{model}")
-async def test_model(
-    model:str
-):
-
-    try:
+    except Exception:
 
-        provider=choose_provider(
-            model
+        logger.info(
+            "Using g4f fallback"
         )
 
-        result=await asyncio.to_thread(
+        content=await asyncio.to_thread(
 
             g4f.ChatCompletion.create,
 
-            model=model,
+            model=body.model,
 
-            provider=provider,
+            messages=messages
+        )
 
-            messages=[
 
-                {
-                    "role":"user",
-                    "content":"hello"
-                }
 
-            ]
-        )
+    return {
 
-        return {
+        "id":"chatcmpl",
 
-            "model":model,
-            "provider":str(provider),
-            "working":True
+        "object":"chat.completion",
 
-        }
+        "created":int(time.time()),
 
-    except Exception as e:
+        "model":body.model,
+
+        "choices":[
+
+            {
+
+                "index":0,
 
-        return {
+                "message":{
 
-            "model":model,
-            "working":False,
-            "error":str(e)
+                    "role":"assistant",
+
+                    "content":str(content)
+
+                },
+
+                "finish_reason":"stop"
+
+            }
+
+        ]
 
-        }
+    }