agent_core.py

# agent_core.py
import asyncio
import base64
import json
import hashlib
import random
import time  # <-- FIXED: missing import added
from typing import Dict, Any, List, Optional
from dataclasses import dataclass, field
from enum import Enum
import httpx
import groq
from openai import AsyncOpenAI  # for OpenRouter compatibility
from playwright.async_api import async_playwright, Browser, Page, BrowserContext
from playwright_stealth import stealth_async
from fake_useragent import UserAgent
from stem import Signal
from stem.control import Controller

# ---------- HARDCODED API KEYS (as provided) ----------
GROQ_KEYS = [
    "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",
    "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"
]
OPENROUTER_KEYS = [
    "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",
    "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",
    "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",
    "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",
    "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",
    "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",
    "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"
]

class LLMRouter:
    def __init__(self):
        self.groq_idx = 0
        self.or_idx = 0
        self.groq_clients = [groq.Groq(api_key=k) for k in GROQ_KEYS]
        self.or_clients = [AsyncOpenAI(base_url="https://openrouter.ai/api/v1", api_key=k) for k in OPENROUTER_KEYS]
    
    async def query_groq_vision(self, prompt: str, image_b64: str) -> str:
        client = self.groq_clients[self.groq_idx % len(self.groq_clients)]
        self.groq_idx += 1
        response = client.chat.completions.create(
            model="llama-3.2-90b-vision-preview",
            messages=[{
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}}
                ]
            }],
            temperature=0.1,
            max_tokens=4096
        )
        return response.choices[0].message.content
    
    async def query_openrouter_text(self, prompt: str, system: str = "") -> str:
        client = self.or_clients[self.or_idx % len(self.or_clients)]
        self.or_idx += 1
        response = await client.chat.completions.create(
            model="anthropic/claude-3.5-sonnet",
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": prompt}
            ],
            temperature=0.3,
            max_tokens=8192
        )
        return response.choices[0].message.content

# ---------- ANONYMITY ENGINE ----------
class AnonymityEngine:
    def __init__(self, tor_socks_port=9050, control_port=9051, control_pass=""):
        self.tor_socks = f"socks5://127.0.0.1:{tor_socks_port}"
        self.ua = UserAgent()
        self.controller = Controller.from_port(port=control_port)
        self.controller.authenticate(control_pass)
        self.proxy_pool = [
            "socks5://res1:pass@geo.ip:1080",
            "socks5://res2:pass@geo.ip:1081",
            "socks5://res3:pass@geo.ip:1082"
        ]  # Rotating residential proxies – expand with your own
    
    def renew_tor_ip(self):
        self.controller.signal(Signal.NEWNYM)
        time.sleep(1)
    
    def get_browser_context_args(self) -> Dict:
        proxy = random.choice(self.proxy_pool) if random.random() < 0.7 else self.tor_socks
        return {
            "proxy": {"server": proxy} if proxy.startswith("socks") else None,
            "user_agent": self.ua.random,
            "viewport": {"width": random.choice([1366, 1440, 1920]), "height": random.choice([768, 900, 1080])},
            "locale": random.choice(["en-US", "en-GB", "de-DE", "fr-FR"]),
            "timezone_id": random.choice(["America/New_York", "Europe/London", "Asia/Tokyo"]),
            "permissions": [],
            "device_scale_factor": random.choice([1, 2]),
            "is_mobile": False,
            "has_touch": False,
            "color_scheme": random.choice(["light", "dark"]),
            "extra_http_headers": {
                "Accept-Language": random.choice(["en-US,en;q=0.9", "en-GB,en;q=0.8"]),
                "DNT": "1"
            }
        }
    
    async def apply_fingerprint_stealth(self, page: Page):
        await stealth_async(page)  # Patches navigator.webdriver, plugins, etc.
        await page.evaluate("""
            Object.defineProperty(navigator, 'hardwareConcurrency', { get: () => 8 });
            Object.defineProperty(navigator, 'deviceMemory', { get: () => 8 });
            Object.defineProperty(navigator, 'platform', { get: () => 'Win32' });
            const originalGetContext = HTMLCanvasElement.prototype.getContext;
            HTMLCanvasElement.prototype.getContext = function(type, ...args) {
                if (type === 'webgl' || type === 'webgl2') {
                    const ctx = originalGetContext.call(this, type, ...args);
                    ctx.getParameter = new Proxy(ctx.getParameter, {
                        apply: (target, thisArg, args) => {
                            if (args[0] === 37445) return 'Intel Inc.';
                            if (args[0] === 37446) return 'Intel Iris Plus Graphics';
                            return Reflect.apply(target, thisArg, args);
                        }
                    });
                    return ctx;
                }
                return originalGetContext.call(this, type, ...args);
            };
        """)

