workers-ai-provider 0.2.0 → 0.2.1

package/src/index.ts

@@ -1,103 +1,105 @@
+import { createRun } from "./utils";
 import { WorkersAIChatLanguageModel } from "./workersai-chat-language-model";
 import type { WorkersAIChatSettings } from "./workersai-chat-settings";
-import type { TextGenerationModels } from "./workersai-models";
-import { createRun } from "./utils";
+import { WorkersAIImageModel } from "./workersai-image-model";
+import type { WorkersAIImageSettings } from "./workersai-image-settings";
+import type { ImageGenerationModels, TextGenerationModels } from "./workersai-models";
 
-export type WorkersAISettings =
-  ({
-  /**
-   * Provide a Cloudflare AI binding.
-   */
-  binding: Ai;
+export type WorkersAISettings = (
+	| {
+			/**
+			 * Provide a Cloudflare AI binding.
+			 */
+			binding: Ai;
 
-  /**
-   * Credentials must be absent when a binding is given.
-   */
-  accountId?: never;
-  apiKey?: never;
-}
-  | {
-  /**
-   * Provide Cloudflare API credentials directly. Must be used if a binding is not specified.
-   */
-  accountId: string;
-  apiKey: string;
-  /**
-   * Both binding must be absent if credentials are used directly.
-   */
-  binding?: never;
-}) & {
-  /**
-   * Optionally specify a gateway.
-   */
-  gateway?: GatewayOptions;
+			/**
+			 * Credentials must be absent when a binding is given.
+			 */
+			accountId?: never;
+			apiKey?: never;
+	  }
+	| {
+			/**
+			 * Provide Cloudflare API credentials directly. Must be used if a binding is not specified.
+			 */
+			accountId: string;
+			apiKey: string;
+			/**
+			 * Both binding must be absent if credentials are used directly.
+			 */
+			binding?: never;
+	  }
+) & {
+	/**
+	 * Optionally specify a gateway.
+	 */
+	gateway?: GatewayOptions;
 };
 
 export interface WorkersAI {
-  (
-    modelId: TextGenerationModels,
-    settings?: WorkersAIChatSettings
-  ): WorkersAIChatLanguageModel;
+	(modelId: TextGenerationModels, settings?: WorkersAIChatSettings): WorkersAIChatLanguageModel;
+	/**
+	 * Creates a model for text generation.
+	 **/
+	chat(
+		modelId: TextGenerationModels,
+		settings?: WorkersAIChatSettings,
+	): WorkersAIChatLanguageModel;
 
-  /**
-   * Creates a model for text generation.
-   **/
-  chat(
-    modelId: TextGenerationModels,
-    settings?: WorkersAIChatSettings
-  ): WorkersAIChatLanguageModel;
+	/**
+	 * Creates a model for image generation.
+	 **/
+	image(modelId: ImageGenerationModels, settings?: WorkersAIImageSettings): WorkersAIImageModel;
 }
 
 /**
  * Create a Workers AI provider instance.
  */
 export function createWorkersAI(options: WorkersAISettings): WorkersAI {
-  // Use a binding if one is directly provided. Otherwise use credentials to create
-  // a `run` method that calls the Cloudflare REST API.
-  let binding: Ai | undefined;
+	// Use a binding if one is directly provided. Otherwise use credentials to create
+	// a `run` method that calls the Cloudflare REST API.
+	let binding: Ai | undefined;
 
-  if (options.binding) {
-    binding = options.binding;
-  } else {
-    const { accountId, apiKey } = options;
-    binding = {
-      run: createRun(accountId, apiKey),
-    } as Ai;
-  }
+	if (options.binding) {
+		binding = options.binding;
+	} else {
+		const { accountId, apiKey } = options;
+		binding = {
+			run: createRun({ accountId, apiKey }),
+		} as Ai;
+	}
 
-  if (!binding) {
-    throw new Error(
-      "Either a binding or credentials must be provided."
-    );
-  }
+	if (!binding) {
+		throw new Error("Either a binding or credentials must be provided.");
+	}
 
-  /**
-   * Helper function to create a chat model instance.
-   */
-  const createChatModel = (
-    modelId: TextGenerationModels,
-    settings: WorkersAIChatSettings = {}
-  ) =>
-    new WorkersAIChatLanguageModel(modelId, settings, {
-      provider: "workersai.chat",
-      binding,
-      gateway: options.gateway
-    });
+	const createChatModel = (modelId: TextGenerationModels, settings: WorkersAIChatSettings = {}) =>
+		new WorkersAIChatLanguageModel(modelId, settings, {
+			provider: "workersai.chat",
+			binding,
+			gateway: options.gateway,
+		});
 
-  const provider = function (
-    modelId: TextGenerationModels,
-    settings?: WorkersAIChatSettings
-  ) {
-    if (new.target) {
-      throw new Error(
-        "The WorkersAI model function cannot be called with the new keyword."
-      );
-    }
-    return createChatModel(modelId, settings);
-  };
+	const createImageModel = (
+		modelId: ImageGenerationModels,
+		settings: WorkersAIImageSettings = {},
+	) =>
+		new WorkersAIImageModel(modelId, settings, {
+			provider: "workersai.image",
+			binding,
+			gateway: options.gateway,
+		});
 
-  provider.chat = createChatModel;
+	const provider = (modelId: TextGenerationModels, settings?: WorkersAIChatSettings) => {
+		if (new.target) {
+			throw new Error("The WorkersAI model function cannot be called with the new keyword.");
+		}
+		return createChatModel(modelId, settings);
+	};
 
-  return provider;
-}
+	provider.chat = createChatModel;
+	provider.image = createImageModel;
+	provider.imageModel = createImageModel;
 
+	return provider;
+}

