@ai-sdk/amazon-bedrock 4.0.27 → 4.0.29

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

@@ -1,148 +0,0 @@
-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-event-stream-response-handler.test.ts

@@ -1,233 +0,0 @@
-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' } },
-    });
-  });
-});