Gemini API with JSON schema

Gists

Overview

These are sample scripts in Python and Node.js for controlling the output format of the Gemini API using JSON schemas.

Description

In a previous report, “Taming the Wild Output: Effective Control of Gemini API Response Formats with response_mime_type,” I presented sample scripts created with Google Apps Script. Ref Following its publication, I received requests for sample scripts using Python and Node.js. This report addresses those requests by providing sample scripts in both languages.

Step of this approach

In this approach, in order to control the output format, JSON schema and response_mime_type are used. By this, the complicated JSON object can be returned.

Sample scripts for Python

Sample 1

import google.generativeai as genai
import google.ai.generativelanguage as glm
import json

apiKey = "###" # Please set your API key.


jsonSchema = {
    "title": "Current population of Kyoto, Osaka, Aichi, Fukuoka, Tokyo in Japan",
    "description": "Return the current population of Kyoto, Osaka, Aichi, Fukuoka, Tokyo in Japan",
    "type": "object",
    "properties": {
        "propertyNames": {"description": "Prefecture names"},
        "patternProperties": {"": {"type": "number", "description": "Population"}},
    },
}
prompt = f"Follow JSON schema.<JSONSchema>{json.dumps(jsonSchema)}</JSONSchema>"
genai.configure(api_key=apiKey)
model = genai.GenerativeModel(
    model_name="gemini-1.5-pro-latest",
    generation_config=glm.GenerationConfig(response_mime_type="application/json"),
)
response = model.generate_content(prompt)
print(response.text)

When this script is run, the following result is obtained.

{
  "Kyoto": 2539740,
  "Osaka": 8837681,
  "Aichi": 7552873,
  "Fukuoka": 5138026,
  "Tokyo": 14064563
}

Sample 2

import google.generativeai as genai
import google.ai.generativelanguage as glm
import json

apiKey = "###" # Please set your API key.


jsonSchema = {
    "title": "5 popular cookie recipes",
    "description": "List 5 popular cookie recipes by including the following properties.",
    "type": "array",
    "items": {
        "type": "object",
        "properties": {
            "recipe_name": {"description": "Names of recipe.", "type": "string"},
            "materials": {
                "description": "Requirement materials for running the recipe.",
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "material": {
                            "description": "Requirement material for running the recipe.",
                            "type": "string",
                        },
                        "amount": {
                            "description": "Requirement amount of material for running the recipe. Unit is grams.",
                            "type": "number",
                        },
                        "cost": {
                            "description": "Cost of requirement material for running the recipe. Unit is dollar.",
                            "type": "number",
                        },
                    },
                    "required": ["material", "material", "cost"],
                    "additionalProperties": False,
                },
            },
        },
        "additionalProperties": False,
    },
}

prompt = f"Follow JSON schema.<JSONSchema>{json.dumps(jsonSchema)}</JSONSchema>"
genai.configure(api_key=apiKey)
model = genai.GenerativeModel(
    model_name="gemini-1.5-pro-latest",
    generation_config=glm.GenerationConfig(response_mime_type="application/json"),
)
response = model.generate_content(prompt)
print(response.text)

When this script is run, the following result is obtained.

[
  {
    "recipe_name": "Chocolate Chip Cookies",
    "materials": [
      { "material": "Butter", "amount": 226, "cost": 2.5 },
      { "material": "Sugar", "amount": 200, "cost": 0.5 },
      { "material": "Flour", "amount": 320, "cost": 0.8 }
    ]
  },
  {
    "recipe_name": "Oatmeal Raisin Cookies",
    "materials": [
      { "material": "Butter", "amount": 113, "cost": 1.25 },
      { "material": "Sugar", "amount": 150, "cost": 0.4 },
      { "material": "Flour", "amount": 230, "cost": 0.58 }
    ]
  },
  {
    "recipe_name": "Peanut Butter Cookies",
    "materials": [
      { "material": "Butter", "amount": 113, "cost": 1.25 },
      { "material": "Peanut Butter", "amount": 120, "cost": 1.2 },
      { "material": "Sugar", "amount": 150, "cost": 0.4 }
    ]
  },
  {
    "recipe_name": "Sugar Cookies",
    "materials": [
      { "material": "Butter", "amount": 226, "cost": 2.5 },
      { "material": "Sugar", "amount": 200, "cost": 0.5 },
      { "material": "Flour", "amount": 320, "cost": 0.8 }
    ]
  },
  {
    "recipe_name": "Shortbread Cookies",
    "materials": [
      { "material": "Butter", "amount": 226, "cost": 2.5 },
      { "material": "Sugar", "amount": 100, "cost": 0.25 },
      { "material": "Flour", "amount": 320, "cost": 0.8 }
    ]
  }
]

