Merge pull request #159 from MahlerTom/main

[BUG] `CodeInterpreterTool` cannot handle mutli-line code

crewai_tools/tools/code_interpreter_tool/code_interpreter_tool.py

@@ -2,7 +2,9 @@
 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 @@ def _verify_docker_image(self) -> None:
         """
         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,25 +75,25 @@ def _run(self, **kwargs) -> str:
             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
         try:
             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 @@ def run_code_in_docker(self, code: str, libraries_used: List[str]) -> str:
         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()

tests/tools/test_code_interpreter_tool.py

@@ -7,32 +7,47 @@
 
 
 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)