feat: update code

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