diff --git a/src/crewai/project/annotations.py b/src/crewai/project/annotations.py index 3327296a4..e402b8cfe 100644 --- a/src/crewai/project/annotations.py +++ b/src/crewai/project/annotations.py @@ -1,14 +1,28 @@ tasks_order = [] +def memoize(func): + cache = {} + + def memoized_func(*args, **kwargs): + key = (args, tuple(kwargs.items())) + if key not in cache: + cache[key] = func(*args, **kwargs) + return cache[key] + + return memoized_func + + def task(func): func.is_task = True tasks_order.append(func.__name__) + func = memoize(func) return func def agent(func): func.is_agent = True + func = memoize(func) return func diff --git a/tests/project_test.py b/tests/project_test.py new file mode 100644 index 000000000..d487bdb33 --- /dev/null +++ b/tests/project_test.py @@ -0,0 +1,35 @@ +from crewai.agent import Agent +from crewai.project import agent, task +from crewai.task import Task + + +class SimpleCrew: + @agent + def simple_agent(self): + return Agent( + role="Simple Agent", goal="Simple Goal", backstory="Simple Backstory" + ) + + @task + def simple_task(self): + return Task(description="Simple Description", expected_output="Simple Output") + + +def test_agent_memoization(): + crew = SimpleCrew() + first_call_result = crew.simple_agent() + second_call_result = crew.simple_agent() + + assert ( + first_call_result is second_call_result + ), "Agent memoization is not working as expected" + + +def test_task_memoization(): + crew = SimpleCrew() + first_call_result = crew.simple_task() + second_call_result = crew.simple_task() + + assert ( + first_call_result is second_call_result + ), "Task memoization is not working as expected"