@ai-sdk/amazon-bedrock 4.0.24 → 4.0.26

package/src/bedrock-chat-options.ts

@@ -0,0 +1,114 @@
+import { z } from 'zod/v4';
+
+// https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html
+export type BedrockChatModelId =
+  | 'amazon.titan-tg1-large'
+  | 'amazon.titan-text-express-v1'
+  | 'anthropic.claude-v2'
+  | 'anthropic.claude-v2:1'
+  | 'anthropic.claude-instant-v1'
+  | 'anthropic.claude-haiku-4-5-20251001-v1:0'
+  | 'anthropic.claude-sonnet-4-20250514-v1:0'
+  | 'anthropic.claude-sonnet-4-5-20250929-v1:0'
+  | 'anthropic.claude-opus-4-20250514-v1:0'
+  | 'anthropic.claude-opus-4-1-20250805-v1:0'
+  | 'anthropic.claude-3-7-sonnet-20250219-v1:0'
+  | 'anthropic.claude-3-5-sonnet-20240620-v1:0'
+  | 'anthropic.claude-3-5-sonnet-20241022-v2:0'
+  | 'anthropic.claude-3-5-haiku-20241022-v1:0'
+  | 'anthropic.claude-3-sonnet-20240229-v1:0'
+  | 'anthropic.claude-3-haiku-20240307-v1:0'
+  | 'anthropic.claude-3-opus-20240229-v1:0'
+  | 'cohere.command-text-v14'
+  | 'cohere.command-light-text-v14'
+  | 'cohere.command-r-v1:0'
+  | 'cohere.command-r-plus-v1:0'
+  | 'meta.llama3-70b-instruct-v1:0'
+  | 'meta.llama3-8b-instruct-v1:0'
+  | 'meta.llama3-1-405b-instruct-v1:0'
+  | 'meta.llama3-1-70b-instruct-v1:0'
+  | 'meta.llama3-1-8b-instruct-v1:0'
+  | 'meta.llama3-2-11b-instruct-v1:0'
+  | 'meta.llama3-2-1b-instruct-v1:0'
+  | 'meta.llama3-2-3b-instruct-v1:0'
+  | 'meta.llama3-2-90b-instruct-v1:0'
+  | 'mistral.mistral-7b-instruct-v0:2'
+  | 'mistral.mixtral-8x7b-instruct-v0:1'
+  | 'mistral.mistral-large-2402-v1:0'
+  | 'mistral.mistral-small-2402-v1:0'
+  | 'openai.gpt-oss-120b-1:0'
+  | 'openai.gpt-oss-20b-1:0'
+  | 'amazon.titan-text-express-v1'
+  | 'amazon.titan-text-lite-v1'
+  | 'us.amazon.nova-premier-v1:0'
+  | 'us.amazon.nova-pro-v1:0'
+  | 'us.amazon.nova-micro-v1:0'
+  | 'us.amazon.nova-lite-v1:0'
+  | 'us.anthropic.claude-3-sonnet-20240229-v1:0'
+  | 'us.anthropic.claude-3-opus-20240229-v1:0'
+  | 'us.anthropic.claude-3-haiku-20240307-v1:0'
+  | 'us.anthropic.claude-3-5-sonnet-20240620-v1:0'
+  | 'us.anthropic.claude-3-5-haiku-20241022-v1:0'
+  | 'us.anthropic.claude-3-5-sonnet-20241022-v2:0'
+  | 'us.anthropic.claude-3-7-sonnet-20250219-v1:0'
+  | 'us.anthropic.claude-sonnet-4-20250514-v1:0'
+  | 'us.anthropic.claude-sonnet-4-5-20250929-v1:0'
+  | 'us.anthropic.claude-opus-4-20250514-v1:0'
+  | 'us.anthropic.claude-opus-4-1-20250805-v1:0'
+  | 'us.meta.llama3-2-11b-instruct-v1:0'
+  | 'us.meta.llama3-2-3b-instruct-v1:0'
+  | 'us.meta.llama3-2-90b-instruct-v1:0'
+  | 'us.meta.llama3-2-1b-instruct-v1:0'
+  | 'us.meta.llama3-1-8b-instruct-v1:0'
+  | 'us.meta.llama3-1-70b-instruct-v1:0'
+  | 'us.meta.llama3-3-70b-instruct-v1:0'
+  | 'us.deepseek.r1-v1:0'
+  | 'us.mistral.pixtral-large-2502-v1:0'
+  | 'us.meta.llama4-scout-17b-instruct-v1:0'
+  | 'us.meta.llama4-maverick-17b-instruct-v1:0'
+  | (string & {});
+
+/**
+ * Bedrock file part provider options for document-specific features.
+ * These options apply to individual file parts (documents).
+ */
+export const bedrockFilePartProviderOptions = z.object({
+  /**
+   * Citation configuration for this document.
+   * When enabled, this document will generate citations in the response.
+   */
+  citations: z
+    .object({
+      /**
+       * Enable citations for this document
+       */
+      enabled: z.boolean(),
+    })
+    .optional(),
+});
+
+export type BedrockFilePartProviderOptions = z.infer<
+  typeof bedrockFilePartProviderOptions
+>;
+
+export const bedrockProviderOptions = z.object({
+  /**
+   * Additional inference parameters that the model supports,
+   * beyond the base set of inference parameters that Converse
+   * supports in the inferenceConfig field
+   */
+  additionalModelRequestFields: z.record(z.string(), z.any()).optional(),
+  reasoningConfig: z
+    .object({
+      type: z.union([z.literal('enabled'), z.literal('disabled')]).optional(),
+      budgetTokens: z.number().optional(),
+      maxReasoningEffort: z.enum(['low', 'medium', 'high']).optional(),
+    })
+    .optional(),
+  /**
+   * Anthropic beta features to enable
+   */
+  anthropicBeta: z.array(z.string()).optional(),
+});
+
+export type BedrockProviderOptions = z.infer<typeof bedrockProviderOptions>;

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

