Update app.py

app.py

@@ -107,11 +107,11 @@ def save_screenshot(memory_step: ActionStep, agent: CodeAgent) -> Image.Image:
         image.save(screenshot_path)
         logger.info(f"Saved screenshot to: {screenshot_path}")
         # Update observations
-        url_info = f"Current url: {driver.current_url}"
+        url_info = f"Current url: {driver.current_url}\nScreenshot saved at: {screenshot_path}"
         memory_step.observations = (
             url_info if memory_step.observations is None else memory_step.observations + "\n" + url_info
         )
-        memory_step.observations_images = screenshot_path
+        memory_step.observations_images = image
         return image
 
 # Initialize model and agent
@@ -175,12 +175,17 @@ def run_agent_chat(user_input: str, history: list):
         search_request = f"Please go to {url}. {request}"
         agent_output = agent.run(search_request + helium_instructions)
         
-        # Collect the latest screenshot
+        # Collect the latest screenshot path from observations
         latest_screenshot = None
         for step in reversed(agent.memory.steps):
-            if isinstance(step, ActionStep) and step.observations_images:
-                latest_screenshot = step.observations_images
-                break
+            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}"