Mosaic AI Gateway: Unified AI Governance at Scale

Table of Contents

Databricks Deep Dive - This article is part of a series.

Part : Agent Bricks: Building Enterprise AI Agents on Databricks

Part : Lakebase: Postgres Meets the Lakehouse — Hands-On with Databricks' OLTP Play

Part : Lakeflow GA: The Unified Data Engineering Platform You've Been Waiting For

Part : MLflow 3: Rebuilt for the GenAI Era

Part : This Article

Part : Serverless GPU & AI Runtime: Running GPU Workloads Without Cluster Management on Databricks

Part : Unity Catalog Deep Dive: Fine-Grained Access Control

Every enterprise deploying LLMs hits the same wall: governance at scale. You start with one OpenAI endpoint. Then someone adds Anthropic. A team fine-tunes Llama on provisioned throughput. Another team spins up a custom model. Before you know it, you have five different API keys scattered across three clouds, no unified logging, no rate limiting, and a CFO asking why the AI bill tripled last month.

Mosaic AI Gateway is Databricks’ answer to this chaos. It’s a centralized proxy layer that sits between your applications and every AI model — whether hosted on Databricks, routed through OpenAI, Anthropic, Bedrock, or any custom provider. One API surface. One governance model. One place to enforce guardrails, track costs, and prove compliance.

In this article, we’ll go beyond the overview and get hands-on: configure multi-model endpoints, set up rate limiting, enable guardrails, wire up cost attribution, and query the unified audit trail — all with real code you can run today.

Why AI Gateway Matters
#

Without a gateway, the typical enterprise AI architecture looks like this:

1
2
3
4
5
6
7
8
9


[App A] --> [OpenAI API]      (API key in env var)
[App B] --> [Anthropic API]   (API key in Vault)
[App C] --> [Bedrock]         (IAM role)
[App D] --> [Custom Model]    (mTLS cert)
                |
        No unified logging
        No rate limiting
        No cost attribution
        No guardrails

Every model provider has its own auth, its own rate limits, its own logging format. Your security team can’t audit what prompts went where. Your finance team can’t attribute costs to projects. Your compliance team can’t prove PII isn’t leaking into third-party models.

With AI Gateway, this collapses:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


[App A] ─┐
[App B] ─┤
[App C] ─┼──> [AI Gateway] ──> [Any Model Provider]
[App D] ─┘        │
              ┌────┴────┐
              │ Unified  │
              │ Logging  │
              │ Rate     │
              │ Limits   │
              │ Guards   │
              │ Costs    │
              └──────────┘

One API. One token. One governance layer. Every request — whether it’s going to GPT-4o, Claude Sonnet, or a fine-tuned Llama 3.1 — flows through the same control plane.

AI Gateway supports 300K+ QPS per endpoint with route optimization enabled, adding less than 20ms of overhead. This isn’t a bottleneck — it’s infrastructure that scales beyond what most organizations will ever need.

Prerequisites
#

Before you start:

A Databricks workspace on AWS, Azure, or GCP
Workspace admin privileges (to create serving endpoints)
A Databricks personal access token (PAT) or service principal token
Python 3.8+ with databricks-sdk and openai packages installed
Basic familiarity with REST APIs and the Databricks UI
A Unity Catalog metastore attached to your workspace (for logging and audit)

Architecture: How AI Gateway Works
#

AI Gateway is built into Databricks Model Serving infrastructure. It’s not a separate service you deploy — it’s a configuration layer on top of serving endpoints.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25


┌─────────────────────────────────────────────────┐
│                  AI Gateway                      │
│                                                  │
│  ┌──────────┐ ┌──────────┐ ┌──────────────────┐ │
│  │ Auth &   │ │ Rate     │ │ Guardrails       │ │
│  │ Routing  │ │ Limiter  │ │ (Safety + PII)   │ │
│  └────┬─────┘ └────┬─────┘ └───────┬──────────┘ │
│       └─────────────┼───────────────┘            │
│                     │                            │
│              ┌──────▼──────┐                     │
│              │ Inference   │                     │
│              │ Logging     │                     │
│              │ (→ UC Delta)│                     │
│              └──────┬──────┘                     │
│                     │                            │
│  ┌──────────────────▼────────────────────────┐   │
│  │         Traffic Router                     │  │
│  │  (Splitting / Fallbacks / Load Balancing)  │  │
│  └──┬──────────┬──────────┬──────────┬───────┘  │
│     │          │          │          │           │
└─────┼──────────┼──────────┼──────────┼───────────┘
      ▼          ▼          ▼          ▼
  [Databricks  [OpenAI]  [Anthropic] [Custom
   Foundation                        Provider]
   Models]

