diff --git a/lib/crewai/src/crewai/cli/cli.py b/lib/crewai/src/crewai/cli/cli.py index c3c3fac39..a25fb41d8 100644 --- a/lib/crewai/src/crewai/cli/cli.py +++ b/lib/crewai/src/crewai/cli/cli.py @@ -139,16 +139,29 @@ def train(n_iterations: int, filename: str) -> None: type=str, help="Replay the crew from this task ID, including all subsequent tasks.", ) -def replay(task_id: str) -> None: - """ - Replay the crew execution from a specific task. +@click.option( + "-f", + "--filename", + "trained_agents_file", + type=str, + default=None, + help=( + "Path to a trained-agents pickle (produced by `crewai train -f`). " + "When set, agents load suggestions from this file instead of the " + "default trained_agents_data.pkl. Equivalent to setting " + "CREWAI_TRAINED_AGENTS_FILE." + ), +) +def replay(task_id: str, trained_agents_file: str | None) -> None: + """Replay the crew execution from a specific task. Args: - task_id (str): The ID of the task to replay from. + task_id: The ID of the task to replay from. + trained_agents_file: Optional trained-agents pickle path. """ try: click.echo(f"Replaying the crew from task {task_id}") - replay_task_command(task_id) + replay_task_command(task_id, trained_agents_file=trained_agents_file) except Exception as e: click.echo(f"An error occurred while replaying: {e}", err=True) @@ -332,10 +345,23 @@ def memory( default="gpt-4o-mini", help="LLM Model to run the tests on the Crew. For now only accepting only OpenAI models.", ) -def test(n_iterations: int, model: str) -> None: +@click.option( + "-f", + "--filename", + "trained_agents_file", + type=str, + default=None, + help=( + "Path to a trained-agents pickle (produced by `crewai train -f`). " + "When set, agents load suggestions from this file instead of the " + "default trained_agents_data.pkl. Equivalent to setting " + "CREWAI_TRAINED_AGENTS_FILE." + ), +) +def test(n_iterations: int, model: str, trained_agents_file: str | None) -> None: """Test the crew and evaluate the results.""" click.echo(f"Testing the crew for {n_iterations} iterations with model {model}") - evaluate_crew(n_iterations, model) + evaluate_crew(n_iterations, model, trained_agents_file=trained_agents_file) @crewai.command( diff --git a/lib/crewai/src/crewai/cli/evaluate_crew.py b/lib/crewai/src/crewai/cli/evaluate_crew.py index a158eeaa7..834c3c636 100644 --- a/lib/crewai/src/crewai/cli/evaluate_crew.py +++ b/lib/crewai/src/crewai/cli/evaluate_crew.py @@ -2,22 +2,33 @@ import subprocess import click +from crewai.cli.utils import build_env_with_all_tool_credentials +from crewai.utilities.constants import CREWAI_TRAINED_AGENTS_FILE_ENV -def evaluate_crew(n_iterations: int, model: str) -> None: - """ - Test and Evaluate the crew by running a command in the UV environment. + +def evaluate_crew( + n_iterations: int, model: str, trained_agents_file: str | None = None +) -> None: + """Test and Evaluate the crew by running a command in the UV environment. Args: - n_iterations (int): The number of iterations to test the crew. - model (str): The model to test the crew with. + n_iterations: The number of iterations to test the crew. + model: The model to test the crew with. + trained_agents_file: Optional trained-agents pickle path forwarded to + the subprocess via the ``CREWAI_TRAINED_AGENTS_FILE`` env var. """ command = ["uv", "run", "test", str(n_iterations), model] + env = build_env_with_all_tool_credentials() + if trained_agents_file: + env[CREWAI_TRAINED_AGENTS_FILE_ENV] = trained_agents_file try: if n_iterations <= 0: raise ValueError("The number of iterations must be a positive integer.") - result = subprocess.run(command, capture_output=False, text=True, check=True) # noqa: S603 + result = subprocess.run( # noqa: S603 + command, capture_output=False, text=True, check=True, env=env + ) if result.stderr: click.echo(result.stderr, err=True) diff --git a/lib/crewai/src/crewai/cli/replay_from_task.py b/lib/crewai/src/crewai/cli/replay_from_task.py index f3c8ae557..f97b22d8a 100644 --- a/lib/crewai/src/crewai/cli/replay_from_task.py +++ b/lib/crewai/src/crewai/cli/replay_from_task.py @@ -2,18 +2,27 @@ import subprocess import click +from crewai.cli.utils import build_env_with_all_tool_credentials +from crewai.utilities.constants import CREWAI_TRAINED_AGENTS_FILE_ENV -def replay_task_command(task_id: str) -> None: - """ - Replay the crew execution from a specific task. + +def replay_task_command(task_id: str, trained_agents_file: str | None = None) -> None: + """Replay the crew execution from a specific task. Args: - task_id (str): The ID of the task to replay from. + task_id: The ID of the task to replay from. + trained_agents_file: Optional trained-agents pickle path forwarded to + the subprocess via the ``CREWAI_TRAINED_AGENTS_FILE`` env var. """ command = ["uv", "run", "replay", task_id] + env = build_env_with_all_tool_credentials() + if trained_agents_file: + env[CREWAI_TRAINED_AGENTS_FILE_ENV] = trained_agents_file try: - result = subprocess.run(command, capture_output=False, text=True, check=True) # noqa: S603 + result = subprocess.run( # noqa: S603 + command, capture_output=False, text=True, check=True, env=env + ) if result.stderr: click.echo(result.stderr, err=True) diff --git a/lib/crewai/tests/cli/test_cli.py b/lib/crewai/tests/cli/test_cli.py index b324294b1..decb15e70 100644 --- a/lib/crewai/tests/cli/test_cli.py +++ b/lib/crewai/tests/cli/test_cli.py @@ -307,7 +307,7 @@ def test_version_command_with_tools(runner): def test_test_default_iterations(evaluate_crew, runner): result = runner.invoke(test) - evaluate_crew.assert_called_once_with(3, "gpt-4o-mini") + evaluate_crew.assert_called_once_with(3, "gpt-4o-mini", trained_agents_file=None) assert result.exit_code == 0 assert "Testing the crew for 3 iterations with model gpt-4o-mini" in result.output @@ -316,7 +316,7 @@ def test_test_default_iterations(evaluate_crew, runner): def test_test_custom_iterations(evaluate_crew, runner): result = runner.invoke(test, ["--n_iterations", "5", "--model", "gpt-4o"]) - evaluate_crew.assert_called_once_with(5, "gpt-4o") + evaluate_crew.assert_called_once_with(5, "gpt-4o", trained_agents_file=None) assert result.exit_code == 0 assert "Testing the crew for 5 iterations with model gpt-4o" in result.output diff --git a/lib/crewai/tests/cli/test_crew_test.py b/lib/crewai/tests/cli/test_crew_test.py index 83bcd55cc..3ebe0c49a 100644 --- a/lib/crewai/tests/cli/test_crew_test.py +++ b/lib/crewai/tests/cli/test_crew_test.py @@ -27,6 +27,7 @@ def test_crew_success(mock_subprocess_run, n_iterations, model): capture_output=False, text=True, check=True, + env=mock.ANY, ) assert result is None @@ -66,6 +67,7 @@ def test_test_crew_called_process_error(mock_subprocess_run, click): capture_output=False, text=True, check=True, + env=mock.ANY, ) click.echo.assert_has_calls( [ @@ -91,7 +93,30 @@ def test_test_crew_unexpected_exception(mock_subprocess_run, click): capture_output=False, text=True, check=True, + env=mock.ANY, ) click.echo.assert_called_once_with( "An unexpected error occurred: Unexpected error", err=True ) + + +@mock.patch("crewai.cli.evaluate_crew.subprocess.run") +def test_evaluate_crew_sets_trained_agents_env_var(mock_subprocess_run): + mock_subprocess_run.return_value = subprocess.CompletedProcess( + args=["uv", "run", "test", "1", "gpt-4o"], returncode=0 + ) + evaluate_crew.evaluate_crew(1, "gpt-4o", trained_agents_file="my_custom.pkl") + + _, kwargs = mock_subprocess_run.call_args + assert kwargs["env"]["CREWAI_TRAINED_AGENTS_FILE"] == "my_custom.pkl" + + +@mock.patch("crewai.cli.evaluate_crew.subprocess.run") +def test_evaluate_crew_omits_env_var_without_filename(mock_subprocess_run): + mock_subprocess_run.return_value = subprocess.CompletedProcess( + args=["uv", "run", "test", "1", "gpt-4o"], returncode=0 + ) + evaluate_crew.evaluate_crew(1, "gpt-4o") + + _, kwargs = mock_subprocess_run.call_args + assert "CREWAI_TRAINED_AGENTS_FILE" not in kwargs["env"] diff --git a/lib/crewai/tests/cli/test_replay_from_task.py b/lib/crewai/tests/cli/test_replay_from_task.py new file mode 100644 index 000000000..c1752c4f1 --- /dev/null +++ b/lib/crewai/tests/cli/test_replay_from_task.py @@ -0,0 +1,61 @@ +"""Tests for ``crewai replay`` and the trained-agents file plumbing.""" + +import subprocess +from unittest import mock + +from click.testing import CliRunner +import pytest + +from crewai.cli import replay_from_task +from crewai.cli.cli import replay + + +@pytest.fixture +def runner() -> CliRunner: + return CliRunner() + + +@mock.patch("crewai.cli.cli.replay_task_command") +def test_replay_passes_filename(replay_task_command_mock: mock.Mock, runner: CliRunner) -> None: + result = runner.invoke(replay, ["-t", "abc123", "-f", "my_custom.pkl"]) + + replay_task_command_mock.assert_called_once_with( + "abc123", trained_agents_file="my_custom.pkl" + ) + assert result.exit_code == 0 + + +@mock.patch("crewai.cli.cli.replay_task_command") +def test_replay_without_filename_passes_none( + replay_task_command_mock: mock.Mock, runner: CliRunner +) -> None: + result = runner.invoke(replay, ["-t", "abc123"]) + + replay_task_command_mock.assert_called_once_with( + "abc123", trained_agents_file=None + ) + assert result.exit_code == 0 + + +@mock.patch("crewai.cli.replay_from_task.subprocess.run") +def test_replay_task_command_sets_env_var(mock_subprocess_run: mock.Mock) -> None: + mock_subprocess_run.return_value = subprocess.CompletedProcess( + args=["uv", "run", "replay", "abc123"], returncode=0 + ) + replay_from_task.replay_task_command("abc123", trained_agents_file="my_custom.pkl") + + _, kwargs = mock_subprocess_run.call_args + assert kwargs["env"]["CREWAI_TRAINED_AGENTS_FILE"] == "my_custom.pkl" + + +@mock.patch("crewai.cli.replay_from_task.subprocess.run") +def test_replay_task_command_omits_env_var_without_filename( + mock_subprocess_run: mock.Mock, +) -> None: + mock_subprocess_run.return_value = subprocess.CompletedProcess( + args=["uv", "run", "replay", "abc123"], returncode=0 + ) + replay_from_task.replay_task_command("abc123") + + _, kwargs = mock_subprocess_run.call_args + assert "CREWAI_TRAINED_AGENTS_FILE" not in kwargs["env"] \ No newline at end of file