Skip to content

feat(gemini): route Gemini CLI through a Model Provider Service - #417

Open
variablenerd wants to merge 2 commits into
databricks:mainfrom
variablenerd:feat/gemini-model-provider-service
Open

feat(gemini): route Gemini CLI through a Model Provider Service#417
variablenerd wants to merge 2 commits into
databricks:mainfrom
variablenerd:feat/gemini-model-provider-service

Conversation

@variablenerd

@variablenerd variablenerd commented Aug 30, 2026

Copy link
Copy Markdown

Users need to access Model Provider Services in Unity Gateway via ucode and Gemini CLI. --provider routing existed for Claude and Codex but not Gemini, so a Gemini Enterprise Model Provider Service registered in Unity Gateway was unreachable.

This extends the existing --provider path to Gemini:

  • databricks.py — add gemini_enterprise to the tool→provider-type map so the service resolves for Gemini.
  • agents/gemini.py — emit the Databricks-Model-Provider-Service header (via GEMINI_CLI_CUSTOM_HEADERS) and pin the service's target model; re-emitted on token refresh.
  • agents/__init__.pyresolve_gemini_provider_model picks the target (--model when the service declares several); pass the provider through configure_tool.
  • cli.py--provider/--model on ucode gemini; ucode configure offers Gemini the provider picker.

Commands:

# route Gemini CLI through a Gemini MPS (uses its target model)
ucode gemini --provider {catalog}.{schema}.{mps-name}

# pick a model when the service declares several targets
ucode gemini --provider {catalog}.{schema}.{mps-name} --model gemini-3.5-flash

# persist it as the default
ucode configure          # choose Gemini, then select the provider service
ucode gemini

Verified end-to-end against a Gemini Enterprise MPS: the gateway returns HTTP 200 for POST /ai-gateway/gemini/v1beta/models/<target>:generateContent with the Databricks-Model-Provider-Service header.

Tests: uv run pytest (added focused tests in test_agent_gemini.py, test_agents_init.py, test_databricks.py), uv run ruff check ., uv run ruff format --check.

`ucode gemini` could only run Databricks-hosted models. Provider routing
(`--provider`) was wired for claude and codex but not gemini, so a Gemini
Enterprise Model Provider Service registered in Unity Catalog was unreachable.

Extend the existing `--provider` path to gemini: add `gemini_enterprise` to the
tool/provider-type map, emit the `Databricks-Model-Provider-Service` header via
`GEMINI_CLI_CUSTOM_HEADERS`, and pin the service's target model (chosen with
`--model` when it declares several). The routing header is re-emitted on token
refresh, and `ucode configure` offers gemini the provider picker.

Co-authored-by: Isaac <no-reply@databricks.com>
@variablenerd
variablenerd force-pushed the feat/gemini-model-provider-service branch from 156a672 to e852953 Compare August 30, 2026 18:01
Keep one test per new code path: the MPS routing header, target-model
resolution (happy + not-a-target), and gemini_enterprise resolving for gemini.

Co-authored-by: Isaac <no-reply@databricks.com>
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