Add execution_image parameter to Agent for custom Docker images

- Add execution_image field to Agent class for specifying custom Docker images
- Modify get_code_execution_tools() to pass custom image to CodeInterpreterTool
- Add comprehensive tests for the new functionality
- Update documentation with usage examples
- Fixes #2933: Allow customizing CrewAI code executor Docker image

Co-Authored-By: João <joao@crewai.com>
This commit is contained in:
Devin AI
2025-06-02 18:18:14 +00:00
parent c045399d6b
commit 229bbd9bbe
4 changed files with 143 additions and 1 deletions

View File

@@ -0,0 +1,82 @@
import pytest
from unittest.mock import patch, MagicMock
from crewai import Agent
def test_agent_with_custom_execution_image():
"""Test that Agent can be created with custom execution image."""
agent = Agent(
role="Test Agent",
goal="Test goal",
backstory="Test backstory",
allow_code_execution=True,
execution_image="my-custom-image:latest"
)
assert agent.execution_image == "my-custom-image:latest"
def test_agent_without_custom_execution_image():
"""Test that Agent works without custom execution image (default behavior)."""
agent = Agent(
role="Test Agent",
goal="Test goal",
backstory="Test backstory",
allow_code_execution=True
)
assert agent.execution_image is None
@patch('crewai.agent.CodeInterpreterTool')
def test_get_code_execution_tools_with_custom_image(mock_code_interpreter):
"""Test that get_code_execution_tools passes custom image to CodeInterpreterTool."""
agent = Agent(
role="Test Agent",
goal="Test goal",
backstory="Test backstory",
allow_code_execution=True,
execution_image="my-custom-image:latest"
)
tools = agent.get_code_execution_tools()
mock_code_interpreter.assert_called_once_with(
unsafe_mode=False,
default_image_tag="my-custom-image:latest"
)
@patch('crewai.agent.CodeInterpreterTool')
def test_get_code_execution_tools_without_custom_image(mock_code_interpreter):
"""Test that get_code_execution_tools works without custom image."""
agent = Agent(
role="Test Agent",
goal="Test goal",
backstory="Test backstory",
allow_code_execution=True
)
tools = agent.get_code_execution_tools()
mock_code_interpreter.assert_called_once_with(unsafe_mode=False)
@patch('crewai.agent.CodeInterpreterTool')
def test_get_code_execution_tools_with_unsafe_mode_and_custom_image(mock_code_interpreter):
"""Test that both unsafe_mode and custom image work together."""
agent = Agent(
role="Test Agent",
goal="Test goal",
backstory="Test backstory",
allow_code_execution=True,
code_execution_mode="unsafe",
execution_image="my-custom-image:latest"
)
tools = agent.get_code_execution_tools()
mock_code_interpreter.assert_called_once_with(
unsafe_mode=True,
default_image_tag="my-custom-image:latest"
)

View File

@@ -0,0 +1,36 @@
import pytest
from unittest.mock import patch, MagicMock
from crewai import Agent, Task, Crew
@patch('crewai_tools.CodeInterpreterTool')
def test_crew_with_custom_execution_image_integration(mock_code_interpreter_class):
"""Integration test for custom execution image in a Crew workflow."""
mock_tool_instance = MagicMock()
mock_code_interpreter_class.return_value = mock_tool_instance
agent = Agent(
role="Python Developer",
goal="Execute Python code",
backstory="Expert in Python programming",
allow_code_execution=True,
execution_image="python:3.11-slim"
)
task = Task(
description="Calculate 2 + 2",
expected_output="The result of 2 + 2",
agent=agent
)
crew = Crew(
agents=[agent],
tasks=[task]
)
tools = crew._prepare_tools(task, agent)
mock_code_interpreter_class.assert_called_with(
unsafe_mode=False,
default_image_tag="python:3.11-slim"
)