""" Ultimate Web Agent – Fully Optimized for Hugging Face Spaces - Forced browser path to match Docker build - Headless mode always on - Memory‑saving viewport and launch args - Concurrency limited to 1 browser at a time - All API keys read from environment (with secure fallbacks) - Comprehensive error logging - No startup browser test (prevents early crashes) """ import os import sys import json import asyncio import base64 import re import uuid import logging from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import HTMLResponse, JSONResponse from pydantic import BaseModel from playwright.async_api import async_playwright, Page import httpx from groq import AsyncGroq # ===================================================================== # CRITICAL FIX: Force Playwright to use the shared browser path # This MUST match the ENV set in Dockerfile # ===================================================================== os.environ["PLAYWRIGHT_BROWSERS_PATH"] = "/home/user/.cache/ms-playwright" # Set up logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger("WebAgent") # ===================================================================== # API KEYS – read from environment variables (secure) # Fallbacks are provided only for initial testing – replace with your own. # ===================================================================== GROQ_KEYS = [ os.getenv("GROQ_KEY_1", "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"), os.getenv("GROQ_KEY_2", "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") ] OPENROUTER_KEYS = [ os.getenv("OPENROUTER_KEY", "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") ] # ===================================================================== # LLM CLIENT – rotates keys and falls back to OpenRouter / local # ===================================================================== class LLMClient: _idx = 0 _lock = asyncio.Lock() @classmethod async def chat(cls, messages, temp=0.1, max_tokens=1024): # Try Groq first try: async with cls._lock: key = GROQ_KEYS[cls._idx % len(GROQ_KEYS)] cls._idx += 1 client = AsyncGroq(api_key=key) completion = await client.chat.completions.create( model="llama-3.1-8b-instant", messages=messages, temperature=temp, max_tokens=max_tokens, response_format={"type": "json_object"} ) return completion.choices[0].message.content except Exception as e: logger.warning(f"Groq failed: {e}") # Fallback to OpenRouter try: async with cls._lock: key = OPENROUTER_KEYS[cls._idx % len(OPENROUTER_KEYS)] cls._idx += 1 async with httpx.AsyncClient(timeout=30) as client: resp = await client.post( "https://openrouter.ai/api/v1/chat/completions", headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"}, json={ "model": "meta-llama/llama-3.1-8b-instruct", "messages": messages, "temperature": temp, "max_tokens": max_tokens } ) resp.raise_for_status() data = resp.json() if "choices" in data and data["choices"]: return data["choices"][0]["message"]["content"] except Exception as e: logger.warning(f"OpenRouter failed: {e}") # Ultimate fallback (never fails) return LLMClient._fallback(messages) @staticmethod def _fallback(messages): prompt = " ".join([m["content"] for m in messages if m["role"] == "user"]) known = {"wikipedia": "https://www.wikipedia.org", "google": "https://google.com"} match = re.search(r'(?:go to|open|visit|search)\s+([a-zA-Z0-9.-]+)', prompt, re.I) if match: domain = match.group(1).lower() url = known.get(domain, f"https://{domain}.com") return json.dumps({"thought": "Fallback navigate", "action": {"tool": "navigate", "params": {"url": url}}}) return json.dumps({"thought": "Fallback done", "action": {"tool": "done", "params": {"result": "No action"}}}) # ===================================================================== # TOOL REGISTRY – the actions the agent can perform # ===================================================================== class Tools: def __init__(self, page: Page, state): self.page = page self.state = state async def navigate(self, url: str): await self.page.goto(url, timeout=30000, wait_until="networkidle") return {"result": f"Navigated to {url}"} async def click(self, selector: str): try: await self.page.click(selector, timeout=5000) return {"result": f"Clicked {selector}"} except: try: await self.page.click(f"text='{selector}'", timeout=5000) return {"result": f"Clicked by text: {selector}"} except Exception as e: return {"error": str(e)} async def fill(self, selector: str, value: str): await self.page.fill(selector, value, timeout=5000) return {"result": f"Filled {selector}"} async def get_text(self, selector: str): text = await self.page.text_content(selector, timeout=5000) return {"result": text.strip() if text else ""} async def screenshot(self): ss = await self.page.screenshot(full_page=False) return {"result": base64.b64encode(ss).decode("utf-8")} async def done(self, result: str = "Task completed."): self.state.done = True self.state.result = result return {"result": result} # ===================================================================== # AGENT – main loop that calls the LLM and executes actions # ===================================================================== class Agent: def __init__(self, tools: Tools, state): self.tools = tools self.state = state self.history = [] async def run(self, task: str): self.state.task = task self.history = [] self.state.done = False self.state.result = None max_steps = 8 for _ in range(max_steps): if self.state.done: break last = self.history[-1]['result'] if self.history else "None" context = f"Task: {task}\nLast result: {last}" sys_prompt = ( "You are a web agent. Output JSON with 'thought' and 'action'. " "Tools: navigate, click, fill, get_text, screenshot, done. " "Example: {'action':{'tool':'navigate','params':{'url':'https://www.wikipedia.org'}}}" ) raw = await LLMClient.chat([ {"role": "system", "content": sys_prompt}, {"role": "user", "content": context} ]) try: data = json.loads(raw) action = data.get("action", {"tool": "done", "params": {"result": "No action"}}) thought = data.get("thought", "") except: action = {"tool": "done", "params": {"result": "Invalid LLM output"}} thought = "Invalid" tool = action.get("tool") params = action.get("params", {}) if tool == "done": self.state.done = True self.state.result = params.get("result", "Task completed") try: ss = await self.tools.screenshot() result = {"result": self.state.result, "screenshot": ss["result"]} except: result = {"result": self.state.result} else: try: result = await getattr(self.tools, tool)(**params) ss = await self.tools.screenshot() result["screenshot"] = ss["result"] except Exception as e: result = {"error": str(e)} self.history.append({"thought": thought, "action": action, "result": result}) return {"history": self.history, "result": self.state.result, "steps": len(self.history)} # ===================================================================== # FASTAPI APPLICATION with WebSocket support # ===================================================================== app = FastAPI() app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # In-memory stores for sessions, status, and WebSocket connections sessions = {} status = {} ws_connections = {} # Limit to only 1 concurrent browser to avoid OOM on Hugging Face free tier browser_semaphore = asyncio.Semaphore(1) # ===================================================================== # BACKGROUND TASK: runs the agent in a separate coroutine # ===================================================================== async def run_agent(sid, task, headless): async with browser_semaphore: playwright = None browser = None try: logger.info(f"Starting agent for session {sid}, task: {task[:50]}...") playwright = await async_playwright().start() # Launch Chromium with memory‑saving flags browser = await playwright.chromium.launch( headless=True, # always headless on HF args=[ "--disable-dev-shm-usage", # prevent /dev/shm issues "--no-sandbox", # required in containers "--disable-gpu", # reduces memory "--disable-setuid-sandbox" ] ) # Smaller viewport = less memory context = await browser.new_context(viewport={"width": 1024, "height": 600}) page = await context.new_page() await page.goto("about:blank") state = type('State', (), {'variables': {}, 'done': False, 'result': None, 'task': task})() tools = Tools(page, state) agent = Agent(tools, state) result = await agent.run(task) status[sid] = {"status": "completed", "result": result, "progress": 100} # Notify any WebSocket listeners if sid in ws_connections: for ws in ws_connections[sid]: try: await ws.send_json({"type": "completed", "result": result}) except: pass logger.info(f"Agent completed for session {sid}") except Exception as e: logger.exception(f"Agent failed for session {sid}: {e}") status[sid] = {"status": "error", "error": str(e)} if sid in ws_connections: for ws in ws_connections[sid]: try: await ws.send_json({"type": "error", "error": str(e)}) except: pass finally: if browser: await browser.close() if playwright: await playwright.stop() # ===================================================================== # API ENDPOINTS # ===================================================================== class StartRequest(BaseModel): task: str headless: bool = True @app.post("/start") async def start(req: StartRequest): sid = str(uuid.uuid4()) status[sid] = {"status": "running", "progress": 0} asyncio.create_task(run_agent(sid, req.task, req.headless)) return {"session_id": sid, "status": "started"} @app.get("/status/{sid}") async def get_status(sid: str): if sid not in status: raise HTTPException(404, "Session not found") return JSONResponse(status[sid]) @app.post("/stop/{sid}") async def stop(sid: str): if sid in status: status[sid]["status"] = "stopped" return {"status": "stopped"} return {"status": "not found"} @app.websocket("/ws/{sid}") async def websocket_endpoint(websocket: WebSocket, sid: str): await websocket.accept() if sid not in ws_connections: ws_connections[sid] = [] ws_connections[sid].append(websocket) try: while True: data = await websocket.receive_json() if data.get("type") == "ping": await websocket.send_json({"type": "pong"}) elif data.get("type") == "stop": if sid in status: status[sid]["status"] = "stopped" await websocket.send_json({"type": "stopped"}) except WebSocketDisconnect: if sid in ws_connections: ws_connections[sid].remove(websocket) # ===================================================================== # UI – complete HTML page (same as before, but no headless checkbox) # ===================================================================== HTML = """ Web Agent

🤖 Web Agent

Ready
×
""" @app.get("/", response_class=HTMLResponse) async def root(): return HTMLResponse(HTML) @app.get("/ping") async def ping(): return {"status": "ok"} @app.get("/health") async def health(): return {"status": "healthy", "running": True} # ===================================================================== # MAIN ENTRY POINT – no browser test on startup (avoids early crashes) # ===================================================================== if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)