diff --git a/langchain/utilities/bash.py b/langchain/utilities/bash.py index 50bd688c..d4bcf73d 100644 --- a/langchain/utilities/bash.py +++ b/langchain/utilities/bash.py @@ -6,9 +6,10 @@ from typing import List, Union class BashProcess: """Executes bash commands and returns the output.""" - def __init__(self, strip_newlines: bool = False): + def __init__(self, strip_newlines: bool = False, return_err_output: bool = False): """Initialize with stripping newlines.""" self.strip_newlines = strip_newlines + self.return_err_output = return_err_output def run(self, commands: Union[str, List[str]]) -> str: """Run commands and return final output.""" @@ -16,8 +17,16 @@ class BashProcess: commands = [commands] commands = ";".join(commands) try: - output = subprocess.check_output(commands, shell=True).decode() + output = subprocess.run( + commands, + shell=True, + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ).stdout.decode() except subprocess.CalledProcessError as error: + if self.return_err_output: + return error.stdout.decode() return str(error) if self.strip_newlines: output = output.strip() diff --git a/tests/unit_tests/test_bash.py b/tests/unit_tests/test_bash.py index 5c4dd785..157b1c9c 100644 --- a/tests/unit_tests/test_bash.py +++ b/tests/unit_tests/test_bash.py @@ -21,6 +21,13 @@ def test_incorrect_command() -> None: assert output == "Command 'invalid_command' returned non-zero exit status 127." +def test_incorrect_command_return_err_output() -> None: + """Test optional returning of shell output on incorrect command.""" + session = BashProcess(return_err_output=True) + output = session.run(["invalid_command"]) + assert output == "/bin/sh: 1: invalid_command: not found\n" + + def test_create_directory_and_files(tmp_path: Path) -> None: """Test creation of a directory and files in a temporary directory.""" session = BashProcess(strip_newlines=True)