Skip to content

Commit

Permalink
Add converter_cls option to Task (#800)
Browse files Browse the repository at this point in the history
* 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>
  • Loading branch information
eelkevdbos and joaomdmoura authored Jul 6, 2024
1 parent 58558a1 commit 7edacf6
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 1 deletion.
17 changes: 16 additions & 1 deletion src/crewai/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
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
Expand Down Expand Up @@ -97,6 +98,10 @@ class Config:
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
Expand Down Expand Up @@ -305,6 +310,16 @@ def get_agent_by_role(role: str) -> Union["BaseAgent", None]:

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."
Expand Down Expand Up @@ -335,7 +350,7 @@ def _export_output(self, result: str) -> Any:
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
)

Expand Down
32 changes: 32 additions & 0 deletions tests/task_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@

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():
Expand Down Expand Up @@ -393,6 +394,37 @@ class ScoreOutput(BaseModel):
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
Expand Down

0 comments on commit 7edacf6

Please sign in to comment.