Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
Needs to set required fields as required
  • Loading branch information
Mattias Wigh authored and Mattias Wigh committed Oct 2, 2024
1 parent 8f77cc6 commit f30c36d
Showing 1 changed file with 14 additions and 3 deletions.
17 changes: 14 additions & 3 deletions google/generativeai/types/content_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,8 @@
import mimetypes
import pathlib
import typing
from typing import Any, Callable, Union
from typing_extensions import TypedDict
from typing import Any, Callable, Union, get_type_hints, get_origin, get_args
from typing_extensions import TypedDict, is_typeddict

import pydantic

Expand Down Expand Up @@ -336,7 +336,18 @@ def to_contents(contents: ContentsType) -> list[protos.Content]:

def _schema_for_class(cls: TypedDict) -> dict[str, Any]:
schema = _build_schema("dummy", {"dummy": (cls, pydantic.Field())})
return schema["properties"]["dummy"]
properties = schema["properties"]["dummy"]
if is_typeddict(cls):
required_keys = []
type_hints = get_type_hints(cls)
for key, type_hint in type_hints.items():
if key in cls.__required_keys__:
# Check if the type is Optional, i.e., Union[..., NoneType]
if get_origin(type_hint) is Union and type(None) in get_args(type_hint):
continue
required_keys.append(key)
properties["required"] = required_keys
return properties


def _schema_for_function(
Expand Down

0 comments on commit f30c36d

Please sign in to comment.