From b579cb2cce135d5f6e89c0c9d669e181a0e4bbb6 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Tue, 25 Feb 2025 13:36:50 +0000
Subject: [PATCH] Address code review comments and fix lint error

Co-Authored-By: Joe Moura <joao@crewai.com>
---
 src/crewai/agents/agent_builder/base_agent.py | 10 +++++++
 src/crewai/project/annotations.py             | 29 ++++++++++++++++---
 src/crewai/task.py                            | 14 +++++++--
 tests/test_crewbase_agent_method.py           | 19 ++++++++++++
 4 files changed, 66 insertions(+), 6 deletions(-)

diff --git a/src/crewai/agents/agent_builder/base_agent.py b/src/crewai/agents/agent_builder/base_agent.py
index c0d8041c1..6f90c69b3 100644
--- a/src/crewai/agents/agent_builder/base_agent.py
+++ b/src/crewai/agents/agent_builder/base_agent.py
@@ -134,6 +134,16 @@ class BaseAgent(ABC, BaseModel):
     @model_validator(mode="before")
     @classmethod
     def process_model_config(cls, values):
+        """
+        Process model configuration values.
+        
+        Args:
+            values: Configuration values or callable agent
+                When using CrewBase decorator, this can be a callable that returns an agent
+        
+        Returns:
+            Processed configuration or callable agent
+        """
         # Handle case where values is a function (can happen with CrewBase decorator)
         if callable(values) and not isinstance(values, dict):
             return values
diff --git a/src/crewai/project/annotations.py b/src/crewai/project/annotations.py
index 86a77bd63..d8981f799 100644
--- a/src/crewai/project/annotations.py
+++ b/src/crewai/project/annotations.py
@@ -65,6 +65,27 @@ def cache_handler(func):
     return memoize(func)
 
 
+def _resolve_agent(task_instance):
+    """
+    Resolve an agent from a task instance.
+    
+    If the agent is a callable (e.g., a method from CrewBase), call it to get the agent instance.
+    
+    Args:
+        task_instance: The task instance containing the agent
+    
+    Returns:
+        The resolved agent instance or None if no agent is present
+    """
+    if not hasattr(task_instance, 'agent') or not task_instance.agent:
+        return None
+        
+    if callable(task_instance.agent) and not isinstance(task_instance.agent, type):
+        return task_instance.agent()
+        
+    return task_instance.agent
+
+
 def crew(func) -> Callable[..., Crew]:
 
     @wraps(func)
@@ -82,10 +103,10 @@ def crew(func) -> Callable[..., Crew]:
             # Get the task instance
             task_instance = task_method(self)
             
-            # Handle case where agent is a method (function) from CrewBase
-            if hasattr(task_instance, 'agent') and task_instance.agent and callable(task_instance.agent) and not isinstance(task_instance.agent, type):
-                # Call the agent method to get the agent instance
-                task_instance.agent = task_instance.agent()
+            # Resolve the agent
+            agent = _resolve_agent(task_instance)
+            if agent:
+                task_instance.agent = agent
                 
             instantiated_tasks.append(task_instance)
             agent_instance = getattr(task_instance, "agent", None)
diff --git a/src/crewai/task.py b/src/crewai/task.py
index f20270ff0..f31dd7380 100644
--- a/src/crewai/task.py
+++ b/src/crewai/task.py
@@ -65,8 +65,18 @@ class Task(BaseModel):
     def __init__(self, **data):
         # Handle case where agent is a callable (can happen with CrewBase decorator)
         if 'agent' in data and callable(data['agent']) and not isinstance(data['agent'], type):
-            # Call the agent method to get the agent instance
-            data['agent'] = data['agent']()
+            try:
+                # Call the agent method to get the agent instance
+                agent = data['agent']()
+                
+                # Verify that the agent is a valid instance
+                from crewai.agents.agent_builder.base_agent import BaseAgent
+                if agent is not None and not isinstance(agent, BaseAgent):
+                    raise ValueError(f"Expected BaseAgent instance, got {type(agent)}")
+                    
+                data['agent'] = agent
+            except Exception as e:
+                raise ValueError(f"Failed to initialize agent from callable: {e}")
         
         # Call the parent class __init__ method
         super().__init__(**data)
diff --git a/tests/test_crewbase_agent_method.py b/tests/test_crewbase_agent_method.py
index cfb14c523..a2da0a92e 100644
--- a/tests/test_crewbase_agent_method.py
+++ b/tests/test_crewbase_agent_method.py
@@ -1,4 +1,5 @@
 import unittest
+
 from crewai import Agent, Task
 
 
@@ -30,3 +31,21 @@ class TestTaskInitFix(unittest.TestCase):
         self.assertIsInstance(task.agent, Agent)
         self.assertEqual(task.agent.role, "Test Agent")
         self.assertIs(task.agent, agent_instance)
+        
+    def test_task_init_handles_invalid_callable_agent(self):
+        """Test that the Task.__init__ method correctly handles invalid callable agents."""
+        
+        # Create a callable that returns an invalid agent (not an Agent instance)
+        def invalid_callable_agent():
+            return "Not an agent"
+        
+        # Create a task with the invalid callable agent
+        with self.assertRaises(ValueError) as context:
+            task = Task(
+                description="Test Task",
+                expected_output="Test Output",
+                agent=invalid_callable_agent
+            )
+        
+        # Verify that the error message is correct
+        self.assertIn("Expected BaseAgent instance", str(context.exception))