From 1a4ac76b1eb11ec70a11784a4f6b1d33948e9171 Mon Sep 17 00:00:00 2001 From: Eduardo Chiarotti Date: Thu, 20 Jun 2024 20:24:47 -0300 Subject: [PATCH] feat: update code --- .../code_interpreter_tool.py | 52 +++++++++---------- 1 file changed, 26 insertions(+), 26 deletions(-) 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 f497a7c96..caafd44e3 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 @@ -30,7 +30,7 @@ class CodeInterpreterTool(BaseTool): args_schema: Type[BaseModel] = CodeInterpreterSchema code: Optional[str] = None - def __init__(self, code: Optional[str] = None, **kwargs): + def __init__(self, code: Optional[str] = None, **kwargs) -> None: super().__init__(**kwargs) if code is not None: self.code = code @@ -38,37 +38,37 @@ class CodeInterpreterTool(BaseTool): self.args_schema = FixedCodeInterpreterSchemaSchema self._generate_description() - def _run(self, **kwargs): + def _run(self, **kwargs) -> str: code = kwargs.get("code", self.code) libraries_used = kwargs.get("libraries_used", None) return self.run_code_in_docker(code, libraries_used) - def run_code_in_docker(self, code, libraries_used): - client = docker.from_env() - - def run_code(container, code): - cmd_to_run = f'python3 -c "{code}"' - exec_result = container.exec_run(cmd_to_run) - return exec_result - - container = client.containers.run( - "code-interpreter", detach=True, tty=True, working_dir="/workspace" - ) - if libraries_used: - self._install_libraries(container, libraries_used.split(",")) - - exec_result = run_code(container, code) - - container.stop() - container.remove() - if exec_result.exit_code != 0: - return f"Something went wrong while running the code: \n{exec_result.output.decode('utf-8')}" - - return exec_result.output.decode("utf-8") - - def _install_libraries(self, container, libraries): + def _install_libraries( + self, container: docker.models.containers.Container, libraries: list[str] + ) -> None: """ Install missing libraries in the Docker container """ for library in libraries: container.exec_run(f"pip install {library}") + + def _init_docker_container(self) -> docker.models.containers.Container: + client = docker.from_env() + return client.containers.run( + "code-interpreter", detach=True, tty=True, working_dir="/workspace" + ) + + def run_code_in_docker(self, code: str, libraries_used: str) -> str: + container = self._init_docker_container() + + if libraries_used: + self._install_libraries(container, libraries_used.split(",")) + + cmd_to_run = f'python3 -c "{code}"' + exec_result = container.exec_run(cmd_to_run) + + container.stop().remove() + + if exec_result.exit_code != 0: + return f"Something went wrong while running the code: \n{exec_result.output.decode('utf-8')}" + return exec_result.output.decode("utf-8")