Update app.py

app.py

@@ -1,18 +1,24 @@
 import os
-import gradio as gr
-from smolagents import CodeAgent, LiteLLMModel, tool
-from smolagents.agents import ActionStep
-import helium
+import json
+import yaml
+import logging
+from dotenv import load_dotenv
+from huggingface_hub import login
 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
-import logging
+import helium
+from smolagents import CodeAgent, LiteLLMModel
+from smolagents.agents import ActionStep
+from tools.search_item_ctrl_f import SearchItemCtrlFTool
+from tools.go_back import GoBackTool
+from tools.close_popups import ClosePopupsTool
+from tools.final_answer import FinalAnswerTool
+from GRADIO_UI import GradioUI
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -30,44 +36,6 @@ if not gemini_api_key:
 
 login(hf_token, add_to_git_credential=False)
 
-# Debug ChromeDriver path
-chromedriver_path = '/usr/bin/chromedriver'
-logger.info(f"Checking ChromeDriver at: {chromedriver_path}")
-logger.info(f"ChromeDriver exists: {os.path.exists(chromedriver_path)}")
-logger.info(f"ChromeDriver executable: {os.access(chromedriver_path, os.X_OK)}")
-logger.info(f"System PATH: {os.environ.get('PATH')}")
-
-# Define tools
-@tool
-def search_item_ctrl_f(text: str, nth_result: int = 1) -> str:
-    """
-    Searches for text on the current page via Ctrl + F and jumps to the nth occurrence.
-    Args:
-        text: The text to search for
-        nth_result: Which occurrence to jump to (default: 1)
-    """
-    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:
-    """Goes back to previous page."""
-    driver.back()
-
-@tool
-def close_popups() -> str:
-    """
-    Closes any visible modal or pop-up on the page. Use this to dismiss pop-up windows!
-    This does not work on cookie consent banners.
-    """
-    webdriver.ActionChains(driver).send_keys(Keys.ESCAPE).perform()
-
 # Initialize Chrome driver
 try:
     chrome_options = webdriver.ChromeOptions()
@@ -113,108 +81,53 @@ def save_screenshot(memory_step: ActionStep, agent: CodeAgent) -> str:
         )
         return screenshot_path
 
-# Initialize model and agent
-model = LiteLLMModel("gemini/gemini-2.0-flash")
+# Load prompt templates
+try:
+    with open("prompts.yaml", 'r') as stream:
+        prompt_templates = yaml.safe_load(stream)
+except FileNotFoundError:
+    prompt_templates = {}
+
+# Initialize tools
+tools = [
+    SearchItemCtrlFTool(driver=driver),
+    GoBackTool(driver=driver),
+    ClosePopupsTool(driver=driver),
+    FinalAnswerTool()
+]
+
+# Initialize model
+model = LiteLLMModel(model_name="gemini/gemini-2.0-flash", api_key=gemini_api_key, max_tokens=2096, temperature=0.5)
+
+# Initialize agent
 agent = CodeAgent(
-    tools=[go_back, close_popups, search_item_ctrl_f],
     model=model,
-    additional_authorized_imports=["helium"],
-    step_callbacks=[save_screenshot],
+    tools=tools,
     max_steps=20,
     verbosity_level=2,
+    prompt_templates=prompt_templates,
+    step_callbacks=[save_screenshot],
+    authorized_imports=[
+        "helium",
+        "unicodedata",
+        "stat",
+        "datetime",
+        "random",
+        "pandas",
+        "itertools",
+        "math",
+        "statistics",
+        "queue",
+        "time",
+        "collections",
+        "re"
+    ]
 )
 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)
-```<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.
-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>
-"""
-
-# Chatbot interface function
-def run_agent_chat(user_input: str, history: list):
-    try:
-        # Extract URL and request from user input or history
-        if "http" in user_input:
-            url = user_input.split()[0] if user_input.startswith("http") else next((w for w in user_input.split() if w.startswith("http")), "")
-            request = user_input.replace(url, "").strip() or "Navigate to the URL and describe the page."
-        else:
-            url = "[invalid url, do not cite]"
-            request = user_input
-
-        search_request = f"Please go to {url}. {request}"
-        agent_output = agent.run(search_request + helium_instructions)
-        
-        # Collect the latest screenshot path from observations
-        latest_screenshot = None
-        for step in reversed(agent.memory.steps):
-            if isinstance(step, ActionStep) and step.observations:
-                # Extract screenshot path from observations
-                for line in step.observations.split("\n"):
-                    if line.startswith("Screenshot saved at:"):
-                        latest_screenshot = line.replace("Screenshot saved at: ", "").strip()
-                        break
-                if latest_screenshot:
-                    break
-        
-        # Format output for chatbot
-        output = f"**Agent Output:**\n{agent_output}"
-        if latest_screenshot:
-            output += f"\n\n**Latest Screenshot:**"
-        
-        return output, latest_screenshot
-    except Exception as e:
-        logger.error(f"Agent execution failed: {str(e)}")
-        return f"Error: {str(e)}", None
-
-# Custom Gradio interface
-def process_input(user_input, history):
-    if not user_input.strip():
-        return history, None
-    output, latest_screenshot = run_agent_chat(user_input, history)
-    new_history = history + [[user_input, output]]
-    return new_history, latest_screenshot
-
-if __name__ == "__main__":
-    with gr.Blocks() as demo:
-        gr.Markdown("# Web Navigation Agent")
-        chatbot = gr.Chatbot(label="Chat")
-        msg = gr.Textbox(placeholder="Enter URL and request (e.g., [invalid url, do not cite] Click on Developers)")
-        btn = gr.Button("Send")
-        image = gr.Image(label="Latest Screenshot")
-        btn.click(process_input, inputs=[msg, chatbot], outputs=[chatbot, image])
-        msg.submit(process_input, inputs=[msg, chatbot], outputs=[chatbot, image])
-    try:
-        demo.launch()
-    except KeyboardInterrupt:
-        driver.quit()
-        logger.info("Chrome driver closed on exit.")
+# Launch Gradio UI
+try:
+    GradioUI(agent, file_upload_folder=os.path.join(tempfile.gettempdir(), "uploads")).launch()
+except KeyboardInterrupt:
+    driver.quit()
+    logger.info("Chrome driver closed on exit.")