Skip to content

Commit

Permalink
adding meomization to crewai project annotations
Browse files Browse the repository at this point in the history
  • Loading branch information
joaomdmoura committed May 3, 2024
1 parent 2edc88e commit b264eba
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 0 deletions.
14 changes: 14 additions & 0 deletions src/crewai/project/annotations.py
Original file line number Diff line number Diff line change
@@ -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


Expand Down
35 changes: 35 additions & 0 deletions tests/project_test.py
Original file line number Diff line number Diff line change
@@ -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"

0 comments on commit b264eba

Please sign in to comment.