Skip to content

Add support for querying Unity Catalog model services - #1584

Open
billiam-wang wants to merge 1 commit into
databricks:mainfrom
billiam-wang:add-model-services-query-helper
Open

billiam-wang wants to merge 1 commit into
databricks:mainfrom
billiam-wang:add-model-services-query-helper

Conversation

@billiam-wang

@billiam-wang billiam-wang commented Sep 25, 2026 •

Copy link
Copy Markdown

What

Adds a ModelServicesExt mixin exposing w.model_services.get_open_ai_client() and get_langchain_chat_open_ai_client(), so users can query model services (model APIs) through Unity Gateway. A model service is a Unity Catalog securable that represents a governed LLM endpoint. The returned client targets {host}/ai-gateway/mlflow/v1 with refreshed Databricks bearer auth; the model service's Unity Catalog three-level name (e.g. system.ai.claude-sonnet-4-5) is passed as model. This mirrors the existing ServingEndpointsExt query surface.

Usage

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
client = w.model_services.get_open_ai_client()
resp = client.chat.completions.create(
    model="system.ai.claude-sonnet-4-5",
    messages=[{"role": "user", "content": "Hello!"}],
    max_tokens=256,
)
print(resp.choices[0].message.content)

Note for maintainers

w.model_services has no backing generated API, so the wiring in the generated databricks/sdk/__init__.py was added by hand. It likely needs to move into the SDK generator config to survive regeneration.

Testing

Added tests/test_model_services_mixin.py mirroring test_open_ai_mixin.py (base URL, reserved-param guard, bearer-auth wiring). Verified locally that the client resolves to /ai-gateway/mlflow/v1 and the guard raises.

This pull request and its description were written by Isaac.

Adds a ModelServicesExt mixin so w.model_services.get_open_ai_client()
queries Unity Catalog v3 system.ai.* model services via the AI Gateway,
mirroring the existing ServingEndpointsExt query surface.

Co-authored-by: Isaac <no-reply@databricks.com>
Signed-off-by: Bill Wang <bill.wang@databricks.com>
@billiam-wang
billiam-wang force-pushed the add-model-services-query-helper branch from cb19f34 to 4fa28d2 Compare September 28, 2026 21:34

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant