Skip to main content
Last updated on

Query model provider services

Beta

this feature is in Beta. Account admins can manage access to this feature from the account console Previews page. See Manage 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:

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

Authenticate with your 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)

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:

Bash
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:

Bash
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:

Bash
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:

JSON
{ "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)},
)

Next steps​

On this page