diff --git a/langchain/tools/python/tool.py b/langchain/tools/python/tool.py index 0a1adb46..0b2ac846 100644 --- a/langchain/tools/python/tool.py +++ b/langchain/tools/python/tool.py @@ -25,9 +25,12 @@ class PythonREPLTool(BaseTool): "with `print(...)`." ) python_repl: PythonREPL = Field(default_factory=_get_default_python_repl) + sanitize_input: bool = True def _run(self, query: str) -> str: """Use the tool.""" + if self.sanitize_input: + query = query.strip().strip("```") return self.python_repl.run(query) async def _arun(self, query: str) -> str: diff --git a/tests/unit_tests/test_python.py b/tests/unit_tests/test_python.py index d1eb3ae9..fab0c88d 100644 --- a/tests/unit_tests/test_python.py +++ b/tests/unit_tests/test_python.py @@ -1,6 +1,15 @@ """Test functionality of Python REPL.""" from langchain.python import PythonREPL +from langchain.tools.python.tool import PythonREPLTool + +_SAMPLE_CODE = """ +``` +def multiply(): + print(5*6) +multiply() +``` +""" def test_python_repl() -> None: @@ -43,6 +52,14 @@ def test_functionality() -> None: assert output == "2\n" +def test_functionality_multiline() -> None: + """Test correct functionality for ChatGPT multiline commands.""" + chain = PythonREPL() + tool = PythonREPLTool(python_repl=chain) + output = tool.run(_SAMPLE_CODE) + assert output == "30\n" + + def test_function() -> None: """Test correct functionality.""" chain = PythonREPL()