TusharPatel/webagentos-brain

1import os2import json3import hashlib4import logging5import time6import uuid7from urllib.parse import urlparse8from fastapi import FastAPI, Request9from fastapi.middleware.cors import CORSMiddleware10from groq import Groq, APIStatusError11import google.generativeai as genai12from dotenv import load_dotenv13import httpx14from agent_runtime import parse_plan, prefer_add_to_cart15from config import settings16from api_models import (17    PlanRequest,18    PlanResponse,19    CoordinatesRequest,20    CoordinatesResponse,21    VerifyRequest,22    VerifyResponse,23    MemoryExtractRequest,24    MemoryExtractResponse,25)26 27load_dotenv()28app = FastAPI(title=settings.service_name)29log = logging.getLogger("wao.brain")30logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO").upper())31 32@app.middleware("http")33async def _req_timing(request: Request, call_next):34    start = time.time()35    req_id = request.headers.get("x-request-id") or str(uuid.uuid4())36    try:37        resp = await call_next(request)38    except Exception as e:39        log.exception("unhandled_error", extra={"request_id": req_id, "path": str(request.url.path)})40        raise41    dur_ms = int((time.time() - start) * 1000)42    resp.headers["x-response-time-ms"] = str(dur_ms)43    resp.headers["x-request-id"] = req_id44    return resp45 46app.add_middleware(47    CORSMiddleware,48    allow_origins=settings.allowed_origins,49    allow_methods=["*"],50    allow_headers=["*"],51)52 53# ── Strategy 3: Key Rotation — GROQ_API_KEY = comma-separated list ────────────54_groq_keys = [k.strip() for k in os.getenv("GROQ_API_KEY", "").split(",") if k.strip()]55_groq_idx  = 056 57def _groq_client():58    if not _groq_keys:59        raise RuntimeError("No GROQ_API_KEY configured")60    return Groq(api_key=_groq_keys[_groq_idx % len(_groq_keys)])61 62def _rotate_groq():63    global _groq_idx64    _groq_idx = (_groq_idx + 1) % max(len(_groq_keys), 1)65 66# Gemini key rotation — GEMINI_API_KEY = comma-separated list67_gemini_keys = [k.strip() for k in os.getenv("GEMINI_API_KEY", "").split(",") if k.strip()]68_gemini_idx  = 069 70def _get_gemini_model(model_name: str):71    global _gemini_idx72    if not _gemini_keys:73        raise RuntimeError("No GEMINI_API_KEY configured")74    key = _gemini_keys[_gemini_idx % len(_gemini_keys)]75    genai.configure(api_key=key)76    return genai.GenerativeModel(model_name)77 78 79# ── Strategy 2: Semantic Cache via Upstash Redis ──────────────────────────────80_REDIS_URL   = os.getenv("UPSTASH_REDIS_REST_URL", "")81_REDIS_TOKEN = os.getenv("UPSTASH_REDIS_REST_TOKEN", "")82 83async def _redis_get(key: str):84    if not _REDIS_URL:85        return None86    try:87        async with httpx.AsyncClient() as c:88            r = await c.post(89                _REDIS_URL,90                headers={"Authorization": f"Bearer {_REDIS_TOKEN}"},91                json=["GET", key],92                timeout=3,93            )94            return r.json().get("result")95    except Exception:96        return None97 98async def _redis_set(key: str, value: str, ex: int = 14400):99    if not _REDIS_URL:100        return101    try:102        async with httpx.AsyncClient() as c:103            await c.post(104                _REDIS_URL,105                headers={"Authorization": f"Bearer {_REDIS_TOKEN}"},106                json=["SET", key, value, "EX", str(ex)],107                timeout=3,108            )109    except Exception:110        pass111 112def _cache_key(context: str, goal: str) -> str:113    domain = ""114    for line in context.split("\n"):115        if line.startswith("URL:"):116            try:117                domain = urlparse(line[4:].strip()).netloc118            except Exception:119                pass120            break121    return "wao:v1:" + hashlib.sha256(f"{domain}:{goal.lower().strip()}".encode()).hexdigest()[:16]122 123# ── Groq text planner with key rotation ──────────────────────────────────────124async def _groq_plan(messages: list, max_tokens: int = 512) -> str:125    for _ in range(max(len(_groq_keys), 1)):126        try:127            resp = _groq_client().chat.completions.create(128                model="llama-3.3-70b-versatile",129                messages=messages,130                response_format={"type": "json_object"},131                temperature=0.1,132                max_tokens=max_tokens,133            )134            return resp.choices[0].message.content135        except APIStatusError as e:136            if e.status_code == 429:137                _rotate_groq()138                continue139            if e.status_code == 400:140                # Groq sometimes returns json_validate_failed even with response_format=json_object.141                # Retry once in a lenient mode and let our server parse/repair JSON.142                try:143                    resp = _groq_client().chat.completions.create(144                        model="llama-3.3-70b-versatile",145                        messages=messages,146                        temperature=0.0,147                        max_tokens=max_tokens,148                    )149                    return resp.choices[0].message.content150                except Exception:151                    pass152            raise153    raise RuntimeError(f"All {len(_groq_keys)} Groq key(s) rate-limited — add more to GROQ_API_KEY")154 155def _extract_json_object(text: str):156    if not text:157        return None158    s = text.strip()159    if s.startswith("```"):160        s = "\n".join(s.split("\n")[1:-1]).strip()161    try:162        return json.loads(s)163    except Exception:164        pass165    start = s.find("{")166    end = s.rfind("}")167    if start >= 0 and end > start:168        try:169            return json.loads(s[start:end + 1])170        except Exception:171            return None172    return None173 174async def _repair_plan_json(raw_text: str) -> dict:175    fix_prompt = (176        "Convert the following into VALID JSON ONLY (no markdown). "177        "It MUST follow this schema exactly:\n"178        "{\n"179        '  \"steps\": [\n'180        "    {\n"181        '      \"action\": \"click\" | \"type\" | \"scroll\" | \"navigate\" | \"done\" | \"fail\",\n'182        '      \"selector\": string,\n'183        '      \"value\": string,\n'184        '      \"reason\": string\n'185        "    }\n"186        "  ]\n"187        "}\n\n"188        "If the content is unusable, return:\n"189        '{\"steps\":[{\"action\":\"fail\",\"selector\":\"\",\"value\":\"\",\"reason\":\"Could not produce a valid plan JSON.\"}]}\n\n'190        f"CONTENT:\n{(raw_text or '')[:3500]}"191    )192    fixed = await _groq_plan([{"role": "user", "content": fix_prompt}], max_tokens=350)193    data = _extract_json_object(fixed)194    if isinstance(data, dict):195        return data196    return {"steps": [{"action": "fail", "selector": "", "value": "", "reason": "Could not repair plan JSON."}]}197 198# ── Groq vision — image analysis using Llama vision models ───────────────────199GROQ_VISION_MODELS = [200    "meta-llama/llama-4-scout-17b-16e-instruct",201    "llama-3.2-90b-vision-preview",202    "llama-3.2-11b-vision-preview",203]204 205async def _groq_vision(prompt: str, image_b64: str, max_tokens: int = 300) -> str:206    messages = [{207        "role": "user",208        "content": [209            {"type": "text", "text": prompt},210            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}},211        ],212    }]213    for model in GROQ_VISION_MODELS:214        for _ in range(max(len(_groq_keys), 1)):215            try:216                resp = _groq_client().chat.completions.create(217                    model=model,218                    messages=messages,219                    temperature=0.1,220                    max_tokens=max_tokens,221                )222                return resp.choices[0].message.content223            except APIStatusError as e:224                if e.status_code == 429:225                    _rotate_groq()226                    continue227                break  # non-429 error — try next model228            except Exception:229                break230    raise RuntimeError("All Groq vision models exhausted")231 232# ── Phase 5: User Memory (Persistent Vault via Redis) ────────────────────────233MEMORY_EXTRACT_PROMPT = """You are a memory extraction agent. A browser task just completed.234Extract 0-3 user preference facts worth remembering for FUTURE tasks.235 236Save only personal/preference facts:237- Personal: name, address, phone, city238- Preferences: brand, size, color, budget, payment method239- Account: saved username (NEVER passwords)240 241Do NOT save: clicked buttons, scrolled pages, navigation steps.242 243Return JSON only:244{"facts": ["fact 1", "fact 2"]}245 246If nothing worth saving, return {"facts": []}247"""248 249async def _memory_get(user_id: str) -> list:250    raw = await _redis_get(f"user:{user_id}:memory")251    if not raw:252        return []253    try:254        return json.loads(raw)255    except Exception:256        return []257 258async def _memory_save(user_id: str, facts: list):259    key      = f"user:{user_id}:memory"260    existing = await _memory_get(user_id)261    for fact in facts:262        if fact and fact not in existing:263            existing.append(fact)264    await _redis_set(key, json.dumps(existing[-20:]), ex=30 * 24 * 3600)  # 30-day TTL265 266# ── System prompts ────────────────────────────────────────────────────────────267# Strategy 1: Action Chunking — prompt now asks for 1-3 steps at once268PLAN_PROMPT = """You are an AI Browser Agent. Return the next 1-3 steps as a JSON object.269Only chain steps when the sequence is predictable without seeing the intermediate page state.270 271JSON schema:272{273  "steps": [274    {275      "action": "click" | "type" | "scroll" | "navigate" | "done" | "fail",276      "selector": "sel value EXACTLY as shown (CSS or xpath:... string). Empty for navigate/scroll/done/fail.",277      "value": "text to type | url for navigate | 'up'/'down' for scroll | empty otherwise",278      "reason": "one sentence"279    }280  ]281}282 283Page context format:284  [idx] tag[type] sel="SELECTOR" label="VISIBLE TEXT"285 286Rules:287- Return 1 step when uncertain what the page will look like after the action.288- Return 2-3 steps only for predictable sequences (e.g. click input → type text → press Enter).289- Never chain more than 3 steps.290- "done" or "fail" must be the LAST step in the array.291- Copy sel="..." values EXACTLY — do not modify them.292- If sel starts with xpath: keep the full string including the xpath: prefix.293- For "scroll": selector empty, value = "down" or "up".294- For "navigate": selector empty, value = full URL.295- Elements sorted HIGH-PRIORITY first (cart/buy/checkout).296- If target is not in list, first step should be scroll to reveal it.297- Safety: If the user goal says "add to cart", prefer "Add to cart" over "Buy now". Do NOT purchase unless explicitly asked.298"""299 300VERIFY_PROMPT = """You are a browser QA agent verifying whether an action worked correctly.301Look at the screenshot carefully and answer with valid JSON only — no markdown.302 303JSON schema:304{305  "action_ok": true | false,306  "goal_done": true | false,307  "observation": "one sentence describing what you see on screen"308}309 310Rules:311- action_ok = true if the last action visibly changed the page as expected.312- goal_done = true ONLY if the user's complete goal is fully achieved and visible on screen.313- Be conservative — only set goal_done=true if you are certain.314"""315 316# Strategy 3: Model Routing — supervisor uses Gemini Flash, not Groq317SUPERVISOR_PROMPT = """You are the Supervisor of an AI agent team. Decide which specialist should handle the user's goal.318 319Specialists:320- EXECUTOR: clicks, types, scrolls on the current browser page321- RESEARCHER: needs to look up info, prices, reviews before acting322- VERIFIER: checks if the last action achieved the goal323 324Return ONLY valid JSON:325{326  "specialist": "EXECUTOR" | "RESEARCHER" | "VERIFIER",327  "reason": "one sentence"328}329"""330 331# ── Routes ────────────────────────────────────────────────────────────────────332@app.get("/")333def home():334    return {335        "status": "online",336        "service": f"{settings.service_name} {settings.service_version}",337        "groq_keys": len(_groq_keys),338        "redis": bool(_REDIS_URL),339    }340 341@app.get("/health")342def health():343    return {"ok": True}344 345# ── Memory endpoints ──────────────────────────────────────────────────────────346@app.post("/memory/get")347async def memory_get(request: Request):348    body    = await request.json()349    user_id = body.get("user_id", "")350    if not user_id:351        return {"facts": []}352    return {"facts": await _memory_get(user_id)}353 354@app.post("/memory/save")355async def memory_save(request: Request):356    body    = await request.json()357    user_id = body.get("user_id", "")358    fact    = body.get("fact", "").strip()359    if not user_id:360        return {"ok": False}361    if fact == "__clear__":362        await _redis_set(f"user:{user_id}:memory", "[]", ex=30 * 24 * 3600)363        return {"ok": True, "cleared": True}364    if not fact:365        return {"ok": False}366    await _memory_save(user_id, [fact])367    return {"ok": True}368 369@app.post("/memory/extract")370async def memory_extract(request: Request):371    """After task completes, Gemini extracts saveable preference facts from the session."""372    body = MemoryExtractRequest.model_validate(await request.json())373    goal = body.goal374    history = body.history375    if not goal:376        return MemoryExtractResponse(saved=0, facts=[])377 378    history_text = "\n".join(f"Step {i+1}: action={h.get('action')} value={h.get('value','')}" for i, h in enumerate(history[-10:]))379 380    for model_name in ["gemini-2.0-flash", "gemini-1.5-flash"]:381        try:382            model    = _get_gemini_model(model_name)383            response = model.generate_content(384                f"{MEMORY_EXTRACT_PROMPT}\n\nGoal: {goal}\nSteps taken:\n{history_text}",385                generation_config=genai.types.GenerationConfig(386                    response_mime_type="application/json",387                    temperature=0.1,388                    max_output_tokens=200,389                ),390            )391            raw   = response.text.strip()392            if raw.startswith("```"):393                raw = "\n".join(raw.split("\n")[1:-1])394            facts = json.loads(raw).get("facts", [])395            return MemoryExtractResponse(saved=len(facts), facts=facts)396        except Exception as e:397            if "429" in str(e) or "quota" in str(e).lower():398                continue399            break400    return MemoryExtractResponse(saved=0, facts=[])401 402@app.post("/supervisor")403async def supervisor(request: Request):404    """Model Routing: uses Gemini Flash (high RPM free tier) — preserves Groq quota."""405    data = await request.json()406    goal = data.get("goal", "")407    if not goal:408        return {"error": "goal is required"}409 410    for model_name in ["gemini-2.0-flash", "gemini-1.5-flash"]:411        try:412            model    = _get_gemini_model(model_name)413            response = model.generate_content(414                f"{SUPERVISOR_PROMPT}\n\nGoal: {goal}",415                generation_config=genai.types.GenerationConfig(416                    response_mime_type="application/json",417                    temperature=0.1,418                    max_output_tokens=100,419                ),420            )421            raw = response.text.strip()422            if raw.startswith("```"):423                raw = "\n".join(raw.split("\n")[1:-1])424            return json.loads(raw)425        except Exception as e:426            if "429" in str(e) or "quota" in str(e).lower():427                continue428            break429 430    # Groq fallback if all Gemini models fail431    try:432        raw = await _groq_plan([433            {"role": "system", "content": SUPERVISOR_PROMPT},434            {"role": "user",   "content": f"Goal: {goal}"},435        ], max_tokens=100)436        return json.loads(raw)437    except Exception as e:438        return {"specialist": "EXECUTOR", "reason": f"fallback — {str(e)[:100]}"}439 440@app.post("/plan")441async def plan_action(request: Request):442    body = PlanRequest.model_validate(await request.json())443    user_goal        = body.goal444    page_context     = body.context445    history          = body.history446    last_observation = body.last_observation447    user_id          = body.user_id448    user_memory_in   = body.user_memory449 450    if not user_goal:451        return {"error": "goal is required"}452 453    # Strategy 4: DOM Pruning — cap context before it reaches Groq454    page_context = (page_context or "")[: settings.max_context_chars]455    history = (history or [])[-settings.max_history_items :]456 457    is_healing = bool(last_observation and "FAILED" in last_observation)458 459    # Strategy 2: Semantic Cache — skip on self-heal (known-bad page state)460    ck = _cache_key(page_context, user_goal)461    if not is_healing and not history:462        cached = await _redis_get(ck)463        if cached:464            try:465                steps = json.loads(cached)466                parsed = parse_plan({"steps": steps})467                return PlanResponse(steps=[s.model_dump() for s in parsed.steps], cache_hit=True)468            except Exception:469                pass470 471    messages = [{"role": "system", "content": PLAN_PROMPT}]472 473    # Phase 5: User Memory (Local-first)474    # Prefer memory facts provided by the extension (stored locally), fall back to Redis only for backward compat.475    user_memory = []476    if isinstance(user_memory_in, list) and user_memory_in:477        user_memory = [str(f).strip() for f in user_memory_in if str(f).strip()]478    elif user_id:479        user_memory = await _memory_get(user_id)480 481    if user_memory:482        mem_text = "\n".join(f"- {f}" for f in user_memory[-10:])483        messages.append({"role": "user", "content": f"Known user preferences (use automatically, never ask the user):\n{mem_text}"})484 485    if history:486        history_text = "\n".join(f"Step {i+1}: {json.dumps(h)}" for i, h in enumerate(history[-5:]))487        messages.append({"role": "user", "content": f"Previous actions:\n{history_text}"})488 489    if last_observation:490        messages.append({"role": "user", "content": f"Gemini feedback: {last_observation}\nAdjust your next step(s) accordingly."})491 492    messages.append({"role": "user", "content": f"Goal: {user_goal}\n\nPage Context:\n{page_context}"})493 494    try:495        raw   = await _groq_plan(messages, max_tokens=700)496        data = _extract_json_object(raw)497        if not isinstance(data, dict):498            data = await _repair_plan_json(raw)499        parsed = parse_plan(data)500        steps = [s.model_dump() for s in parsed.steps]501        steps = prefer_add_to_cart(steps, user_goal)502 503        # Cache multi-step plans that don't contain navigate (selectors are reusable)504        if not is_healing and len(steps) > 1 and not any(s.get("action") == "navigate" for s in steps):505            await _redis_set(ck, json.dumps(steps))506 507        return PlanResponse(steps=steps, cache_hit=False)508    except json.JSONDecodeError:509        return PlanResponse(steps=[{"action": "fail", "selector": "", "value": "", "reason": "Brain returned invalid JSON. Please retry."}], cache_hit=False)510    except RuntimeError as e:511        return PlanResponse(steps=[{"action": "fail", "selector": "", "value": "", "reason": str(e)[:180]}], cache_hit=False)512    except Exception as e:513        return PlanResponse(steps=[{"action": "fail", "selector": "", "value": "", "reason": f"Brain error: {str(e)[:160]}"}], cache_hit=False)514 515 516RESEARCH_PROMPT = """You are an expert research assistant. Given a user's question and page content, provide a comprehensive answer.517 518Return ONLY valid JSON:519{520  "answer": "Your detailed answer (use \\n\\n for paragraphs)",521  "key_points": ["point 1", "point 2", "point 3"],522  "confidence": "high | medium | low"523}524 525Be factual. Reference the page content when relevant. Be concise but thorough."""526 527TAVILY_URL = "https://api.tavily.com/search"528_TAVILY_KEY = os.getenv("TAVILY_API_KEY", "")529 530async def _tavily_search(query: str, max_results: int = 5) -> list:531    if not _TAVILY_KEY:532        return []533    try:534        async with httpx.AsyncClient() as c:535            r = await c.post(536                TAVILY_URL,537                json={"api_key": _TAVILY_KEY, "query": query, "search_depth": "basic", "max_results": max_results},538                timeout=10,539            )540            data = r.json()541            return [{"title": x.get("title",""), "url": x.get("url",""), "snippet": x.get("content","")[:300]} for x in data.get("results", [])]542    except Exception:543        return []544 545@app.post("/research")546async def research(request: Request):547    body         = await request.json()548    goal         = body.get("goal", "").strip()549    page_content = body.get("page_content", "")[:3000]550    page_url     = body.get("url", "")551 552    if not goal:553        return {"error": "goal is required"}554 555    # 1. Tavily web search for external sources556    sources = await _tavily_search(goal)557 558    # 2. Build context from page + search results559    search_context = ""560    if sources:561        search_context = "\n\nWeb search results:\n" + "\n".join(562            f"- [{s['title']}]({s['url']}): {s['snippet']}" for s in sources563        )564 565    messages = [566        {"role": "system", "content": RESEARCH_PROMPT},567        {"role": "user",   "content": f"Question: {goal}\n\nCurrent page ({page_url}):\n{page_content}{search_context}"},568    ]569 570    try:571        raw  = await _groq_plan(messages, max_tokens=800)572        data = _extract_json_object(raw)573        if not isinstance(data, dict):574            data = {"answer": raw[:1000], "key_points": [], "confidence": "low"}575        return {"answer": data.get("answer", ""), "key_points": data.get("key_points", []), "confidence": data.get("confidence", "medium"), "sources": sources}576    except Exception as e:577        return {"error": str(e)[:200]}578 579 580@app.post("/coordinates")581async def get_coordinates(request: Request):582    data = CoordinatesRequest.model_validate(await request.json())583    screenshot = data.screenshot584    target = data.target585 586    if not screenshot:587        return CoordinatesResponse(x=None, y=None, error="no screenshot")588 589    prompt = (590        f"Look at this screenshot carefully. Find the '{target}' element.\n"591        f"Return ONLY valid JSON — no markdown, no explanation:\n"592        f'{{ "x": <center_x_number>, "y": <center_y_number> }}\n'593        f"If not visible: {{ \"x\": null, \"y\": null }}"594    )595 596    try:597        raw = await _groq_vision(prompt, screenshot, max_tokens=60)598        raw = raw.strip()599        if raw.startswith("```"): raw = "\n".join(raw.split("\n")[1:-1])600        result = json.loads(raw)601        return CoordinatesResponse(**result, model_used="groq-vision")602    except Exception as e:603        return CoordinatesResponse(x=None, y=None, error=str(e)[:150])604 605 606@app.post("/verify")607async def verify_action(request: Request):608    data = VerifyRequest.model_validate(await request.json())609    screenshot  = data.screenshot610    goal        = data.goal611    last_action = data.last_action612 613    if not screenshot:614        return VerifyResponse(action_ok=True, goal_done=False, observation="no screenshot provided")615 616    prompt = (617        f"{VERIFY_PROMPT}\n\n"618        f"User Goal: {goal}\n"619        f"Last action taken: {json.dumps(last_action)}\n\n"620        "Look at the screenshot and respond with JSON only — no markdown."621    )622 623    try:624        raw = await _groq_vision(prompt, screenshot, max_tokens=150)625        raw = raw.strip()626        if raw.startswith("```"):627            raw = "\n".join(raw.split("\n")[1:-1])628        result = json.loads(raw)629        return VerifyResponse(**result, model_used="groq-vision")630    except Exception as e:631        return VerifyResponse(action_ok=True, goal_done=False, observation=f"verify error: {str(e)[:100]}")632