@@ -0,0 +1,148 @@
+import { createTestServer } from '@ai-sdk/test-server/with-vitest';
+import { BedrockEmbeddingModel } from './bedrock-embedding-model';
+import { injectFetchHeaders } from './inject-fetch-headers';
+import { beforeEach, describe, expect, it } from 'vitest';
+
+const mockEmbeddings = [
+  [-0.09, 0.05, -0.02, 0.01, 0.04],
+  [-0.08, 0.06, -0.03, 0.02, 0.03],
+];
+
+const fakeFetchWithAuth = injectFetchHeaders({ 'x-amz-auth': 'test-auth' });
+
+const testValues = ['sunny day at the beach', 'rainy day in the city'];
+
+const embedUrl = `https://bedrock-runtime.us-east-1.amazonaws.com/model/${encodeURIComponent(
+  'amazon.titan-embed-text-v2:0',
+)}/invoke`;
+
+describe('doEmbed', () => {
+  const mockConfigHeaders = {
+    'config-header': 'config-value',
+    'shared-header': 'config-shared',
+  };
+
+  const server = createTestServer({
+    [embedUrl]: {
+      response: {
+        type: 'binary',
+        headers: {
+          'content-type': 'application/json',
+        },
+        body: Buffer.from(
+          JSON.stringify({
+            embedding: mockEmbeddings[0],
+            inputTextTokenCount: 8,
+          }),
+        ),
+      },
+    },
+  });
+
+  const model = new BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', {
+    baseUrl: () => 'https://bedrock-runtime.us-east-1.amazonaws.com',
+    headers: mockConfigHeaders,
+    fetch: fakeFetchWithAuth,
+  });
+
+  let callCount = 0;
+
+  beforeEach(() => {
+    callCount = 0;
+    server.urls[embedUrl].response = {
+      type: 'binary',
+      headers: {
+        'content-type': 'application/json',
+      },
+      body: Buffer.from(
+        JSON.stringify({
+          embedding: mockEmbeddings[0],
+          inputTextTokenCount: 8,
+        }),
+      ),
+    };
+  });
+
+  it('should handle single input value and return embeddings', async () => {
+    const { embeddings } = await model.doEmbed({
+      values: [testValues[0]],
+    });
+
+    expect(embeddings.length).toBe(1);
+    expect(embeddings[0]).toStrictEqual(mockEmbeddings[0]);
+
+    const body = await server.calls[0].requestBodyJson;
+    expect(body).toEqual({
+      inputText: testValues[0],
+      dimensions: undefined,
+      normalize: undefined,
+    });
+  });
+
+  it('should handle single input value and extract usage', async () => {
+    const { usage } = await model.doEmbed({
+      values: [testValues[0]],
+    });
+
+    expect(usage?.tokens).toStrictEqual(8);
+  });
+
+  it('should properly combine headers from all sources', async () => {
+    const optionsHeaders = {
+      'options-header': 'options-value',
+      'shared-header': 'options-shared',
+    };
+
+    const modelWithHeaders = new BedrockEmbeddingModel(
+      'amazon.titan-embed-text-v2:0',
+      {
+        baseUrl: () => 'https://bedrock-runtime.us-east-1.amazonaws.com',
+        headers: {
+          'model-header': 'model-value',
+          'shared-header': 'model-shared',
+        },
+        fetch: injectFetchHeaders({
+          'signed-header': 'signed-value',
+          authorization: 'AWS4-HMAC-SHA256...',
+        }),
+      },
+    );
+
+    await modelWithHeaders.doEmbed({
+      values: [testValues[0]],
+      headers: optionsHeaders,
+    });
+
+    const requestHeaders = server.calls[0].requestHeaders;
+    expect(requestHeaders['options-header']).toBe('options-value');
+    expect(requestHeaders['model-header']).toBe('model-value');
+    expect(requestHeaders['signed-header']).toBe('signed-value');
+    expect(requestHeaders['authorization']).toBe('AWS4-HMAC-SHA256...');
+    expect(requestHeaders['shared-header']).toBe('options-shared');
+  });
+
+  it('should work with partial headers', async () => {
+    const modelWithPartialHeaders = new BedrockEmbeddingModel(
+      'amazon.titan-embed-text-v2:0',
+      {
+        baseUrl: () => 'https://bedrock-runtime.us-east-1.amazonaws.com',
+        headers: {
+          'model-header': 'model-value',
+        },
+        fetch: injectFetchHeaders({
+          'signed-header': 'signed-value',
+          authorization: 'AWS4-HMAC-SHA256...',
+        }),
+      },
+    );
+
+    await modelWithPartialHeaders.doEmbed({
+      values: [testValues[0]],
+    });
+
+    const requestHeaders = server.calls[0].requestHeaders;
+    expect(requestHeaders['model-header']).toBe('model-value');
+    expect(requestHeaders['signed-header']).toBe('signed-value');
+    expect(requestHeaders['authorization']).toBe('AWS4-HMAC-SHA256...');
+  });
+});

