mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-09 08:08:32 +00:00
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:
@@ -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}"
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
Reference in New Issue
Block a user