app.py

import os
import gradio as gr
from smolagents import CodeAgent, LiteLLMModel, tool
from smolagents.agents import ActionStep
import helium
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from io import BytesIO
from PIL import Image
from datetime import datetime
from dotenv import load_dotenv
from huggingface_hub import login
import tempfile

# Load environment variables
load_dotenv()
hf_token = os.getenv("HF_TOKEN")
gemini_api_key = os.getenv("GEMINI_API_KEY")

if hf_token:
    login(hf_token, add_to_git_credential=False)
else:
    raise ValueError("HF_TOKEN environment variable not set.")

if not gemini_api_key:
    raise ValueError("GEMINI_API_KEY environment variable not set.")

# Define tools
@tool
def search_item_ctrl_f(text: str, nth_result: int = 1) -> str:
    elements = driver.find_elements(By.XPATH, f"//*[contains(text(), '{text}')]")
    if nth_result > len(elements):
        raise Exception(f"Match n°{nth_result} not found (only {len(elements)} matches found)")
    result = f"Found {len(elements)} matches for '{text}'."
    elem = elements[nth_result - 1]
    driver.execute_script("arguments[0].scrollIntoView(true);", elem)
    result += f"Focused on element {nth_result} of {len(elements)}"
    return result

@tool
def go_back() -> None:
    driver.back()

@tool
def close_popups() -> str:
    webdriver.ActionChains(driver).send_keys(Keys.ESCAPE).perform()

# Initialize Chrome driver
chrome_options = webdriver.ChromeOptions()
chrome_options.add_argument("--force-device-scale-factor=1")
chrome_options.add_argument("--window-size=1000,1350")
chrome_options.add_argument("--disable-pdf-viewer")
chrome_options.add_argument("--no-sandbox")
chrome_options.add_argument("--disable-dev-shm-usage")
chrome_options.add_argument("--window-position=0,0")
chrome_options.add_argument("--headless")

driver = helium.start_chrome(headless=True, options=chrome_options)

# Screenshot callback
def save_screenshot(memory_step: ActionStep, agent: CodeAgent) -> Image.Image:
    from time import sleep
    sleep(1.0)
    driver = helium.get_driver()
    current_step = memory_step.step_number
    if driver is not None:
        for previous_memory_step in agent.memory.steps:
            if isinstance(previous_memory_step, ActionStep) and previous_memory_step.step_number <= current_step - 2:
                previous_memory_step.observations_images = None
        png_bytes = driver.get_screenshot_as_png()
        image = Image.open(BytesIO(png_bytes))
        screenshot_dir = tempfile.gettempdir()
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        screenshot_path = f"{screenshot_dir}/screenshot_step_{current_step}_{timestamp}.png"
        image.save(screenshot_path)
        url_info = f"Current url: {driver.current_url}"
        memory_step.observations = (
            url_info if memory_step.observations is None else memory_step.observations + "\n" + url_info
        )
        return image

# Initialize model and agent
model = LiteLLMModel("gemini/gemini-2.0-flash")
agent = CodeAgent(
    tools=[go_back, close_popups, search_item_ctrl_f],
    model=model,
    additional_authorized_imports=["helium"],
    step_callbacks=[save_screenshot],
    max_steps=20,
    verbosity_level=2,
)
agent.python_executor("from helium import *")

# Helium instructions
helium_instructions = """
You can use helium to access websites. Don't bother about the helium driver, it's already managed.
We've already ran "from helium import *"
Then you can go to pages!
Code:
go_to('github.com/trending')
```<end_code>

You can directly click clickable elements by inputting the text that appears on them.
Code:
click("Top products")
```<end_code>

If it's a link:
Code:
click(Link("Top products"))
```<end_code>

If you try to interact with an element and it's not found, you'll get a LookupError.
In general stop your action after each button click to see what happens on your screenshot.
Never try to login in a page.

To scroll up or down, use scroll_down or scroll_up with as an argument the number of pixels to scroll from.
Code:
scroll_down(num_pixels=1200) # This will scroll one viewport down
```<end_code>

When you have pop-ups with a cross icon to close, don't try to click the close icon by finding its element or targeting an 'X' element (this most often fails).
Just use your built-in tool `close_popups` to close them:
Code:
close_popups()
```<end_code>

You can use .exists() to check for the existence of an element. For example:
Code:
if Text('Accept cookies?').exists():
    click('I accept')
```<end_code>
"""

# Gradio interface function
def run_agent(url: str, request: str):
    try:
        search_request = f"Please go to {url}. {request}"
        agent_output = agent.run(search_request + helium_instructions)
        screenshot_path = next(
            (f"{tempfile.gettempdir()}/screenshot_step_{step.step_number}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
             for step in agent.memory.steps if isinstance(step, ActionStep) and step.observations_images),
            None
        )
        return agent_output, screenshot_path
    except Exception as e:
        return f"Error: {str(e)}", None
    finally:
        driver.quit()

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Web Navigation Agent")
    url_input = gr.Textbox(label="Enter URL", placeholder="https://example.com")
    request_input = gr.Textbox(label="Enter Request", placeholder="Describe what to do on the website")
    submit_button = gr.Button("Run Agent")
    output_text = gr.Textbox(label="Agent Output")
    output_image = gr.Image(label="Screenshot")
    submit_button.click(
        fn=run_agent,
        inputs=[url_input, request_input],
        outputs=[output_text, output_image]
    )

demo.launch()