package/src/bedrock-embedding-model.ts

@@ -0,0 +1,104 @@
+import {
+  EmbeddingModelV3,
+  TooManyEmbeddingValuesForCallError,
+} from '@ai-sdk/provider';
+import {
+  FetchFunction,
+  Resolvable,
+  combineHeaders,
+  createJsonErrorResponseHandler,
+  createJsonResponseHandler,
+  parseProviderOptions,
+  postJsonToApi,
+  resolve,
+} from '@ai-sdk/provider-utils';
+import {
+  BedrockEmbeddingModelId,
+  bedrockEmbeddingProviderOptions,
+} from './bedrock-embedding-options';
+import { BedrockErrorSchema } from './bedrock-error';
+import { z } from 'zod/v4';
+
+type BedrockEmbeddingConfig = {
+  baseUrl: () => string;
+  headers: Resolvable<Record<string, string | undefined>>;
+  fetch?: FetchFunction;
+};
+
+type DoEmbedResponse = Awaited<ReturnType<EmbeddingModelV3['doEmbed']>>;
+
+export class BedrockEmbeddingModel implements EmbeddingModelV3 {
+  readonly specificationVersion = 'v3';
+  readonly provider = 'amazon-bedrock';
+  readonly maxEmbeddingsPerCall = 1;
+  readonly supportsParallelCalls = true;
+
+  constructor(
+    readonly modelId: BedrockEmbeddingModelId,
+    private readonly config: BedrockEmbeddingConfig,
+  ) {}
+
+  private getUrl(modelId: string): string {
+    const encodedModelId = encodeURIComponent(modelId);
+    return `${this.config.baseUrl()}/model/${encodedModelId}/invoke`;
+  }
+
+  async doEmbed({
+    values,
+    headers,
+    abortSignal,
+    providerOptions,
+  }: Parameters<EmbeddingModelV3['doEmbed']>[0]): Promise<DoEmbedResponse> {
+    if (values.length > this.maxEmbeddingsPerCall) {
+      throw new TooManyEmbeddingValuesForCallError({
+        provider: this.provider,
+        modelId: this.modelId,
+        maxEmbeddingsPerCall: this.maxEmbeddingsPerCall,
+        values,
+      });
+    }
+
+    // Parse provider options
+    const bedrockOptions =
+      (await parseProviderOptions({
+        provider: 'bedrock',
+        providerOptions,
+        schema: bedrockEmbeddingProviderOptions,
+      })) ?? {};
+
+    // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InvokeModel.html
+    const args = {
+      inputText: values[0],
+      dimensions: bedrockOptions.dimensions,
+      normalize: bedrockOptions.normalize,
+    };
+    const url = this.getUrl(this.modelId);
+    const { value: response } = await postJsonToApi({
+      url,
+      headers: await resolve(
+        combineHeaders(await resolve(this.config.headers), headers),
+      ),
+      body: args,
+      failedResponseHandler: createJsonErrorResponseHandler({
+        errorSchema: BedrockErrorSchema,
+        errorToMessage: error => `${error.type}: ${error.message}`,
+      }),
+      successfulResponseHandler: createJsonResponseHandler(
+        BedrockEmbeddingResponseSchema,
+      ),
+      fetch: this.config.fetch,
+      abortSignal,
+    });
+
+    return {
+      warnings: [],
+      embeddings: [response.embedding],
+      usage: { tokens: response.inputTextTokenCount },
+    };
+  }
+}
+
+const BedrockEmbeddingResponseSchema = z.object({
+  embedding: z.array(z.number()),
+  inputTextTokenCount: z.number(),
+});