Key architectural properties:

Horizontally scalable — the inference server, auth layer, proxy, and rate limiter all scale independently
Credential isolation — API keys are stored via Databricks Secrets, never exposed to application code
OpenAI-compatible API — your existing OpenAI SDK code works unchanged, just point it at the gateway
Unity Catalog native — all logs flow into governed Delta tables

Step 1: Create a Multi-Model Endpoint
#

Let’s start by setting up an endpoint that routes to an external model. This is how you bring OpenAI, Anthropic, or any provider under the AI Gateway umbrella.

External Model — Anthropic
#

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22


import mlflow.deployments

client = mlflow.deployments.get_deploy_client("databricks")

client.create_endpoint(
    name="claude-gateway",
    config={
        "served_entities": [
            {
                "name": "claude-sonnet",
                "external_model": {
                    "name": "claude-sonnet-4-5",
                    "provider": "anthropic",
                    "task": "llm/v1/chat",
                    "anthropic_config": {
                        "anthropic_api_key": "{{secrets/ai-keys/anthropic}}"
                    }
                }
            }
        ]
    }
)

External Model — OpenAI (Azure)
#

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22


client.create_endpoint(
    name="gpt-gateway",
    config={
        "served_entities": [
            {
                "name": "gpt-4o",
                "external_model": {
                    "name": "gpt-4o",
                    "provider": "openai",
                    "task": "llm/v1/chat",
                    "openai_config": {
                        "openai_api_type": "azure",
                        "openai_api_key": "{{secrets/ai-keys/azure-openai}}",
                        "openai_api_base": "https://my-org.openai.azure.com",
                        "openai_deployment_name": "gpt-4o-deployment",
                        "openai_api_version": "2024-02-15-preview"
                    }
                }
            }
        ]
    }
)

External Model — Amazon Bedrock
#

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20


client.create_endpoint(
    name="bedrock-gateway",
    config={
        "served_entities": [
            {
                "name": "bedrock-claude",
                "external_model": {
                    "name": "anthropic.claude-sonnet-4-v2:0",
                    "provider": "amazon-bedrock",
                    "task": "llm/v1/chat",
                    "amazon_bedrock_config": {
                        "aws_region": "us-east-1",
                        "uc_service_credential_name": "bedrock-credential",
                        "bedrock_provider": "anthropic"
                    }
                }
            }
        ]
    }
)

Custom Provider
#

Got a model running on your own infrastructure? No problem:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


client.create_endpoint(
    name="custom-gateway",
    config={
        "served_entities": [
            {
                "name": "internal-model",
                "external_model": {
                    "name": "our-fine-tuned-llama",
                    "provider": "custom",
                    "task": "llm/v1/chat",
                    "custom_provider_config": {
                        "custom_provider_url": "https://ml.internal.company.com/v1/chat",
                        "bearer_token_auth": {
                            "token": "{{secrets/ai-keys/internal-ml}}"
                        }
                    }
                }
            }
        ]
    }
)

Notice the pattern: every provider’s API key is stored in Databricks Secrets — never in code, never in environment variables. The {{secrets/scope/key}} notation is resolved server-side.

Seven built-in providers are supported: OpenAI, Anthropic, Cohere, Amazon Bedrock, Google Cloud Vertex AI, AI21 Labs, and Databricks Model Serving. The custom provider covers everything else with a simple URL + auth configuration.

Step 2: Query via the OpenAI SDK
#

Here’s the part that makes AI Gateway seamless: it speaks the OpenAI API. Your existing code works with a two-line change.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17


from openai import OpenAI

client = OpenAI(
    api_key="dapi-your-databricks-token",
    base_url="https://your-workspace.cloud.databricks.com/serving-endpoints"
)

response = client.chat.completions.create(
    model="claude-gateway",  # your AI Gateway endpoint name
    messages=[
        {"role": "system", "content": "You are a helpful data engineering assistant."},
        {"role": "user", "content": "Explain the difference between Z-Order and Liquid Clustering in Delta Lake."}
    ],
    temperature=0.7
)

