update app

app.py

@@ -3,6 +3,8 @@ import os
 import shutil
 import time
 import uuid
+import tempfile
+import atexit
 import unicodedata
 from io import BytesIO
 from threading import Timer
@@ -13,11 +15,17 @@ import gradio as gr
 import torch
 import spaces
 from dotenv import load_dotenv
-from e2b_desktop import Sandbox
-from gradio_modal import Modal
-from huggingface_hub import login, upload_folder
 from PIL import Image, ImageDraw
 
+# Selenium Imports
+from selenium import webdriver
+from selenium.webdriver.chrome.service import Service as ChromeService
+from selenium.webdriver.chrome.options import Options as ChromeOptions
+from selenium.webdriver.common.action_chains import ActionChains
+from selenium.webdriver.common.by import By
+from selenium.webdriver.common.keys import Keys
+from webdriver_manager.chrome import ChromeDriverManager
+
 # Smolagents imports
 from smolagents import CodeAgent, tool, AgentImage
 from smolagents.memory import ActionStep, TaskStep
@@ -38,14 +46,12 @@ load_dotenv(override=True)
 # CONFIGURATION & CONSTANTS
 # -----------------------------------------------------------------------------
 
-E2B_API_KEY = os.getenv("E2B_API")
 HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_API_KEY")
 if HF_TOKEN:
+    from huggingface_hub import login
     login(token=HF_TOKEN)
 
-SANDBOXES = {}
-SANDBOX_METADATA = {}
-SANDBOX_TIMEOUT = 600
+# Browser Sandbox Config
 WIDTH = 1024
 HEIGHT = 768
 TMP_DIR = "./tmp/"
@@ -58,7 +64,6 @@ if not os.path.exists(TMP_DIR):
 
 print("Loading Fara Model... This may take a moment.")
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-# Using the Microsoft Fara model as requested
 MODEL_ID_F = "microsoft/Fara-7B" 
 
 # Global model variables
@@ -76,7 +81,7 @@ try:
     print(f"Fara Model loaded successfully on {DEVICE}")
 except Exception as e:
     print(f"Error loading Fara Model: {e}")
-    print("Falling back to Qwen/Qwen2.5-VL-7B-Instruct for demonstration if Fara is unavailable...")
+    print("Falling back to Qwen/Qwen2.5-VL-7B-Instruct...")
     try:
         MODEL_ID_F = "Qwen/Qwen2.5-VL-7B-Instruct" 
         processor_f = AutoProcessor.from_pretrained(MODEL_ID_F, trust_remote_code=True)
@@ -98,15 +103,13 @@ except Exception as e:
 @spaces.GPU(duration=120)
 def run_model_inference(formatted_messages, max_tokens=1024, stop_sequences=None):
     """
-    This function runs on the GPU worker. 
-    It receives simple python objects (lists/dicts), not the complex Agent object.
+    Runs inference on the GPU worker.
     """
     global model_f, processor_f
     
     if model_f is None:
         raise ValueError("Model is not loaded.")
 
-    # Process Inputs (Tokenization happens here to ensure tensors are on correct device)
     text = processor_f.apply_chat_template(
         formatted_messages, tokenize=False, add_generation_prompt=True
     )
@@ -121,10 +124,8 @@ def run_model_inference(formatted_messages, max_tokens=1024, stop_sequences=None
         return_tensors="pt",
     )
     
-    # Move inputs to the model's device (GPU)
     inputs = inputs.to(model_f.device)
 
