workers-ai-provider 0.4.0 → 0.5.0

package/src/streaming.ts

@@ -0,0 +1,37 @@
+import { events } from "fetch-event-stream";
+
+import type { LanguageModelV1StreamPart } from "@ai-sdk/provider";
+import { mapWorkersAIUsage } from "./map-workersai-usage";
+
+export function getMappedStream(response: Response) {
+	const chunkEvent = events(response);
+	let usage = { promptTokens: 0, completionTokens: 0 };
+
+	return new ReadableStream<LanguageModelV1StreamPart>({
+		async start(controller) {
+			for await (const event of chunkEvent) {
+				if (!event.data) {
+					continue;
+				}
+				if (event.data === "[DONE]") {
+					break;
+				}
+				const chunk = JSON.parse(event.data);
+				if (chunk.usage) {
+					usage = mapWorkersAIUsage(chunk);
+				}
+				chunk.response?.length &&
+					controller.enqueue({
+						type: "text-delta",
+						textDelta: chunk.response,
+					});
+			}
+			controller.enqueue({
+				type: "finish",
+				finishReason: "stop",
+				usage: usage,
+			});
+			controller.close();
+		},
+	});
+}

package/src/utils.ts

@@ -1,3 +1,5 @@
+import type { LanguageModelV1 } from "@ai-sdk/provider";
+
 /**
  * General AI run interface with overloads to handle distinct return types.
  *
@@ -83,7 +85,9 @@ export function createRun(config: CreateRunConfig): AiRun {
 			}
 		}
 
-		const url = `https://api.cloudflare.com/client/v4/accounts/${accountId}/ai/run/${model}${urlParams ? `?${urlParams}` : ""}`;
+		const url = `https://api.cloudflare.com/client/v4/accounts/${accountId}/ai/run/${model}${
+			urlParams ? `?${urlParams}` : ""
+		}`;
 
 		// Merge default and custom headers.
 		const headers = {
@@ -120,3 +124,91 @@ export function createRun(config: CreateRunConfig): AiRun {
 		return data.result;
 	};
 }
+
+export function prepareToolsAndToolChoice(
+	mode: Parameters<LanguageModelV1["doGenerate"]>[0]["mode"] & {
+		type: "regular";
+	},
+) {
+	// when the tools array is empty, change it to undefined to prevent errors:
+	const tools = mode.tools?.length ? mode.tools : undefined;
+
+	if (tools == null) {
+		return { tools: undefined, tool_choice: undefined };
+	}
+
+	const mappedTools = tools.map((tool) => ({
+		type: "function",
+		function: {
+			name: tool.name,
+			// @ts-expect-error - description is not a property of tool
+			description: tool.description,
+			// @ts-expect-error - parameters is not a property of tool
+			parameters: tool.parameters,
+		},
+	}));
+
+	const toolChoice = mode.toolChoice;
+
+	if (toolChoice == null) {
+		return { tools: mappedTools, tool_choice: undefined };
+	}
+
+	const type = toolChoice.type;
+
+	switch (type) {
+		case "auto":
+			return { tools: mappedTools, tool_choice: type };
+		case "none":
+			return { tools: mappedTools, tool_choice: type };
+		case "required":
+			return { tools: mappedTools, tool_choice: "any" };
+
+		// workersAI does not support tool mode directly,
+		// so we filter the tools and force the tool choice through 'any'
+		case "tool":
+			return {
+				tools: mappedTools.filter((tool) => tool.function.name === toolChoice.toolName),
+				tool_choice: "any",
+			};
+		default: {
+			const exhaustiveCheck = type satisfies never;
+			throw new Error(`Unsupported tool choice type: ${exhaustiveCheck}`);
+		}
+	}
+}
+
+export function lastMessageWasUser<T extends { role: string }>(messages: T[]) {
+	return messages.length > 0 && messages[messages.length - 1]!.role === "user";
+}
+
+export function processToolCalls(output: any) {
+	// Check for OpenAI format tool calls first
+	if (output.tool_calls && Array.isArray(output.tool_calls)) {
+		return output.tool_calls.map((toolCall: any) => {
+			// Handle new format
+			if (toolCall.function && toolCall.id) {
+				return {
+					toolCallType: "function",
+					toolCallId: toolCall.id,
+					toolName: toolCall.function.name,
+					args:
+						typeof toolCall.function.arguments === "string"
+							? toolCall.function.arguments
+							: JSON.stringify(toolCall.function.arguments || {}),
+				};
+			}
+			return {
+				toolCallType: "function",
+				toolCallId: toolCall.name,
+				toolName: toolCall.name,
+				args:
+					typeof toolCall.arguments === "string"
+						? toolCall.arguments
+						: JSON.stringify(toolCall.arguments || {}),
+			};
+		});
+	}
+
+	return [];
+}

package/src/workers-ai-embedding-model.ts

@@ -0,0 +1,87 @@
+import { TooManyEmbeddingValuesForCallError, type EmbeddingModelV1 } from "@ai-sdk/provider";
+import type { StringLike } from "./utils";
+import type { EmbeddingModels } from "./workersai-models";
+
+export type WorkersAIEmbeddingConfig = {
+	provider: string;
+	binding: Ai;
+	gateway?: GatewayOptions;
+};
+
+export type WorkersAIEmbeddingSettings = {
+	gateway?: GatewayOptions;
+	maxEmbeddingsPerCall?: number;
+	supportsParallelCalls?: boolean;
+} & {
+	/**
+	 * Arbitrary provider-specific options forwarded unmodified.
+	 */
+	[key: string]: StringLike;
+};
+
+export class WorkersAIEmbeddingModel implements EmbeddingModelV1<string> {
+	/**
+	 * Semantic version of the {@link EmbeddingModelV1} specification implemented
+	 * by this class. It never changes.
+	 */
+	readonly specificationVersion = "v1";
+	readonly modelId: EmbeddingModels;
+	private readonly config: WorkersAIEmbeddingConfig;
+	private readonly settings: WorkersAIEmbeddingSettings;
+
+	/**
+	 * Provider name exposed for diagnostics and error reporting.
+	 */
+	get provider(): string {
+		return this.config.provider;
+	}
+
+	get maxEmbeddingsPerCall(): number {
+		// https://developers.cloudflare.com/workers-ai/platform/limits/#text-embeddings
+		const maxEmbeddingsPerCall = this.modelId === "@cf/baai/bge-large-en-v1.5" ? 1500 : 3000;
+		return this.settings.maxEmbeddingsPerCall ?? maxEmbeddingsPerCall;
+	}
+
+	get supportsParallelCalls(): boolean {
+		return this.settings.supportsParallelCalls ?? true;
+	}
+
+	constructor(
+		modelId: EmbeddingModels,
+		settings: WorkersAIEmbeddingSettings,
+		config: WorkersAIEmbeddingConfig,
+	) {
+		this.modelId = modelId;
+		this.settings = settings;
+		this.config = config;
+	}
+
+	async doEmbed({
+		values,
+	}: Parameters<EmbeddingModelV1<string>["doEmbed"]>[0]): Promise<
+		Awaited<ReturnType<EmbeddingModelV1<string>["doEmbed"]>>
+	> {
+		if (values.length > this.maxEmbeddingsPerCall) {
+			throw new TooManyEmbeddingValuesForCallError({
+				provider: this.provider,
+				modelId: this.modelId,
+				maxEmbeddingsPerCall: this.maxEmbeddingsPerCall,
+				values,
+			});
+		}
+
+		const { gateway, ...passthroughOptions } = this.settings;
+
+		const response = await this.config.binding.run(
+			this.modelId,
+			{
+				text: values,
+			},
+			{ gateway: this.config.gateway ?? gateway, ...passthroughOptions },
+		);
+
+		return {
+			embeddings: response.data,
+		};
+	}
+}

