Structured outputs on Azure Databricks

Article
10/30/2024

Important

This feature is in Public Preview and is supported on both Foundation Model APIs pay-per-token and provisioned throughput endpoints.

This article describes structured outputs on Azure Databricks and how to use them as part of your generative AI application workflows. Structured outputs is OpenAI-compatible and is only available during model serving as part of Foundation Model APIs.

What are structured outputs?

Structured outputs provide a way to generate structured data in the form of JSON objects from your input data. You can choose to generate text, unstructured JSON objects, and JSON objects that adhere to a specific JSON schema. Structured outputs are supported for chat models served using Foundation Model APIs pay-per-token and provisioned throughput endpoints.

Databricks recommends using structured outputs for the following scenarios:

Extracting data from large amounts of documents. For example, identifying and classifying product review feedback as negative, positive or neutral.
Batch inference tasks that require outputs to be in a specified format.
Data processing, like turning unstructured data into structured data.

Use structured outputs

Specify your structured outputs using response_format in your chat request. See Foundation model REST API reference.

The following is an example of data extraction of research papers to a specific JSON schema.

import os
import json
from openai import OpenAI

DATABRICKS_TOKEN = os.environ.get('YOUR_DATABRICKS_TOKEN')
DATABRICKS_BASE_URL = os.environ.get('YOUR_DATABRICKS_BASE_URL')

client = OpenAI(
  api_key=DATABRICKS_TOKEN,
  base_url=DATABRICKS_BASE_URL
  )

response_format = {
      "type": "json_schema",
      "json_schema": {
        "name": "research_paper_extraction",
        "schema": {
          "type": "object",
          "properties": {
            "title": { "type": "string" },
            "authors": {
              "type": "array",
              "items": { "type": "string" }
            },
            "abstract": { "type": "string" },
            "keywords": {
              "type": "array",
              "items": { "type": "string" }
            }
          },
        },
        "strict": True
      }
    }

messages = [{
        "role": "system",
        "content": "You are an expert at structured data extraction. You will be given unstructured text from a research paper and should convert it into the given structure."
      },
      {
        "role": "user",
        "content": "..."
      }]

response = client.chat.completions.create(
    model="databricks-meta-llama-3-1-70b-instruct",
    messages=messages,
    response_format=response_format
)

print(json.dumps(response.choices[0].message.model_dump()['content'], indent=2))

The following is an example of JSON extraction, but the JSON schema is not known before hand.

import os
import json
from openai import OpenAI

DATABRICKS_TOKEN = os.environ.get('YOUR_DATABRICKS_TOKEN')
DATABRICKS_BASE_URL = os.environ.get('YOUR_DATABRICKS_BASE_URL')

client = OpenAI(
  api_key=DATABRICKS_TOKEN,
  base_url=DATABRICKS_BASE_URL
  )

response_format = {
      "type": "json_object",
    }

messages = [
      {
        "role": "user",
        "content": "Extract the name, size, price, and color from this product description as a JSON object:\n<description>\nThe SmartHome Mini is a compact smart home assistant available in black or white for only $49.99. It's 5 inches wide.\n</description>"
      }]

response = client.chat.completions.create(
    model="databricks-meta-llama-3-1-70b-instruct",
    messages=messages,
    response_format=response_format
)

print(json.dumps(response.choices[0].message.model_dump()['content'], indent=2))

JSON schema

Foundation Model APIs broadly support structured outputs accepted by OpenAI. However, using a simpler JSON schema for JSON schema definitions results in higher quality JSON generation. To promote higher quality generation, Foundation Model APIs only support a subset of JSON schema specifications.

The following function call definition keys are not supported:

Regular expressions using pattern.
Complex nested or schema composition and validation using: anyOf, oneOf, allOf, prefixItems, or $ref.
Lists of types except for the special case of [type, “null”] where one type in the list is a valid JSON type and the other is "null"

Token usage

Prompt injection and other techniques are used to enhance the quality of structured outputs. Doing so impacts the number of input and output tokens consumed by the model, which in turn results in billing implications.

Limitations

The maximum number of keys specified in the JSON schema is 64.
Foundation Model APIs does not enforce length or size constraints for objects and arrays.
- This includes keywords like maxProperties, minProperties, and maxLength.
Heavily nested JSON schemas result in lower quality generation. If possible, try flattening the JSON schema for better results.

Share via