fix: hash callback args correctly to ensure caching works

Author: greysonlalonde

lib/crewai/src/crewai/project/utils.py

@@ -1,21 +1,75 @@
 """Utility functions for the crewai project module."""
 
 from collections.abc import Callable
-from functools import lru_cache
-from typing import ParamSpec, TypeVar, cast
+from functools import wraps
+from typing import Any, ParamSpec, TypeVar, cast
+
+from pydantic import BaseModel
+
+from crewai.agents.cache.cache_handler import CacheHandler
 
 
 P = ParamSpec("P")
 R = TypeVar("R")
+cache = CacheHandler()
+
+
+def _make_hashable(arg: Any) -> Any:
+    """Convert argument to hashable form for caching.
+
+    Args:
+        arg: The argument to convert.
+
+    Returns:
+        Hashable representation of the argument.
+    """
+    if isinstance(arg, BaseModel):
+        return arg.model_dump_json()
+    if isinstance(arg, dict):
+        return tuple(sorted((k, _make_hashable(v)) for k, v in arg.items()))
+    if isinstance(arg, list):
+        return tuple(_make_hashable(item) for item in arg)
+    if hasattr(arg, "__dict__"):
+        return ("__instance__", id(arg))
+    return arg
 
 
 def memoize(meth: Callable[P, R]) -> Callable[P, R]:
     """Memoize a method by caching its results based on arguments.
 
+    Handles Pydantic BaseModel instances by converting them to JSON strings
+    before hashing for cache lookup.
+
     Args:
         meth: The method to memoize.
 
     Returns:
         A memoized version of the method that caches results.
     """
-    return cast(Callable[P, R], lru_cache(typed=True)(meth))
+
+    @wraps(meth)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        """Wrapper that converts arguments to hashable form before caching.
+
+        Args:
+            *args: Positional arguments to the memoized method.
+            **kwargs: Keyword arguments to the memoized method.
+
+        Returns:
+            The result of the memoized method call.
+        """
+        hashable_args = tuple(_make_hashable(arg) for arg in args)
+        hashable_kwargs = tuple(
+            sorted((k, _make_hashable(v)) for k, v in kwargs.items())
+        )
+        cache_key = str((hashable_args, hashable_kwargs))
+
+        cached_result: R | None = cache.read(tool=meth.__name__, input=cache_key)
+        if cached_result is not None:
+            return cached_result
+
+        result = meth(*args, **kwargs)
+        cache.add(tool=meth.__name__, input=cache_key, output=result)
+        return result
+
+    return cast(Callable[P, R], wrapper)

lib/crewai/tests/experimental/evaluation/test_agent_evaluator.py

@@ -62,18 +62,23 @@ def test_evaluate_current_iteration(self, mock_crew):
             agents=mock_crew.agents, evaluators=[GoalAlignmentEvaluator()]
         )
 
-        task_completed_event = threading.Event()
+        task_completed_condition = threading.Condition()
+        task_completed = False
 
         @crewai_event_bus.on(TaskCompletedEvent)
         async def on_task_completed(source, event):
             # TaskCompletedEvent fires AFTER evaluation results are stored
-            task_completed_event.set()
+            nonlocal task_completed
+            with task_completed_condition:
+                task_completed = True
+                task_completed_condition.notify()
 
         mock_crew.kickoff()
 
-        assert task_completed_event.wait(timeout=5), (
-            "Timeout waiting for task completion"
-        )
+        with task_completed_condition:
+            assert task_completed_condition.wait_for(
+                lambda: task_completed, timeout=5
+            ), "Timeout waiting for task completion"
 
         results = agent_evaluator.get_evaluation_results()
 

lib/crewai/tests/project/test_callback_with_taskoutput.py

@@ -0,0 +1,94 @@
+"""Test callback decorator with TaskOutput arguments."""
+
+from unittest.mock import MagicMock, patch
+
+from crewai import Agent, Crew, Task
+from crewai.project import CrewBase, callback, task
+from crewai.tasks.output_format import OutputFormat
+from crewai.tasks.task_output import TaskOutput
+
+
+def test_callback_decorator_with_taskoutput() -> None:
+    """Test that @callback decorator works with TaskOutput arguments."""
+
+    @CrewBase
+    class TestCrew:
+        """Test crew with callback."""
+
+        callback_called = False
+        callback_output = None
+
+        @callback
+        def task_callback(self, output: TaskOutput) -> None:
+            """Test callback that receives TaskOutput."""
+            self.callback_called = True
+            self.callback_output = output
+
+        @task
+        def test_task(self) -> Task:
+            """Test task with callback."""
+            return Task(
+                description="Test task",
+                expected_output="Test output",
+                callback=self.task_callback,
+            )
+
+    test_crew = TestCrew()
+    task_instance = test_crew.test_task()
+
+    test_output = TaskOutput(
+        description="Test task",
+        agent="Test Agent",
+        raw="test result",
+        output_format=OutputFormat.RAW,
+    )
+
+    task_instance.callback(test_output)
+
+    assert test_crew.callback_called
+    assert test_crew.callback_output == test_output
+
+
+def test_callback_decorator_with_taskoutput_integration() -> None:
+    """Integration test for callback with actual task execution."""
+
+    @CrewBase
+    class TestCrew:
+        """Test crew with callback integration."""
+
+        callback_called = False
+        received_output: TaskOutput | None = None
+
+        @callback
+        def task_callback(self, output: TaskOutput) -> None:
+            """Callback executed after task completion."""
+            self.callback_called = True
+            self.received_output = output
+
+        @task
+        def test_task(self) -> Task:
+            """Test task."""
+            return Task(
+                description="Test task",
+                expected_output="Test output",
+                callback=self.task_callback,
+            )
+
+    test_crew = TestCrew()
+
+    agent = Agent(
+        role="Test Agent",
+        goal="Test goal",
+        backstory="Test backstory",
+    )
+
+    task_instance = test_crew.test_task()
+    task_instance.agent = agent
+
+    with patch.object(Agent, "execute_task") as mock_execute:
+        mock_execute.return_value = "test result"
+        task_instance.execute_sync()
+
+        assert test_crew.callback_called
+        assert test_crew.received_output is not None
+        assert test_crew.received_output.raw == "test result"