From 8007938d6b7c68c1314e35260306af1b614c089a Mon Sep 17 00:00:00 2001
From: Eduardo Chiarotti <dudumelgaco@hotmail.com>
Date: Tue, 27 Aug 2024 08:35:56 -0300
Subject: [PATCH] feat: ADd volume option

---
 .../tools/code_interpreter_tool/README.md     | 11 +++++
 .../code_interpreter_tool.py                  | 42 ++++++++++++++-----
 2 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/src/crewai_tools/tools/code_interpreter_tool/README.md b/src/crewai_tools/tools/code_interpreter_tool/README.md
index e66a82e39..035f910a4 100644
--- a/src/crewai_tools/tools/code_interpreter_tool/README.md
+++ b/src/crewai_tools/tools/code_interpreter_tool/README.md
@@ -27,3 +27,14 @@ Agent(
     tools=[CodeInterpreterTool()],
 )
 ```
+
+Or if you need to pass your own Dockerfile just do this
+
+```python
+from crewai_tools import CodeInterpreterTool
+
+Agent(
+    ...
+    tools=[CodeInterpreterTool(user_docker_file_path="<Dockerfile_path>")],
+)
+```
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 f341e52d0..2bde9a54b 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
@@ -25,7 +25,9 @@ class CodeInterpreterTool(BaseTool):
     name: str = "Code Interpreter"
     description: str = "Interprets Python3 code strings with a final print statement."
     args_schema: Type[BaseModel] = CodeInterpreterSchema
+    default_image_tag: str = "code-interpreter:latest"
     code: Optional[str] = None
+    user_docker_file_path: Optional[str] = None
 
     @staticmethod
     def _get_installed_package_path():
@@ -34,23 +36,31 @@ class CodeInterpreterTool(BaseTool):
 
     def _verify_docker_image(self) -> None:
         """
-        Verify if the Docker image is available
+        Verify if the Docker image is available. Optionally use a user-provided Dockerfile.
         """
-        image_tag = "code-interpreter:latest"
         client = docker.from_env()
 
         try:
-            client.images.get(image_tag)
+            client.images.get(self.default_image_tag)
 
         except docker.errors.ImageNotFound:
-            package_path = self._get_installed_package_path()
-            dockerfile_path = os.path.join(package_path, "tools/code_interpreter_tool")
-            if not os.path.exists(dockerfile_path):
-                raise FileNotFoundError(f"Dockerfile not found in {dockerfile_path}")
+            if self.user_docker_file_path and os.path.exists(
+                self.user_docker_file_path
+            ):
+                dockerfile_path = self.user_docker_file_path
+            else:
+                package_path = self._get_installed_package_path()
+                dockerfile_path = os.path.join(
+                    package_path, "tools/code_interpreter_tool"
+                )
+                if not os.path.exists(dockerfile_path):
+                    raise FileNotFoundError(
+                        f"Dockerfile not found in {dockerfile_path}"
+                    )
 
             client.images.build(
                 path=dockerfile_path,
-                tag=image_tag,
+                tag=self.default_image_tag,
                 rm=True,
             )
 
@@ -69,13 +79,25 @@ class CodeInterpreterTool(BaseTool):
             container.exec_run(f"pip install {library}")
 
     def _init_docker_container(self) -> docker.models.containers.Container:
+        container_name = "code-interpreter"
         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:
+            pass  # Container does not exist, no need to remove
+
         return client.containers.run(
-            "code-interpreter",
+            self.default_image_tag,
             detach=True,
             tty=True,
             working_dir="/workspace",
-            name="code-interpreter",
+            name=container_name,
+            volumes={current_path: {"bind": "/workspace", "mode": "rw"}},  # type: ignore
         )
 
     def run_code_in_docker(self, code: str, libraries_used: List[str]) -> str: