unity-ai-gateway-model-service

name

unity-ai-gateway-model-service

description

Create, configure, query, and govern Databricks Unity AI Gateway "model services" (UC securables at catalog.schema.service that front FMAPI models), and wire a deployed Mosaic AI agent endpoint to call one. Covers the live 2.1 REST API (create-or-get, routing/traffic-split, rate limits, inference tables, guardrails/service-policies, usage tracking), OpenAI-compatible + ChatDatabricks querying, observability via system tables, and the dedicated-service-principal "dance" required for deployed agents to authenticate to a model service. Use when working with Unity AI Gateway, model services, /ai-gateway/mlflow/v1, use_ai_gateway=True, or a deployed agent that 404s calling a gateway model service.

Unity AI Gateway — model service management & agent auth

Verified live on a Beta workspace, 2026-07-06. Code samples are in resources/.

What a "model service" is

A Unity Catalog securable at catalog.schema.<service> (a "Service" object alongside tables/functions/models) that fronts one or more FMAPI foundation models behind governed routing, rate limits, inference logging, guardrails, and usage tracking. Callers invoke it by its fully-qualified name over an OpenAI-compatible API and inherit all governance. Governed with normal UC privileges (USE CATALOG/USE SCHEMA/EXECUTE). system.ai.* ships one pre-provisioned service per Databricks-hosted model (non-deletable).

API surface — 2.1 (live) vs 2.2 (docs)

• 2.1 is what's live on Beta rings today — AIP/Google-style: parent + model_service_id query params, a config{} body, and PATCH with update_mask. Use this. No dedicated Python SDK / Terraform / CREATE SERVICE SQL yet → drive it with WorkspaceClient().api_client.do(...).
• 2.2 appears in public AWS docs (flat body: catalog_name/schema_name/routes/rate_limits{tpm,qpm}) but 404s on Beta workspaces — it's the forward/Public-Preview shape. Don't target it until it's on your ring.
• SQL supports only GRANT EXECUTE ON MODEL SERVICE … / DROP MODEL SERVICE ….

Base path: /api/2.1/unity-catalog/model-services. See resources/manage_model_service.py for a full CRUD + config module.

CRUD (all verified)

Op	Call
List	GET /api/2.1/unity-catalog/model-services
Get	GET .../model-services/{cat}.{sch}.{name}
Create	POST .../model-services?parent=schemas/{cat}.{sch}&model_service_id={name} body {"config":{"destinations":[…]}}
Update	PATCH .../model-services/{cat}.{sch}.{name}?update_mask=config.<field>
Delete	DELETE .../model-services/{cat}.{sch}.{name}

parent and model_service_id are query params (not body). model_service_id can't contain ./space///%.

Config fields (exact shapes)

• Destination (a fronted model): {"name":"system.ai.<model>","type":"DESTINATION_TYPE_PAY_PER_TOKEN_FOUNDATION_MODEL","traffic_percentage":N,"pay_per_token_config":{"model":"models/system.ai.<model>"}}. Note the resource-name form models/system.ai.<model>.
• Traffic split — multiple destinations whose traffic_percentage sum to 100; PATCH update_mask=config.destinations. Settable via API. Max 5 destinations.
• Fallback chain — nested config.routing.fallback.destinations[] (each traffic_percentage:0), used on 429/5xx. UI-only in Beta — update_mask=config.routing.fallback / config.fallback / config.routing all reject with "not yet supported by UpdateModelService". Configure on the service's Routing tab.
• Rate limits — config.rate_limits[] (array): {"key":"RATE_LIMIT_KEY_SERVICE","renewal_period":"RATE_LIMIT_RENEWAL_PERIOD_MINUTE","requests":"100"}. Separate entries for requests (QPM) vs tokens (TPM); string values. RATE_LIMIT_KEY_USER requires a principal. Max 20/service, 5 group-specific.
• Inference table (object, not a string): {"parent":"schemas/{cat}.{sch}","table_name_prefix":"<prefix>","enabled":true} → auto-creates <prefix>_payload. Needs CREATE TABLE. Read-back adds table/table_id.
• Guardrails — config.service_policies[]: {"name":"…","policy_type":"POLICY_TYPE_BUILTIN","handler":"system.ai.block_jailbreak","rank":1,"options":{"model_service":"model-services/system.ai.gpt-5-nano","phases":"pre_call","dry_run":"false"}}. A small judge model evaluates each request. POLICY_TYPE_CUSTOM handler = a UC UDF functions/{cat}.{sch}.{fn}. Built-ins: jailbreak (system.ai.block_jailbreak verified), PII redact/block, unsafe content, hallucination.
• config.usage_tracking.enabled is mandatory — PATCHing it false → 400. config.tracing.enabled on by default.

Granting access — use the permissions REST API (avoids DDL ambiguity)

