@ai-sdk/amazon-bedrock 4.0.24 → 4.0.26

package/src/normalize-tool-call-id.test.ts

@@ -0,0 +1,72 @@
+import { describe, it, expect } from 'vitest';
+import { isMistralModel, normalizeToolCallId } from './normalize-tool-call-id';
+
+describe('isMistralModel', () => {
+  it('should return true for mistral models', () => {
+    expect(isMistralModel('mistral.mistral-7b-instruct-v0:2')).toBe(true);
+    expect(isMistralModel('mistral.mixtral-8x7b-instruct-v0:1')).toBe(true);
+    expect(isMistralModel('mistral.mistral-large-2402-v1:0')).toBe(true);
+    expect(isMistralModel('mistral.mistral-small-2402-v1:0')).toBe(true);
+    expect(isMistralModel('mistral.mistral-large-2407-v1:0')).toBe(true);
+    expect(isMistralModel('mistral.ministral-3-14b-instruct')).toBe(true);
+    expect(isMistralModel('mistral.ministral-3-8b-instruct')).toBe(true);
+  });
+
+  it('should return true for region-prefixed mistral models', () => {
+    expect(isMistralModel('us.mistral.pixtral-large-2502-v1:0')).toBe(true);
+    expect(isMistralModel('eu.mistral.mistral-large-2407-v1:0')).toBe(true);
+  });
+
+  it('should return false for non-mistral models', () => {
+    expect(isMistralModel('anthropic.claude-3-5-sonnet-20241022-v2:0')).toBe(
+      false,
+    );
+    expect(isMistralModel('amazon.nova-pro-v1:0')).toBe(false);
+    expect(isMistralModel('openai.gpt-4o')).toBe(false);
+    expect(isMistralModel('meta.llama3-70b-instruct-v1:0')).toBe(false);
+  });
+});
+
+describe('normalizeToolCallId', () => {
+  it('should return the original ID when not a Mistral model', () => {
+    const originalId = 'tooluse_bpe71yCfRu2b5i-nKGDr5g';
+    expect(normalizeToolCallId(originalId, false)).toBe(originalId);
+  });
+
+  it('should extract first 9 alphanumeric characters for Mistral models', () => {
+    // Bedrock format: tooluse_bpe71yCfRu2b5i-nKGDr5g
+    // After removing non-alphanumeric: toolusebpe71yCfRu2b5inKGDr5g
+    // First 9 chars: toolusebp
+    expect(normalizeToolCallId('tooluse_bpe71yCfRu2b5i-nKGDr5g', true)).toBe(
+      'toolusebp',
+    );
+  });
+
+  it('should handle IDs with various special characters', () => {
+    expect(normalizeToolCallId('tool-use_123ABC456', true)).toBe('tooluse12');
+    expect(normalizeToolCallId('___abc123DEF___', true)).toBe('abc123DEF');
+  });
+
+  it('should handle IDs that are already alphanumeric', () => {
+    expect(normalizeToolCallId('abcdefghi', true)).toBe('abcdefghi');
+    expect(normalizeToolCallId('abc123XYZ', true)).toBe('abc123XYZ');
+  });
+
+  it('should handle short IDs', () => {
+    expect(normalizeToolCallId('abc', true)).toBe('abc');
+    expect(normalizeToolCallId('12345', true)).toBe('12345');
+  });
+
+  it('should handle IDs with only special characters', () => {
+    expect(normalizeToolCallId('___---___', true)).toBe('');
+  });
+
+  it('should produce valid Mistral tool call IDs (9 alphanumeric chars)', () => {
+    const normalizedId = normalizeToolCallId(
+      'tooluse_bpe71yCfRu2b5i-nKGDr5g',
+      true,
+    );
+    // Verify the ID matches Mistral's requirements: ^[a-zA-Z0-9]{1,9}$
+    expect(normalizedId).toMatch(/^[a-zA-Z0-9]{1,9}$/);
+  });
+});

package/src/normalize-tool-call-id.ts

