diff --git a/gptme/tools/__init__.py b/gptme/tools/__init__.py index 66008831..fd401c78 100644 --- a/gptme/tools/__init__.py +++ b/gptme/tools/__init__.py @@ -17,6 +17,7 @@ from .tmux import tool as tmux_tool from .vision import tool as vision_tool from .youtube import tool as youtube_tool +from .screenshot import tool as screenshot_tool logger = logging.getLogger(__name__) @@ -40,6 +41,7 @@ gh_tool, chats_tool, youtube_tool, + screenshot_tool, vision_tool, # python tool is loaded last to ensure all functions are registered python_tool, diff --git a/gptme/tools/screenshot.py b/gptme/tools/screenshot.py new file mode 100644 index 00000000..770a89ce --- /dev/null +++ b/gptme/tools/screenshot.py @@ -0,0 +1,53 @@ +""" +A simple screenshot tool, using `screencapture` on macOS and `scrot` on Linux. +""" + +import os +import subprocess +from collections.abc import Generator +from datetime import datetime +from pathlib import Path + +from ..message import Message +from .base import ToolSpec + + +def _screenshot(path: Path) -> Path: + path = Path(path).resolve() + + if os.name == "posix": + if os.uname().sysname == "Darwin": # macOS + subprocess.run(["screencapture", str(path)], check=True) + else: # Linux + subprocess.run(["scrot", "-s", str(path)], check=True) + else: + raise NotImplementedError( + "Screenshot functionality is only available on macOS and Linux." + ) + + return path + + +def screenshot(path: Path | None = None) -> Generator[Message, None, None]: + """ + Take a screenshot and save it to a file. + """ + if path is None: + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + path = Path(f"screenshot_{timestamp}.png") + + try: + path = _screenshot(path) + yield Message("system", f"Screenshot saved to: {path}") + except NotImplementedError as e: + yield Message("system", str(e)) + except subprocess.CalledProcessError: + yield Message("system", "Failed to capture screenshot.") + + +tool = ToolSpec( + name="screenshot", + desc="Take a screenshot", + instructions="Use this tool to capture a screenshot. You can optionally specify a filename.", + functions=[screenshot], +)