PATCH /api/2.1/unity-catalog/permissions/{securable_type}/{full_name} with {"changes":[{"principal":"<sp-app-id-or-user>","add":["<PRIV>"]}]}. Securable types: catalog (USE_CATALOG), schema (USE_SCHEMA), model_service (EXECUTE). This sidesteps the GRANT EXECUTE ON SERVICE vs ON MODEL SERVICE keyword uncertainty.

Querying a model service

• OpenAI-compatible: OpenAI(api_key=<token>, base_url="https://<host>/ai-gateway/mlflow/v1"), model="<catalog.schema.service>". Tool-calling works through the gateway.
• LangChain: ChatDatabricks(model="<fqn>", use_ai_gateway=True). Add workspace_client=<SP client> to authenticate as a specific SP (see below). ChatDatabricks also has use_ai_gateway_native_api and a profile field.
• Do NOT pass temperature to Claude/Sonnet-5 through the gateway — it 400s (does not support the temperature parameter).

Observability

• system.ai_gateway.usage — tokens, latency, TTFT, requester, status_code, endpoint_name (the model-service FQN). Account-admin read only. Rejected/NOT_FOUND calls are not logged here.
• system.billing.usage — DBU cost with ai_gateway.* metadata.
• The inference _payload table — request/response payloads (aggregated every few minutes).
• Budgets live in the Account Console (not a service API field). Lineage isn't in Beta.

---

The dedicated-SP "dance" (deployed agent → model service auth)

Problem: a Mosaic AI agent deployed via agents.deploy() runs under a hidden, Databricks-managed system service principal whose automatic-auth credentials are scoped to declared mlflow.models.resources. There is no DatabricksModelService resource type, so that SP can't see the model service and the deployed agent's LLM call fails with 404 NOT_FOUND: '<fqn>' does not exist' (UC hides securables the caller lacks USE CATALOG/USE SCHEMA/EXECUTE on). The hidden SP is undiscoverable/ungrantable, and its rejected call leaves no system.ai_gateway.usage row. (Deployed agents aren't yet first-class Unity AI Gateway callers — "coming next quarter" per the internal FAQ.)

Fix (verified end-to-end): authenticate the LLM hop as a dedicated, grantable SP — scoped to only that hop.

1. Create/reuse a dedicated SP and store its OAuth creds in a secret scope. Workspace-admin SCIM is sufficient (no account admin): w.service_principals.create(...) + w.service_principal_secrets_proxy.create(service_principal_id=scim_id).secret. See resources/agent_sp_auth.py.
2. Grant it USE_CATALOG + USE_SCHEMA + EXECUTE on the model service via the permissions REST API (above).
3. Pass an SP-authed client only to the LLM: ChatDatabricks(model=fqn, use_ai_gateway=True, workspace_client=WorkspaceClient(host=…, client_id=SP_CLIENT_ID, client_secret=SP_CLIENT_SECRET, auth_type="oauth-m2m")). Auth is per-client, so only the LLM runs as the SP; UCFunctionToolkit tools keep the endpoint's automatic passthrough (leave their DatabricksFunction resources declared).
4. Deploy with the creds injected as secret-backed environment_vars, under SP_CLIENT_ID/SP_CLIENT_SECRET — NOT DATABRICKS_CLIENT_ID/SECRET (those are auto-detected by the SDK default auth chain and would make the whole agent run as the SP, disabling tool passthrough).
5. Probe with retry after granting — grants + service reconfiguration take a few seconds to propagate to the gateway; a fresh grant + immediate invoke can transiently 404. Retry a tiny SP invoke until it succeeds before relying on it.

Deploy gotchas (all hit during validation)

• pip_requirements must be explicit. MLflow models-from-code auto-inference drops databricks-langchain/langgraph/unitycatalog-ai → "missing Python dependency" load failure. List them.
• Pin the langgraph pair. An unpinned --upgrade/inference leaves langgraph core and langgraph-prebuilt skewed → ImportError: cannot import name 'ExecutionInfo' from 'langgraph.runtime'. Pin langgraph==1.2.6 + langgraph-prebuilt==1.1.0 (a verified-compatible pair) in BOTH the notebook %pip and the served model's pip_requirements.
• deploy_feedback_model=False — the auto-deployed feedback model is an extra served entity that can intermittently block a multi-model endpoint update (known gunicorn/mlflow-3.10+ deadlock family; platform-fixed 2026-04-24 — re-log rather than pin mlflow to an old 3.9).
• Stuck subsequent deploys: a first deploy to a fresh endpoint usually succeeds; a subsequent in-place update to a wedged endpoint can hang → delete + recreate, or deploy under a fresh endpoint name.

Environment / Beta notes

• Enable via Account Console → Previews. Not on GovCloud/CSP/HIPAA yet; GCP has no inference tables yet.
• Destinations are Databricks-hosted pay-per-token FMAPI models only in Beta (no PT / external providers).