print(response.choices[0].message.content)

Or query directly from a Databricks notebook using SQL:

1
2
3
4


SELECT ai_query(
    "claude-gateway",
    "What are the benefits of Unity Catalog for data governance?"
) AS response;

The application doesn’t know or care whether it’s talking to Anthropic, OpenAI, or a self-hosted model. It just calls the gateway endpoint. That’s the point.

Step 3: Traffic Splitting and Fallbacks
#

Real production deployments need more than a single model behind an endpoint. AI Gateway supports traffic splitting for A/B testing and fallbacks for reliability.

Traffic Splitting
#

Route 70% of traffic to one model and 30% to another — useful for gradual rollouts or comparing model quality:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31


client.create_endpoint(
    name="chat-production",
    config={
        "served_entities": [
            {
                "name": "primary-model",
                "external_model": {
                    "name": "claude-sonnet-4-5",
                    "provider": "anthropic",
                    "task": "llm/v1/chat",
                    "anthropic_config": {
                        "anthropic_api_key": "{{secrets/ai-keys/anthropic}}"
                    }
                },
                "traffic_percentage": 70
            },
            {
                "name": "challenger-model",
                "external_model": {
                    "name": "gpt-4o",
                    "provider": "openai",
                    "task": "llm/v1/chat",
                    "openai_config": {
                        "openai_api_key": "{{secrets/ai-keys/openai}}"
                    }
                },
                "traffic_percentage": 30
            }
        ]
    }
)

Fallbacks
#

If the primary model returns a 429 (rate limited) or 5XX error, the gateway automatically tries the next model. Set traffic percentage to 0% for fallback-only models:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


# Primary model gets 100% of traffic; fallback activates on errors
config = {
    "served_entities": [
        {
            "name": "primary",
            "external_model": { "name": "gpt-4o", "provider": "openai", ... },
            "traffic_percentage": 100
        },
        {
            "name": "fallback",
            "external_model": { "name": "claude-sonnet-4-5", "provider": "anthropic", ... },
            "traffic_percentage": 0  # fallback only
        }
    ]
}

Maximum two fallback models per endpoint. The gateway tries them sequentially — if the first fallback also fails, it tries the second before returning an error.

Step 4: Rate Limiting
#

Rate limiting is the simplest and most effective governance control. AI Gateway supports a three-tier structure:

Level	Scope	Description
Endpoint	Global	Max QPM/TPM ceiling for all traffic to this endpoint
User Default	Per-user	Default limit applied to every user unless overridden
Custom	Specific	Limits for specific users, service principals, or groups

Configure via the Databricks SDK
#

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27


from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import (
    AiGatewayConfig,
    AiGatewayRateLimit,
    AiGatewayRateLimitRenewalPeriod,
    AiGatewayRateLimitKey
)

w = WorkspaceClient()

w.serving_endpoints.put_ai_gateway(
    name="claude-gateway",
    rate_limits=[
        # Global endpoint limit: 10,000 queries per minute
        AiGatewayRateLimit(
            calls=10000,
            renewal_period=AiGatewayRateLimitRenewalPeriod.MINUTE,
            key=AiGatewayRateLimitKey.ENDPOINT
        ),
        # Default per-user limit: 100 queries per minute
        AiGatewayRateLimit(
            calls=100,
            renewal_period=AiGatewayRateLimitRenewalPeriod.MINUTE,
            key=AiGatewayRateLimitKey.USER
        )
    ]
)

Why this matters
#

Without rate limiting, a single runaway script can:

Burn through your entire OpenAI quota in minutes
Trigger provider-side throttling that affects all users
Generate a five-figure invoice from a weekend experiment

With AI Gateway, you set the guardrails once and every request — from notebooks, jobs, apps, or APIs — respects them.

Rate limiting is free — no additional DBU charges. You can set up to 20 rate limits per endpoint, including up to 5 group-specific limits. Both QPM (queries per minute) and TPM (tokens per minute) are supported.

Step 5: AI Guardrails
#

Rate limits control volume. Guardrails control content. AI Gateway offers two built-in guardrail types, both in Public Preview.

Safety Filtering
#

Powered by Meta Llama Guard 2-8B, safety filtering prevents models from interacting with unsafe content. It applies to both inputs (what users send) and outputs (what models return).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17


