Update app.py

app.py

@@ -71,9 +71,8 @@ def save_screenshot(memory_step: ActionStep, agent: CodeAgent) -> str:
         for previous_memory_step in agent.memory.steps:
             if isinstance(previous_memory_step, ActionStep) and previous_memory_step.step_number < current_step:
                 previous_memory_step.observations_images = None
+        # Capture only the viewport
         original_size = driver.get_window_size()
-        total_height = driver.execute_script("return document.body.scrollHeight")
-        driver.set_window_size(original_size['width'], total_height)
         png_bytes = driver.get_screenshot_as_png()
         image = Image.open(BytesIO(png_bytes))
         screenshot_dir = os.path.join(tempfile.gettempdir(), "web_agent_screenshots")
@@ -82,7 +81,6 @@ def save_screenshot(memory_step: ActionStep, agent: CodeAgent) -> str:
         screenshot_filename = f"screenshot_step_{current_step}_{timestamp}.png"
         screenshot_path = os.path.join(screenshot_dir, screenshot_filename)
         image.save(screenshot_path)
-        driver.set_window_size(original_size['width'], original_size['height'])
         logger.info(f"Saved screenshot to: {screenshot_path}")
         url_info = f"Current url: {driver.current_url}\nScreenshot saved at: {screenshot_path}"
         memory_step.observations = (