workers-ai-provider 0.2.0 → 0.2.1

package/src/workersai-chat-prompt.ts

@@ -1,33 +1,33 @@
 export type WorkersAIChatPrompt = Array<WorkersAIChatMessage>;
 
 export type WorkersAIChatMessage =
-  | WorkersAISystemMessage
-  | WorkersAIUserMessage
-  | WorkersAIAssistantMessage
-  | WorkersAIToolMessage;
+	| WorkersAISystemMessage
+	| WorkersAIUserMessage
+	| WorkersAIAssistantMessage
+	| WorkersAIToolMessage;
 
 export interface WorkersAISystemMessage {
-  role: "system";
-  content: string;
+	role: "system";
+	content: string;
 }
 
 export interface WorkersAIUserMessage {
-  role: "user";
-  content: string;
+	role: "user";
+	content: string;
 }
 
 export interface WorkersAIAssistantMessage {
-  role: "assistant";
-  content: string;
-  tool_calls?: Array<{
-    id: string;
-    type: "function";
-    function: { name: string; arguments: string };
-  }>;
+	role: "assistant";
+	content: string;
+	tool_calls?: Array<{
+		id: string;
+		type: "function";
+		function: { name: string; arguments: string };
+	}>;
 }
 
 export interface WorkersAIToolMessage {
-  role: "tool";
-  name: string;
-  content: string;
+	role: "tool";
+	name: string;
+	content: string;
 }

package/src/workersai-chat-settings.ts

@@ -1,13 +1,13 @@
 export interface WorkersAIChatSettings {
-  /**
-  Whether to inject a safety prompt before all conversations.
-  
-  Defaults to `false`.
-     */
-  safePrompt?: boolean;
-  /**
-   * Optionally set Cloudflare AI Gateway options.
-   * @deprecated
-   */
-  gateway?: GatewayOptions;
+	/**
+	 * Whether to inject a safety prompt before all conversations.
+	 * Defaults to `false`.
+	 */
+	safePrompt?: boolean;
+
+	/**
+	 * Optionally set Cloudflare AI Gateway options.
+	 * @deprecated
+	 */
+	gateway?: GatewayOptions;
 }

package/src/workersai-error.ts

@@ -1,18 +1,15 @@
-import {
-  createJsonErrorResponseHandler,
-} from "@ai-sdk/provider-utils";
+import { createJsonErrorResponseHandler } from "@ai-sdk/provider-utils";
 import { z } from "zod";
 
 const workersAIErrorDataSchema = z.object({
-  object: z.literal("error"),
-  message: z.string(),
-  type: z.string(),
-  param: z.string().nullable(),
-  code: z.string().nullable(),
+	object: z.literal("error"),
+	message: z.string(),
+	type: z.string(),
+	param: z.string().nullable(),
+	code: z.string().nullable(),
 });
 
-export const workersAIFailedResponseHandler =
-  createJsonErrorResponseHandler({
-    errorSchema: workersAIErrorDataSchema,
-    errorToMessage: (data) => data.message,
-  });
+export const workersAIFailedResponseHandler = createJsonErrorResponseHandler({
+	errorSchema: workersAIErrorDataSchema,
+	errorToMessage: (data) => data.message,
+});

package/src/workersai-image-config.ts

@@ -0,0 +1,5 @@
+export type WorkersAIImageConfig = {
+	provider: string;
+	binding: Ai;
+	gateway?: GatewayOptions;
+};

package/src/workersai-image-model.ts

@@ -0,0 +1,114 @@
+import type { ImageModelV1, ImageModelV1CallWarning } from "@ai-sdk/provider";
+import type { WorkersAIImageConfig } from "./workersai-image-config";
+import type { WorkersAIImageSettings } from "./workersai-image-settings";
+import type { ImageGenerationModels } from "./workersai-models";
+
+export class WorkersAIImageModel implements ImageModelV1 {
+	readonly specificationVersion = "v1";
+
+	get maxImagesPerCall(): number {
+		return this.settings.maxImagesPerCall ?? 1;
+	}
+
+	get provider(): string {
+		return this.config.provider;
+	}
+	constructor(
+		readonly modelId: ImageGenerationModels,
+		readonly settings: WorkersAIImageSettings,
+		readonly config: WorkersAIImageConfig,
+	) {}
+
+	async doGenerate({
+		prompt,
+		n,
+		size,
+		aspectRatio,
+		seed,
+		// headers,
+		// abortSignal,
+	}: Parameters<ImageModelV1["doGenerate"]>[0]): Promise<
+		Awaited<ReturnType<ImageModelV1["doGenerate"]>>
+	> {
+		const { width, height } = getDimensionsFromSizeString(size);
+
+		const warnings: Array<ImageModelV1CallWarning> = [];
+
+		if (aspectRatio != null) {
+			warnings.push({
+				type: "unsupported-setting",
+				setting: "aspectRatio",
+				details: "This model does not support aspect ratio. Use `size` instead.",
+			});
+		}
+
+		const generateImage = async () => {
+			const outputStream: ReadableStream<Uint8Array> = await this.config.binding.run(
+				this.modelId,
+				{
+					prompt,
+					seed,
+					width,
+					height,
+				},
+			);
+
+			// Convert the output stream to a Uint8Array.
+			return streamToUint8Array(outputStream);
+		};
+
+		const images: Uint8Array[] = await Promise.all(
+			Array.from({ length: n }, () => generateImage()),
+		);
+
+		// type AiTextToImageOutput = ReadableStream<Uint8Array>;
+
+		return {
+			images,
+			warnings,
+			response: {
+				timestamp: new Date(),
+				modelId: this.modelId,
+				headers: {},
+			},
+		};
+	}
+}
+
+function getDimensionsFromSizeString(size: string | undefined) {
+	const [width, height] = size?.split("x") ?? [undefined, undefined];
+
+	return {
+		width: parseInteger(width),
+		height: parseInteger(height),
+	};
+}
+
+function parseInteger(value?: string) {
+	if (value === "" || !value) return undefined;
+	const number = Number(value);
+	return Number.isInteger(number) ? number : undefined;
+}
+
+async function streamToUint8Array(stream: ReadableStream<Uint8Array>): Promise<Uint8Array> {
+	const reader = stream.getReader();
+	const chunks: Uint8Array[] = [];
+	let totalLength = 0;
+
+	// Read the stream until it is finished.
+	while (true) {
+		const { done, value } = await reader.read();
+		if (done) break;
+		chunks.push(value);
+		totalLength += value.length;
+	}
+
+	// Allocate a new Uint8Array to hold all the data.
+	const result = new Uint8Array(totalLength);
+	let offset = 0;
+	for (const chunk of chunks) {
+		result.set(chunk, offset);
+		offset += chunk.length;
+	}
+	return result;
+}

package/src/workersai-image-settings.ts

@@ -0,0 +1,3 @@
+export type WorkersAIImageSettings = {
+	maxImagesPerCall?: number;
+};

package/src/workersai-models.ts

@@ -2,8 +2,13 @@
  * The names of the BaseAiTextGeneration models.
  */
 export type TextGenerationModels = Exclude<
-  value2key<AiModels, BaseAiTextGeneration>,
-  value2key<AiModels, BaseAiTextToImage>
+	value2key<AiModels, BaseAiTextGeneration>,
+	value2key<AiModels, BaseAiTextToImage>
 >;
 
+/*
+ * The names of the BaseAiTextToImage models.
+ */
+export type ImageGenerationModels = value2key<AiModels, BaseAiTextToImage>;
+
 type value2key<T, V> = { [K in keyof T]: T[K] extends V ? K : never }[keyof T];