@@ -0,0 +1,36 @@
+/**
+ * Checks if the given model ID is a Mistral model.
+ * Mistral models on Bedrock are prefixed with 'mistral.' or region-prefixed like 'us.mistral.'.
+ */
+export function isMistralModel(modelId: string): boolean {
+  return modelId.includes('mistral.');
+}
+
+/**
+ * Normalizes a tool call ID for Mistral models.
+ *
+ * Mistral models require tool call IDs to match the regex `^[a-zA-Z0-9]{9}$`:
+ * - Exactly 9 characters
+ * - Alphanumeric only (no underscores, hyphens, or other characters)
+ *
+ * Bedrock generates tool call IDs in formats like `tooluse_bpe71yCfRu2b5i-nKGDr5g`,
+ * which are incompatible with Mistral's requirements.
+ *
+ * This function extracts the first 9 alphanumeric characters from the ID.
+ *
+ * @param toolCallId - The original tool call ID from Bedrock
+ * @param isMistral - Whether the model is a Mistral model
+ * @returns The normalized tool call ID (9 alphanumeric chars) if Mistral, otherwise the original ID
+ */
+export function normalizeToolCallId(
+  toolCallId: string,
+  isMistral: boolean,
+): string {
+  if (!isMistral) {
+    return toolCallId;
+  }
+
+  // Extract only alphanumeric characters and take first 9
+  const alphanumericChars = toolCallId.replace(/[^a-zA-Z0-9]/g, '');
+  return alphanumericChars.slice(0, 9);
+}

package/src/reranking/__fixtures__/bedrock-reranking.1.json

@@ -0,0 +1,12 @@
+{
+  "results": [
+    {
+      "index": 0,
+      "relevanceScore": 0.5110583305358887
+    },
+    {
+      "index": 5,
+      "relevanceScore": 0.30241215229034424
+    }
+  ]
+}

package/src/reranking/bedrock-reranking-api.ts

@@ -0,0 +1,44 @@
+import { lazySchema, zodSchema } from '@ai-sdk/provider-utils';
+import { z } from 'zod/v4';
+
+// https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_Rerank.html
+export type BedrockRerankingInput = {
+  nextToken?: string;
+  queries: [{ type: 'TEXT'; textQuery: { text: string } }];
+  rerankingConfiguration: {
+    type: 'BEDROCK_RERANKING_MODEL';
+    bedrockRerankingConfiguration: {
+      modelConfiguration: {
+        modelArn: string;
+        additionalModelRequestFields?: Record<string, unknown>;
+      };
+      numberOfResults?: number;
+    };
+  };
+  sources: {
+    type: 'INLINE';
+    inlineDocumentSource:
+      | {
+          type: 'TEXT';
+          textDocument: { text: string };
+        }
+      | {
+          type: 'JSON';
+          jsonDocument: unknown;
+        };
+  }[];
+};
+
+export const bedrockRerankingResponseSchema = lazySchema(() =>
+  zodSchema(
+    z.object({
+      results: z.array(
+        z.object({
+          index: z.number(),
+          relevanceScore: z.number(),
+        }),
+      ),
+      nextToken: z.string().optional(),
+    }),
+  ),
+);

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

