feat (provider/gateway): add support for image models

.changeset/friendly-peaches-hammer.md

@@ -0,0 +1,5 @@
+---
+'@ai-sdk/gateway': patch
+---
+
+feat (provider/gateway): add support for image models

packages/gateway/src/gateway-image-model-settings.ts

@@ -0,0 +1 @@
+export type GatewayImageModelId = 'google/imagen-4.0-generate' | (string & {});

packages/gateway/src/gateway-image-model.ts

@@ -0,0 +1,128 @@
+import type {
+  ImageModelV2,
+  ImageModelV2CallWarning,
+  ImageModelV2ProviderMetadata,
+} from '@ai-sdk/provider';
+import {
+  combineHeaders,
+  createJsonResponseHandler,
+  createJsonErrorResponseHandler,
+  postJsonToApi,
+  resolve,
+  type Resolvable,
+} from '@ai-sdk/provider-utils';
+import { z } from 'zod/v4';
+import type { GatewayConfig } from './gateway-config';
+import { asGatewayError } from './errors';
+import { parseAuthMethod } from './errors/parse-auth-method';
+
+export class GatewayImageModel implements ImageModelV2 {
+  readonly specificationVersion = 'v2';
+  readonly maxImagesPerCall = 4;
+
+  constructor(
+    readonly modelId: string,
+    private readonly config: GatewayConfig & {
+      provider: string;
+      o11yHeaders: Resolvable<Record<string, string>>;
+    },
+  ) {}
+
+  get provider(): string {
+    return this.config.provider;
+  }
+
+  async doGenerate({
+    prompt,
+    n,
+    size,
+    aspectRatio,
+    seed,
+    providerOptions,
+    headers,
+    abortSignal,
+  }: Parameters<ImageModelV2['doGenerate']>[0]): Promise<
+    Awaited<ReturnType<ImageModelV2['doGenerate']>>
+  > {
+    const resolvedHeaders = await resolve(this.config.headers());
+    try {
+      const {
+        responseHeaders,
+        value: responseBody,
+        rawValue,
+      } = await postJsonToApi({
+        url: this.getUrl(),
+        headers: combineHeaders(
+          resolvedHeaders,
+          headers ?? {},
+          this.getModelConfigHeaders(),
+          await resolve(this.config.o11yHeaders),
+        ),
+        body: {
+          prompt,
+          n,
+          ...(size && { size }),
+          ...(aspectRatio && { aspectRatio }),
+          ...(seed && { seed }),
+          ...(providerOptions ?? {}),
+        },
+        successfulResponseHandler: createJsonResponseHandler(
+          gatewayImageResponseSchema,
+        ),
+        failedResponseHandler: createJsonErrorResponseHandler({
+          errorSchema: z.any(),
+          errorToMessage: data => data,
+        }),
+        ...(abortSignal && { abortSignal }),
+        fetch: this.config.fetch,
+      });
+
+      return {
+        images: responseBody.images, // Always base64 strings from server
+        warnings: responseBody.warnings ?? [],
+        providerMetadata:
+          responseBody.providerMetadata as ImageModelV2ProviderMetadata,
+        response: {
+          timestamp: new Date(),
+          modelId: this.modelId,
+          headers: responseHeaders,
+        },
+      };
+    } catch (error) {
+      throw asGatewayError(error, parseAuthMethod(resolvedHeaders));
+    }
+  }
+
+  private getUrl() {
+    return `${this.config.baseURL}/image-model`;
+  }
+
+  private getModelConfigHeaders() {
+    return {
+      'ai-image-model-specification-version': '2',
+      'ai-model-id': this.modelId,
+    };
+  }
+}
+
+const gatewayImageResponseSchema = z.object({
+  images: z.array(z.string()), // Always base64 strings over the wire
+  warnings: z
+    .array(
+      z.object({
+        type: z.literal('other'),
+        message: z.string(),
+      }),
+    )
+    .optional(),
+  providerMetadata: z
+    .record(
+      z.string(),
+      z
+        .object({
+          images: z.array(z.unknown()),
+        })
+        .and(z.record(z.string(), z.unknown())),
+    )
+    .optional(),
+});

packages/gateway/src/gateway-provider.ts

@@ -1,4 +1,3 @@
-import { NoSuchModelError } from '@ai-sdk/provider';
 import {
   loadOptionalSetting,
   withoutTrailingSlash,
@@ -15,12 +14,15 @@ import {
 } from './gateway-fetch-metadata';
 import { GatewayLanguageModel } from './gateway-language-model';
 import { GatewayEmbeddingModel } from './gateway-embedding-model';
+import { GatewayImageModel } from './gateway-image-model';
 import type { GatewayEmbeddingModelId } from './gateway-embedding-model-settings';
+import type { GatewayImageModelId } from './gateway-image-model-settings';
 import { getVercelOidcToken, getVercelRequestId } from './vercel-environment';
 import type { GatewayModelId } from './gateway-language-model-settings';
 import type {
   LanguageModelV2,
   EmbeddingModelV2,
+  ImageModelV2,
   ProviderV2,
 } from '@ai-sdk/provider';
 
@@ -43,6 +45,11 @@ Creates a model for generating text embeddings.
   textEmbeddingModel(
     modelId: GatewayEmbeddingModelId,
   ): EmbeddingModelV2<string>;
+
+  /**
+Creates a model for generating images.
+*/
+  imageModel(modelId: GatewayImageModelId): ImageModelV2;
 }
 
 export interface GatewayProviderSettings {
@@ -185,8 +192,14 @@ export function createGatewayProvider(
   };
 
   provider.getAvailableModels = getAvailableModels;
-  provider.imageModel = (modelId: string) => {
-    throw new NoSuchModelError({ modelId, modelType: 'imageModel' });
+  provider.imageModel = (modelId: GatewayImageModelId) => {
+    return new GatewayImageModel(modelId, {
+      provider: 'gateway',
+      baseURL,
+      headers: getHeaders,
+      fetch: options.fetch,
+      o11yHeaders: createO11yHeaders(),
+    });
   };
   provider.languageModel = createLanguageModel;
   provider.textEmbeddingModel = (modelId: GatewayEmbeddingModelId) => {