Update app.py

app.py

@@ -12,20 +12,25 @@ from datetime import datetime
 from dotenv import load_dotenv
 from huggingface_hub import login
 import tempfile
+import logging
+
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # 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:
+if not hf_token:
     raise ValueError("HF_TOKEN environment variable not set.")
-
 if not gemini_api_key:
     raise ValueError("GEMINI_API_KEY environment variable not set.")
 
+login(hf_token, add_to_git_credential=False)
+
 # Define tools
 @tool
 def search_item_ctrl_f(text: str, nth_result: int = 1) -> str:
@@ -57,17 +62,21 @@ 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)
+# Initialize Chrome driver with error handling
+try:
+    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=new")
+    driver = helium.start_chrome(headless=True, options=chrome_options)
+    logger.info("Chrome driver initialized successfully.")
+except Exception as e:
+    logger.error(f"Failed to initialize Chrome driver: {str(e)}")
+    raise
 
 # Screenshot callback
 def save_screenshot(memory_step: ActionStep, agent: CodeAgent) -> Image.Image:
@@ -85,6 +94,7 @@ def save_screenshot(memory_step: ActionStep, agent: CodeAgent) -> Image.Image:
         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)
+        logger.info(f"Saved screenshot to: {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
@@ -156,9 +166,14 @@ def run_agent(url: str, request: str):
         )
         return agent_output, screenshot_path
     except Exception as e:
+        logger.error(f"Agent execution failed: {str(e)}")
         return f"Error: {str(e)}", None
     finally:
-        driver.quit()
+        try:
+            driver.quit()
+            logger.info("Chrome driver closed.")
+        except:
+            logger.warning("Failed to close Chrome driver.")
 
 # Gradio interface
 with gr.Blocks() as demo:
@@ -174,4 +189,5 @@ with gr.Blocks() as demo:
         outputs=[output_text, output_image]
     )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()