package/src/map-workersai-finish-reason.ts

@@ -1,17 +1,17 @@
 import type { LanguageModelV1FinishReason } from "@ai-sdk/provider";
 
 export function mapWorkersAIFinishReason(
-  finishReason: string | null | undefined
+	finishReason: string | null | undefined,
 ): LanguageModelV1FinishReason {
-  switch (finishReason) {
-    case "stop":
-      return "stop";
-    case "length":
-    case "model_length":
-      return "length";
-    case "tool_calls":
-      return "tool-calls";
-    default:
-      return "other";
-  }
+	switch (finishReason) {
+		case "stop":
+			return "stop";
+		case "length":
+		case "model_length":
+			return "length";
+		case "tool_calls":
+			return "tool-calls";
+		default:
+			return "other";
+	}
 }

package/src/map-workersai-usage.ts

@@ -1,11 +1,15 @@
-export function mapWorkersAIUsage(output:  AiTextGenerationOutput | AiTextToImageOutput) {
-	const usage = (output as { usage: { prompt_tokens: number, completion_tokens: number} }).usage ?? {
+export function mapWorkersAIUsage(output: AiTextGenerationOutput | AiTextToImageOutput) {
+	const usage = (
+		output as {
+			usage: { prompt_tokens: number; completion_tokens: number };
+		}
+	).usage ?? {
 		prompt_tokens: 0,
 		completion_tokens: 0,
 	};
 
 	return {
 		promptTokens: usage.prompt_tokens,
-		completionTokens: usage.completion_tokens
-	}
+		completionTokens: usage.completion_tokens,
+	};
 }

package/src/utils.ts

@@ -1,64 +1,101 @@
 /**
- * Creates a run method that mimics the Cloudflare Workers AI binding,
- * but uses the Cloudflare REST API under the hood.
+ * General AI run interface with overloads to handle distinct return types.
  *
- * @param accountId - Your Cloudflare account identifier.
- * @param apiKey - Your Cloudflare API token/key with appropriate permissions.
- * @returns An function matching `Ai['run']`.
+ * The behaviour depends on the combination of parameters:
+ * 1. `returnRawResponse: true` => returns the raw Response object.
+ * 2. `stream: true`           => returns a ReadableStream (if available).
+ * 3. Otherwise                => returns post-processed AI results.
  */
-export function createRun(accountId: string, apiKey: string): AiRun {
-  return async <Name extends keyof AiModels>(model: Name, inputs: AiModels[Name]["inputs"], options?: AiOptions | undefined) => {
-    const url = `https://api.cloudflare.com/client/v4/accounts/${accountId}/ai/run/${model}`;
-    const body = JSON.stringify(inputs);
+export interface AiRun {
+	// (1) Return raw Response if `options.returnRawResponse` is `true`.
+	<Name extends keyof AiModels>(
+		model: Name,
+		inputs: AiModels[Name]["inputs"],
+		options: AiOptions & { returnRawResponse: true },
+	): Promise<Response>;
 
-    const headers = {
-      "Content-Type": "application/json",
-      Authorization: `Bearer ${apiKey}`,
-    };
+	// (2) Return a stream if the input has `stream: true`.
+	<Name extends keyof AiModels>(
+		model: Name,
+		inputs: AiModels[Name]["inputs"] & { stream: true },
+		options?: AiOptions,
+	): Promise<ReadableStream<Uint8Array>>;
 
-    const response = await fetch(url, {
-      method: "POST",
-      headers,
-      body,
-    }) as Response;
-
-    if (options?.returnRawResponse) {
-      return response;
-    }
+	// (3) Return post-processed outputs by default.
+	<Name extends keyof AiModels>(
+		model: Name,
+		inputs: AiModels[Name]["inputs"],
+		options?: AiOptions,
+	): Promise<AiModels[Name]["postProcessedOutputs"]>;
+}
 
-    if ((inputs as AiTextGenerationInput).stream === true) {
-      // If there's a stream, return the raw body so the caller can process it
-      if (response.body) {
-        return response.body;
-      }
-      throw new Error("No readable body available for streaming.");
-    }
+/**
+ * Parameters for configuring the Cloudflare-based AI runner.
+ */
+export interface CreateRunConfig {
+	/** Your Cloudflare account identifier. */
+	accountId: string;
 
-    // Otherwise, parse JSON and return the data.result
-    const data = await response.json<{ result: AiModels[Name]["postProcessedOutputs"] }>();
-    return data.result;
-  };
+	/** Cloudflare API token/key with appropriate permissions. */
+	apiKey: string;
 }
 
-interface AiRun {
-  // (1) Return raw response if `options.returnRawResponse` is `true`.
-  <Name extends keyof AiModels>(
-    model: Name,
-    inputs: AiModels[Name]["inputs"],
-    options: AiOptions & { returnRawResponse: true }
-  ): Promise<Response>;
+/**
+ * Creates a run method that emulates the Cloudflare Workers AI binding,
+ * but uses the Cloudflare REST API under the hood. Headers and abort
+ * signals are configured at creation time, rather than per-request.
+ *
+ * @param config An object containing:
+ *   - `accountId`: Cloudflare account identifier.
+ *   - `apiKey`: Cloudflare API token/key with suitable permissions.
+ *   - `headers`: Optional custom headers to merge with defaults.
+ *   - `signal`: Optional AbortSignal for request cancellation.
+ *
+ * @returns A function matching the AiRun interface.
+ */
+export function createRun(config: CreateRunConfig): AiRun {
+	const { accountId, apiKey } = config;
+
+	// Return the AiRun-compatible function.
+	return async function run<Name extends keyof AiModels>(
+		model: Name,
+		inputs: AiModels[Name]["inputs"],
+		options?: AiOptions,
+	): Promise<Response | ReadableStream<Uint8Array> | AiModels[Name]["postProcessedOutputs"]> {
+		const url = `https://api.cloudflare.com/client/v4/accounts/${accountId}/ai/run/${model}`;
+
+		// Merge default and custom headers.
+		const headers = {
+			"Content-Type": "application/json",
+			Authorization: `Bearer ${apiKey}`,
+		};
+
+		const body = JSON.stringify(inputs);
+
+		// Execute the POST request. The optional AbortSignal is applied here.
+		const response = await fetch(url, {
+			method: "POST",
+			headers,
+			body,
+		});
+
+		// (1) If the user explicitly requests the raw Response, return it as-is.
+		if (options?.returnRawResponse) {
+			return response;
+		}
 
-  // (2) Return a stream if the input has `stream: true`.
-  <Name extends keyof AiModels>(
-    model: Name,
-    inputs: AiModels[Name]["inputs"] & { stream: true },
-    options?: AiOptions
-  ): Promise<ReadableStream<Uint8Array>>;
+		// (2) If the AI input requests streaming, return the ReadableStream if available.
+		if ((inputs as AiTextGenerationInput).stream === true) {
+			if (response.body) {
+				return response.body;
+			}
+			throw new Error("No readable body available for streaming.");
+		}
 
-  // (3) Return the post-processed outputs by default.
-  <Name extends keyof AiModels>(
-    model: Name,
-    inputs: AiModels[Name]["inputs"],
-    options?: AiOptions
-  ): Promise<AiModels[Name]["postProcessedOutputs"]>;
+		// (3) In all other cases, parse JSON and return the result field.
+		const data = await response.json<{
+			result: AiModels[Name]["postProcessedOutputs"];
+		}>();
+		return data.result;
+	};
 }