04 AI Gateway

AI Gateway (Central Governance for LLM + MCP Endpoints)

Databricks AI Gateway is the central control plane between your agents (Apps, notebooks, serving endpoints) and LLM / MCP providers. It gives you a single URL to hit from agent code, and behind that URL you get:

• Usage tracking — per-user, per-token, per-endpoint counts in a UC table.
• Inference tables — full request/response Delta logs for audit.
• Guardrails — content safety, PII detection, topic restrictions.
• Rate limits — per-user / per-endpoint QPS ceilings.
• Fallbacks — route to a second provider on primary failure.

Think of it as the network-level governance layer; MLflow traces + UC OTEL tables are the application-level observability layer. Run both.

When to Add AI Gateway

Add AI Gateway when any of these are true for your agent:

• Multiple teams or apps share LLM endpoints and you need chargeback.
• Compliance / security requires a durable audit trail of model I/O.
• You want rate-limit protection against runaway or abusive callers.
• You need uniform guardrails (PII, safety) applied regardless of which agent is calling.
• You want provider fallback (e.g. route to a secondary model if primary fails).

If the agent is a solo prototype used by its author, MLflow tracing alone is fine — skip this skill until you have ≥ 2 consumers.

Architecture

┌────────────┐        ┌──────────────────────────────┐        ┌──────────────┐
│ Agent App  │───▶────│ Databricks AI Gateway        │───▶────│ LLM / MCP    │
│ (Module 5) │        │  - usage tracking             │        │ Provider     │
│ Serving    │        │  - inference tables (UC Δ)    │        │ (Foundation, │
│ endpoint   │        │  - guardrails / rate limits   │        │  3P, MCP)    │
│ (Module 7) │        │  - fallbacks                  │        └──────────────┘
└────────────┘        └──────────────────────────────┘
      │                           │
      ▼                           ▼
 MLflow traces           UC inference_tables
 (OTEL, app-level)       (provider audit, I/O)

Agents talk to a gateway URL instead of the raw model URL. The gateway transparently proxies OpenAI-compatible and Anthropic-compatible APIs, so existing client code (OpenAI SDK, Anthropic SDK, openai.Agent) keeps working by changing only the base URL.

Consuming the endpoint audit

AI Gateway is downstream of the role-based endpoint binding in vibecoding-state (Phase 1.2). Before creating any gateway endpoint, F4 reads state://endpoint_guardrail_audit (populated by vibecoding-state.endpoint_guardrail_audit) and the bound endpoint under each state://llm_role_endpoints.<role>. The gateway's served_entities[].entity_name MUST equal a candidate that has already passed its role probe — never a name picked by hand.

import yaml

state = yaml.safe_load(open("state.yaml"))
audit  = state["endpoint_guardrail_audit"]
roles  = state["llm_role_endpoints"]

primary = roles["agent_chat"]["endpoint"]
assert audit[primary]["streaming_ok"] is True, (
    f"AI Gateway refusing to bind {primary} as agent_chat primary: "
    "endpoint_guardrail_audit reports streaming_ok != true. "
    "Re-run vibecoding-state.endpoint_guardrail_audit; if no candidate "
    "passes, halt and ask the operator for a workspace-specific FMAPI "
    "candidate list (this workshop has no admin-ticket path)."
)

Halt rule: no admin-ticket fallback

If endpoint_guardrail_audit does not have a passing candidate for a role this gateway needs to serve, halt the F4 step and surface an explicit prompt to the operator:

No FMAPI endpoint passed the role probe for <role>. Provide a workspace-specific candidate list (workspace name + endpoint names) so vibecoding-state.llm_role_endpoint_probe can re-bind. The workshop deliberately has no admin-ticket pathway — gateway creation does not silently downgrade or pick "any healthy" endpoint.

Do not create an admin ticket, do not fall back to a hard-coded default, and do not create the gateway with an unaudited served_entity. The halt is the contract.

