@ai-sdk/gateway 0.0.0-70e0935a-20260114150030 → 0.0.0-98261322-20260122142521

package/src/gateway-config.ts

@@ -0,0 +1,7 @@
+import type { FetchFunction, Resolvable } from '@ai-sdk/provider-utils';
+
+export type GatewayConfig = {
+  baseURL: string;
+  headers: () => Resolvable<Record<string, string | undefined>>;
+  fetch?: FetchFunction;
+};

package/src/gateway-embedding-model-settings.ts

@@ -0,0 +1,22 @@
+export type GatewayEmbeddingModelId =
+  | 'alibaba/qwen3-embedding-0.6b'
+  | 'alibaba/qwen3-embedding-4b'
+  | 'alibaba/qwen3-embedding-8b'
+  | 'amazon/titan-embed-text-v2'
+  | 'cohere/embed-v4.0'
+  | 'google/gemini-embedding-001'
+  | 'google/text-embedding-005'
+  | 'google/text-multilingual-embedding-002'
+  | 'mistral/codestral-embed'
+  | 'mistral/mistral-embed'
+  | 'openai/text-embedding-3-large'
+  | 'openai/text-embedding-3-small'
+  | 'openai/text-embedding-ada-002'
+  | 'voyage/voyage-3-large'
+  | 'voyage/voyage-3.5'
+  | 'voyage/voyage-3.5-lite'
+  | 'voyage/voyage-code-2'
+  | 'voyage/voyage-code-3'
+  | 'voyage/voyage-finance-2'
+  | 'voyage/voyage-law-2'
+  | (string & {});

package/src/gateway-embedding-model.test.ts

@@ -0,0 +1,213 @@
+import { describe, it, expect } from 'vitest';
+import { createTestServer } from '@ai-sdk/test-server/with-vitest';
+import { GatewayEmbeddingModel } from './gateway-embedding-model';
+import type { GatewayConfig } from './gateway-config';
+import {
+  GatewayInvalidRequestError,
+  GatewayInternalServerError,
+} from './errors';
+
+const dummyEmbeddings = [
+  [0.1, 0.2, 0.3],
+  [0.4, 0.5, 0.6],
+];
+const testValues = ['sunny day at the beach', 'rainy afternoon in the city'];
+
+const server = createTestServer({
+  'https://api.test.com/embedding-model': {},
+});
+
+const createTestModel = (
+  config: Partial<
+    GatewayConfig & { o11yHeaders?: Record<string, string> }
+  > = {},
+) =>
+  new GatewayEmbeddingModel('openai/text-embedding-3-small', {
+    provider: 'gateway',
+    baseURL: 'https://api.test.com',
+    headers: () => ({
+      Authorization: 'Bearer test-token',
+      'ai-gateway-auth-method': 'api-key',
+    }),
+    fetch: globalThis.fetch,
+    o11yHeaders: config.o11yHeaders || {},
+    ...config,
+  });
+
+describe('GatewayEmbeddingModel', () => {
+  function prepareJsonResponse({
+    embeddings = dummyEmbeddings,
+    usage = { tokens: 8 },
+    headers,
+  }: {
+    embeddings?: number[][];
+    usage?: { tokens: number };
+    headers?: Record<string, string>;
+  } = {}) {
+    server.urls['https://api.test.com/embedding-model'].response = {
+      type: 'json-value',
+      headers,
+      body: {
+        embeddings,
+        usage,
+      },
+    };
+  }
+
+  describe('doEmbed', () => {
+    it('should pass headers correctly', async () => {
+      prepareJsonResponse();
+
+      await createTestModel().doEmbed({
+        values: testValues,
+        headers: { 'Custom-Header': 'test-value' },
+      });
+
+      const headers = server.calls[0].requestHeaders;
+      expect(headers).toMatchObject({
+        authorization: 'Bearer test-token',
+        'custom-header': 'test-value',
+        'ai-embedding-model-specification-version': '3',
+        'ai-model-id': 'openai/text-embedding-3-small',
+      });
+    });
+
+    it('should include o11y headers', async () => {
+      prepareJsonResponse();
+
+      const o11yHeaders = {
+        'ai-o11y-deployment-id': 'deployment-1',
+        'ai-o11y-environment': 'production',
+        'ai-o11y-region': 'iad1',
+      } as const;
+
+      await createTestModel({ o11yHeaders }).doEmbed({ values: testValues });
+
+      const headers = server.calls[0].requestHeaders;
+      expect(headers).toMatchObject(o11yHeaders);
+    });
+
+    it('should extract embeddings and usage', async () => {
+      prepareJsonResponse({
+        embeddings: dummyEmbeddings,
+        usage: { tokens: 42 },
+      });
+
+      const { embeddings, usage } = await createTestModel().doEmbed({
+        values: testValues,
+      });
+
+      expect(embeddings).toStrictEqual(dummyEmbeddings);
+      expect(usage).toStrictEqual({ tokens: 42 });
+    });
+
+    it('should send value as array', async () => {
+      prepareJsonResponse();
+
+      await createTestModel().doEmbed({ values: testValues });
+      expect(await server.calls[0].requestBodyJson).toStrictEqual({
+        values: testValues,
+      });
+    });
+
+    it('should pass providerOptions into request body', async () => {
+      prepareJsonResponse();
+
+      await createTestModel().doEmbed({
+        values: testValues,
+        providerOptions: { openai: { dimensions: 64 } },
+      });
+
+      expect(await server.calls[0].requestBodyJson).toStrictEqual({
+        values: testValues,
+        providerOptions: { openai: { dimensions: 64 } },
+      });
+    });
+
+    it('should not include providerOptions when not provided', async () => {
+      prepareJsonResponse();
+
+      await createTestModel().doEmbed({ values: testValues });
+
+      const body = await server.calls[0].requestBodyJson;
+      expect(body).toStrictEqual({ values: testValues });
+      expect('providerOptions' in body).toBe(false);
+    });
+
+    it('should convert gateway error responses', async () => {
+      server.urls['https://api.test.com/embedding-model'].response = {
+        type: 'error',
+        status: 400,
+        body: JSON.stringify({
+          error: {
+            message: 'Invalid input',
+            type: 'invalid_request_error',
+          },
+        }),
+      };
+
+      await expect(
+        createTestModel().doEmbed({ values: testValues }),
+      ).rejects.toSatisfy(
+        err =>
+          GatewayInvalidRequestError.isInstance(err) && err.statusCode === 400,
+      );
+
+      server.urls['https://api.test.com/embedding-model'].response = {
+        type: 'error',
+        status: 500,
+        body: JSON.stringify({
+          error: {
+            message: 'Server blew up',
+            type: 'internal_server_error',
+          },
+        }),
+      };
+
+      await expect(
+        createTestModel().doEmbed({ values: testValues }),
+      ).rejects.toSatisfy(
+        err =>
+          GatewayInternalServerError.isInstance(err) && err.statusCode === 500,
+      );
+    });
+
+    it('should include providerMetadata in response body', async () => {
+      server.urls['https://api.test.com/embedding-model'].response = {
+        type: 'json-value',
+        body: {
+          embeddings: dummyEmbeddings,
+          usage: { tokens: 5 },
+          providerMetadata: { gateway: { routing: { test: true } } },
+        },
+      };
+
+      const { response } = await createTestModel().doEmbed({
+        values: testValues,
+      });
+
+      expect(response?.body).toMatchObject({
+        providerMetadata: { gateway: { routing: { test: true } } },
+      });
+    });
+
+    it('should extract providerMetadata to top level', async () => {
+      server.urls['https://api.test.com/embedding-model'].response = {
+        type: 'json-value',
+        body: {
+          embeddings: dummyEmbeddings,
+          usage: { tokens: 5 },
+          providerMetadata: { gateway: { routing: { test: true } } },
+        },
+      };
+
+      const result = await createTestModel().doEmbed({
+        values: testValues,
+      });
+
+      expect(result.providerMetadata).toStrictEqual({
+        gateway: { routing: { test: true } },
+      });
+    });
+  });
+});

