gemini_response_schema_size.py

from pprint import pprint
import random
import string
from vertexai.generative_models import GenerationConfig, GenerativeModel
from google.api_core.exceptions import InvalidArgument

random.seed(1337)

prompt = "Respond according to the JSON schema."
num_properties = 40
property_name_length = 40

properties = [
    "".join(random.choices(string.ascii_lowercase, k=property_name_length))
    for _ in range(num_properties)
]

json_schema = {
    "type": "object",
    "properties": {name: {"type": "string"} for name in properties},
}

model = GenerativeModel("gemini-1.5-pro-001")

try:
    response = model.generate_content(
        contents=prompt,
        generation_config=GenerationConfig(
            temperature=0.0,
            response_mime_type="application/json",
            response_schema=json_schema,
        ),
    )
except InvalidArgument as e:
    print("Request failed as expected with InvalidArgument error:")
    print(e)
    print(
        f"generation_config.response_schema had {num_properties} properties, {property_name_length} characters each:"
    )
    pprint(json_schema)