-    # Generate
     with torch.no_grad():
         generated_ids = model_f.generate(
             **inputs,
@@ -133,7 +134,6 @@ def run_model_inference(formatted_messages, max_tokens=1024, stop_sequences=None
             tokenizer=processor_f.tokenizer, 
         )
 
-    # Decode
     generated_ids_trimmed = [
         out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
     ]
@@ -145,9 +145,6 @@ def run_model_inference(formatted_messages, max_tokens=1024, stop_sequences=None
 
 
 class FaraLocalModel(Model):
-    """
-    Wrapper for the local Fara (Qwen2.5-VL) model to work with SmolAgents.
-    """
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
@@ -159,13 +156,9 @@ class FaraLocalModel(Model):
     ) -> ChatMessage:
         
         formatted_messages = []
-        
-        # Convert SmolAgents messages to Qwen/Transformers format
-        # We perform this conversion here (CPU side) to create simple dicts/lists
         for msg in messages:
             role = msg["role"]
             content = msg["content"]
-            
             new_content = []
             
             if isinstance(content, str):
@@ -177,7 +170,6 @@ class FaraLocalModel(Model):
                     elif isinstance(item, dict):
                         if "type" in item:
                             if item["type"] == "image":
-                                # Handle path or url - extract value to ensure serializability
                                 val = item.get("image") or item.get("url") or item.get("path")
                                 new_content.append({"type": "image", "image": val})
                             else:
@@ -185,9 +177,6 @@ class FaraLocalModel(Model):
             
             formatted_messages.append({"role": role, "content": new_content})
 
-        # Call the decorated global function
-        # This crosses the boundary to the GPU worker safely because
-        # formatted_messages contains only standard Python types (str, list, dict, PIL.Image)
         output_text = run_model_inference(
             formatted_messages=formatted_messages,
             max_tokens=kwargs.get("max_tokens", 1024),
@@ -200,14 +189,138 @@ class FaraLocalModel(Model):
         )
 
 # -----------------------------------------------------------------------------
-# E2B AGENT & TOOLS
+# SELENIUM CHROME SANDBOX
 # -----------------------------------------------------------------------------
 
-E2B_SYSTEM_PROMPT_TEMPLATE = """You are a desktop automation assistant that can control a remote desktop environment. The current date is <<current_date>>.
+class SeleniumSandbox:
+    def __init__(self, width=1024, height=768):
+        self.width = width
+        self.height = height
+        self.tmp_dir = tempfile.mkdtemp(prefix="chrome_sandbox_")
+        
+        # Setup Chrome Options
+        chrome_opts = ChromeOptions()
+        chrome_opts.add_argument("--headless=new") # Run headless for Space compatibility
+        chrome_opts.add_argument(f"--user-data-dir={self.tmp_dir}")
+        chrome_opts.add_argument(f"--window-size={width},{height}")
+        chrome_opts.add_argument("--no-sandbox") 
+        chrome_opts.add_argument("--disable-dev-shm-usage")
+        chrome_opts.add_argument("--disable-extensions")
+        chrome_opts.add_argument("--disable-gpu")
+        chrome_opts.add_argument("--force-device-scale-factor=1")
+        chrome_opts.add_argument("--hide-scrollbars")
+        
+        # Initialize Driver
+        try:
+            self.driver = webdriver.Chrome(
+                service=ChromeService(ChromeDriverManager().install()), 
+                options=chrome_opts
+            )
+            # Ensure precise viewport size (excluding chrome UI overhead if any)
+            self.driver.set_window_size(width, height)
+            
+            # Start with a blank page
+            self.driver.get("about:blank")
+            print(f"Selenium Chrome Driver started. Data dir: {self.tmp_dir}")
+        except Exception as e:
+            print(f"Failed to initialize Selenium: {e}")
+            self.cleanup()
+            raise e
+
+    def get_screenshot(self):
+        """Returns screenshot as PIL Image"""
+        png_data = self.driver.get_screenshot_as_png()
+        return Image.open(BytesIO(png_data))
+
+    def move_mouse_and_click(self, x, y, click_type="left"):
+        """
+        Simulate mouse movement and click using ActionChains.
+        Note: Selenium tracks state, so we move relative to the 'body' tag to ensure absolute positioning simulation.
+        """
+        try:
+            body = self.driver.find_element(By.TAG_NAME, "body")
+            actions = ActionChains(self.driver)
+            
+            # Move to 0,0 of body, then offset
+            actions.move_to_element_with_offset(body, 0, 0)
+            actions.move_by_offset(x, y)
+            
+            if click_type == "left":
+                actions.click()
+            elif click_type == "right":
+                actions.context_click()
+            elif click_type == "double":
+                actions.double_click()
+                
+            actions.perform()
+        except Exception as e:
+            print(f"Error in move_mouse_and_click: {e}")
+
+    def drag_and_drop(self, x1, y1, x2, y2):
+        try:
+            body = self.driver.find_element(By.TAG_NAME, "body")
+            actions = ActionChains(self.driver)
+            
+            actions.move_to_element_with_offset(body, 0, 0)
+            actions.move_by_offset(x1, y1)
+            actions.click_and_hold()
+            
+            # Move relative from x1,y1 to x2,y2
+            actions.move_by_offset(x2 - x1, y2 - y1)
+            actions.release()
+            actions.perform()
+        except Exception as e:
+            print(f"Error in drag_and_drop: {e}")
+
+    def type_text(self, text):
+        actions = ActionChains(self.driver)
+        actions.send_keys(text)
+        actions.perform()
+        
+    def press_key(self, key_name):
+        try:
+            k = getattr(Keys, key_name.upper(), None)
+            if not k:
+                # Handle common overrides
+                if key_name.lower() == "enter": k = Keys.ENTER
+                elif key_name.lower() == "space": k = Keys.SPACE
+                elif key_name.lower() == "backspace": k = Keys.BACK_SPACE
+                elif key_name.lower() == "esc": k = Keys.ESCAPE
+                else: k = key_name # Fallback to literal
+
+            actions = ActionChains(self.driver)
+            actions.send_keys(k)
+            actions.perform()
+        except Exception as e:
+            print(f"Error pressing key: {e}")
+
+    def scroll(self, amount, direction="down"):
+        # Selenium doesn't have a great scroll wheel primitive, use JS
+        try:
+            scroll_y = amount * 100 # Arbitrary multiplier to match "notches"
+            if direction == "up":
+                scroll_y = -scroll_y
+            self.driver.execute_script(f"window.scrollBy(0, {scroll_y});")
+        except Exception as e:
+            print(f"Error scrolling: {e}")
+
+    def cleanup(self):
+        try:
+            if hasattr(self, 'driver'):
+                self.driver.quit()
+        except:
+            pass
+        shutil.rmtree(self.tmp_dir, ignore_errors=True)
+
+# -----------------------------------------------------------------------------
+# AGENT SETUP
+# -----------------------------------------------------------------------------
+
+SYSTEM_PROMPT_TEMPLATE = """You are a browser automation assistant controlling a Google Chrome web browser. The current date is <<current_date>>.
 
 <action process>
 You will be given a task to solve in several steps. At each step you will perform an action.
-After each action, you'll receive an updated screenshot. 
+After each action, you'll receive an updated screenshot of the browser. 
 Then you will proceed as follows, with these sections: don't skip any!
 
 Short term goal: ...
@@ -222,7 +335,7 @@ Always format your action ('Action:' part) as Python code blocks as shown above.
 </action_process>
 
 <tools>
-On top of performing computations in the Python code snippets that you create, you only have access to these tools to interact with the desktop, no additional ones:
+On top of performing computations in the Python code snippets that you create, you only have access to these tools to interact with the browser:
 {%- for tool in tools.values() %}
 - {{ tool.name }}: {{ tool.description }}
     Takes inputs: {{tool.inputs}}
@@ -231,29 +344,23 @@ On top of performing computations in the Python code snippets that you create, y
 </tools>
 
 <click_guidelines>
-Look at elements on the screen to determine what to click or interact with.
-The desktop has a resolution of <<resolution_x>>x<<resolution_y>> pixels, take it into account to decide clicking coordinates. NEVER USE HYPOTHETIC OR ASSUMED COORDINATES, USE TRUE COORDINATES that you can see from the screenshot.
-Use precise coordinates based on the current screenshot for mouse movements and clicks. 
-Whenever you click, MAKE SURE to click in the middle of the button, text, link or any other clickable element. Not under, not on the side. IN THE MIDDLE, else you risk to miss it.
-In menus it is always better to click in the middle of the text rather than in the tiny icon. Calculate extremelly well the coordinates. A mistake here can make the full task fail.
-Sometimes you may have missed a click, so never assume that you're on the right page, always make sure that your previous action worked.
-In the screenshot you will see a green crosshair displayed over the position of your last click: this way can inspect if the mouse pointer is off of the targeted element, pay special attention to it.
+The browser has a resolution of <<resolution_x>>x<<resolution_y>> pixels.
+NEVER USE HYPOTHETIC OR ASSUMED COORDINATES, USE TRUE COORDINATES that you can see from the screenshot.
+Use precise coordinates based on the current screenshot.
+Whenever you click, MAKE SURE to click in the middle of the button, text, link or any other clickable element.
+In the screenshot you will see a green crosshair displayed over the position of your last click.
 </click_guidelines>
 
 <general_guidelines>
-Always analyze the latest screenshot carefully before performing actions.
-You can wait for appropriate loading times using the wait() tool. But don't wait forever, sometimes you've just misclicked and the process didn't launch.
-Execute one action at a time: don't try to pack a click and typing in one action.
-On each step, look at the last screenshot and action to validate if previous steps worked and decide the next action. If you repeated an action already without effect, it means that this action is useless: don't repeat it and try something else.
-Use click to move through menus on the desktop and scroll for web and specific applications.
-Always analyze the latest screenshot carefully before performing actions.
-Desktop menus usually expand with more options, the tiny triangle next to some text in a menu means that menu expands. For example in Office in the Applications menu expands showing presentation or writing applications. 
-NEVER CLICK THE WEB BROWSER ICON TO OPEN THE WEB BROWSER: use open_url directly.
-In browser, ignore any sign-in popups while they don't interfere with the elements you want to interact with.
+Execute one action at a time.
+Use `open_url` to navigate to websites.
+Use `click` to navigate links and interface elements.
+Use `type_text` to input into forms.
+Use `scroll` to see more content.
+If you get stuck, try using `open_url` to search on Google.
 </general_guidelines>
 """.replace("<<current_date>>", datetime.now().strftime("%A, %d-%B-%Y"))
 
-
 def draw_marker_on_image(image_copy, click_coordinates):
     x, y = click_coordinates
     draw = ImageDraw.Draw(image_copy)
@@ -261,385 +368,225 @@ def draw_marker_on_image(image_copy, click_coordinates):
     # Draw cross
     draw.line((x - cross_size, y, x + cross_size, y), fill="green", width=linewidth)
     draw.line((x, y - cross_size, x, y + cross_size), fill="green", width=linewidth)
-    # Add a circle around it for better visibility
     draw.ellipse(
-        (
-            x - cross_size * 2,
-            y - cross_size * 2,
-            x + cross_size * 2,
-            y + cross_size * 2,
-        ),
+        (x - cross_size * 2, y - cross_size * 2, x + cross_size * 2, y + cross_size * 2),
         outline="green",
         width=linewidth,
     )
     return image_copy
 
-def get_agent_summary_erase_images(agent):
-    for memory_step in agent.memory.steps:
-        if hasattr(memory_step, "observations_images"):
-            memory_step.observations_images = None
-        if hasattr(memory_step, "task_images"):
-            memory_step.task_images = None
-    return agent.write_memory_to_messages()
-
-
-class E2BVisionAgent(CodeAgent):
-    """Agent for e2b desktop automation with Vision capabilities"""
+class SeleniumVisionAgent(CodeAgent):
+    """Agent for Browser automation with Selenium and Vision"""
 
     def __init__(
         self,
         model: Model,
         data_dir: str,
-        desktop: Sandbox,
-        tools: List[tool] = None,
-        max_steps: int = 200,
+        sandbox: SeleniumSandbox,
+        max_steps: int = 20,
         verbosity_level: LogLevel = 2,
-        planning_interval: int = None,
-        use_v1_prompt: bool = False,
         **kwargs,
     ):
-        self.desktop = desktop
+        self.sandbox = sandbox
         self.data_dir = data_dir
-        self.planning_interval = planning_interval
-        # Initialize Desktop
-        self.width, self.height = self.desktop.get_screen_size()
-        print(f"Screen size: {self.width}x{self.height}")
-
-        # Set up temp directory
+        
+        # Initialize
+        print(f"Browser size: {self.sandbox.width}x{self.sandbox.height}")
         os.makedirs(self.data_dir, exist_ok=True)
-        print(f"Screenshots and steps will be saved to: {self.data_dir}")
 
-        self.use_v1_prompt = use_v1_prompt
-        # Initialize base agent
         super().__init__(
-            tools=tools or [],
+            tools=[],
             model=model,
             max_steps=max_steps,
             verbosity_level=verbosity_level,
-            planning_interval=self.planning_interval,
             **kwargs,
         )
-        self.prompt_templates["system_prompt"] = E2B_SYSTEM_PROMPT_TEMPLATE.replace(
-            "<<resolution_x>>", str(self.width)
-        ).replace("<<resolution_y>>", str(self.height))
-
-        # Add screen info to state
-        self.state["screen_width"] = self.width
-        self.state["screen_height"] = self.height
+        
+        self.prompt_templates["system_prompt"] = SYSTEM_PROMPT_TEMPLATE.replace(
+            "<<resolution_x>>", str(self.sandbox.width)
+        ).replace("<<resolution_y>>", str(self.sandbox.height))
 
-        # Add default tools
-        self.logger.log("Setting up agent tools...")
-        self._setup_desktop_tools()
+        self.register_tools()
         self.step_callbacks.append(self.take_screenshot_callback)
 
-    def _setup_desktop_tools(self):
-        """Register all desktop tools"""
+    def register_tools(self):
 
         @tool
         def click(x: int, y: int) -> str:
             """
-            Performs a left-click at the specified coordinates
+            Performs a left-click at the specified coordinates.
             Args:
-                x: The x coordinate (horizontal position)
-                y: The y coordinate (vertical position)
+                x: The x coordinate (horizontal position).
+                y: The y coordinate (vertical position).
             """
-            self.desktop.move_mouse(x, y)
-            self.desktop.left_click()
+            self.sandbox.move_mouse_and_click(x, y, "left")
             self.click_coordinates = [x, y]
-            self.logger.log(f"Clicked at coordinates ({x}, {y})")
-            return f"Clicked at coordinates ({x}, {y})"
+            return f"Clicked at ({x}, {y})"
 
         @tool
         def right_click(x: int, y: int) -> str:
             """
-            Performs a right-click at the specified coordinates
+            Performs a right-click at the specified coordinates.
             Args:
-                x: The x coordinate (horizontal position)
-                y: The y coordinate (vertical position)
+                x: The x coordinate.
+                y: The y coordinate.
             """
-            self.desktop.move_mouse(x, y)
-            self.desktop.right_click()
+            self.sandbox.move_mouse_and_click(x, y, "right")
             self.click_coordinates = [x, y]
-            self.logger.log(f"Right-clicked at coordinates ({x}, {y})")
-            return f"Right-clicked at coordinates ({x}, {y})"
+            return f"Right-clicked at ({x}, {y})"
 
         @tool
         def double_click(x: int, y: int) -> str:
             """
-            Performs a double-click at the specified coordinates
+            Performs a double-click at the specified coordinates.
             Args:
-                x: The x coordinate (horizontal position)
-                y: The y coordinate (vertical position)
+                x: The x coordinate.
+                y: The y coordinate.
             """
-            self.desktop.move_mouse(x, y)
-            self.desktop.double_click()
+            self.sandbox.move_mouse_and_click(x, y, "double")
             self.click_coordinates = [x, y]
-            self.logger.log(f"Double-clicked at coordinates ({x}, {y})")
-            return f"Double-clicked at coordinates ({x}, {y})"
-
-        @tool
-        def move_mouse(x: int, y: int) -> str:
-            """
-            Moves the mouse cursor to the specified coordinates
-            Args:
-                x: The x coordinate (horizontal position)
-                y: The y coordinate (vertical position)
-            """
-            self.desktop.move_mouse(x, y)
-            self.logger.log(f"Moved mouse to coordinates ({x}, {y})")
-            return f"Moved mouse to coordinates ({x}, {y})"
-
-        def normalize_text(text):
-            return "".join(
-                c
-                for c in unicodedata.normalize("NFD", text)
-                if not unicodedata.combining(c)
-            )
+            return f"Double-clicked at ({x}, {y})"
 
         @tool
         def type_text(text: str) -> str:
             """
-            Types the specified text at the current cursor position.
+            Types the specified text.
             Args:
-                text: The text to type
+                text: The text to type.
             """
-            clean_text = normalize_text(text)
-            self.desktop.write(clean_text, delay_in_ms=75)
-            self.logger.log(f"Typed text: '{clean_text}'")
+            clean_text = unicodedata.normalize("NFD", text)
+            self.sandbox.type_text(clean_text)
             return f"Typed text: '{clean_text}'"
 
         @tool
         def press_key(key: str) -> str:
             """
-            Presses a keyboard key
+            Presses a keyboard key (e.g., 'enter', 'backspace', 'esc').
             Args:
-                key: The key to press (e.g. "enter", "space", "backspace", etc.).
+                key: The key name.
             """
-            self.desktop.press(key)
-            self.logger.log(f"Pressed key: {key}")
+            self.sandbox.press_key(key)
             return f"Pressed key: {key}"
 
-        @tool
-        def go_back() -> str:
-            """
-            Goes back to the previous page in the browser.
-            """
-            self.desktop.press(["alt", "left"])
-            self.logger.log("Went back one page")
-            return "Went back one page"
-
         @tool
         def drag_and_drop(x1: int, y1: int, x2: int, y2: int) -> str:
             """
-            Clicks [x1, y1], drags mouse to [x2, y2], then release click.
+            Drags from (x1, y1) and drops at (x2, y2).
             Args:
-                x1: The x coordinate of the start position.
-                y1: The y coordinate of the start position.
-                x2: The x coordinate of the end position.
-                y2: The y coordinate of the end position.
+                x1: Start x coordinate.
+                y1: Start y coordinate.
+                x2: End x coordinate.
+                y2: End y coordinate.
             """
-            self.desktop.drag([x1, y1], [x2, y2])
-            message = f"Dragged and dropped from [{x1}, {y1}] to [{x2}, {y2}]"
-            self.logger.log(message)
-            return message
+            self.sandbox.drag_and_drop(x1, y1, x2, y2)
+            return f"Dragged from [{x1}, {y1}] to [{x2}, {y2}]"
 
         @tool
-        def scroll(x: int, y: int, direction: str = "down", amount: int = 2) -> str:
+        def scroll(amount: int, direction: str = "down") -> str:
             """
-            Moves the mouse to selected coordinates, then uses the scroll button.
+            Scrolls the page.
             Args:
-                x: The x coordinate
-                y: The y coordinate
-                direction: "up" or "down"
-                amount: The amount to scroll.
+                amount: The amount to scroll (1-10).
+                direction: "up" or "down".
             """
-            self.desktop.move_mouse(x, y)
-            self.desktop.scroll(direction=direction, amount=amount)
-            message = f"Scrolled {direction} by {amount}"
-            self.logger.log(message)
-            return message
+            self.sandbox.scroll(amount, direction)
+            return f"Scrolled {direction} by {amount}"
 
         @tool
         def wait(seconds: float) -> str:
             """
             Waits for the specified number of seconds.
             Args:
-                seconds: The duration to wait in seconds.
+                seconds: The duration to wait.
             """
             time.sleep(seconds)
-            self.logger.log(f"Waited for {seconds} seconds")
             return f"Waited for {seconds} seconds"
 
         @tool
         def open_url(url: str) -> str:
             """
-            Directly opens a browser with the specified url.
+            Navigates the browser to the specified URL.
             Args:
-                url: The website URL to open.
+                url: The URL to open.
             """
             if not url.startswith(("http://", "https://")):
                 url = "https://" + url
-
-            self.desktop.open(url)
-            time.sleep(2)
-            self.logger.log(f"Opening URL: {url}")
-            return f"Opened URL: {url}"
+            try:
+                self.sandbox.driver.get(url)
+                time.sleep(2)
+                title = self.sandbox.driver.title
+                return f"Opened URL: {url}. Page Title: {title}"
+            except Exception as e:
+                return f"Failed to open URL: {e}"
 
         @tool
-        def find_on_page_ctrl_f(search_string: str) -> str:
+        def go_back() -> str:
             """
-            Scroll the browser viewport to the first occurrence of the search string (Ctrl+F).
-            Args:
-                search_string: The text to search for on the page.
+            Goes back to the previous page in history.
             """
-            self.desktop.press(["ctrl", "f"])
-            time.sleep(0.3)
-            clean_text = normalize_text(search_string)
-            self.desktop.write(clean_text, delay_in_ms=75)
-            time.sleep(0.3)
-            self.desktop.press("enter")
-            time.sleep(0.3)
-            self.desktop.press("esc")
-            output_message = f"Scrolled to the first occurrence of '{clean_text}'"
-            self.logger.log(output_message)
-            return output_message
-
-        # Register the tools
+            self.sandbox.driver.back()
+            return "Went back one page"
+
         self.tools["click"] = click
         self.tools["right_click"] = right_click
         self.tools["double_click"] = double_click
-        self.tools["move_mouse"] = move_mouse
         self.tools["type_text"] = type_text
         self.tools["press_key"] = press_key
+        self.tools["drag_and_drop"] = drag_and_drop
         self.tools["scroll"] = scroll
         self.tools["wait"] = wait
         self.tools["open_url"] = open_url
         self.tools["go_back"] = go_back
-        self.tools["drag_and_drop"] = drag_and_drop
-        self.tools["find_on_page_ctrl_f"] = find_on_page_ctrl_f
 
-    def take_screenshot_callback(self, memory_step: ActionStep, agent=None) -> None:
-        """Callback that takes a screenshot + memory snapshot after a step completes"""
-        self.logger.log("Analyzing screen content...")
 
+    def take_screenshot_callback(self, memory_step: ActionStep, agent=None) -> None:
+        """Takes a screenshot and saves it to memory"""
         current_step = memory_step.step_number
+        time.sleep(1.0) # Wait for renders
 
-        time.sleep(2.5)  # Let things happen on the desktop
-        screenshot_bytes = self.desktop.screenshot(format="bytes")
-        image = Image.open(BytesIO(screenshot_bytes))
+        image = self.sandbox.get_screenshot()
 
-        # Create a filename with step number
+        # Save to disk
         screenshot_path = os.path.join(self.data_dir, f"step_{current_step:03d}.png")
         image.save(screenshot_path)
 
         image_copy = image.copy()
-
         if getattr(self, "click_coordinates", None):
             image_copy = draw_marker_on_image(image_copy, self.click_coordinates)
 
         self.last_marked_screenshot = AgentImage(screenshot_path)
-        print(f"Saved screenshot for step {current_step} to {screenshot_path}")
-
-        # Optimization: remove previous raw images from memory to save context/speed
+        
+        # Cleanup old images in memory to save RAM
         for previous_memory_step in agent.memory.steps:
-            if (
-                isinstance(previous_memory_step, ActionStep)
-                and previous_memory_step.step_number <= current_step - 1
-            ):
+            if isinstance(previous_memory_step, ActionStep) and previous_memory_step.step_number <= current_step - 1:
                 previous_memory_step.observations_images = None
             elif isinstance(previous_memory_step, TaskStep):
                 previous_memory_step.task_images = None
 
-        # Add the marker-edited image to the current memory step
         memory_step.observations_images = [image_copy]
-        self.click_coordinates = None  # Reset click marker
-
-
-# -----------------------------------------------------------------------------
-# SANDBOX MANAGEMENT & HELPERS
-# -----------------------------------------------------------------------------
-
-def upload_to_hf_and_remove(folder_path):
-    repo_id = "smolagents/computer-agent-logs"
-    try:
-        folder_name = os.path.basename(os.path.normpath(folder_path))
-        print(f"Uploading {folder_path} to {repo_id}/{folder_name}...")
-        url = upload_folder(
-            folder_path=folder_path,
-            repo_id=repo_id,
-            repo_type="dataset",
-            path_in_repo=folder_name,
-            ignore_patterns=[".git/*", ".gitignore"],
-        )
-        print(f"Upload complete. Removing local folder {folder_path}...")
-        shutil.rmtree(folder_path)
-        return url
-    except Exception as e:
-        print(f"Error during upload or cleanup: {str(e)}")
-        # Don't raise, just log, to keep app running
-        return None
+        self.click_coordinates = None
 
-def cleanup_sandboxes():
-    current_time = time.time()
-    sandboxes_to_remove = []
 
-    for session_id, metadata in SANDBOX_METADATA.items():
-        if current_time - metadata["last_accessed"] > SANDBOX_TIMEOUT:
-            sandboxes_to_remove.append(session_id)
-
-    for session_id in sandboxes_to_remove:
-        if session_id in SANDBOXES:
-            try:
-                data_dir = os.path.join(TMP_DIR, session_id)
-                if os.path.exists(data_dir):
-                    shutil.rmtree(data_dir) # Just local cleanup for this demo
-                
-                SANDBOXES[session_id].kill()
-                del SANDBOXES[session_id]
-                del SANDBOX_METADATA[session_id]
-                print(f"Cleaned up sandbox for session {session_id}")
-            except Exception as e:
-                print(f"Error cleaning up sandbox {session_id}: {str(e)}")
-
-def get_or_create_sandbox(session_uuid):
-    current_time = time.time()
-
-    if (
-        session_uuid in SANDBOXES
-        and session_uuid in SANDBOX_METADATA
-        and current_time - SANDBOX_METADATA[session_uuid]["created_at"]
-        < SANDBOX_TIMEOUT
-    ):
-        print(f"Reusing Sandbox for  {session_uuid}")
-        SANDBOX_METADATA[session_uuid]["last_accessed"] = current_time
-        return SANDBOXES[session_uuid]
-    else:
-        print("No sandbox found, creating a new one")
-
-    if session_uuid in SANDBOXES:
-        try:
-            SANDBOXES[session_uuid].kill()
-        except Exception:
-            pass
-
-    print(f"Creating new sandbox for session {session_uuid}")
-    desktop = Sandbox(
-        api_key=E2B_API_KEY,
-        resolution=(WIDTH, HEIGHT),
-        dpi=96,
-        timeout=SANDBOX_TIMEOUT,
-        template="k0wmnzir0zuzye6dndlw",
+def create_agent(data_dir, sandbox):
+    model = FaraLocalModel()
+    return SeleniumVisionAgent(
+        model=model,
+        data_dir=data_dir,
+        sandbox=sandbox,
+        max_steps=30,
+        verbosity_level=2
     )
-    desktop.stream.start(require_auth=True)
-    setup_cmd = """sudo mkdir -p /usr/lib/firefox-esr/distribution && echo '{"policies":{"OverrideFirstRunPage":"","OverridePostUpdatePage":"","DisableProfileImport":true,"DontCheckDefaultBrowser":true}}' | sudo tee /usr/lib/firefox-esr/distribution/policies.json > /dev/null"""
-    desktop.commands.run(setup_cmd)
 
-    SANDBOXES[session_uuid] = desktop
-    SANDBOX_METADATA[session_uuid] = {
-        "created_at": current_time,
-        "last_accessed": current_time,
-    }
-    return desktop
+def generate_interaction_id(session_uuid):
+    return f"{session_uuid}_{int(time.time())}"
+
+def get_agent_summary_erase_images(agent):
+    for memory_step in agent.memory.steps:
+        if hasattr(memory_step, "observations_images"):
+            memory_step.observations_images = None
+        if hasattr(memory_step, "task_images"):
+            memory_step.task_images = None
+    return agent.write_memory_to_messages()
 
 def save_final_status(folder, status: str, summary, error_message=None) -> None:
     try:
@@ -653,142 +600,24 @@ def save_final_status(folder, status: str, summary, error_message=None) -> None:
     except Exception as e:
         print(f"Failed to save metadata: {e}")
 
-def create_agent(data_dir, desktop):
-    # Initialize the wrapper that calls the global GPU function
-    model = FaraLocalModel()
-    
-    return E2BVisionAgent(
-        model=model,
-        data_dir=data_dir,
-        desktop=desktop,
-        max_steps=200,
-        verbosity_level=2,
-        use_v1_prompt=True,
-    )
-
-def generate_interaction_id(session_uuid):
-    return f"{session_uuid}_{int(time.time())}"
-
 # -----------------------------------------------------------------------------
-# GRADIO UI & INTERACTION
+# UI & APP
 # -----------------------------------------------------------------------------
 
+# CSS
 custom_css = """
 .modal-container { margin: var(--size-16) auto!important; }
-.sandbox-container { position: relative; width: 910px; height: 800px; overflow: hidden; margin: auto; }
-.sandbox-frame { display: none; position: absolute; top: 0; left: 0; width: 910px; height: 800px; pointer-events:none; }
-.sandbox-iframe, .bsod-image { position: absolute; width: <<WIDTH>>px; height: <<HEIGHT>>px; border: 4px solid #444444; transform-origin: 0 0; }
-.primary-color-label label span { font-weight: bold; color: var(--color-accent); }
-.status-bar { display: flex; flex-direction: row; align-items: center; z-index: 100; }
-.status-indicator { width: 15px; height: 15px; border-radius: 50%; }
-.status-text { font-size: 16px; font-weight: bold; padding-left: 8px; text-shadow: none; }
-.status-interactive { background-color: #2ecc71; animation: blink 2s infinite; }
-.status-view-only { background-color: #e74c3c; }
-.status-error { background-color: #e74c3c; animation: blink-error 1s infinite; }
-@keyframes blink-error { 0% { background-color: rgba(231, 76, 60, 1); } 50% { background-color: rgba(231, 76, 60, 0.4); } 100% { background-color: rgba(231, 76, 60, 1); } }
-@keyframes blink { 0% { background-color: rgba(46, 204, 113, 1); } 50% { background-color: rgba(46, 204, 113, 0.4); } 100% { background-color: rgba(46, 204, 113, 1); } }
-#chatbot { height:1000px!important; }
-#chatbot .role { max-width:95% }
-.logo-container { display: flex; flex-direction: column; align-items: flex-start; gap: 5px; }
-.logo-item { display: flex; align-items: center; padding: 0 30px; gap: 10px; text-decoration: none!important; color: #f59e0b; font-size:17px; }
-""".replace("<<WIDTH>>", str(WIDTH + 15)).replace("<<HEIGHT>>", str(HEIGHT + 10))
-
-sandbox_html_template = """
-<style>
-@import url('https://fonts.googleapis.com/css2?family=Oxanium:wght@200..800&display=swap');
-</style>
-    <h1 style="color:var(--color-accent);margin:0;">Fara CUA - <i>Powered by <a href="https://github.com/huggingface/smolagents">smolagents</a></i><h1>
-<div class="sandbox-container" style="margin:0;">
-    <div class="status-bar">
-        <div class="status-indicator {status_class}"></div>
-        <div class="status-text">{status_text}</div>
-    </div>
-    <iframe id="sandbox-iframe"
-        src="{stream_url}" 
-        class="sandbox-iframe"
-        style="display: block;"
-        allowfullscreen>
-    </iframe>
-    <img src="https://huggingface.co/datasets/mfarre/servedfiles/resolve/main/blue_screen_of_death.gif" class="bsod-image" style="display: none;"/>
-    <img src="https://huggingface.co/datasets/m-ric/images/resolve/main/HUD_thom.png" class="sandbox-frame" />
+.browser-container { position: relative; width: 100%; height: 600px; border: 1px solid #444; background: #222; display: flex; align-items: center; justify-content: center; overflow: hidden; }
+.browser-image { max-width: 100%; max-height: 100%; object-fit: contain; }
+#chatbot { height: 800px!important; }
+"""
+
+# HTML Template for the output area (Replaced IFrame with simple image holder logic handled by Gradio Image)
+browser_html = """
+<div class="browser-container">
+    <p style="color: #888;">Browser Screenshot will appear here after steps.</p>
 </div>
-""".replace("<<WIDTH>>", str(WIDTH + 15)).replace("<<HEIGHT>>", str(HEIGHT + 10))
-
-custom_js = """function() {
-    document.body.classList.add('dark');
-    // Function to check if sandbox is timing out
-    const checkSandboxTimeout = function() {
-        const timeElement = document.getElementById('sandbox-creation-time');
-        if (timeElement) {
-            const creationTime = parseFloat(timeElement.getAttribute('data-time'));
-            const timeoutValue = parseFloat(timeElement.getAttribute('data-timeout'));
-            const currentTime = Math.floor(Date.now() / 1000);
-            if (currentTime - creationTime >= timeoutValue) {
-                showBSOD('Error');
-                return;
-            }
-        }
-        setTimeout(checkSandboxTimeout, 5000);
-    };
-    const showBSOD = function(statusText = 'Error') {
-        const iframe = document.getElementById('sandbox-iframe');
-        const bsod = document.getElementById('bsod-image');
-        if (iframe && bsod) {
-            iframe.style.display = 'none';
-            bsod.style.display = 'block';
-            document.querySelector('.status-indicator').className = 'status-indicator status-error';
-            document.querySelector('.status-text').innerText = statusText;
-        }
-    };
-    const resetBSOD = function() {
-        const iframe = document.getElementById('sandbox-iframe');
-        const bsod = document.getElementById('bsod-image');
-        if (iframe && bsod && bsod.style.display === 'block') {
-            iframe.style.display = 'block';
-            bsod.style.display = 'none';
-        }
-    };
-    document.addEventListener('click', function(e) {
-        if (e.target.tagName === 'BUTTON' && e.target.innerText === "Let's go!") {
-            resetBSOD();
-        }
-    });
-    checkSandboxTimeout();
-    const params = new URLSearchParams(window.location.search);
-    if (!params.has('__theme')) {
-        params.set('__theme', 'dark');
-        window.location.search = params.toString();
-    }
-}"""
-
-def update_html(interactive_mode: bool, session_uuid):
-    desktop = get_or_create_sandbox(session_uuid)
-    auth_key = desktop.stream.get_auth_key()
-    base_url = desktop.stream.get_url(auth_key=auth_key)
-    stream_url = base_url if interactive_mode else f"{base_url}&view_only=true"
-
-    status_class = "status-interactive" if interactive_mode else "status-view-only"
-    status_text = "Interactive" if interactive_mode else "Agent running..."
-    creation_time = (
-        SANDBOX_METADATA[session_uuid]["created_at"]
-        if session_uuid in SANDBOX_METADATA
-        else time.time()
-    )
-
-    sandbox_html_content = sandbox_html_template.format(
-        stream_url=stream_url,
-        status_class=status_class,
-        status_text=status_text,
-    )
-    sandbox_html_content += f'<div id="sandbox-creation-time" style="display:none;" data-time="{creation_time}" data-timeout="{SANDBOX_TIMEOUT}"></div>'
-    return sandbox_html_content
-
-def initialize_session(interactive_mode, browser_uuid):
-    if not browser_uuid:
-        new_uuid = str(uuid.uuid4())
-        return update_html(interactive_mode, new_uuid), new_uuid
-    else:
-        return update_html(interactive_mode, browser_uuid), browser_uuid
+"""
 
 class EnrichedGradioUI(GradioUI):
     def interact_with_agent(
@@ -801,35 +630,30 @@ class EnrichedGradioUI(GradioUI):
         request: gr.Request,
     ):
         interaction_id = generate_interaction_id(session_uuid)
-        desktop = get_or_create_sandbox(session_uuid)
-
         data_dir = os.path.join(TMP_DIR, interaction_id)
-        if not os.path.exists(data_dir):
-            os.makedirs(data_dir)
-
-        # Create fresh agent. 
-        # Note: We do NOT store the full agent in session_state passed between Gradio events
-        # if possible, or if we do, we ensure this function isn't wrapped in @spaces.GPU
-        agent = create_agent(data_dir=data_dir, desktop=desktop)
-        session_state["agent"] = agent # Storing in state is fine if this function runs on CPU
+        
+        # Initialize Sandbox per run (clean slate)
+        sandbox = SeleniumSandbox(width=WIDTH, height=HEIGHT)
+        agent = create_agent(data_dir=data_dir, sandbox=sandbox)
+        session_state["agent"] = agent
 
         try:
             stored_messages.append(gr.ChatMessage(role="user", content=task_input))
-            yield stored_messages
-
-            screenshot_bytes = agent.desktop.screenshot(format="bytes")
-            initial_screenshot = Image.open(BytesIO(screenshot_bytes))
+            yield stored_messages, None
 
+            # Initial screenshot
+            screenshot = sandbox.get_screenshot()
+            
+            # Run Agent
             for msg in stream_to_gradio(
                 agent,
                 task=task_input,
-                task_images=[initial_screenshot],
+                task_images=[screenshot],
                 reset_agent_memory=False,
             ):
-                if (
-                    hasattr(agent, "last_marked_screenshot")
-                    and msg.content == "-----"
-                ):
+                # Update Chat
+                if hasattr(agent, "last_marked_screenshot") and msg.content == "-----":
+                     # Add image to chat
                     stored_messages.append(
                         gr.ChatMessage(
                             role="assistant",
@@ -839,13 +663,17 @@ class EnrichedGradioUI(GradioUI):
                             },
                         )
                     )
-                stored_messages.append(msg)
-                yield stored_messages
+                    # Yield updated chat AND the latest screenshot for the side panel
+                    yield stored_messages, agent.last_marked_screenshot.to_string()
+                else:
+                    stored_messages.append(msg)
+                    yield stored_messages, None # Keep previous side image
 
             if consent_storage:
                 summary = get_agent_summary_erase_images(agent)
                 save_final_status(data_dir, "completed", summary=summary)
-            yield stored_messages
+            
+            yield stored_messages, None
 
         except Exception as e:
             error_message = f"Error in interaction: {str(e)}"
@@ -853,121 +681,87 @@ class EnrichedGradioUI(GradioUI):
             stored_messages.append(
                 gr.ChatMessage(role="assistant", content="Run failed:\n" + error_message)
             )
-            if consent_storage:
-                save_final_status(data_dir, "failed", summary=None, error_message=error_message)
-            yield stored_messages
-
-# -----------------------------------------------------------------------------
-# MAIN APP CONSTRUCTION
-# -----------------------------------------------------------------------------
+            yield stored_messages, None
+        finally:
+            # Important: Cleanup Sandbox
+            sandbox.cleanup()
 
+# Gradio Block Construction
 theme = gr.themes.Default(
     font=["Oxanium", "sans-serif"], primary_hue="amber", secondary_hue="blue"
 )
 
-with gr.Blocks(theme=theme, css=custom_css, js=custom_js) as demo:
-    session_uuid_state = gr.State(None)
+with gr.Blocks(theme=theme, css=custom_css) as demo:
+    session_uuid_state = gr.State(lambda: str(uuid.uuid4()))
+    session_state = gr.State({})
+    stored_messages = gr.State([])
     
     with gr.Row():
-        sandbox_html = gr.HTML(
-            value=sandbox_html_template.format(
-                stream_url="",
-                status_class="status-interactive",
-                status_text="Interactive",
-            ),
-            label="Output",
-        )
-        with gr.Sidebar(position="left"):
-            with Modal(visible=True) as modal:
-                gr.Markdown("""### Welcome to Fara CUA Demo 🖥️
-This agent uses **microsoft/Fara-7B** (running locally via ZeroGPU) and **smolagents** to control a remote computer.
-
-👉 Type a task, click 'Let's go!', and watch the agent work.
-""")
+        # Left Sidebar: Inputs
+        with gr.Column(scale=1):
+            gr.Markdown("### Fara CUA - Chrome Agent 🌐")
+            
             task_input = gr.Textbox(
-                value="Find me pictures of cute puppies",
-                label="Enter your task below:",
-                elem_classes="primary-color-label",
+                value="Go to google.com and search for 'Hugging Face'",
+                label="Task",
+                lines=3
             )
-
-            run_btn = gr.Button("Let's go!", variant="primary")
+            run_btn = gr.Button("Start Task", variant="primary")
+            stop_btn = gr.Button("Stop", variant="secondary")
             
-            # Simple controls
-            stop_btn = gr.Button("Stop the agent!", variant="secondary")
-            consent_storage = gr.Checkbox(label="Store logs locally?", value=True)
+            consent_storage = gr.Checkbox(label="Save logs locally?", value=True)
             
             gr.Examples(
                 examples=[
-                    "Use Google Maps to find the Hugging Face HQ in Paris",
-                    "Go to Wikipedia and find what happened on April 4th",
-                    "Find out the travel time by train from Bern to Basel on Google Maps",
+                    "Go to google.com and search for 'Hugging Face', then click the first link.",
+                    "Go to wikipedia.org, type 'Python' in search, and click the search button.",
+                    "Open huggingface.co and find the 'Spaces' link."
                 ],
-                inputs=task_input,
+                inputs=task_input
             )
 
-            session_state = gr.State({})
-            stored_messages = gr.State([])
-
-    chatbot_display = gr.Chatbot(
-        elem_id="chatbot",
-        label="Agent's execution logs",
-        type="messages",
-        avatar_images=(None, "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png"),
-        resizable=True,
-    )
+        # Right Main: Chat & Live View
+        with gr.Column(scale=3):
+            with gr.Row():
+                # Side-by-side: Chat and Latest Screenshot
+                with gr.Column(scale=1):
+                    chatbot_display = gr.Chatbot(
+                        label="Agent Trace",
+                        type="messages",
+                        height=800,
+                        avatar_images=(None, "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png"),
+                    )
+                with gr.Column(scale=1):
+                    gr.Markdown("### Latest Browser View")
+                    live_browser_view = gr.Image(
+                        label="Browser View", 
+                        type="filepath", 
+                        interactive=False,
+                        height=600
+                    )
 
-    # Dummy agent init for UI wrapper (actual agent created in interaction loop)
-    # We pass a dummy CodeAgent just to initialize the UI class
+    # UI Handler
     agent_ui = EnrichedGradioUI(CodeAgent(tools=[], model=Model(), name="init"))
 
-    is_interactive = gr.Checkbox(value=True, visible=False)
-
-    def clear_and_set_view_only(task_input, session_uuid):
-        return update_html(False, session_uuid)
-
-    def set_interactive(session_uuid):
-        return update_html(True, session_uuid)
-    
     def interrupt_agent(session_state):
-        if "agent" in session_state and hasattr(session_state["agent"], "interrupt_switch") and not session_state["agent"].interrupt_switch:
-            session_state["agent"].interrupt()
-            return "Stopped"
-        return "Stop"
-
-    # Event Wiring
-    run_event = (
-        run_btn.click(
-            fn=clear_and_set_view_only,
-            inputs=[task_input, session_uuid_state],
-            outputs=[sandbox_html],
-        )
-        .then(
-            agent_ui.interact_with_agent,
-            inputs=[
-                task_input,
-                stored_messages,
-                session_state,
-                session_uuid_state,
-                consent_storage,
-            ],
-            outputs=[chatbot_display],
-        )
-        .then(fn=set_interactive, inputs=[session_uuid_state], outputs=[sandbox_html])
+        if "agent" in session_state and hasattr(session_state["agent"], "interrupt_switch"):
+            session_state["agent"].interrupt_switch = True
+            return "Interrupted"
+    
+    # Run Logic
+    run_event = run_btn.click(
+        fn=agent_ui.interact_with_agent,
+        inputs=[
+            task_input,
+            stored_messages,
+            session_state,
+            session_uuid_state,
+            consent_storage,
+        ],
+        outputs=[chatbot_display, live_browser_view]
     )
     
     stop_btn.click(fn=interrupt_agent, inputs=[session_state], outputs=[])
 
-    # Initialization on load
-    demo.load(
-        fn=lambda: True,
-        outputs=[is_interactive],
-    ).then(
-        fn=initialize_session,
-        js="() => localStorage.getItem('gradio-session-uuid') || (() => { const id = self.crypto.randomUUID(); localStorage.setItem('gradio-session-uuid', id); return id })()",
-        inputs=[is_interactive],
-        outputs=[sandbox_html, session_uuid_state],
-    )
-
 if __name__ == "__main__":
-    Timer(60, cleanup_sandboxes).start()
     demo.launch()