Sample scripts for Node.js

Sample 1

const { GoogleGenerativeAI } = require("@google/generative-ai");

const apiKey = "###"; // Please set your API key.

const genAI = new GoogleGenerativeAI(apiKey);

async function main() {
  const model = genAI.getGenerativeModel({
    model: "gemini-1.5-pro-latest",
    generationConfig: { response_mime_type: "application/json" },
  });
  const jsonSchema = {
    title: "Current population of Kyoto, Osaka, Aichi, Fukuoka, Tokyo in Japan",
    description:
      "Return the current population of Kyoto, Osaka, Aichi, Fukuoka, Tokyo in Japan",
    type: "object",
    properties: {
      propertyNames: { description: "Prefecture names" },
      patternProperties: { "": { type: "number", description: "Population" } },
    },
  };
  const prompt = `Follow JSON schema.<JSONSchema>${JSON.stringify(
    jsonSchema
  )}</JSONSchema>`;
  const result = await model.generateContent(prompt);
  const text = await result.response.text();
  console.log(text);
}

main();

When this script is run, the following result is obtained.

{
  "Kyoto": 1446575,
  "Osaka": 8808183,
  "Aichi": 7552112,
  "Fukuoka": 5138809,
  "Tokyo": 14047594
}

Sample 2

const { GoogleGenerativeAI } = require("@google/generative-ai");

const apiKey = "###"; // Please set your API key.

const genAI = new GoogleGenerativeAI(apiKey);

async function main() {
  const model = genAI.getGenerativeModel({
    model: "gemini-1.5-pro-latest",
    generationConfig: { response_mime_type: "application/json" },
  });
  const jsonSchema = {
    title: "5 popular cookie recipes",
    description:
      "List 5 popular cookie recipes by including the following properties.",
    type: "array",
    items: {
      type: "object",
      properties: {
        recipe_name: { description: "Names of recipe.", type: "string" },
        materials: {
          description: "Requirement materials for running the recipe.",
          type: "array",
          items: {
            type: "object",
            properties: {
              material: {
                description: "Requirement material for running the recipe.",
                type: "string",
              },
              amount: {
                description:
                  "Requirement amount of material for running the recipe. Unit is grams.",
                type: "number",
              },
              cost: {
                description:
                  "Cost of requirement material for running the recipe. Unit is dollar.",
                type: "number",
              },
            },
            required: ["material", "material", "cost"],
            additionalProperties: false,
          },
        },
      },
      additionalProperties: false,
    },
  };
  const prompt = `Follow JSON schema.<JSONSchema>${JSON.stringify(
    jsonSchema
  )}</JSONSchema>`;
  const result = await model.generateContent(prompt);
  const text = await result.response.text();
  console.log(text);
}

main();

When this script is run, the following result is obtained.

[
  {
    "recipe_name": "Chocolate Chip Cookies",
    "materials": [
      { "material": "Butter", "amount": 115, "cost": 2 },
      { "material": "Sugar", "amount": 100, "cost": 1 },
      { "material": "Flour", "amount": 150, "cost": 0.5 }
    ]
  },
  {
    "recipe_name": "Peanut Butter Cookies",
    "materials": [
      { "material": "Peanut Butter", "amount": 120, "cost": 3 },
      { "material": "Sugar", "amount": 100, "cost": 1 },
      { "material": "Flour", "amount": 150, "cost": 0.5 }
    ]
  },
  {
    "recipe_name": "Oatmeal Raisin Cookies",
    "materials": [
      { "material": "Oatmeal", "amount": 100, "cost": 2 },
      { "material": "Raisin", "amount": 80, "cost": 2.5 },
      { "material": "Flour", "amount": 150, "cost": 0.5 }
    ]
  },
  {
    "recipe_name": "Sugar Cookies",
    "materials": [
      { "material": "Butter", "amount": 115, "cost": 2 },
      { "material": "Sugar", "amount": 150, "cost": 1.5 },
      { "material": "Flour", "amount": 150, "cost": 0.5 }
    ]
  },
  {
    "recipe_name": "Shortbread Cookies",
    "materials": [
      { "material": "Butter", "amount": 120, "cost": 2 },
      { "material": "Sugar", "amount": 60, "cost": 0.5 },
      { "material": "Flour", "amount": 180, "cost": 0.7 }
    ]
  }
]

 Share!