Merge in main to bugfix/kickoff-for-each-usage-metrics

src/crewai/task.py

@@ -3,17 +3,22 @@ import re
 import threading
 import uuid
 from copy import copy
-from typing import Any, Dict, List, Optional, Type
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
 
 from langchain_openai import ChatOpenAI
+from opentelemetry.trace import Span
 from pydantic import UUID4, BaseModel, Field, field_validator, model_validator
 from pydantic_core import PydanticCustomError
 
-from crewai.agent import Agent
+from crewai.agents.agent_builder.base_agent import BaseAgent
 from crewai.tasks.task_output import TaskOutput
-from crewai.utilities import I18N, Converter, ConverterError, Printer
+from crewai.telemetry.telemetry import Telemetry
+from crewai.utilities import I18N, ConverterError, Printer
 from crewai.utilities.pydantic_schema_parser import PydanticSchemaParser
 
+if TYPE_CHECKING:
+    from crewai.agents import Agent
+
 
 class Task(BaseModel):
     """Class that represents a task to be executed.
@@ -42,7 +47,6 @@ class Task(BaseModel):
     tools_errors: int = 0
     delegations: int = 0
     i18n: I18N = I18N()
-    thread: Optional[threading.Thread] = None
     prompt_context: Optional[str] = None
     description: str = Field(description="Description of the actual task.")
     expected_output: str = Field(
@@ -55,7 +59,7 @@ class Task(BaseModel):
     callback: Optional[Any] = Field(
         description="Callback to be executed after the task is completed.", default=None
     )
-    agent: Optional[Agent] = Field(
+    agent: Optional[BaseAgent] = Field(
         description="Agent responsible for execution the task.", default=None
     )
     context: Optional[List["Task"]] = Field(
@@ -95,8 +99,11 @@ class Task(BaseModel):
         default=False,
     )
 
+    _telemetry: Telemetry
+    _execution_span: Span | None = None
     _original_description: str | None = None
     _original_expected_output: str | None = None
+    _thread: threading.Thread | None = None
 
     def __init__(__pydantic_self__, **data):
         config = data.pop("config", {})
@@ -118,6 +125,12 @@ class Task(BaseModel):
             return value[1:]
         return value
 
+    @model_validator(mode="after")
+    def set_private_attrs(self) -> "Task":
+        """Set private attributes."""
+        self._telemetry = Telemetry()
+        return self
+
     @model_validator(mode="after")
     def set_attributes_based_on_config(self) -> "Task":
         """Set attributes based on the agent configuration."""
@@ -145,9 +158,21 @@ class Task(BaseModel):
             )
         return self
 
+    def wait_for_completion(self) -> str | BaseModel:
+        """Wait for asynchronous task completion and return the output."""
+        assert self.async_execution, "Task is not set to be executed asynchronously."
+
+        if self._thread:
+            self._thread.join()
+            self._thread = None
+
+        assert self.output, "Task output is not set."
+
+        return self.output.exported_output
+
     def execute(  # type: ignore # Missing return statement
         self,
-        agent: Agent | None = None,
+        agent: BaseAgent | None = None,
         context: Optional[str] = None,
         tools: Optional[List[Any]] = None,
     ) -> str:
@@ -157,6 +182,8 @@ class Task(BaseModel):
             Output of the task.
         """
 
+        self._execution_span = self._telemetry.task_started(self)
+
         agent = agent or self.agent
         if not agent:
             raise Exception(
@@ -168,8 +195,8 @@ class Task(BaseModel):
             context = []
             for task in self.context:
                 if task.async_execution:
-                    task.thread.join()  # type: ignore # Item "None" of "Thread | None" has no attribute "join"
-                if task and task.output:
+                    task.wait_for_completion()
+                if task.output:
                     # type: ignore # Item "str" of "str | None" has no attribute "append"
                     context.append(task.output.raw_output)
             # type: ignore # Argument 1 to "join" of "str" has incompatible type "str | None"; expected "Iterable[str]"
@@ -179,10 +206,10 @@ class Task(BaseModel):
         tools = tools or self.tools
 
         if self.async_execution:
-            self.thread = threading.Thread(
+            self._thread = threading.Thread(
                 target=self._execute, args=(agent, self, context, tools)
             )
-            self.thread.start()
+            self._thread.start()
         else:
             result = self._execute(
                 task=self,
@@ -192,15 +219,15 @@ class Task(BaseModel):
             )
             return result
 
-    def _execute(self, agent: Agent, task, context, tools):
+    def _execute(self, agent: "Agent", task, context, tools):
         result = agent.execute_task(
             task=task,
             context=context,
             tools=tools,
         )
-
         exported_output = self._export_output(result)
 
+        # type: ignore # the responses are usually str but need to figure out a more elegant solution here
         self.output = TaskOutput(
             description=self.description,
             exported_output=exported_output,
@@ -211,6 +238,10 @@ class Task(BaseModel):
         if self.callback:
             self.callback(self.output)
 
+        if self._execution_span:
+            self._telemetry.task_ended(self._execution_span, self)
+            self._execution_span = None
+
         return exported_output
 
     def prompt(self) -> str:
@@ -246,7 +277,7 @@ class Task(BaseModel):
         """Increment the delegations counter."""
         self.delegations += 1
 
-    def copy(self, agents: Optional[List[Agent]] = None) -> "Task":
+    def copy(self, agents: Optional[List["Agent"]] = None) -> "Task":
         """Create a deep copy of the Task."""
         exclude = {
             "id",
@@ -262,15 +293,11 @@ class Task(BaseModel):
             [task.copy() for task in self.context] if self.context else None
         )
 
-        # TODO: Make sure this clone approach is correct.
-        def get_agent_by_role(role: str) -> Agent | None:
+        def get_agent_by_role(role: str) -> Union["Agent", None]:
             return next((agent for agent in agents if agent.role == role), None)
 
         cloned_agent = get_agent_by_role(self.agent.role) if self.agent else None
-        # cloned_agent = self.agent.copy() if self.agent else None
-        print("TOOLS BEFORE COPY", self.tools)
         cloned_tools = copy(self.tools) if self.tools else []
-        print("TOOLS AFTER COPY", cloned_tools)
 
         copied_task = Task(
             **copied_data,
@@ -311,14 +338,13 @@ class Task(BaseModel):
                         pass
 
             # type: ignore # Item "None" of "Agent | None" has no attribute "function_calling_llm"
-            llm = self.agent.function_calling_llm or self.agent.llm
-
+            llm = getattr(self.agent, "function_calling_llm", None) or self.agent.llm
             if not self._is_gpt(llm):
                 # type: ignore # Argument "model" to "PydanticSchemaParser" has incompatible type "type[BaseModel] | None"; expected "type[BaseModel]"
                 model_schema = PydanticSchemaParser(model=model).get_schema()
                 instructions = f"{instructions}\n\nThe json should have the following structure, with the following keys:\n{model_schema}"
 
-            converter = Converter(
+            converter = self.agent.get_output_converter(
                 llm=llm, text=result, model=model, instructions=instructions
             )