diff --git a/src/crewai/task.py b/src/crewai/task.py index 292f178ff..6d90fcfa6 100644 --- a/src/crewai/task.py +++ b/src/crewai/task.py @@ -14,6 +14,7 @@ from crewai.agents.agent_builder.base_agent import BaseAgent from crewai.tasks.task_output import TaskOutput from crewai.telemetry.telemetry import Telemetry from crewai.utilities.converter import ConverterError +from crewai.utilities.converter import Converter from crewai.utilities.i18n import I18N from crewai.utilities.printer import Printer 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", default=False, ) + converter_cls: Optional[Type[Converter]] = Field( + description="A converter class used to export structured output", + default=None, + ) _telemetry: Telemetry _execution_span: Span | None = None @@ -305,6 +310,16 @@ class Task(BaseModel): 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: exported_result = result 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]" 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 ) diff --git a/tests/task_test.py b/tests/task_test.py index 76dcaa4a9..b5ac6970f 100644 --- a/tests/task_test.py +++ b/tests/task_test.py @@ -9,6 +9,7 @@ from pydantic_core import ValidationError from crewai import Agent, Crew, Process, Task from crewai.tasks.task_output import TaskOutput +from crewai.utilities.converter import Converter 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}') +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"]) def test_increment_delegations_for_hierarchical_process(): from langchain_openai import ChatOpenAI