Skip to content

Commit 3f1ca5e

Browse files
committed
Add Databricks provider (AI Gateway / Foundation Model APIs)
Catalog providers/databricks for Foundation Model APIs on AI Gateway: default mlflow/v1 via @databricks/ai-sdk-provider; per-model overrides for Anthropic Messages, Gemini generateContent, and Openai Responses (codex). Schema allows @databricks/ai-sdk-provider where api is required. Adds discovery and HTTP E2E scripts, README, and root README pointer.
1 parent d99a158 commit 3f1ca5e

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

47 files changed

+1537
-346
lines changed

README.md

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,7 @@ Models must conform to the following schema, as defined in `app/schemas.ts`.
146146
- `npm`: String - AI SDK Package name
147147
- `env`: String[] - Environment variable keys used for auth
148148
- `doc`: String - Link to the provider's documentation
149-
- `api` _(optional)_: String - OpenAI-compatible API endpoint. Required only when using `@ai-sdk/openai-compatible` as the npm package
149+
- `api` _(optional)_: String - OpenAI-compatible API endpoint. Required when using `@ai-sdk/openai-compatible` or `@databricks/ai-sdk-provider` as the npm package
150150

151151
**Model Schema:**
152152

@@ -186,6 +186,7 @@ See existing providers in the `providers/` directory for reference:
186186
- `providers/anthropic/` - Anthropic Claude models
187187
- `providers/openai/` - OpenAI GPT models
188188
- `providers/google/` - Google Gemini models
189+
- `providers/databricks/` - Databricks Foundation Model APIs on **AI Gateway**: default **`mlflow/v1`** for chat and embeddings via **`@databricks/ai-sdk-provider`**; per-model **`[provider]`** for **Claude** (Anthropic Messages), **Gemini** (native API), and **OpenAI Responses**. See [providers/databricks/README.md](providers/databricks/README.md) for discovery, authentication, and validation scripts.
189190

190191
### Working on frontend
191192

bun.lock

Lines changed: 0 additions & 341 deletions
This file was deleted.

