Update main.py

main.py

@@ -13,8 +13,9 @@ from fastapi.responses import JSONResponse
 # Configuration
 # ====================================================================
 API_SECRET_KEY = os.getenv("API_SECRET_KEY", "change-me-secret")
-POOL_SIZE      = int(os.getenv("POOL_SIZE", "3"))
-MAX_REQUESTS   = int(os.getenv("MAX_REQUESTS", "50"))   # تجديد context كل N طلب
+POOL_SIZE      = int(os.getenv("POOL_SIZE", "5"))
+MAX_REQUESTS   = int(os.getenv("MAX_REQUESTS", "30"))
+QUEUE_TIMEOUT  = int(os.getenv("QUEUE_TIMEOUT", "300"))  # ثواني انتظار Queue
 
 DUCK_MODELS = {
     "gpt-5-mini":       "GPT-5 mini",
@@ -135,10 +136,10 @@ class BrowserWorker:
                 self.busy = False
 
     async def _do_chat(self, model_label: str, prompt: str) -> str:
-        # ── تأخير عشوائي لتبدو الطلبات طبيعية ──────────────
+        # تأخير عشوائي لتبدو الطلبات طبيعية
         await asyncio.sleep(random.uniform(0.5, 2.0))
 
-        # ── تجديد context دورياً ────────────────────────────
+        # تجديد context دورياً
         self._request_count += 1
         self._total_count   += 1
         if self._request_count >= MAX_REQUESTS:
@@ -148,11 +149,9 @@ class BrowserWorker:
         try:
             page.set_default_timeout(180000)
 
-            # ── 1. فتح الصفحة ──────────────────────────────
             await page.goto("https://duckduckgo.com/aichat", wait_until="domcontentloaded")
             await asyncio.sleep(3)
 
-            # ── 2. Popup احتياطي ────────────────────────────
             try:
                 agree = page.locator('button:has-text("Agree and Continue")')
                 if await agree.count() > 0 and await agree.first.is_visible():
@@ -161,17 +160,11 @@ class BrowserWorker:
             except Exception:
                 pass
 
-            # ── 3. انتظار textarea ──────────────────────────
             await page.wait_for_selector('textarea[name="user-prompt"]')
             print(self.tag(f"Input ready (req #{self._total_count}) ✓"))
 
-            # ── 4. تغيير النموذج ────────────────────────────
             await self._select_model(page, model_label)
-
-            # ── 5. إرسال ────────────────────────────────────
             await self._send(page, prompt)
-
-            # ── 6. استخراج الرد ─────────────────────────────
             return await self._extract(page)
 
         except Exception as e:
@@ -325,16 +318,18 @@ class BrowserWorker:
 class BrowserPool:
 
     def __init__(self):
-        self.loop        = asyncio.new_event_loop()
+        self.loop            = asyncio.new_event_loop()
         self.workers: list[BrowserWorker] = []
-        self.ready_event = threading.Event()
-        self._thread     = threading.Thread(target=self._run, daemon=True)
+        self.ready_event     = threading.Event()
+        self._thread         = threading.Thread(target=self._run, daemon=True)
         self._queue: asyncio.Queue | None = None
         self._total_requests = 0
+        self._rejected       = 0   # عدد الطلبات المرفوضة بسبب timeout
 
     def start(self):
         self._thread.start()
-        print(f"[POOL] Starting {POOL_SIZE} browsers (rotation every {MAX_REQUESTS} req)...")
+        print(f"[POOL] Starting {POOL_SIZE} browsers "
+              f"(rotation every {MAX_REQUESTS} req, queue timeout {QUEUE_TIMEOUT}s)...")
 
     def _run(self):
         asyncio.set_event_loop(self.loop)
@@ -355,8 +350,23 @@ class BrowserPool:
 
     async def _process(self, model_label: str, prompt: str) -> str:
         self._total_requests += 1
-        worker: BrowserWorker = await self._queue.get()
-        print(f"[POOL] Assigned W{worker.id} (total req: {self._total_requests}) ✓")
+
+        # انتظر worker حر — بحد أقصى QUEUE_TIMEOUT ثانية
+        try:
+            worker: BrowserWorker = await asyncio.wait_for(
+                self._queue.get(), timeout=QUEUE_TIMEOUT
+            )
+        except asyncio.TimeoutError:
+            self._rejected += 1
+            print(f"[POOL] ⚠️ Queue timeout! All {POOL_SIZE} workers busy. "
+                  f"Rejected: {self._rejected}")
+            raise RuntimeError(
+                f"All {POOL_SIZE} workers are busy. "
+                f"Please retry in a moment. (rejected total: {self._rejected})"
+            )
+
+        print(f"[POOL] Assigned W{worker.id} "
+              f"(total req: {self._total_requests}) ✓")
         try:
             return await worker.chat(model_label, prompt)
         finally:
@@ -369,7 +379,7 @@ class BrowserPool:
         future = asyncio.run_coroutine_threadsafe(
             self._process(model_label, prompt), self.loop
         )
-        return future.result(timeout=240)
+        return future.result(timeout=QUEUE_TIMEOUT + 240)
 
 
 pool = BrowserPool()
@@ -399,7 +409,9 @@ def _build_prompt(messages: list) -> str:
         if not content.strip():
             continue
         if role == "system":
-            parts.append(f"=== SYSTEM INSTRUCTIONS ===\n{content}\n=== END INSTRUCTIONS ===")
+            parts.append(
+                f"=== SYSTEM INSTRUCTIONS ===\n{content}\n=== END INSTRUCTIONS ==="
+            )
         elif role == "assistant":
             parts.append(f"[Assistant]: {content}")
         else:
@@ -506,7 +518,8 @@ async def chat_completions(request: Request):
         return _make_completion(start_time, model, text, messages, tools)
     except Exception as e:
         print(f"[API] ERROR: {e}")
-        return JSONResponse(status_code=500, content={"error": {"message": str(e)}})
+        status = 503 if "busy" in str(e).lower() else 500
+        return JSONResponse(status_code=status, content={"error": {"message": str(e)}})
 
 
 @app.post("/v1/responses")
@@ -565,7 +578,8 @@ async def responses(request: Request):
         }
     except Exception as e:
         print(f"[API] ERROR: {e}")
