workersai-chat-language-model.ts

import type { LanguageModelV3, SharedV3Warning, LanguageModelV3StreamPart } from "@ai-sdk/provider";
import { generateId } from "ai";
import { convertToWorkersAIChatMessages } from "./convert-to-workersai-chat-messages";
import { mapWorkersAIFinishReason } from "./map-workersai-finish-reason";
import { mapWorkersAIUsage } from "./map-workersai-usage";
import { getMappedStream, prependStreamStart } from "./streaming";
import {
	normalizeMessagesForBinding,
	prepareToolsAndToolChoice,
	processText,
	processToolCalls,
} from "./utils";
import type { WorkersAIChatSettings } from "./workersai-chat-settings";
import type { TextGenerationModels } from "./workersai-models";

type WorkersAIChatConfig = {
	provider: string;
	binding: Ai;
	gateway?: GatewayOptions;
	/** True when using a real Workers AI binding (not the REST shim). */
	isBinding: boolean;
};

export class WorkersAIChatLanguageModel implements LanguageModelV3 {
	readonly specificationVersion = "v3";
	readonly defaultObjectGenerationMode = "json";

	readonly supportedUrls: Record<string, RegExp[]> | PromiseLike<Record<string, RegExp[]>> = {};

	readonly modelId: TextGenerationModels;
	readonly settings: WorkersAIChatSettings;

	private readonly config: WorkersAIChatConfig;

	constructor(
		modelId: TextGenerationModels,
		settings: WorkersAIChatSettings,
		config: WorkersAIChatConfig,
	) {
		this.modelId = modelId;
		this.settings = settings;
		this.config = config;
	}

	get provider(): string {
		return this.config.provider;
	}

	private getArgs({
		responseFormat,
		tools,
		toolChoice,
		maxOutputTokens,
		temperature,
		topP,
		frequencyPenalty,
		presencePenalty,
		seed,
	}: Parameters<LanguageModelV3["doGenerate"]>[0]) {
		const type = responseFormat?.type ?? "text";

		const warnings: SharedV3Warning[] = [];

		if (frequencyPenalty != null) {
			warnings.push({ feature: "frequencyPenalty", type: "unsupported" });
		}

		if (presencePenalty != null) {
			warnings.push({ feature: "presencePenalty", type: "unsupported" });
		}

		const baseArgs = {
			max_tokens: maxOutputTokens,
			model: this.modelId,
			random_seed: seed,
			safe_prompt: this.settings.safePrompt,
			temperature,
			top_p: topP,
		};

		switch (type) {
			case "text": {
				return {
					args: {
						...baseArgs,
						response_format: undefined as
							| { type: string; json_schema?: unknown }
							| undefined,
						...prepareToolsAndToolChoice(tools, toolChoice),
					},
					warnings,
				};
			}

			case "json": {
				return {
					args: {
						...baseArgs,
						response_format: {
							type: "json_schema",
							json_schema:
								responseFormat?.type === "json" ? responseFormat.schema : undefined,
						},
						tools: undefined,
						tool_choice: undefined,
					},
					warnings,
				};
			}

			default: {
				const exhaustiveCheck = type satisfies never;
				throw new Error(`Unsupported type: ${exhaustiveCheck}`);
			}
		}
	}

	/**
	 * Build the inputs object for `binding.run()`, shared by doGenerate and doStream.
	 *
	 * Images are embedded inline in messages as OpenAI-compatible content
	 * arrays with `image_url` parts. Both the REST API and the binding
	 * accept this format at runtime.
	 *
	 * The binding path additionally normalises null content to empty strings.
	 *
	 * Reasoning controls (`reasoning_effort`, `chat_template_kwargs`) are
	 * forwarded here from settings. These belong on the INPUTS object, not on
	 * the 3rd-arg options / REST query string — see
	 * https://github.com/cloudflare/ai/issues/501. Per-call values from
	 * `providerOptions["workers-ai"]` override settings.
	 *
	 * `reasoning_effort: null` is a valid value ("disable reasoning"), so we
	 * check `!== undefined` rather than truthiness.
	 */
	private buildRunInputs(
		args: ReturnType<typeof this.getArgs>["args"],
		messages: ReturnType<typeof convertToWorkersAIChatMessages>["messages"],
		options?: { stream?: boolean; providerOptions?: Record<string, unknown> },
	) {
		// The AI SDK types this as `Record<string, JSONObject>` but we defensively
		// accept anything and only treat it as a lookup if it's a plain object.
		// `"key" in x` throws for primitives, so we can't skip the typeof guard.
		const rawPerCall = options?.providerOptions?.["workers-ai"];
		const perCall: Record<string, unknown> =
			rawPerCall !== null && typeof rawPerCall === "object" && !Array.isArray(rawPerCall)
				? (rawPerCall as Record<string, unknown>)
				: {};
		const reasoningEffort =
			"reasoning_effort" in perCall ? perCall.reasoning_effort : this.settings.reasoning_effort;
		const chatTemplateKwargs =
			"chat_template_kwargs" in perCall
				? perCall.chat_template_kwargs
				: this.settings.chat_template_kwargs;

		return {
			max_tokens: args.max_tokens,
			messages: this.config.isBinding ? normalizeMessagesForBinding(messages) : messages,
			temperature: args.temperature,
			tools: args.tools,
			...(args.tool_choice ? { tool_choice: args.tool_choice } : {}),
			top_p: args.top_p,
			...(args.response_format ? { response_format: args.response_format } : {}),
			...(options?.stream ? { stream: true } : {}),
			...(reasoningEffort !== undefined ? { reasoning_effort: reasoningEffort } : {}),
			...(chatTemplateKwargs !== undefined
				? { chat_template_kwargs: chatTemplateKwargs }
				: {}),
		};
	}

