Update app.py

app.py

@@ -1,69 +1,67 @@
+import os
+import time
+import base64
+import httpx
 import asyncio
-import uvicorn
 from fastapi import FastAPI
 from pydantic import BaseModel
 from playwright.async_api import async_playwright
 from contextlib import asynccontextmanager
 
 # --- Configuration ---
-# Limit simultaneous browser instances to prevent memory crashes
-MAX_CONCURRENT_REQUESTS = 3
-semaphore = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)
+IMGBB_API_KEY = os.getenv("IMGBB_API_KEY", "5765c99f7a977515ba320a689330e508")
+SAVE_FOLDER = "/tmp/downloaded_images"
+os.makedirs(SAVE_FOLDER, exist_ok=True)
 
 class GeminiBot:
     def __init__(self):
         self.playwright = None
-        self.browser = None
+        self.context = None
 
     async def start(self):
         self.playwright = await async_playwright().start()
-        self.browser = await self.playwright.chromium.launch(
+        # Persistent context maintains login state in ./user_data
+        self.context = await self.playwright.chromium.launch_persistent_context(
+            user_data_dir="./user_data",
             headless=True,
-            args=["--no-sandbox", "--disable-blink-features=AutomationControlled"]
+            args=["--no-sandbox", "--disable-dev-shm-usage"]
         )
-        print("--- Gemini Bot (Concurrent Stateless Mode) Ready ---")
+
+    async def upload_to_imgbb(self, image_path):
+        if not IMGBB_API_KEY or not image_path or not os.path.exists(image_path):
+            return None
+        try:
+            with open(image_path, "rb") as file:
+                img_base64 = base64.b64encode(file.read()).decode('utf-8')
+            async with httpx.AsyncClient() as client:
+                response = await client.post("https://api.imgbb.com/1/upload", data={"key": IMGBB_API_KEY, "image": img_base64})
+                return response.json()['data']['url'] if response.status_code == 200 else None
+        except:
+            return None
 
     async def process_chat(self, message: str):
-        # The semaphore ensures only N requests run at once
-        async with semaphore:
-            context = await self.browser.new_context(
-                user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
-            )
-            page = await context.new_page()
+        # Create a new page for parallel execution
+        page = await self.context.new_page()
+        try:
+            await page.goto("https://gemini.google.com")
+            # Wait for login state
+            await page.wait_for_selector("div[contenteditable='true']", timeout=60000)
             
-            try:
-                # Block heavy assets for speed
-                await page.route("**/*", lambda route: route.abort() if route.request.resource_type in 
-                                 ["image", "font", "media", "stylesheet"] else route.continue_())
-                
-                await page.goto("https://gemini.google.com", wait_until="networkidle")
-                
-                # Logic: If not logged in, this will wait until timeout
-                chat_box = page.locator("div[contenteditable='true']")
-                await chat_box.wait_for(timeout=30000)
-                await chat_box.fill(message)
-                await page.keyboard.press("Enter")
-                
-                # Wait for the response container
-                await page.wait_for_selector("message-content", timeout=45000)
-                
-                # Extract the last message
-                current_text = await page.evaluate("""() => {
-                    const msgs = document.querySelectorAll('message-content');
-                    return msgs.length > 0 ? msgs[msgs.length - 1].innerText : "No response";
-                }""")
-                
-                return {"text": current_text}
-                
-            except Exception as e:
-                return {"text": f"Error: {str(e)}"}
-            finally:
-                # Session is wiped here, regardless of success or failure
-                await context.close()
-
-    async def stop(self):
-        if self.browser: await self.browser.close()
-        if self.playwright: await self.playwright.stop()
+            chat_box = page.locator("div[contenteditable='true']").first
+            await chat_box.fill(message)
+            await page.keyboard.press("Enter")
+            
+            # Wait for response
+            await page.wait_for_selector("message-content", timeout=60000)
+            await asyncio.sleep(3) # Wait for final rendering
+            
+            # Extract text
+            text = await page.evaluate("document.querySelectorAll('message-content')[document.querySelectorAll('message-content').length-1].innerText")
+            return {"text": text, "image_url": None} # Image extraction simplified for parallel safety
+        except Exception as e:
+            return {"text": f"Error: {str(e)}", "image_url": None}
+        finally:
+            await page.close()
 
 bot = GeminiBot()
 
@@ -71,17 +69,12 @@ bot = GeminiBot()
 async def lifespan(app: FastAPI):
     await bot.start()
     yield
-    await bot.stop()
+    if bot.context: await bot.context.close()
+    if bot.playwright: await bot.playwright.stop()
 
 app = FastAPI(lifespan=lifespan)
 
-class ChatRequest(BaseModel):
-    message: str
-
 @app.post("/chat")
-async def chat_api(request: ChatRequest):
-    # This call is non-blocking to other concurrent requests
-    return await bot.process_chat(request.message)
-
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+async def chat_api(request: BaseModel):
+    msg = request.dict().get("message")
+    return {"status": "success", "data": await bot.process_chat(msg)}