diff --git a/packages/ai/src/images.ts b/packages/ai/src/images.ts index 089c5f31..8595f6b5 100644 --- a/packages/ai/src/images.ts +++ b/packages/ai/src/images.ts @@ -1,3 +1,5 @@ +import "./providers/register-images-builtins.js"; + import { getImagesApiProvider } from "./images-api-registry.js"; import type { AssistantImages, diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 7235ecc7..d3c808e9 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -3,6 +3,8 @@ export { Type } from "typebox"; export * from "./api-registry.js"; export * from "./env-api-keys.js"; +export * from "./images.js"; +export * from "./images-api-registry.js"; export * from "./models.js"; export type { BedrockOptions, BedrockThinkingDisplay } from "./providers/amazon-bedrock.js"; export type { AnthropicEffort, AnthropicOptions, AnthropicThinkingDisplay } from "./providers/anthropic.js"; @@ -16,6 +18,7 @@ export type { OpenAICodexResponsesOptions } from "./providers/openai-codex-respo export type { OpenAICompletionsOptions } from "./providers/openai-completions.js"; export type { OpenAIResponsesOptions } from "./providers/openai-responses.js"; export * from "./providers/register-builtins.js"; +export * from "./providers/register-images-builtins.js"; export * from "./stream.js"; export * from "./types.js"; export * from "./utils/event-stream.js"; diff --git a/packages/ai/src/providers/openrouter-images.ts b/packages/ai/src/providers/openrouter-images.ts new file mode 100644 index 00000000..4dafc196 --- /dev/null +++ b/packages/ai/src/providers/openrouter-images.ts @@ -0,0 +1,202 @@ +import OpenAI from "openai"; +import type { + ChatCompletion, + ChatCompletionContentPart, + ChatCompletionContentPartImage, + ChatCompletionContentPartText, + ChatCompletionCreateParamsNonStreaming, +} from "openai/resources/chat/completions.js"; +import { getEnvApiKey } from "../env-api-keys.js"; +import type { + AssistantImages, + ImageContent, + ImagesContext, + ImagesFunction, + ImagesModel, + ImagesOptions, + TextContent, +} from "../types.js"; +import { AssistantImagesEventStream } from "../utils/event-stream.js"; +import { headersToRecord } from "../utils/headers.js"; +import { sanitizeSurrogates } from "../utils/sanitize-unicode.js"; + +interface OpenRouterGeneratedImage { + image_url?: string | { url?: string }; +} + +type OpenRouterImageGenerationMessage = ChatCompletion["choices"][number]["message"] & { + images?: OpenRouterGeneratedImage[]; +}; + +type OpenRouterImageGenerationChoice = ChatCompletion["choices"][number] & { + message: OpenRouterImageGenerationMessage; +}; + +type OpenRouterImageGenerationResponse = ChatCompletion & { + choices: OpenRouterImageGenerationChoice[]; +}; + +export const imagesOpenRouter: ImagesFunction<"openrouter-images", ImagesOptions> = ( + model: ImagesModel<"openrouter-images">, + context: ImagesContext, + options?: ImagesOptions, +) => { + const stream = new AssistantImagesEventStream(); + + (async () => { + const output: AssistantImages = { + api: model.api, + provider: model.provider, + model: model.id, + output: [], + stopReason: "stop", + timestamp: Date.now(), + }; + + try { + const apiKey = options?.apiKey || getEnvApiKey(model.provider) || ""; + const client = createClient(model, apiKey, options?.headers); + let params = buildParams(model, context); + const nextParams = await options?.onPayload?.(params, model); + if (nextParams !== undefined) { + params = nextParams as typeof params; + } + const requestOptions = { + ...(options?.signal ? { signal: options.signal } : {}), + ...(options?.timeoutMs !== undefined ? { timeout: options.timeoutMs } : {}), + ...(options?.maxRetries !== undefined ? { maxRetries: options.maxRetries } : {}), + }; + const { data: response, response: rawResponse } = await client.chat.completions + .create(params as unknown as ChatCompletionCreateParamsNonStreaming, requestOptions) + .withResponse(); + await options?.onResponse?.( + { status: rawResponse.status, headers: headersToRecord(rawResponse.headers) }, + model, + ); + + stream.push({ type: "start", partial: output }); + + const imageResponse = response as OpenRouterImageGenerationResponse; + output.responseId = imageResponse.id; + if (imageResponse.usage) { + output.usage = parseUsage(imageResponse.usage, model); + } + + const choice = imageResponse.choices[0]; + if (choice) { + const content = choice.message.content; + if (typeof content === "string" && content.length > 0) { + output.output.push({ type: "text", text: content } satisfies TextContent); + } + + for (const image of choice.message.images ?? []) { + const imageUrl = typeof image.image_url === "string" ? image.image_url : image.image_url?.url; + if (!imageUrl?.startsWith("data:")) continue; + const matches = imageUrl.match(/^data:([^;]+);base64,(.+)$/); + if (!matches) continue; + const imageBlock: ImageContent = { + type: "image", + mimeType: matches[1], + data: matches[2], + }; + output.output.push(imageBlock); + const contentIndex = output.output.length - 1; + stream.push({ type: "image_start", contentIndex, partial: output }); + stream.push({ type: "image_end", contentIndex, image: imageBlock, partial: output }); + } + } + + stream.push({ type: "done", reason: "stop", images: output }); + stream.end(); + } catch (error) { + output.stopReason = options?.signal?.aborted ? "aborted" : "error"; + output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error); + stream.push({ type: "error", reason: output.stopReason, error: output }); + stream.end(output); + } + })(); + + return stream; +}; + +function createClient( + model: ImagesModel<"openrouter-images">, + apiKey: string, + optionsHeaders?: Record, +): OpenAI { + return new OpenAI({ + apiKey, + baseURL: model.baseUrl, + dangerouslyAllowBrowser: true, + defaultHeaders: { + ...model.headers, + ...optionsHeaders, + }, + }); +} + +type OpenRouterImagesCreateParams = Omit & { + modalities: Array<"image" | "text">; +}; + +function buildParams(model: ImagesModel<"openrouter-images">, context: ImagesContext): OpenRouterImagesCreateParams { + const content: ChatCompletionContentPart[] = context.input.map((item): ChatCompletionContentPart => { + if (item.type === "text") { + return { + type: "text", + text: sanitizeSurrogates(item.text), + } satisfies ChatCompletionContentPartText; + } + return { + type: "image_url", + image_url: { + url: `data:${item.mimeType};base64,${item.data}`, + }, + } satisfies ChatCompletionContentPartImage; + }); + + return { + model: model.id, + messages: [ + { + role: "user" as const, + content, + }, + ], + stream: false, + modalities: model.output.includes("text") ? ["image", "text"] : ["image"], + }; +} + +function parseUsage( + rawUsage: { + prompt_tokens?: number; + completion_tokens?: number; + prompt_tokens_details?: { cached_tokens?: number; cache_write_tokens?: number }; + }, + model: ImagesModel<"openrouter-images">, +) { + const promptTokens = rawUsage.prompt_tokens || 0; + const reportedCachedTokens = rawUsage.prompt_tokens_details?.cached_tokens || 0; + const cacheWriteTokens = rawUsage.prompt_tokens_details?.cache_write_tokens || 0; + const cacheReadTokens = + cacheWriteTokens > 0 ? Math.max(0, reportedCachedTokens - cacheWriteTokens) : reportedCachedTokens; + const input = Math.max(0, promptTokens - cacheReadTokens - cacheWriteTokens); + const output = rawUsage.completion_tokens || 0; + const usage = { + input, + output, + cacheRead: cacheReadTokens, + cacheWrite: cacheWriteTokens, + totalTokens: input + output + cacheReadTokens + cacheWriteTokens, + cost: { + input: (model.cost.input / 1000000) * input, + output: (model.cost.output / 1000000) * output, + cacheRead: (model.cost.cacheRead / 1000000) * cacheReadTokens, + cacheWrite: (model.cost.cacheWrite / 1000000) * cacheWriteTokens, + total: 0, + }, + }; + usage.cost.total = usage.cost.input + usage.cost.output + usage.cost.cacheRead + usage.cost.cacheWrite; + return usage; +} diff --git a/packages/ai/src/providers/register-images-builtins.ts b/packages/ai/src/providers/register-images-builtins.ts new file mode 100644 index 00000000..9c5c04ef --- /dev/null +++ b/packages/ai/src/providers/register-images-builtins.ts @@ -0,0 +1,80 @@ +import { clearImagesApiProviders, registerImagesApiProvider } from "../images-api-registry.js"; +import type { + AssistantImages, + AssistantImagesEvent, + ImagesContext, + ImagesFunction, + ImagesModel, + ImagesOptions, +} from "../types.js"; +import { AssistantImagesEventStream } from "../utils/event-stream.js"; +import type { imagesOpenRouter as imagesOpenRouterFunction } from "./openrouter-images.js"; + +interface OpenRouterImagesProviderModule { + imagesOpenRouter: typeof imagesOpenRouterFunction; +} + +let openRouterImagesProviderModulePromise: Promise | undefined; + +function forwardImagesStream(target: AssistantImagesEventStream, source: AsyncIterable): void { + (async () => { + for await (const event of source) { + target.push(event); + } + target.end(); + })(); +} + +function createLazyLoadErrorImages(model: ImagesModel<"openrouter-images">, error: unknown): AssistantImages { + return { + api: model.api, + provider: model.provider, + model: model.id, + output: [], + stopReason: "error", + errorMessage: error instanceof Error ? error.message : String(error), + timestamp: Date.now(), + }; +} + +function loadOpenRouterImagesProviderModule(): Promise { + openRouterImagesProviderModulePromise ||= import("./openrouter-images.js").then( + (module) => module as OpenRouterImagesProviderModule, + ); + return openRouterImagesProviderModulePromise; +} + +export const imagesOpenRouter: ImagesFunction<"openrouter-images", ImagesOptions> = ( + model: ImagesModel<"openrouter-images">, + context: ImagesContext, + options?: ImagesOptions, +) => { + const outer = new AssistantImagesEventStream(); + + loadOpenRouterImagesProviderModule() + .then((module) => { + const inner = module.imagesOpenRouter(model, context, options); + forwardImagesStream(outer, inner); + }) + .catch((error) => { + const images = createLazyLoadErrorImages(model, error); + outer.push({ type: "error", reason: "error", error: images }); + outer.end(images); + }); + + return outer; +}; + +export function registerBuiltInImagesApiProviders(): void { + registerImagesApiProvider({ + api: "openrouter-images", + images: imagesOpenRouter, + }); +} + +export function resetImagesApiProviders(): void { + clearImagesApiProviders(); + registerBuiltInImagesApiProviders(); +} + +registerBuiltInImagesApiProviders(); diff --git a/packages/ai/test/openrouter-images.test.ts b/packages/ai/test/openrouter-images.test.ts new file mode 100644 index 00000000..5926fbf4 --- /dev/null +++ b/packages/ai/test/openrouter-images.test.ts @@ -0,0 +1,111 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { completeImages, images } from "../src/images.js"; +import type { ImagesContext, ImagesModel } from "../src/types.js"; + +const mockState = vi.hoisted(() => ({ + lastParams: undefined as unknown, +})); + +vi.mock("openai", () => { + class FakeOpenAI { + chat = { + completions: { + create: (params: unknown) => { + mockState.lastParams = params; + const response = { + id: "img-1", + usage: { + prompt_tokens: 12, + completion_tokens: 34, + prompt_tokens_details: { cached_tokens: 0 }, + }, + choices: [ + { + message: { + content: "Here is your image.", + images: [{ image_url: "data:image/png;base64,ZmFrZS1wbmc=" }], + }, + }, + ], + }; + const promise = Promise.resolve(response) as Promise & { + withResponse: () => Promise<{ + data: typeof response; + response: { status: number; headers: Headers }; + }>; + }; + promise.withResponse = async () => ({ + data: response, + response: { status: 200, headers: new Headers() }, + }); + return promise; + }, + }, + }; + } + + return { default: FakeOpenAI }; +}); + +describe("openrouter images", () => { + beforeEach(() => { + mockState.lastParams = undefined; + }); + + it("emits image events and returns text plus images in final output", async () => { + const model: ImagesModel<"openrouter-images"> = { + id: "google/gemini-3.1-flash-image-preview", + name: "Gemini 3.1 Flash Image Preview", + api: "openrouter-images", + provider: "openrouter", + baseUrl: "https://openrouter.ai/api/v1", + input: ["text", "image"], + output: ["text", "image"], + cost: { input: 0.015, output: 0.03, cacheRead: 0, cacheWrite: 0 }, + headers: { "HTTP-Referer": "https://example.com" }, + }; + const context: ImagesContext = { + input: [{ type: "text", text: "Generate a cat" }], + }; + + const result = images(model, context, { apiKey: "test" }); + const eventTypes: string[] = []; + for await (const event of result) { + eventTypes.push(event.type); + } + + const output = await result.result(); + expect(eventTypes).toEqual(["start", "image_start", "image_end", "done"]); + expect(output.responseId).toBe("img-1"); + expect(output.output[0]).toMatchObject({ type: "text", text: "Here is your image." }); + expect(output.output[1]).toMatchObject({ type: "image", mimeType: "image/png", data: "ZmFrZS1wbmc=" }); + + const params = mockState.lastParams as { + stream?: boolean; + modalities?: string[]; + messages?: [{ content?: [{ type: string; text?: string }] }]; + }; + expect(params.stream).toBe(false); + expect(params.modalities).toEqual(["image", "text"]); + expect(params.messages?.[0]?.content?.[0]).toMatchObject({ type: "text", text: "Generate a cat" }); + }); + + it("completeImages resolves the final assistant images result", async () => { + const model: ImagesModel<"openrouter-images"> = { + id: "black-forest-labs/flux.2-pro", + name: "FLUX.2 Pro", + api: "openrouter-images", + provider: "openrouter", + baseUrl: "https://openrouter.ai/api/v1", + input: ["text", "image"], + output: ["image"], + cost: { input: 0.015, output: 0.03, cacheRead: 0, cacheWrite: 0 }, + }; + const context: ImagesContext = { + input: [{ type: "text", text: "Generate a cat" }], + }; + + const output = await completeImages(model, context, { apiKey: "test" }); + expect(output.output.some((item) => item.type === "image")).toBe(true); + }); +});