fix(amazon-bedrock): support Cohere embedding models (#11213)

## Background

Cohere embedding model IDs (e.g. `cohere.embed-english-v3`,
`cohere.embed-multilingual-v3`) are listed as supported by
`@ai-sdk/amazon-bedrock`, but embedding requests fail at runtime because
Cohere-on-Bedrock requires a different request schema than Titan
(notably a required `input_type`). This prevents Cohere embedding models
from being usable through the Bedrock provider and triggers a 400
“Malformed input request” error. Fixes
[#5055](#5055).

## Summary

- Updated `BedrockEmbeddingModel` to detect `cohere.embed-*` model IDs
and:
- Send Cohere-compatible request payload including required `input_type`
(defaulting to `search_query`) and optional `truncate`.
- Parse Cohere-style embedding responses (`{ embeddings: number[][] }`)
in addition to Titan-style responses.
- Extended `bedrockEmbeddingProviderOptions` with Cohere-specific
options:
- `inputType` (`search_document | search_query | classification |
clustering`)
  - `truncate` (`NONE | START | END`)
- Added/updated unit tests for Cohere embedding request/response
handling.
- Updated Bedrock provider docs to mention the required Cohere
`inputType`.
- Added a **patch changeset** for `@ai-sdk/amazon-bedrock`.

## Manual Verification

- Verified end-to-end by running `embed()` against Amazon Bedrock using
`cohere.embed-english-v3` and confirming the API call succeeds and
returns a numeric embedding vector (previously returned a 400 validation
error complaining about missing `input_type`).

## Checklist

- [x] Tests have been added / updated (for bug fixes / features)
- [x] Documentation has been added / updated (for bug fixes / features)
- [x] A _patch_ changeset for relevant packages has been added (for bug
fixes / features - run `pnpm changeset` in the project root)
- [x] I have reviewed this pull request (self-review)

## Future Work

- If/when Bedrock exposes additional Cohere embedding model IDs (e.g.
newer versions), we can consider documenting them explicitly, but the
provider already supports arbitrary model IDs via `string` model IDs.

## Related Issues

Fixes
- [#5055](#5055)
- [#9390](#9390)

---------

Co-authored-by: Aayush Kapoor <83492835+aayush-kapoor@users.noreply.github.com>

Author: turisanapo

.changeset/bedrock-cohere-embeddings.md

@@ -0,0 +1,5 @@
+---
+'@ai-sdk/amazon-bedrock': patch
+---
+
+Fix Cohere embedding model request format on Bedrock by sending the required `input_type` and parsing Cohere-style responses.

content/providers/01-ai-sdk-providers/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)

examples/ai-functions/src/embed-many/amazon-bedrock-cohere-v4.ts

@@ -0,0 +1,18 @@
+import { bedrock } from '@ai-sdk/amazon-bedrock';
+import { embedMany } from 'ai';
+import { run } from '../lib/run';
+
+run(async () => {
+  const { embeddings, usage, warnings } = await embedMany({
+    model: bedrock.embedding('cohere.embed-v4:0'),
+    values: [
+      'sunny day at the beach',
+      'rainy afternoon in the city',
+      'snowy night in the mountains',
+    ],
+  });
+
+  console.log(embeddings);
+  console.log(usage);
+  console.log(warnings);
+});

examples/ai-functions/src/embed-many/amazon-bedrock-cohere.ts

@@ -0,0 +1,18 @@
+import { bedrock } from '@ai-sdk/amazon-bedrock';
+import { embedMany } from 'ai';
+import { run } from '../lib/run';
+
+run(async () => {
+  const { embeddings, usage, warnings } = await embedMany({
+    model: bedrock.embedding('cohere.embed-english-v3'),
+    values: [
+      'sunny day at the beach',
+      'rainy afternoon in the city',
+      'snowy night in the mountains',
+    ],
+  });
+
+  console.log(embeddings);
+  console.log(usage);
+  console.log(warnings);
+});

examples/ai-functions/src/embed/amazon-bedrock-cohere-input-type.ts

@@ -0,0 +1,21 @@
+import { bedrock } from '@ai-sdk/amazon-bedrock';
+import { embed } from 'ai';
+import { run } from '../lib/run';
+
+run(async () => {
+  // Use 'search_document' for documents to be searched,
+  // and 'search_query' for search queries (default).
+  const { embedding, usage, warnings } = await embed({
+    model: bedrock.embedding('cohere.embed-english-v3'),
+    value: 'sunny day at the beach',
+    providerOptions: {
+      bedrock: {
+        inputType: 'search_document',
+      },
+    },
+  });
+
+  console.log(embedding);
+  console.log(usage);
+  console.log(warnings);
+});

examples/ai-functions/src/embed/amazon-bedrock-cohere-v4.ts

@@ -0,0 +1,14 @@
+import { bedrock } from '@ai-sdk/amazon-bedrock';
+import { embed } from 'ai';
+import { run } from '../lib/run';
+
+run(async () => {
+  const { embedding, usage, warnings } = await embed({
+    model: bedrock.embedding('cohere.embed-v4:0'),
+    value: 'sunny day at the beach',
+  });
+
+  console.log(embedding);
+  console.log(usage);
+  console.log(warnings);
+});

examples/ai-functions/src/embed/amazon-bedrock-cohere.ts

@@ -0,0 +1,14 @@
+import { bedrock } from '@ai-sdk/amazon-bedrock';
+import { embed } from 'ai';
+import { run } from '../lib/run';
+
+run(async () => {
+  const { embedding, usage, warnings } = await embed({
+    model: bedrock.embedding('cohere.embed-english-v3'),
+    value: 'sunny day at the beach',
+  });
+
+  console.log(embedding);
+  console.log(usage);
+  console.log(warnings);
+});

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

@@ -16,6 +16,14 @@ const embedUrl = `https://bedrock-runtime.us-east-1.amazonaws.com/model/${encode
   'amazon.titan-embed-text-v2:0',
 )}/invoke`;
 
+const cohereEmbedUrl = `https://bedrock-runtime.us-east-1.amazonaws.com/model/${encodeURIComponent(
+  'cohere.embed-english-v3',
+)}/invoke`;
+
+const cohereV4EmbedUrl = `https://bedrock-runtime.us-east-1.amazonaws.com/model/${encodeURIComponent(
+  'cohere.embed-v4:0',
+)}/invoke`;
+
 describe('doEmbed', () => {
   const mockConfigHeaders = {
     'config-header': 'config-value',
@@ -37,6 +45,32 @@ describe('doEmbed', () => {
         ),
       },
     },
+    [cohereEmbedUrl]: {
+      response: {
+        type: 'binary',
+        headers: {
+          'content-type': 'application/json',
+        },
+        body: Buffer.from(
+          JSON.stringify({
+            embeddings: [mockEmbeddings[0]],
+          }),
+        ),
+      },
+    },
+    [cohereV4EmbedUrl]: {
+      response: {
+        type: 'binary',
+        headers: {
+          'content-type': 'application/json',
+        },
+        body: Buffer.from(
+          JSON.stringify({
+            embeddings: { float: [mockEmbeddings[0]] },
+          }),
+        ),
+      },
+    },
   });
 
   const model = new BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', {
@@ -87,6 +121,52 @@ describe('doEmbed', () => {
     expect(usage?.tokens).toStrictEqual(8);
   });
 
+  it('should support Cohere embedding models', async () => {
+    const cohereModel = new BedrockEmbeddingModel('cohere.embed-english-v3', {
+      baseUrl: () => 'https://bedrock-runtime.us-east-1.amazonaws.com',
+      headers: mockConfigHeaders,
+      fetch: fakeFetchWithAuth,
+    });
+
+    const { embeddings, usage } = await cohereModel.doEmbed({
+      values: [testValues[0]],
+    });
+
+    expect(embeddings.length).toBe(1);
+    expect(embeddings[0]).toStrictEqual(mockEmbeddings[0]);
+    expect(Number.isNaN(usage?.tokens)).toBe(true);
+
+    const body = await server.calls[0].requestBodyJson;
+    expect(body).toEqual({
+      input_type: 'search_query',
+      texts: [testValues[0]],
+      truncate: undefined,
+    });
+  });
+
+  it('should support Cohere v4 embedding models', async () => {
+    const cohereV4Model = new BedrockEmbeddingModel('cohere.embed-v4:0', {
+      baseUrl: () => 'https://bedrock-runtime.us-east-1.amazonaws.com',
+      headers: mockConfigHeaders,
+      fetch: fakeFetchWithAuth,
+    });
+
+    const { embeddings, usage } = await cohereV4Model.doEmbed({
+      values: [testValues[0]],
+    });
+
+    expect(embeddings.length).toBe(1);
+    expect(embeddings[0]).toStrictEqual(mockEmbeddings[0]);
+    expect(Number.isNaN(usage?.tokens)).toBe(true);
+
+    const body = await server.calls[0].requestBodyJson;
+    expect(body).toEqual({
+      input_type: 'search_query',
+      texts: [testValues[0]],
+      truncate: undefined,
+    });
+  });
+
   it('should properly combine headers from all sources', async () => {
     const optionsHeaders = {
       'options-header': 'options-value',

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

@@ -67,11 +67,23 @@ 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). We keep the public interface stable and
+    // adapt here based on the modelId.
+    const args = this.modelId.startsWith('cohere.embed-')
+      ? {
+          // 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 +102,40 @@ export class BedrockEmbeddingModel implements EmbeddingModelV3 {
       abortSignal,
     });
 
+    const embedding =
+      'embedding' in response
+        ? response.embedding
+        : Array.isArray(response.embeddings)
+          ? response.embeddings[0]
+          : response.embeddings.float[0];
+
     return {
       warnings: [],
-      embeddings: [response.embedding],
-      usage: { tokens: response.inputTextTokenCount },
+      embeddings: [embedding],
+      usage: {
+        tokens:
+          'inputTextTokenCount' in response
+            ? response.inputTextTokenCount
+            : NaN,
+      },
     };
   }
 }
 
-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(),
+  }),
+  // 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())),
+    }),
+  }),
+]);

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

@@ -21,4 +21,18 @@ 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(),
+
+  /**
+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 for Cohere embedding models on Bedrock.
+   */
+  truncate: z.enum(['NONE', 'START', 'END']).optional(),
 });