from databricks.sdk.service.serving import (
    AiGatewayConfig,
    AiGatewayGuardrails,
    AiGatewayGuardrailParameters
)

w.serving_endpoints.put_ai_gateway(
    name="claude-gateway",
    guardrails=AiGatewayGuardrails(
        input=AiGatewayGuardrailParameters(
            safety=True  # Enable safety filtering on inputs
        ),
        output=AiGatewayGuardrailParameters(
            safety=True  # Enable safety filtering on outputs
        )
    )
)

When safety filtering blocks a request, the gateway returns HTTP 400 — your application gets a clear signal that the content was rejected, not a model hallucination about why it can’t help.

PII Detection
#

Three modes: None (disabled), Block (reject requests containing PII), or Mask (redact PII before forwarding to the model).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19


from databricks.sdk.service.serving import AiGatewayGuardrailPiiBehavior

w.serving_endpoints.put_ai_gateway(
    name="claude-gateway",
    guardrails=AiGatewayGuardrails(
        input=AiGatewayGuardrailParameters(
            safety=True,
            pii=AiGatewayGuardrailPiiBehavior(
                behavior="BLOCK"  # or "MASK" or "NONE"
            )
        ),
        output=AiGatewayGuardrailParameters(
            safety=True,
            pii=AiGatewayGuardrailPiiBehavior(
                behavior="MASK"  # Mask PII in model responses
            )
        )
    )
)

Detected PII categories (U.S.): credit card numbers, Social Security numbers, phone numbers, email addresses, names, and addresses.

This is where AI Gateway and Unity Catalog’s security model converge. You’ve already locked down data access with column masking and row-level security. Now you’re locking down the AI layer too — preventing sensitive data from leaking through prompts or model responses.

Step 6: Payload Logging and Usage Tracking
#

Every request and response can be logged to inference tables — Delta tables managed by Unity Catalog. This gives you a complete audit trail of what was asked, what was answered, and who asked it.

Enable Inference Table Logging
#

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


from databricks.sdk.service.serving import AiGatewayInferenceTableConfig

w.serving_endpoints.put_ai_gateway(
    name="claude-gateway",
    inference_table_config=AiGatewayInferenceTableConfig(
        catalog_name="ml_governance",
        schema_name="ai_gateway_logs",
        enabled=True
    )
)

Once enabled, every request generates a row in the inference table with:

Full request payload (prompt, parameters)
Full response payload (completion, token counts)
Caller identity and timestamp
Latency metrics

Usage Tracking via System Tables
#

For aggregate usage analytics without storing full payloads, query the system table:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


-- Token usage by user over the last 30 days
SELECT
    eu.created_by AS user_email,
    se.served_entity_name AS model,
    SUM(eu.input_token_count) AS total_input_tokens,
    SUM(eu.output_token_count) AS total_output_tokens,
    COUNT(*) AS total_requests
FROM system.serving.endpoint_usage AS eu
JOIN system.serving.served_entities AS se
    ON eu.served_entity_id = se.served_entity_id
WHERE eu.request_date >= current_date() - INTERVAL 30 DAYS
GROUP BY eu.created_by, se.served_entity_name
ORDER BY total_output_tokens DESC;

Cost Attribution with Usage Context
#

The real power for finance teams: usage context lets you tag every request with project, team, or user identifiers for precise cost attribution.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


response = client.chat.completions.create(
    model="claude-gateway",
    messages=[{"role": "user", "content": "Summarize Q3 revenue trends."}],
    extra_body={
        "usage_context": {
            "project": "financial-reporting",
            "team": "data-analytics",
            "end_user_to_charge": "analyst-jane-doe"
        }
    }
)

Then query costs by project:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


-- Cost attribution by project and team
SELECT
    usage_context:project AS project,
    usage_context:team AS team,
    SUM(input_token_count + output_token_count) AS total_tokens,
    COUNT(*) AS request_count
FROM system.serving.endpoint_usage
WHERE request_date >= current_date() - INTERVAL 30 DAYS
GROUP BY usage_context:project, usage_context:team
ORDER BY total_tokens DESC;

No more arguing about which team ran up the bill. Every token is accounted for.

Step 7: Putting It All Together
#