package/src/bedrock-embedding-options.ts

@@ -0,0 +1,24 @@
+import { z } from 'zod/v4';
+
+export type BedrockEmbeddingModelId =
+  | 'amazon.titan-embed-text-v1'
+  | 'amazon.titan-embed-text-v2:0'
+  | 'cohere.embed-english-v3'
+  | 'cohere.embed-multilingual-v3'
+  | (string & {});
+
+export const bedrockEmbeddingProviderOptions = z.object({
+  /**
+The number of dimensions the resulting output embeddings should have (defaults to 1024).
+Only supported in amazon.titan-embed-text-v2:0.
+   */
+  dimensions: z
+    .union([z.literal(1024), z.literal(512), z.literal(256)])
+    .optional(),
+
+  /**
+Flag indicating whether or not to normalize the output embeddings. Defaults to true
+Only supported in amazon.titan-embed-text-v2:0.
+ */
+  normalize: z.boolean().optional(),
+});

package/src/bedrock-error.ts

@@ -0,0 +1,6 @@
+import { z } from 'zod/v4';
+
+export const BedrockErrorSchema = z.object({
+  message: z.string(),
+  type: z.string().nullish(),
+});

package/src/bedrock-event-stream-decoder.ts

@@ -0,0 +1,59 @@
+import { EventStreamCodec } from '@smithy/eventstream-codec';
+import { toUtf8, fromUtf8 } from '@smithy/util-utf8';
+
+export interface DecodedEvent {
+  messageType: string;
+  eventType: string;
+  data: string;
+}
+
+export function createBedrockEventStreamDecoder<T>(
+  body: ReadableStream<Uint8Array>,
+  processEvent: (
+    event: DecodedEvent,
+    controller: TransformStreamDefaultController<T>,
+  ) => void | Promise<void>,
+): ReadableStream<T> {
+  const codec = new EventStreamCodec(toUtf8, fromUtf8);
+  let buffer = new Uint8Array(0);
+  const textDecoder = new TextDecoder();
+
+  return body.pipeThrough(
+    new TransformStream<Uint8Array, T>({
+      async transform(chunk, controller) {
+        const newBuffer = new Uint8Array(buffer.length + chunk.length);
+        newBuffer.set(buffer);
+        newBuffer.set(chunk, buffer.length);
+        buffer = newBuffer;
+
+        while (buffer.length >= 4) {
+          const totalLength = new DataView(
+            buffer.buffer,
+            buffer.byteOffset,
+            buffer.byteLength,
+          ).getUint32(0, false);
+
+          if (buffer.length < totalLength) {
+            break;
+          }
+
+          try {
+            const subView = buffer.subarray(0, totalLength);
+            const decoded = codec.decode(subView);
+
+            buffer = buffer.slice(totalLength);
+
+            const messageType = decoded.headers[':message-type']
+              ?.value as string;
+            const eventType = decoded.headers[':event-type']?.value as string;
+            const data = textDecoder.decode(decoded.body);
+
+            await processEvent({ messageType, eventType, data }, controller);
+          } catch {
+            break;
+          }
+        }
+      },
+    }),
+  );
+}

