@dbx-tools/cli-model-proxy 0.3.21
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +143 -0
- package/bin/dbx-tools-model-proxy.mjs +11 -0
- package/bin/dbx-tools-model-proxy.ts +12 -0
- package/index.ts +11 -0
- package/package.json +59 -0
- package/src/backend.ts +146 -0
- package/src/cli.ts +200 -0
- package/src/defaults.ts +18 -0
- package/src/responses.ts +391 -0
- package/src/server.ts +416 -0
- package/tsconfig.json +43 -0
package/README.md
ADDED
|
@@ -0,0 +1,143 @@
|
|
|
1
|
+
# dbx-tools-model-proxy
|
|
2
|
+
|
|
3
|
+
Local OpenAI-compatible proxy for Databricks Model Serving.
|
|
4
|
+
|
|
5
|
+
Import this package or run its CLI when a tool expects the OpenAI API shape but
|
|
6
|
+
you want Databricks Model Serving auth, endpoint discovery, and fuzzy model
|
|
7
|
+
names. The proxy does not translate the OpenAI wire format; it resolves the
|
|
8
|
+
requested model, mints/refreshes Databricks auth through the SDK, and streams the
|
|
9
|
+
upstream response back to the caller.
|
|
10
|
+
|
|
11
|
+
Key features:
|
|
12
|
+
|
|
13
|
+
- OpenAI-compatible `/v1/*` forwarding for local tools that already know how to
|
|
14
|
+
call chat/completions endpoints.
|
|
15
|
+
- Databricks SDK auth, including profile selection, token refresh, and workspace
|
|
16
|
+
host resolution.
|
|
17
|
+
- Fuzzy model names and model-class requests powered by
|
|
18
|
+
[`@dbx-tools/model`](../../node/model).
|
|
19
|
+
- Optional local API-key enforcement for loopback safety.
|
|
20
|
+
- One-shot terminal chat mode that injects `OPENAI_BASE_URL`, `OPENAI_API_KEY`,
|
|
21
|
+
and `OPENAI_MODEL` into a child process.
|
|
22
|
+
- Programmatic Express app/server creation for tests and custom developer tools.
|
|
23
|
+
|
|
24
|
+
## Why Not Just AppKit Serving?
|
|
25
|
+
|
|
26
|
+
Use native AppKit Serving routes inside a Databricks App. They preserve AppKit's
|
|
27
|
+
plugin lifecycle, OBO request context, generated types, and React hooks.
|
|
28
|
+
|
|
29
|
+
Use this proxy for local tools that already speak the OpenAI API shape and know
|
|
30
|
+
nothing about AppKit:
|
|
31
|
+
|
|
32
|
+
- terminal chat clients and IDE integrations that only accept `OPENAI_BASE_URL`;
|
|
33
|
+
- local experiments where Databricks SDK auth should mint the upstream token;
|
|
34
|
+
- loose model names resolved through `@dbx-tools/model`;
|
|
35
|
+
- test harnesses that need an Express server with Databricks-backed `/v1/*`
|
|
36
|
+
routes.
|
|
37
|
+
|
|
38
|
+
## Run The Proxy
|
|
39
|
+
|
|
40
|
+
```sh
|
|
41
|
+
dbx-tools-model-proxy --profile my-workspace --port 4000
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
The package installs two equivalent commands: `dbx-tools-model-proxy` and the
|
|
45
|
+
shorter `dbxt-model-proxy`. Neither matches the package name, so a one-off run
|
|
46
|
+
has to name the command explicitly:
|
|
47
|
+
|
|
48
|
+
```sh
|
|
49
|
+
npx --package @dbx-tools/cli-model-proxy dbx-tools-model-proxy
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
Then point any OpenAI-compatible client at `http://127.0.0.1:4000/v1`:
|
|
53
|
+
|
|
54
|
+
```sh
|
|
55
|
+
curl http://127.0.0.1:4000/v1/chat/completions \
|
|
56
|
+
-H 'content-type: application/json' \
|
|
57
|
+
-d '{"model":"claude sonnet","messages":[{"role":"user","content":"hi"}]}'
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
The response includes `x-resolved-model`, showing which Databricks serving
|
|
61
|
+
endpoint the loose request snapped to.
|
|
62
|
+
|
|
63
|
+
The proxy is intentionally local-first. Bind it to `127.0.0.1` unless you are
|
|
64
|
+
putting another trusted access-control layer in front of it.
|
|
65
|
+
|
|
66
|
+
## Use A Terminal Chat Client
|
|
67
|
+
|
|
68
|
+
```sh
|
|
69
|
+
dbx-tools-model-proxy chat --profile my-workspace --model "claude sonnet"
|
|
70
|
+
dbx-tools-model-proxy chat --client "aichat" --model "chat fast"
|
|
71
|
+
```
|
|
72
|
+
|
|
73
|
+
`chat` starts the proxy, sets `OPENAI_BASE_URL`, `OPENAI_API_KEY`, and
|
|
74
|
+
`OPENAI_MODEL` for the child process, then shuts the proxy down when the child
|
|
75
|
+
exits. Use it to try Databricks-hosted models in any OpenAI-compatible terminal
|
|
76
|
+
client without editing that client's config.
|
|
77
|
+
|
|
78
|
+
## Inspect Model Resolution
|
|
79
|
+
|
|
80
|
+
```sh
|
|
81
|
+
dbx-tools-model-proxy models --profile my-workspace
|
|
82
|
+
dbx-tools-model-proxy resolve claude sonnet --profile my-workspace
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
These commands are useful when a client request resolves unexpectedly. They use
|
|
86
|
+
the same backend and resolver as the proxy server (`@dbx-tools/model`
|
|
87
|
+
`rankModels`: Fuse match, then class, then within-class version - so `opus`
|
|
88
|
+
prefers `opus-5` over `opus-4-7`).
|
|
89
|
+
|
|
90
|
+
## Require A Client API Key
|
|
91
|
+
|
|
92
|
+
```sh
|
|
93
|
+
dbx-tools-model-proxy --api-key "$LOCAL_PROXY_KEY"
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
With `--api-key` or `PROXY_API_KEY`, callers must send
|
|
97
|
+
`Authorization: Bearer <key>`. This protects the loopback proxy when another
|
|
98
|
+
local process may be able to reach it.
|
|
99
|
+
|
|
100
|
+
## Start Programmatically
|
|
101
|
+
|
|
102
|
+
```ts
|
|
103
|
+
import { backend, server } from "@dbx-tools/cli-model-proxy";
|
|
104
|
+
|
|
105
|
+
const db = await backend.DatabricksBackend.create({
|
|
106
|
+
profile: "my-workspace",
|
|
107
|
+
fuzzyThreshold: 0.35,
|
|
108
|
+
});
|
|
109
|
+
|
|
110
|
+
const running = await server.startProxyServer(db, {
|
|
111
|
+
host: "127.0.0.1",
|
|
112
|
+
port: 4000,
|
|
113
|
+
apiKey: process.env.LOCAL_PROXY_KEY,
|
|
114
|
+
});
|
|
115
|
+
|
|
116
|
+
console.log(running.url);
|
|
117
|
+
```
|
|
118
|
+
|
|
119
|
+
Use this when tests or local developer tools need a managed proxy lifecycle.
|
|
120
|
+
`server.createProxyServer()` returns the Express app without binding a port.
|
|
121
|
+
|
|
122
|
+
## How Requests Flow
|
|
123
|
+
|
|
124
|
+
1. `backend.DatabricksBackend` reads the OpenAI request body and resolves
|
|
125
|
+
`body.model` through [`@dbx-tools/model`](../../node/model).
|
|
126
|
+
2. The Databricks SDK supplies a fresh authorization header for the workspace.
|
|
127
|
+
3. The proxy forwards the body to the resolved serving endpoint's
|
|
128
|
+
`/invocations` route.
|
|
129
|
+
4. JSON or SSE response bodies are piped back unchanged.
|
|
130
|
+
|
|
131
|
+
This keeps the package small: Databricks already speaks the OpenAI schema, so
|
|
132
|
+
the useful work is auth and endpoint resolution.
|
|
133
|
+
|
|
134
|
+
## Modules
|
|
135
|
+
|
|
136
|
+
- `cli` - Commander program and `runCli()`.
|
|
137
|
+
- `backend` - `DatabricksBackend`, auth, model resolution, and upstream request
|
|
138
|
+
forwarding.
|
|
139
|
+
- `server` - Express proxy app and `startProxyServer()`.
|
|
140
|
+
- `defaults` - bind host, port, and invocation path constants.
|
|
141
|
+
|
|
142
|
+
Endpoint ranking and fuzzy matching come from
|
|
143
|
+
[`@dbx-tools/model`](../../node/model).
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
#!/usr/bin/env node
|
|
2
|
+
// GENERATED by projen synth (cli tag) - DO NOT EDIT.
|
|
3
|
+
// Regenerated from ./dbx-tools-model-proxy.ts.
|
|
4
|
+
// Hand edits are overwritten on the next watch; this file is read-only.
|
|
5
|
+
|
|
6
|
+
// Resolved as a bare specifier so Node looks in THIS file's package rather than the
|
|
7
|
+
// caller's cwd - the only way a globally installed CLI finds its own tsx.
|
|
8
|
+
import { register } from "tsx/esm/api";
|
|
9
|
+
|
|
10
|
+
register();
|
|
11
|
+
await import(new URL("./dbx-tools-model-proxy.ts", import.meta.url).href);
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
#!/usr/bin/env -S npx tsx
|
|
2
|
+
/**
|
|
3
|
+
* `dbx-tools-model-proxy` entry: a local OpenAI-compatible proxy in front of Databricks
|
|
4
|
+
* Model Serving. Delegates to the commander program in `../src/cli`.
|
|
5
|
+
*/
|
|
6
|
+
import { CommanderError, runCli } from "../src/cli";
|
|
7
|
+
|
|
8
|
+
runCli(process.argv).catch((err: unknown) => {
|
|
9
|
+
if (err instanceof CommanderError) process.exit(err.exitCode);
|
|
10
|
+
process.stderr.write(`${err instanceof Error ? (err.stack ?? err.message) : String(err)}\n`);
|
|
11
|
+
process.exit(1);
|
|
12
|
+
});
|
package/index.ts
ADDED
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
// GENERATED by projen watch - DO NOT EDIT.
|
|
2
|
+
// Regenerated from the exporting modules in ./src.
|
|
3
|
+
// Hand edits are overwritten on the next watch; this file is read-only.
|
|
4
|
+
|
|
5
|
+
export * as backend from "./src/backend";
|
|
6
|
+
export * as cli from "./src/cli";
|
|
7
|
+
export * as defaults from "./src/defaults";
|
|
8
|
+
export * as responses from "./src/responses";
|
|
9
|
+
export * as server from "./src/server";
|
|
10
|
+
export type { BackendOptions } from "./src/backend";
|
|
11
|
+
export type { ProxyServerOptions, StartProxyOptions } from "./src/server";
|
package/package.json
ADDED
|
@@ -0,0 +1,59 @@
|
|
|
1
|
+
{
|
|
2
|
+
"name": "@dbx-tools/cli-model-proxy",
|
|
3
|
+
"repository": {
|
|
4
|
+
"type": "git",
|
|
5
|
+
"url": "git+https://github.com/reggie-db/dbx-tools.git",
|
|
6
|
+
"directory": "workspaces/cli/model-proxy"
|
|
7
|
+
},
|
|
8
|
+
"bin": {
|
|
9
|
+
"dbx-tools-model-proxy": "./bin/dbx-tools-model-proxy.mjs",
|
|
10
|
+
"dbxt-model-proxy": "./bin/dbx-tools-model-proxy.mjs"
|
|
11
|
+
},
|
|
12
|
+
"devDependencies": {
|
|
13
|
+
"@types/node": "^24.6.0",
|
|
14
|
+
"typescript": "^5.9.3"
|
|
15
|
+
},
|
|
16
|
+
"dependencies": {
|
|
17
|
+
"@clack/prompts": "^1.7.0",
|
|
18
|
+
"@databricks/sdk-experimental": "^0.17.0",
|
|
19
|
+
"commander": "^15.0.0",
|
|
20
|
+
"tsx": "^4.23.0",
|
|
21
|
+
"@dbx-tools/model": "0.3.21",
|
|
22
|
+
"@dbx-tools/shared-core": "0.3.21",
|
|
23
|
+
"@dbx-tools/shared-model": "0.3.21"
|
|
24
|
+
},
|
|
25
|
+
"main": "index.ts",
|
|
26
|
+
"license": "UNLICENSED",
|
|
27
|
+
"publishConfig": {
|
|
28
|
+
"access": "public"
|
|
29
|
+
},
|
|
30
|
+
"version": "0.3.21",
|
|
31
|
+
"types": "index.ts",
|
|
32
|
+
"type": "module",
|
|
33
|
+
"exports": {
|
|
34
|
+
".": "./index.ts",
|
|
35
|
+
"./backend": "./src/backend.ts",
|
|
36
|
+
"./cli": "./src/cli.ts",
|
|
37
|
+
"./defaults": "./src/defaults.ts",
|
|
38
|
+
"./responses": "./src/responses.ts",
|
|
39
|
+
"./server": "./src/server.ts",
|
|
40
|
+
"./package.json": "./package.json"
|
|
41
|
+
},
|
|
42
|
+
"dbxToolsConfig": {
|
|
43
|
+
"tags": [
|
|
44
|
+
"cli"
|
|
45
|
+
]
|
|
46
|
+
},
|
|
47
|
+
"//": "~~ Generated by projen. To modify, edit .projenrc.js and run \"pnpm exec projen\".",
|
|
48
|
+
"scripts": {
|
|
49
|
+
"build": "projen build",
|
|
50
|
+
"compile": "projen compile",
|
|
51
|
+
"default": "projen default",
|
|
52
|
+
"package": "projen package",
|
|
53
|
+
"post-compile": "projen post-compile",
|
|
54
|
+
"pre-compile": "projen pre-compile",
|
|
55
|
+
"test": "projen test",
|
|
56
|
+
"watch": "projen watch",
|
|
57
|
+
"projen": "projen"
|
|
58
|
+
}
|
|
59
|
+
}
|
package/src/backend.ts
ADDED
|
@@ -0,0 +1,146 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* Databricks backend for the local model proxy.
|
|
3
|
+
*
|
|
4
|
+
* Wraps a single default-auth {@link WorkspaceClient} and exposes the three
|
|
5
|
+
* things the proxy server needs: the workspace serving-endpoint list, fuzzy
|
|
6
|
+
* name resolution (reusing `@dbx-tools/model`'s {@link resolve.rankModels} so a
|
|
7
|
+
* loose `"opus"` / `"claude sonnet"` snaps to the best live endpoint id -
|
|
8
|
+
* match score, then class, then within-class version), and a fresh set of auth
|
|
9
|
+
* headers per upstream request.
|
|
10
|
+
*
|
|
11
|
+
* Auth is delegated entirely to the Databricks SDK: `config.authenticate`
|
|
12
|
+
* re-runs the configured credential provider on every call and refreshes the
|
|
13
|
+
* underlying OAuth / PAT token when it is close to expiry, so the proxy never
|
|
14
|
+
* manages token lifetimes itself - each request is signed with a
|
|
15
|
+
* currently-valid bearer token.
|
|
16
|
+
*
|
|
17
|
+
* The endpoint catalogue is listed once and reused for the process, and
|
|
18
|
+
* re-listed on a resolve miss so a model deployed after start-up still resolves
|
|
19
|
+
* on first use - no cache layer, just one lazy load.
|
|
20
|
+
*
|
|
21
|
+
* @module
|
|
22
|
+
*/
|
|
23
|
+
|
|
24
|
+
import { log } from "@dbx-tools/shared-core";
|
|
25
|
+
import { type ServingEndpointSummary } from "@dbx-tools/shared-model";
|
|
26
|
+
import { resolve, serving, type ResolvedModel } from "@dbx-tools/model";
|
|
27
|
+
import { WorkspaceClient } from "@databricks/sdk-experimental";
|
|
28
|
+
|
|
29
|
+
import { INVOCATIONS_SUFFIX } from "./defaults";
|
|
30
|
+
|
|
31
|
+
const logger = log.logger("model-proxy/backend");
|
|
32
|
+
|
|
33
|
+
/** Options for {@link DatabricksBackend.create}. */
|
|
34
|
+
export interface BackendOptions {
|
|
35
|
+
/** Databricks config profile (`~/.databrickscfg`). Defaults to SDK auth resolution. */
|
|
36
|
+
profile?: string;
|
|
37
|
+
/** Override the workspace host; otherwise resolved from SDK auth (env / profile). */
|
|
38
|
+
host?: string;
|
|
39
|
+
/** Fuse.js fuzzy threshold (0 = exact, 1 = anything). Defaults to the model package default. */
|
|
40
|
+
threshold?: number;
|
|
41
|
+
}
|
|
42
|
+
|
|
43
|
+
export class DatabricksBackend {
|
|
44
|
+
private readonly client: WorkspaceClient;
|
|
45
|
+
/** Resolved workspace host, e.g. `https://my-workspace.cloud.databricks.com/`. */
|
|
46
|
+
readonly host: string;
|
|
47
|
+
private readonly threshold: number | undefined;
|
|
48
|
+
/** Lazily loaded endpoint catalogue, reused for the process lifetime. */
|
|
49
|
+
private endpoints: ServingEndpointSummary[] | undefined;
|
|
50
|
+
|
|
51
|
+
private constructor(client: WorkspaceClient, host: string, threshold: number | undefined) {
|
|
52
|
+
this.client = client;
|
|
53
|
+
this.host = host;
|
|
54
|
+
this.threshold = threshold;
|
|
55
|
+
}
|
|
56
|
+
|
|
57
|
+
/**
|
|
58
|
+
* Build a backend: construct a default-auth workspace client (optionally
|
|
59
|
+
* pinned to a profile / host) and resolve the workspace host once, so a bad
|
|
60
|
+
* profile fails at start-up rather than on the first proxied request.
|
|
61
|
+
*/
|
|
62
|
+
static async create(options: BackendOptions = {}): Promise<DatabricksBackend> {
|
|
63
|
+
const client = new WorkspaceClient({
|
|
64
|
+
...(options.host ? { host: options.host } : {}),
|
|
65
|
+
...(options.profile ? { profile: options.profile } : {}),
|
|
66
|
+
});
|
|
67
|
+
const host = (await client.config.getHost()).toString();
|
|
68
|
+
logger.info("connected", { host });
|
|
69
|
+
return new DatabricksBackend(client, host, options.threshold);
|
|
70
|
+
}
|
|
71
|
+
|
|
72
|
+
/**
|
|
73
|
+
* The workspace's serving-endpoint catalogue, as the minimal
|
|
74
|
+
* {@link ServingEndpointSummary} the resolver needs. Loaded lazily and
|
|
75
|
+
* reused; pass `force` to re-list (used by `/v1/models` and the
|
|
76
|
+
* resolve-on-miss path).
|
|
77
|
+
*/
|
|
78
|
+
async models(force = false): Promise<ServingEndpointSummary[]> {
|
|
79
|
+
if (this.endpoints && !force) return this.endpoints;
|
|
80
|
+
const out = await serving.listServingEndpointsUncached(this.client);
|
|
81
|
+
this.endpoints = out;
|
|
82
|
+
logger.debug("listed endpoints", { count: out.length });
|
|
83
|
+
return out;
|
|
84
|
+
}
|
|
85
|
+
|
|
86
|
+
/**
|
|
87
|
+
* Snap a (possibly loose) OpenAI-style model name to the best real serving
|
|
88
|
+
* endpoint via {@link resolve.rankModels} (the same ranking
|
|
89
|
+
* {@link resolve.resolveModel} uses): Fuse match (bucketed so version
|
|
90
|
+
* siblings tie), then class, then within-class version - so `"opus"` prefers
|
|
91
|
+
* `opus-5` over `opus-4-7`. On a miss the catalogue is re-listed once and
|
|
92
|
+
* retried, so a freshly deployed model resolves without a restart. Returns
|
|
93
|
+
* the input unchanged with `matched: false` when nothing scores within the
|
|
94
|
+
* threshold, so a deliberate endpoint id is never silently rewritten and
|
|
95
|
+
* Databricks surfaces a clean 404.
|
|
96
|
+
*/
|
|
97
|
+
async resolve(model: string): Promise<ResolvedModel> {
|
|
98
|
+
let resolved = this.resolveAgainst(model, await this.models());
|
|
99
|
+
if (!resolved.matched) {
|
|
100
|
+
resolved = this.resolveAgainst(model, await this.models(true));
|
|
101
|
+
}
|
|
102
|
+
return resolved;
|
|
103
|
+
}
|
|
104
|
+
|
|
105
|
+
/**
|
|
106
|
+
* Rank `model` against a catalogue snapshot and collapse to a single
|
|
107
|
+
* {@link ResolvedModel}. Uses {@link resolve.rankModels} with `limit: 1`
|
|
108
|
+
* rather than Fuse-only {@link serving.resolveModelId}, so equal-score
|
|
109
|
+
* siblings are broken by class / version.
|
|
110
|
+
*/
|
|
111
|
+
private resolveAgainst(
|
|
112
|
+
model: string,
|
|
113
|
+
endpoints: readonly ServingEndpointSummary[],
|
|
114
|
+
): ResolvedModel {
|
|
115
|
+
const [top] = resolve.rankModels(endpoints, {
|
|
116
|
+
search: model,
|
|
117
|
+
limit: 1,
|
|
118
|
+
...(this.threshold !== undefined ? { threshold: this.threshold } : {}),
|
|
119
|
+
});
|
|
120
|
+
if (!top) return { modelId: model, matched: false };
|
|
121
|
+
return { modelId: top.endpoint.name, matched: true, score: top.score };
|
|
122
|
+
}
|
|
123
|
+
|
|
124
|
+
/**
|
|
125
|
+
* Mint auth headers for one upstream request. The SDK refreshes the
|
|
126
|
+
* underlying token when needed, so every call gets a valid `Authorization`
|
|
127
|
+
* header without the proxy tracking expiry.
|
|
128
|
+
*/
|
|
129
|
+
async authHeaders(): Promise<Record<string, string>> {
|
|
130
|
+
const headers = new Headers();
|
|
131
|
+
await this.client.config.authenticate(headers);
|
|
132
|
+
const out: Record<string, string> = {};
|
|
133
|
+
headers.forEach((value, key) => {
|
|
134
|
+
out[key] = value;
|
|
135
|
+
});
|
|
136
|
+
return out;
|
|
137
|
+
}
|
|
138
|
+
|
|
139
|
+
/** OpenAI-compatible invocations URL for a resolved endpoint id. */
|
|
140
|
+
invocationsUrl(endpoint: string): string {
|
|
141
|
+
return new URL(
|
|
142
|
+
`serving-endpoints/${encodeURIComponent(endpoint)}/${INVOCATIONS_SUFFIX}`,
|
|
143
|
+
this.host,
|
|
144
|
+
).toString();
|
|
145
|
+
}
|
|
146
|
+
}
|
package/src/cli.ts
ADDED
|
@@ -0,0 +1,200 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* `dbx-tools-model-proxy` CLI.
|
|
3
|
+
*
|
|
4
|
+
* Run bare, it serves: a loopback OpenAI-compatible endpoint that fronts
|
|
5
|
+
* Databricks Model Serving with fuzzy model names and per-request auth. `chat`
|
|
6
|
+
* starts that same proxy and hands off to an off-the-shelf terminal client
|
|
7
|
+
* wired to it. `models` lists the resolvable endpoints; `resolve` shows what a
|
|
8
|
+
* fuzzy name snaps to. Auth comes from the standard Databricks SDK resolution
|
|
9
|
+
* (env vars, `--profile`, or `databricks auth login`).
|
|
10
|
+
*
|
|
11
|
+
* @module
|
|
12
|
+
*/
|
|
13
|
+
|
|
14
|
+
import { spawn } from "node:child_process";
|
|
15
|
+
import type { Server } from "node:http";
|
|
16
|
+
import { Command, CommanderError } from "commander";
|
|
17
|
+
|
|
18
|
+
import { type ServingEndpointSummary } from "@dbx-tools/shared-model";
|
|
19
|
+
|
|
20
|
+
import { DatabricksBackend, type BackendOptions } from "./backend";
|
|
21
|
+
import { DEFAULT_BIND_HOST, DEFAULT_PORT } from "./defaults";
|
|
22
|
+
import { startProxyServer } from "./server";
|
|
23
|
+
|
|
24
|
+
/** Capability flags an endpoint advertises, derived from its task/class. */
|
|
25
|
+
interface EndpointCapabilities {
|
|
26
|
+
/** OpenAI chat/completions + Responses (what a chat agent like Codex needs). */
|
|
27
|
+
chat: boolean;
|
|
28
|
+
/** Embedding (vector) endpoint. */
|
|
29
|
+
embedding: boolean;
|
|
30
|
+
/** Function / tool calling. Chat endpoints here support it. */
|
|
31
|
+
tools: boolean;
|
|
32
|
+
}
|
|
33
|
+
|
|
34
|
+
/**
|
|
35
|
+
* Derive an endpoint's capabilities from its Databricks task hint and the
|
|
36
|
+
* classifier's model class. `llm/v1/chat` (and chat-classed endpoints) are
|
|
37
|
+
* chat- and tool-capable; `llm/v1/embeddings` (and embedding-classed) are
|
|
38
|
+
* embedding-only. Keeps capability logic in one place so consumers don't
|
|
39
|
+
* re-derive it from raw strings.
|
|
40
|
+
*/
|
|
41
|
+
function capabilitiesFor(endpoint: ServingEndpointSummary): EndpointCapabilities {
|
|
42
|
+
const task = endpoint.task;
|
|
43
|
+
const cls = endpoint.class;
|
|
44
|
+
const embedding = task === "llm/v1/embeddings" || cls === "embedding";
|
|
45
|
+
const chat =
|
|
46
|
+
!embedding && (task === "llm/v1/chat" || (typeof cls === "string" && cls.startsWith("chat")));
|
|
47
|
+
return { chat, embedding, tools: chat };
|
|
48
|
+
}
|
|
49
|
+
|
|
50
|
+
/**
|
|
51
|
+
* Default terminal client for `chat`, launched via `bunx`. OpenHarness is an
|
|
52
|
+
* OpenAI-compatible TUI that reads the endpoint straight from the env we inject
|
|
53
|
+
* (`LLM_PROVIDER=openai-compat` + `OPENAI_BASE_URL`), so no config step is
|
|
54
|
+
* needed. Override with `--client` / `PROXY_CHAT_CLIENT`.
|
|
55
|
+
*/
|
|
56
|
+
const DEFAULT_CHAT_CLIENT = "bunx @alhazmiai/openharness";
|
|
57
|
+
|
|
58
|
+
/** Shared `--profile` / `--workspace-host` / `--threshold` flags. */
|
|
59
|
+
interface CommonOpts {
|
|
60
|
+
profile?: string;
|
|
61
|
+
workspaceHost?: string;
|
|
62
|
+
threshold?: string;
|
|
63
|
+
}
|
|
64
|
+
|
|
65
|
+
/** Shared listen + auth flags for the proxy-starting commands. */
|
|
66
|
+
interface ServeOpts extends CommonOpts {
|
|
67
|
+
port: string;
|
|
68
|
+
host: string;
|
|
69
|
+
apiKey?: string;
|
|
70
|
+
}
|
|
71
|
+
|
|
72
|
+
/** Map shared CLI flags onto {@link BackendOptions}. */
|
|
73
|
+
function backendOptions(opts: CommonOpts): BackendOptions {
|
|
74
|
+
return {
|
|
75
|
+
...(opts.profile ? { profile: opts.profile } : {}),
|
|
76
|
+
...(opts.workspaceHost ? { host: opts.workspaceHost } : {}),
|
|
77
|
+
...(opts.threshold !== undefined ? { threshold: Number(opts.threshold) } : {}),
|
|
78
|
+
};
|
|
79
|
+
}
|
|
80
|
+
|
|
81
|
+
/** Create the backend and start the proxy from the shared listen flags. */
|
|
82
|
+
async function startProxy(
|
|
83
|
+
opts: ServeOpts,
|
|
84
|
+
): Promise<{ backend: DatabricksBackend; server: Server; url: string }> {
|
|
85
|
+
const backend = await DatabricksBackend.create(backendOptions(opts));
|
|
86
|
+
const apiKey = opts.apiKey ?? process.env.PROXY_API_KEY;
|
|
87
|
+
const { server, url } = await startProxyServer(backend, {
|
|
88
|
+
host: opts.host,
|
|
89
|
+
port: Number(opts.port),
|
|
90
|
+
...(apiKey ? { apiKey } : {}),
|
|
91
|
+
});
|
|
92
|
+
return { backend, server, url };
|
|
93
|
+
}
|
|
94
|
+
|
|
95
|
+
/** Build the `dbx-tools-model-proxy` commander program (no side effects until parsed). */
|
|
96
|
+
export function buildProgram(): Command {
|
|
97
|
+
// Serving is the root action, not a subcommand: the bare command runs the
|
|
98
|
+
// proxy, and `chat`/`models`/`resolve` are the named detours off it.
|
|
99
|
+
const program = new Command()
|
|
100
|
+
.name("dbx-tools-model-proxy")
|
|
101
|
+
.description("Local OpenAI-compatible proxy to Databricks Model Serving.")
|
|
102
|
+
.option("-p, --port <port>", "port to listen on", String(DEFAULT_PORT))
|
|
103
|
+
.option("-H, --host <host>", "address to bind", DEFAULT_BIND_HOST)
|
|
104
|
+
.option("--profile <profile>", "Databricks config profile")
|
|
105
|
+
.option("--workspace-host <url>", "override the Databricks workspace host")
|
|
106
|
+
.option("-t, --threshold <n>", "fuzzy match threshold (0..1)")
|
|
107
|
+
.option("-k, --api-key <key>", "require this bearer token from local clients")
|
|
108
|
+
.action(async (opts: ServeOpts) => {
|
|
109
|
+
const { backend, url } = await startProxy(opts);
|
|
110
|
+
process.stderr.write(`model-proxy -> ${backend.host}\n`);
|
|
111
|
+
process.stderr.write(` OpenAI base URL: ${url}/v1\n`);
|
|
112
|
+
});
|
|
113
|
+
|
|
114
|
+
program
|
|
115
|
+
.command("chat")
|
|
116
|
+
.description("Start the proxy and launch a terminal chat client wired to it.")
|
|
117
|
+
.option("-p, --port <port>", "proxy port", String(DEFAULT_PORT))
|
|
118
|
+
.option("-H, --host <host>", "proxy bind host", DEFAULT_BIND_HOST)
|
|
119
|
+
.option("--profile <profile>", "Databricks config profile")
|
|
120
|
+
.option("--workspace-host <url>", "override the Databricks workspace host")
|
|
121
|
+
.option("-t, --threshold <n>", "fuzzy match threshold (0..1)")
|
|
122
|
+
.option("-m, --model <name>", "default model (fuzzy name ok)")
|
|
123
|
+
.option(
|
|
124
|
+
"--client <cmd>",
|
|
125
|
+
"terminal chat CLI to launch (run via your shell)",
|
|
126
|
+
process.env.PROXY_CHAT_CLIENT ?? DEFAULT_CHAT_CLIENT,
|
|
127
|
+
)
|
|
128
|
+
.action(async (opts: ServeOpts & { model?: string; client: string }) => {
|
|
129
|
+
const { backend, server, url } = await startProxy(opts);
|
|
130
|
+
const baseUrl = `${url}/v1`;
|
|
131
|
+
process.stderr.write(
|
|
132
|
+
`model-proxy -> ${backend.host}\n OpenAI base URL: ${baseUrl}\n launching: ${opts.client}\n`,
|
|
133
|
+
);
|
|
134
|
+
// Hand off to an off-the-shelf OpenAI-compatible client, pointing it at
|
|
135
|
+
// the proxy via the standard env vars (plus the provider switches a
|
|
136
|
+
// couple of popular CLIs read). `shell: true` lets `--client` carry its
|
|
137
|
+
// own args, e.g. `--client "bunx merlion"`.
|
|
138
|
+
const child = spawn(opts.client, {
|
|
139
|
+
stdio: "inherit",
|
|
140
|
+
shell: true,
|
|
141
|
+
env: {
|
|
142
|
+
...process.env,
|
|
143
|
+
OPENAI_BASE_URL: baseUrl,
|
|
144
|
+
OPENAI_API_KEY: process.env.OPENAI_API_KEY ?? "model-proxy",
|
|
145
|
+
...(opts.model ? { OPENAI_MODEL: opts.model } : {}),
|
|
146
|
+
LLM_PROVIDER: "openai-compat",
|
|
147
|
+
CLAUDE_CODE_USE_OPENAI: "1",
|
|
148
|
+
},
|
|
149
|
+
});
|
|
150
|
+
child.on("exit", (code) => {
|
|
151
|
+
server.close();
|
|
152
|
+
process.exit(code ?? 0);
|
|
153
|
+
});
|
|
154
|
+
process.on("SIGINT", () => child.kill("SIGINT"));
|
|
155
|
+
});
|
|
156
|
+
|
|
157
|
+
program
|
|
158
|
+
.command("models")
|
|
159
|
+
.description("List resolvable Databricks serving endpoints (as JSON).")
|
|
160
|
+
.option("--profile <profile>", "Databricks config profile")
|
|
161
|
+
.option("--workspace-host <url>", "override the Databricks workspace host")
|
|
162
|
+
.option("--chat", "only list chat-capable endpoints")
|
|
163
|
+
.action(async (opts: CommonOpts & { chat?: boolean }) => {
|
|
164
|
+
const backend = await DatabricksBackend.create(backendOptions(opts));
|
|
165
|
+
const endpoints = await backend.models();
|
|
166
|
+
// Enrich each endpoint with a derived `capabilities` object so consumers
|
|
167
|
+
// filter on capability rather than re-deriving it from raw `task`/`class`
|
|
168
|
+
// strings. `chat` = OpenAI chat/completions + Responses (what an agent
|
|
169
|
+
// like Codex needs); `embedding` = vector endpoints; `tools` = whether the
|
|
170
|
+
// endpoint supports function/tool calls (all chat endpoints here do).
|
|
171
|
+
const enriched = endpoints.map((endpoint) => {
|
|
172
|
+
const caps = capabilitiesFor(endpoint);
|
|
173
|
+
return { ...endpoint, capabilities: caps };
|
|
174
|
+
});
|
|
175
|
+
const out = opts.chat ? enriched.filter((e) => e.capabilities.chat) : enriched;
|
|
176
|
+
process.stdout.write(`${JSON.stringify(out, null, 2)}\n`);
|
|
177
|
+
});
|
|
178
|
+
|
|
179
|
+
program
|
|
180
|
+
.command("resolve")
|
|
181
|
+
.description("Show what a fuzzy model name resolves to (as JSON).")
|
|
182
|
+
.argument("<query...>", "model name / fuzzy search terms")
|
|
183
|
+
.option("--profile <profile>", "Databricks config profile")
|
|
184
|
+
.option("--workspace-host <url>", "override the Databricks workspace host")
|
|
185
|
+
.option("-t, --threshold <n>", "fuzzy match threshold (0..1)")
|
|
186
|
+
.action(async (query: string[], opts: CommonOpts) => {
|
|
187
|
+
const backend = await DatabricksBackend.create(backendOptions(opts));
|
|
188
|
+
const resolved = await backend.resolve(query.join(" "));
|
|
189
|
+
process.stdout.write(`${JSON.stringify(resolved, null, 2)}\n`);
|
|
190
|
+
});
|
|
191
|
+
|
|
192
|
+
return program;
|
|
193
|
+
}
|
|
194
|
+
|
|
195
|
+
/** Parse `argv` and run the matching command. Throws {@link CommanderError} on flag errors. */
|
|
196
|
+
export async function runCli(argv: string[]): Promise<void> {
|
|
197
|
+
await buildProgram().parseAsync(argv);
|
|
198
|
+
}
|
|
199
|
+
|
|
200
|
+
export { CommanderError };
|
package/src/defaults.ts
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
// Default values for the local Databricks model proxy.
|
|
2
|
+
|
|
3
|
+
/**
|
|
4
|
+
* Loopback address the proxy binds to by default. Keeps the OpenAI-compatible
|
|
5
|
+
* endpoint private to the machine unless the operator explicitly opts into a
|
|
6
|
+
* wider bind (e.g. `0.0.0.0`).
|
|
7
|
+
*/
|
|
8
|
+
export const DEFAULT_BIND_HOST = "127.0.0.1";
|
|
9
|
+
|
|
10
|
+
/** Default TCP port for the local proxy. */
|
|
11
|
+
export const DEFAULT_PORT = 4000;
|
|
12
|
+
|
|
13
|
+
/**
|
|
14
|
+
* Path segment Databricks Model Serving exposes for OpenAI-compatible requests.
|
|
15
|
+
* Every serving endpoint answers chat / completion / embedding payloads under
|
|
16
|
+
* `serving-endpoints/<name>/invocations`.
|
|
17
|
+
*/
|
|
18
|
+
export const INVOCATIONS_SUFFIX = "invocations";
|