@@ -0,0 +1,299 @@
+import { createTestServer } from '@ai-sdk/test-server/with-vitest';
+import fs from 'node:fs';
+import { beforeEach, describe, expect, it } from 'vitest';
+import { injectFetchHeaders } from '../inject-fetch-headers';
+import { BedrockRerankingModel } from './bedrock-reranking-model';
+import { BedrockRerankingOptions } from './bedrock-reranking-options';
+
+const fakeFetchWithAuth = injectFetchHeaders({ 'x-amz-auth': 'test-auth' });
+
+const model = new BedrockRerankingModel('cohere.rerank-v3-5:0', {
+  baseUrl: () => 'https://bedrock-agent-runtime.us-east-1.amazonaws.com',
+  region: 'us-west-2',
+  headers: {
+    'config-header': 'config-value',
+    'shared-header': 'config-shared',
+  },
+  fetch: fakeFetchWithAuth,
+});
+
+describe('doRerank', () => {
+  const server = createTestServer({
+    'https://bedrock-agent-runtime.us-east-1.amazonaws.com/rerank': {},
+  });
+
+  function prepareJsonFixtureResponse(filename: string) {
+    server.urls[
+      'https://bedrock-agent-runtime.us-east-1.amazonaws.com/rerank'
+    ].response = {
+      type: 'binary',
+      headers: { 'content-type': 'application/json' },
+      body: Buffer.from(
+        fs.readFileSync(`src/reranking/__fixtures__/${filename}.json`, 'utf8'),
+      ),
+    };
+  }
+
+  describe('json documents', () => {
+    let result: Awaited<ReturnType<typeof model.doRerank>>;
+
+    beforeEach(async () => {
+      prepareJsonFixtureResponse('bedrock-reranking.1');
+
+      result = await model.doRerank({
+        documents: {
+          type: 'object',
+          values: [
+            { example: 'sunny day at the beach' },
+            { example: 'rainy day in the city' },
+          ],
+        },
+        query: 'rainy day',
+        topN: 2,
+        providerOptions: {
+          bedrock: {
+            nextToken: 'test-token',
+            additionalModelRequestFields: {
+              test: 'test-value',
+            },
+          } satisfies BedrockRerankingOptions,
+        },
+      });
+    });
+
+    it('should send request with stringified json documents', async () => {
+      expect(await server.calls[0].requestBodyJson).toMatchInlineSnapshot(`
+        {
+          "nextToken": "test-token",
+          "queries": [
+            {
+              "textQuery": {
+                "text": "rainy day",
+              },
+              "type": "TEXT",
+            },
+          ],
+          "rerankingConfiguration": {
+            "bedrockRerankingConfiguration": {
+              "modelConfiguration": {
+                "additionalModelRequestFields": {
+                  "test": "test-value",
+                },
+                "modelArn": "arn:aws:bedrock:us-west-2::foundation-model/cohere.rerank-v3-5:0",
+              },
+              "numberOfResults": 2,
+            },
+            "type": "BEDROCK_RERANKING_MODEL",
+          },
+          "sources": [
+            {
+              "inlineDocumentSource": {
+                "jsonDocument": {
+                  "example": "sunny day at the beach",
+                },
+                "type": "JSON",
+              },
+              "type": "INLINE",
+            },
+            {
+              "inlineDocumentSource": {
+                "jsonDocument": {
+                  "example": "rainy day in the city",
+                },
+                "type": "JSON",
+              },
+              "type": "INLINE",
+            },
+          ],
+        }
+      `);
+    });
+
+    it('should send request with the correct headers', async () => {
+      expect(server.calls[0].requestHeaders).toMatchInlineSnapshot(`
+        {
+          "config-header": "config-value",
+          "content-type": "application/json",
+          "shared-header": "config-shared",
+          "x-amz-auth": "test-auth",
+        }
+      `);
+    });
+
+    it('should return result with warnings', async () => {
+      expect(result.warnings).toMatchInlineSnapshot(`undefined`);
+    });
+
+    it('should return result with the correct ranking', async () => {
+      expect(result.ranking).toMatchInlineSnapshot(`
+        [
+          {
+            "index": 0,
+            "relevanceScore": 0.5110583305358887,
+          },
+          {
+            "index": 5,
+            "relevanceScore": 0.30241215229034424,
+          },
+        ]
+      `);
+    });
+
+    it('should not return provider metadata (use response body instead)', async () => {
+      expect(result.providerMetadata).toMatchInlineSnapshot(`undefined`);
+    });
+
+    it('should return result with the correct response', async () => {
+      expect(result.response).toMatchInlineSnapshot(`
+        {
+          "body": {
+            "results": [
+              {
+                "index": 0,
+                "relevanceScore": 0.5110583305358887,
+              },
+              {
+                "index": 5,
+                "relevanceScore": 0.30241215229034424,
+              },
+            ],
+          },
+          "headers": {
+            "content-length": "171",
+            "content-type": "application/json",
+          },
+        }
+      `);
+    });
+  });
+
+  describe('text documents', () => {
+    let result: Awaited<ReturnType<typeof model.doRerank>>;
+
+    beforeEach(async () => {
+      prepareJsonFixtureResponse('bedrock-reranking.1');
+
+      result = await model.doRerank({
+        documents: {
+          type: 'text',
+          values: ['sunny day at the beach', 'rainy day in the city'],
+        },
+        query: 'rainy day',
+        topN: 2,
+        providerOptions: {
+          bedrock: {
+            nextToken: 'test-token',
+            additionalModelRequestFields: {
+              test: 'test-value',
+            },
+          } satisfies BedrockRerankingOptions,
+        },
+      });
+    });
+
+    it('should send request with text documents', async () => {
+      expect(await server.calls[0].requestBodyJson).toMatchInlineSnapshot(`
+        {
+          "nextToken": "test-token",
+          "queries": [
+            {
+              "textQuery": {
+                "text": "rainy day",
+              },
+              "type": "TEXT",
+            },
+          ],
+          "rerankingConfiguration": {
+            "bedrockRerankingConfiguration": {
+              "modelConfiguration": {
+                "additionalModelRequestFields": {
+                  "test": "test-value",
+                },
+                "modelArn": "arn:aws:bedrock:us-west-2::foundation-model/cohere.rerank-v3-5:0",
+              },
+              "numberOfResults": 2,
+            },
+            "type": "BEDROCK_RERANKING_MODEL",
+          },
+          "sources": [
+            {
+              "inlineDocumentSource": {
+                "textDocument": {
+                  "text": "sunny day at the beach",
+                },
+                "type": "TEXT",
+              },
+              "type": "INLINE",
+            },
+            {
+              "inlineDocumentSource": {
+                "textDocument": {
+                  "text": "rainy day in the city",
+                },
+                "type": "TEXT",
+              },
+              "type": "INLINE",
+            },
+          ],
+        }
+      `);
+    });
+
+    it('should send request with the correct headers', async () => {
+      expect(server.calls[0].requestHeaders).toMatchInlineSnapshot(`
+        {
+          "config-header": "config-value",
+          "content-type": "application/json",
+          "shared-header": "config-shared",
+          "x-amz-auth": "test-auth",
+        }
+      `);
+    });
+
+    it('should return result without warnings', async () => {
+      expect(result.warnings).toMatchInlineSnapshot(`undefined`);
+    });
+
+    it('should return result with the correct ranking', async () => {
+      expect(result.ranking).toMatchInlineSnapshot(`
+        [
+          {
+            "index": 0,
+            "relevanceScore": 0.5110583305358887,
+          },
+          {
+            "index": 5,
+            "relevanceScore": 0.30241215229034424,
+          },
+        ]
+      `);
+    });
+
+    it('should not return provider metadata (use response body instead)', async () => {
+      expect(result.providerMetadata).toMatchInlineSnapshot(`undefined`);
+    });
+
+    it('should return result with the correct response', async () => {
+      expect(result.response).toMatchInlineSnapshot(`
+        {
+          "body": {
+            "results": [
+              {
+                "index": 0,
+                "relevanceScore": 0.5110583305358887,
+              },
+              {
+                "index": 5,
+                "relevanceScore": 0.30241215229034424,
+              },
+            ],
+          },
+          "headers": {
+            "content-length": "171",
+            "content-type": "application/json",
+          },
+        }
+      `);
+    });
+  });
+});

