Update app.py

app.py

@@ -1,221 +1,397 @@
-import time
-import random
-import traceback
-
-from fastapi import FastAPI, Header
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
-import uvicorn
+from typing import List, Optional
+
+import json
+import time
+import uuid
+import logging
 
+import g4f
 from g4f.client import Client
-from g4f.Provider import __providers__
 
-MODEL = "gpt-4o-mini"
 
-MAX_RETRIES = 10
-TIMEOUT = 60
+# =====================================================
+# LOGGING
+# =====================================================
 
-OPENAI_API_KEY = "sk-your-key"
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-SAFE_PROVIDERS = []
 
-for provider in __providers__:
+# =====================================================
+# CONFIG
+# =====================================================
 
-    try:
+API_KEY = "sk-your-secret-key"
 
-        name = provider.__name__.lower()
 
-        if not getattr(provider, "working", False):
-            continue
+# =====================================================
+# FASTAPI
+# =====================================================
 
-        if getattr(provider, "needs_auth", False):
-            continue
+app = FastAPI(
+    title="Universal AI Gateway",
+    version="5.3.0"
+)
 
-        if getattr(provider, "use_nodriver", False):
-            continue
 
-        blocked = [
-            "openai",
-            "qwen",
-            "copilot",
-            "gemini",
-            "claude",
-        ]
+# =====================================================
+# CORS
+# =====================================================
 
-        if any(x in name for x in blocked):
-            continue
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-        SAFE_PROVIDERS.append(provider)
 
-    except:
-        pass
+# =====================================================
+# MODELS
+# =====================================================
 
-random.shuffle(SAFE_PROVIDERS)
+class Message(BaseModel):
+    role: str
+    content: str
 
-print(f"[+] SAFE PROVIDERS: {len(SAFE_PROVIDERS)}")
 
-client = Client()
+class ChatRequest(BaseModel):
+    model: str
+    messages: List[Message]
+    stream: bool = False
+    temperature: Optional[float] = 0.7
+    max_tokens: Optional[int] = 4096
 
 
-class SmartG4F:
+# =====================================================
+# AUTH
+# =====================================================
 
-    def __init__(self):
+def verify_api_key(req: Request):
 
-        self.good = []
-        self.bad = {}
+    auth = req.headers.get("Authorization")
 
-    def mark_bad(self, provider, cooldown=300):
+    # السماح بدون مفتاح للاختبار
+    if not auth:
+        return True
 
-        self.bad[provider.__name__] = time.time() + cooldown
+    if not auth.startswith("Bearer "):
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid Authorization Format"
+        )
 
-    def is_bad(self, provider):
+    token = auth.replace("Bearer ", "").strip()
 
-        expire = self.bad.get(provider.__name__)
+    if token != API_KEY:
+        raise HTTPException(
+            status_code=403,
+            detail="Invalid API Key"
+        )
 
-        if not expire:
-            return False
+    return True
 
-        if time.time() > expire:
-            del self.bad[provider.__name__]
-            return False
 
-        return True
+# =====================================================
+# ROOT
+# =====================================================
 
-    def provider_pool(self):
+@app.get("/")
+async def root():
 
-        providers = []
+    return {
+        "status": "online",
+        "service": "Universal AI Gateway",
+        "version": "5.3.0"
+    }
 
-        providers.extend(self.good)
 
-        for p in SAFE_PROVIDERS:
+# =====================================================
+# MODELS
+# =====================================================
 
-            if p not in providers and not self.is_bad(p):
-                providers.append(p)
+@app.get("/v1/models")
+async def get_models():
 
-        random.shuffle(providers)
+    models_data = []
 
-        return providers
+    try:
 
-    def ask(self, prompt):
+        if hasattr(g4f.models, "_all_models"):
 
-        errors = []
+            all_models = list(g4f.models._all_models)
 
-        for attempt in range(MAX_RETRIES):
+            for model in all_models:
 
-            pool = self.provider_pool()
+                models_data.append({
+                    "id": str(model),
+                    "object": "model",
+                    "created": int(time.time()),
+                    "owned_by": "g4f"
+                })
 
-            if not pool:
+    except Exception as e:
 
-                return {
-                    "success": False,
-                    "error": "No providers available"
-                }
+        logger.error(f"Models error: {e}")
+
+    return {
+        "object": "list",
+        "data": models_data
+    }
+
+
+# =====================================================
+# CHAT COMPLETIONS
+# =====================================================
+
+@app.post("/v1/chat/completions")
+async def chat_completions(
+    req: Request,
+    body: ChatRequest
+):
+
+    verify_api_key(req)
+
+    messages = [
+        {
+            "role": m.role,
+            "content": m.content
+        }
+        for m in body.messages
+    ]
+
+    logger.info(
+        f"Request model={body.model} stream={body.stream}"
+    )
+
+    # =================================================
+    # STREAMING
+    # =================================================
+
+    if body.stream:
 
-            provider = random.choice(pool)
+        def generate_stream():
 
-            print(f"[TRY {attempt+1}] {provider.__name__}")
+            chunk_id = f"chatcmpl-{uuid.uuid4().hex}"
 
             try:
 
