"""
Screenshot storage service for browser-use agents.
"""
import base64
from pathlib import Path
import anyio
class ScreenshotService:
"""Simple screenshot storage service that saves screenshots to disk"""
def __init__(self, agent_directory: str | Path):
"""Initialize with agent directory path"""
self.agent_directory = Path(agent_directory) if isinstance(agent_directory, str) else agent_directory
# Create screenshots subdirectory
self.screenshots_dir = self.agent_directory / 'screenshots'
self.screenshots_dir.mkdir(parents=True, exist_ok=True)
async def store_screenshot(self, screenshot_b64: str, step_number: int) -> str:
"""Store screenshot to disk and return the full path as string"""
screenshot_filename = f'step_{step_number}.png'
screenshot_path = self.screenshots_dir / screenshot_filename
# Decode base64 and save to disk
screenshot_data = base64.b64decode(screenshot_b64)
async with await anyio.open_file(screenshot_path, 'wb') as f:
await f.write(screenshot_data)
return str(screenshot_path)
async def get_screenshot(self, screenshot_path: str) -> str | None:
"""Load screenshot from disk path and return as base64"""
if not screenshot_path:
return None
path = Path(screenshot_path)
if not path.exists():
return None
# Load from disk and encode to base64
async with await anyio.open_file(path, 'rb') as f:
screenshot_data = await f.read()
return base64.b64encode(screenshot_data).decode('utf-8')
