feat(amazon-bedrock): binary embeddings

.changeset/thick-wolves-bake.md

@@ -0,0 +1,5 @@
+---
+'@ai-sdk/amazon-bedrock': patch
+---
+
+Allow generating binary embeddings using Amazon Titan Text Embeddings V2.

content/providers/01-ai-sdk-providers/08-amazon-bedrock.mdx

@@ -453,6 +453,10 @@ The following optional settings are available for Bedrock Titan embedding models
 
   Flag indicating whether or not to normalize the output embeddings. Defaults to true.
 
+- **embeddingType** _'float' | 'binary'_
+
+  The type of embedding to return. Defaults to 'float'. Binary embeddings are only supported in amazon.titan-embed-text-v2:0.
+
 ### Model Capabilities
 
 | Model                          | Default Dimensions | Custom Dimensions   |

examples/ai-core/src/embed/amazon-bedrock.ts

@@ -10,6 +10,16 @@ async function main() {
 
   console.log(embedding);
   console.log(usage);
+
+  const { embedding: binaryEmbedding, usage: binaryUsage } = await embed({
+    model: bedrock.embedding('amazon.titan-embed-text-v2:0', {
+      embeddingType: 'binary',
+    }),
+    value: 'sunny day at the beach',
+  });
+
+  console.log(binaryEmbedding);
+  console.log(binaryUsage);
 }
 
 main().catch(console.error);

packages/amazon-bedrock/src/bedrock-embedding-model.test.ts

@@ -7,6 +7,11 @@ const mockEmbeddings = [
   [-0.08, 0.06, -0.03, 0.02, 0.03],
 ];
 
+const mockBinaryEmbeddings = [
+  [0, 1, 0, 1, 0],
+  [1, 0, 1, 0, 1],
+];
+
 const fakeFetchWithAuth = injectFetchHeaders({ 'x-amz-auth': 'test-auth' });
 
 const testValues = ['sunny day at the beach', 'rainy day in the city'];
@@ -158,4 +163,46 @@ describe('doEmbed', () => {
     expect(requestHeaders['signed-header']).toBe('signed-value');
     expect(requestHeaders['authorization']).toBe('AWS4-HMAC-SHA256...');
   });
+
+  it('should work with binary embeddings', async () => {
+    const modelWithBinaryEmbeddings = new BedrockEmbeddingModel(
+      'amazon.titan-embed-text-v2:0',
+      { embeddingType: 'binary' },
+      {
+        baseUrl: () => 'https://bedrock-runtime.us-east-1.amazonaws.com',
+        headers: mockConfigHeaders,
+        fetch: fakeFetchWithAuth,
+      },
+    );
+
+    server.urls[embedUrl].response = {
+      type: 'binary',
+      headers: {
+        'content-type': 'application/json',
+      },
+      body: Buffer.from(
+        JSON.stringify({
+          embeddingsByType: {
+            binary: mockBinaryEmbeddings[0],
+          },
+          inputTextTokenCount: 8,
+        }),
+      ),
+    };
+
+    const { embeddings } = await modelWithBinaryEmbeddings.doEmbed({
+      values: [testValues[0]],
+    });
+
+    expect(embeddings.length).toBe(1);
+    expect(embeddings[0]).toStrictEqual(mockBinaryEmbeddings[0]);
+
+    const body = await server.calls[0].requestBody;
+    expect(body).toEqual({
+      inputText: testValues[0],
+      dimensions: undefined,
+      normalize: undefined,
+      embeddingTypes: ['binary'],
+    });
+  });
 });

packages/amazon-bedrock/src/bedrock-embedding-model.ts

@@ -48,11 +48,16 @@ export class BedrockEmbeddingModel implements EmbeddingModelV1<string> {
     EmbeddingModelV1<string>['doEmbed']
   >[0]): Promise<DoEmbedResponse> {
     const embedSingleText = async (inputText: string) => {
+      const isBinaryEmbedding = this.settings.embeddingType === 'binary';
+
       // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InvokeModel.html
       const args = {
         inputText,
         dimensions: this.settings.dimensions,
         normalize: this.settings.normalize,
+        ...(isBinaryEmbedding && {
+          embeddingTypes: ['binary'],
+        }),
       };
       const url = this.getUrl(this.modelId);
       const { value: response } = await postJsonToApi({
@@ -66,14 +71,19 @@ export class BedrockEmbeddingModel implements EmbeddingModelV1<string> {
           errorToMessage: error => `${error.type}: ${error.message}`,
         }),
         successfulResponseHandler: createJsonResponseHandler(
-          BedrockEmbeddingResponseSchema,
+          z.union([
+            BedrockBinaryEmbeddingResponseSchema,
+            BedrockEmbeddingResponseSchema,
+          ]),
         ),
         fetch: this.config.fetch,
         abortSignal,
       });
 
       return {
-        embedding: response.embedding,
+        embedding: isBinaryEmbedding
+          ? (response as BedrockBinaryEmbeddingResponse).embeddingsByType.binary
+          : (response as BedrockEmbeddingResponse).embedding,
         inputTextTokenCount: response.inputTextTokenCount,
       };
     };
@@ -97,3 +107,16 @@ const BedrockEmbeddingResponseSchema = z.object({
   embedding: z.array(z.number()),
   inputTextTokenCount: z.number(),
 });
+
+type BedrockEmbeddingResponse = z.infer<typeof BedrockEmbeddingResponseSchema>;
+
+const BedrockBinaryEmbeddingResponseSchema = z.object({
+  embeddingsByType: z.object({
+    binary: z.array(z.number()),
+  }),
+  inputTextTokenCount: z.number(),
+});
+
+type BedrockBinaryEmbeddingResponse = z.infer<
+  typeof BedrockBinaryEmbeddingResponseSchema
+>;

packages/amazon-bedrock/src/bedrock-embedding-settings.ts

@@ -17,4 +17,10 @@ Flag indicating whether or not to normalize the output embeddings. Defaults to t
 Only supported in amazon.titan-embed-text-v2:0.
    */
   normalize?: boolean;
+
+  /**
+The type of embedding to return. Defaults to float.
+Binary embeddings are only supported in amazon.titan-embed-text-v2:0.
+  */
+  embeddingType?: 'float' | 'binary';
 }