Here’s a complete configuration that combines everything — multi-model routing with fallbacks, rate limiting, guardrails, and logging:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69


from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import *

w = WorkspaceClient()

# Create the endpoint with traffic splitting and fallbacks
w.serving_endpoints.create(
    name="enterprise-chat",
    config=EndpointCoreConfigInput(
        served_entities=[
            ServedEntityInput(
                name="primary-claude",
                external_model=ExternalModel(
                    name="claude-sonnet-4-5",
                    provider="anthropic",
                    task="llm/v1/chat",
                    anthropic_config=AnthropicConfig(
                        anthropic_api_key="{{secrets/ai-keys/anthropic}}"
                    )
                ),
                traffic_percentage=100
            ),
            ServedEntityInput(
                name="fallback-gpt",
                external_model=ExternalModel(
                    name="gpt-4o",
                    provider="openai",
                    task="llm/v1/chat",
                    openai_config=OpenAiConfig(
                        openai_api_key="{{secrets/ai-keys/openai}}"
                    )
                ),
                traffic_percentage=0  # fallback only
            )
        ]
    )
)

# Configure AI Gateway features
w.serving_endpoints.put_ai_gateway(
    name="enterprise-chat",
    rate_limits=[
        AiGatewayRateLimit(
            calls=10000,
            renewal_period=AiGatewayRateLimitRenewalPeriod.MINUTE,
            key=AiGatewayRateLimitKey.ENDPOINT
        ),
        AiGatewayRateLimit(
            calls=200,
            renewal_period=AiGatewayRateLimitRenewalPeriod.MINUTE,
            key=AiGatewayRateLimitKey.USER
        )
    ],
    guardrails=AiGatewayGuardrails(
        input=AiGatewayGuardrailParameters(
            safety=True,
            pii=AiGatewayGuardrailPiiBehavior(behavior="BLOCK")
        ),
        output=AiGatewayGuardrailParameters(
            safety=True,
            pii=AiGatewayGuardrailPiiBehavior(behavior="MASK")
        )
    ),
    inference_table_config=AiGatewayInferenceTableConfig(
        catalog_name="ml_governance",
        schema_name="ai_gateway_logs",
        enabled=True
    )
)

That’s a production-ready AI governance layer — deployed in under 50 lines of Python.

Databricks Foundation Models: Pay-Per-Token
#

AI Gateway isn’t just for external providers. Databricks hosts its own Foundation Model APIs — pre-deployed models you can query immediately with no infrastructure to manage.

Available models include:

Provider	Models
Meta	Llama 4 Maverick, Llama 3.3 70B, Llama 3.1 405B/8B
Anthropic	Claude Opus 4, Claude Sonnet 4.5, Claude Haiku 4.5
Google	Gemini 2.5 Pro/Flash
OpenAI	GPT-4o, GPT-4o mini

Two pricing models:

Pay-per-token: Billed in DBUs per token. Ideal for experimentation and intermittent workloads. No minimum commitment.
Provisioned throughput: Reserved capacity with guaranteed performance. Recommended for production workloads with predictable traffic.

Query a foundation model the same way as an external endpoint:

1
2
3
4


response = client.chat.completions.create(
    model="databricks-meta-llama-3-3-70b-instruct",
    messages=[{"role": "user", "content": "Explain Delta Lake time travel."}]
)

All AI Gateway features — rate limiting, guardrails, logging, cost attribution — apply uniformly to foundation models. Same governance, whether you’re calling Claude via Anthropic’s API or via Databricks’ hosted instance.

When to Use AI Gateway vs. Direct API Calls
#

Use AI Gateway when…	Use direct API calls when…
Multiple teams share AI model access	You have a single, isolated use case
You need unified cost attribution	Cost tracking per provider is sufficient
Compliance requires prompt/response logging	Logging isn’t a regulatory requirement
You want provider-agnostic failover	You’re committed to a single provider
Security policy mandates PII filtering	Your data pipeline already strips PII
You’re already on Databricks	You have no Databricks footprint

Architecture Pattern: Enterprise AI Governance Stack
#

Here’s how AI Gateway fits into the broader Databricks governance story:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22


┌──────────────────────────────────────────────────────────┐
│                    Applications                           │
│  Chatbots · Agents · Internal Tools · Customer-Facing AI  │
└──────────────────────┬───────────────────────────────────┘
                       │