-        return JSONResponse(status_code=500, content={"error": {"message": str(e)}})
+        status = 503 if "busy" in str(e).lower() else 500
+        return JSONResponse(status_code=status, content={"error": {"message": str(e)}})
 
 
 @app.get("/v1/models")
@@ -583,19 +597,25 @@ async def list_models(request: Request):
 async def health():
     busy  = sum(1 for w in pool.workers if w.busy)
     stats = [
-        {"id": w.id, "busy": w.busy, "total_requests": w._total_count,
-         "requests_until_rotation": MAX_REQUESTS - w._request_count}
+        {
+            "id":                    w.id,
+            "busy":                  w.busy,
+            "total_requests":        w._total_count,
+            "requests_until_rotation": MAX_REQUESTS - w._request_count,
+        }
         for w in pool.workers
     ]
     return {
-        "status":         "running",
-        "message":        "Duck.ai API Pool Server is active!",
-        "models":         ALL_MODELS,
-        "pool_size":      POOL_SIZE,
-        "workers_busy":   busy,
-        "workers_free":   POOL_SIZE - busy,
-        "total_requests": pool._total_requests,
-        "workers":        stats,
+        "status":           "running",
+        "message":          "Duck.ai API Pool Server is active!",
+        "models":           ALL_MODELS,
+        "pool_size":        POOL_SIZE,
+        "workers_busy":     busy,
+        "workers_free":     POOL_SIZE - busy,
+        "total_requests":   pool._total_requests,
+        "rejected_requests": pool._rejected,
+        "queue_timeout_sec": QUEUE_TIMEOUT,
+        "workers":          stats,
     }