update app

app.py

@@ -0,0 +1,928 @@
+import json
+import os
+import shutil
+import time
+import uuid
+import unicodedata
+from io import BytesIO
+from threading import Timer
+from typing import Any, Dict, List, Optional
+from datetime import datetime
+
+import gradio as gr
+import torch
+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
+
+# Smolagents imports
+from smolagents import CodeAgent, tool, AgentImage
+from smolagents.memory import ActionStep, TaskStep
+from smolagents.models import ChatMessage, Model, MessageRole
+from smolagents.gradio_ui import GradioUI, stream_to_gradio
+from smolagents.monitoring import LogLevel
+
+# Transformers for Fara Model
+from transformers import (
+    Qwen2_5_VLForConditionalGeneration,
+    AutoProcessor,
+)
+from qwen_vl_utils import process_vision_info
+
+load_dotenv(override=True)
+
+# -----------------------------------------------------------------------------
+# CONFIGURATION & CONSTANTS
+# -----------------------------------------------------------------------------
+
+E2B_API_KEY = os.getenv("E2B_API_KEY")
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_API_KEY")
+if HF_TOKEN:
+    login(token=HF_TOKEN)
+
+SANDBOXES = {}
+SANDBOX_METADATA = {}
+SANDBOX_TIMEOUT = 600
+WIDTH = 1024
+HEIGHT = 768
+TMP_DIR = "./tmp/"
+if not os.path.exists(TMP_DIR):
+    os.makedirs(TMP_DIR)
+
+# -----------------------------------------------------------------------------
+# MODEL INITIALIZATION (Fara-7B / Qwen2.5-VL)
+# -----------------------------------------------------------------------------
+
+print("Loading Fara Model... This may take a moment.")
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+MODEL_ID_F = "microsoft/Fara-7B"  # Ensure this ID is accessible or point to local path
+
+try:
+    processor_f = AutoProcessor.from_pretrained(MODEL_ID_F, trust_remote_code=True)
+    model_f = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+        MODEL_ID_F,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
+        device_map="auto" if DEVICE == "cuda" else None,
+    )
+    if DEVICE == "cpu":
+        model_f.to(DEVICE)
+    
+    model_f.eval()
+    print(f"Fara Model loaded successfully on {DEVICE}")
+except Exception as e:
+    print(f"Error loading Fara Model: {e}")
+    print("Please ensure you have access to the model and enough GPU memory.")
+    # Fallback to prevent crash during import, though app won't work without model
+    model_f = None
+    processor_f = None
+
+class FaraLocalModel(Model):
+    """
+    Wrapper for the local Fara (Qwen2.5-VL) model to work with SmolAgents.
+    """
+    def __init__(self, model, processor, **kwargs):
+        super().__init__(**kwargs)
+        self.model = model
+        self.processor = processor
+
+    def __call__(
+        self,
+        messages: List[Dict[str, Any]],
+        stop_sequences: Optional[List[str]] = None,
+        **kwargs,
+    ) -> ChatMessage:
+        if self.model is None:
+            raise ValueError("Fara Model is not loaded.")
+
+        # Convert SmolAgents messages to Qwen/Transformers format
+        # SmolAgents uses a specific dict structure for content. 
+        # We need to normalize it for process_vision_info / apply_chat_template
+        
+        formatted_messages = []
+        
+        for msg in messages:
+            role = msg["role"]
+            content = msg["content"]
+            
+            new_content = []
+            
+            if isinstance(content, str):
+                new_content.append({"type": "text", "text": content})
+            elif isinstance(content, list):
+                for item in content:
+                    if isinstance(item, str):
+                        new_content.append({"type": "text", "text": item})
+                    elif isinstance(item, dict):
+                        if "type" in item:
+                            if item["type"] == "image":
+                                # Handle path or url
+                                val = item.get("image") or item.get("url") or item.get("path")
+                                new_content.append({"type": "image", "image": val})
+                            else:
+                                new_content.append(item)
+            
+            formatted_messages.append({"role": role, "content": new_content})
+
+        # Process Inputs
+        text = self.processor.apply_chat_template(
+            formatted_messages, tokenize=False, add_generation_prompt=True
+        )
+        
+        image_inputs, video_inputs = process_vision_info(formatted_messages)
+        
+        inputs = self.processor(
+            text=[text],
+            images=image_inputs,
+            videos=video_inputs,
+            padding=True,
+            return_tensors="pt",
+        )
+        
+        inputs = inputs.to(self.model.device)
+
+        # Generate
+        with torch.no_grad():
+            generated_ids = self.model.generate(
+                **inputs,
+                max_new_tokens=kwargs.get("max_tokens", 1024),
+                stop_strings=stop_sequences,
+                tokenizer=self.processor.tokenizer, # Specific for stop_strings in modern transformers
+            )
+
+        # Decode
+        generated_ids_trimmed = [
+            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        output_text = self.processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )[0]
+
+        return ChatMessage(
+            role=MessageRole.ASSISTANT,
+            content=output_text,
+        )
+
+# -----------------------------------------------------------------------------
+# E2B AGENT & TOOLS
+# -----------------------------------------------------------------------------
+
+E2B_SYSTEM_PROMPT_TEMPLATE = """You are a desktop automation assistant that can control a remote desktop environment. 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. 
+Then you will proceed as follows, with these sections: don't skip any!
+
+Short term goal: ...
+What I see: ...
+Reflection: ...
+Action:
+```python
+click(254, 308)
+```<end_code>
+
+Akways 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:
+{%- for tool in tools.values() %}
+- {{ tool.name }}: {{ tool.description }}
+    Takes inputs: {{tool.inputs}}
+    Returns an output of type: {{tool.output_type}}
+{%- endfor %}
+</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.
+</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.
+</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)
+    cross_size, linewidth = 10, 3
+    # 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,
+        ),
+        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"""
+
+    def __init__(
+        self,
+        model: Model,
+        data_dir: str,
+        desktop: Sandbox,
+        tools: List[tool] = None,
+        max_steps: int = 200,
+        verbosity_level: LogLevel = 2,
+        planning_interval: int = None,
+        use_v1_prompt: bool = False,
+        **kwargs,
+    ):
+        self.desktop = desktop
+        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
+        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 [],
+            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
+
+        # Add default tools
+        self.logger.log("Setting up agent tools...")
+        self._setup_desktop_tools()
+        self.step_callbacks.append(self.take_screenshot_callback)
+
+    def _setup_desktop_tools(self):
+        """Register all desktop tools"""
+
+        @tool
+        def click(x: int, y: int) -> str:
+            """
+            Performs a left-click at the specified coordinates
+            Args:
+                x: The x coordinate (horizontal position)
+                y: The y coordinate (vertical position)
+            """
+            self.desktop.move_mouse(x, y)
+            self.desktop.left_click()
+            self.click_coordinates = [x, y]
+            self.logger.log(f"Clicked at coordinates ({x}, {y})")
+            return f"Clicked at coordinates ({x}, {y})"
+
+        @tool
+        def right_click(x: int, y: int) -> str:
+            """
+            Performs a right-click at the specified coordinates
+            Args:
+                x: The x coordinate (horizontal position)
+                y: The y coordinate (vertical position)
+            """
+            self.desktop.move_mouse(x, y)
+            self.desktop.right_click()
+            self.click_coordinates = [x, y]
+            self.logger.log(f"Right-clicked at coordinates ({x}, {y})")
+            return f"Right-clicked at coordinates ({x}, {y})"
+
+        @tool
+        def double_click(x: int, y: int) -> str:
+            """
+            Performs a double-click at the specified coordinates
+            Args:
+                x: The x coordinate (horizontal position)
+                y: The y coordinate (vertical position)
+            """
+            self.desktop.move_mouse(x, y)
+            self.desktop.double_click()
+            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)
+            )
+
+        @tool
+        def type_text(text: str) -> str:
+            """
+            Types the specified text at the current cursor position.
+            Args:
+                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}'")
+            return f"Typed text: '{clean_text}'"
+
+        @tool
+        def press_key(key: str) -> str:
+            """
+            Presses a keyboard key
+            Args:
+                key: The key to press (e.g. "enter", "space", "backspace", etc.).
+            """
+            self.desktop.press(key)
+            self.logger.log(f"Pressed 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.
+            """
+            self.desktop.drag([x1, y1], [x2, y2])
+            message = f"Dragged and dropped from [{x1}, {y1}] to [{x2}, {y2}]"
+            self.logger.log(message)
+            return message
+
+        @tool
+        def scroll(x: int, y: int, direction: str = "down", amount: int = 2) -> str:
+            """
+            Moves the mouse to selected coordinates, then uses the scroll button.
+            Args:
+                x: The x coordinate
+                y: The y coordinate
+                direction: "up" or "down"
+                amount: The amount to scroll.
+            """
+            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
+
+        @tool
+        def wait(seconds: float) -> str:
+            """
+            Waits for the specified number of seconds.
+            """
+            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.
+            """
+            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}"
+
+        @tool
+        def find_on_page_ctrl_f(search_string: str) -> str:
+            """
+            Scroll the browser viewport to the first occurrence of the search string (Ctrl+F).
+            """
+            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.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["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...")
+
+        current_step = memory_step.step_number
+
+        time.sleep(2.5)  # Let things happen on the desktop
+        screenshot_bytes = self.desktop.screenshot(format="bytes")
+        image = Image.open(BytesIO(screenshot_bytes))
+
+        # Create a filename with step number
+        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
+        for previous_memory_step in agent.memory.steps:
+            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
+
+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",
+    )
+    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 save_final_status(folder, status: str, summary, error_message=None) -> None:
+    try:
+        with open(os.path.join(folder, "metadata.json"), "w") as output_file:
+            output_file.write(
+                json.dumps(
+                    {"status": status, "summary": summary, "error_message": error_message},
+                    default=str
+                )
+            )
+    except Exception as e:
+        print(f"Failed to save metadata: {e}")
+
+def create_agent(data_dir, desktop):
+    # Instantiate the local model wrapper
+    if model_f is None:
+        raise RuntimeError("Fara model was not loaded successfully.")
+    
+    model = FaraLocalModel(model=model_f, processor=processor_f)
+    
+    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
+# -----------------------------------------------------------------------------
+
+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" />
+</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(
+        self,
+        task_input,
+        stored_messages,
+        session_state,
+        session_uuid,
+        consent_storage,
+        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)
+
+        # Re-create agent to ensure fresh context with the Fara model
+        session_state["agent"] = create_agent(data_dir=data_dir, desktop=desktop)
+
+        try:
+            stored_messages.append(gr.ChatMessage(role="user", content=task_input))
+            yield stored_messages
+
+            screenshot_bytes = session_state["agent"].desktop.screenshot(format="bytes")
+            initial_screenshot = Image.open(BytesIO(screenshot_bytes))
+
+            for msg in stream_to_gradio(
+                session_state["agent"],
+                task=task_input,
+                task_images=[initial_screenshot],
+                reset_agent_memory=False,
+            ):
+                if (
+                    hasattr(session_state["agent"], "last_marked_screenshot")
+                    and msg.content == "-----"
+                ):
+                    stored_messages.append(
+                        gr.ChatMessage(
+                            role="assistant",
+                            content={
+                                "path": session_state["agent"].last_marked_screenshot.to_string(),
+                                "mime_type": "image/png",
+                            },
+                        )
+                    )
+                stored_messages.append(msg)
+                yield stored_messages
+
+            if consent_storage:
+                summary = get_agent_summary_erase_images(session_state["agent"])
+                save_final_status(data_dir, "completed", summary=summary)
+            yield stored_messages
+
+        except Exception as e:
+            error_message = f"Error in interaction: {str(e)}"
+            print(error_message)
+            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
+# -----------------------------------------------------------------------------
+
+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.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) and **smolagents** to control a remote computer.
+
+👉 Type a task, click 'Let's go!', and watch the agent work.
+""")
+            task_input = gr.Textbox(
+                value="Find me pictures of cute puppies",
+                label="Enter your task below:",
+                elem_classes="primary-color-label",
+            )
+
+            run_btn = gr.Button("Let's go!", variant="primary")
+            
+            # Simple controls
+            stop_btn = gr.Button("Stop the agent!", variant="secondary")
+            consent_storage = gr.Checkbox(label="Store 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",
+                ],
+                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,
+    )
+
+    # Dummy agent init for UI wrapper (actual agent created in interaction loop)
+    # We pass a dummy CodeAgent just to initialize the UI class
+    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 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])
+    )
+    
+    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()