hadnling pydantic obejct with Optional fields

This commit is contained in:
João Moura
2024-10-01 14:18:11 -07:00
parent ba55160d6b
commit 71c5972fc7
2 changed files with 12 additions and 3 deletions

View File

@@ -103,10 +103,12 @@ def convert_to_model(
return handle_partial_json(
result, model, bool(output_json), agent, converter_cls
)
except ValidationError:
return handle_partial_json(
result, model, bool(output_json), agent, converter_cls
)
except Exception as e:
Printer().print(
content=f"Unexpected error during model conversion: {type(e).__name__}: {e}. Returning original result.",

View File

@@ -1,4 +1,4 @@
from typing import Type, get_args, get_origin
from typing import Type, get_args, get_origin, Union
from pydantic import BaseModel
@@ -36,7 +36,14 @@ class PydanticSchemaParser(BaseModel):
return f"List[\n{nested_schema}\n{' ' * 4 * depth}]"
else:
return f"List[{list_item_type.__name__}]"
elif issubclass(field_type, BaseModel):
elif get_origin(field_type) is Union:
union_args = get_args(field_type)
if type(None) in union_args:
non_none_type = next(arg for arg in union_args if arg is not type(None))
return f"Optional[{self._get_field_type(field.__class__(annotation=non_none_type), depth)}]"
else:
return f"Union[{', '.join(arg.__name__ for arg in union_args)}]"
elif isinstance(field_type, type) and issubclass(field_type, BaseModel):
return self._get_model_schema(field_type, depth)
else:
return field_type.__name__
return getattr(field_type, "__name__", str(field_type))