diff --git a/packages/ai/README.md b/packages/ai/README.md index 042c1bb5..6228dfb3 100644 --- a/packages/ai/README.md +++ b/packages/ai/README.md @@ -18,7 +18,6 @@ Unified LLM API with automatic model discovery, provider configuration, token an - [Image Input](#image-input) - [Image Generation](#image-generation) - [Basic Image Generation](#basic-image-generation) - - [Streaming Image Events](#streaming-image-events) - [Notes and Limitations](#notes-and-limitations) - [Thinking/Reasoning](#thinkingreasoning) - [Unified Interface](#unified-interface-streamsimplecompletesimple) @@ -425,9 +424,9 @@ for (const block of response.content) { ## Image Generation -Image generation uses a separate API surface from text/chat generation. Use `getImageModel()` / `getImageModels()` / `getImageProviders()` to discover image-generation models, `images()` to stream events, and `generateImages()` to get the final result. +Image generation uses a separate API surface from text/chat generation. Use `getImageModel()` / `getImageModels()` / `getImageProviders()` to discover image-generation models, and `generateImages()` to get the final result. -Do not use `stream()` or `complete()` for image generation. Request options, error handling, and event/result flow follow the same overall pattern as the streaming APIs, but use image-specific types and events. +Do not use `stream()` or `complete()` for image generation. Image generation is a one-shot API: `generateImages()` waits for the provider response and returns the final `AssistantImages` result. ### Basic Image Generation @@ -475,59 +474,12 @@ console.log(model.input); // ['text', 'image'] console.log(model.output); // ['image'] or ['image', 'text'] ``` -### Streaming Image Events - -`images()` returns an `AssistantImagesEventStream` with image-specific events: - -```typescript -import { getImageModel, images } from '@mariozechner/pi-ai'; - -const model = getImageModel('openrouter', 'google/gemini-2.5-flash-image'); -const s = images(model, { - input: [{ type: 'text', text: 'Generate a simple green triangle.' }] -}, { - apiKey: process.env.OPENROUTER_API_KEY -}); - -for await (const event of s) { - switch (event.type) { - case 'start': - console.log(`Starting image generation with ${event.partial.model}`); - break; - case 'image_start': - console.log(`Image block started at index ${event.contentIndex}`); - break; - case 'image_end': - console.log(`Image block complete: ${event.image.mimeType}`); - break; - case 'done': - console.log(`Finished: ${event.reason}`); - break; - case 'error': - console.error(`Error: ${event.error.errorMessage}`); - break; - } -} - -const finalResult = await s.result(); -``` - -| Event Type | Description | Key Properties | -|------------|-------------|----------------| -| `start` | Stream begins | `partial`: Initial `AssistantImages` structure | -| `image_start` | Image block starts | `contentIndex`: Position in output array | -| `image_end` | Image block complete | `image`: Final `ImageContent`, `contentIndex`: Position | -| `done` | Stream complete | `reason`: Stop reason (`"stop"`), `images`: Final `AssistantImages` | -| `error` | Error occurred | `reason`: Error type (`"error"` or `"aborted"`), `error`: Final `AssistantImages` | - -Text output, if the model returns any, is available on the final `AssistantImages.output` array. - ### Notes and Limitations - Use `getImageModel(...)`, not `getModel(...)`. -- Use `images()` / `generateImages()`, not `stream()` / `complete()`. +- Use `generateImages()`, not `stream()` / `complete()`. - Image-generation models do not participate in tool calling. -- Outputs are returned as base64-encoded `ImageContent` blocks in `AssistantImages.output`. +- Outputs are returned in `AssistantImages.output` and can include both base64-encoded `ImageContent` blocks and `TextContent` blocks. - Some models return only images, others return images plus text. Check `model.output`. - Some models accept image input, others are text-to-image only. Check `model.input`. - Like the streaming APIs, image generation supports options such as `apiKey`, `signal`, `headers`, `onPayload`, and `onResponse`, and results may include `stopReason`, `responseId`, and `usage`. diff --git a/packages/ai/src/images-api-registry.ts b/packages/ai/src/images-api-registry.ts index 395a060d..2a0c0fd9 100644 --- a/packages/ai/src/images-api-registry.ts +++ b/packages/ai/src/images-api-registry.ts @@ -1,26 +1,19 @@ -import type { - AssistantImagesEventStream, - ImagesApi, - ImagesContext, - ImagesFunction, - ImagesModel, - ImagesOptions, -} from "./types.js"; +import type { AssistantImages, ImagesApi, ImagesContext, ImagesFunction, ImagesModel, ImagesOptions } from "./types.js"; export type ImagesApiFunction = ( model: ImagesModel, context: ImagesContext, options?: ImagesOptions, -) => AssistantImagesEventStream; +) => Promise; export interface ImagesApiProvider { api: TApi; - images: ImagesFunction; + generateImages: ImagesFunction; } interface ImagesApiProviderInternal { api: ImagesApi; - images: ImagesApiFunction; + generateImages: ImagesApiFunction; } type RegisteredImagesApiProvider = { @@ -30,15 +23,15 @@ type RegisteredImagesApiProvider = { const imagesApiProviderRegistry = new Map(); -function wrapImages( +function wrapGenerateImages( api: TApi, - images: ImagesFunction, + generateImages: ImagesFunction, ): ImagesApiFunction { return (model, context, options) => { if (model.api !== api) { throw new Error(`Mismatched api: ${model.api} expected ${api}`); } - return images(model as ImagesModel, context, options as TOptions); + return generateImages(model as ImagesModel, context, options as TOptions); }; } @@ -49,7 +42,7 @@ export function registerImagesApiProvider( - model: ImagesModel, - context: ImagesContext, - options?: ProviderImagesOptions, -): AssistantImagesEventStream { - const provider = resolveImagesApiProvider(model.api); - return provider.images(model, context, options as ImagesOptions); -} - export async function generateImages( model: ImagesModel, context: ImagesContext, options?: ProviderImagesOptions, ): Promise { - const s = images(model, context, options); - return s.result(); + const provider = resolveImagesApiProvider(model.api); + return provider.generateImages(model, context, options); } diff --git a/packages/ai/src/providers/images/openrouter.ts b/packages/ai/src/providers/images/openrouter.ts index c1a5d966..8a7f91d3 100644 --- a/packages/ai/src/providers/images/openrouter.ts +++ b/packages/ai/src/providers/images/openrouter.ts @@ -16,7 +16,6 @@ import type { 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"; @@ -36,87 +35,70 @@ type OpenRouterImageGenerationResponse = ChatCompletion & { choices: OpenRouterImageGenerationChoice[]; }; -export const imagesOpenRouter: ImagesFunction<"openrouter-images", ImagesOptions> = ( +export const generateImagesOpenRouter: ImagesFunction<"openrouter-images", ImagesOptions> = async ( model: ImagesModel<"openrouter-images">, context: ImagesContext, options?: ImagesOptions, ) => { - const stream = new AssistantImagesEventStream(); + const output: AssistantImages = { + api: model.api, + provider: model.provider, + model: model.id, + output: [], + stopReason: "stop", + timestamp: Date.now(), + }; - (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); + 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); - return stream; + 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; + output.output.push({ + type: "image", + mimeType: matches[1], + data: matches[2], + } satisfies ImageContent); + } + } + + return output; + } catch (error) { + output.stopReason = options?.signal?.aborted ? "aborted" : "error"; + output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error); + return output; + } }; function createClient( diff --git a/packages/ai/src/providers/images/register-builtins.ts b/packages/ai/src/providers/images/register-builtins.ts index 7a874b2a..f7d1adc3 100644 --- a/packages/ai/src/providers/images/register-builtins.ts +++ b/packages/ai/src/providers/images/register-builtins.ts @@ -1,74 +1,10 @@ import { 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.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.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; -}; +import { generateImagesOpenRouter } from "./openrouter.js"; export function registerBuiltInImagesApiProviders(): void { registerImagesApiProvider({ api: "openrouter-images", - images: imagesOpenRouter, + generateImages: generateImagesOpenRouter, }); } diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 29bc6bb8..db6ef893 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -1,10 +1,7 @@ import type { AssistantMessageDiagnostic } from "./utils/diagnostics.js"; -import type { AssistantImagesEventStream, AssistantMessageEventStream } from "./utils/event-stream.js"; +import type { AssistantMessageEventStream } from "./utils/event-stream.js"; -export type { - AssistantImagesEventStream, - AssistantMessageEventStream, -} from "./utils/event-stream.js"; +export type { AssistantMessageEventStream } from "./utils/event-stream.js"; export type KnownApi = | "openai-completions" @@ -215,7 +212,7 @@ export type ImagesFunction, context: ImagesContext, options?: TOptions, -) => AssistantImagesEventStream; +) => Promise; export interface TextSignatureV1 { v: 1; @@ -360,13 +357,6 @@ export type AssistantMessageEvent = | { type: "done"; reason: Extract; message: AssistantMessage } | { type: "error"; reason: Extract; error: AssistantMessage }; -export type AssistantImagesEvent = - | { type: "start"; partial: AssistantImages } - | { type: "image_start"; contentIndex: number; partial: AssistantImages } - | { type: "image_end"; contentIndex: number; image: ImageContent; partial: AssistantImages } - | { type: "done"; reason: Extract; images: AssistantImages } - | { type: "error"; reason: Extract; error: AssistantImages }; - /** * Compatibility settings for OpenAI-compatible completions APIs. * Use this to override URL-based auto-detection for custom providers. diff --git a/packages/ai/src/utils/event-stream.ts b/packages/ai/src/utils/event-stream.ts index 9fb1c114..f4a7ceba 100644 --- a/packages/ai/src/utils/event-stream.ts +++ b/packages/ai/src/utils/event-stream.ts @@ -1,4 +1,4 @@ -import type { AssistantImages, AssistantImagesEvent, AssistantMessage, AssistantMessageEvent } from "../types.js"; +import type { AssistantMessage, AssistantMessageEvent } from "../types.js"; // Generic event stream class for async iteration export class EventStream implements AsyncIterable { @@ -81,28 +81,7 @@ export class AssistantMessageEventStream extends EventStream { - constructor() { - super( - (event) => event.type === "done" || event.type === "error", - (event) => { - if (event.type === "done") { - return event.images; - } else if (event.type === "error") { - return event.error; - } - throw new Error("Unexpected event type for final result"); - }, - ); - } -} - /** Factory function for AssistantMessageEventStream (for use in extensions) */ export function createAssistantMessageEventStream(): AssistantMessageEventStream { return new AssistantMessageEventStream(); } - -/** Factory function for AssistantImagesEventStream (for use in extensions) */ -export function createAssistantImagesEventStream(): AssistantImagesEventStream { - return new AssistantImagesEventStream(); -} diff --git a/packages/ai/test/images.test.ts b/packages/ai/test/images.test.ts index 84aac37a..59912c70 100644 --- a/packages/ai/test/images.test.ts +++ b/packages/ai/test/images.test.ts @@ -3,7 +3,7 @@ import { dirname, join } from "node:path"; import { fileURLToPath } from "node:url"; import { describe, expect, it } from "vitest"; import { getImageModel } from "../src/image-models.js"; -import { generateImages, images } from "../src/images.js"; +import { generateImages } from "../src/images.js"; import type { ImageContent, ImagesContext, ImagesModel, ProviderImagesOptions } from "../src/types.js"; const __filename = fileURLToPath(import.meta.url); @@ -24,36 +24,6 @@ async function basicImageGeneration(model: ImagesModel(model: ImagesModel, options?: ImagesOptionsWithExtras) { - const context: ImagesContext = { - input: [{ type: "text", text: "Generate a simple blue square on a plain white background. No text." }], - }; - - const s = images(model, context, options); - let started = false; - let imageStarted = false; - let imageCompleted = false; - - for await (const event of s) { - if (event.type === "start") { - started = true; - } else if (event.type === "image_start") { - imageStarted = true; - } else if (event.type === "image_end") { - imageCompleted = true; - expect(event.image.type).toBe("image"); - expect(event.image.data.length).toBeGreaterThan(0); - } - } - - const response = await s.result(); - expect(response.stopReason, `Error: ${response.errorMessage}`).toBe("stop"); - expect(started).toBe(true); - expect(imageStarted).toBe(true); - expect(imageCompleted).toBe(true); - expect(response.output.some((item) => item.type === "image")).toBe(true); -} - async function handleTextAndImageOutput( model: ImagesModel, options?: ImagesOptionsWithExtras, @@ -108,10 +78,6 @@ describe("Images E2E Tests", () => { await basicImageGeneration(model); }); - it("should handle image streaming", { retry: 3 }, async () => { - await handleStreaming(model); - }); - it("should handle text plus image output", { retry: 3 }, async () => { await handleTextAndImageOutput(model); }); diff --git a/packages/ai/test/openrouter-images.test.ts b/packages/ai/test/openrouter-images.test.ts index b4a79127..de2ddbba 100644 --- a/packages/ai/test/openrouter-images.test.ts +++ b/packages/ai/test/openrouter-images.test.ts @@ -1,17 +1,28 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; -import { generateImages, images } from "../src/images.js"; +import { generateImages } from "../src/images.js"; import type { ImagesContext, ImagesModel } from "../src/types.js"; const mockState = vi.hoisted(() => ({ lastParams: undefined as unknown, + lastRequestOptions: undefined as unknown, })); vi.mock("openai", () => { class FakeOpenAI { chat = { completions: { - create: (params: unknown) => { + create: (params: unknown, requestOptions?: unknown) => { mockState.lastParams = params; + mockState.lastRequestOptions = requestOptions; + const signal = (requestOptions as { signal?: AbortSignal } | undefined)?.signal; + if (signal?.aborted) { + const error = new Error("Request aborted"); + return { + withResponse: async () => { + throw error; + }, + }; + } const response = { id: "img-1", usage: { @@ -50,9 +61,10 @@ vi.mock("openai", () => { describe("openrouter images", () => { beforeEach(() => { mockState.lastParams = undefined; + mockState.lastRequestOptions = undefined; }); - it("emits image events and returns text plus images in final output", async () => { + it("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", @@ -68,14 +80,8 @@ describe("openrouter images", () => { input: [{ type: "text", text: "Generate a dog" }], }; - 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"]); + const output = await generateImages(model, context, { apiKey: "test" }); + expect(output.stopReason).toBe("stop"); 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=" }); @@ -90,6 +96,29 @@ describe("openrouter images", () => { expect(params.messages?.[0]?.content?.[0]).toMatchObject({ type: "text", text: "Generate a dog" }); }); + it("passes through abort signal and returns aborted 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 dog" }], + }; + const controller = new AbortController(); + controller.abort(); + + const output = await generateImages(model, context, { apiKey: "test", signal: controller.signal }); + expect(output.stopReason).toBe("aborted"); + expect(output.errorMessage).toBe("Request aborted"); + expect(mockState.lastRequestOptions).toMatchObject({ signal: controller.signal }); + }); + it("generateImages resolves the final assistant images result", async () => { const model: ImagesModel<"openrouter-images"> = { id: "black-forest-labs/flux.2-pro",