2022-10-31 01:09:04 +00:00
|
|
|
"""Test functionality of Python REPL."""
|
2023-04-11 04:25:03 +00:00
|
|
|
import sys
|
|
|
|
|
|
|
|
import pytest
|
2022-10-31 01:09:04 +00:00
|
|
|
|
|
|
|
from langchain.python import PythonREPL
|
2023-04-11 04:25:03 +00:00
|
|
|
from langchain.tools.python.tool import PythonAstREPLTool, PythonREPLTool
|
2023-04-01 19:54:06 +00:00
|
|
|
|
|
|
|
_SAMPLE_CODE = """
|
|
|
|
```
|
|
|
|
def multiply():
|
|
|
|
print(5*6)
|
|
|
|
multiply()
|
|
|
|
```
|
|
|
|
"""
|
2022-10-31 01:09:04 +00:00
|
|
|
|
2023-04-11 04:25:03 +00:00
|
|
|
_AST_SAMPLE_CODE = """
|
|
|
|
```
|
|
|
|
def multiply():
|
|
|
|
return(5*6)
|
|
|
|
multiply()
|
|
|
|
```
|
|
|
|
"""
|
|
|
|
|
2022-10-31 01:09:04 +00:00
|
|
|
|
|
|
|
def test_python_repl() -> None:
|
|
|
|
"""Test functionality when globals/locals are not provided."""
|
|
|
|
repl = PythonREPL()
|
|
|
|
|
|
|
|
# Run a simple initial command.
|
|
|
|
repl.run("foo = 1")
|
2023-03-01 03:44:39 +00:00
|
|
|
assert repl.locals is not None
|
|
|
|
assert repl.locals["foo"] == 1
|
2022-10-31 01:09:04 +00:00
|
|
|
|
|
|
|
# Now run a command that accesses `foo` to make sure it still has it.
|
|
|
|
repl.run("bar = foo * 2")
|
2023-03-01 03:44:39 +00:00
|
|
|
assert repl.locals is not None
|
|
|
|
assert repl.locals["bar"] == 2
|
2022-10-31 01:09:04 +00:00
|
|
|
|
|
|
|
|
|
|
|
def test_python_repl_no_previous_variables() -> None:
|
|
|
|
"""Test that it does not have access to variables created outside the scope."""
|
|
|
|
foo = 3 # noqa: F841
|
|
|
|
repl = PythonREPL()
|
2022-12-19 02:51:23 +00:00
|
|
|
output = repl.run("print(foo)")
|
|
|
|
assert output == "name 'foo' is not defined"
|
2022-10-31 01:09:04 +00:00
|
|
|
|
|
|
|
|
|
|
|
def test_python_repl_pass_in_locals() -> None:
|
|
|
|
"""Test functionality when passing in locals."""
|
|
|
|
_locals = {"foo": 4}
|
|
|
|
repl = PythonREPL(_locals=_locals)
|
|
|
|
repl.run("bar = foo * 2")
|
2023-03-01 03:44:39 +00:00
|
|
|
assert repl.locals is not None
|
|
|
|
assert repl.locals["bar"] == 8
|
2022-12-04 23:57:36 +00:00
|
|
|
|
|
|
|
|
|
|
|
def test_functionality() -> None:
|
|
|
|
"""Test correct functionality."""
|
|
|
|
chain = PythonREPL()
|
|
|
|
code = "print(1 + 1)"
|
|
|
|
output = chain.run(code)
|
|
|
|
assert output == "2\n"
|
2023-03-01 03:44:39 +00:00
|
|
|
|
|
|
|
|
2023-04-01 19:54:06 +00:00
|
|
|
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"
|
|
|
|
|
|
|
|
|
2023-04-11 04:25:03 +00:00
|
|
|
def test_python_ast_repl_multiline() -> None:
|
|
|
|
"""Test correct functionality for ChatGPT multiline commands."""
|
|
|
|
if sys.version_info < (3, 9):
|
|
|
|
pytest.skip("Python 3.9+ is required for this test")
|
|
|
|
tool = PythonAstREPLTool()
|
|
|
|
output = tool.run(_AST_SAMPLE_CODE)
|
|
|
|
assert output == 30
|
|
|
|
|
|
|
|
|
2023-03-01 03:44:39 +00:00
|
|
|
def test_function() -> None:
|
|
|
|
"""Test correct functionality."""
|
|
|
|
chain = PythonREPL()
|
|
|
|
code = "def add(a, b): " " return a + b"
|
|
|
|
output = chain.run(code)
|
|
|
|
assert output == ""
|
|
|
|
|
|
|
|
code = "print(add(1, 2))"
|
|
|
|
output = chain.run(code)
|
|
|
|
assert output == "3\n"
|