Update app.py

app.py

@@ -1,64 +1,166 @@
+import os
 import gradio as gr
-from huggingface_hub import InferenceClient
+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
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+# 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.")
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+if not gemini_api_key:
+    raise ValueError("GEMINI_API_KEY environment variable not set.")
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+# 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
 
-    messages.append({"role": "user", "content": message})
+@tool
+def go_back() -> None:
+    driver.back()
 
-    response = ""
+@tool
+def close_popups() -> str:
+    webdriver.ActionChains(driver).send_keys(Keys.ESCAPE).perform()
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+# 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")
 
-        response += token
-        yield response
+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
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
+# 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]
+    )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()