+                client = Client()
+
                 response = client.chat.completions.create(
-                    model=MODEL,
-                    provider=provider,
-                    messages=[
-                        {
-                            "role": "user",
-                            "content": prompt
-                        }
-                    ],
-                    timeout=TIMEOUT,
+                    model=body.model,
+                    messages=messages,
+                    stream=True
                 )
 
-                text = response.choices[0].message.content
+                has_content = False
+
+                for chunk in response:
+
+                    try:
+
+                        content = ""
+
+                        if (
+                            hasattr(chunk, "choices")
+                            and chunk.choices
+                            and len(chunk.choices) > 0
+                            and hasattr(chunk.choices[0], "delta")
+                            and chunk.choices[0].delta
+                            and chunk.choices[0].delta.content
+                        ):
+                            content = chunk.choices[0].delta.content
+
+                        if not content:
+                            continue
+
+                        has_content = True
+
+                        payload = {
+                            "id": chunk_id,
+                            "object": "chat.completion.chunk",
+                            "created": int(time.time()),
+                            "model": body.model,
+                            "choices": [
+                                {
+                                    "index": 0,
+                                    "delta": {
+                                        "content": content
+                                    },
+                                    "finish_reason": None
+                                }
+                            ]
+                        }
 
-                if not text:
-                    raise Exception("Empty response")
+                        yield (
+                            "data: "
+                            + json.dumps(
+                                payload,
+                                ensure_ascii=False
+                            )
+                            + "\n\n"
+                        )
 
-                if provider not in self.good:
-                    self.good.append(provider)
+                    except Exception as chunk_error:
 
-                return {
-                    "success": True,
-                    "provider": provider.__name__,
-                    "response": text
+                        logger.error(
+                            f"Chunk error: {chunk_error}"
+                        )
+
+                if not has_content:
+
+                    error_payload = {
+                        "error": {
+                            "message": "Provider returned empty stream",
+                            "type": "empty_stream"
+                        }
+                    }
+
+                    yield (
+                        "data: "
+                        + json.dumps(error_payload)
+                        + "\n\n"
+                    )
+
+                final_payload = {
+                    "id": chunk_id,
+                    "object": "chat.completion.chunk",
+                    "created": int(time.time()),
+                    "model": body.model,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {},
+                            "finish_reason": "stop"
+                        }
+                    ]
                 }
 
+                yield (
+                    "data: "
+                    + json.dumps(final_payload)
+                    + "\n\n"
+                )
+
+                yield "data: [DONE]\n\n"
+
             except Exception as e:
 
-                err = str(e).lower()
+                logger.error(
+                    f"Streaming error: {e}"
+                )
+
+                error_payload = {
+                    "error": {
+                        "message": str(e),
+                        "type": "server_error"
+                    }
+                }
 
-                traceback.print_exc()
+                yield (
+                    "data: "
+                    + json.dumps(error_payload)
+                    + "\n\n"
+                )
 
-                errors.append({
-                    "provider": provider.__name__,
-                    "error": str(e)
-                })
+                yield "data: [DONE]\n\n"
 
-                if "429" in err:
+        return StreamingResponse(
+            generate_stream(),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no"
+            }
+        )
 
-                    self.mark_bad(provider, 600)
-                    time.sleep(random.randint(10, 20))
+    # =================================================
+    # NORMAL RESPONSE
+    # =================================================
 
-                elif "cloudflare" in err:
+    try:
 
-                    self.mark_bad(provider, 1200)
-                    time.sleep(20)
+        client = Client()
 
-                elif "timeout" in err:
+        response = client.chat.completions.create(
+            model=body.model,
+            messages=messages,
+            stream=False
+        )
 
-                    self.mark_bad(provider, 300)
+        assistant_message = ""
 
-                else:
+        try:
 
-                    self.mark_bad(provider, 180)
+            assistant_message = (
+                response
+                .choices[0]
+                .message
+                .content
+            )
 
-                continue
+        except Exception:
 
-        return {
-            "success": False,
-            "errors": errors
-        }
+            assistant_message = str(response)
 
+        return JSONResponse({
 
-app = FastAPI()
+            "id": (
+                f"chatcmpl-{uuid.uuid4().hex}"
+            ),
 
-ai = SmartG4F()
+            "object": "chat.completion",
 
+            "created": int(time.time()),
 
-class Query(BaseModel):
-    prompt: str
+            "model": body.model,
 
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": assistant_message
+                    },
+                    "finish_reason": "stop"
+                }
+            ],
 
-@app.get("/")
-async def home():
+            "usage": {
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "total_tokens": 0
+            }
 
-    return {
-        "status": "running",
-        "providers": len(SAFE_PROVIDERS)
-    }
+        })
 
+    except Exception as e:
 
-@app.post("/ask")
-async def ask(
-    query: Query,
-    authorization: str = Header(None)
-):
+        logger.error(f"Chat error: {e}")
 
-    if authorization != f"Bearer {OPENAI_API_KEY}":
-        return {
-            "success": False,
-            "error": "Invalid API Key"
-        }
+        raise HTTPException(
+            status_code=500,
+            detail=str(e)
+        )
 
-    return ai.ask(query.prompt)
 
+# =====================================================
+# RUN
+# =====================================================
 
 if __name__ == "__main__":
 
+    import uvicorn
+
     uvicorn.run(
         app,
         host="0.0.0.0",
-        port=7860
+        port=7860,
+        log_level="info"
     )