-
Notifications
You must be signed in to change notification settings - Fork 41
feat: add Databricks AI Gateway as LLM provider #649
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
262f9b5
fd082ce
4fd436c
822a92c
15f779a
f3ec481
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,125 @@ | ||
| import type { Hooks, PluginInput } from "@opencode-ai/plugin" | ||
| import { Auth, OAUTH_DUMMY_KEY } from "@/auth" | ||
|
|
||
| /** | ||
| * Databricks workspace host regex. | ||
| * Matches patterns like: myworkspace.cloud.databricks.com, adb-1234567890.12.azuredatabricks.net | ||
| */ | ||
| export const VALID_HOST_RE = /^[a-zA-Z0-9._-]+\.(cloud\.databricks\.com|azuredatabricks\.net|gcp\.databricks\.com)$/ | ||
|
|
||
| /** Parse a `host::token` credential string for Databricks PAT auth. */ | ||
| export function parseDatabricksPAT(code: string): { host: string; token: string } | null { | ||
| const sep = code.indexOf("::") | ||
| if (sep === -1) return null | ||
| const host = code.substring(0, sep).trim() | ||
| const token = code.substring(sep + 2).trim() | ||
| if (!host || !token) return null | ||
| if (!VALID_HOST_RE.test(host)) return null | ||
| return { host, token } | ||
| } | ||
|
|
||
| /** | ||
| * Transform a Databricks request body string. | ||
| * Databricks Foundation Model APIs use max_tokens (OpenAI-compatible), | ||
| * but some endpoints may prefer max_completion_tokens. | ||
| */ | ||
| export function transformDatabricksBody(bodyText: string): { body: string } { | ||
| const parsed = JSON.parse(bodyText) | ||
|
|
||
| // Databricks uses max_tokens for most endpoints, but some newer ones | ||
| // expect max_completion_tokens. Normalize to max_tokens for compatibility. | ||
| if ("max_completion_tokens" in parsed && !("max_tokens" in parsed)) { | ||
| parsed.max_tokens = parsed.max_completion_tokens | ||
| delete parsed.max_completion_tokens | ||
| } | ||
|
|
||
| return { body: JSON.stringify(parsed) } | ||
| } | ||
|
|
||
| export async function DatabricksAuthPlugin(_input: PluginInput): Promise<Hooks> { | ||
| return { | ||
| auth: { | ||
| provider: "databricks", | ||
| async loader(getAuth, provider) { | ||
| const auth = await getAuth() | ||
| if (auth.type !== "oauth") return {} | ||
|
|
||
| for (const model of Object.values(provider.models)) { | ||
| model.cost = { input: 0, output: 0, cache: { read: 0, write: 0 } } | ||
| } | ||
|
|
||
| return { | ||
| apiKey: OAUTH_DUMMY_KEY, | ||
| async fetch(requestInput: RequestInfo | URL, init?: RequestInit) { | ||
| const currentAuth = await getAuth() | ||
| if (currentAuth.type !== "oauth") return fetch(requestInput, init) | ||
|
|
||
| const headers = new Headers() | ||
| if (init?.headers) { | ||
| if (init.headers instanceof Headers) { | ||
| init.headers.forEach((value, key) => headers.set(key, value)) | ||
| } else if (Array.isArray(init.headers)) { | ||
| for (const [key, value] of init.headers) { | ||
| if (value !== undefined) headers.set(key, String(value)) | ||
| } | ||
| } else { | ||
| for (const [key, value] of Object.entries(init.headers)) { | ||
| if (value !== undefined) headers.set(key, String(value)) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| headers.set("authorization", `Bearer ${currentAuth.access}`) | ||
|
|
||
| let body = init?.body | ||
| if (body) { | ||
| try { | ||
| let text: string | ||
| if (typeof body === "string") { | ||
| text = body | ||
| } else if (body instanceof Uint8Array || body instanceof ArrayBuffer) { | ||
| text = new TextDecoder().decode(body) | ||
| } else { | ||
| text = "" | ||
| } | ||
| if (text) { | ||
| const result = transformDatabricksBody(text) | ||
| body = result.body | ||
| headers.delete("content-length") | ||
| } | ||
| } catch { | ||
| // JSON parse error — pass original body through untransformed | ||
| } | ||
| } | ||
|
|
||
| return fetch(requestInput, { ...init, headers, body }) | ||
| }, | ||
| } | ||
| }, | ||
| methods: [ | ||
| { | ||
| label: "Databricks PAT", | ||
| type: "oauth", | ||
| authorize: async () => ({ | ||
| url: "https://accounts.cloud.databricks.com", | ||
| instructions: | ||
| "Enter your credentials as: <workspace-host>::<PAT-token>\n e.g. myworkspace.cloud.databricks.com::dapi1234567890abcdef\n Create a PAT in Databricks: Settings → Developer → Access Tokens → Generate New Token", | ||
| method: "code" as const, | ||
| callback: async (code: string) => { | ||
| const parsed = parseDatabricksPAT(code) | ||
| if (!parsed) return { type: "failed" as const } | ||
| return { | ||
| type: "success" as const, | ||
| access: parsed.token, | ||
| refresh: "", | ||
| // Databricks PATs can be configured with custom TTLs; use 90-day default | ||
| expires: Date.now() + 90 * 24 * 60 * 60 * 1000, | ||
| accountId: parsed.host, | ||
| } | ||
| }, | ||
| }), | ||
| }, | ||
| ], | ||
| }, | ||
| } | ||
| } |
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -50,6 +50,9 @@ import { ModelID, ProviderID } from "./schema" | |||||||||||||||||||||||||
| // altimate_change start — snowflake cortex account validation | ||||||||||||||||||||||||||
| import { VALID_ACCOUNT_RE } from "../altimate/plugin/snowflake" | ||||||||||||||||||||||||||
| // altimate_change end | ||||||||||||||||||||||||||
| // altimate_change start — databricks host validation | ||||||||||||||||||||||||||
| import { VALID_HOST_RE } from "../altimate/plugin/databricks" | ||||||||||||||||||||||||||
| // altimate_change end | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| const DEFAULT_CHUNK_TIMEOUT = 120_000 | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
|
|
@@ -733,6 +736,32 @@ export namespace Provider { | |||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||
| }, | ||||||||||||||||||||||||||
| // altimate_change end | ||||||||||||||||||||||||||
| // altimate_change start — databricks provider loader | ||||||||||||||||||||||||||
| 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, | ||||||||||||||||||||||||||
| }, | ||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||
| const host = auth.accountId ?? Env.get("DATABRICKS_HOST") | ||||||||||||||||||||||||||
| if (!host || !VALID_HOST_RE.test(host)) return { autoload: false } | ||||||||||||||||||||||||||
| return { | ||||||||||||||||||||||||||
| autoload: true, | ||||||||||||||||||||||||||
| options: { | ||||||||||||||||||||||||||
| baseURL: `https://${host}/serving-endpoints`, | ||||||||||||||||||||||||||
| }, | ||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||
| }, | ||||||||||||||||||||||||||
| // altimate_change end | ||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| export const Model = z | ||||||||||||||||||||||||||
|
|
@@ -1019,6 +1048,70 @@ export namespace Provider { | |||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||
| // altimate_change end | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| // altimate_change start — databricks provider models | ||||||||||||||||||||||||||
| function makeDatabricksModel( | ||||||||||||||||||||||||||
| id: string, | ||||||||||||||||||||||||||
| name: string, | ||||||||||||||||||||||||||
| limits: { context: number; output: number }, | ||||||||||||||||||||||||||
| caps?: { reasoning?: boolean; attachment?: boolean; toolcall?: boolean; image?: boolean }, | ||||||||||||||||||||||||||
| ): Model { | ||||||||||||||||||||||||||
| const m: Model = { | ||||||||||||||||||||||||||
| id: ModelID.make(id), | ||||||||||||||||||||||||||
| providerID: ProviderID.databricks, | ||||||||||||||||||||||||||
| api: { | ||||||||||||||||||||||||||
| id, | ||||||||||||||||||||||||||
| url: "", | ||||||||||||||||||||||||||
| npm: "@ai-sdk/openai-compatible", | ||||||||||||||||||||||||||
| }, | ||||||||||||||||||||||||||
| name, | ||||||||||||||||||||||||||
| capabilities: { | ||||||||||||||||||||||||||
| temperature: true, | ||||||||||||||||||||||||||
| reasoning: caps?.reasoning ?? false, | ||||||||||||||||||||||||||
| attachment: caps?.attachment ?? false, | ||||||||||||||||||||||||||
| toolcall: caps?.toolcall ?? true, | ||||||||||||||||||||||||||
| input: { text: true, audio: false, image: caps?.image ?? false, video: false, pdf: false }, | ||||||||||||||||||||||||||
| output: { text: true, audio: false, image: false, video: false, pdf: false }, | ||||||||||||||||||||||||||
| interleaved: false, | ||||||||||||||||||||||||||
| }, | ||||||||||||||||||||||||||
| cost: { input: 0, output: 0, cache: { read: 0, write: 0 } }, | ||||||||||||||||||||||||||
| limit: { context: limits.context, output: limits.output }, | ||||||||||||||||||||||||||
| status: "active" as const, | ||||||||||||||||||||||||||
| options: {}, | ||||||||||||||||||||||||||
| headers: {}, | ||||||||||||||||||||||||||
| release_date: "2024-01-01", | ||||||||||||||||||||||||||
| variants: {}, | ||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||
| m.variants = mapValues(ProviderTransform.variants(m), (v) => v) | ||||||||||||||||||||||||||
| return m | ||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| database["databricks"] = { | ||||||||||||||||||||||||||
| id: ProviderID.databricks, | ||||||||||||||||||||||||||
| source: "custom", | ||||||||||||||||||||||||||
| name: "Databricks", | ||||||||||||||||||||||||||
| env: ["DATABRICKS_TOKEN"], | ||||||||||||||||||||||||||
| options: {}, | ||||||||||||||||||||||||||
|
Comment on lines
+1088
to
+1093
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Don't opt Databricks into the generic Line 1271 treats Suggested fix database["databricks"] = {
id: ProviderID.databricks,
source: "custom",
name: "Databricks",
- env: ["DATABRICKS_TOKEN"],
+ env: [],
options: {},📝 Committable suggestion
Suggested change
🤖 Prompt for AI Agents |
||||||||||||||||||||||||||
| models: { | ||||||||||||||||||||||||||
| // Meta Llama models — tool calling supported | ||||||||||||||||||||||||||
| "databricks-meta-llama-3-1-405b-instruct": makeDatabricksModel("databricks-meta-llama-3-1-405b-instruct", "Meta Llama 3.1 405B Instruct", { context: 128000, output: 4096 }), | ||||||||||||||||||||||||||
| "databricks-meta-llama-3-1-70b-instruct": makeDatabricksModel("databricks-meta-llama-3-1-70b-instruct", "Meta Llama 3.1 70B Instruct", { context: 128000, output: 4096 }), | ||||||||||||||||||||||||||
| "databricks-meta-llama-3-1-8b-instruct": makeDatabricksModel("databricks-meta-llama-3-1-8b-instruct", "Meta Llama 3.1 8B Instruct", { context: 128000, output: 4096 }), | ||||||||||||||||||||||||||
| // Claude models via Databricks AI Gateway | ||||||||||||||||||||||||||
| "databricks-claude-sonnet-4-6": makeDatabricksModel("databricks-claude-sonnet-4-6", "Claude Sonnet 4.6", { context: 200000, output: 64000 }), | ||||||||||||||||||||||||||
| "databricks-claude-opus-4-6": makeDatabricksModel("databricks-claude-opus-4-6", "Claude Opus 4.6", { context: 200000, output: 32000 }), | ||||||||||||||||||||||||||
| // GPT models via Databricks AI Gateway | ||||||||||||||||||||||||||
| "databricks-gpt-5-4": makeDatabricksModel("databricks-gpt-5-4", "GPT-5-4", { context: 128000, output: 16384 }), | ||||||||||||||||||||||||||
| "databricks-gpt-5-mini": makeDatabricksModel("databricks-gpt-5-mini", "GPT-5 Mini", { context: 128000, output: 16384 }), | ||||||||||||||||||||||||||
| // Gemini models via Databricks AI Gateway | ||||||||||||||||||||||||||
| "databricks-gemini-3-1-pro": makeDatabricksModel("databricks-gemini-3-1-pro", "Gemini 3.1 Pro", { context: 1000000, output: 8192 }), | ||||||||||||||||||||||||||
| // DBRX — Databricks native model | ||||||||||||||||||||||||||
| "databricks-dbrx-instruct": makeDatabricksModel("databricks-dbrx-instruct", "DBRX Instruct", { context: 32768, output: 4096 }), | ||||||||||||||||||||||||||
| // Mixtral via Databricks | ||||||||||||||||||||||||||
| "databricks-mixtral-8x7b-instruct": makeDatabricksModel("databricks-mixtral-8x7b-instruct", "Mixtral 8x7B Instruct", { context: 32768, output: 4096 }, { toolcall: false }), | ||||||||||||||||||||||||||
| }, | ||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||
| // altimate_change end | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| // altimate_change start — register altimate-backend as an OpenAI-compatible provider | ||||||||||||||||||||||||||
| if (!database["altimate-backend"]) { | ||||||||||||||||||||||||||
| const backendModels: Record<string, Model> = { | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Validate
DATABRICKS_HOSTin the env fallback.The OAuth path applies
VALID_HOST_RE, but this branch interpolates rawDATABRICKS_HOSTintohttps://${host}/serving-endpoints. A value likehttps://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
🤖 Prompt for AI Agents