app.py

"""
Ultimate Web Agent – Fully Functional
"""
import os
import json
import asyncio
import base64
import re
import uuid
import logging
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, JSONResponse
from pydantic import BaseModel
from playwright.async_api import async_playwright
from groq import AsyncGroq
import httpx

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("WebAgent")

# ---------- API Keys ----------
GROQ_KEYS = [
    "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",
    "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"
]
OPENROUTER_KEYS = ["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"]

# ---------- LLM Client ----------
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
        return json.dumps({
            "thought": "LLM unavailable, using fallback",
            "action": {"tool": "done", "params": {"result": "LLM error"}}
        })

# ---------- Tool Registry ----------
class Tools:
    def __init__(self, page, state):
        self.page = page
        self.state = state

    async def navigate(self, url):
        await self.page.goto(url, timeout=30000)
        return {"result": f"Navigated to {url}"}

    async def click(self, selector):
        await self.page.click(selector, timeout=5000)
        return {"result": f"Clicked {selector}"}

    async def fill(self, selector, value):
        await self.page.fill(selector, value, timeout=5000)
        return {"result": f"Filled {selector}"}

    async def get_text(self, selector):
        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()}

    async def done(self, result="Task completed."):
        self.state.done = True
        self.state.result = result
        return {"result": result}

# ---------- Agent ----------
class Agent:
    def __init__(self, tools, state):
        self.tools = tools
        self.state = state
        self.history = []

    async def run(self, task):
        self.state.task = task
        for step in range(8):  # max 8 steps
            if self.state.done:
                break

            # Build context
            last_result = self.history[-1]['result'] if self.history else 'None'
            context = f"Task: {task}\nLast result: {last_result}"

            sys_prompt = (
                "You are a web automation 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 output"

            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 ----------
app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

sessions = {}
status = {}

async def run_agent(sid, task, headless):
    try:
        playwright = await async_playwright().start()
        browser = await playwright.chromium.launch(headless=headless)
        context = await browser.new_context(viewport={"width": 1280, "height": 720})
        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}
        await browser.close()
    except Exception as e:
        status[sid] = {"status": "error", "error": str(e)}

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}

@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"}

# ---------- UI ----------
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>
<label><input type="checkbox" id="headless" checked> Headless</label>
<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;
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('');
}
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){
    resultbox.style.display='block';
    resultcontent.textContent=data.result.result||JSON.stringify(data.result,null,2);
    if(data.result.history){
     const last=data.result.history[data.result.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=data.result.steps||0;
     addTimeline(data.result.history);
    }
   }
  } 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;}
 const headless=document.getElementById('headless').checked;
 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})});
  const data=await res.json();
  if(!res.ok) throw new Error(data.detail||'Error');
  sid=data.session_id;startPolling();
 }catch(e){setStatus('Error: '+e.message,'error');runBtn.disabled=false;stopBtn.disabled=true;}
};
stopBtn.onclick=async ()=>{
 if(!sid) return;
 try{await fetch(`/stop/${sid}`,{method:'POST'});setStatus('⏹️ Stopping...','info');stopPolling();runBtn.disabled=false;stopBtn.disabled=true;}catch(e){setStatus('Stop error: '+e.message,'error');}
};
</script>
</body>
</html>
"""

@app.get("/", response_class=HTMLResponse)
async def root():
    return HTMLResponse(HTML)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)