Add converter_cls option to Task (#800)

* Add converter_cls option to Task

Fixes #799

* Update task_test.py

* Update task.py

* Update task.py

* Update task_test.py

* Update task.py

* Update task.py

* Update task.py

* Update task.py

---------

Co-authored-by: João Moura <joaomdmoura@gmail.com>
This commit is contained in:
Eelke van den Bos
2024-07-06 01:01:39 -04:00
committed by GitHub
parent 58558a1950
commit 7edacf6e24
2 changed files with 48 additions and 1 deletions

View File

@@ -14,6 +14,7 @@ from crewai.agents.agent_builder.base_agent import BaseAgent
from crewai.tasks.task_output import TaskOutput from crewai.tasks.task_output import TaskOutput
from crewai.telemetry.telemetry import Telemetry from crewai.telemetry.telemetry import Telemetry
from crewai.utilities.converter import ConverterError from crewai.utilities.converter import ConverterError
from crewai.utilities.converter import Converter
from crewai.utilities.i18n import I18N from crewai.utilities.i18n import I18N
from crewai.utilities.printer import Printer from crewai.utilities.printer import Printer
from crewai.utilities.pydantic_schema_parser import PydanticSchemaParser from crewai.utilities.pydantic_schema_parser import PydanticSchemaParser
@@ -97,6 +98,10 @@ class Task(BaseModel):
description="Whether the task should have a human review the final answer of the agent", description="Whether the task should have a human review the final answer of the agent",
default=False, default=False,
) )
converter_cls: Optional[Type[Converter]] = Field(
description="A converter class used to export structured output",
default=None,
)
_telemetry: Telemetry _telemetry: Telemetry
_execution_span: Span | None = None _execution_span: Span | None = None
@@ -305,6 +310,16 @@ class Task(BaseModel):
return copied_task return copied_task
def _create_converter(self, *args, **kwargs) -> Converter: # type: ignore
converter = self.agent.get_output_converter( # type: ignore # Item "None" of "BaseAgent | None" has no attribute "get_output_converter"
*args, **kwargs
)
if self.converter_cls:
converter = self.converter_cls( # type: ignore # Item "None" of "BaseAgent | None" has no attribute "get_output_converter"
*args, **kwargs
)
return converter
def _export_output(self, result: str) -> Any: def _export_output(self, result: str) -> Any:
exported_result = result exported_result = result
instructions = "I'm gonna convert this raw text into valid JSON." instructions = "I'm gonna convert this raw text into valid JSON."
@@ -335,7 +350,7 @@ class Task(BaseModel):
model_schema = PydanticSchemaParser(model=model).get_schema() # type: ignore # Argument "model" to "PydanticSchemaParser" has incompatible type "type[BaseModel] | None"; expected "type[BaseModel]" model_schema = PydanticSchemaParser(model=model).get_schema() # type: ignore # Argument "model" to "PydanticSchemaParser" has incompatible type "type[BaseModel] | None"; expected "type[BaseModel]"
instructions = f"{instructions}\n\nThe json should have the following structure, with the following keys:\n{model_schema}" instructions = f"{instructions}\n\nThe json should have the following structure, with the following keys:\n{model_schema}"
converter = self.agent.get_output_converter( # type: ignore # Item "None" of "BaseAgent | None" has no attribute "get_output_converter" converter = self._create_converter( # type: ignore # Item "None" of "BaseAgent | None" has no attribute "get_output_converter"
llm=llm, text=result, model=model, instructions=instructions llm=llm, text=result, model=model, instructions=instructions
) )

View File

@@ -9,6 +9,7 @@ from pydantic_core import ValidationError
from crewai import Agent, Crew, Process, Task from crewai import Agent, Crew, Process, Task
from crewai.tasks.task_output import TaskOutput from crewai.tasks.task_output import TaskOutput
from crewai.utilities.converter import Converter
def test_task_tool_reflect_agent_tools(): def test_task_tool_reflect_agent_tools():
@@ -393,6 +394,37 @@ def test_save_task_pydantic_output():
save_file.assert_called_once_with('{"score":4}') save_file.assert_called_once_with('{"score":4}')
def test_custom_converter_cls():
class ScoreOutput(BaseModel):
score: int
class ScoreConverter(Converter):
pass
scorer = Agent(
role="Scorer",
goal="Score the title",
backstory="You're an expert scorer, specialized in scoring titles.",
allow_delegation=False,
)
task = Task(
description="Give me an integer score between 1-5 for the following title: 'The impact of AI in the future of work'",
expected_output="The score of the title.",
output_file="score.json",
output_pydantic=ScoreOutput,
converter_cls=ScoreConverter,
agent=scorer,
)
crew = Crew(agents=[scorer], tasks=[task])
with patch.object(ScoreConverter, "__new__", ScoreConverter.__new__) as converter_constructor:
crew.kickoff()
converter_constructor.assert_called_once
@pytest.mark.vcr(filter_headers=["authorization"]) @pytest.mark.vcr(filter_headers=["authorization"])
def test_increment_delegations_for_hierarchical_process(): def test_increment_delegations_for_hierarchical_process():
from langchain_openai import ChatOpenAI from langchain_openai import ChatOpenAI