From 39409e2338dfc2824e878611579ff3934c849ec1 Mon Sep 17 00:00:00 2001 From: demouo <2081510953@qq.com> Date: Tue, 6 May 2025 17:11:39 +0800 Subject: [PATCH 1/4] Added --fix to rechat on the last command and output (#698) --- sgpt/app.py | 8 ++++++++ sgpt/command.py | 31 +++++++++++++++++++++++++++++++ sgpt/config.py | 2 ++ sgpt/utils.py | 28 +++++++++++++++++++++++++++- tests/test_shell.py | 28 ++++++++++++++++++++++------ 5 files changed, 90 insertions(+), 7 deletions(-) create mode 100644 sgpt/command.py diff --git a/sgpt/app.py b/sgpt/app.py index 9c711e66..62656651 100644 --- a/sgpt/app.py +++ b/sgpt/app.py @@ -17,6 +17,7 @@ from sgpt.role import DefaultRoles, SystemRole from sgpt.utils import ( get_edited_prompt, + get_fixed_prompt, get_sgpt_version, install_shell_integration, run_command, @@ -84,6 +85,10 @@ def main( False, help="Open $EDITOR to provide a prompt.", ), + fix: bool = typer.Option( + False, + help="Fix the wrong last command.", + ), cache: bool = typer.Option( True, help="Cache completion results.", @@ -199,6 +204,9 @@ def main( if editor: prompt = get_edited_prompt() + if fix: + prompt = get_fixed_prompt() + role_class = ( DefaultRoles.check_get(shell, describe_shell, code) if not role diff --git a/sgpt/command.py b/sgpt/command.py new file mode 100644 index 00000000..a98ffc25 --- /dev/null +++ b/sgpt/command.py @@ -0,0 +1,31 @@ +from pathlib import Path +from click import UsageError +import json + + +class Command: + def __init__(self, command_path: Path): + self.command_path = command_path + self.command_path.mkdir(parents=True, exist_ok=True) + + def get_last_command(self) -> tuple[str, str]: + """ + get the last command and output from the command path + """ + last_command_file = self.command_path / "last_command.json" + if not last_command_file.exists(): + raise UsageError("No last command and output found.") + with open(last_command_file, "r", encoding="utf-8") as file: + data = json.load(file) + command = data.get("command", "") + output = data.get("output", "") + return command, output + + def set_last_command(self, command: str, output: str) -> None: + """ + set the last command and output to the command path + """ + last_command_file = self.command_path / "last_command.json" + with open(last_command_file, "w", encoding="utf-8") as file: + data = {"command": command, "output": output} + json.dump(data, file, ensure_ascii=False) diff --git a/sgpt/config.py b/sgpt/config.py index bc083733..b0ee308b 100644 --- a/sgpt/config.py +++ b/sgpt/config.py @@ -13,6 +13,7 @@ FUNCTIONS_PATH = SHELL_GPT_CONFIG_FOLDER / "functions" CHAT_CACHE_PATH = Path(gettempdir()) / "chat_cache" CACHE_PATH = Path(gettempdir()) / "cache" +COMMAND_PATH = Path(gettempdir()) / "commands" # TODO: Refactor ENV variables with SGPT_ prefix. DEFAULT_CONFIG = { @@ -37,6 +38,7 @@ "SHELL_INTERACTION": os.getenv("SHELL_INTERACTION ", "true"), "OS_NAME": os.getenv("OS_NAME", "auto"), "SHELL_NAME": os.getenv("SHELL_NAME", "auto"), + "COMMAND_PATH": os.getenv("COMMAND_PATH", str(COMMAND_PATH)), # New features might add their own config variables here. } diff --git a/sgpt/utils.py b/sgpt/utils.py index d49af9a3..a628ccdc 100644 --- a/sgpt/utils.py +++ b/sgpt/utils.py @@ -1,6 +1,8 @@ import os +from pathlib import Path import platform import shlex +import subprocess from tempfile import NamedTemporaryFile from typing import Any, Callable @@ -8,6 +10,8 @@ from click import BadParameter, UsageError from sgpt.__version__ import __version__ +from sgpt.command import Command +from sgpt.config import cfg from sgpt.integration import bash_integration, zsh_integration @@ -33,6 +37,17 @@ def get_edited_prompt() -> str: return output +command_helper = Command(command_path=Path(cfg.get("COMMAND_PATH"))) + + +def get_fixed_prompt() -> str: + """ + get the last command and output then return a PROMPT + """ + command, output = command_helper.get_last_command() + return f"The last command `{command}` failed with error:\n{output}\nPlease fix it." + + def run_command(command: str) -> None: """ Runs a command in the user's shell. @@ -50,7 +65,18 @@ def run_command(command: str) -> None: shell = os.environ.get("SHELL", "/bin/sh") full_command = f"{shell} -c {shlex.quote(command)}" - os.system(full_command) + # os.system(full_command) + process = subprocess.Popen( + args=full_command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + shell=True, + ) + output, _ = process.communicate() + print(output) + + command_helper.set_last_command(command, output) def option_callback(func: Callable) -> Callable: # type: ignore diff --git a/tests/test_shell.py b/tests/test_shell.py index f7ad3caa..b0f0c4c7 100644 --- a/tests/test_shell.py +++ b/tests/test_shell.py @@ -1,4 +1,5 @@ import os +import subprocess from pathlib import Path from unittest.mock import patch @@ -81,15 +82,22 @@ def test_describe_shell_stdin(completion): assert "lists" in result.stdout -@patch("os.system") +@patch("subprocess.Popen") @patch("sgpt.handlers.handler.completion") -def test_shell_run_description(completion, system): +def test_shell_run_description(completion, mock_popen): + mock_popen.return_value.communicate.return_value = ("stdout", None) completion.side_effect = [mock_comp("echo hello"), mock_comp("prints hello")] args = {"prompt": "echo hello", "--shell": True} inputs = "__sgpt__eof__\nd\ne\n" result = runner.invoke(app, cmd_args(**args), input=inputs) shell = os.environ.get("SHELL", "/bin/sh") - system.assert_called_once_with(f"{shell} -c 'echo hello'") + mock_popen.assert_called_once_with( + args=f"{shell} -c 'echo hello'", + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + shell=True, + ) assert result.exit_code == 0 assert "echo hello" in result.stdout assert "prints hello" in result.stdout @@ -133,9 +141,10 @@ def test_shell_chat(completion): # TODO: Shell chat can be recalled without --shell option. -@patch("os.system") +@patch("subprocess.Popen") @patch("sgpt.handlers.handler.completion") -def test_shell_repl(completion, mock_system): +def test_shell_repl(completion, mock_popen): + mock_popen.return_value.communicate.return_value = ("stdout", None) completion.side_effect = [mock_comp("ls"), mock_comp("ls | sort")] role = SystemRole.get(DefaultRoles.SHELL.value) chat_name = "_test" @@ -146,7 +155,14 @@ def test_shell_repl(completion, mock_system): inputs = ["__sgpt__eof__", "list folder", "sort by name", "e", "exit()"] result = runner.invoke(app, cmd_args(**args), input="\n".join(inputs)) shell = os.environ.get("SHELL", "/bin/sh") - mock_system.assert_called_once_with(f"{shell} -c 'ls | sort'") + + mock_popen.assert_called_once_with( + args=f"{shell} -c 'ls | sort'", + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + shell=True, + ) expected_messages = [ {"role": "system", "content": role.role}, From b3563f064c5e93a134fccf6405c33b5b1ec4fc8f Mon Sep 17 00:00:00 2001 From: demouo <2081510953@qq.com> Date: Wed, 7 May 2025 22:22:10 +0800 Subject: [PATCH 2/4] Add command extraction, processing empty prompt, and fix last command for REPL mode. --- sgpt/app.py | 12 ++++++++++-- sgpt/handlers/repl_handler.py | 11 +++++++---- sgpt/utils.py | 10 ++++++++++ 3 files changed, 27 insertions(+), 6 deletions(-) diff --git a/sgpt/app.py b/sgpt/app.py index 62656651..3187e63f 100644 --- a/sgpt/app.py +++ b/sgpt/app.py @@ -16,6 +16,7 @@ from sgpt.llm_functions.init_functions import install_functions as inst_funcs from sgpt.role import DefaultRoles, SystemRole from sgpt.utils import ( + extract_command_from_completion, get_edited_prompt, get_fixed_prompt, get_sgpt_version, @@ -226,6 +227,11 @@ def main( functions=function_schemas, ) + if prompt == "": + raise BadArgumentUsage( + "Prompt cant be empty. Use `sgpt ` to get started." + ) + if chat: full_completion = ChatHandler(chat, role_class, md).handle( prompt=prompt, @@ -245,6 +251,8 @@ def main( functions=function_schemas, ) + command = extract_command_from_completion(full_completion) + while shell and interaction: option = typer.prompt( text="[E]xecute, [D]escribe, [A]bort", @@ -255,10 +263,10 @@ def main( ) if option in ("e", "y"): # "y" option is for keeping compatibility with old version. - run_command(full_completion) + run_command(command) elif option == "d": DefaultHandler(DefaultRoles.DESCRIBE_SHELL.get_role(), md).handle( - full_completion, + command, model=model, temperature=temperature, top_p=top_p, diff --git a/sgpt/handlers/repl_handler.py b/sgpt/handlers/repl_handler.py index 04019fd8..10a6ec7a 100644 --- a/sgpt/handlers/repl_handler.py +++ b/sgpt/handlers/repl_handler.py @@ -5,7 +5,7 @@ from rich.rule import Rule from ..role import DefaultRoles, SystemRole -from ..utils import run_command +from ..utils import run_command, get_fixed_prompt, extract_command_from_completion from .chat_handler import ChatHandler from .default_handler import DefaultHandler @@ -32,7 +32,7 @@ def handle(self, init_prompt: str, **kwargs: Any) -> None: # type: ignore if not self.role.name == DefaultRoles.SHELL.value else ( "Entering shell REPL mode, type [e] to execute commands " - "or [d] to describe the commands, press Ctrl+C to exit." + "or [d] to describe the commands, or [f] to fix the last command, press Ctrl+C to exit." ) ) typer.secho(info_message, fg="yellow") @@ -53,14 +53,17 @@ def handle(self, init_prompt: str, **kwargs: Any) -> None: # type: ignore if init_prompt: prompt = f"{init_prompt}\n\n\n{prompt}" init_prompt = "" + if self.role.name == DefaultRoles.SHELL.value and prompt == "f": + prompt = get_fixed_prompt() + command = extract_command_from_completion(full_completion) if self.role.name == DefaultRoles.SHELL.value and prompt == "e": typer.echo() - run_command(full_completion) + run_command(command) typer.echo() rich_print(Rule(style="bold magenta")) elif self.role.name == DefaultRoles.SHELL.value and prompt == "d": DefaultHandler( DefaultRoles.DESCRIBE_SHELL.get_role(), self.markdown - ).handle(prompt=full_completion, **kwargs) + ).handle(prompt=command, **kwargs) else: full_completion = super().handle(prompt=prompt, **kwargs) diff --git a/sgpt/utils.py b/sgpt/utils.py index a628ccdc..776e7407 100644 --- a/sgpt/utils.py +++ b/sgpt/utils.py @@ -1,4 +1,5 @@ import os +import re from pathlib import Path import platform import shlex @@ -48,6 +49,15 @@ def get_fixed_prompt() -> str: return f"The last command `{command}` failed with error:\n{output}\nPlease fix it." +def extract_command_from_completion(completion: str) -> str: + """ + using regex to extract the command from the completion + """ + if match := re.search(r"```(.*sh)?(.*?)```", completion, re.DOTALL): + return match[2].strip() + return completion + + def run_command(command: str) -> None: """ Runs a command in the user's shell. From b06ea65b529eaad8af96409a3df928c920612695 Mon Sep 17 00:00:00 2001 From: demouo <2081510953@qq.com> Date: Wed, 7 May 2025 22:47:24 +0800 Subject: [PATCH 3/4] update empty prompt logic. --- sgpt/app.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/sgpt/app.py b/sgpt/app.py index 3187e63f..d7499de0 100644 --- a/sgpt/app.py +++ b/sgpt/app.py @@ -227,10 +227,9 @@ def main( functions=function_schemas, ) - if prompt == "": - raise BadArgumentUsage( - "Prompt cant be empty. Use `sgpt ` to get started." - ) + if not prompt: + print("Prompt cant be empty. Use `sgpt ` to get started.") + return if chat: full_completion = ChatHandler(chat, role_class, md).handle( @@ -251,9 +250,8 @@ def main( functions=function_schemas, ) - command = extract_command_from_completion(full_completion) - while shell and interaction: + command = extract_command_from_completion(full_completion) option = typer.prompt( text="[E]xecute, [D]escribe, [A]bort", type=Choice(("e", "d", "a", "y"), case_sensitive=False), From cb481e1998f28d5ec0816b9877f3e932e7153334 Mon Sep 17 00:00:00 2001 From: demouo <2081510953@qq.com> Date: Wed, 7 May 2025 23:01:57 +0800 Subject: [PATCH 4/4] Update processing empty prompt with --show-chat. --- sgpt/app.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sgpt/app.py b/sgpt/app.py index d7499de0..4b7646c3 100644 --- a/sgpt/app.py +++ b/sgpt/app.py @@ -228,7 +228,8 @@ def main( ) if not prompt: - print("Prompt cant be empty. Use `sgpt ` to get started.") + if not show_chat: + print("Prompt cant be empty. Use `sgpt ` to get started.") return if chat: