wrangler: add AI related commands (cloudflare#3986)

* wrangler: add AI related commands * wrangler: truncate based on terminal size
OilyLime · Sep 22, 2023 · 00247a8 · 00247a8
1 parent f099a4f
commit 00247a8
Show file tree

Hide file tree

Showing 13 changed files with 996 additions and 3 deletions.
diff --git a/.changeset/orange-lizards-joke.md b/.changeset/orange-lizards-joke.md
@@ -0,0 +1,5 @@
+---
+"wrangler": patch
+---
+
+Added AI related CLI commands
diff --git a/packages/wrangler/src/__tests__/ai.test.ts b/packages/wrangler/src/__tests__/ai.test.ts
diff --git a/packages/wrangler/src/__tests__/deploy.test.ts b/packages/wrangler/src/__tests__/deploy.test.ts
@@ -143,7 +143,7 @@ describe("deploy", () => {
 
 			expect(std.out).toMatchInlineSnapshot(`
 			"Attempting to login via OAuth...
-			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
+			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
 			Successfully logged in.
 			Total Upload: xx KiB / gzip: xx KiB
 			Uploaded test-name (TIMINGS)
@@ -183,7 +183,7 @@ describe("deploy", () => {
 
 				expect(std.out).toMatchInlineSnapshot(`
 			"Attempting to login via OAuth...
-			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
+			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
 			Successfully logged in.
 			Total Upload: xx KiB / gzip: xx KiB
 			Uploaded test-name (TIMINGS)

diff --git a/packages/wrangler/src/__tests__/index.test.ts b/packages/wrangler/src/__tests__/index.test.ts
@@ -49,6 +49,7 @@ describe("wrangler", () => {
 			  wrangler r2                          📦 Interact with an R2 store
 			  wrangler dispatch-namespace          📦 Interact with a dispatch namespace
 			  wrangler d1                          🗄  Interact with a D1 database
+			  wrangler ai                          🤖 Interact with AI models
 			  wrangler constellation               🤖 Interact with Constellation models
 			  wrangler pubsub                      📮 Interact and manage Pub/Sub Brokers
 			  wrangler mtls-certificate            🪪 Manage certificates used for mTLS connections
@@ -101,6 +102,7 @@ describe("wrangler", () => {
 			  wrangler r2                          📦 Interact with an R2 store
 			  wrangler dispatch-namespace          📦 Interact with a dispatch namespace
 			  wrangler d1                          🗄  Interact with a D1 database
+			  wrangler ai                          🤖 Interact with AI models
 			  wrangler constellation               🤖 Interact with Constellation models
 			  wrangler pubsub                      📮 Interact and manage Pub/Sub Brokers
 			  wrangler mtls-certificate            🪪 Manage certificates used for mTLS connections

diff --git a/packages/wrangler/src/__tests__/mtls-certificates.test.ts b/packages/wrangler/src/__tests__/mtls-certificates.test.ts
@@ -379,6 +379,7 @@ describe("wrangler", () => {
 				  wrangler r2                          📦 Interact with an R2 store
 				  wrangler dispatch-namespace          📦 Interact with a dispatch namespace
 				  wrangler d1                          🗄  Interact with a D1 database
+				  wrangler ai                          🤖 Interact with AI models
 				  wrangler constellation               🤖 Interact with Constellation models
 				  wrangler pubsub                      📮 Interact and manage Pub/Sub Brokers
 				  wrangler mtls-certificate            🪪 Manage certificates used for mTLS connections

diff --git a/packages/wrangler/src/__tests__/user.test.ts b/packages/wrangler/src/__tests__/user.test.ts
@@ -60,7 +60,7 @@ describe("User", () => {
 			expect(counter).toBe(1);
 			expect(std.out).toMatchInlineSnapshot(`
 			"Attempting to login via OAuth...
-			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
+			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
 			Successfully logged in."
 		`);
 			expect(readAuthConfigFile()).toEqual<UserAuthConfig>({

diff --git a/packages/wrangler/src/ai/index.ts b/packages/wrangler/src/ai/index.ts
@@ -0,0 +1,11 @@
+import * as ListCatalog from "./listCatalog";
+import type { CommonYargsArgv } from "../yargs-types";
+
+export function ai(yargs: CommonYargsArgv) {
+	return yargs.command(
+		"models",
+		"List catalog models",
+		ListCatalog.options,
+		ListCatalog.handler
+	);
+}
diff --git a/packages/wrangler/src/ai/listCatalog.tsx b/packages/wrangler/src/ai/listCatalog.tsx
@@ -0,0 +1,55 @@
+import { withConfig } from "../config";
+import { logger } from "../logger";
+import { requireAuth } from "../user";
+import { asJson } from "./options";
+import { listCatalogEntries, truncate } from "./utils";
+import type {
+	CommonYargsArgv,
+	StrictYargsOptionsToInterface,
+} from "../yargs-types";
+
+export function options(yargs: CommonYargsArgv) {
+	return asJson(yargs);
+}
+
+function truncateDescription(
+	description: string | undefined,
+	alreadyUsed: number
+): string {
+	if (description === undefined || description === null) {
+		return "";
+	}
+
+	if (process.stdout.columns === undefined) {
+		return truncate(description, 100);
+	}
+
+	return truncate(description, process.stdout.columns - alreadyUsed);
+}
+
+type HandlerOptions = StrictYargsOptionsToInterface<typeof options>;
+export const handler = withConfig<HandlerOptions>(
+	async ({ json, config }): Promise<void> => {
+		const accountId = await requireAuth(config);
+		const entries = await listCatalogEntries(accountId);
+
+		if (json) {
+			logger.log(JSON.stringify(entries, null, 2));
+		} else {
+			logger.table(
+				entries.map((entry) => ({
+					model: entry.id,
+					name: entry.name,
+					description: truncateDescription(
+						entry.description,
+						entry.id.length +
+							entry.name.length +
+							(entry.task ? entry.task.name.length : 0) +
+							10
+					),
+					task: entry.task ? entry.task.name : "",
+				}))
+			);
+		}
+	}
+);
diff --git a/packages/wrangler/src/ai/options.ts b/packages/wrangler/src/ai/options.ts
@@ -0,0 +1,17 @@
+import type { CommonYargsArgv } from "../yargs-types";
+
+export function takeName(yargs: CommonYargsArgv) {
+	return yargs.positional("name", {
+		describe: "The name of the project",
+		type: "string",
+		demandOption: true,
+	});
+}
+
+export function asJson(yargs: CommonYargsArgv) {
+	return yargs.option("json", {
+		describe: "return output as clean JSON",
+		type: "boolean",
+		default: false,
+	});
+}
diff --git a/packages/wrangler/src/ai/types.ts b/packages/wrangler/src/ai/types.ts
@@ -0,0 +1,14 @@
+export type Task = {
+	id: string;
+	name: string;
+	description: string;
+};
+
+export type Model = {
+	id: string;
+	source: number;
+	task?: Task;
+	tags: string[];
+	name: string;
+	description?: string;
+};
diff --git a/packages/wrangler/src/ai/utils.ts b/packages/wrangler/src/ai/utils.ts
@@ -0,0 +1,37 @@
+import { fetchResult } from "../cfetch";
+import type { Model } from "./types";
+
+export async function aiList<ResponseType>(
+	accountId: string,
+	partialUrl: string
+): Promise<Array<ResponseType>> {
+	const pageSize = 50;
+	let page = 1;
+	const results = [];
+	while (results.length % pageSize === 0) {
+		const json: Array<ResponseType> = await fetchResult(
+			`/accounts/${accountId}/ai/${partialUrl}`,
+			{},
+			new URLSearchParams({
+				per_page: pageSize.toString(),
+				page: page.toString(),
+			})
+		);
+		page++;
+		results.push(...json);
+		if (json.length < pageSize) {
+			break;
+		}
+	}
+	return results;
+}
+
+export const listCatalogEntries = async (
+	accountId: string
+): Promise<Array<Model>> => {
+	return await aiList(accountId, "models/search");
+};
+
+export function truncate(str: string, maxLen: number) {
+	return str.slice(0, maxLen) + (str.length > maxLen ? "..." : "");
+}
diff --git a/packages/wrangler/src/index.ts b/packages/wrangler/src/index.ts
@@ -5,6 +5,7 @@ import supportsColor from "supports-color";
 import { ProxyAgent, setGlobalDispatcher } from "undici";
 import makeCLI from "yargs";
 import { version as wranglerVersion } from "../package.json";
+import { ai } from "./ai";
 import { loadDotEnv, readConfig } from "./config";
 import { constellation } from "./constellation";
 import { d1 } from "./d1";
@@ -451,6 +452,11 @@ export function createCLIParser(argv: string[]) {
 	});
 
 	// ai
+	wrangler.command("ai", "🤖 Interact with AI models", (aiYargs) => {
+		return ai(aiYargs.command(subHelp));
+	});
+
+	// constellation
 	wrangler.command(
 		"constellation",
 		"🤖 Interact with Constellation models",

diff --git a/packages/wrangler/src/user/user.ts b/packages/wrangler/src/user/user.ts
@@ -348,6 +348,7 @@ const Scopes = {
 	"zone:read": "Grants read level access to account zone.",
 	"ssl_certs:write": "See and manage mTLS certificates for your account",
 	"constellation:write": "Manage Constellation projects/models",
+	"ai:read": "List AI models",
 } as const;
 
 /**