andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-14 18:48:29 +00:00
Code Issues Actions 3 Packages Projects Releases Wiki Activity

Merge pull request #159 from MahlerTom/main

Browse Source 
[BUG] `CodeInterpreterTool` cannot handle mutli-line code
This commit is contained in:
João Moura
2024-12-27 18:03:47 -03:00
committed by GitHub
parent b7a132db89 ba8f95964f
commit 078120e548
2 changed files with 34 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
src/crewai_tools/tools/code_interpreter_tool/code_interpreter_tool.py
View File

@@ -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()

31
tests/tools/test_code_interpreter_tool.py
View File

@@ -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)