Skip to content
Draft
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ Models must conform to the following schema, as defined in `app/schemas.ts`.
- `npm`: String - AI SDK Package name
- `env`: String[] - Environment variable keys used for auth
- `doc`: String - Link to the provider's documentation
- `api` _(optional)_: String - OpenAI-compatible API endpoint. Required only when using `@ai-sdk/openai-compatible` as the npm package
- `api` _(optional)_: String - OpenAI-compatible API endpoint. Required when using `@ai-sdk/openai-compatible` or `@databricks/ai-sdk-provider` as the npm package

**Model Schema:**

Expand Down Expand Up @@ -186,6 +186,7 @@ See existing providers in the `providers/` directory for reference:
- `providers/anthropic/` - Anthropic Claude models
- `providers/openai/` - OpenAI GPT models
- `providers/google/` - Google Gemini models
- `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.

### Working on frontend

Expand Down
331 changes: 189 additions & 142 deletions bun.lock

Large diffs are not rendered by default.

4 changes: 3 additions & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,9 @@
"helicone:generate": "bun ./packages/core/script/generate-helicone.ts",
"venice:generate": "bun ./packages/core/script/generate-venice.ts",
"vercel:generate": "bun ./packages/core/script/generate-vercel.ts",
"wandb:generate": "bun ./packages/core/script/generate-wandb.ts"
"wandb:generate": "bun ./packages/core/script/generate-wandb.ts",
"databricks:list-gateway": "bun ./packages/core/script/list-databricks-ai-gateway.ts",
"databricks:test-inference": "bun ./packages/core/script/test-databricks.ts"
},
"dependencies": {
"@cloudflare/workers-types": "^4.20250801.0",
Expand Down
1 change: 1 addition & 0 deletions packages/core/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
"$schema": "https://json.schemastore.org/package.json",
"type": "module",
"dependencies": {
"@databricks/sdk-experimental": "^0.16.0",
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i'm unsure how packages/core is used, but we prob wouldn't want to add databricks as a dependency for the core package right? wouldn't this be installed for all customers? this additional dependency probably belongs in a provider-specific folder

"zod": "catalog:"
},
"main": "./src/index.ts",
Expand Down
77 changes: 77 additions & 0 deletions packages/core/script/databricks-ai-gateway-shared.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
/**
* Shared types and helpers for Databricks AI Gateway discovery (system.ai FMA routes).
*/

import { WorkspaceClient } from "@databricks/sdk-experimental";

export interface Destination {
name?: string;
type?: string;
}

export interface Endpoint {
name?: string;
ai_gateway_url?: string;
config?: { destinations?: Destination[] };
}

export interface EndpointsResponse {
endpoints?: Endpoint[];
}

export interface FilteredGatewayRoute {
gateway_name: string;
system_ai_destinations: string[];
ai_gateway_url?: string;
}

export function isSystemAiFma(dest: Destination | undefined): boolean {
if (!dest) return false;
const t = dest.type ?? "";
const name = dest.name ?? "";
return t === "PAY_PER_TOKEN_FOUNDATION_MODEL" && name.startsWith("system.ai.");
}

export function filterEndpoints(endpoints: Endpoint[]): FilteredGatewayRoute[] {
const out: FilteredGatewayRoute[] = [];

for (const ep of endpoints) {
const name = ep.name ?? "";
if (!name.startsWith("databricks-")) continue;
const dests = ep.config?.destinations ?? [];
const sysAi = dests
.filter(isSystemAiFma)
.map((d) => d.name!)
.filter(Boolean);
if (sysAi.length === 0) continue;
out.push({
gateway_name: name,
system_ai_destinations: sysAi,
ai_gateway_url: ep.ai_gateway_url,
});
}
out.sort((a, b) => a.gateway_name.localeCompare(b.gateway_name));
return out;
}

export async function fetchFilteredGatewayRoutes(
client: WorkspaceClient,
): Promise<FilteredGatewayRoute[]> {
const raw = (await client.apiClient.request(
{
path: "/api/ai-gateway/v2/endpoints",
method: "GET",
headers: new Headers(),
raw: false,
},
undefined,
)) as EndpointsResponse;

return filterEndpoints(raw.endpoints ?? []);
}

/** OpenAI-compatible base URL for chat/embeddings (no trailing slash). */
export function mlflowOpenAiBaseUrl(aiGatewayUrl: string): string {
const u = aiGatewayUrl.replace(/\/$/, "");
return `${u}/mlflow/v1`;
}
63 changes: 63 additions & 0 deletions packages/core/script/list-databricks-ai-gateway.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
#!/usr/bin/env bun
/**
* List Databricks AI Gateway routes aligned with Unity Catalog system.ai foundation models.
*
* Uses Databricks WorkspaceClient (JavaScript SDK; same auth patterns as ~/.databrickscfg, profiles, env).
*
* Usage (from repo root):
* bun run databricks:list-gateway -- --profile YOUR_PROFILE
* bun run databricks:list-gateway -- --profile YOUR_PROFILE --json
*/

import { WorkspaceClient } from "@databricks/sdk-experimental";
import { fetchFilteredGatewayRoutes } from "./databricks-ai-gateway-shared.js";

function parseArgs() {
const argv = process.argv.slice(2);
let profile = process.env.DATABRICKS_CONFIG_PROFILE;
let json = false;
for (let i = 0; i < argv.length; i++) {
const a = argv[i];
if (a === "--profile" && argv[i + 1]) {
profile = argv[++i];
continue;
}
if (a === "--json") {
json = true;
continue;
}
if (a === "--help" || a === "-h") {
console.log(`Usage: list-databricks-ai-gateway.ts [--profile NAME] [--json]

--profile Databricks config profile (~/.databrickscfg). Default: DATABRICKS_CONFIG_PROFILE or SDK default chain.
--json Print JSON array instead of TSV lines.
`);
process.exit(0);
}
}
return { profile, json };
}

async function main() {
const { profile, json } = parseArgs();

const client = new WorkspaceClient(profile ? { profile } : {});
const rows = await fetchFilteredGatewayRoutes(client);

if (json) {
console.log(JSON.stringify(rows, null, 2));
return;
}

const hostUrl = (await client.apiClient.host).toString();
console.log(`# host: ${hostUrl}`);
console.log(`# count: ${rows.length}\n`);
for (const r of rows) {
console.log(`${r.gateway_name}\t${r.system_ai_destinations.join(", ")}`);
}
}

main().catch((e) => {
console.error(e);
process.exit(1);
});
Loading