feat: added screenshot tool (#92)

* feat: started working on screenshot tool

* fix: updated the screenshot tool

* fix: completed screenshot tool

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,

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],
+)