@tanstack/ai-bedrock 0.1.5 → 0.2.0

package/src/embedding/embedding-provider-options.ts

@@ -0,0 +1,63 @@
+/**
+ * Provider options for the Bedrock embedding models.
+ *
+ * `dimensions` is deliberately absent from every options shape: it's a
+ * first-class top-level option on `embed()`. The adapter maps it onto each
+ * model's native field (Titan Text `dimensions`, Titan Multimodal
+ * `embeddingConfig.outputEmbeddingLength`) and rejects it for the Cohere
+ * models, whose output size is fixed.
+ */
+
+/** Options for `amazon.titan-embed-text-v2:0`. */
+export interface BedrockTitanTextEmbeddingProviderOptions {
+  /**
+   * Normalize the output embedding to unit length. Titan's service-side
+   * default is `true`; leave unset to use it.
+   */
+  normalize?: boolean
+}
+
+/**
+ * Options for `amazon.titan-embed-image-v1` (Titan Multimodal Embeddings).
+ * The model has no per-request options beyond the top-level `dimensions`
+ * (mapped onto `embeddingConfig.outputEmbeddingLength`).
+ */
+export type BedrockTitanImageEmbeddingProviderOptions = Record<string, never>
+
+/**
+ * Cohere embed input types. Cohere REQUIRES the caller to say what the
+ * embeddings are for; there is no service-side default.
+ */
+export type BedrockCohereEmbeddingInputType =
+  | 'search_document'
+  | 'search_query'
+  | 'classification'
+  | 'clustering'
+
+/** Options for `cohere.embed-english-v3` / `cohere.embed-multilingual-v3`. */
+export interface BedrockCohereEmbeddingProviderOptions {
+  /**
+   * What the embeddings will be used for (required by the Cohere API; sent
+   * as `input_type`). Use `search_document` when indexing, `search_query`
+   * when querying, `classification` / `clustering` for those downstream
+   * tasks. Because this field is required, `modelOptions` is required at the
+   * `embed()` call site for the Cohere models.
+   */
+  inputType: BedrockCohereEmbeddingInputType
+  /**
+   * How to handle inputs longer than the model's maximum token length.
+   * `NONE` (the API default) returns an error, `START` / `END` truncate
+   * from that side.
+   */
+  truncate?: 'NONE' | 'START' | 'END'
+}
+
+/**
+ * Broad union used as the adapter's base `TProviderOptions` fallback when
+ * the model isn't statically known. Per-model narrowing happens via
+ * `BedrockEmbeddingModelProviderOptionsByName`.
+ */
+export type BedrockEmbeddingProviderOptions =
+  | BedrockTitanTextEmbeddingProviderOptions
+  | BedrockTitanImageEmbeddingProviderOptions
+  | BedrockCohereEmbeddingProviderOptions

package/src/index.ts

@@ -161,6 +161,19 @@ export {
   createBedrockConverse,
   type BedrockConverseConfig,
 } from './adapters/converse-text'
+export {
+  BedrockEmbeddingAdapter,
+  bedrockEmbedding,
+  createBedrockEmbedding,
+  type BedrockEmbeddingConfig,
+} from './adapters/embedding'
+export type {
+  BedrockCohereEmbeddingInputType,
+  BedrockCohereEmbeddingProviderOptions,
+  BedrockEmbeddingProviderOptions,
+  BedrockTitanImageEmbeddingProviderOptions,
+  BedrockTitanTextEmbeddingProviderOptions,
+} from './embedding/embedding-provider-options'
 export type { BedrockConverseProviderOptions } from './converse/provider-options'
 export {
   resolveBedrockAuth,
@@ -173,6 +186,10 @@ export {
   BEDROCK_CHAT_MODELS,
   BEDROCK_RESPONSES_MODELS,
   BEDROCK_CONVERSE_MODELS,
+  BEDROCK_EMBEDDING_MODELS,
+  type BedrockEmbeddingModel,
+  type BedrockEmbeddingModelProviderOptionsByName,
+  type BedrockEmbeddingModelInputModalitiesByName,
   type BedrockChatModels,
   type BedrockResponsesModels,
   type BedrockConverseModels,

package/src/model-meta.ts

@@ -1,6 +1,12 @@
 import { GENERATED_BEDROCK_MODELS } from './model-catalog.generated.js'
 import type { BedrockTextProviderOptions } from './text/text-provider-options'
 import type { BedrockConverseProviderOptions } from './converse/provider-options'
+import type {
+  BedrockCohereEmbeddingProviderOptions,
+  BedrockEmbeddingProviderOptions,
+  BedrockTitanImageEmbeddingProviderOptions,
+  BedrockTitanTextEmbeddingProviderOptions,
+} from './embedding/embedding-provider-options'
 
 type Entry = (typeof GENERATED_BEDROCK_MODELS)[number]
 
@@ -67,3 +73,50 @@ export type ResolveInputModalities<TModel extends string> =
   TModel extends keyof BedrockModelInputModalitiesByName
     ? BedrockModelInputModalitiesByName[TModel]
     : readonly ['text']
+
+// ============================================================================
+// Embedding models
+// ============================================================================
+
+/**
+ * Embedding models reachable through Bedrock's `InvokeModel` API. These are
+ * not part of the generated Converse catalog (embedding models have no
+ * conversational surface), so they're maintained by hand here.
+ */
+export const BEDROCK_EMBEDDING_MODELS = [
+  'amazon.titan-embed-text-v2:0',
+  'amazon.titan-embed-image-v1',
+  'cohere.embed-english-v3',
+  'cohere.embed-multilingual-v3',
+] as const
+
+export type BedrockEmbeddingModel = (typeof BEDROCK_EMBEDDING_MODELS)[number]
+
+/**
+ * Type-only map from embedding model name to its provider options type.
+ * The Cohere models make `modelOptions` REQUIRED at the `embed()` call site
+ * because `inputType` is a required field.
+ */
+export type BedrockEmbeddingModelProviderOptionsByName = {
+  'amazon.titan-embed-text-v2:0': BedrockTitanTextEmbeddingProviderOptions
+  'amazon.titan-embed-image-v1': BedrockTitanImageEmbeddingProviderOptions
+  'cohere.embed-english-v3': BedrockCohereEmbeddingProviderOptions
+  'cohere.embed-multilingual-v3': BedrockCohereEmbeddingProviderOptions
+}
+
+/**
+ * Per-model input modalities for embedding models. Titan Multimodal accepts
+ * text and/or images (including fused text+image items embedded into one
+ * vector); the rest are text-only, so image inputs fail at compile time.
+ */
+export type BedrockEmbeddingModelInputModalitiesByName = {
+  'amazon.titan-embed-text-v2:0': readonly ['text']
+  'amazon.titan-embed-image-v1': readonly ['text', 'image']
+  'cohere.embed-english-v3': readonly ['text']
+  'cohere.embed-multilingual-v3': readonly ['text']
+}
+
+export type ResolveEmbeddingProviderOptions<TModel extends string> =
+  TModel extends keyof BedrockEmbeddingModelProviderOptionsByName
+    ? BedrockEmbeddingModelProviderOptionsByName[TModel]
+    : BedrockEmbeddingProviderOptions