save final screenshot

.gitignore

@@ -171,4 +171,5 @@ cython_debug/
 .pypirc
 
 logs/
-local_trajectories/
+local_trajectories/
+screenshots/

src/proxy_lite/cli.py

@@ -1,5 +1,6 @@
 import argparse
 import asyncio
+import base64
 import os
 from pathlib import Path
 from typing import Optional
@@ -35,7 +36,15 @@ def do_command(args):
     if args.viewport_height:
         config.viewport_height = args.viewport_height
     o = Runner(config=config)
-    asyncio.run(o.run(do_text))
+    result = asyncio.run(o.run(do_text))
+
+    final_screenshot = result.observations[-1].info["original_image"]
+    folder_path = Path(__file__).parent.parent.parent / "screenshots"
+    folder_path.mkdir(parents=True, exist_ok=True)
+    path = folder_path / f"{result.run_id}.png"
+    with open(path, "wb") as f:
+        f.write(base64.b64decode(final_screenshot))
+    logger.info(f"🤖 Screenshot saved to {path}")
 
 
 def main():

src/proxy_lite/configs/default.yaml

@@ -7,6 +7,7 @@ environment:
   include_poi_text: true
   headless: false
   homepage: https://www.google.co.uk
+  keep_original_image: true
 solver:
   name: simple
   agent:

src/proxy_lite/recorder.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import datetime
 import json
 import os
-import sys
 import uuid
 from pathlib import Path
 from typing import Any, Optional, Self
@@ -80,7 +79,7 @@ class DataRecorder:
         self.local_folder = local_folder
 
     def initialise_run(self, task: str) -> Run:
-        self.local_folder = Path(os.path.abspath(".")) / "local_trajectories"
+        self.local_folder = Path(__file__).parent.parent.parent / "local_trajectories"
         os.makedirs(self.local_folder, exist_ok=True)
         return Run.initialise(task)