fix: fixes pydantic parser to support nested objects.

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
    }
}
This commit is contained in:
Jason Vertrees
2024-12-09 15:41:19 -06:00
parent da4f356fab
commit af0ed9ad53
2 changed files with 122 additions and 32 deletions

View File

@@ -0,0 +1,73 @@
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()