Query model provider services

Important

this feature is in Beta. Account admins can manage access to this feature from the account console Previews page. See Manage Azure Databricks previews.

This page describes how to query a model provider service through Unity AI Gateway. Unity AI Gateway supplies the stored credentials and routes the request to the external provider, so callers don't handle the provider secret.

Requirements

Identify a model provider service

You select a model provider service for a request with the Databricks-Model-Provider-Service header, set to the service's three-part name:

Databricks-Model-Provider-Service: main.default.openai_prod

Authenticate with your Azure Databricks token, not the provider's credential. The base URL is your workspace URL followed by /ai-gateway.

Query Supported APIs

Managed paths make available each provider's API under a stable Unity AI Gateway path. Unity AI Gateway translates between the request and the provider, applies governance such as guardrails and rate limits, and records usage. This is the recommended way to query a model provider service.

The following example sends a chat completion through an OpenAI model provider service using the managed OpenAI path. Because the request uses the OpenAI Chat Completions API, you can point the OpenAI client at the Unity AI Gateway base URL.

Python

from openai import OpenAI

client = OpenAI(
    api_key="<databricks-token>",
    base_url="https://<workspace-url>/ai-gateway/openai/v1",
    default_headers={"Databricks-Model-Provider-Service": "main.default.openai_prod"},
)

response = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[{"role": "user", "content": "Say hello in exactly 3 words."}],
)
print(response.choices[0].message.content)

REST

curl https://<workspace-url>/ai-gateway/openai/v1/chat/completions \
  -H "Authorization: Bearer $DATABRICKS_TOKEN" \
  -H "Databricks-Model-Provider-Service: main.default.openai_prod" \
  -H "Content-Type: application/json" \
  -d '{
    "model": "gpt-4o-mini",
    "messages": [{"role": "user", "content": "Say hello in exactly 3 words."}]
  }'

The managed path you call depends on the provider's API:

Provider API Managed path
OpenAI (chat completions) /ai-gateway/openai/v1/chat/completions
OpenAI (responses) /ai-gateway/openai/v1/responses
OpenAI (embeddings) /ai-gateway/openai/v1/embeddings
Anthropic (messages) /ai-gateway/anthropic/v1/messages
Gemini (generate content) /ai-gateway/gemini/v1beta/models/<model>:generateContent

The model in the request body (or the Gemini path segment) must be a model the model provider service allows. See Create and manage model provider services for allow_all_targets and targets.

Query Other APIs (Passthrough)

If a managed path does not cover a provider endpoint — for example, an OpenAI file or batch endpoint — you can pass the request through to the provider unchanged. Unity AI Gateway strips the /ai-gateway prefix, attaches the stored credential, and forwards the remaining path to the provider.

Unmanaged passthrough is off by default. Enable it on the model provider service by setting forward_unmanaged_paths to true in the service configuration:

curl https://<workspace-url>/api/2.1/unity-catalog/model-provider-services/main.default.openai_prod \
  -X PATCH \
  -H "Authorization: Bearer $DATABRICKS_TOKEN" \
  -H "Content-Type: application/json" \
  -G \
  --data-urlencode "update_mask=config.forward_unmanaged_paths" \
  --data '{ "config": { "forward_unmanaged_paths": true } }'

After you enable passthrough, call the provider's native path under /ai-gateway:

curl https://<workspace-url>/ai-gateway/v1/files \
  -H "Authorization: Bearer $DATABRICKS_TOKEN" \
  -H "Databricks-Model-Provider-Service: main.default.openai_prod"

Headers and query parameters forwarding

By default, Unity AI Gateway does not pass the client's request headers or query parameters to the upstream provider. Two service-configuration flags change this, and they apply to both managed and unmanaged paths:

  • forward_headers — when true, Unity AI Gateway forwards client request headers to the provider. Enable this when a provider requires a header that Unity AI Gateway does not set for you, such as OpenAI-Organization.
  • forward_query_parameters — when true, Unity AI Gateway forwards client query parameters to the provider.

Set them on the model provider service like any other configuration field:

curl https://<workspace-url>/api/2.1/unity-catalog/model-provider-services/main.default.openai_prod \
  -X PATCH \
  -H "Authorization: Bearer $DATABRICKS_TOKEN" \
  -H "Content-Type: application/json" \
  -G \
  --data-urlencode "update_mask=config.forward_headers,config.forward_query_parameters" \
  --data '{ "config": { "forward_headers": true, "forward_query_parameters": true } }'

Tag requests for usage tracking

You can attach custom key-value tags to individual requests using the Databricks-Ai-Gateway-Request-Tags HTTP header. Unity AI Gateway logs request tags to the request_tags column in both the usage tracking system table and inference tables. You can then track costs, attribute usage, and filter analytics by project, team, environment, or another dimension.

The header value must be a JSON object mapping string keys to string values. For example:

{ "project": "chatbot", "team": "ml-platform", "environment": "production" }

Send the tags header alongside the Databricks-Model-Provider-Service header. Use the extra_headers parameter (Python) or pass the header directly (REST API):

Python

from openai import OpenAI
import json

client = OpenAI(
    api_key="<databricks-token>",
    base_url="https://<workspace-url>/ai-gateway/openai/v1",
    default_headers={"Databricks-Model-Provider-Service": "main.default.openai_prod"},
)

request_tags = {"project": "chatbot", "team": "ml-platform"}

response = client.chat.completions.create(
    model="gpt-5.5",
    messages=[{"role": "user", "content": "What is Databricks?"}],
    extra_headers={"Databricks-Ai-Gateway-Request-Tags": json.dumps(request_tags)},
)

REST

curl https://<workspace-url>/ai-gateway/openai/v1/chat/completions \
  -H "Authorization: Bearer $DATABRICKS_TOKEN" \
  -H "Databricks-Model-Provider-Service: main.default.openai_prod" \
  -H "Content-Type: application/json" \
  -H 'Databricks-Ai-Gateway-Request-Tags: {"project": "chatbot", "team": "ml-platform"}' \
  -d '{
    "model": "gpt-5.5",
    "messages": [{"role": "user", "content": "What is Databricks?"}]
  }'

Next steps

  • Last updated on