Issue google-gemini#560

Needs to set required fields as required

Handle cases dataclasses, TypedDict and Pydantic models

Author: Mattias Wigh

google/generativeai/types/content_types.py

@@ -23,6 +23,7 @@
 import typing
 from typing import Any, Callable, Union, get_type_hints, get_origin, get_args
 from typing_extensions import TypedDict, is_typeddict
+import dataclasses
 
 import pydantic
 
@@ -334,9 +335,11 @@ def to_contents(contents: ContentsType) -> list[protos.Content]:
     return contents
 
 
-def _schema_for_class(cls: TypedDict) -> dict[str, Any]:
+def _schema_for_class(cls: type) -> dict[str, Any]:
     schema = _build_schema("dummy", {"dummy": (cls, pydantic.Field())})
     properties = schema["properties"]["dummy"]
+
+    # Handling TypedDict
     if is_typeddict(cls):
         required_keys = []
         type_hints = get_type_hints(cls)
@@ -347,6 +350,26 @@ def _schema_for_class(cls: TypedDict) -> dict[str, Any]:
                     continue
                 required_keys.append(key)
         properties["required"] = required_keys
+
+    # Handling dataclasses
+    elif dataclasses.is_dataclass(cls):
+        required_keys = []
+        for field in dataclasses.fields(cls):
+            if field.default is dataclasses.MISSING and field.default_factory is dataclasses.MISSING:
+                required_keys.append(field.name)  # Field is required if it has no default value
+        properties["required"] = required_keys
+
+    # Handling Pydantic models
+    elif issubclass(cls, pydantic.BaseModel):
+        required_keys = [name for name, field in cls.__fields__.items() if field.is_required()]
+        properties["required"] = required_keys
+
+    # Bug that it sets default values in case default exists
+    # TODO: Should be handled in the schema generation or not be allowed
+
+    for key in properties["properties"]:
+        if 'default' in properties["properties"][key]:
+            properties["properties"][key].pop('default')
     return properties