Inference-table prefix and grants are F2 territory

The gateway's inference_table_config.table_name_prefix is a provider-level prefix (e.g. gw, no trailing underscore — Databricks appends the per-named-stream suffix). It is not the same as state://otel_table_prefix (which is owned by F2). Capture it under state://Foundation.f4_gateway.inference_table_prefix and grant the same MODIFY, SELECT privileges to the agent SP using the F2 pattern (explicit grants, not ALL_PRIVILEGES).

Validator: gw, not gw_. The same trailing-underscore failure mode that bites OTel table_prefix= also bites the gateway table_name_prefix. Validate that the captured value does not end in _; the platform appends _<name>_payload itself, so a value of gw_ produces gw__<name>_payload (double underscore) and breaks downstream queries.

Creating an AI Gateway Endpoint (CLI)

Use the Databricks CLI to create a gateway in front of an existing serving endpoint:

databricks serving-endpoints create \
  --json '{
    "name": "skyloyalty-ai-gateway",
    "ai_gateway": {
      "usage_tracking_config": {"enabled": true},
      "inference_table_config": {
        "enabled": true,
        "catalog_name": "main",
        "schema_name": "skyloyalty_ops",
        "table_name_prefix": "gw"
      },
      "rate_limits": [
        {"key": "endpoint", "renewal_period": "minute", "calls": 120}
      ],
      "guardrails": {
        "input": {"pii": {"behavior": "BLOCK"}, "safety": true},
        "output": {"pii": {"behavior": "BLOCK"}, "safety": true}
      }
    },
    "config": {
      "served_entities": [{
        "name": "claude-sonnet-46",
        "entity_name": "databricks-claude-sonnet-4-6",
        "entity_version": "1",
        "workload_size": "Small",
        "scale_to_zero_enabled": true
      }]
    }
  }'

After creation, the gateway is reachable at:

https://<workspace-host>/serving-endpoints/skyloyalty-ai-gateway/invocations

or via OpenAI-compatible base:

https://<workspace-host>/serving-endpoints/skyloyalty-ai-gateway/

Inference tables land at main.skyloyalty_ops.gw_<name>_payload (Delta, UC governed) within ~30 min of first traffic.

Pointing Your Agent at the Gateway

OpenAI SDK (most common, e.g. OpenAI Agents SDK, LangChain, custom):

from openai import OpenAI

client = OpenAI(
    base_url="https://<workspace-host>/serving-endpoints/skyloyalty-ai-gateway/",
    api_key=os.environ["DATABRICKS_TOKEN"],
)

response = client.chat.completions.create(
    model="claude-sonnet-46",  # served_entity name from the gateway config
    messages=[...],
)

Model serving client (Databricks SDK):

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
response = w.serving_endpoints.query(
    name="skyloyalty-ai-gateway",
    messages=[...],
)

MLflow tracing continues to work — the gateway is transparent. Your @mlflow.trace decorators and OTEL spans capture the application-level view; inference tables capture the provider-level view. Both are valuable.

Inference Tables (UC Delta Audit)

Inference tables log every request/response through the gateway. Schema highlights:

Correlate with MLflow traces by setting the databricks_request_id header from your agent to the MLflow trace request id:

import mlflow

trace_id = mlflow.get_current_active_span().request_id if mlflow.get_current_active_span() else None
response = client.chat.completions.create(
    model="claude-sonnet-46",
    messages=[...],
    extra_headers={"databricks_request_id": trace_id} if trace_id else None,
)

Now main.skyloyalty_ops.gw_*_payload.databricks_request_id joins to MLflow's trace request_id for end-to-end analysis.

Load references/inference-tables-queries.md if you need SQL recipes for daily spend, top users, error rate analysis, or join-to-MLflow patterns.

Usage Tracking

usage_tracking_config.enabled = true populates system.ai_gateway.usage (or a workspace-scoped UC table, depending on configuration) with per-request token counts, model, and user identity. Use for:

