@ai-sdk/amazon-bedrock 4.0.29 → 4.0.31

package/docs/08-amazon-bedrock.mdx

@@ -758,6 +758,13 @@ The following optional provider options are available for Bedrock Titan embeddin
 | `cohere.embed-english-v3`      | 1024               | <Cross size={18} /> |
 | `cohere.embed-multilingual-v3` | 1024               | <Cross size={18} /> |
 
+<Note>
+  Cohere embedding models on Bedrock require an <code>input_type</code>. Set it
+  via
+  <code>providerOptions.bedrock.inputType</code> (defaults to{' '}
+  <code>search_query</code>).
+</Note>
+
 ## Reranking Models
 
 You can create models that call the [Bedrock Rerank API](https://docs.aws.amazon.com/bedrock/latest/userguide/rerank-api.html)

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@ai-sdk/amazon-bedrock",
-  "version": "4.0.29",
+  "version": "4.0.31",
   "license": "Apache-2.0",
   "sideEffects": false,
   "main": "./dist/index.js",

package/src/bedrock-embedding-model.ts

@@ -67,11 +67,41 @@ export class BedrockEmbeddingModel implements EmbeddingModelV3 {
       })) ?? {};
 
     // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InvokeModel.html
-    const args = {
-      inputText: values[0],
-      dimensions: bedrockOptions.dimensions,
-      normalize: bedrockOptions.normalize,
-    };
+    //
+    // Note: Different embedding model families expect different request/response
+    // payloads (e.g. Titan vs Cohere vs Nova). We keep the public interface stable and
+    // adapt here based on the modelId.
+    const isNovaModel =
+      this.modelId.startsWith('amazon.nova-') && this.modelId.includes('embed');
+    const isCohereModel = this.modelId.startsWith('cohere.embed-');
+
+    const args = isNovaModel
+      ? {
+          taskType: 'SINGLE_EMBEDDING',
+          singleEmbeddingParams: {
+            embeddingPurpose:
+              bedrockOptions.embeddingPurpose ?? 'GENERIC_INDEX',
+            embeddingDimension: bedrockOptions.embeddingDimension ?? 1024,
+            text: {
+              truncationMode: bedrockOptions.truncate ?? 'END',
+              value: values[0],
+            },
+          },
+        }
+      : isCohereModel
+        ? {
+            // Cohere embedding models on Bedrock require `input_type`.
+            // Without it, the service attempts other schema branches and rejects the request.
+            input_type: bedrockOptions.inputType ?? 'search_query',
+            texts: [values[0]],
+            truncate: bedrockOptions.truncate,
+          }
+        : {
+            inputText: values[0],
+            dimensions: bedrockOptions.dimensions,
+            normalize: bedrockOptions.normalize,
+          };
+
     const url = this.getUrl(this.modelId);
     const { value: response } = await postJsonToApi({
       url,
@@ -90,15 +120,69 @@ export class BedrockEmbeddingModel implements EmbeddingModelV3 {
       abortSignal,
     });
 
+    // Extract embedding based on response format
+    let embedding: number[];
+    if ('embedding' in response) {
+      // Titan response
+      embedding = response.embedding;
+    } else if (Array.isArray(response.embeddings)) {
+      const firstEmbedding = response.embeddings[0];
+      if (
+        typeof firstEmbedding === 'object' &&
+        firstEmbedding !== null &&
+        'embeddingType' in firstEmbedding
+      ) {
+        // Nova response
+        embedding = firstEmbedding.embedding;
+      } else {
+        // Cohere v3 response
+        embedding = firstEmbedding as number[];
+      }
+    } else {
+      // Cohere v4 response
+      embedding = response.embeddings.float[0];
+    }
+
+    // Extract token count based on response format
+    const tokens =
+      'inputTextTokenCount' in response
+        ? response.inputTextTokenCount // Titan response
+        : 'inputTokenCount' in response
+          ? (response.inputTokenCount ?? 0) // Nova response
+          : NaN; // Cohere doesn't return token count
+
     return {
+      embeddings: [embedding],
+      usage: { tokens },
       warnings: [],
-      embeddings: [response.embedding],
-      usage: { tokens: response.inputTextTokenCount },
     };
   }
 }
 
-const BedrockEmbeddingResponseSchema = z.object({
-  embedding: z.array(z.number()),
-  inputTextTokenCount: z.number(),
-});
+const BedrockEmbeddingResponseSchema = z.union([
+  // Titan-style response
+  z.object({
+    embedding: z.array(z.number()),
+    inputTextTokenCount: z.number(),
+  }),
+  // Nova-style response
+  z.object({
+    embeddings: z.array(
+      z.object({
+        embeddingType: z.string(),
+        embedding: z.array(z.number()),
+      }),
+    ),
+    inputTokenCount: z.number().optional(),
+  }),
+  // Cohere v3-style response
+  z.object({
+    embeddings: z.array(z.array(z.number())),
+  }),
+  // Cohere v4-style response
+  z.object({
+    embeddings: z.object({
+      float: z.array(z.array(z.number())),
+    }),
+  }),
+]);

package/src/bedrock-embedding-options.ts

@@ -21,4 +21,46 @@ Flag indicating whether or not to normalize the output embeddings. Defaults to t
 Only supported in amazon.titan-embed-text-v2:0.
  */
   normalize: z.boolean().optional(),
+
+  /**
+The number of dimensions for Nova embedding models (defaults to 1024).
+Supported values: 256, 384, 1024, 3072.
+Only supported in amazon.nova-* embedding models.
+ */
+  embeddingDimension: z
+    .union([z.literal(256), z.literal(384), z.literal(1024), z.literal(3072)])
+    .optional(),
+
+  /**
+The purpose of the embedding. Defaults to 'GENERIC_INDEX'.
+Only supported in amazon.nova-* embedding models.
+ */
+  embeddingPurpose: z
+    .enum([
+      'GENERIC_INDEX',
+      'TEXT_RETRIEVAL',
+      'IMAGE_RETRIEVAL',
+      'VIDEO_RETRIEVAL',
+      'DOCUMENT_RETRIEVAL',
+      'AUDIO_RETRIEVAL',
+      'GENERIC_RETRIEVAL',
+      'CLASSIFICATION',
+      'CLUSTERING',
+    ])
+    .optional(),
+
+  /**
+Input type for Cohere embedding models on Bedrock.
+Common values: `search_document`, `search_query`, `classification`, `clustering`.
+If not set, the provider defaults to `search_query`.
+   */
+  inputType: z
+    .enum(['search_document', 'search_query', 'classification', 'clustering'])
+    .optional(),
+
+  /**
+Truncation behavior when input exceeds the model's context length.
+Supported in Cohere and Nova embedding models. Defaults to 'END' for Nova models.
+   */
+  truncate: z.enum(['NONE', 'START', 'END']).optional(),
 });