Skip to content

Commit 00247a8

Browse files
authored
wrangler: add AI related commands (#3986)
* wrangler: add AI related commands * wrangler: truncate based on terminal size
1 parent f099a4f commit 00247a8

13 files changed

+996
-3
lines changed

.changeset/orange-lizards-joke.md

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"wrangler": patch
3+
---
4+
5+
Added AI related CLI commands

packages/wrangler/src/__tests__/ai.test.ts

+844
Large diffs are not rendered by default.

packages/wrangler/src/__tests__/deploy.test.ts

+2-2
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,7 @@ describe("deploy", () => {
143143

144144
expect(std.out).toMatchInlineSnapshot(`
145145
"Attempting to login via OAuth...
146-
Opening a link in your default browser: https://dash.cloudflare.com/oauth2/auth?response_type=code&client_id=54d11594-84e4-41aa-b438-e81b8fa78ee7&redirect_uri=http%3A%2F%2Flocalhost%3A8976%2Foauth%2Fcallback&scope=account%3Aread%20user%3Aread%20workers%3Awrite%20workers_kv%3Awrite%20workers_routes%3Awrite%20workers_scripts%3Awrite%20workers_tail%3Aread%20d1%3Awrite%20pages%3Awrite%20zone%3Aread%20ssl_certs%3Awrite%20constellation%3Awrite%20offline_access&state=MOCK_STATE_PARAM&code_challenge=MOCK_CODE_CHALLENGE&code_challenge_method=S256
146+
Opening a link in your default browser: https://dash.cloudflare.com/oauth2/auth?response_type=code&client_id=54d11594-84e4-41aa-b438-e81b8fa78ee7&redirect_uri=http%3A%2F%2Flocalhost%3A8976%2Foauth%2Fcallback&scope=account%3Aread%20user%3Aread%20workers%3Awrite%20workers_kv%3Awrite%20workers_routes%3Awrite%20workers_scripts%3Awrite%20workers_tail%3Aread%20d1%3Awrite%20pages%3Awrite%20zone%3Aread%20ssl_certs%3Awrite%20constellation%3Awrite%20ai%3Aread%20offline_access&state=MOCK_STATE_PARAM&code_challenge=MOCK_CODE_CHALLENGE&code_challenge_method=S256
147147
Successfully logged in.
148148
Total Upload: xx KiB / gzip: xx KiB
149149
Uploaded test-name (TIMINGS)
@@ -183,7 +183,7 @@ describe("deploy", () => {
183183

184184
expect(std.out).toMatchInlineSnapshot(`
185185
"Attempting to login via OAuth...
186-
Opening a link in your default browser: https://dash.staging.cloudflare.com/oauth2/auth?response_type=code&client_id=54d11594-84e4-41aa-b438-e81b8fa78ee7&redirect_uri=http%3A%2F%2Flocalhost%3A8976%2Foauth%2Fcallback&scope=account%3Aread%20user%3Aread%20workers%3Awrite%20workers_kv%3Awrite%20workers_routes%3Awrite%20workers_scripts%3Awrite%20workers_tail%3Aread%20d1%3Awrite%20pages%3Awrite%20zone%3Aread%20ssl_certs%3Awrite%20constellation%3Awrite%20offline_access&state=MOCK_STATE_PARAM&code_challenge=MOCK_CODE_CHALLENGE&code_challenge_method=S256
186+
Opening a link in your default browser: https://dash.staging.cloudflare.com/oauth2/auth?response_type=code&client_id=54d11594-84e4-41aa-b438-e81b8fa78ee7&redirect_uri=http%3A%2F%2Flocalhost%3A8976%2Foauth%2Fcallback&scope=account%3Aread%20user%3Aread%20workers%3Awrite%20workers_kv%3Awrite%20workers_routes%3Awrite%20workers_scripts%3Awrite%20workers_tail%3Aread%20d1%3Awrite%20pages%3Awrite%20zone%3Aread%20ssl_certs%3Awrite%20constellation%3Awrite%20ai%3Aread%20offline_access&state=MOCK_STATE_PARAM&code_challenge=MOCK_CODE_CHALLENGE&code_challenge_method=S256
187187
Successfully logged in.
188188
Total Upload: xx KiB / gzip: xx KiB
189189
Uploaded test-name (TIMINGS)

packages/wrangler/src/__tests__/index.test.ts

+2
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ describe("wrangler", () => {
4949
wrangler r2 📦 Interact with an R2 store
5050
wrangler dispatch-namespace 📦 Interact with a dispatch namespace
5151
wrangler d1 🗄 Interact with a D1 database
52+
wrangler ai 🤖 Interact with AI models
5253
wrangler constellation 🤖 Interact with Constellation models
5354
wrangler pubsub 📮 Interact and manage Pub/Sub Brokers
5455
wrangler mtls-certificate 🪪 Manage certificates used for mTLS connections
@@ -101,6 +102,7 @@ describe("wrangler", () => {
101102
wrangler r2 📦 Interact with an R2 store
102103
wrangler dispatch-namespace 📦 Interact with a dispatch namespace
103104
wrangler d1 🗄 Interact with a D1 database
105+
wrangler ai 🤖 Interact with AI models
104106
wrangler constellation 🤖 Interact with Constellation models
105107
wrangler pubsub 📮 Interact and manage Pub/Sub Brokers
106108
wrangler mtls-certificate 🪪 Manage certificates used for mTLS connections

packages/wrangler/src/__tests__/mtls-certificates.test.ts

+1
Original file line numberDiff line numberDiff line change
@@ -379,6 +379,7 @@ describe("wrangler", () => {
379379
wrangler r2 📦 Interact with an R2 store
380380
wrangler dispatch-namespace 📦 Interact with a dispatch namespace
381381
wrangler d1 🗄 Interact with a D1 database
382+
wrangler ai 🤖 Interact with AI models
382383
wrangler constellation 🤖 Interact with Constellation models
383384
wrangler pubsub 📮 Interact and manage Pub/Sub Brokers
384385
wrangler mtls-certificate 🪪 Manage certificates used for mTLS connections

packages/wrangler/src/__tests__/user.test.ts

+1-1
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ describe("User", () => {
6060
expect(counter).toBe(1);
6161
expect(std.out).toMatchInlineSnapshot(`
6262
"Attempting to login via OAuth...
63-
Opening a link in your default browser: https://dash.cloudflare.com/oauth2/auth?response_type=code&client_id=54d11594-84e4-41aa-b438-e81b8fa78ee7&redirect_uri=http%3A%2F%2Flocalhost%3A8976%2Foauth%2Fcallback&scope=account%3Aread%20user%3Aread%20workers%3Awrite%20workers_kv%3Awrite%20workers_routes%3Awrite%20workers_scripts%3Awrite%20workers_tail%3Aread%20d1%3Awrite%20pages%3Awrite%20zone%3Aread%20ssl_certs%3Awrite%20constellation%3Awrite%20offline_access&state=MOCK_STATE_PARAM&code_challenge=MOCK_CODE_CHALLENGE&code_challenge_method=S256
63+
Opening a link in your default browser: https://dash.cloudflare.com/oauth2/auth?response_type=code&client_id=54d11594-84e4-41aa-b438-e81b8fa78ee7&redirect_uri=http%3A%2F%2Flocalhost%3A8976%2Foauth%2Fcallback&scope=account%3Aread%20user%3Aread%20workers%3Awrite%20workers_kv%3Awrite%20workers_routes%3Awrite%20workers_scripts%3Awrite%20workers_tail%3Aread%20d1%3Awrite%20pages%3Awrite%20zone%3Aread%20ssl_certs%3Awrite%20constellation%3Awrite%20ai%3Aread%20offline_access&state=MOCK_STATE_PARAM&code_challenge=MOCK_CODE_CHALLENGE&code_challenge_method=S256
6464
Successfully logged in."
6565
`);
6666
expect(readAuthConfigFile()).toEqual<UserAuthConfig>({

packages/wrangler/src/ai/index.ts

+11
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
import * as ListCatalog from "./listCatalog";
2+
import type { CommonYargsArgv } from "../yargs-types";
3+
4+
export function ai(yargs: CommonYargsArgv) {
5+
return yargs.command(
6+
"models",
7+
"List catalog models",
8+
ListCatalog.options,
9+
ListCatalog.handler
10+
);
11+
}
+55
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
import { withConfig } from "../config";
2+
import { logger } from "../logger";
3+
import { requireAuth } from "../user";
4+
import { asJson } from "./options";
5+
import { listCatalogEntries, truncate } from "./utils";
6+
import type {
7+
CommonYargsArgv,
8+
StrictYargsOptionsToInterface,
9+
} from "../yargs-types";
10+
11+
export function options(yargs: CommonYargsArgv) {
12+
return asJson(yargs);
13+
}
14+
15+
function truncateDescription(
16+
description: string | undefined,
17+
alreadyUsed: number
18+
): string {
19+
if (description === undefined || description === null) {
20+
return "";
21+
}
22+
23+
if (process.stdout.columns === undefined) {
24+
return truncate(description, 100);
25+
}
26+
27+
return truncate(description, process.stdout.columns - alreadyUsed);
28+
}
29+
30+
type HandlerOptions = StrictYargsOptionsToInterface<typeof options>;
31+
export const handler = withConfig<HandlerOptions>(
32+
async ({ json, config }): Promise<void> => {
33+
const accountId = await requireAuth(config);
34+
const entries = await listCatalogEntries(accountId);
35+
36+
if (json) {
37+
logger.log(JSON.stringify(entries, null, 2));
38+
} else {
39+
logger.table(
40+
entries.map((entry) => ({
41+
model: entry.id,
42+
name: entry.name,
43+
description: truncateDescription(
44+
entry.description,
45+
entry.id.length +
46+
entry.name.length +
47+
(entry.task ? entry.task.name.length : 0) +
48+
10
49+
),
50+
task: entry.task ? entry.task.name : "",
51+
}))
52+
);
53+
}
54+
}
55+
);

packages/wrangler/src/ai/options.ts

+17
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
import type { CommonYargsArgv } from "../yargs-types";
2+
3+
export function takeName(yargs: CommonYargsArgv) {
4+
return yargs.positional("name", {
5+
describe: "The name of the project",
6+
type: "string",
7+
demandOption: true,
8+
});
9+
}
10+
11+
export function asJson(yargs: CommonYargsArgv) {
12+
return yargs.option("json", {
13+
describe: "return output as clean JSON",
14+
type: "boolean",
15+
default: false,
16+
});
17+
}

packages/wrangler/src/ai/types.ts

+14
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
export type Task = {
2+
id: string;
3+
name: string;
4+
description: string;
5+
};
6+
7+
export type Model = {
8+
id: string;
9+
source: number;
10+
task?: Task;
11+
tags: string[];
12+
name: string;
13+
description?: string;
14+
};

packages/wrangler/src/ai/utils.ts

+37
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
import { fetchResult } from "../cfetch";
2+
import type { Model } from "./types";
3+
4+
export async function aiList<ResponseType>(
5+
accountId: string,
6+
partialUrl: string
7+
): Promise<Array<ResponseType>> {
8+
const pageSize = 50;
9+
let page = 1;
10+
const results = [];
11+
while (results.length % pageSize === 0) {
12+
const json: Array<ResponseType> = await fetchResult(
13+
`/accounts/${accountId}/ai/${partialUrl}`,
14+
{},
15+
new URLSearchParams({
16+
per_page: pageSize.toString(),
17+
page: page.toString(),
18+
})
19+
);
20+
page++;
21+
results.push(...json);
22+
if (json.length < pageSize) {
23+
break;
24+
}
25+
}
26+
return results;
27+
}
28+
29+
export const listCatalogEntries = async (
30+
accountId: string
31+
): Promise<Array<Model>> => {
32+
return await aiList(accountId, "models/search");
33+
};
34+
35+
export function truncate(str: string, maxLen: number) {
36+
return str.slice(0, maxLen) + (str.length > maxLen ? "..." : "");
37+
}

packages/wrangler/src/index.ts

+6
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import supportsColor from "supports-color";
55
import { ProxyAgent, setGlobalDispatcher } from "undici";
66
import makeCLI from "yargs";
77
import { version as wranglerVersion } from "../package.json";
8+
import { ai } from "./ai";
89
import { loadDotEnv, readConfig } from "./config";
910
import { constellation } from "./constellation";
1011
import { d1 } from "./d1";
@@ -451,6 +452,11 @@ export function createCLIParser(argv: string[]) {
451452
});
452453

453454
// ai
455+
wrangler.command("ai", "🤖 Interact with AI models", (aiYargs) => {
456+
return ai(aiYargs.command(subHelp));
457+
});
458+
459+
// constellation
454460
wrangler.command(
455461
"constellation",
456462
"🤖 Interact with Constellation models",

packages/wrangler/src/user/user.ts

+1
Original file line numberDiff line numberDiff line change
@@ -348,6 +348,7 @@ const Scopes = {
348348
"zone:read": "Grants read level access to account zone.",
349349
"ssl_certs:write": "See and manage mTLS certificates for your account",
350350
"constellation:write": "Manage Constellation projects/models",
351+
"ai:read": "List AI models",
351352
} as const;
352353

353354
/**

0 commit comments

Comments
 (0)