Update main.py

main.py

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