WIP. Needing team to review change

This commit is contained in:
Brandon Hancock
2024-07-03 11:09:19 -04:00
parent e745094d73
commit 55af7e0f15
3 changed files with 22 additions and 7 deletions

View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, List, Optional, Union
from typing import Any, Dict, List, Union
from pydantic import BaseModel, Field
@@ -16,20 +16,28 @@ class CrewOutput(BaseModel):
)
# TODO: Ask @joao what is the desired behavior here
def result(self) -> Union[str, BaseModel, Dict[str, Any]]:
def result(
self,
) -> List[str | BaseModel | Dict[str, Any]]]:
"""Return the result of the task based on the available output."""
results = [output.result() for output in self.output]
return results if len(results) > 1 else results[0]
return results
def raw_output(self) -> str:
"""Return the raw output of the task."""
return aggregate_raw_outputs_from_task_outputs(self.output)
def to_output_dict(self) -> Dict[str, Any]:
self.output.to_output_dict()
def to_output_dict(self) -> List[Dict[str, Any]]:
output_dict = [output.to_output_dict() for output in self.output]
return output_dict
def __getitem__(self, key: str) -> Any:
self.output[key]
if len(self.output) == 0:
return None
elif len(self.output) == 1:
return self.output[0][key]
else:
return [output[key] for output in self.output]
# TODO: Confirm with Joao that we want to print the raw output and not the object
def __str__(self):

View File

@@ -49,7 +49,7 @@ class TaskOutput(BaseModel):
if self.json_output:
output_dict.update(self.json_output)
if self.pydantic_output:
output_dict.update(self.pydantic_output.dict())
output_dict.update(self.pydantic_output.model_dump())
return output_dict
def __str__(self) -> str:

View File

@@ -7,6 +7,7 @@ from unittest.mock import patch
import pydantic_core
import pytest
from crewai.agent import Agent
from crewai.agents.cache import CacheHandler
from crewai.crew import Crew
@@ -146,6 +147,7 @@ def test_crew_creation():
assert result.raw_output() == expected_string_output
assert isinstance(result, CrewOutput)
assert len(result.tasks_output) == len(tasks)
assert result.result() == [expected_string_output]
@pytest.mark.vcr(filter_headers=["authorization"])
@@ -1736,3 +1738,8 @@ def test__setup_for_training():
for agent in agents:
assert agent.allow_delegation is False
# TODO: TEST EXPORT OUTPUT TASK WITH PYDANTIC
# TODO: TEST EXPORT OUTPUT TASK WITH JSON
# TODO: TEST EXPORT OUTPUT TASK CALLBACK