# ---------- PLAYWRIGHT AGENT ----------
class AIONAgent:
    def __init__(self):
        self.llm = LLMRouter()
        self.anonymity = AnonymityEngine()
        self.playwright = None
        self.browser: Browser = None
        self.context: BrowserContext = None
        self.page: Page = None
        self.memory = []  # episodic memory (jsonl)
        self.tool_schemas = self._load_tools()
    
    async def __aenter__(self):
        self.playwright = await async_playwright().start()
        await self._launch_browser()
        return self
    
    async def __aexit__(self, *args):
        if self.browser:
            await self.browser.close()
        if self.playwright:
            await self.playwright.stop()
    
    async def _launch_browser(self):
        args = self.anonymity.get_browser_context_args()
        # Launch a fresh Chromium with stealth flags
        browser_args = [
            "--disable-blink-features=AutomationControlled",
            "--disable-dev-shm-usage",
            "--no-sandbox",
            "--disable-setuid-sandbox",
            "--disable-gpu",
            "--window-size=1920,1080",
            "--disable-web-security",  # for CORS bypass if needed
            "--disable-features=IsolateOrigins,site-per-process",
            "--disable-extensions",
            "--disable-plugins",
            "--disable-images"  # optional – set to False for full rendering
        ]
        if args.get("proxy"):
            proxy_str = args["proxy"]["server"]
            if proxy_str.startswith("socks"):
                browser_args.append(f"--proxy-server={proxy_str}")
        
        self.browser = await self.playwright.chromium.launch(
            headless=False,  # Set True for Hugging Face – but we need Xvfb
            args=browser_args,
            chromium_sandbox=False
        )
        self.context = await self.browser.new_context(
            user_agent=args["user_agent"],
            viewport=args["viewport"],
            locale=args["locale"],
            timezone_id=args["timezone_id"],
            extra_http_headers=args["extra_http_headers"],
            device_scale_factor=args["device_scale_factor"],
            color_scheme=args["color_scheme"]
        )
        self.page = await self.context.new_page()
        await self.anonymity.apply_fingerprint_stealth(self.page)
        # Rotate IP on every new context
        self.anonymity.renew_tor_ip()
    
    # ---------- CORE TOOLS ----------
    async def navigate(self, url: str, wait_until: str = "networkidle"):
        await self.page.goto(url, wait_until=wait_until, timeout=60000)
        return await self.page.title()
    
    async def screenshot(self, full_page: bool = True, selector: str = None) -> str:
        if selector:
            element = await self.page.query_selector(selector)
            screenshot_bytes = await element.screenshot()
        else:
            screenshot_bytes = await self.page.screenshot(full_page=full_page, type="png")
        b64 = base64.b64encode(screenshot_bytes).decode("utf-8")
        # Optionally upload to external host for persistence
        return b64
    
    async def click(self, selector: str, force: bool = False):
        await self.page.click(selector, force=force)
        await self.page.wait_for_load_state("networkidle", timeout=30000)
    
    async def type_text(self, selector: str, text: str, delay: int = 50):
        await self.page.fill(selector, text)  # or type() with delay
    
    async def extract_text(self, selector: str = "body") -> str:
        return await self.page.inner_text(selector)
    
    async def execute_js(self, script: str) -> Any:
        return await self.page.evaluate(script)
    
    async def wait_for_selector(self, selector: str, timeout: int = 30000):
        await self.page.wait_for_selector(selector, timeout=timeout)
    
    async def handle_popup(self) -> str:
        async with self.page.expect_popup() as popup_info:
            popup = await popup_info.value
            await popup.wait_for_load_state()
            return await popup.title()
    
    # ---------- MULTIMODAL REASONING LOOP ----------
    async def plan_and_act(self, objective: str, max_steps: int = 10):
        """ReAct loop: observe -> reason -> act -> screenshot -> reflect."""
        for step in range(max_steps):
            # 1. Capture current state
            page_text = await self.extract_text()
            screenshot_b64 = await self.screenshot(full_page=False)
            
            # 2. Build prompt for GROK vision (decide next action)
            prompt = f"""
            Objective: {objective}
            Current page text (truncated): {page_text[:2000]}
            Available tools: {json.dumps(self.tool_schemas, indent=2)}
            You are a web automation agent. Based on the screenshot and text, output a JSON with:
            - "thought": brief analysis
            - "action": one of [navigate, click, type, extract, scroll, wait, eval, screenshot, done]
            - "params": parameters for that action (e.g., {{"selector":"#login", "text":"password"}})
            - "confidence": 0-1
            Return ONLY valid JSON.
            """
            decision_json = await self.llm.query_groq_vision(prompt, screenshot_b64)
            # Parse JSON (with fallback)
            try:
                decision = json.loads(decision_json)
            except:
                # Use OpenRouter as fallback parser
                decision = json.loads(await self.llm.query_openrouter_text(
                    f"Fix this malformed JSON: {decision_json}", 
                    system="Return only corrected JSON."
                ))
            
            # 3. Execute action
            action = decision.get("action")
            params = decision.get("params", {})
            if action == "done":
                return await self.screenshot(full_page=True)
            elif action == "navigate":
                await self.navigate(params["url"])
            elif action == "click":
                await self.click(params["selector"])
            elif action == "type":
                await self.type_text(params["selector"], params["text"])
            elif action == "extract":
                result = await self.extract_text(params.get("selector", "body"))
                self.memory.append({"step": step, "extracted": result})
            elif action == "eval":
                result = await self.execute_js(params["script"])
                self.memory.append({"step": step, "js_result": result})
            elif action == "screenshot":
                # Capture and store externally
                img_b64 = await self.screenshot(full_page=params.get("full", True))
                # upload to ImgBB or similar (fake endpoint)
                # await self._upload_image(img_b64)
                self.memory.append({"step": step, "screenshot_b64": img_b64[:100] + "..."})
            
            # 4. Reflect and store in memory
            reflection = await self.llm.query_openrouter_text(
                f"Given action '{action}' with params {params}, what did you observe? Briefly summarize.",
                system="You are an AI reflecting on browser automation."
            )
            self.memory.append({"step": step, "reflection": reflection})
        
        return {"status": "max_steps_reached", "memory": self.memory}
    
    def _load_tools(self) -> List[Dict]:
        return [
            {"name": "navigate", "description": "Go to URL", "params": {"url": "string"}},
            {"name": "click", "description": "Click element by CSS selector", "params": {"selector": "string"}},
            {"name": "type", "description": "Input text into field", "params": {"selector": "string", "text": "string"}},
            {"name": "extract", "description": "Get inner text from selector", "params": {"selector": "string (optional)"}},
            {"name": "eval", "description": "Execute JS in page", "params": {"script": "string"}},
            {"name": "screenshot", "description": "Capture full or partial page", "params": {"full": "boolean"}},
            {"name": "done", "description": "Finish task", "params": {}}
        ]

# ---------- ENTRY POINT FOR HUGGING FACE SPACE ----------
async def run_agent(objective: str):
    async with AIONAgent() as agent:
        result = await agent.plan_and_act(objective, max_steps=8)
        # Save memory to persistent volume
        with open("/data/memory.jsonl", "a") as f:
            for entry in agent.memory:
                f.write(json.dumps(entry) + "\n")
        final_screenshot = await agent.screenshot(full_page=True)
        return {
            "final_screenshot": final_screenshot,
            "memory": agent.memory,
            "status": result.get("status", "completed")
        }