[mason] Add a model picker to the local chat UI (Unity Catalog AI Gateway) - #547
Open
ankit-db wants to merge 3 commits into
Open
[mason] Add a model picker to the local chat UI (Unity Catalog AI Gateway)#547ankit-db wants to merge 3 commits into
ankit-db wants to merge 3 commits into
Conversation
…eway) The demo chat app hardcoded the agent's model, so trying a different model locally meant editing code and restarting. Add a model picker to the chat header for both templates (langgraph + openai), backed by the Unity Catalog AI Gateway. - New `databricks_mason.runtime.models.list_ai_gateway_models(client)` lists the workspace's chat-capable AI Gateway models — the `system.ai.*` model services from `GET /api/2.1/unity-catalog/model-services?parent=schemas/system.ai` (embeddings-only services filtered out). Re-exported from `databricks_mason` and `databricks_mason.runtime`; unit-tested in the CI-covered CLI suite. - `GET /api/demo/models` in each template's `runtime/ui.py` returns the agent's default model (always first, always present) plus those gateway models, falling back to just the default when `system.ai` isn't readable. - Agents call through the AI Gateway (`<host>/ai-gateway/mlflow/v1`): langgraph via `ChatDatabricks(..., use_ai_gateway=True)`, openai via `AsyncDatabricksOpenAI(use_ai_gateway=True)`. The default `MODEL` is now a `system.ai.*` model; `create_agent_graph`/`create_agent` take an optional model override and `stream_handler` forwards the request body's `model`, so sync, streaming, and background invocations all honor the picked model. - Frontend: a `#model-select` in the chat header, populated from `/api/demo/models`, sends the chosen `system.ai.*` model as `model` in each invocation body; disabled while a run is in flight; default option labelled `(default)`. - Docs (`CHAT_APP.md`) and tests updated for both templates. Reported in the Custom Agents bug bash. Co-authored-by: Isaac <no-reply@databricks.com>
Point the default MODEL (and the picker's default option) at the Opus 4.8 AI Gateway model in both templates. The picker still lists whatever system.ai chat models the workspace exposes, so this only changes the out-of-the-box default. Co-authored-by: Isaac <no-reply@databricks.com>
| # gateway (see `use_ai_gateway=True` below). Swap for any `system.ai.*` model your workspace exposes | ||
| # — the demo chat app's picker lists what's available. `mason dev`'s chat UI can also override this | ||
| # per request. | ||
| MODEL = "system.ai.claude-opus-4-8" |
Contributor
There was a problem hiding this comment.
When I tested claude with langgraph, I saw some strange interaction with claude's reasoning and the UI rendering, with errors like
PydanticSerializationUnexpectedValue(Expected `str` - serialized value may not be as expected [field_name='content', input_value=[{'type': 'reasoning', 's... '', 'signature': ''}]}], input_type=list])
return self.__pydantic_serializer__.to_python(
and some rendering issues in the chat UI too.
Did you end up seeing this too?
Author
There was a problem hiding this comment.
Thanks for flagging — this is directly relevant now: after merging main + switching to the UC AI Gateway, the default MODEL is system.ai.claude-opus-4-8, so Claude is what a fresh scaffold hits first. That PydanticSerializationUnexpectedValue comes from Claude returning content as a list of blocks (incl. reasoning), which _serialize_events' msg.model_dump() and the UI's text extraction don't special-case. I'd rather fix it than ship a Claude default that renders oddly: I'm planning a focused follow-up that strips/handles reasoning blocks in _serialize_events (and the UI's extractText) so only user-visible text renders. Did your testing show a clean way to keep the reasoning out of the rendered turn, or just drop those blocks? Leaving this thread open until that's in.
Contributor
|
whoops my PR #538 got merged |
Main independently added a serving-endpoint model picker (alongside the actor refactor). Reconciled by keeping main's picker structure (actor plumbing, config-embedded `models`, frontend picker, `_discover_chat_models`) and routing it through the Unity Catalog AI Gateway instead of serving endpoints: - agent.py (both): default MODEL = system.ai.claude-opus-4-8; ChatDatabricks / AsyncDatabricksOpenAI built with use_ai_gateway=True. - ui.py (both): _discover_chat_models now lists the workspace's system.ai models via list_ai_gateway_models (UC model-services) instead of serving_endpoints.list(). - Kept the new databricks_mason.runtime.models helper (+ re-exports, unit tests); added transient-failure retry to its list call (review comment). - Updated demo-UI / models unit tests and CHAT_APP docs for the gateway. Co-authored-by: Isaac <no-reply@databricks.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds a model picker to the local chat UI (both templates) so you can switch models without editing code and restarting.
system.ai.*) and routes the agent through the gateway (use_ai_gateway=True).system.ai.claude-opus-4-8; pick any listed model per turn from the chat header.