feat: add code-interpreter tool

This commit is contained in:
Eduardo Chiarotti
2024-06-19 20:45:04 -03:00
parent f3bf13e1b1
commit 2b47377a78
3 changed files with 78 additions and 0 deletions

View File

@@ -0,0 +1,21 @@
# Use an official Ubuntu as a parent image
FROM ubuntu:20.04
# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
# Install common utilities
RUN apt-get update && apt-get install -y \
build-essential \
curl \
wget \
software-properties-common
# Install Python
RUN apt-get install -y python3 python3-pip
# Clean up
RUN apt-get clean && rm -rf /var/lib/apt/lists/*
# Set the working directory
WORKDIR /workspace

View File

@@ -0,0 +1,57 @@
from typing import Optional, Type
import docker
from crewai_tools.tools.base_tool import BaseTool
from pydantic.v1 import BaseModel, Field
class FixedCodeInterpreterSchemaSchema(BaseModel):
"""Input for DirectoryReadTool."""
pass
class CodeInterpreterSchema(FixedCodeInterpreterSchemaSchema):
"""Input for DirectoryReadTool."""
code: str = Field(
...,
description="Python3 code used to be interpreted in the Docker container and output the result",
)
class CodeInterpreterTool(BaseTool):
name: str = "Code Interpreter"
description: str = "Interprets Python code in a Docker container"
args_schema: Type[BaseModel] = CodeInterpreterSchema
code: Optional[str] = None
def __init__(self, code: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
if code is not None:
self.code = code
self.description = (
"A tool that can be used to run Python code in a Docker container"
)
self.args_schema = FixedCodeInterpreterSchemaSchema
self._generate_description()
def _run(self, **kwargs):
code = kwargs.get("code", self.code)
return self.run_code_in_docker(code)
def run_code_in_docker(self, code):
client = docker.from_env()
container = client.containers.run(
"code-interpreter",
command=f'python3 -c "{code}"',
detach=True,
working_dir="/workspace",
)
result = container.logs().decode("utf-8")
container.stop()
container.remove()
return result