Controlled Generations

Controlled Generations ensure that the model always follows your supplied JSON schema. Portkey supports Vertex AI’s Controlled Generations feature out of the box with our SDKs & APIs.

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.

Pydantic with Python

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)

Zod with NodeJS

To use Zod with VerteX AI you will also need to import { zodResponseFormat } from openai/helpers/zod

import { Portkey } from 'portkey-ai';
import { z } from 'zod';
import { zodResponseFormat } from "openai/helpers/zod";

const MathReasoning = z.object({
  steps: z.array(z.object({ explanation: z.string(), output: z.string() })),
  final_answer: z.string()
});

const portkey = new Portkey({
    apiKey: "YOUR_API_KEY",
    virtualKey: "YOUR_GEMINI_VIRTUAL_KEY"
});

async function runMathTutor() {
  try {
    const completion = await portkey.chat.completions.create({
      model: "gemini-1.5-pro-002",
      messages: [
        { role: "system", content: "You are a helpful math tutor." },
        { role: "user", content: "Solve 8x + 7 = -23" }
      ],
      response_format: zodResponseFormat(MathReasoning, "MathReasoning")
    });

    console.log(completion.choices[0].message.content);
  } catch (error) {
    console.error("Error:", error);
  }
}

runMathTutor();

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.

Python with Enums

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.

NodeJS

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();

Python

from portkey_ai import Portkey

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

completion = portkey.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)

cURL

curl https://api.portkey.ai/v1/chat/completions \
  -H "x-portkey-api-key: $PORTKEY_API_KEY" \
  -H "x-portkey-virtual-key: $VERTEX_VIRTUAL_KEY" \
  -H "Content-Type: application/json" \
  -d '{
    "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": {
      "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
      }
    }
  }'

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