"response_schema" results in KeyError

[ganjigajanan]

Description of the bug:

Hello, I'm trying to recreate the example at: https://github.com/google-gemini/cookbook/blob/main/examples/json_capabilities/Text_Summarization.ipynb. After specifying the classes and creating the generation_config, while invoking the genai.GenerativeModel model with generation_config results in a KeyError.

In [22]: from google import generativeai as genai
         from typing_extensions import TypedDict

In [23]: genai.__version__
Out[23]: '0.7.1'

In [24]: class Character(TypedDict):
            name: str
            description: str
            alignment: str

        class Location(TypedDict):
            name: str
            description: str

        class TextSummary(TypedDict):
            synopsis: str
            genres: list[str]
            locations: list[Location]
            characters: list[Character]

In [25]: generation_config = {
            "response_mime_type": "application/json",
            "response_schema": TextSummary
        }

In [26]: model = Genai.GenerativeModel(
            model_name=model_name,
            generation_config=generation_config,
        )

Traceback (most recent call last):

  Cell In[26], line 1
    model = Genai.GenerativeModel(

  File ~\AppData\Local\anaconda3\Lib\site-packages\google\generativeai\generative_models.py:86 in __init__
    self._generation_config = generation_types.to_generation_config_dict(generation_config)

  File ~\AppData\Local\anaconda3\Lib\site-packages\google\generativeai\types\generation_types.py:213 in to_generation_config_dict
    _normalize_schema(generation_config)

  File ~\AppData\Local\anaconda3\Lib\site-packages\google\generativeai\types\generation_types.py:184 in _normalize_schema
    response_schema = content_types._schema_for_class(response_schema)

  File ~\AppData\Local\anaconda3\Lib\site-packages\google\generativeai\types\content_types.py:298 in _schema_for_class
    schema = _build_schema("dummy", {"dummy": (cls, pydantic.Field())})

  File ~\AppData\Local\anaconda3\Lib\site-packages\google\generativeai\types\content_types.py:383 in _build_schema
    unpack_defs(parameters, defs)

  File ~\AppData\Local\anaconda3\Lib\site-packages\google\generativeai\types\content_types.py:404 in unpack_defs
    ref = defs[ref_key.split("defs/")[-1]]

KeyError: '#/definitions/TextSummary'

Actual vs expected behavior:

No response

Any other information you'd like to share?

No response

[sineeli]

Hi @ganjigajanan,

I have tried in the similar way passing the generation_config while declaring the model and it works for me as expected, can you take a look at the gist try and let me know if still the error persists.

model = genai.GenerativeModel(
    model_name="gemini-1.5-pro-latest",
    generation_config={"response_mime_type": "application/json", "response_schema": TextSummary},
)

Thanks