package.json

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,9 @@
1919
"helicone:generate": "bun ./packages/core/script/generate-helicone.ts",
2020
"venice:generate": "bun ./packages/core/script/generate-venice.ts",
2121
"vercel:generate": "bun ./packages/core/script/generate-vercel.ts",
22-
"wandb:generate": "bun ./packages/core/script/generate-wandb.ts"
22+
"wandb:generate": "bun ./packages/core/script/generate-wandb.ts",
23+
"databricks:list-gateway": "bun ./packages/core/script/list-databricks-ai-gateway.ts",
24+
"databricks:test-inference": "bun ./packages/core/script/test-databricks.ts"
2325
},
2426
"dependencies": {
2527
"@cloudflare/workers-types": "^4.20250801.0",

packages/core/package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
"$schema": "https://json.schemastore.org/package.json",
55
"type": "module",
66
"dependencies": {
7+
"@databricks/sdk-experimental": "^0.16.0",
78
"zod": "catalog:"
89
},
910
"main": "./src/index.ts",
Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
/**
2+
* Shared types and helpers for Databricks AI Gateway discovery (system.ai FMA routes).
3+
*/
4+
5+
import { WorkspaceClient } from "@databricks/sdk-experimental";
6+
7+
export interface Destination {
8+
name?: string;
9+
type?: string;
10+
}
11+
12+
export interface Endpoint {
13+
name?: string;
14+
ai_gateway_url?: string;
15+
config?: { destinations?: Destination[] };
16+
}
17+
18+
export interface EndpointsResponse {
19+
endpoints?: Endpoint[];
20+
}
21+
22+
export interface FilteredGatewayRoute {
23+
gateway_name: string;
24+
system_ai_destinations: string[];
25+
ai_gateway_url?: string;
26+
}
27+
28+
export function isSystemAiFma(dest: Destination | undefined): boolean {
29+
if (!dest) return false;
30+
const t = dest.type ?? "";
31+
const name = dest.name ?? "";
32+
return t === "PAY_PER_TOKEN_FOUNDATION_MODEL" && name.startsWith("system.ai.");
33+
}
34+
35+
export function filterEndpoints(endpoints: Endpoint[]): FilteredGatewayRoute[] {
36+
const out: FilteredGatewayRoute[] = [];
37+
38+
for (const ep of endpoints) {
39+
const name = ep.name ?? "";
40+
if (!name.startsWith("databricks-")) continue;
41+
const dests = ep.config?.destinations ?? [];
42+
const sysAi = dests
43+
.filter(isSystemAiFma)
44+
.map((d) => d.name!)
45+
.filter(Boolean);
46+
if (sysAi.length === 0) continue;
47+
out.push({
48+
gateway_name: name,
49+
system_ai_destinations: sysAi,
50+
ai_gateway_url: ep.ai_gateway_url,
51+
});
52+
}
53+
out.sort((a, b) => a.gateway_name.localeCompare(b.gateway_name));
54+
return out;
55+
}
56+
57+
export async function fetchFilteredGatewayRoutes(
58+
client: WorkspaceClient,
59+
): Promise<FilteredGatewayRoute[]> {
60+
const raw = (await client.apiClient.request(
61+
{
62+
path: "/api/ai-gateway/v2/endpoints",
63+
method: "GET",
64+
headers: new Headers(),
65+
raw: false,
66+
},
67+
undefined,
68+
)) as EndpointsResponse;
69+
70+
return filterEndpoints(raw.endpoints ?? []);
71+
}
72+
73+
/** OpenAI-compatible base URL for chat/embeddings (no trailing slash). */
74+
export function mlflowOpenAiBaseUrl(aiGatewayUrl: string): string {
75+
const u = aiGatewayUrl.replace(/\/$/, "");
76+
return `${u}/mlflow/v1`;
77+
}
Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
#!/usr/bin/env bun
2+
/**
3+
* List Databricks AI Gateway routes aligned with Unity Catalog system.ai foundation models.
4+
*
5+
* Uses Databricks WorkspaceClient (JavaScript SDK; same auth patterns as ~/.databrickscfg, profiles, env).
6+
*
7+
* Usage (from repo root):
8+
* bun run databricks:list-gateway -- --profile YOUR_PROFILE
9+
* bun run databricks:list-gateway -- --profile YOUR_PROFILE --json
10+
*/
11+
12+
import { WorkspaceClient } from "@databricks/sdk-experimental";
13+
import { fetchFilteredGatewayRoutes } from "./databricks-ai-gateway-shared.js";
14+
15+
function parseArgs() {
16+
const argv = process.argv.slice(2);
17+
let profile = process.env.DATABRICKS_CONFIG_PROFILE;
18+
let json = false;
19+
for (let i = 0; i < argv.length; i++) {
20+
const a = argv[i];
21+
if (a === "--profile" && argv[i + 1]) {
22+
profile = argv[++i];
23+
continue;
24+
}
25+
if (a === "--json") {
26+
json = true;
27+
continue;
28+
}
29+
if (a === "--help" || a === "-h") {
30+
console.log(`Usage: list-databricks-ai-gateway.ts [--profile NAME] [--json]
31+
32+
--profile Databricks config profile (~/.databrickscfg). Default: DATABRICKS_CONFIG_PROFILE or SDK default chain.
33+
--json Print JSON array instead of TSV lines.
34+
`);
35+
process.exit(0);
36+
}
37+
}
38+
return { profile, json };
39+
}
40+
41+
async function main() {
42+
const { profile, json } = parseArgs();
43+
44+
const client = new WorkspaceClient(profile ? { profile } : {});
45+
const rows = await fetchFilteredGatewayRoutes(client);
46+
47+
if (json) {
48+
console.log(JSON.stringify(rows, null, 2));
49+
return;
50+
}
51+
52+
const hostUrl = (await client.apiClient.host).toString();
53+
console.log(`# host: ${hostUrl}`);
54+
console.log(`# count: ${rows.length}\n`);
55+
for (const r of rows) {
56+
console.log(`${r.gateway_name}\t${r.system_ai_destinations.join(", ")}`);
57+
}
58+
}
59+
60+
main().catch((e) => {
61+
console.error(e);
62+
process.exit(1);
63+
});

0 commit comments

Comments
 (0)