app.py

"""
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 = """<!DOCTYPE html>
<html>
<head><title>Web Agent</title>
<style>
body{background:#0b0e14;color:#e8edf5;font-family:sans-serif;padding:20px}
.container{max-width:1200px;margin:auto;background:#141a24;border-radius:20px;padding:30px}
.card{background:#1a212d;border-radius:16px;padding:20px;margin-bottom:20px;border:1px solid #2a3340}
.btn{background:#4f7cff;border:none;color:#fff;padding:10px 20px;border-radius:30px;cursor:pointer}
.btn:disabled{opacity:.5}
.btn-danger{background:#e74c3c}
.btn-gold{background:#f1c40f;color:#000}
textarea{width:100%;background:#0f141d;border:1px solid #2a3340;border-radius:12px;padding:12px;color:#fff}
.progress-bar{width:100%;height:6px;background:#2a3340;border-radius:4px;margin:10px 0;overflow:hidden}
.progress-bar .fill{height:100%;background:linear-gradient(90deg,#4f7cff,#a78bfa);transition:width .5s}
.timeline-item{display:flex;gap:16px;background:#0f141d;padding:14px;border-radius:12px;margin-bottom:14px;border-left:4px solid #4f7cff}
.timeline-item img{max-width:180px;border-radius:8px;cursor:pointer}
.agent-card{background:#0f141d;padding:14px;border-radius:12px;border-left:4px solid #4f7cff}
.modal{display:none;position:fixed;inset:0;background:rgba(0,0,0,.9);justify-content:center;align-items:center;z-index:1000}
.modal img{max-width:95%;max-height:95%;border-radius:12px}
.modal-close{position:absolute;top:20px;right:40px;color:#fff;font-size:40px;cursor:pointer}
</style>
</head>
<body>
<div class="container">
<h1>🤖 Web Agent</h1>
<div class="card">
<textarea id="task" rows="2">Go to wikipedia and search for AI</textarea>
<button class="btn btn-gold" id="run">🚀 Run</button>
<button class="btn btn-danger" id="stop" disabled>⏹ Stop</button>
<div id="status" class="status">Ready</div>
<div class="progress-bar"><div class="fill" id="progress" style="width:0%"></div></div>
<div id="progresstext"></div>
</div>
<div class="card" id="resultbox" style="display:none"><h3>✅ Result</h3><pre id="resultcontent"></pre></div>
<div class="card" id="agentcard" style="display:none"><h3>🤖 Agent</h3><div class="agent-card"><div><strong>Thought:</strong> <span id="thought">—</span></div><div><strong>Action:</strong> <span id="action">—</span></div><div><strong>Result:</strong> <span id="result">—</span></div><div><strong>Steps:</strong> <span id="steps">0</span></div></div></div>
<div class="card" id="timelinebox" style="display:none"><h3>📸 Timeline</h3><div id="timeline"></div></div>
</div>
<div class="modal" onclick="this.style.display='none'"><span class="modal-close">&times;</span><img id="modalimg" src=""></div>
<script>
let sid=null, interval=null, ws=null;
const runBtn=document.getElementById('run'), stopBtn=document.getElementById('stop');
const statusDiv=document.getElementById('status'), progress=document.getElementById('progress'), progresstext=document.getElementById('progresstext');
const resultbox=document.getElementById('resultbox'), resultcontent=document.getElementById('resultcontent');
const agentcard=document.getElementById('agentcard'), thought=document.getElementById('thought'), action=document.getElementById('action'), result=document.getElementById('result'), steps=document.getElementById('steps');
const timeline=document.getElementById('timeline'), timelinebox=document.getElementById('timelinebox');
const modal=document.getElementById('modal'), modalimg=document.getElementById('modalimg');

function setStatus(msg,type='info'){statusDiv.textContent=msg;statusDiv.className='status '+type}
function setProgress(p){progress.style.width=Math.min(100,p)+'%';progresstext.textContent=`Progress: ${Math.round(p)}%`}
function addTimeline(history){
 if(!history||!history.length) return;
 timelinebox.style.display='block';
 timeline.innerHTML=history.map((item,i)=>{
   const img=item.result?.screenshot ? `<img src="data:image/png;base64,${item.result.screenshot}" onclick="document.getElementById('modalimg').src=this.src;document.getElementById('modal').style.display='flex'">` : '<div style="color:#8899b0;">No screenshot</div>';
   return `<div class="timeline-item"><div>${img}</div><div><div>💭 ${item.thought||''}</div><div>🔧 ${item.action?.tool||''}</div><div>✅ ${item.result?.result||item.result?.error||''}</div></div></div>`;
 }).join('');
}
function connectWS(){
 if(!sid) return;
 const protocol=location.protocol==='https:'?'wss:':'ws:';
 ws=new WebSocket(`${protocol}//${location.host}/ws/${sid}`);
 ws.onmessage=(e)=>{
  const data=JSON.parse(e.data);
  if(data.type==='progress') setProgress(data.progress);
  else if(data.type==='completed'){setStatus('✅ Completed','success');setProgress(100);stopPolling();runBtn.disabled=false;stopBtn.disabled=true;if(data.result) updateResult(data.result);}
  else if(data.type==='error'){setStatus('❌ Error: '+data.error,'error');stopPolling();runBtn.disabled=false;stopBtn.disabled=true;}
  else if(data.type==='stopped'){setStatus('⏹️ Stopped','info');stopPolling();runBtn.disabled=false;stopBtn.disabled=true;}
 };
 ws.onclose=()=>{ ws=null; };
}
function updateResult(res){
 resultbox.style.display='block';
 resultcontent.textContent=res.result||JSON.stringify(res,null,2);
 if(res.history){
  const last=res.history[res.history.length-1];
  agentcard.style.display='block';
  thought.textContent=last.thought||'—';
  action.textContent=last.action?.tool||'—';
  result.textContent=last.result?.result||last.result?.error||'—';
  steps.textContent=res.steps||0;
  addTimeline(res.history);
 }
}
async function poll(){
 if(!sid) return;
 try{
  const res=await fetch(`/status/${sid}`);
  const data=await res.json();
  if(data.status==='running'){setStatus('🔄 Running...','running');setProgress(data.progress||0);}
  else if(data.status==='completed'){setStatus('✅ Completed','success');setProgress(100);stopPolling();runBtn.disabled=false;stopBtn.disabled=true;if(data.result) updateResult(data.result);}
  else if(data.status==='error'){setStatus('❌ Error: '+data.error,'error');stopPolling();runBtn.disabled=false;stopBtn.disabled=true;}
  else if(data.status==='stopped'){setStatus('⏹️ Stopped','info');stopPolling();runBtn.disabled=false;stopBtn.disabled=true;}
 }catch(e){console.error(e)}
}
function startPolling(){if(interval)clearInterval(interval);interval=setInterval(poll,1500);}
function stopPolling(){if(interval){clearInterval(interval);interval=null;}}
runBtn.onclick=async ()=>{
 const task=document.getElementById('task').value.trim();
 if(!task){setStatus('Please enter a task.','error');return;}
 runBtn.disabled=true;stopBtn.disabled=false;
 setStatus('🚀 Starting...','info');setProgress(0);
 resultbox.style.display='none';agentcard.style.display='none';timelinebox.style.display='none';
 try{
  const res=await fetch('/start',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({task, headless: true})});
  const data=await res.json();
  if(!res.ok) throw new Error(data.detail||'Error');
  sid=data.session_id;
  connectWS();
  startPolling();
 }catch(e){setStatus('Error: '+e.message,'error');runBtn.disabled=false;stopBtn.disabled=true;}
};
stopBtn.onclick=async ()=>{
 if(!sid) return;
 if(ws && ws.readyState===WebSocket.OPEN){ws.send(JSON.stringify({type:'stop'}));}
 else{try{await fetch(`/stop/${sid}`,{method:'POST'});setStatus('⏹️ Stopping...','info');}catch(e){setStatus('Stop error: '+e.message,'error');}}
 stopPolling();runBtn.disabled=false;stopBtn.disabled=true;
};
</script>
</body>
</html>"""

@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)