Skip to content

Add Databricks provider#15

Open
blairj09 wants to merge 2 commits into
mainfrom
feature/databricks-provider
Open

Add Databricks provider#15
blairj09 wants to merge 2 commits into
mainfrom
feature/databricks-provider

Conversation

@blairj09

@blairj09 blairj09 commented Jul 15, 2026

Copy link
Copy Markdown
Collaborator

Adds Databricks as a built-in provider across all three packages, giving Posit Assistant access to models on Databricks Model Serving (foundation models, external models, and custom serving endpoints).

Part of posit-dev/assistant#1110.

ai-provider-bridge

  • providers/databricks-provider.ts — model discovery and chat:
    • Discovery: lists chat-capable endpoints from GET /api/2.0/serving-endpoints (READY + llm/v1/chat, including external-model entities). When the workspace has Unity AI Gateway enabled, discovery switches to GET /api/2.0/serving-endpoints:foundation-models, filtered to gateway-v2 chat-capable models.
    • Gateway detection: a probe of /api/ai-gateway/v2/endpoints is cached per workspace host (definitive 200/404/403 results only; transient failures fall back to model serving without poisoning the cache). The cache is shared between discovery and chat so both route consistently, and clears with the model cache.
    • Chat: OpenAI-compatible chat completions through the existing OpenAIClient against {host}/serving-endpoints, with a fetch wrapper that rewrites requests to {host}/ai-gateway/mlflow/v1 when the gateway is available. Credentials are bearer-token apikey credentials (PAT or OAuth access token — the host application decides which) with the workspace host as baseUrl.
  • model-capabilities/databricks-helpers.ts — capability inference from the underlying model identity (foundation-model / external-model / endpoint name), reusing the Anthropic and OpenAI helpers after stripping databricks- / system.ai. prefixes. Thinking-effort levels are deliberately dropped for v1; unrecognized models get conservative defaults.
  • Registered in PROVIDER_IDS, PROVIDER_MAP (auth provider databricks, apikey), PROVIDER_REGISTRARS, and the public exports. customHeaders are supported on both discovery (additive) and chat paths.

ai-credentials

  • CredentialConfig.getDatabricks() + a Databricks branch in shapeCredentials that resolves the workspace host into the credential baseUrl via a new normalizeDatabricksHost helper (accepts scheme-less hosts and trailing slashes).
  • The Positron backend reads authentication.databricks.credentials.DATABRICKS_HOST with a process.env.DATABRICKS_HOST fallback for managed environments (e.g. Posit Workbench injecting it into sessions).
  • DATABRICKS_TOKEN added to the secret env mappings (store-backend).

ai-config

  • databricks added to BUILTIN_PROVIDER_IDS and CLIENT_KIND_VALUES (identity mapping; not offered as a custom-provider type), keeping the cross-package shape guard green.
  • The workspace host lives in a dedicated strict databricks: { host } connection section — not baseUrl. Per-model endpoint resolution falls back to the provider baseUrl, so storing the host there routed chat to https://<host>/chat/completions (404) and bypassed the gateway rewrite; the section keeps the host available to the credential shaper without entering endpoint resolution. This mirrors how Snowflake stores host/account.
  • The authentication.* host fragment gains a "databricks" read kind (host + customHeaders → the databricks section), and the DATABRICKS_HOST env var overlays the section in the catalog builder.

Testing

  • New unit tests: provider discovery/routing/client factory (23 tests), capability inference, credential shaping, Positron backend host resolution, and the host-fragment read kind; existing ai-config provider counts updated (14 → 15).
  • npm run build, check-types (including the shape guard), test, lint, and format:check all pass.
  • Verified end-to-end in a dev Positron build against a live Databricks workspace: model discovery, and streamed chat through model serving (companion PRs: posit-dev/assistant#1809, Add Databricks authentication provider positron#14898).

Merge ordering

This PR should land before posit-dev/assistant#1809, whose packages/ai-lib submodule pin references this branch.

blairj09 added 2 commits July 15, 2026 10:46
…covery + OpenAI-compatible chat

Ported from posit-dev/ai-provider-bridge feature/databricks-provider and
adapted to the ai-lib monorepo split:

- ai-provider-bridge: databricks provider (gateway probe with per-host
  caching, serving-endpoints and foundation-models discovery, chat URL
  rewriting), capability inference helpers, PROVIDER_IDS/PROVIDER_MAP/
  registrar entries
- ai-credentials: normalizeDatabricksHost, CredentialConfig.getDatabricks
  + apikey shaping branch, vscode credential config reader, DATABRICKS_TOKEN
  env mapping
- ai-config: databricks built-in provider id + client kind (shape guard),
  DATABRICKS_HOST connection env mapping, authentication.databricks host
  fragment read kind
…, not baseUrl

Per-model endpoint resolution falls back to the provider baseUrl, so putting
the workspace host there routed chat requests to https://<host>/chat/completions
(bypassing the serving-endpoints path and the AI Gateway rewrite) and returned
404. Follow the snowflake pattern instead: a strict databricks: { host } section
on the databricks provider block, folded in from
authentication.databricks.credentials and the DATABRICKS_HOST env var, and read
by consumers via connection.databricks.host. Also update the hard-coded
built-in provider counts in ai-config tests (14 -> 15).
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