andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-11 00:58:30 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

improved imports from docker for better type hinting

Browse Source
This commit is contained in:
Tom Mahler
2024-12-24 12:17:57 +02:00
parent bb19f1c74c
commit 331840e6cc
1 changed files with 9 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -2,7 +2,9 @@ import importlib.util
import os import os
from typing import List, Optional, Type 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 crewai.tools import BaseTool
from pydantic import BaseModel, Field 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. Verify if the Docker image is available. Optionally use a user-provided Dockerfile.
""" """
client = docker.from_env() client = docker_from_env()
try: try:
client.images.get(self.default_image_tag) 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): if self.user_dockerfile_path and os.path.exists(self.user_dockerfile_path):
dockerfile_path = self.user_dockerfile_path dockerfile_path = self.user_dockerfile_path
else: else:
@@ -73,7 +75,7 @@ class CodeInterpreterTool(BaseTool):
return self.run_code_in_docker(code, libraries_used) return self.run_code_in_docker(code, libraries_used)
def _install_libraries( def _install_libraries(
self, container: docker.models.containers.Container, libraries: List[str] self, container: Container, libraries: List[str]
) -> None: ) -> None:
""" """
Install missing libraries in the Docker container Install missing libraries in the Docker container
@@ -81,9 +83,9 @@ class CodeInterpreterTool(BaseTool):
for library in libraries: for library in libraries:
container.exec_run(["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" container_name = "code-interpreter"
client = docker.from_env() client = docker_from_env()
current_path = os.getcwd() current_path = os.getcwd()
# Check if the container is already running # Check if the container is already running
@@ -91,7 +93,7 @@ class CodeInterpreterTool(BaseTool):
existing_container = client.containers.get(container_name) existing_container = client.containers.get(container_name)
existing_container.stop() existing_container.stop()
existing_container.remove() existing_container.remove()
except docker.errors.NotFound: except NotFound:
pass # Container does not exist, no need to remove pass # Container does not exist, no need to remove
return client.containers.run( return client.containers.run(