Update main.py

main.py

@@ -14,15 +14,19 @@ from typing import List, Optional
 
 # ================= Logging =================
 
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s %(levelname)s %(message)s"
+)
+
 logger=logging.getLogger(__name__)
 
 
 # ================= App =================
 
 app=FastAPI(
-    title="G4F OpenAI Compatible",
-    version="1.0"
+    title="G4F Smart Gateway",
+    version="2.0"
 )
 
 app.add_middleware(
@@ -51,6 +55,8 @@ class ChatRequest(BaseModel):
     messages:List[ChatMessage]
     stream:bool=False
     provider:Optional[str]=None
+    temperature:Optional[float]=0.7
+    max_tokens:Optional[int]=2048
 
 
 # ================= Auth =================
@@ -84,10 +90,72 @@ def verify_key(auth:str):
         )
 
 
+# ================= Provider Selection =================
+
+def choose_provider(model:str):
+
+    model=model.lower()
+
+    try:
+
+        if "qwen" in model:
+
+            return getattr(
+                g4f.Provider,
+                "DeepInfra",
+                None
+            )
+
+        elif "perplexity" in model:
+
+            return getattr(
+                g4f.Provider,
+                "PerplexityLabs",
+                None
+            )
+
+        elif "llama" in model:
+
+            return getattr(
+                g4f.Provider,
+                "DeepInfra",
+                None
+            )
+
+        elif "claude" in model:
+
+            return getattr(
+                g4f.Provider,
+                "OpenaiChat",
+                None
+            )
+
+        elif "gemini" in model:
+
+            return getattr(
+                g4f.Provider,
+                "OpenaiChat",
+                None
+            )
+
+        elif "gpt" in model:
+
+            return getattr(
+                g4f.Provider,
+                "OpenaiChat",
+                None
+            )
+
+    except:
+        pass
+
+    return None
+
+
 # ================= Health =================
 
 @app.get("/")
-async def health():
+async def root():
 
     return {
         "status":"online"
@@ -101,28 +169,38 @@ async def models(
     authorization:str=Header(None)
 ):
 
-    verify_key(authorization)
+    verify_key(
+        authorization
+    )
 
     try:
 
-        models=[]
+        found=[]
 
-        if hasattr(g4f.models,"ModelUtils"):
+        if hasattr(
+            g4f.models,
+            "ModelUtils"
+        ):
 
-            models=list(
+            found=list(
                 g4f.models.ModelUtils.convert.keys()
             )
 
-        models=sorted(
-            list(set(models))
+        found=sorted(
+            list(set(found))
         )
 
-        if not models:
+        if not found:
+
+            found=[
 
-            models=[
                 "gpt-4o-mini",
                 "gpt-4",
-                "gpt-3.5-turbo"
+                "gpt-3.5-turbo",
+                "qwen-2.5-72b",
+                "llama-3-70b",
+                "perplexity"
+
             ]
 
         return {
@@ -130,12 +208,15 @@ async def models(
             "object":"list",
 
             "data":[
+
                 {
                     "id":m,
                     "object":"model",
                     "owned_by":"g4f"
                 }
-                for m in models
+
+                for m in found
+
             ]
         }
 
@@ -148,10 +229,12 @@ async def models(
             "object":"list",
 
             "data":[
+
                 {
                     "id":"gpt-4o-mini",
                     "object":"model"
                 }
+
             ]
         }
 
@@ -164,7 +247,9 @@ async def chat(
     authorization:str=Header(None)
 ):
 
-    verify_key(authorization)
+    verify_key(
+        authorization
+    )
 
     try:
 
@@ -189,8 +274,19 @@ async def chat(
                 None
             )
 
+        else:
+
+            provider=choose_provider(
+                body.model
+            )
+
+
+        logger.info(
+            f"Model={body.model} Provider={provider}"
+        )
+
 
-        # ===== STREAM =====
+        # ========= STREAM =========
 
         if body.stream:
 
@@ -199,10 +295,12 @@ async def chat(
                 try:
 
                     response=g4f.ChatCompletion.create(
+
                         model=body.model,
                         messages=messages,
                         provider=provider,
                         stream=True
+
                     )
 
                     for chunk in response:
@@ -242,7 +340,7 @@ async def chat(
                     logger.exception(e)
 
                     yield (
-                        f"data: {json.dumps({'error':str(e)})}\n\n"
+                        f"data:{json.dumps({'error':str(e)})}\n\n"
                     )
 
 
@@ -252,7 +350,7 @@ async def chat(
             )
 
 
-        # ===== NORMAL =====
+        # ========= NORMAL =========
 
         response=await asyncio.to_thread(
 
@@ -261,17 +359,18 @@ async def chat(
             model=body.model,
             messages=messages,
             provider=provider
+
         )
 
+
         return {
 
             "id":"chatcmpl",
 
             "object":"chat.completion",
 
-            "created":int(
-                time.time()
-            ),
+            "created":
+            int(time.time()),
 
             "model":
             body.model,
@@ -284,7 +383,8 @@ async def chat(
 
                     "message":{
 
-                        "role":"assistant",
+                        "role":
+                        "assistant",
 
                         "content":
                         str(response)
@@ -293,12 +393,12 @@ async def chat(
 
                     "finish_reason":
                     "stop"
+
                 }
 
             ]
         }
 
-
     except Exception as e:
 
         logger.exception(e)
@@ -306,4 +406,54 @@ async def chat(
         raise HTTPException(
             status_code=500,
             detail=str(e)
-        )
+        )
+
+
+# ================= Test =================
+
+@app.get("/test/{model}")
+async def test_model(
+    model:str
+):
+
+    try:
+
+        provider=choose_provider(
+            model
+        )
+
+        result=await asyncio.to_thread(
+
+            g4f.ChatCompletion.create,
+
+            model=model,
+
+            provider=provider,
+
+            messages=[
+
+                {
+                    "role":"user",
+                    "content":"hello"
+                }
+
+            ]
+        )
+
+        return {
+
+            "model":model,
+            "provider":str(provider),
+            "working":True
+
+        }
+
+    except Exception as e:
+
+        return {
+
+            "model":model,
+            "working":False,
+            "error":str(e)
+
+        }