Implement major fixes from yesterdays group conversation. Now working on tests.

This commit is contained in:
Brandon Hancock
2024-07-09 10:27:39 -04:00
parent fffe4df8c3
commit ecc3d913da
7 changed files with 144 additions and 130 deletions

View File

@@ -1,24 +1,26 @@
from typing import Any, Dict, Optional, Union
from typing import Any, Dict, Optional
from pydantic import BaseModel, Field, model_validator
from crewai.tasks.output_format import OutputFormat
# TODO: This is a breaking change. Confirm with @joao
class TaskOutput(BaseModel):
"""Class that represents the result of a task."""
# TODO: MAKE SURE TO FULLY SUPPORT OUTPUT FILE
description: str = Field(description="Description of the task")
summary: Optional[str] = Field(description="Summary of the task", default=None)
raw_output: str = Field(description="Result of the task")
pydantic_output: Optional[BaseModel] = Field(
raw: str = Field(
description="Result of the task"
) # TODO: @joao: breaking change, by renaming raw_output to raw, but now consistent with CrewOutput
pydantic: Optional[BaseModel] = Field(
description="Pydantic model output", default=None
)
json_output: Optional[Dict[str, Any]] = Field(
description="JSON output", default=None
)
json: Optional[Dict[str, Any]] = Field(description="JSON output", default=None)
agent: str = Field(description="Agent that executed the task")
output_format: OutputFormat = Field(
description="Output format of the task", default=OutputFormat.RAW
)
@model_validator(mode="after")
def set_summary(self):
@@ -27,32 +29,18 @@ class TaskOutput(BaseModel):
self.summary = f"{excerpt}..."
return self
# TODO: Ask @joao what is the desired behavior here
def result(self) -> Union[str, BaseModel, Dict[str, Any]]:
"""Return the result of the task based on the available output."""
if self.pydantic_output:
return self.pydantic_output
elif self.json_output:
return self.json_output
else:
return self.raw_output
def __getitem__(self, key: str) -> Any:
"""Retrieve a value from the pydantic_output or json_output based on the key."""
if self.pydantic_output and hasattr(self.pydantic_output, key):
return getattr(self.pydantic_output, key)
if self.json_output and key in self.json_output:
return self.json_output[key]
raise KeyError(f"Key '{key}' not found in pydantic_output or json_output")
def to_dict(self) -> Dict[str, Any]:
"""Convert json_output and pydantic_output to a dictionary."""
output_dict = {}
if self.json_output:
output_dict.update(self.json_output)
if self.pydantic_output:
output_dict.update(self.pydantic_output.model_dump())
if self.json:
output_dict.update(self.json)
if self.pydantic:
output_dict.update(self.pydantic.model_dump())
return output_dict
def __str__(self) -> str:
return self.raw_output
if self.pydantic:
return str(self.pydantic)
if self.json:
return str(self.json)
return self.raw