package/src/bedrock-event-stream-response-handler.test.ts

@@ -0,0 +1,233 @@
+import { EmptyResponseBodyError } from '@ai-sdk/provider';
+import { createBedrockEventStreamResponseHandler } from './bedrock-event-stream-response-handler';
+import { EventStreamCodec } from '@smithy/eventstream-codec';
+import { z } from 'zod/v4';
+import { describe, it, expect, vi, MockInstance } from 'vitest';
+
+// Helper that constructs a properly framed message.
+// The first 4 bytes will contain the frame total length (big-endian).
+const createFrame = (payload: Uint8Array): Uint8Array => {
+  const totalLength = 4 + payload.length;
+  const frame = new Uint8Array(totalLength);
+  new DataView(frame.buffer).setUint32(0, totalLength, false);
+  frame.set(payload, 4);
+  return frame;
+};
+
+// Mock EventStreamCodec
+vi.mock('@smithy/eventstream-codec', () => ({
+  EventStreamCodec: vi.fn(),
+}));
+
+describe('createEventSourceResponseHandler', () => {
+  // Define a sample schema for testing
+  const testSchema = z.object({
+    chunk: z.object({
+      content: z.string(),
+    }),
+  });
+
+  it('throws EmptyResponseBodyError when response body is null', async () => {
+    const response = new Response(null);
+    const handler = createBedrockEventStreamResponseHandler(testSchema);
+
+    await expect(
+      handler({
+        response,
+        url: 'test-url',
+        requestBodyValues: {},
+      }),
+    ).rejects.toThrow(EmptyResponseBodyError);
+  });
+
+  it('successfully processes valid event stream data', async () => {
+    // Prepare the message we wish to simulate.
+    // Our decoded message will contain headers and a body that is valid JSON.
+    const message = {
+      headers: {
+        ':message-type': { value: 'event' },
+        ':event-type': { value: 'chunk' },
+      },
+      body: new TextEncoder().encode(
+        JSON.stringify({ content: 'test message' }),
+      ),
+    };
+
+    // Create a frame that properly encapsulates the message.
+    const dummyPayload = new Uint8Array([1, 2, 3, 4]); // arbitrary payload that makes the length check pass
+    const frame = createFrame(dummyPayload);
+
+    const mockDecode = vi.fn().mockReturnValue(message);
+    (EventStreamCodec as unknown as MockInstance).mockImplementation(() => ({
+      decode: mockDecode,
+    }));
+
+    // Create a stream that enqueues the complete frame.
+    const stream = new ReadableStream({
+      start(controller) {
+        controller.enqueue(frame);
+        controller.close();
+      },
+    });
+
+    const response = new Response(stream);
+    const handler = createBedrockEventStreamResponseHandler(testSchema);
+    const result = await handler({
+      response,
+      url: 'test-url',
+      requestBodyValues: {},
+    });
+
+    const reader = result.value.getReader();
+    const { done, value } = await reader.read();
+
+    expect(done).toBe(false);
+    expect(value).toEqual({
+      success: true,
+      value: { chunk: { content: 'test message' } },
+      rawValue: { chunk: { content: 'test message' } },
+    });
+  });
+
+  it('handles invalid JSON data', async () => {
+    // Our mock decode returns a body that is not valid JSON.
+    const message = {
+      headers: {
+        ':message-type': { value: 'event' },
+        ':event-type': { value: 'chunk' },
+      },
+      body: new TextEncoder().encode('invalid json'),
+    };
+
+    const dummyPayload = new Uint8Array([5, 6, 7, 8]);
+    const frame = createFrame(dummyPayload);
+
+    const mockDecode = vi.fn().mockReturnValue(message);
+    (EventStreamCodec as unknown as MockInstance).mockImplementation(() => ({
+      decode: mockDecode,
+    }));
+
+    const stream = new ReadableStream({
+      start(controller) {
+        controller.enqueue(frame);
+        controller.close();
+      },
+    });
+
+    const response = new Response(stream);
+    const handler = createBedrockEventStreamResponseHandler(testSchema);
+    const result = await handler({
+      response,
+      url: 'test-url',
+      requestBodyValues: {},
+    });
+
+    const reader = result.value.getReader();
+    const { done, value } = await reader.read();
+
+    expect(done).toBe(false);
+    // When JSON is invalid, safeParseJSON returns a result with success: false.
+    expect(value?.success).toBe(false);
+    expect((value as { success: false; error: Error }).error).toBeDefined();
+  });
+
+  it('handles schema validation failures', async () => {
+    // The decoded message returns valid JSON but that does not meet our schema.
+    const message = {
+      headers: {
+        ':message-type': { value: 'event' },
+        ':event-type': { value: 'chunk' },
+      },
+      body: new TextEncoder().encode(JSON.stringify({ invalid: 'data' })),
+    };
+
+    const dummyPayload = new Uint8Array([9, 10, 11, 12]);
+    const frame = createFrame(dummyPayload);
+
+    const mockDecode = vi.fn().mockReturnValue(message);
+    (EventStreamCodec as unknown as MockInstance).mockImplementation(() => ({
+      decode: mockDecode,
+    }));
+
+    const stream = new ReadableStream({
+      start(controller) {
+        controller.enqueue(frame);
+        controller.close();
+      },
+    });
+
+    const response = new Response(stream);
+    const handler = createBedrockEventStreamResponseHandler(testSchema);
+    const result = await handler({
+      response,
+      url: 'test-url',
+      requestBodyValues: {},
+    });
+
+    const reader = result.value.getReader();
+    const { done, value } = await reader.read();
+
+    expect(done).toBe(false);
+    // The schema does not match so safeParseJSON with the schema should yield success: false.
+    expect(value?.success).toBe(false);
+    expect((value as { success: false; error: Error }).error).toBeDefined();
+  });
+
+  it('handles partial messages correctly', async () => {
+    // In this test, we simulate a partial message followed by a complete one.
+    // The first invocation of decode will throw an error (simulated incomplete message),
+    // and the subsequent invocation returns a valid event.
+    const message = {
+      headers: {
+        ':message-type': { value: 'event' },
+        ':event-type': { value: 'chunk' },
+      },
+      body: new TextEncoder().encode(
+        JSON.stringify({ content: 'complete message' }),
+      ),
+    };
+
+    const dummyPayload1 = new Uint8Array([13, 14]); // too short, part of a frame
+    const frame1 = createFrame(dummyPayload1);
+    const dummyPayload2 = new Uint8Array([15, 16, 17, 18]);
+    const frame2 = createFrame(dummyPayload2);
+
+    const mockDecode = vi
+      .fn()
+      .mockImplementationOnce(() => {
+        throw new Error('Incomplete data');
+      })
+      .mockReturnValue(message);
+    (EventStreamCodec as unknown as MockInstance).mockImplementation(() => ({
+      decode: mockDecode,
+    }));
+
+    const stream = new ReadableStream({
+      start(controller) {
+        // Send first, incomplete frame (decode will throw error).
+        controller.enqueue(frame1);
+        // Then send a proper frame.
+        controller.enqueue(frame2);
+        controller.close();
+      },
+    });
+
+    const response = new Response(stream);
+    const handler = createBedrockEventStreamResponseHandler(testSchema);
+    const result = await handler({
+      response,
+      url: 'test-url',
+      requestBodyValues: {},
+    });
+
+    const reader = result.value.getReader();
+    const { done, value } = await reader.read();
+
+    expect(done).toBe(false);
+    expect(value).toEqual({
+      success: true,
+      value: { chunk: { content: 'complete message' } },
+      rawValue: { chunk: { content: 'complete message' } },
+    });
+  });
+});