diff --git a/src/crewai/task.py b/src/crewai/task.py index d0d070805..579a144ae 100644 --- a/src/crewai/task.py +++ b/src/crewai/task.py @@ -15,19 +15,22 @@ class Task(BaseModel): __hash__ = object.__hash__ # type: ignore i18n: I18N = I18N() description: str = Field(description="Description of the actual task.") + callback: Optional[Any] = Field( + description="Callback to be executed after the task is completed.", default=None + ) agent: Optional[Agent] = Field( - description="Agent responsible for the task.", default=None + description="Agent responsible for executiong the task.", default=None ) - tools: List[Any] = Field( - default_factory=list, - description="Tools the agent is limited to use for this task.", - ) - expected_output: str = Field( + expected_output: Optional[str] = Field( description="Clear definition of expected output for the task.", default=None, ) output: Optional[TaskOutput] = Field( - description="Task output, it's final result.", default=None + description="Task output, it's final result after being executed", default=None + ) + tools: List[Any] = Field( + default_factory=list, + description="Tools the agent is limited to use for this task.", ) id: UUID4 = Field( default_factory=uuid.uuid4, @@ -66,6 +69,7 @@ class Task(BaseModel): ) self.output = TaskOutput(description=self.description, result=result) + self.callback(self.output) if self.callback else None return result def _prompt(self) -> str: diff --git a/tests/task_test.py b/tests/task_test.py index f1793feb1..9b8a872d7 100644 --- a/tests/task_test.py +++ b/tests/task_test.py @@ -75,3 +75,31 @@ def test_task_prompt_includes_expected_output(): execute.return_value = "ok" task.execute() execute.assert_called_once_with(task=task._prompt(), context=None, tools=[]) + + +def test_task_callback(): + from unittest.mock import MagicMock + + researcher = Agent( + role="Researcher", + goal="Make the best research and analysis on content about AI and AI agents", + backstory="You're an expert researcher, specialized in technology, software engineering, AI and startups. You work as a freelancer and is now working on doing research and analysis for a new customer.", + allow_delegation=False, + ) + + task_completed = MagicMock(return_value="done") + + task = Task( + description="Give me a list of 5 interesting ideas to explore for na article, what makes them unique and interesting.", + expected_output="Bullet point list of 5 interesting ideas.", + agent=researcher, + callback=task_completed, + allow_delegation=False, + ) + + from unittest.mock import patch + + with patch.object(Agent, "execute_task") as execute: + execute.return_value = "ok" + task.execute() + task_completed.assert_called_once_with(task.output)