unity-ai-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).