Add workers ai staging flag (#5705)

* Add workers ai staging flag * add `staging` option to ai binding config --------- Co-authored-by: Rahul Sethi <5822355+RamIdeas@users.noreply.github.com>
cloudflare · Apr 29, 2024 · 4097759 · 4097759
1 parent 76a456a
commit 4097759
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 4 deletions.
diff --git a/.changeset/thin-wasps-bathe.md b/.changeset/thin-wasps-bathe.md
@@ -0,0 +1,5 @@
+---
+"wrangler": patch
+---
+
+Add `staging` flag to AI binding
diff --git a/packages/wrangler/src/config/environment.ts b/packages/wrangler/src/config/environment.ts
@@ -632,6 +632,7 @@ export interface EnvironmentNonInheritable {
 	ai:
 		| {
 				binding: string;
+				staging?: boolean;
 		  }
 		| undefined;
 

diff --git a/packages/wrangler/src/config/index.ts b/packages/wrangler/src/config/index.ts
@@ -183,8 +183,10 @@ export function printBindings(bindings: CfWorkerInit["bindings"]) {
 		return `${s.substring(0, maxLength - 3)}...`;
 	};
 
-	const output: { type: string; entries: { key: string; value: string }[] }[] =
-		[];
+	const output: {
+		type: string;
+		entries: { key: string; value: string | boolean }[];
+	}[] = [];
 
 	const {
 		data_blobs,
@@ -419,9 +421,14 @@ export function printBindings(bindings: CfWorkerInit["bindings"]) {
 	}
 
 	if (ai !== undefined) {
+		const entries: [{ key: string; value: string | boolean }] = [
+			{ key: "Name", value: ai.binding },
+		];
+		if (ai.staging) entries.push({ key: "Staging", value: ai.staging });
+
 		output.push({
 			type: "AI",
-			entries: [{ key: "Name", value: ai.binding }],
+			entries: entries,
 		});
 	}
 

diff --git a/packages/wrangler/src/deployment-bundle/create-worker-upload-form.ts b/packages/wrangler/src/deployment-bundle/create-worker-upload-form.ts
@@ -42,7 +42,7 @@ export type WorkerMetadataBinding =
 	| { type: "wasm_module"; name: string; part: string }
 	| { type: "text_blob"; name: string; part: string }
 	| { type: "browser"; name: string }
-	| { type: "ai"; name: string }
+	| { type: "ai"; name: string; staging?: boolean }
 	| { type: "version_metadata"; name: string }
 	| { type: "data_blob"; name: string; part: string }
 	| { type: "kv_namespace"; name: string; namespace_id: string }
@@ -329,6 +329,7 @@ export function createWorkerUploadForm(worker: CfWorkerInit): FormData {
 	if (bindings.ai !== undefined) {
 		metadataBindings.push({
 			name: bindings.ai.binding,
+			staging: bindings.ai.staging,
 			type: "ai",
 		});
 	}

diff --git a/packages/wrangler/src/deployment-bundle/worker.ts b/packages/wrangler/src/deployment-bundle/worker.ts
@@ -111,6 +111,7 @@ export interface CfBrowserBinding {
 
 export interface CfAIBinding {
 	binding: string;
+	staging?: boolean;
 }
 
 /**