package/src/workersai-chat-language-model.ts

@@ -8,9 +8,9 @@ import { convertToWorkersAIChatMessages } from "./convert-to-workersai-chat-mess
 import type { WorkersAIChatSettings } from "./workersai-chat-settings";
 import type { TextGenerationModels } from "./workersai-models";
 
-import { events } from "fetch-event-stream";
 import { mapWorkersAIUsage } from "./map-workersai-usage";
-import type { WorkersAIChatPrompt } from "./workersai-chat-prompt";
+import { getMappedStream } from "./streaming";
+import { lastMessageWasUser, prepareToolsAndToolChoice, processToolCalls } from "./utils";
 
 type WorkersAIChatConfig = {
 	provider: string;
@@ -138,9 +138,7 @@ export class WorkersAIChatLanguageModel implements LanguageModelV1 {
 		const { gateway, safePrompt, ...passthroughOptions } = this.settings;
 
 		// Extract image from messages if present
-		const { messages, images } = convertToWorkersAIChatMessages(
-			options.prompt,
-		);
+		const { messages, images } = convertToWorkersAIChatMessages(options.prompt);
 
 		// TODO: support for multiple images
 		if (images.length !== 0 && images.length !== 1) {
@@ -189,9 +187,7 @@ export class WorkersAIChatLanguageModel implements LanguageModelV1 {
 		const { args, warnings } = this.getArgs(options);
 
 		// Extract image from messages if present
-		const { messages, images } = convertToWorkersAIChatMessages(
-			options.prompt,
-		);
+		const { messages, images } = convertToWorkersAIChatMessages(options.prompt);
 
 		// [1] When the latest message is not a tool response, we use the regular generate function
 		// and simulate it as a streamed response in order to satisfy the AI SDK's interface for
@@ -265,129 +261,10 @@ export class WorkersAIChatLanguageModel implements LanguageModelV1 {
 			throw new Error("This shouldn't happen");
 		}
 
-		const chunkEvent = events(new Response(response));
-		let usage = { promptTokens: 0, completionTokens: 0 };
-
 		return {
-			stream: new ReadableStream<LanguageModelV1StreamPart>({
-				async start(controller) {
-					for await (const event of chunkEvent) {
-						if (!event.data) {
-							continue;
-						}
-						if (event.data === "[DONE]") {
-							break;
-						}
-						const chunk = JSON.parse(event.data);
-						if (chunk.usage) {
-							usage = mapWorkersAIUsage(chunk);
-						}
-						chunk.response?.length &&
-							controller.enqueue({
-								type: "text-delta",
-								textDelta: chunk.response,
-							});
-					}
-					controller.enqueue({
-						type: "finish",
-						finishReason: "stop",
-						usage: usage,
-					});
-					controller.close();
-				},
-			}),
+			stream: getMappedStream(new Response(response)),
 			rawCall: { rawPrompt: messages, rawSettings: args },
 			warnings,
 		};
 	}
 }
-
-function processToolCalls(output: any) {
-	// Check for OpenAI format tool calls first
-	if (output.tool_calls && Array.isArray(output.tool_calls)) {
-		return output.tool_calls.map((toolCall: any) => {
-			// Handle new format
-			if (toolCall.function && toolCall.id) {
-				return {
-					toolCallType: "function",
-					toolCallId: toolCall.id,
-					toolName: toolCall.function.name,
-					args:
-						typeof toolCall.function.arguments === "string"
-							? toolCall.function.arguments
-							: JSON.stringify(toolCall.function.arguments || {}),
-				};
-			}
-			return {
-				toolCallType: "function",
-				toolCallId: toolCall.name,
-				toolName: toolCall.name,
-				args:
-					typeof toolCall.arguments === "string"
-						? toolCall.arguments
-						: JSON.stringify(toolCall.arguments || {}),
-			};
-		});
-	}
-
-	return [];
-}
-
-function prepareToolsAndToolChoice(
-	mode: Parameters<LanguageModelV1["doGenerate"]>[0]["mode"] & {
-		type: "regular";
-	},
-) {
-	// when the tools array is empty, change it to undefined to prevent errors:
-	const tools = mode.tools?.length ? mode.tools : undefined;
-
-	if (tools == null) {
-		return { tools: undefined, tool_choice: undefined };
-	}
-
-	const mappedTools = tools.map((tool) => ({
-		type: "function",
-		function: {
-			name: tool.name,
-			// @ts-expect-error - description is not a property of tool
-			description: tool.description,
-			// @ts-expect-error - parameters is not a property of tool
-			parameters: tool.parameters,
-		},
-	}));
-
-	const toolChoice = mode.toolChoice;
-
-	if (toolChoice == null) {
-		return { tools: mappedTools, tool_choice: undefined };
-	}
-
-	const type = toolChoice.type;
-
-	switch (type) {
-		case "auto":
-			return { tools: mappedTools, tool_choice: type };
-		case "none":
-			return { tools: mappedTools, tool_choice: type };
-		case "required":
-			return { tools: mappedTools, tool_choice: "any" };
-
-		// workersAI does not support tool mode directly,
-		// so we filter the tools and force the tool choice through 'any'
-		case "tool":
-			return {
-				tools: mappedTools.filter(
-					(tool) => tool.function.name === toolChoice.toolName,
-				),
-				tool_choice: "any",
-			};
-		default: {
-			const exhaustiveCheck = type satisfies never;
-			throw new Error(`Unsupported tool choice type: ${exhaustiveCheck}`);
-		}
-	}
-}
-
-function lastMessageWasUser(messages: WorkersAIChatPrompt) {
-	return messages.length > 0 && messages[messages.length - 1].role === "user";
-}

package/src/workersai-models.ts

@@ -11,4 +11,9 @@ export type TextGenerationModels = Exclude<
  */
 export type ImageGenerationModels = value2key<AiModels, BaseAiTextToImage>;
 
+/**
+ * The names of the BaseAiTextToEmbeddings models.
+ */
+export type EmbeddingModels = value2key<AiModels, BaseAiTextEmbeddings>;
+
 type value2key<T, V> = { [K in keyof T]: T[K] extends V ? K : never }[keyof T];