feat: add Databricks AI Gateway as LLM provider#649
feat: add Databricks AI Gateway as LLM provider#649VJ-yadav wants to merge 7 commits intoAltimateAI:mainfrom
Conversation
Implements Databricks serving endpoints support with PAT auth, workspace URL resolution, and OpenAI-compatible request handling. - Add databricks ProviderID to schema - Create auth plugin with PAT parsing and host validation - Add custom loader with env var fallback (DATABRICKS_HOST + DATABRICKS_TOKEN) - Register 11 foundation models (Llama, Claude, GPT, Gemini, DBRX, Mixtral) - Add 24 unit tests for host validation, PAT parsing, body transforms - E2E tests included (skipped without credentials) Closes AltimateAI#602 Co-Authored-By: Vijay Yadav <vjyadav194@gmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
This PR doesn't fully meet our contributing guidelines and PR template. What needs to be fixed:
Please edit this PR description to address the above within 2 hours, or it will be automatically closed. If you believe this was flagged incorrectly, please let a maintainer know. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds Databricks support: a Databricks OAuth/PAT auth plugin, provider loader and model registrations, provider ID schema entry, request-body normalization for Databricks APIs, and unit + optional E2E tests for validation. Changes
Sequence Diagram(s)sequenceDiagram
participant App as App/Client
participant PluginLoader as Plugin Loader
participant AuthSvc as Auth Provider (internal)
participant Provider as Databricks Provider / Fetch
participant Databricks as Databricks API
App->>PluginLoader: initialize plugins
PluginLoader->>AuthSvc: register DatabricksAuthPlugin
AuthSvc-->>PluginLoader: auth provider "databricks" available
App->>Provider: request model call
Provider->>AuthSvc: getAuth("databricks")
alt OAuth present
AuthSvc-->>Provider: { type: "oauth", access, accountId }
Provider->>Provider: set model cost -> 0, use dummy apiKey
Provider->>Provider: wrap fetch to add Authorization header and transform body
Provider->>Databricks: POST /serving-endpoints/... with Bearer token (transformed body)
Databricks-->>Provider: response (stream or json)
Provider-->>App: forwarded response
else No OAuth
Provider->>Env: read DATABRICKS_HOST/DATABRICKS_TOKEN
Env-->>Provider: host + token (or missing)
Provider->>Databricks: POST with API key header
Databricks-->>Provider: response
Provider-->>App: forwarded response
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
packages/opencode/test/altimate/databricks-provider.test.ts (1)
181-227: These E2E tests bypass the opencode provider path.They validate the raw Databricks API, but they won't catch regressions in
CUSTOM_LOADERS.databricks, plugin registration, or thecreateOpenAICompatiblewiring. One smoke test that resolves a Databricks model throughProvider.getLanguage()would cover the actual feature surface.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/opencode/test/altimate/databricks-provider.test.ts` around lines 181 - 227, The tests currently call Databricks directly and don't exercise the opencode provider wiring; add a smoke E2E test that resolves a Databricks model through the provider layer (e.g., call Provider.getLanguage(...) for the databricks model id) to validate CUSTOM_LOADERS.databricks registration and the createOpenAICompatible wiring; specifically, replace or add a test that initializes the provider/registry used in production, ensures the databricks loader/plugin is registered, calls Provider.getLanguage or equivalent factory to obtain a language model client for "databricks-meta-llama-3-1-8b-instruct", then performs a minimal invocation (or stream) and asserts responses—this will catch regressions in CUSTOM_LOADERS.databricks, plugin registration, and createOpenAICompatible wiring.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/opencode/src/provider/provider.ts`:
- Around line 1088-1093: The Databricks provider is being auto-activated via the
generic provider.env bootstrap because database["databricks"] includes env:
["DATABRICKS_TOKEN"]; remove the env entry from the Databricks provider
definition (database["databricks"]) so it no longer participates in the generic
provider.env activation path, and let CUSTOM_LOADERS.databricks (the custom
loader and its autoload flag) be solely responsible for enabling/registering the
provider; ensure no other bootstrap logic treats Databricks as present based
solely on environment variables.
- Around line 740-752: The env-fallback in the databricks provider uses raw
DATABRICKS_HOST without validation; update the databricks branch (the async
databricks provider logic that calls Auth.get("databricks") and
Env.get("DATABRICKS_HOST")) to validate and normalize the host using the same
VALID_HOST_RE used on the OAuth path: reject or strip protocol if present,
ensure it matches VALID_HOST_RE, and only construct
baseURL=`https://${host}/serving-endpoints` when validation passes; if
validation fails, return { autoload: false } (or equivalent) to preserve the
workspace-host guard.
---
Nitpick comments:
In `@packages/opencode/test/altimate/databricks-provider.test.ts`:
- Around line 181-227: The tests currently call Databricks directly and don't
exercise the opencode provider wiring; add a smoke E2E test that resolves a
Databricks model through the provider layer (e.g., call
Provider.getLanguage(...) for the databricks model id) to validate
CUSTOM_LOADERS.databricks registration and the createOpenAICompatible wiring;
specifically, replace or add a test that initializes the provider/registry used
in production, ensures the databricks loader/plugin is registered, calls
Provider.getLanguage or equivalent factory to obtain a language model client for
"databricks-meta-llama-3-1-8b-instruct", then performs a minimal invocation (or
stream) and asserts responses—this will catch regressions in
CUSTOM_LOADERS.databricks, plugin registration, and createOpenAICompatible
wiring.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 2dad8fd4-cedc-468a-a284-6da0b6323232
📒 Files selected for processing (5)
packages/opencode/src/altimate/plugin/databricks.tspackages/opencode/src/plugin/index.tspackages/opencode/src/provider/provider.tspackages/opencode/src/provider/schema.tspackages/opencode/test/altimate/databricks-provider.test.ts
| databricks: async () => { | ||
| const auth = await Auth.get("databricks") | ||
| if (auth?.type !== "oauth") { | ||
| // Fall back to env-based config | ||
| const host = Env.get("DATABRICKS_HOST") | ||
| const token = Env.get("DATABRICKS_TOKEN") | ||
| if (!host || !token) return { autoload: false } | ||
| return { | ||
| autoload: true, | ||
| options: { | ||
| baseURL: `https://${host}/serving-endpoints`, | ||
| apiKey: token, | ||
| }, |
There was a problem hiding this comment.
Validate DATABRICKS_HOST in the env fallback.
The OAuth path applies VALID_HOST_RE, but this branch interpolates raw DATABRICKS_HOST into https://${host}/serving-endpoints. A value like https://adb... or a mistyped domain will silently produce a broken URL and skip the workspace-host guard.
Suggested fix
databricks: async () => {
const auth = await Auth.get("databricks")
if (auth?.type !== "oauth") {
// Fall back to env-based config
const host = Env.get("DATABRICKS_HOST")
const token = Env.get("DATABRICKS_TOKEN")
- if (!host || !token) return { autoload: false }
+ if (!host || !token || !VALID_HOST_RE.test(host)) return { autoload: false }
return {
autoload: true,
options: {
baseURL: `https://${host}/serving-endpoints`,
apiKey: token,📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| databricks: async () => { | |
| const auth = await Auth.get("databricks") | |
| if (auth?.type !== "oauth") { | |
| // Fall back to env-based config | |
| const host = Env.get("DATABRICKS_HOST") | |
| const token = Env.get("DATABRICKS_TOKEN") | |
| if (!host || !token) return { autoload: false } | |
| return { | |
| autoload: true, | |
| options: { | |
| baseURL: `https://${host}/serving-endpoints`, | |
| apiKey: token, | |
| }, | |
| databricks: async () => { | |
| const auth = await Auth.get("databricks") | |
| if (auth?.type !== "oauth") { | |
| // Fall back to env-based config | |
| const host = Env.get("DATABRICKS_HOST") | |
| const token = Env.get("DATABRICKS_TOKEN") | |
| if (!host || !token || !VALID_HOST_RE.test(host)) return { autoload: false } | |
| return { | |
| autoload: true, | |
| options: { | |
| baseURL: `https://${host}/serving-endpoints`, | |
| apiKey: token, | |
| }, |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@packages/opencode/src/provider/provider.ts` around lines 740 - 752, The
env-fallback in the databricks provider uses raw DATABRICKS_HOST without
validation; update the databricks branch (the async databricks provider logic
that calls Auth.get("databricks") and Env.get("DATABRICKS_HOST")) to validate
and normalize the host using the same VALID_HOST_RE used on the OAuth path:
reject or strip protocol if present, ensure it matches VALID_HOST_RE, and only
construct baseURL=`https://${host}/serving-endpoints` when validation passes; if
validation fails, return { autoload: false } (or equivalent) to preserve the
workspace-host guard.
| database["databricks"] = { | ||
| id: ProviderID.databricks, | ||
| source: "custom", | ||
| name: "Databricks", | ||
| env: ["DATABRICKS_TOKEN"], | ||
| options: {}, |
There was a problem hiding this comment.
Don't opt Databricks into the generic provider.env bootstrap.
Line 1271 treats provider.env as “any populated variable is enough”, so DATABRICKS_TOKEN alone materializes this provider. Line 1361 then keeps it alive even when CUSTOM_LOADERS.databricks returns autoload: false, which leaves Databricks registered without any baseURL. Let the custom loader own env activation for this provider.
Suggested fix
database["databricks"] = {
id: ProviderID.databricks,
source: "custom",
name: "Databricks",
- env: ["DATABRICKS_TOKEN"],
+ env: [],
options: {},📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| database["databricks"] = { | |
| id: ProviderID.databricks, | |
| source: "custom", | |
| name: "Databricks", | |
| env: ["DATABRICKS_TOKEN"], | |
| options: {}, | |
| database["databricks"] = { | |
| id: ProviderID.databricks, | |
| source: "custom", | |
| name: "Databricks", | |
| env: [], | |
| options: {}, |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@packages/opencode/src/provider/provider.ts` around lines 1088 - 1093, The
Databricks provider is being auto-activated via the generic provider.env
bootstrap because database["databricks"] includes env: ["DATABRICKS_TOKEN"];
remove the env entry from the Databricks provider definition
(database["databricks"]) so it no longer participates in the generic
provider.env activation path, and let CUSTOM_LOADERS.databricks (the custom
loader and its autoload flag) be solely responsible for enabling/registering the
provider; ensure no other bootstrap logic treats Databricks as present based
solely on environment variables.
|
Thanks @anandgupta42! Will do both:
Joining the Slack now — appreciate the invite! |
|
@VJ-yadav let me know when this PR is ready for review |
|
@anandgupta42 Ready for review. 24 unit tests cover host validation (AWS/Azure/GCP), PAT parsing, credential format, model registration, and request body transforms. E2E tests are structured to skip gracefully without credentials, matching the existing provider pattern. Happy to add Databricks setup docs as a follow-up once this lands. |
2 participants
Summary
Adds Databricks serving endpoints as a first-class LLM provider, closing #602.
src/altimate/plugin/databricks.ts) — PAT-based auth withhost::tokencredential format, workspace host validation for AWS (cloud.databricks.com), Azure (azuredatabricks.net), and GCP (gcp.databricks.com) instancesprovider.ts— resolves workspace URL from PAT auth or env vars (DATABRICKS_HOST+DATABRICKS_TOKEN)max_completion_tokens↔max_tokensfor Databricks OpenAI-compatible APIFollows the same pattern as the Snowflake Cortex provider — uses
@ai-sdk/openai-compatibleSDK since Databricks serving endpoints are OpenAI-compatible.Usage
Test Plan
DATABRICKS_HOST+DATABRICKS_TOKEN)Checklist
@ai-sdk/openai-compatible(no new npm dependencies)// altimate_change start/endcommentsSummary by CodeRabbit
New Features
Tests