mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-11 17:18:29 +00:00
e16606672a
git-subtree-dir: packages/tools git-subtree-split: 78317b9c127f18bd040c1d77e3c0840cdc9a5b38
543 lines
20 KiB
Python
543 lines
20 KiB
Python
"""Toolkit for working with AWS Bedrock Code Interpreter."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from typing import TYPE_CHECKING, Dict, List, Tuple, Optional, Type, Any
|
|
|
|
from crewai.tools import BaseTool
|
|
from pydantic import BaseModel, Field
|
|
|
|
if TYPE_CHECKING:
|
|
from bedrock_agentcore.tools.code_interpreter_client import CodeInterpreter
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def extract_output_from_stream(response):
|
|
"""
|
|
Extract output from code interpreter response stream
|
|
|
|
Args:
|
|
response: Response from code interpreter execution
|
|
|
|
Returns:
|
|
Extracted output as string
|
|
"""
|
|
output = []
|
|
for event in response["stream"]:
|
|
if "result" in event:
|
|
result = event["result"]
|
|
for content_item in result["content"]:
|
|
if content_item["type"] == "text":
|
|
output.append(content_item["text"])
|
|
if content_item["type"] == "resource":
|
|
resource = content_item["resource"]
|
|
if "text" in resource:
|
|
file_path = resource["uri"].replace("file://", "")
|
|
file_content = resource["text"]
|
|
output.append(f"==== File: {file_path} ====\n{file_content}\n")
|
|
else:
|
|
output.append(json.dumps(resource))
|
|
|
|
return "\n".join(output)
|
|
|
|
|
|
# Input schemas
|
|
class ExecuteCodeInput(BaseModel):
|
|
"""Input for ExecuteCode."""
|
|
code: str = Field(description="The code to execute")
|
|
language: str = Field(default="python", description="The programming language of the code")
|
|
clear_context: bool = Field(default=False, description="Whether to clear execution context")
|
|
thread_id: str = Field(default="default", description="Thread ID for the code interpreter session")
|
|
|
|
|
|
class ExecuteCommandInput(BaseModel):
|
|
"""Input for ExecuteCommand."""
|
|
command: str = Field(description="The command to execute")
|
|
thread_id: str = Field(default="default", description="Thread ID for the code interpreter session")
|
|
|
|
|
|
class ReadFilesInput(BaseModel):
|
|
"""Input for ReadFiles."""
|
|
paths: List[str] = Field(description="List of file paths to read")
|
|
thread_id: str = Field(default="default", description="Thread ID for the code interpreter session")
|
|
|
|
|
|
class ListFilesInput(BaseModel):
|
|
"""Input for ListFiles."""
|
|
directory_path: str = Field(default="", description="Path to the directory to list")
|
|
thread_id: str = Field(default="default", description="Thread ID for the code interpreter session")
|
|
|
|
|
|
class DeleteFilesInput(BaseModel):
|
|
"""Input for DeleteFiles."""
|
|
paths: List[str] = Field(description="List of file paths to delete")
|
|
thread_id: str = Field(default="default", description="Thread ID for the code interpreter session")
|
|
|
|
|
|
class WriteFilesInput(BaseModel):
|
|
"""Input for WriteFiles."""
|
|
files: List[Dict[str, str]] = Field(description="List of dictionaries with path and text fields")
|
|
thread_id: str = Field(default="default", description="Thread ID for the code interpreter session")
|
|
|
|
|
|
class StartCommandInput(BaseModel):
|
|
"""Input for StartCommand."""
|
|
command: str = Field(description="The command to execute asynchronously")
|
|
thread_id: str = Field(default="default", description="Thread ID for the code interpreter session")
|
|
|
|
|
|
class GetTaskInput(BaseModel):
|
|
"""Input for GetTask."""
|
|
task_id: str = Field(description="The ID of the task to check")
|
|
thread_id: str = Field(default="default", description="Thread ID for the code interpreter session")
|
|
|
|
|
|
class StopTaskInput(BaseModel):
|
|
"""Input for StopTask."""
|
|
task_id: str = Field(description="The ID of the task to stop")
|
|
thread_id: str = Field(default="default", description="Thread ID for the code interpreter session")
|
|
|
|
|
|
# Tool classes
|
|
class ExecuteCodeTool(BaseTool):
|
|
"""Tool for executing code in various languages."""
|
|
name: str = "execute_code"
|
|
description: str = "Execute code in various languages (primarily Python)"
|
|
args_schema: Type[BaseModel] = ExecuteCodeInput
|
|
toolkit: Any = Field(default=None, exclude=True)
|
|
|
|
def __init__(self, toolkit):
|
|
super().__init__()
|
|
self.toolkit = toolkit
|
|
|
|
def _run(self, code: str, language: str = "python", clear_context: bool = False, thread_id: str = "default") -> str:
|
|
try:
|
|
# Get or create code interpreter
|
|
code_interpreter = self.toolkit._get_or_create_interpreter(thread_id=thread_id)
|
|
|
|
# Execute code
|
|
response = code_interpreter.invoke(
|
|
method="executeCode",
|
|
params={"code": code, "language": language, "clearContext": clear_context},
|
|
)
|
|
|
|
return extract_output_from_stream(response)
|
|
except Exception as e:
|
|
return f"Error executing code: {str(e)}"
|
|
|
|
async def _arun(self, code: str, language: str = "python", clear_context: bool = False, thread_id: str = "default") -> str:
|
|
# Use _run as we're working with a synchronous API that's thread-safe
|
|
return self._run(code=code, language=language, clear_context=clear_context, thread_id=thread_id)
|
|
|
|
|
|
class ExecuteCommandTool(BaseTool):
|
|
"""Tool for running shell commands in the code interpreter environment."""
|
|
name: str = "execute_command"
|
|
description: str = "Run shell commands in the code interpreter environment"
|
|
args_schema: Type[BaseModel] = ExecuteCommandInput
|
|
toolkit: Any = Field(default=None, exclude=True)
|
|
|
|
def __init__(self, toolkit):
|
|
super().__init__()
|
|
self.toolkit = toolkit
|
|
|
|
def _run(self, command: str, thread_id: str = "default") -> str:
|
|
try:
|
|
# Get or create code interpreter
|
|
code_interpreter = self.toolkit._get_or_create_interpreter(thread_id=thread_id)
|
|
|
|
# Execute command
|
|
response = code_interpreter.invoke(
|
|
method="executeCommand", params={"command": command}
|
|
)
|
|
|
|
return extract_output_from_stream(response)
|
|
except Exception as e:
|
|
return f"Error executing command: {str(e)}"
|
|
|
|
async def _arun(self, command: str, thread_id: str = "default") -> str:
|
|
# Use _run as we're working with a synchronous API that's thread-safe
|
|
return self._run(command=command, thread_id=thread_id)
|
|
|
|
|
|
class ReadFilesTool(BaseTool):
|
|
"""Tool for reading content of files in the environment."""
|
|
name: str = "read_files"
|
|
description: str = "Read content of files in the environment"
|
|
args_schema: Type[BaseModel] = ReadFilesInput
|
|
toolkit: Any = Field(default=None, exclude=True)
|
|
|
|
def __init__(self, toolkit):
|
|
super().__init__()
|
|
self.toolkit = toolkit
|
|
|
|
def _run(self, paths: List[str], thread_id: str = "default") -> str:
|
|
try:
|
|
# Get or create code interpreter
|
|
code_interpreter = self.toolkit._get_or_create_interpreter(thread_id=thread_id)
|
|
|
|
# Read files
|
|
response = code_interpreter.invoke(method="readFiles", params={"paths": paths})
|
|
|
|
return extract_output_from_stream(response)
|
|
except Exception as e:
|
|
return f"Error reading files: {str(e)}"
|
|
|
|
async def _arun(self, paths: List[str], thread_id: str = "default") -> str:
|
|
# Use _run as we're working with a synchronous API that's thread-safe
|
|
return self._run(paths=paths, thread_id=thread_id)
|
|
|
|
|
|
class ListFilesTool(BaseTool):
|
|
"""Tool for listing files in directories in the environment."""
|
|
name: str = "list_files"
|
|
description: str = "List files in directories in the environment"
|
|
args_schema: Type[BaseModel] = ListFilesInput
|
|
toolkit: Any = Field(default=None, exclude=True)
|
|
|
|
def __init__(self, toolkit):
|
|
super().__init__()
|
|
self.toolkit = toolkit
|
|
|
|
def _run(self, directory_path: str = "", thread_id: str = "default") -> str:
|
|
try:
|
|
# Get or create code interpreter
|
|
code_interpreter = self.toolkit._get_or_create_interpreter(thread_id=thread_id)
|
|
|
|
# List files
|
|
response = code_interpreter.invoke(
|
|
method="listFiles", params={"directoryPath": directory_path}
|
|
)
|
|
|
|
return extract_output_from_stream(response)
|
|
except Exception as e:
|
|
return f"Error listing files: {str(e)}"
|
|
|
|
async def _arun(self, directory_path: str = "", thread_id: str = "default") -> str:
|
|
# Use _run as we're working with a synchronous API that's thread-safe
|
|
return self._run(directory_path=directory_path, thread_id=thread_id)
|
|
|
|
|
|
class DeleteFilesTool(BaseTool):
|
|
"""Tool for removing files from the environment."""
|
|
name: str = "delete_files"
|
|
description: str = "Remove files from the environment"
|
|
args_schema: Type[BaseModel] = DeleteFilesInput
|
|
toolkit: Any = Field(default=None, exclude=True)
|
|
|
|
def __init__(self, toolkit):
|
|
super().__init__()
|
|
self.toolkit = toolkit
|
|
|
|
def _run(self, paths: List[str], thread_id: str = "default") -> str:
|
|
try:
|
|
# Get or create code interpreter
|
|
code_interpreter = self.toolkit._get_or_create_interpreter(thread_id=thread_id)
|
|
|
|
# Remove files
|
|
response = code_interpreter.invoke(
|
|
method="removeFiles", params={"paths": paths}
|
|
)
|
|
|
|
return extract_output_from_stream(response)
|
|
except Exception as e:
|
|
return f"Error deleting files: {str(e)}"
|
|
|
|
async def _arun(self, paths: List[str], thread_id: str = "default") -> str:
|
|
# Use _run as we're working with a synchronous API that's thread-safe
|
|
return self._run(paths=paths, thread_id=thread_id)
|
|
|
|
|
|
class WriteFilesTool(BaseTool):
|
|
"""Tool for creating or updating files in the environment."""
|
|
name: str = "write_files"
|
|
description: str = "Create or update files in the environment"
|
|
args_schema: Type[BaseModel] = WriteFilesInput
|
|
toolkit: Any = Field(default=None, exclude=True)
|
|
|
|
def __init__(self, toolkit):
|
|
super().__init__()
|
|
self.toolkit = toolkit
|
|
|
|
def _run(self, files: List[Dict[str, str]], thread_id: str = "default") -> str:
|
|
try:
|
|
# Get or create code interpreter
|
|
code_interpreter = self.toolkit._get_or_create_interpreter(thread_id=thread_id)
|
|
|
|
# Write files
|
|
response = code_interpreter.invoke(
|
|
method="writeFiles", params={"content": files}
|
|
)
|
|
|
|
return extract_output_from_stream(response)
|
|
except Exception as e:
|
|
return f"Error writing files: {str(e)}"
|
|
|
|
async def _arun(self, files: List[Dict[str, str]], thread_id: str = "default") -> str:
|
|
# Use _run as we're working with a synchronous API that's thread-safe
|
|
return self._run(files=files, thread_id=thread_id)
|
|
|
|
|
|
class StartCommandTool(BaseTool):
|
|
"""Tool for starting long-running commands asynchronously."""
|
|
name: str = "start_command_execution"
|
|
description: str = "Start long-running commands asynchronously"
|
|
args_schema: Type[BaseModel] = StartCommandInput
|
|
toolkit: Any = Field(default=None, exclude=True)
|
|
|
|
def __init__(self, toolkit):
|
|
super().__init__()
|
|
self.toolkit = toolkit
|
|
|
|
def _run(self, command: str, thread_id: str = "default") -> str:
|
|
try:
|
|
# Get or create code interpreter
|
|
code_interpreter = self.toolkit._get_or_create_interpreter(thread_id=thread_id)
|
|
|
|
# Start command execution
|
|
response = code_interpreter.invoke(
|
|
method="startCommandExecution", params={"command": command}
|
|
)
|
|
|
|
return extract_output_from_stream(response)
|
|
except Exception as e:
|
|
return f"Error starting command: {str(e)}"
|
|
|
|
async def _arun(self, command: str, thread_id: str = "default") -> str:
|
|
# Use _run as we're working with a synchronous API that's thread-safe
|
|
return self._run(command=command, thread_id=thread_id)
|
|
|
|
|
|
class GetTaskTool(BaseTool):
|
|
"""Tool for checking status of async tasks."""
|
|
name: str = "get_task"
|
|
description: str = "Check status of async tasks"
|
|
args_schema: Type[BaseModel] = GetTaskInput
|
|
toolkit: Any = Field(default=None, exclude=True)
|
|
|
|
def __init__(self, toolkit):
|
|
super().__init__()
|
|
self.toolkit = toolkit
|
|
|
|
def _run(self, task_id: str, thread_id: str = "default") -> str:
|
|
try:
|
|
# Get or create code interpreter
|
|
code_interpreter = self.toolkit._get_or_create_interpreter(thread_id=thread_id)
|
|
|
|
# Get task status
|
|
response = code_interpreter.invoke(method="getTask", params={"taskId": task_id})
|
|
|
|
return extract_output_from_stream(response)
|
|
except Exception as e:
|
|
return f"Error getting task status: {str(e)}"
|
|
|
|
async def _arun(self, task_id: str, thread_id: str = "default") -> str:
|
|
# Use _run as we're working with a synchronous API that's thread-safe
|
|
return self._run(task_id=task_id, thread_id=thread_id)
|
|
|
|
|
|
class StopTaskTool(BaseTool):
|
|
"""Tool for stopping running tasks."""
|
|
name: str = "stop_task"
|
|
description: str = "Stop running tasks"
|
|
args_schema: Type[BaseModel] = StopTaskInput
|
|
toolkit: Any = Field(default=None, exclude=True)
|
|
|
|
def __init__(self, toolkit):
|
|
super().__init__()
|
|
self.toolkit = toolkit
|
|
|
|
def _run(self, task_id: str, thread_id: str = "default") -> str:
|
|
try:
|
|
# Get or create code interpreter
|
|
code_interpreter = self.toolkit._get_or_create_interpreter(thread_id=thread_id)
|
|
|
|
# Stop task
|
|
response = code_interpreter.invoke(
|
|
method="stopTask", params={"taskId": task_id}
|
|
)
|
|
|
|
return extract_output_from_stream(response)
|
|
except Exception as e:
|
|
return f"Error stopping task: {str(e)}"
|
|
|
|
async def _arun(self, task_id: str, thread_id: str = "default") -> str:
|
|
# Use _run as we're working with a synchronous API that's thread-safe
|
|
return self._run(task_id=task_id, thread_id=thread_id)
|
|
|
|
|
|
class CodeInterpreterToolkit:
|
|
"""Toolkit for working with AWS Bedrock code interpreter environment.
|
|
|
|
This toolkit provides a set of tools for working with a remote code interpreter environment:
|
|
|
|
* execute_code - Run code in various languages (primarily Python)
|
|
* execute_command - Run shell commands
|
|
* read_files - Read content of files in the environment
|
|
* list_files - List files in directories
|
|
* delete_files - Remove files from the environment
|
|
* write_files - Create or update files
|
|
* start_command_execution - Start long-running commands asynchronously
|
|
* get_task - Check status of async tasks
|
|
* stop_task - Stop running tasks
|
|
|
|
The toolkit lazily initializes the code interpreter session on first use.
|
|
It supports multiple threads by maintaining separate code interpreter sessions for each thread ID.
|
|
|
|
Example:
|
|
```python
|
|
from crewai import Agent, Task, Crew
|
|
from crewai_tools.aws.bedrock.code_interpreter import create_code_interpreter_toolkit
|
|
|
|
# Create the code interpreter toolkit
|
|
toolkit, code_tools = create_code_interpreter_toolkit(region="us-west-2")
|
|
|
|
# Create a CrewAI agent that uses the code interpreter tools
|
|
developer_agent = Agent(
|
|
role="Python Developer",
|
|
goal="Create and execute Python code to solve problems",
|
|
backstory="You're a skilled Python developer with expertise in data analysis.",
|
|
tools=code_tools
|
|
)
|
|
|
|
# Create a task for the agent
|
|
coding_task = Task(
|
|
description="Write a Python function that calculates the factorial of a number and test it.",
|
|
agent=developer_agent
|
|
)
|
|
|
|
# Create and run the crew
|
|
crew = Crew(
|
|
agents=[developer_agent],
|
|
tasks=[coding_task]
|
|
)
|
|
result = crew.kickoff()
|
|
|
|
# Clean up resources when done
|
|
import asyncio
|
|
asyncio.run(toolkit.cleanup())
|
|
```
|
|
"""
|
|
|
|
def __init__(self, region: str = "us-west-2"):
|
|
"""
|
|
Initialize the toolkit
|
|
|
|
Args:
|
|
region: AWS region for the code interpreter
|
|
"""
|
|
self.region = region
|
|
self._code_interpreters: Dict[str, CodeInterpreter] = {}
|
|
self.tools: List[BaseTool] = []
|
|
self._setup_tools()
|
|
|
|
def _setup_tools(self) -> None:
|
|
"""Initialize tools without creating any code interpreter sessions."""
|
|
self.tools = [
|
|
ExecuteCodeTool(self),
|
|
ExecuteCommandTool(self),
|
|
ReadFilesTool(self),
|
|
ListFilesTool(self),
|
|
DeleteFilesTool(self),
|
|
WriteFilesTool(self),
|
|
StartCommandTool(self),
|
|
GetTaskTool(self),
|
|
StopTaskTool(self)
|
|
]
|
|
|
|
def _get_or_create_interpreter(
|
|
self, thread_id: str = "default"
|
|
) -> CodeInterpreter:
|
|
"""Get or create a code interpreter for the specified thread.
|
|
|
|
Args:
|
|
thread_id: Thread ID for the code interpreter session
|
|
|
|
Returns:
|
|
CodeInterpreter instance
|
|
"""
|
|
if thread_id in self._code_interpreters:
|
|
return self._code_interpreters[thread_id]
|
|
|
|
# Create a new code interpreter for this thread
|
|
from bedrock_agentcore.tools.code_interpreter_client import CodeInterpreter
|
|
code_interpreter = CodeInterpreter(region=self.region)
|
|
code_interpreter.start()
|
|
logger.info(
|
|
f"Started code interpreter with session_id:{code_interpreter.session_id} for thread:{thread_id}"
|
|
)
|
|
|
|
# Store the interpreter
|
|
self._code_interpreters[thread_id] = code_interpreter
|
|
return code_interpreter
|
|
|
|
|
|
def get_tools(self) -> List[BaseTool]:
|
|
"""
|
|
Get the list of code interpreter tools
|
|
|
|
Returns:
|
|
List of CrewAI tools
|
|
"""
|
|
return self.tools
|
|
|
|
def get_tools_by_name(self) -> Dict[str, BaseTool]:
|
|
"""
|
|
Get a dictionary of tools mapped by their names
|
|
|
|
Returns:
|
|
Dictionary of {tool_name: tool}
|
|
"""
|
|
return {tool.name: tool for tool in self.tools}
|
|
|
|
async def cleanup(self, thread_id: Optional[str] = None) -> None:
|
|
"""Clean up resources
|
|
|
|
Args:
|
|
thread_id: Optional thread ID to clean up. If None, cleans up all sessions.
|
|
"""
|
|
if thread_id:
|
|
# Clean up a specific thread's session
|
|
if thread_id in self._code_interpreters:
|
|
try:
|
|
self._code_interpreters[thread_id].stop()
|
|
del self._code_interpreters[thread_id]
|
|
logger.info(
|
|
f"Code interpreter session for thread {thread_id} cleaned up"
|
|
)
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Error stopping code interpreter for thread {thread_id}: {e}"
|
|
)
|
|
else:
|
|
# Clean up all sessions
|
|
thread_ids = list(self._code_interpreters.keys())
|
|
for tid in thread_ids:
|
|
try:
|
|
self._code_interpreters[tid].stop()
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Error stopping code interpreter for thread {tid}: {e}"
|
|
)
|
|
|
|
self._code_interpreters = {}
|
|
logger.info("All code interpreter sessions cleaned up")
|
|
|
|
|
|
def create_code_interpreter_toolkit(
|
|
region: str = "us-west-2",
|
|
) -> Tuple[CodeInterpreterToolkit, List[BaseTool]]:
|
|
"""
|
|
Create a CodeInterpreterToolkit
|
|
|
|
Args:
|
|
region: AWS region for code interpreter
|
|
|
|
Returns:
|
|
Tuple of (toolkit, tools)
|
|
"""
|
|
toolkit = CodeInterpreterToolkit(region=region)
|
|
tools = toolkit.get_tools()
|
|
return toolkit, tools |