mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 15:48:29 +00:00
58fb717ab2
- Add ExecutionStep and ExecutionTrace models to track crew execution steps - Add ExecutionTraceCollector to capture events and build execution traces - Add trace_execution parameter to Crew class (disabled by default) - Add execution_trace field to CrewOutput to return trace data - Integrate trace collection into crew.kickoff() method - Add comprehensive tests covering execution tracing functionality - Add example demonstrating how to use execution tracing - Export new classes in __init__.py Addresses issue #3268: Users can now track the sequence of steps/actions that a crew takes to complete a goal, including agent thoughts, tool calls, and intermediate results, similar to LangGraph's conversation state. Co-Authored-By: João <joao@crewai.com>
311 lines
9.0 KiB
Python
311 lines
9.0 KiB
Python
import pytest
|
|
from unittest.mock import Mock, patch
|
|
from datetime import datetime
|
|
|
|
from crewai import Agent, Crew, Task, Process
|
|
from crewai.crews.execution_trace import ExecutionStep, ExecutionTrace
|
|
from crewai.utilities.execution_trace_collector import ExecutionTraceCollector
|
|
from crewai.tools.base_tool import BaseTool
|
|
|
|
|
|
class MockTool(BaseTool):
|
|
name: str = "mock_tool"
|
|
description: str = "A mock tool for testing"
|
|
|
|
def _run(self, query: str) -> str:
|
|
return f"Mock result for: {query}"
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_llm():
|
|
llm = Mock()
|
|
llm.call.return_value = "Test response"
|
|
llm.supports_stop_words.return_value = True
|
|
llm.stop = []
|
|
return llm
|
|
|
|
|
|
@pytest.fixture
|
|
def test_agent(mock_llm):
|
|
return Agent(
|
|
role="Test Agent",
|
|
goal="Test goal",
|
|
backstory="Test backstory",
|
|
llm=mock_llm,
|
|
tools=[MockTool()]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def test_task(test_agent):
|
|
return Task(
|
|
description="Test task description",
|
|
expected_output="Test expected output",
|
|
agent=test_agent
|
|
)
|
|
|
|
|
|
def test_execution_step_creation():
|
|
"""Test creating an ExecutionStep."""
|
|
step = ExecutionStep(
|
|
timestamp=datetime.now(),
|
|
step_type="agent_thought",
|
|
agent_role="Test Agent",
|
|
task_description="Test task",
|
|
content={"thought": "I need to think about this"},
|
|
metadata={"iteration": 1}
|
|
)
|
|
|
|
assert step.step_type == "agent_thought"
|
|
assert step.agent_role == "Test Agent"
|
|
assert step.content["thought"] == "I need to think about this"
|
|
assert step.metadata["iteration"] == 1
|
|
|
|
|
|
def test_execution_trace_creation():
|
|
"""Test creating an ExecutionTrace."""
|
|
trace = ExecutionTrace()
|
|
|
|
step1 = ExecutionStep(
|
|
timestamp=datetime.now(),
|
|
step_type="task_started",
|
|
content={"task_id": "1"}
|
|
)
|
|
|
|
step2 = ExecutionStep(
|
|
timestamp=datetime.now(),
|
|
step_type="agent_thought",
|
|
agent_role="Test Agent",
|
|
content={"thought": "Starting work"}
|
|
)
|
|
|
|
trace.add_step(step1)
|
|
trace.add_step(step2)
|
|
|
|
assert trace.total_steps == 2
|
|
assert len(trace.steps) == 2
|
|
assert trace.get_steps_by_type("task_started") == [step1]
|
|
assert trace.get_steps_by_agent("Test Agent") == [step2]
|
|
|
|
|
|
def test_execution_trace_collector():
|
|
"""Test the ExecutionTraceCollector."""
|
|
collector = ExecutionTraceCollector()
|
|
|
|
collector.start_collecting()
|
|
assert collector.is_collecting is True
|
|
assert collector.trace.start_time is not None
|
|
|
|
trace = collector.stop_collecting()
|
|
assert collector.is_collecting is False
|
|
assert trace.end_time is not None
|
|
assert isinstance(trace, ExecutionTrace)
|
|
|
|
|
|
@patch('crewai.crew.crewai_event_bus')
|
|
def test_crew_with_execution_trace_enabled(mock_event_bus, test_agent, test_task, mock_llm):
|
|
"""Test crew execution with trace_execution=True."""
|
|
crew = Crew(
|
|
agents=[test_agent],
|
|
tasks=[test_task],
|
|
process=Process.sequential,
|
|
trace_execution=True
|
|
)
|
|
|
|
with patch.object(test_task, 'execute_sync') as mock_execute:
|
|
from crewai.tasks.task_output import TaskOutput
|
|
mock_output = TaskOutput(
|
|
description="Test task description",
|
|
raw="Test output",
|
|
agent="Test Agent"
|
|
)
|
|
mock_execute.return_value = mock_output
|
|
|
|
result = crew.kickoff()
|
|
|
|
assert result.execution_trace is not None
|
|
assert isinstance(result.execution_trace, ExecutionTrace)
|
|
assert result.execution_trace.start_time is not None
|
|
assert result.execution_trace.end_time is not None
|
|
|
|
|
|
@patch('crewai.crew.crewai_event_bus')
|
|
def test_crew_without_execution_trace(mock_event_bus, test_agent, test_task, mock_llm):
|
|
"""Test crew execution with trace_execution=False (default)."""
|
|
crew = Crew(
|
|
agents=[test_agent],
|
|
tasks=[test_task],
|
|
process=Process.sequential,
|
|
trace_execution=False
|
|
)
|
|
|
|
with patch.object(test_task, 'execute_sync') as mock_execute:
|
|
from crewai.tasks.task_output import TaskOutput
|
|
mock_output = TaskOutput(
|
|
description="Test task description",
|
|
raw="Test output",
|
|
agent="Test Agent"
|
|
)
|
|
mock_execute.return_value = mock_output
|
|
|
|
result = crew.kickoff()
|
|
|
|
assert result.execution_trace is None
|
|
|
|
|
|
def test_execution_trace_with_multiple_agents_and_tasks(mock_llm):
|
|
"""Test execution trace with multiple agents and tasks."""
|
|
agent1 = Agent(
|
|
role="Agent 1",
|
|
goal="Goal 1",
|
|
backstory="Backstory 1",
|
|
llm=mock_llm
|
|
)
|
|
|
|
agent2 = Agent(
|
|
role="Agent 2",
|
|
goal="Goal 2",
|
|
backstory="Backstory 2",
|
|
llm=mock_llm
|
|
)
|
|
|
|
task1 = Task(
|
|
description="Task 1",
|
|
expected_output="Output 1",
|
|
agent=agent1
|
|
)
|
|
|
|
task2 = Task(
|
|
description="Task 2",
|
|
expected_output="Output 2",
|
|
agent=agent2
|
|
)
|
|
|
|
crew = Crew(
|
|
agents=[agent1, agent2],
|
|
tasks=[task1, task2],
|
|
process=Process.sequential,
|
|
trace_execution=True
|
|
)
|
|
|
|
with patch.object(task1, 'execute_sync') as mock_execute1, \
|
|
patch.object(task2, 'execute_sync') as mock_execute2:
|
|
|
|
from crewai.tasks.task_output import TaskOutput
|
|
|
|
mock_output1 = TaskOutput(
|
|
description="Task 1",
|
|
raw="Output 1",
|
|
agent="Agent 1"
|
|
)
|
|
|
|
mock_output2 = TaskOutput(
|
|
description="Task 2",
|
|
raw="Output 2",
|
|
agent="Agent 2"
|
|
)
|
|
|
|
mock_execute1.return_value = mock_output1
|
|
mock_execute2.return_value = mock_output2
|
|
|
|
result = crew.kickoff()
|
|
|
|
assert result.execution_trace is not None
|
|
agent1_steps = result.execution_trace.get_steps_by_agent("Agent 1")
|
|
agent2_steps = result.execution_trace.get_steps_by_agent("Agent 2")
|
|
|
|
assert len(agent1_steps) >= 0
|
|
assert len(agent2_steps) >= 0
|
|
|
|
|
|
def test_execution_trace_step_types():
|
|
"""Test that different step types are properly categorized."""
|
|
trace = ExecutionTrace()
|
|
|
|
steps_data = [
|
|
("task_started", "Task 1", {}),
|
|
("agent_thought", "Agent 1", {"thought": "I need to analyze this"}),
|
|
("tool_call_started", "Agent 1", {"tool_name": "search", "args": {"query": "test"}}),
|
|
("tool_call_completed", "Agent 1", {"tool_name": "search", "output": "results"}),
|
|
("agent_execution_completed", "Agent 1", {"output": "Final answer"}),
|
|
("task_completed", "Task 1", {"output": "Task complete"}),
|
|
]
|
|
|
|
for step_type, agent_role, content in steps_data:
|
|
step = ExecutionStep(
|
|
timestamp=datetime.now(),
|
|
step_type=step_type,
|
|
agent_role=agent_role if "agent" in step_type or "tool" in step_type else None,
|
|
task_description="Task 1" if "task" in step_type else None,
|
|
content=content
|
|
)
|
|
trace.add_step(step)
|
|
|
|
assert len(trace.get_steps_by_type("task_started")) == 1
|
|
assert len(trace.get_steps_by_type("agent_thought")) == 1
|
|
assert len(trace.get_steps_by_type("tool_call_started")) == 1
|
|
assert len(trace.get_steps_by_type("tool_call_completed")) == 1
|
|
assert len(trace.get_steps_by_type("agent_execution_completed")) == 1
|
|
assert len(trace.get_steps_by_type("task_completed")) == 1
|
|
|
|
agent_steps = trace.get_steps_by_agent("Agent 1")
|
|
assert len(agent_steps) == 4
|
|
|
|
|
|
def test_execution_trace_with_async_tasks(mock_llm):
|
|
"""Test execution trace with async tasks."""
|
|
agent = Agent(
|
|
role="Async Agent",
|
|
goal="Async goal",
|
|
backstory="Async backstory",
|
|
llm=mock_llm
|
|
)
|
|
|
|
task = Task(
|
|
description="Async task",
|
|
expected_output="Async output",
|
|
agent=agent,
|
|
async_execution=True
|
|
)
|
|
|
|
crew = Crew(
|
|
agents=[agent],
|
|
tasks=[task],
|
|
process=Process.sequential,
|
|
trace_execution=True
|
|
)
|
|
|
|
with patch.object(task, 'execute_async') as mock_execute_async:
|
|
from concurrent.futures import Future
|
|
from crewai.tasks.task_output import TaskOutput
|
|
|
|
future = Future()
|
|
mock_output = TaskOutput(
|
|
description="Async task",
|
|
raw="Async output",
|
|
agent="Async Agent"
|
|
)
|
|
future.set_result(mock_output)
|
|
mock_execute_async.return_value = future
|
|
|
|
result = crew.kickoff()
|
|
|
|
assert result.execution_trace is not None
|
|
assert isinstance(result.execution_trace, ExecutionTrace)
|
|
|
|
|
|
def test_execution_trace_error_handling():
|
|
"""Test execution trace handles errors gracefully."""
|
|
collector = ExecutionTraceCollector()
|
|
|
|
collector.start_collecting()
|
|
|
|
mock_event = Mock()
|
|
mock_event.agent = Mock()
|
|
mock_event.agent.role = "Test Agent"
|
|
|
|
collector._handle_agent_started(mock_event)
|
|
|
|
trace = collector.stop_collecting()
|
|
assert isinstance(trace, ExecutionTrace)
|