┌──────────────────────▼───────────────────────────────────┐
│                   AI Gateway                              │
│  Rate Limits · Guardrails · Logging · Cost Attribution    │
└──────┬──────────────┬──────────────┬─────────────────────┘
       │              │              │
┌──────▼──────┐ ┌─────▼──────┐ ┌────▼─────────────────┐
│ Foundation  │ │  External  │ │ Custom Models         │
│ Model APIs  │ │  Providers │ │ (Fine-tuned, RAG)     │
│ (Pay/Token) │ │  (OpenAI,  │ │                       │
│             │ │  Anthropic)│ │                       │
└─────────────┘ └────────────┘ └───────────────────────┘
       │              │              │
┌──────▼──────────────▼──────────────▼─────────────────────┐
│                  Unity Catalog                            │
│  Inference Tables · System Tables · Audit Logs            │
│  Column Masking · Row Security · Access Control           │
└──────────────────────────────────────────────────────────┘

The governance stack is layered:

Unity Catalog governs data access (who can query what)
AI Gateway governs model access (who can call which model, how much, with what content)
Together they create an end-to-end audit trail from data to AI output

Current Limitations
#

AI Guardrails are in Public Preview — safety filtering and PII detection may evolve
Custom Guardrails (bring your own safety model) are in Private Preview
Fallbacks are limited to 2 per endpoint and only supported for external models
PII detection covers U.S. categories only — international PII patterns are not yet supported
TPM rate limiting is not available for custom provider models
Request batch size cannot exceed 16 when guardrails are active
Streaming is not supported with output guardrails for embeddings models
Configuration updates take 20-40 seconds to propagate; rate limit changes may take up to 60 seconds

Conclusion
#

AI Gateway is one of those platform capabilities that seems optional until you need it — and by the time you need it, you’re already in trouble. The runaway cost incident. The compliance audit. The prompt injection that leaked customer data to a third-party model.

The practical value is clear:

Unified access — one API surface for every model provider, no scattered API keys
Cost control — rate limiting and usage attribution before the bill arrives, not after
Compliance — every prompt and response logged to governed Delta tables in Unity Catalog
Safety — guardrails that filter content at the platform level, not the application level
Resilience — automatic fallbacks across providers when primary models are unavailable

If you’re running AI workloads on Databricks — or even considering consolidating your model access — AI Gateway should be your first configuration step, not an afterthought. The governance layer you set up today prevents the incident you’ll otherwise investigate tomorrow.

AI Gateway is GA for serving endpoints across AWS, Azure, and GCP. The new AI Gateway experience (with coding agent integration and MCP server governance) is in Beta. Check the Databricks documentation for the latest features and availability.

Author

Mariusz Derela

Cyber Security Specialist | DevSecOps | AI/ML

Databricks Deep Dive - This article is part of a series.

Part : Agent Bricks: Building Enterprise AI Agents on Databricks

Part : Lakebase: Postgres Meets the Lakehouse — Hands-On with Databricks' OLTP Play

Part : Lakeflow GA: The Unified Data Engineering Platform You've Been Waiting For

Part : MLflow 3: Rebuilt for the GenAI Era

Part : This Article

Part : Serverless GPU & AI Runtime: Running GPU Workloads Without Cluster Management on Databricks

Part : Unity Catalog Deep Dive: Fine-Grained Access Control

Why AI Gateway Matters #

Prerequisites #

Architecture: How AI Gateway Works #

Step 1: Create a Multi-Model Endpoint #

External Model — Anthropic #

External Model — OpenAI (Azure) #

External Model — Amazon Bedrock #

Custom Provider #

Step 2: Query via the OpenAI SDK #

Step 3: Traffic Splitting and Fallbacks #

Traffic Splitting #

Fallbacks #

Step 4: Rate Limiting #

Configure via the Databricks SDK #

Why this matters #

Step 5: AI Guardrails #

Safety Filtering #

PII Detection #

Step 6: Payload Logging and Usage Tracking #

Enable Inference Table Logging #

Usage Tracking via System Tables #

Cost Attribution with Usage Context #

Step 7: Putting It All Together #

Databricks Foundation Models: Pay-Per-Token #

When to Use AI Gateway vs. Direct API Calls #

Architecture Pattern: Enterprise AI Governance Stack #

Current Limitations #

Conclusion #

Related