package/src/reranking/bedrock-reranking-model.ts

@@ -0,0 +1,115 @@
+import { RerankingModelV3 } from '@ai-sdk/provider';
+import {
+  FetchFunction,
+  Resolvable,
+  combineHeaders,
+  createJsonErrorResponseHandler,
+  createJsonResponseHandler,
+  parseProviderOptions,
+  postJsonToApi,
+  resolve,
+} from '@ai-sdk/provider-utils';
+import { BedrockErrorSchema } from '../bedrock-error';
+import {
+  BedrockRerankingInput,
+  bedrockRerankingResponseSchema,
+} from './bedrock-reranking-api';
+import {
+  BedrockRerankingModelId,
+  bedrockRerankingOptionsSchema,
+} from './bedrock-reranking-options';
+
+type BedrockRerankingConfig = {
+  baseUrl: () => string;
+  region: string;
+  headers: Resolvable<Record<string, string | undefined>>;
+  fetch?: FetchFunction;
+};
+
+type DoRerankResponse = Awaited<ReturnType<RerankingModelV3['doRerank']>>;
+
+export class BedrockRerankingModel implements RerankingModelV3 {
+  readonly specificationVersion = 'v3';
+  readonly provider = 'amazon-bedrock';
+
+  constructor(
+    readonly modelId: BedrockRerankingModelId,
+    private readonly config: BedrockRerankingConfig,
+  ) {}
+
+  async doRerank({
+    documents,
+    headers,
+    query,
+    topN,
+    abortSignal,
+    providerOptions,
+  }: Parameters<RerankingModelV3['doRerank']>[0]): Promise<DoRerankResponse> {
+    const bedrockOptions = await parseProviderOptions({
+      provider: 'bedrock',
+      providerOptions,
+      schema: bedrockRerankingOptionsSchema,
+    });
+
+    const {
+      value: response,
+      responseHeaders,
+      rawValue,
+    } = await postJsonToApi({
+      url: `${this.config.baseUrl()}/rerank`,
+      headers: await resolve(
+        combineHeaders(await resolve(this.config.headers), headers),
+      ),
+      body: {
+        nextToken: bedrockOptions?.nextToken,
+        queries: [
+          {
+            textQuery: { text: query },
+            type: 'TEXT',
+          },
+        ],
+        rerankingConfiguration: {
+          bedrockRerankingConfiguration: {
+            modelConfiguration: {
+              modelArn: `arn:aws:bedrock:${this.config.region}::foundation-model/${this.modelId}`,
+              additionalModelRequestFields:
+                bedrockOptions?.additionalModelRequestFields,
+            },
+            numberOfResults: topN,
+          },
+          type: 'BEDROCK_RERANKING_MODEL',
+        },
+        sources: documents.values.map(value => ({
+          type: 'INLINE' as const,
+          inlineDocumentSource:
+            documents.type === 'text'
+              ? {
+                  type: 'TEXT' as const,
+                  textDocument: { text: value as string },
+                }
+              : {
+                  type: 'JSON' as const,
+                  jsonDocument: value,
+                },
+        })),
+      } satisfies BedrockRerankingInput,
+      failedResponseHandler: createJsonErrorResponseHandler({
+        errorSchema: BedrockErrorSchema,
+        errorToMessage: error => `${error.type}: ${error.message}`,
+      }),
+      successfulResponseHandler: createJsonResponseHandler(
+        bedrockRerankingResponseSchema,
+      ),
+      fetch: this.config.fetch,
+      abortSignal,
+    });
+
+    return {
+      ranking: response.results,
+      response: {
+        headers: responseHeaders,
+        body: rawValue,
+      },
+    };
+  }
+}

