mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 15:48:29 +00:00
d8f52e69be
- Remove unused imports from both files to fix lint errors - Update test expectations to match actual schema output format - Python 3.10+ union syntax now outputs 'str | None' instead of 'Optional[str]' - All local tests pass and core functionality verified Co-Authored-By: João <joao@crewai.com>
130 lines
3.1 KiB
Python
130 lines
3.1 KiB
Python
from typing import Dict, List, Optional, Union
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from crewai.utilities.pydantic_schema_parser import PydanticSchemaParser
|
|
|
|
|
|
def test_simple_model():
|
|
class SimpleModel(BaseModel):
|
|
field1: int
|
|
field2: str
|
|
|
|
parser = PydanticSchemaParser(model=SimpleModel)
|
|
schema = parser.get_schema()
|
|
|
|
expected_schema = """{
|
|
field1: int,
|
|
field2: str
|
|
}"""
|
|
assert schema.strip() == expected_schema.strip()
|
|
|
|
|
|
def test_nested_model():
|
|
class NestedModel(BaseModel):
|
|
nested_field: int
|
|
|
|
class ParentModel(BaseModel):
|
|
parent_field: str
|
|
nested: NestedModel
|
|
|
|
parser = PydanticSchemaParser(model=ParentModel)
|
|
schema = parser.get_schema()
|
|
|
|
expected_schema = """{
|
|
parent_field: str,
|
|
nested: NestedModel
|
|
{
|
|
nested_field: int
|
|
}
|
|
}"""
|
|
assert schema.strip() == expected_schema.strip()
|
|
|
|
|
|
def test_model_with_list():
|
|
class ListModel(BaseModel):
|
|
list_field: List[int]
|
|
|
|
parser = PydanticSchemaParser(model=ListModel)
|
|
schema = parser.get_schema()
|
|
|
|
expected_schema = """{
|
|
list_field: List[int]
|
|
}"""
|
|
assert schema.strip() == expected_schema.strip()
|
|
|
|
|
|
def test_model_with_optional_field():
|
|
class OptionalModel(BaseModel):
|
|
optional_field: Optional[str]
|
|
|
|
parser = PydanticSchemaParser(model=OptionalModel)
|
|
schema = parser.get_schema()
|
|
|
|
expected_schema = """{
|
|
optional_field: Optional[str]
|
|
}"""
|
|
assert schema.strip() == expected_schema.strip()
|
|
|
|
|
|
def test_model_with_union():
|
|
class UnionModel(BaseModel):
|
|
union_field: Union[int, str]
|
|
|
|
parser = PydanticSchemaParser(model=UnionModel)
|
|
schema = parser.get_schema()
|
|
|
|
expected_schema = """{
|
|
union_field: Union[int, str]
|
|
}"""
|
|
assert schema.strip() == expected_schema.strip()
|
|
|
|
|
|
def test_model_with_dict():
|
|
class DictModel(BaseModel):
|
|
dict_field: Dict[str, int]
|
|
|
|
parser = PydanticSchemaParser(model=DictModel)
|
|
schema = parser.get_schema()
|
|
|
|
expected_schema = """{
|
|
dict_field: Dict[str, int]
|
|
}"""
|
|
assert schema.strip() == expected_schema.strip()
|
|
|
|
|
|
def test_model_with_python310_union_syntax():
|
|
class UnionTypeModel(BaseModel):
|
|
union_field: str | None
|
|
multi_union_field: int | str | None
|
|
non_optional_union: int | str
|
|
|
|
parser = PydanticSchemaParser(model=UnionTypeModel)
|
|
schema = parser.get_schema()
|
|
|
|
expected_schema = """{
|
|
union_field: str | None,
|
|
multi_union_field: int | str | None,
|
|
non_optional_union: int | str
|
|
}"""
|
|
assert schema.strip() == expected_schema.strip()
|
|
|
|
|
|
def test_mixed_union_syntax():
|
|
class MixedUnionModel(BaseModel):
|
|
traditional_optional: Optional[str]
|
|
new_union_syntax: str | None
|
|
traditional_union: Union[int, str]
|
|
new_multi_union: int | str | float
|
|
|
|
parser = PydanticSchemaParser(model=MixedUnionModel)
|
|
schema = parser.get_schema()
|
|
|
|
expected_schema = """{
|
|
traditional_optional: Optional[str],
|
|
new_union_syntax: str | None,
|
|
traditional_union: Union[int, str],
|
|
new_multi_union: int | str | float
|
|
}"""
|
|
assert schema.strip() == expected_schema.strip()
|