Update main.py

main.py

@@ -1,9 +1,12 @@
 import os
 import json
 import time
+import inspect
 import asyncio
 import logging
-import inspect
+
+import g4f
+import g4f.Provider as Provider
 
 from fastapi import FastAPI, HTTPException, Header
 from fastapi.responses import StreamingResponse
@@ -11,25 +14,10 @@ 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)
 logger=logging.getLogger(__name__)
 
-# ----------------------------
-# App
-# ----------------------------
-
-app=FastAPI(
-    title="AI Gateway",
-    version="2.0"
-)
+app=FastAPI(title="G4F Dynamic API")
 
 app.add_middleware(
     CORSMiddleware,
@@ -40,48 +28,36 @@ app.add_middleware(
 
 API_KEY=os.getenv(
     "API_KEY",
-    "your_secret"
+    "secret"
 )
 
-DEFAULT_MODEL=os.getenv(
-    "DEFAULT_MODEL",
-    "groq/llama-3.3-70b-versatile"
-)
-
-# ----------------------------
+# ======================
 # Models
-# ----------------------------
+# ======================
 
-class Message(BaseModel):
+class ChatMessage(BaseModel):
     role:str
     content:str
 
 
 class ChatRequest(BaseModel):
 
-    model:str=DEFAULT_MODEL
+    model:str="gpt-4o"
 
-    messages:List[Message]
+    messages:List[ChatMessage]
 
     stream:bool=False
 
     provider:Optional[str]=None
 
 
-# ----------------------------
-# Auth
-# ----------------------------
+# ======================
+# auth
+# ======================
 
 def verify(auth):
 
-    if not auth:
-
-        raise HTTPException(
-            status_code=401,
-            detail="Missing token"
-        )
-
-    if auth != f"Bearer {API_KEY}":
+    if auth!=f"Bearer {API_KEY}":
 
         raise HTTPException(
             status_code=401,
@@ -89,30 +65,41 @@ def verify(auth):
         )
 
 
-# ----------------------------
-# g4f provider discovery
-# ----------------------------
+# ======================
+# provider discovery
+# ======================
 
 SKIP={
 
 "BaseProvider",
 "RetryProvider",
-"AsyncProvider"
+"AsyncProvider",
+"IterListProvider",
+"ProviderType",
+"CreateResult"
 
 }
 
 
+PROVIDERS={}
+MODEL_MAP={}
+
+
 def collect_models(cls):
 
-    result=[]
+    found=[]
 
-    for attr in [
+    attrs=[
 
         "default_model",
+        "model",
         "models",
-        "model"
+        "text_models",
+        "vision_models"
+
+    ]
 
-    ]:
+    for attr in attrs:
 
         v=getattr(
             cls,
@@ -125,137 +112,151 @@ def collect_models(cls):
 
         if isinstance(v,str):
 
-            result.append(v)
+            found.append(v)
 
-        elif isinstance(
-            v,
-            (list,tuple)
-        ):
+        elif isinstance(v,(list,tuple,set)):
 
-            result.extend(
+            found.extend(
                 [str(x) for x in v]
             )
 
-    return list(
-        set(result)
-    )
-
+        elif isinstance(v,dict):
 
-# ----------------------------
-# health
-# ----------------------------
+            found.extend(
+                list(v.keys())
+            )
 
-@app.get("/")
+    return list(
+        set(found)
+    )
 
-async def root():
 
-    return {
+def build():
 
-        "status":"online",
+    global PROVIDERS
+    global MODEL_MAP
 
-        "default":DEFAULT_MODEL
+    PROVIDERS={}
+    MODEL_MAP={}
 
-    }
+    for name in dir(Provider):
 
+        if name.startswith("_"):
+            continue
 
-# ----------------------------
-# models
-# ----------------------------
-
-@app.get("/v1/models")
-
-async def models(
-authorization:str=Header(None)
-):
+        if name in SKIP:
+            continue
 
-    verify(authorization)
+        try:
 
-    data=[]
+            cls=getattr(
+                Provider,
+                name
+            )
 
-    try:
+            if not inspect.isclass(cls):
+                continue
 
-        # LiteLLM models
+            if not bool(
+                getattr(
+                    cls,
+                    "working",
+                    False
+                )
+            ):
+                continue
 
-        ll_models=[
+            if bool(
+                getattr(
+                    cls,
+                    "needs_auth",
+                    False
+                )
+            ):
+                continue
 
-            "groq/llama-3.3-70b-versatile",
+            models=collect_models(
+                cls
+            )
 
-            "groq/llama-3.1-8b-instant",
+            if not models:
+                continue
 
-            "openrouter/qwen/qwen-2.5-72b-instruct",
+            PROVIDERS[name]=models
 
-            "huggingface/Qwen/Qwen2.5-72B-Instruct",
+            for m in models:
 
-            "openrouter/deepseek/deepseek-chat",
+                if m not in MODEL_MAP:
 
-            "openai/gpt-4o",
+                    MODEL_MAP[m]=name
 
-            "openai/gpt-4o-mini"
+        except:
+            pass
 
-        ]
 
-        for m in ll_models:
+build()
 
-            data.append({
 
-                "id":m,
+# ======================
+# health
+# ======================
 
-                "object":"model",
+@app.get("/")
 
-                "owned_by":"litellm"
+async def health():
 
-            })
+    return{
 
-        # g4f dynamic providers
+        "status":"online",
 
-        for name in dir(Provider):
+        "providers":len(PROVIDERS),
 
-            if name.startswith("_"):
-                continue
+        "models":len(MODEL_MAP)
 
-            if name in SKIP:
-                continue
+    }
 
-            cls=getattr(
-                Provider,
-                name
-            )
 
-            if not inspect.isclass(cls):
-                continue
+# ======================
+# models
+# ======================
 
-            models=collect_models(
-                cls
-            )
+@app.get("/v1/models")
 
-            for m in models:
+async def models(
+authorization:str=Header(None)
+):
 
-                data.append({
+    verify(
+        authorization
+    )
 
-                    "id":m,
+    return{
 
-                    "object":"model",
+        "object":"list",
 
-                    "owned_by":name
+        "data":[
 
-                })
+            {
 
-    except Exception as e:
+                "id":m,
 
-        logger.error(e)
+                "object":"model",
 
-    return {
+                "owned_by":MODEL_MAP[m]
 
-        "object":"list",
+            }
 
-        "data":data
+            for m in sorted(
+                MODEL_MAP.keys()
+            )
 
+        ]
     }
 
 
-# ----------------------------
-# Chat
-# ----------------------------
+# ======================
+# chat
+# ======================
 
 @app.post("/v1/chat/completions")
 
@@ -284,23 +285,42 @@ authorization:str=Header(None)
         for m in body.messages
     ]
 
+    provider=None
 
-    # =====================
-    # Streaming
-    # =====================
+    if body.provider:
+
+        provider=getattr(
+            Provider,
+            body.provider,
+            None
+        )
+
+    elif body.model in MODEL_MAP:
+
+        provider=getattr(
+            Provider,
+            MODEL_MAP[
+                body.model
+            ],
+            None
+        )
+
+    # ==================
+    # stream
+    # ==================
 
     if body.stream:
 
-        async def generate():
+        def generate():
 
             try:
 
-                # LiteLLM first
-
-                response=litellm.completion(
+                response=g4f.ChatCompletion.create(
 
                     model=body.model,
 
+                    provider=provider,
+
                     messages=messages,
 
                     stream=True
@@ -309,132 +329,75 @@ authorization:str=Header(None)
 
                 for chunk in response:
 
-                    content=""
-
-                    try:
-
-                        content=chunk.choices[0].delta.content
-                    except:
-                        pass
-
-                    if content:
-
-                        payload={
-
-                            "id":"chatcmpl",
-
-                            "object":"chat.completion.chunk",
-
-                            "created":int(time.time()),
-
-                            "model":body.model,
-
-                            "choices":[{
-
-                                "delta":{
-
-                                    "content":content
-
-                                },
-
-                                "index":0
-
-                            }]
-                        }
-
-                        yield f"data:{json.dumps(payload)}\n\n"
-
-            except:
-
-                logger.info(
-                    "Fallback g4f"
-                )
-
-                response=g4f.ChatCompletion.create(
-
-                    model=body.model,
+                    payload={
 
-                    messages=messages,
+                        "id":"chatcmpl",
 
-                    stream=True
-                )
+                        "object":"chat.completion.chunk",
 
-                for chunk in response:
+                        "created":int(
+                            time.time()
+                        ),
 
-                    payload={
+                        "model":body.model,
 
                         "choices":[{
 
                             "delta":{
 
-                                "content":str(chunk)
-
+                                "content":str(
+                                    chunk
+                                )
                             }
+
                         }]
                     }
 
                     yield f"data:{json.dumps(payload)}\n\n"
 
-            yield "data:[DONE]\n\n"
+                yield "data:[DONE]\n\n"
 
-        return StreamingResponse(
+            except Exception as e:
 
-            generate(),
+                logger.error(e)
 
-            media_type="text/event-stream"
-        )
-
-
-
-    # =====================
-    # Normal
-    # =====================
+                yield f"data:{json.dumps({'error':str(e)})}\n\n"
 
-    try:
-
-        response=await asyncio.to_thread(
 
-            litellm.completion,
+        return StreamingResponse(
 
-            model=body.model,
+            generate(),
 
-            messages=messages
+            media_type="text/event-stream"
 
         )
 
-        content=response.choices[0].message.content
-
-
-    except Exception:
 
-        logger.info(
-            "Using g4f fallback"
-        )
+    try:
 
-        content=await asyncio.to_thread(
+        response=await asyncio.to_thread(
 
             g4f.ChatCompletion.create,
 
             model=body.model,
 
-            messages=messages
-        )
-
+            provider=provider,
 
+            messages=messages
 
-    return {
+        )
 
-        "id":"chatcmpl",
+        return{
 
-        "object":"chat.completion",
+            "id":"chatcmpl",
 
-        "created":int(time.time()),
+            "object":"chat.completion",
 
-        "model":body.model,
+            "created":int(time.time()),
 
-        "choices":[
+            "model":body.model,
 
-            {
+            "choices":[{
 
                 "index":0,
 
@@ -442,14 +405,21 @@ authorization:str=Header(None)
 
                     "role":"assistant",
 
-                    "content":str(content)
+                    "content":str(
+                        response
+                    )
 
-                },
+                }
 
-                "finish_reason":"stop"
+            }]
 
-            }
+        }
 
-        ]
+    except Exception as e:
 
-    }
+        logger.exception(e)
+
+        raise HTTPException(
+            status_code=500,
+            detail=str(e)
+        )