package/src/reranking/bedrock-reranking-options.ts

@@ -0,0 +1,36 @@
+import { lazySchema, zodSchema } from '@ai-sdk/provider-utils';
+import { z } from 'zod/v4';
+
+// https://docs.aws.amazon.com/bedrock/latest/userguide/rerank-supported.html
+export type BedrockRerankingModelId =
+  | 'amazon.rerank-v1:0'
+  | 'cohere.rerank-v3-5:0'
+  | (string & {});
+
+export type BedrockRerankingOptions = {
+  /**
+   * If the total number of results was greater than could fit in a response, a token is returned in the nextToken field. You can enter that token in this field to return the next batch of results.
+   */
+  nextToken?: string;
+
+  /**
+   * Additional model request fields to pass to the model.
+   */
+  additionalModelRequestFields?: Record<string, unknown>;
+};
+
+export const bedrockRerankingOptionsSchema = lazySchema(() =>
+  zodSchema(
+    z.object({
+      /**
+       * If the total number of results was greater than could fit in a response, a token is returned in the nextToken field. You can enter that token in this field to return the next batch of results.
+       */
+      nextToken: z.string().optional(),
+
+      /**
+       * Additional model request fields to pass to the model.
+       */
+      additionalModelRequestFields: z.record(z.string(), z.any()).optional(),
+    }),
+  ),
+);

package/src/version.ts

@@ -0,0 +1,6 @@
+// Version string of this package injected at build time.
+declare const __PACKAGE_VERSION__: string | undefined;
+export const VERSION: string =
+  typeof __PACKAGE_VERSION__ !== 'undefined'
+    ? __PACKAGE_VERSION__
+    : '0.0.0-test';