mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-05-05 09:12:39 +00:00
af0ed9ad53
This fails before the PR and succeeds afterward:
```python
from pydantic import BaseModel
from typing import List, Optional
from crewai.utilities.pydantic_schema_parser import PydanticSchemaParser
class InnerModel(BaseModel):
inner_field: int
class TestModel(BaseModel):
simple_field: str
list_field: List[int]
optional_field: Optional[str]
nested_model: InnerModel
print(PydanticSchemaParser(model=InnerModel).get_schema()) # works
print(PydanticSchemaParser(model=TestModel).get_schema()) # fails
```
**Note.** Because the `main` branch currently doesn't support nested
schemas and the original code against which I made this PR months ago
drifted, I made a judgement call on how to format the nested
structurees. I chose the following, inferred from the current code:
````
>>> print(PydanticSchemaParser(model=InnerModel).get_schema())
{
inner_field: int
}
>>> print(PydanticSchemaParser(model=TestModel).get_schema())
{
simple_field: str,
list_field: List[int],
optional_field: Optional[str],
nested_model: InnerModel
{
inner_field: int
}
}
74 lines
1.7 KiB
Python
74 lines
1.7 KiB
Python
import unittest
|
|
from pydantic import BaseModel
|
|
from typing import List, Optional
|
|
|
|
from crewai.utilities.pydantic_schema_parser import PydanticSchemaParser
|
|
|
|
|
|
# Define test models
|
|
class InnerModel(BaseModel):
|
|
inner_field: int
|
|
|
|
|
|
class OuterModel(BaseModel):
|
|
simple_field: str
|
|
list_field: List[int]
|
|
optional_field: Optional[str]
|
|
nested_model: InnerModel
|
|
|
|
|
|
# Test cases
|
|
class TestPydanticSchemaParser(unittest.TestCase):
|
|
def test_simple_model(self):
|
|
class SimpleModel(BaseModel):
|
|
field1: int
|
|
field2: str
|
|
|
|
parser = PydanticSchemaParser(model=SimpleModel)
|
|
expected_schema = """
|
|
{
|
|
field1: int,
|
|
field2: str
|
|
}""".strip()
|
|
self.assertEqual(parser.get_schema(), expected_schema)
|
|
|
|
def test_model_with_list(self):
|
|
class ListModel(BaseModel):
|
|
field1: List[int]
|
|
|
|
parser = PydanticSchemaParser(model=ListModel)
|
|
expected_schema = """
|
|
{
|
|
field1: List[int]
|
|
}""".strip()
|
|
self.assertEqual(parser.get_schema(), expected_schema)
|
|
|
|
def test_model_with_optional(self):
|
|
class OptionalModel(BaseModel):
|
|
field1: Optional[int]
|
|
|
|
parser = PydanticSchemaParser(model=OptionalModel)
|
|
expected_schema = """
|
|
{
|
|
field1: Optional[int]
|
|
}""".strip()
|
|
self.assertEqual(parser.get_schema(), expected_schema)
|
|
|
|
def test_nested_model(self):
|
|
parser = PydanticSchemaParser(model=OuterModel)
|
|
expected_schema = """
|
|
{
|
|
simple_field: str,
|
|
list_field: List[int],
|
|
optional_field: Optional[str],
|
|
nested_model: InnerModel
|
|
{
|
|
inner_field: int
|
|
}
|
|
}""".strip()
|
|
self.assertEqual(parser.get_schema(), expected_schema)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|