diff --git a/src/crewai_tools/tools/code_interpreter_tool/code_interpreter_tool.py b/src/crewai_tools/tools/code_interpreter_tool/code_interpreter_tool.py index 61c180fe3..1809dcdda 100644 --- a/src/crewai_tools/tools/code_interpreter_tool/code_interpreter_tool.py +++ b/src/crewai_tools/tools/code_interpreter_tool/code_interpreter_tool.py @@ -2,7 +2,9 @@ import importlib.util import os from typing import List, Optional, Type -import docker +from docker import from_env as docker_from_env +from docker.models.containers import Container +from docker.errors import ImageNotFound, NotFound from crewai.tools import BaseTool from pydantic import BaseModel, Field @@ -39,12 +41,12 @@ class CodeInterpreterTool(BaseTool): """ Verify if the Docker image is available. Optionally use a user-provided Dockerfile. """ - client = docker.from_env() + client = docker_from_env() try: client.images.get(self.default_image_tag) - except docker.errors.ImageNotFound: + except ImageNotFound: if self.user_dockerfile_path and os.path.exists(self.user_dockerfile_path): dockerfile_path = self.user_dockerfile_path else: @@ -73,17 +75,17 @@ class CodeInterpreterTool(BaseTool): return self.run_code_in_docker(code, libraries_used) def _install_libraries( - self, container: docker.models.containers.Container, libraries: List[str] + self, container: Container, libraries: List[str] ) -> None: """ Install missing libraries in the Docker container """ for library in libraries: - container.exec_run(f"pip install {library}") + container.exec_run(["pip", "install", library]) - def _init_docker_container(self) -> docker.models.containers.Container: + def _init_docker_container(self) -> Container: container_name = "code-interpreter" - client = docker.from_env() + client = docker_from_env() current_path = os.getcwd() # Check if the container is already running @@ -91,7 +93,7 @@ class CodeInterpreterTool(BaseTool): existing_container = client.containers.get(container_name) existing_container.stop() existing_container.remove() - except docker.errors.NotFound: + except NotFound: pass # Container does not exist, no need to remove return client.containers.run( @@ -108,8 +110,7 @@ class CodeInterpreterTool(BaseTool): container = self._init_docker_container() self._install_libraries(container, libraries_used) - cmd_to_run = f'python3 -c "{code}"' - exec_result = container.exec_run(cmd_to_run) + exec_result = container.exec_run(["python3", "-c", code]) container.stop() container.remove() diff --git a/tests/tools/test_code_interpreter_tool.py b/tests/tools/test_code_interpreter_tool.py index a9ffb9dbc..6470c9dc1 100644 --- a/tests/tools/test_code_interpreter_tool.py +++ b/tests/tools/test_code_interpreter_tool.py @@ -7,32 +7,47 @@ from crewai_tools.tools.code_interpreter_tool.code_interpreter_tool import ( class TestCodeInterpreterTool(unittest.TestCase): - @patch("crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker") + @patch("crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env") def test_run_code_in_docker(self, docker_mock): tool = CodeInterpreterTool() code = "print('Hello, World!')" - libraries_used = "numpy,pandas" + libraries_used = ["numpy", "pandas"] expected_output = "Hello, World!\n" - docker_mock.from_env().containers.run().exec_run().exit_code = 0 - docker_mock.from_env().containers.run().exec_run().output = ( + docker_mock().containers.run().exec_run().exit_code = 0 + docker_mock().containers.run().exec_run().output = ( expected_output.encode() ) result = tool.run_code_in_docker(code, libraries_used) self.assertEqual(result, expected_output) - @patch("crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker") + @patch("crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env") def test_run_code_in_docker_with_error(self, docker_mock): tool = CodeInterpreterTool() code = "print(1/0)" - libraries_used = "numpy,pandas" + libraries_used = ["numpy", "pandas"] expected_output = "Something went wrong while running the code: \nZeroDivisionError: division by zero\n" - docker_mock.from_env().containers.run().exec_run().exit_code = 1 - docker_mock.from_env().containers.run().exec_run().output = ( + docker_mock().containers.run().exec_run().exit_code = 1 + docker_mock().containers.run().exec_run().output = ( b"ZeroDivisionError: division by zero\n" ) result = tool.run_code_in_docker(code, libraries_used) self.assertEqual(result, expected_output) + + @patch("crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env") + def test_run_code_in_docker_with_script(self, docker_mock): + tool = CodeInterpreterTool() + code = """print("This is line 1") +print("This is line 2")""" + libraries_used = [] # No additional libraries needed for this test + expected_output = "This is line 1\nThis is line 2\n" + + # Mock Docker responses + docker_mock().containers.run().exec_run().exit_code = 0 + docker_mock().containers.run().exec_run().output = expected_output.encode() + + result = tool.run_code_in_docker(code, libraries_used) + self.assertEqual(result, expected_output)