From 65bdb8ad3291cd4a4be691255d3533e6413bc676 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Sun, 27 Oct 2024 19:27:46 +0100 Subject: [PATCH] fix: fixed screenshot() not running if generator not consumed (such as when not last statement in codeblock) --- gptme/tools/screenshot.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/gptme/tools/screenshot.py b/gptme/tools/screenshot.py index 770a89ce..7a6ca3d3 100644 --- a/gptme/tools/screenshot.py +++ b/gptme/tools/screenshot.py @@ -4,11 +4,9 @@ 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 @@ -28,7 +26,7 @@ def _screenshot(path: Path) -> Path: return path -def screenshot(path: Path | None = None) -> Generator[Message, None, None]: +def screenshot(path: Path | None = None) -> Path: """ Take a screenshot and save it to a file. """ @@ -36,13 +34,7 @@ def screenshot(path: Path | None = None) -> Generator[Message, None, 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.") + return _screenshot(path) tool = ToolSpec(