package/src/gateway-embedding-model.ts

@@ -0,0 +1,109 @@
+import type {
+  EmbeddingModelV3,
+  SharedV3ProviderMetadata,
+} from '@ai-sdk/provider';
+import {
+  combineHeaders,
+  createJsonErrorResponseHandler,
+  createJsonResponseHandler,
+  lazySchema,
+  postJsonToApi,
+  resolve,
+  zodSchema,
+  type Resolvable,
+} from '@ai-sdk/provider-utils';
+import { z } from 'zod/v4';
+import { asGatewayError } from './errors';
+import { parseAuthMethod } from './errors/parse-auth-method';
+import type { GatewayConfig } from './gateway-config';
+
+export class GatewayEmbeddingModel implements EmbeddingModelV3 {
+  readonly specificationVersion = 'v3';
+  readonly maxEmbeddingsPerCall = 2048;
+  readonly supportsParallelCalls = true;
+
+  constructor(
+    readonly modelId: string,
+    private readonly config: GatewayConfig & {
+      provider: string;
+      o11yHeaders: Resolvable<Record<string, string>>;
+    },
+  ) {}
+
+  get provider(): string {
+    return this.config.provider;
+  }
+
+  async doEmbed({
+    values,
+    headers,
+    abortSignal,
+    providerOptions,
+  }: Parameters<EmbeddingModelV3['doEmbed']>[0]): Promise<
+    Awaited<ReturnType<EmbeddingModelV3['doEmbed']>>
+  > {
+    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: {
+          values,
+          ...(providerOptions ? { providerOptions } : {}),
+        },
+        successfulResponseHandler: createJsonResponseHandler(
+          gatewayEmbeddingResponseSchema,
+        ),
+        failedResponseHandler: createJsonErrorResponseHandler({
+          errorSchema: z.any(),
+          errorToMessage: data => data,
+        }),
+        ...(abortSignal && { abortSignal }),
+        fetch: this.config.fetch,
+      });
+
+      return {
+        embeddings: responseBody.embeddings,
+        usage: responseBody.usage ?? undefined,
+        providerMetadata:
+          responseBody.providerMetadata as unknown as SharedV3ProviderMetadata,
+        response: { headers: responseHeaders, body: rawValue },
+        warnings: [],
+      };
+    } catch (error) {
+      throw await asGatewayError(error, await parseAuthMethod(resolvedHeaders));
+    }
+  }
+
+  private getUrl() {
+    return `${this.config.baseURL}/embedding-model`;
+  }
+
+  private getModelConfigHeaders() {
+    return {
+      'ai-embedding-model-specification-version': '3',
+      'ai-model-id': this.modelId,
+    };
+  }
+}
+
+const gatewayEmbeddingResponseSchema = lazySchema(() =>
+  zodSchema(
+    z.object({
+      embeddings: z.array(z.array(z.number())),
+      usage: z.object({ tokens: z.number() }).nullish(),
+      providerMetadata: z
+        .record(z.string(), z.record(z.string(), z.unknown()))
+        .optional(),
+    }),
+  ),
+);