adding task callback

This commit is contained in:
João Moura
2024-01-30 22:44:37 -03:00
parent 2010702880
commit 8fc0f33dd5
2 changed files with 39 additions and 7 deletions

View File

@@ -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)