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

Portkey 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 portkey_ai import Portkey
from pydantic import BaseModel

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

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

portkey = Portkey(
    apiKey= "PORTKEY_API_KEY",
    virtual_key="VERTEX_VIRTUAL_KEY"
)

completion = portkey.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 portkey_ai import Portkey
from enum import Enum
from pydantic import BaseModel

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

# Initialize Portkey with API details
portkey = Portkey(
    api_key="PORTKEY_API_KEY",
    virtual_key="VERTEX_VIRTUAL_KEY"
)

# Simple enum classification
completion = portkey.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.

import Portkey from "portkey-ai";

const portkey = new Portkey({
  apiKey: "PORTKEY_API_KEY",
  virtualKey: "VERTEX_VIRTUAL_KEY",
});

async function main() {
  const completion = await portkey.chat.completions.create({
    model: "gemini-1.5-pro-002",
    messages: [
      { role: "system", content: "Extract the event information." },
      {
        role: "user",
        content: "Alice and Bob are going to a science fair on Friday.",
      },
    ],
    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,
      },
    },
  });
  const event = completion.choices[0].message?.content;
  console.log(event);
}

main();

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

Was this page helpful?

Google Vertex AIAzure OpenAI