adding meomization to crewai project annotations

This commit is contained in:
João Moura
2024-05-03 00:49:25 -03:00
parent 9b25b4c5bf
commit b0e38c2a2d
2 changed files with 49 additions and 0 deletions

View File

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

35
tests/project_test.py Normal file
View File

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