Fix/async tasks (#877)

* fix: async tasks calls

* fix: some issue along with some type check errors

* fix: some issue along with some type check errors

* fix: async test
This commit is contained in:
Eduardo Chiarotti
2024-07-06 01:30:07 -03:00
committed by GitHub
parent bb64c80964
commit a41bd18599
3 changed files with 50 additions and 46 deletions

View File

@@ -1,7 +1,7 @@
import os
import re
import threading
import uuid
from concurrent.futures import Future, ThreadPoolExecutor
from copy import copy
from typing import Any, Dict, List, Optional, Type, Union
@@ -102,7 +102,7 @@ class Task(BaseModel):
_execution_span: Span | None = None
_original_description: str | None = None
_original_expected_output: str | None = None
_thread: threading.Thread | None = None
_future: Future | None = None
def __init__(__pydantic_self__, **data):
config = data.pop("config", {})
@@ -161,20 +161,20 @@ class Task(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
if self._future:
self._future.result() # Wait for the future to complete
self._future = None
assert self.output, "Task output is not set."
return self.output.exported_output
def execute( # type: ignore # Missing return statement
def execute(
self,
agent: BaseAgent | None = None,
context: Optional[str] = None,
tools: Optional[List[Any]] = None,
) -> str:
) -> str | None:
"""Execute the task.
Returns:
@@ -186,26 +186,28 @@ class Task(BaseModel):
agent = agent or self.agent
if not agent:
raise Exception(
f"The task '{self.description}' has no agent assigned, therefore it can't be executed directly and should be executed in a Crew using a specific process that support that, like hierarchical."
f"The task '{self.description}' has no agent assigned, therefore it can't be executed directly "
"and should be executed in a Crew using a specific process that support that, like hierarchical."
)
if self.context:
context = [] # type: ignore # Incompatible types in assignment (expression has type "list[Never]", variable has type "str | None")
internal_context = []
for task in self.context:
if task.async_execution:
task.wait_for_completion()
if task.output:
context.append(task.output.raw_output) # type: ignore # Item "str" of "str | None" has no attribute "append"
context = "\n".join(context) # type: ignore # Argument 1 to "join" of "str" has incompatible type "str | None"; expected "Iterable[str]"
internal_context.append(task.output.raw_output)
context = "\n".join(internal_context)
self.prompt_context = context
tools = tools or self.tools
if self.async_execution:
self._thread = threading.Thread(
target=self._execute, args=(agent, self, context, tools)
)
self._thread.start()
with ThreadPoolExecutor() as executor:
self._future = executor.submit(
self._execute, agent, self, context, tools
)
return None
else:
result = self._execute(
task=self,
@@ -215,7 +217,7 @@ class Task(BaseModel):
)
return result
def _execute(self, agent: "BaseAgent", task, context, tools):
def _execute(self, agent: "BaseAgent", task, context, tools) -> str | None:
result = agent.execute_task(
task=task,
context=context,
@@ -223,7 +225,7 @@ class Task(BaseModel):
)
exported_output = self._export_output(result)
self.output = TaskOutput( # 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,
raw_output=result,