feat: enhance pydantic output to include field descriptions

- Update generate_model_description to include field descriptions
- Add tests for field description handling
- Maintain backward compatibility for fields without descriptions

Fixes #2188

Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-02-21 11:40:37 +00:00
parent 96a7e8038f
commit 326f406605
2 changed files with 42 additions and 13 deletions

View File

@@ -263,32 +263,41 @@ def generate_model_description(model: Type[BaseModel]) -> str:
models. models.
""" """
def describe_field(field_type): def describe_field(field_type, field_info=None):
origin = get_origin(field_type) origin = get_origin(field_type)
args = get_args(field_type) args = get_args(field_type)
type_desc = ""
if origin is Union or (origin is None and len(args) > 0): if origin is Union or (origin is None and len(args) > 0):
# Handle both Union and the new '|' syntax # Handle both Union and the new '|' syntax
non_none_args = [arg for arg in args if arg is not type(None)] non_none_args = [arg for arg in args if arg is not type(None)]
if len(non_none_args) == 1: if len(non_none_args) == 1:
return f"Optional[{describe_field(non_none_args[0])}]" type_desc = f"Optional[{describe_field(non_none_args[0])}]"
else: else:
return f"Optional[Union[{', '.join(describe_field(arg) for arg in non_none_args)}]]" type_desc = f"Optional[Union[{', '.join(describe_field(arg) for arg in non_none_args)}]]"
elif origin is list: elif origin is list:
return f"List[{describe_field(args[0])}]" type_desc = f"List[{describe_field(args[0])}]"
elif origin is dict: elif origin is dict:
key_type = describe_field(args[0]) key_type = describe_field(args[0])
value_type = describe_field(args[1]) value_type = describe_field(args[1])
return f"Dict[{key_type}, {value_type}]" type_desc = f"Dict[{key_type}, {value_type}]"
elif isinstance(field_type, type) and issubclass(field_type, BaseModel): elif isinstance(field_type, type) and issubclass(field_type, BaseModel):
return generate_model_description(field_type) type_desc = generate_model_description(field_type)
elif hasattr(field_type, "__name__"): elif hasattr(field_type, "__name__"):
return field_type.__name__ type_desc = field_type.__name__
else: else:
return str(field_type) type_desc = str(field_type)
fields = model.__annotations__ if field_info and field_info.description:
field_descriptions = [ return {"type": type_desc, "description": field_info.description}
f'"{name}": {describe_field(type_)}' for name, type_ in fields.items() return type_desc
]
fields = model.model_fields
field_descriptions = []
for name, field in fields.items():
field_desc = describe_field(field.annotation, field)
if isinstance(field_desc, dict):
field_descriptions.append(f'"{name}": {json.dumps(field_desc)}')
else:
field_descriptions.append(f'"{name}": {field_desc}')
return "{\n " + ",\n ".join(field_descriptions) + "\n}" return "{\n " + ",\n ".join(field_descriptions) + "\n}"

View File

@@ -4,7 +4,7 @@ from typing import Dict, List, Optional
from unittest.mock import MagicMock, Mock, patch from unittest.mock import MagicMock, Mock, patch
import pytest import pytest
from pydantic import BaseModel from pydantic import BaseModel, Field
from crewai.llm import LLM from crewai.llm import LLM
from crewai.utilities.converter import ( from crewai.utilities.converter import (
@@ -328,6 +328,26 @@ def test_generate_model_description_dict_field():
assert description == expected_description assert description == expected_description
def test_generate_model_description_with_field_descriptions():
class ModelWithDescriptions(BaseModel):
name: str = Field(..., description="The user's full name")
age: int = Field(..., description="The user's age in years")
description = generate_model_description(ModelWithDescriptions)
expected = '{\n "name": {"type": "str", "description": "The user\'s full name"},\n "age": {"type": "int", "description": "The user\'s age in years"}\n}'
assert description == expected
def test_generate_model_description_mixed_fields():
class MixedModel(BaseModel):
name: str = Field(..., description="The user's name")
age: int # No description
description = generate_model_description(MixedModel)
expected = '{\n "name": {"type": "str", "description": "The user\'s name"},\n "age": int\n}'
assert description == expected
@pytest.mark.vcr(filter_headers=["authorization"]) @pytest.mark.vcr(filter_headers=["authorization"])
def test_convert_with_instructions(): def test_convert_with_instructions():
llm = LLM(model="gpt-4o-mini") llm = LLM(model="gpt-4o-mini")