Update app.py

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
@@ -11,7 +11,6 @@ import time
 import logging
 
 import g4f
-import g4f.models
 from g4f.client import Client
 
 # =====================================================
@@ -19,15 +18,16 @@ from g4f.client import Client
 # =====================================================
 
 logging.basicConfig(level=logging.INFO)
+
 logger = logging.getLogger(__name__)
 
 # =====================================================
-# FASTAPI
+# APP
 # =====================================================
 
 app = FastAPI(
-    title="Universal AI Gateway",
-    version="5.0.0"
+    title="DuckAI Gateway",
+    version="6.0.0"
 )
 
 # =====================================================
@@ -57,6 +57,12 @@ class ChatRequest(BaseModel):
     temperature: Optional[float] = 0.7
     max_tokens: Optional[int] = 4096
 
+# =====================================================
+# CLIENT
+# =====================================================
+
+client = Client()
+
 # =====================================================
 # ROOT
 # =====================================================
@@ -66,63 +72,45 @@ async def root():
 
     return {
         "status": "online",
-        "service": "Universal AI Gateway",
-        "version": "5.0.0"
+        "service": "DuckAI Gateway"
     }
 
 # =====================================================
-# AUTO MODELS IMPORT
+# MODELS
 # =====================================================
 
 @app.get("/v1/models")
-async def get_models():
+async def models():
 
-    models_data = []
+    output = []
 
     try:
 
-        all_models = []
-
-        if hasattr(g4f.models, "_all_models"):
-            all_models = list(g4f.models._all_models)
-
-        unique_models = set()
+        all_models = list(g4f.models._all_models)
 
         for model in all_models:
 
-            model_name = str(model)
-
-            if model_name not in unique_models:
-
-                unique_models.add(model_name)
-
-                models_data.append({
-                    "id": model_name,
-                    "object": "model",
-                    "created": int(time.time()),
-                    "owned_by": "g4f"
-                })
-
-    except Exception as e:
-
-        logger.error(f"Models error: {e}")
+            output.append({
+                "id": str(model),
+                "object": "model",
+                "created": int(time.time()),
+                "owned_by": "g4f"
+            })
 
-    # fallback
-    if not models_data:
+    except Exception:
 
         fallback = [
             "gpt-4o-mini",
             "gpt-4o",
             "gpt-4",
             "claude-3-haiku",
-            "gemini-pro",
-            "llama-3.1-70b"
+            "gemini-pro"
         ]
 
-        for m in fallback:
+        for model in fallback:
 
-            models_data.append({
-                "id": m,
+            output.append({
+                "id": model,
                 "object": "model",
                 "created": int(time.time()),
                 "owned_by": "g4f"
@@ -130,17 +118,15 @@ async def get_models():
 
     return {
         "object": "list",
-        "data": models_data
+        "data": output
     }
 
 # =====================================================
-# CHAT COMPLETIONS
+# CHAT
 # =====================================================
 
 @app.post("/v1/chat/completions")
-async def chat_completions(
-    body: ChatRequest
-):
+async def chat(body: ChatRequest):
 
     messages = [
         {
@@ -150,24 +136,24 @@ async def chat_completions(
         for m in body.messages
     ]
 
-    logger.info(f"Model: {body.model}")
-
     # =================================================
     # STREAM
     # =================================================
 
     if body.stream:
 
-        async def generate_stream():
+        async def event_stream():
 
             try:
 
-                client = Client()
-
-                response = client.chat.completions.create(
-                    model=body.model,
-                    messages=messages,
-                    stream=True
+                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}"
@@ -225,51 +211,51 @@ async def chat_completions(
                 }
 
                 yield f"data: {json.dumps(done_payload)}\n\n"
+
                 yield "data: [DONE]\n\n"
 
             except Exception as e:
 
                 logger.error(e)
 
-                error_payload = {
+                payload = {
                     "error": {
-                        "message": str(e),
-                        "type": "server_error"
+                        "message": str(e)
                     }
                 }
 
-                yield f"data: {json.dumps(error_payload)}\n\n"
+                yield f"data: {json.dumps(payload)}\n\n"
 
         return StreamingResponse(
-            generate_stream(),
+            event_stream(),
             media_type="text/event-stream",
             headers={
                 "Cache-Control": "no-cache",
-                "Connection": "keep-alive",
-                "X-Accel-Buffering": "no"
+                "Connection": "keep-alive"
             }
         )
 
     # =================================================
-    # NORMAL RESPONSE
+    # NORMAL
     # =================================================
 
     try:
 
-        client = Client()
-
-        response = await asyncio.to_thread(
-            client.chat.completions.create,
-            model=body.model,
-            messages=messages
+        response = await asyncio.wait_for(
+            asyncio.to_thread(
+                client.chat.completions.create,
+                model=body.model,
+                messages=messages
+            ),
+            timeout=60
         )
 
-        assistant_message = ""
+        text = ""
 
         try:
-            assistant_message = response.choices[0].message.content
+            text = response.choices[0].message.content
         except:
-            assistant_message = str(response)
+            text = str(response)
 
         return JSONResponse({
             "id": f"chatcmpl-{uuid.uuid4().hex}",
@@ -281,16 +267,11 @@ async def chat_completions(
                     "index": 0,
                     "message": {
                         "role": "assistant",
-                        "content": assistant_message
+                        "content": text
                     },
                     "finish_reason": "stop"
                 }
-            ],
-            "usage": {
-                "prompt_tokens": 0,
-                "completion_tokens": 0,
-                "total_tokens": 0
-            }
+            ]
         })
 
     except Exception as e: