Controlled Generations allows you to constrain model responses to predefined sets of values. This is particularly useful for classification tasks, multiple choice responses, and structured data extraction.

This feature is available for Gemini 1.5 Pro & Gemini 1.5 Flash models.

With Pydantic & Zod

Obiguard SDKs for Python and JavaScript also make it easy to define object schemas using Pydantic and Zod respectively. Below, you can see how to extract information from unstructured text that conforms to a schema defined in code.

from obiguard import Obiguard
from pydantic import BaseModel

class Step(BaseModel):
    explanation: str
    output: str

class MathReasoning(BaseModel):
    steps: list[Step]
    final_answer: str

client = Obiguard(
    obiguard_api_key="OBIGUARD_API_KEY",  # Your Obiguard API key
    virtual_key="VERTEX_VIRTUAL_KEY"
)

completion = client.beta.chat.completions.parse(
    model="gemini-1.5-pro-002",
    messages=[
        {"role": "system", "content": "You are a helpful math tutor. Guide the user through the solution step by step."},
        {"role": "user", "content": "how can I solve 8x + 7 = -23"}
    ],
    response_format=MathReasoning,
)

print(completion.choices[0].message)
print(completion.choices[0].message.parsed)

Using Enums

You can also use enums to constrain the model’s output to a predefined set of values. This is particularly useful for classification tasks and multiple choice responses.

from openai import OpenAI
from enum import Enum
from pydantic import BaseModel

class InstrumentClass(Enum):
    PERCUSSION = "Percussion"
    STRING = "String"
    WOODWIND = "Woodwind"
    BRASS = "Brass"
    KEYBOARD = "Keyboard"

client = Obiguard(
    obiguard_api_key="sk-obg******",  # Your Obiguard API key
    virtual_key="VERTEX_VIRTUAL_KEY"
)

# Simple enum classification
completion = client.chat.completions.create(
    model="gemini-1.5-pro-002",
    messages=[
        {"role": "system", "content": "Classify the musical instrument."},
        {"role": "user", "content": "What type of instrument is a piano?"}
    ],
    response_format={
        "type": "json_schema",
        "json_schema": {
            "type": "string",
            "enum": [e.value for e in InstrumentClass],
            "title": "instrument_classification"
        }
    }
)

print(completion.choices[0].message.content)

Using JSON schema Directly

This method is more portable across different languages and doesn’t require additional libraries, but lacks the integrated type checking of the Pydantic/Zod approach. Choose the method that best fits your project’s needs.

from obiguard import Obiguard

client = Obiguard(
  obiguard_api_key="sk-obg***",  # Your Obiguard API key
  virtual_key="VERTEX_VIRTUAL_KEY"
)

completion = client.chat.completions.create(
    model="gemini-1.5-pro-002",
    messages=[
        {"role": "system", "content": "Extract the event information."},
        {"role": "user", "content": "A meteor the size of 1000 football stadiums will hit earth this Sunday"},
    ],
    response_format={
          "type": "json_schema",
          "json_schema": {
            "name": "math_reasoning",
            "schema": {
              "type": "object",
              "properties": {
                "steps": {
                  "type": "array",
                  "items": {
                    "type": "object",
                    "properties": {
                      "explanation": { "type": "string" },
                      "output": { "type": "string" }
                    },
                    "required": ["explanation", "output"],
                    "additionalProperties": False
                  }
                },
                "final_answer": { "type": "string" }
              },
              "required": ["steps", "final_answer"],
              "additionalProperties": False
            },
            "strict": True
          }
        },
)

print(completion.choices[0].message.content)

For more, refer to Google Vertex AI’s detailed documentation on Controlled Generations here.