"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

[dwilkinssg]

Hi,

I have the same error. Running Python 3.10.2 on Mac.

Was there a resolution to this? Could it be to do with Pydantic version?

Cheers,

David

[dwilkinssg]

Update:

I did a check on my pydantic install, which was as follows: pydantic==1.10.14 pydantic_core==2.14.6

Upgrading pydantic yielded: pydantic==2.8.2 pydantic_core==2.20.1

With this the KeyError issue is removed and the Gemini API performs as expected. (Now I need to check for any dependencies on my project to the Pydantic v1...)

[github-actions[bot]]

Marking this issue as stale since it has been open for 14 days with no activity. This issue will be closed if no further activity occurs.

[MaKTaiL]

I'm also having an issue with response_schema not working as intended. It appears to ignore all keys of a Dict except the first one. This is my class:

class SubtitleObject(typing.TypedDict):
    index: str
    content: str

This is the response I'm getting:

[{'index': '0'}, {'index': '1'}, {'index': '2'}, {'index': '3'}, {'index': '4'}, {'index': '5'}, {'index': '6'}, {'index': '7'}, {'index': '8'}, {'index': '9'}, {'index': '10'}]

Even the official docs show the same bug: