04-ai-gateway

name: 04-ai-gateway description: > Use when you need central governance, observability, guardrails, or rate limits across LLM endpoints and MCP servers used by agents. Covers creating and configuring Databricks AI Gateway endpoints, enabling usage tracking, inference tables (UC Delta audit), rate limits, and guardrails, plus how deployed agent serving endpoints and Apps consume a gateway URL. Foundation Step 4 — ideal as a governance layer BEFORE Module 7 production deployment. Complements in-code guardrails and MLflow tracing; does NOT replace MLflow trace tables but adds a provider-level audit trail. license: Apache-2.0 compatibility: "Requires Databricks AI Gateway (public preview in most AWS/Azure regions as of 2026-04). Requires workspace Unity Catalog and SQL Warehouse for inference tables." clients: [ide_cli, genie_code] bundle_resource: none deploy_verb: bundle_deploy deploy_note: "AI Gateway endpoint + inference/usage tables created via the Databricks SDK/REST (or CLI through runDatabricksCli on Genie Code); inference tables land in the per-user prefixed schema. Not modeled as a DAB resource kind in this workshop. See skills/genie-code-environment." coverage: full metadata: last_verified: "2026-06-05" volatility: high upstream_sources: [] author: "prashanth-subrahmanyam" version: "1.1.0" domain: "genai-agents" pipeline_position: "F4" consumes: "llm_endpoints, mcp_endpoints, uc_catalog_schema, endpoint_guardrail_audit, llm_role_endpoints" produces: "ai_gateway_endpoint, inference_tables, usage_tracking_tables, rate_limits, inference_table_prefix" grounded_in: "https://docs.databricks.com/aws/en/ai-gateway/, https://docs.databricks.com/aws/en/ai-gateway/configure-ai-gateway-endpoints, https://docs.databricks.com/aws/en/ai-gateway/inference-tables, https://docs.databricks.com/aws/en/ai-gateway/usage-tracking"

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