From 53240b91f2bed99fc0725ad09b34a2e17537c142 Mon Sep 17 00:00:00 2001 From: yzlin Date: Sat, 30 Mar 2024 17:11:11 +0800 Subject: [PATCH] allow di to use terminal tool, add tool log fn, expand truncate len --- examples/di/use_github_repo.py | 18 +++++++ metagpt/actions/di/execute_nb_code.py | 2 +- metagpt/logs.py | 15 ++++++ metagpt/prompts/di/write_analysis_code.py | 6 ++- metagpt/tools/libs/__init__.py | 2 + metagpt/tools/libs/terminal.py | 62 +++++++++++++++++++++++ tests/metagpt/tools/libs/test_terminal.py | 15 ++++++ 7 files changed, 118 insertions(+), 2 deletions(-) create mode 100644 examples/di/use_github_repo.py create mode 100644 metagpt/tools/libs/terminal.py create mode 100644 tests/metagpt/tools/libs/test_terminal.py diff --git a/examples/di/use_github_repo.py b/examples/di/use_github_repo.py new file mode 100644 index 000000000..ad541d2d9 --- /dev/null +++ b/examples/di/use_github_repo.py @@ -0,0 +1,18 @@ +import asyncio + +from metagpt.roles.di.data_interpreter import DataInterpreter + +USE_GOT_REPO_REQ = """ +This is a link to the GOT github repo: https://github.com/spcl/graph-of-thoughts.git. +Clone it, read the README to understand the usage, install it, and finally run the quick start example. +**Note the config for LLM is at `config/config_got.json`, use this path directly.** Don't write all codes in one response, each time, just write code for one step. +""" + + +async def main(): + di = DataInterpreter(tools=["Terminal"]) + await di.run(USE_GOT_REPO_REQ) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/metagpt/actions/di/execute_nb_code.py b/metagpt/actions/di/execute_nb_code.py index 0cf16b70f..aab204499 100644 --- a/metagpt/actions/di/execute_nb_code.py +++ b/metagpt/actions/di/execute_nb_code.py @@ -106,7 +106,7 @@ class ExecuteNbCode(Action): else: cell["outputs"].append(new_output(output_type="stream", name="stdout", text=str(output))) - def parse_outputs(self, outputs: list[str], keep_len: int = 2000) -> Tuple[bool, str]: + def parse_outputs(self, outputs: list[str], keep_len: int = 5000) -> Tuple[bool, str]: """Parses the outputs received from notebook execution.""" assert isinstance(outputs, list) parsed_output, is_success = [], True diff --git a/metagpt/logs.py b/metagpt/logs.py index 90bac21aa..4e5e90e3b 100644 --- a/metagpt/logs.py +++ b/metagpt/logs.py @@ -6,6 +6,8 @@ @File : logs.py """ +from __future__ import annotations + import sys from datetime import datetime from functools import partial @@ -34,9 +36,22 @@ def log_llm_stream(msg): _llm_stream_log(msg) +def log_tool_output(output: str, tool_name: str = "", tags: list[str] = None): + """interface for logging tool output, can be set to log tool output in different ways to different places with set_tool_output_logfunc""" + _tool_output_log(output) + + def set_llm_stream_logfunc(func): global _llm_stream_log _llm_stream_log = func +def set_tool_output_logfunc(func): + global _tool_output_log + _tool_output_log = func + + _llm_stream_log = partial(print, end="") + + +_tool_output_log = partial(print, end="") diff --git a/metagpt/prompts/di/write_analysis_code.py b/metagpt/prompts/di/write_analysis_code.py index e5663d498..d2b4f1299 100644 --- a/metagpt/prompts/di/write_analysis_code.py +++ b/metagpt/prompts/di/write_analysis_code.py @@ -1,4 +1,8 @@ -INTERPRETER_SYSTEM_MSG = """As a data scientist, you need to help user to achieve their goal step by step in a continuous Jupyter notebook. Since it is a notebook environment, don't use asyncio.run. Instead, use await if you need to call an async function.""" +INTERPRETER_SYSTEM_MSG = """ +As a data scientist, you need to help user to achieve their goal step by step in a continuous Jupyter notebook. +Since it is a notebook environment, don't use asyncio.run. Instead, use await if you need to call an async function. +If you want to use shell command such as git clone, pip install packages, navigate folders, read file, etc., use Terminal tool if available before trying ! in notebook block. +""" STRUCTUAL_PROMPT = """ # User Requirement diff --git a/metagpt/tools/libs/__init__.py b/metagpt/tools/libs/__init__.py index 91596fd3d..174924385 100644 --- a/metagpt/tools/libs/__init__.py +++ b/metagpt/tools/libs/__init__.py @@ -11,6 +11,7 @@ from metagpt.tools.libs import ( gpt_v_generator, web_scraping, email_login, + terminal, ) _ = ( @@ -20,4 +21,5 @@ _ = ( gpt_v_generator, web_scraping, email_login, + terminal, ) # Avoid pre-commit error diff --git a/metagpt/tools/libs/terminal.py b/metagpt/tools/libs/terminal.py new file mode 100644 index 000000000..b277db8dc --- /dev/null +++ b/metagpt/tools/libs/terminal.py @@ -0,0 +1,62 @@ +import subprocess + +from metagpt.logs import log_tool_output +from metagpt.tools.tool_registry import register_tool + + +@register_tool() +class Terminal: + """A tool for running terminal commands. Don't initialize a new instance of this class if one already exists.""" + + def __init__(self): + self.shell_command = ["bash"] # FIXME: should consider windows support later + self.command_terminator = "\n" + self.end_marker = "#END_MARKER#" + + # Start a persistent shell process + self.process = subprocess.Popen( + self.shell_command, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, # Line buffered + ) + + def run_command(self, cmd: str) -> str: + """ + Run a command in the terminal and return the output. + When the command is being executed, stream the output to the terminal. + Maintains state across commands, such as current directory. + + Args: + cmd (str): The command to run in the terminal. + + Returns: + str: The output of the terminal command. + """ + cmd_output = [] + + # Send the command + self.process.stdin.write(cmd + self.command_terminator) + self.process.stdin.write( + f'echo "{self.end_marker}"' + self.command_terminator + ) # Unique marker to signal command end + self.process.stdin.flush() + log_tool_output(output=cmd + self.command_terminator, tool_name="Terminal") # log the command + + # Read the output until the unique marker is found + while True: + line = self.process.stdout.readline() + if line.strip() == self.end_marker: + break + log_tool_output(output=line, tool_name="Terminal") # log stdout in real-time + cmd_output.append(line) + + return "".join(cmd_output) + + def close(self): + """Close the persistent shell process.""" + self.process.stdin.close() + self.process.terminate() + self.process.wait() diff --git a/tests/metagpt/tools/libs/test_terminal.py b/tests/metagpt/tools/libs/test_terminal.py new file mode 100644 index 000000000..97c33b977 --- /dev/null +++ b/tests/metagpt/tools/libs/test_terminal.py @@ -0,0 +1,15 @@ +from metagpt.const import DATA_PATH, METAGPT_ROOT +from metagpt.tools.libs.terminal import Terminal + + +def test_terminal(): + terminal = Terminal() + + terminal.run_command(f"cd {METAGPT_ROOT}") + output = terminal.run_command("pwd") + assert output.strip() == str(METAGPT_ROOT) + + # pwd now should be METAGPT_ROOT, cd data should land in DATA_PATH + terminal.run_command("cd data") + output = terminal.run_command("pwd") + assert output.strip() == str(DATA_PATH)