• Weekly chargeback reports per team / app.
• Detecting runaway agents (unexpected spike in tokens per user).
• Capacity planning for PT throughput units.

Sample query:

SELECT
  request_date,
  served_entity_id,
  COUNT(*) AS requests,
  SUM(input_tokens)  AS prompt_tokens,
  SUM(output_tokens) AS completion_tokens
FROM main.skyloyalty_ops.gw_skyloyalty_ai_gateway_usage
WHERE request_date >= current_date() - INTERVAL 7 DAYS
GROUP BY 1, 2
ORDER BY 1 DESC, requests DESC;

Rate Limits

Rate-limit keys:

Example: cap all traffic to 120 calls/min globally AND any single user to 20 calls/min:

"rate_limits": [
  {"key": "endpoint", "renewal_period": "minute", "calls": 120},
  {"key": "user",     "renewal_period": "minute", "calls": 20}
]

When a limit trips, the gateway returns HTTP 429. Agents should backoff + retry; the MLflow trace will show a trace.status=ERROR span with the 429 — use that to feed a production monitoring alert.

Guardrails

Gateway-level guardrails are complementary to in-code guardrails. Use the gateway for org-wide policy (e.g. no PII ever) and in-code guardrails for agent-specific domain rules (e.g. SkyLoyalty never books an award without authorization).

Load references/guardrails-setup.md if you need full config options, per-endpoint override patterns, or guardrail testing recipes.

Fallbacks (Provider Failover)

Configure multiple served_entities so traffic automatically fails over:

"config": {
  "served_entities": [
    {"name": "primary",   "entity_name": "databricks-claude-sonnet-4-6", "traffic_percentage": 100},
    {"name": "secondary", "entity_name": "databricks-claude-opus-4-5",   "traffic_percentage": 0}
  ]
}

Set traffic split or leave secondary at 0% until a failover rule triggers. For agent workloads the common pattern is: primary at 100%; secondary as a capacity reserve for 429/5xx.

MCP Servers via AI Gateway

MCP servers registered with the workspace can also be proxied through the gateway. This gives the same audit/rate-limit/guardrail benefits for tool calls. Configure by adding MCP endpoints as served_entities with MCP-specific payloads. See Databricks MCP docs for the latest payload shape.

Load references/mcp-via-gateway.md if you need the MCP-specific config, including auth-on-behalf-of patterns.

Relation to Other Observability Layers

They do not replace each other. Typical debug path: a user complaint → MLflow trace search → identify bad trace → join databricks_request_id to gateway inference table to see raw provider response → triage.

Do's and Don'ts

Common Mistakes

Validation Checklist

• Gateway endpoint created with usage_tracking_config.enabled=true and inference_table_config.enabled=true.
• Every served_entities[].entity_name corresponds to a candidate that passed its role probe in state://endpoint_guardrail_audit — no hand-picked endpoints.
• If no candidate passed for a required role, F4 halted and asked for a workspace-specific FMAPI candidate list (no admin-ticket fallback was attempted).
• inference_table_config.table_name_prefix does not end in an underscore (validator: gw, not gw_).
• Inference tables visible in UC within 30 min of first traffic.
• Agent app / serving endpoint uses the gateway URL as its base URL.
• databricks_request_id header is passed from client code, tied to MLflow trace.
• At least one rate limit is configured (endpoint or user).
• Gateway guardrails enabled for PII + safety, input + output.
• SQL warehouse access granted on inference / usage tables for the ops team.
• An SQL alert in 07-production-monitoring fires on elevated 429 / 5xx from the gateway.

References

Official documentation (Databricks)

• AI Gateway overview
• Configure AI Gateway endpoints
• Inference tables
• Usage tracking

Local deep-dives

Related skills

02-experiment-tracing-and-uc-storage · sdlc/06-deployment-and-automation · sdlc/07-production-monitoring

Version History