	/**
	 * Get passthrough options for binding.run() from settings.
	 *
	 * `reasoning_effort` and `chat_template_kwargs` are explicitly excluded
	 * here — they belong on the `inputs` object (see `buildRunInputs`), not on
	 * the `options` (3rd) arg of binding.run() or the REST query string.
	 */
	private getRunOptions() {
		const {
			gateway,
			safePrompt: _safePrompt,
			sessionAffinity,
			extraHeaders,
			reasoning_effort: _reasoningEffort,
			chat_template_kwargs: _chatTemplateKwargs,
			...passthroughOptions
		} = this.settings;

		const mergedHeaders = {
			...(extraHeaders && typeof extraHeaders === "object"
				? (extraHeaders as Record<string, string>)
				: {}),
			...(sessionAffinity ? { "x-session-affinity": sessionAffinity } : {}),
		};

		return {
			gateway: this.config.gateway ?? gateway,
			...(Object.keys(mergedHeaders).length > 0 ? { extraHeaders: mergedHeaders } : {}),
			...passthroughOptions,
		};
	}

	async doGenerate(
		options: Parameters<LanguageModelV3["doGenerate"]>[0],
	): Promise<Awaited<ReturnType<LanguageModelV3["doGenerate"]>>> {
		const { args, warnings } = this.getArgs(options);
		const { messages } = convertToWorkersAIChatMessages(options.prompt);

		const inputs = this.buildRunInputs(args, messages, {
			providerOptions: options.providerOptions,
		});
		const runOptions = this.getRunOptions();

		const output = await this.config.binding.run(
			args.model as keyof AiModels,
			inputs as AiModels[keyof AiModels]["inputs"],
			{
				...runOptions,
				signal: options.abortSignal,
			} as AiOptions,
		);

		if (output instanceof ReadableStream) {
			throw new Error(
				"Unexpected streaming response from non-streaming request. Check that `stream: true` was not passed.",
			);
		}

		const outputRecord = output as Record<string, unknown>;
		const choices = outputRecord.choices as
			| Array<{
					message?: { reasoning_content?: string; reasoning?: string };
			  }>
			| undefined;
		const reasoningContent =
			choices?.[0]?.message?.reasoning_content ?? choices?.[0]?.message?.reasoning;

		return {
			finishReason: mapWorkersAIFinishReason(outputRecord),
			content: [
				...(reasoningContent
					? [{ type: "reasoning" as const, text: reasoningContent }]
					: []),
				{
					type: "text",
					text: processText(outputRecord) ?? "",
				},
				...processToolCalls(outputRecord),
			],
			usage: mapWorkersAIUsage(output as Record<string, unknown>),
			warnings,
		};
	}

	async doStream(
		options: Parameters<LanguageModelV3["doStream"]>[0],
	): Promise<Awaited<ReturnType<LanguageModelV3["doStream"]>>> {
		const { args, warnings } = this.getArgs(options);
		const { messages } = convertToWorkersAIChatMessages(options.prompt);

		const inputs = this.buildRunInputs(args, messages, {
			stream: true,
			providerOptions: options.providerOptions,
		});
		const runOptions = this.getRunOptions();

		const response = await this.config.binding.run(
			args.model as keyof AiModels,
			inputs as AiModels[keyof AiModels]["inputs"],
			{
				...runOptions,
				signal: options.abortSignal,
			} as AiOptions,
		);

		// If the binding returned a stream, pipe it through the SSE mapper
		if (response instanceof ReadableStream) {
			return {
				stream: prependStreamStart(getMappedStream(response), warnings),
			};
		}

		// Graceful degradation: some models return a non-streaming response even
		// when stream:true is requested. Wrap the complete response as a stream.
		const outputRecord = response as Record<string, unknown>;
		const choices = outputRecord.choices as
			| Array<{
					message?: { reasoning_content?: string; reasoning?: string };
			  }>
			| undefined;
		const reasoningContent =
			choices?.[0]?.message?.reasoning_content ?? choices?.[0]?.message?.reasoning;

		let textId: string | null = null;
		let reasoningId: string | null = null;

		return {
			stream: new ReadableStream<LanguageModelV3StreamPart>({
				start(controller) {
					controller.enqueue({
						type: "stream-start",
						warnings: warnings as SharedV3Warning[],
					});

					if (reasoningContent) {
						reasoningId = generateId();
						controller.enqueue({ type: "reasoning-start", id: reasoningId });
						controller.enqueue({
							type: "reasoning-delta",
							id: reasoningId,
							delta: reasoningContent,
						});
						controller.enqueue({ type: "reasoning-end", id: reasoningId });
					}

					const text = processText(outputRecord);
					if (text) {
						textId = generateId();
						controller.enqueue({ type: "text-start", id: textId });
						controller.enqueue({ type: "text-delta", id: textId, delta: text });
						controller.enqueue({ type: "text-end", id: textId });
					}

					for (const toolCall of processToolCalls(outputRecord)) {
						controller.enqueue(toolCall);
					}

					controller.enqueue({
						type: "finish",
						finishReason: mapWorkersAIFinishReason(outputRecord),
						usage: mapWorkersAIUsage(response as Record<string, unknown>),
					});
					controller.close();
				},
			}),
		};
	}
}