@ai-sdk/amazon-bedrock/src/amazon-bedrock-embedding-model.ts

import { getAmazonBedrockEmbeddingModelFamily } from './amazon-bedrock-model-catalog';
import {
  TooManyEmbeddingValuesForCallError,
  type EmbeddingModelV4,
} from '@ai-sdk/provider';
import {
  lazySchema,
  zodSchema,
  combineHeaders,
  createJsonResponseHandler,
  parseProviderOptions,
  postJsonToApi,
  resolve,
  serializeModelOptions,
  WORKFLOW_SERIALIZE,
  WORKFLOW_DESERIALIZE,
  type FetchFunction,
  type Resolvable,
} from '@ai-sdk/provider-utils';
import {
  amazonBedrockEmbeddingModelOptionsSchema,
  type AmazonBedrockEmbeddingModelId,
  type AmazonBedrockEmbeddingModelSettings,
} from './amazon-bedrock-embedding-model-options';
import { amazonBedrockFailedResponseHandler } from './amazon-bedrock-error';
import { z } from 'zod/v4';

type AmazonBedrockEmbeddingConfig = {
  baseUrl: () => string;
  headers?: Resolvable<Record<string, string | undefined>>;
  fetch?: FetchFunction;
  modelFamily?: AmazonBedrockEmbeddingModelSettings['modelFamily'];
};

type DoEmbedResponse = Awaited<ReturnType<EmbeddingModelV4['doEmbed']>>;

export class AmazonBedrockEmbeddingModel implements EmbeddingModelV4 {
  readonly specificationVersion = 'v4';
  readonly provider = 'amazon-bedrock';
  readonly supportsParallelCalls = true;

  get maxEmbeddingsPerCall() {
    return this.modelFamily === 'cohere' ? 96 : 1;
  }

  private get modelFamily() {
    return (
      this.config.modelFamily ??
      getAmazonBedrockEmbeddingModelFamily(this.modelId)
    );
  }

  static [WORKFLOW_SERIALIZE](model: AmazonBedrockEmbeddingModel) {
    return serializeModelOptions({
      modelId: model.modelId,
      config: model.config,
    });
  }

  static [WORKFLOW_DESERIALIZE](options: {
    modelId: string;
    config: AmazonBedrockEmbeddingConfig;
  }) {
    return new AmazonBedrockEmbeddingModel(options.modelId, options.config);
  }

  constructor(
    readonly modelId: AmazonBedrockEmbeddingModelId,
    private readonly config: AmazonBedrockEmbeddingConfig,
  ) {}

  private getUrl(modelId: string): string {
    const encodedModelId = encodeURIComponent(modelId);
    return `${this.config.baseUrl()}/model/${encodedModelId}/invoke`;
  }

  async doEmbed({
    values,
    dimensions,
    headers,
    abortSignal,
    providerOptions,
  }: Parameters<EmbeddingModelV4['doEmbed']>[0]): Promise<DoEmbedResponse> {
    if (values.length > this.maxEmbeddingsPerCall) {
      throw new TooManyEmbeddingValuesForCallError({
        provider: this.provider,
        modelId: this.modelId,
        maxEmbeddingsPerCall: this.maxEmbeddingsPerCall,
        values,
      });
    }

    // Parse provider options. Prefer `amazonBedrock`; fall back to legacy
    // `bedrock` key for backward compatibility.
    const amazonBedrockOptions =
      (await parseProviderOptions({
        provider: 'amazonBedrock',
        providerOptions,
        schema: amazonBedrockEmbeddingModelOptionsSchema,
      })) ??
      (await parseProviderOptions({
        provider: 'bedrock',
        providerOptions,
        schema: amazonBedrockEmbeddingModelOptionsSchema,
      })) ??
      {};

    // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InvokeModel.html
    //
    // 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 modelFamily = this.modelFamily;

    const args =
      modelFamily === 'nova'
        ? {
            taskType: 'SINGLE_EMBEDDING',
            singleEmbeddingParams: {
              embeddingPurpose:
                amazonBedrockOptions.embeddingPurpose ?? 'GENERIC_INDEX',
              embeddingDimension:
                amazonBedrockOptions.embeddingDimension ?? dimensions ?? 1024,
              text: {
                truncationMode: amazonBedrockOptions.truncate ?? 'END',
                value: values[0],
              },
            },
          }
        : modelFamily === 'cohere'
          ? {
              // Cohere embedding models on Bedrock require `input_type`.
              // Without it, the service attempts other schema branches and rejects the request.
              input_type: amazonBedrockOptions.inputType ?? 'search_query',
              texts: values,
              truncate: amazonBedrockOptions.truncate,
              output_dimension:
                amazonBedrockOptions.outputDimension ?? dimensions,
            }
          : {
              inputText: values[0],
              dimensions: amazonBedrockOptions.dimensions ?? dimensions,
              normalize: amazonBedrockOptions.normalize,
            };

    const url = this.getUrl(this.modelId);
    const { value: response, responseHeaders } = await postJsonToApi({
      url,
      headers: await resolve(
        combineHeaders(
          this.config.headers ? await resolve(this.config.headers) : undefined,
          headers,
        ),
      ),
      body: args,
      failedResponseHandler: amazonBedrockFailedResponseHandler,
      successfulResponseHandler: createJsonResponseHandler(
        AmazonBedrockEmbeddingResponseSchema,
      ),
      fetch: this.config.fetch,
      abortSignal,
    });

    // Extract embeddings based on response format
    let embeddings: number[][];
    if ('embedding' in response) {
      // Titan response
      embeddings = [response.embedding];
    } else if (Array.isArray(response.embeddings)) {
      const firstEmbedding = response.embeddings[0];
      if (
        typeof firstEmbedding === 'object' &&
        firstEmbedding !== null &&
        'embeddingType' in firstEmbedding
      ) {
        // Nova response
        embeddings = [firstEmbedding.embedding];
      } else {
        // Cohere v3 response
        embeddings = response.embeddings as number[][];
      }
    } else {
      // Cohere v4 response
      embeddings = response.embeddings.float;
    }

    // Extract token count based on response format
    const headerTokenCount = Number(
      responseHeaders?.['x-amzn-bedrock-input-token-count'],
    );
    const tokens =
      'inputTextTokenCount' in response
        ? response.inputTextTokenCount // Titan response
        : 'inputTokenCount' in response
          ? (response.inputTokenCount ?? 0) // Nova response
          : headerTokenCount;

    return {
      embeddings,
      usage: { tokens },
      warnings: [],
    };
  }
}

const AmazonBedrockEmbeddingResponseSchema = lazySchema(() =>
  zodSchema(
    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())),
        }),
      }),
    ]),
  ),
);