diff --git a/.changeset/embed-activity-multimodal.md b/.changeset/embed-activity-multimodal.md new file mode 100644 index 000000000..44a595376 --- /dev/null +++ b/.changeset/embed-activity-multimodal.md @@ -0,0 +1,14 @@ +--- +'@tanstack/ai': minor +'@tanstack/ai-event-client': minor +'@tanstack/ai-openai': minor +'@tanstack/ai-gemini': minor +'@tanstack/ai-mistral': minor +'@tanstack/ai-bedrock': minor +'@tanstack/ai-ollama': minor +'@tanstack/ai-cohere': minor +--- + +Add a multimodal `embed()` activity. A single primitive covers one input or a batch — `input` accepts a string, a text part, an image part, or a fused `{ type: 'content' }` text+image item (one vector per item), with the accepted item types narrowed per model at compile time. Top-level `dimensions` requests Matryoshka output sizes where supported. Results carry `embeddings: [{ vector, index }]` plus `usage` when the provider reports it, and `embed()` participates in generation middleware, debug logging, OTel (`gen_ai.operation.name: embeddings`), and devtools events like every other activity. + +Provider adapters: `openaiEmbedding` (text-embedding-3-small/large), `geminiEmbedding` (gemini-embedding-001), `mistralEmbedding` (mistral-embed, codestral-embed), `ollamaEmbedding` (nomic-embed-text and any local model), `bedrockEmbedding` (Titan Text V2, Titan Multimodal G1 with fused text+image, Cohere Embed v3 on Bedrock), and the new `@tanstack/ai-cohere` package's `cohereEmbedding` (embed-v4.0, multimodal text+image with required `inputType`). diff --git a/docs/adapters/bedrock.md b/docs/adapters/bedrock.md index 8fec9f6ff..826f6b535 100644 --- a/docs/adapters/bedrock.md +++ b/docs/adapters/bedrock.md @@ -204,6 +204,70 @@ for await (const chunk of chat({ } ``` +## Embeddings + +Generate embedding vectors with Titan or Cohere embedding models via `InvokeModel`: + +```typescript +import { embed } from "@tanstack/ai"; +import { bedrockEmbedding } from "@tanstack/ai-bedrock"; + +const result = await embed({ + adapter: bedrockEmbedding("amazon.titan-embed-text-v2:0"), + input: ["a red guitar", "a blue drum kit"], + dimensions: 512, // 256 | 512 | 1024 +}); + +console.log(result.embeddings[0]?.vector); +``` + +Titan Multimodal embeds text and images — alone or fused into a single vector: + +```typescript +import { embed } from "@tanstack/ai"; +import { bedrockEmbedding } from "@tanstack/ai-bedrock"; + +const productPhoto = "iVBORw0KGgo..."; // base64 image data + +const result = await embed({ + adapter: bedrockEmbedding("amazon.titan-embed-image-v1"), + input: { + type: "content", + content: [ + { type: "text", content: "a red guitar" }, + { + type: "image", + source: { + type: "data", + value: productPhoto, + mimeType: "image/png", + }, + }, + ], + }, + dimensions: 1024, // 256 | 384 | 1024 +}); +``` + +Cohere Embed v3 on Bedrock is also supported (text-only, batched, requires `inputType`): + +```typescript +import { embed } from "@tanstack/ai"; +import { bedrockEmbedding } from "@tanstack/ai-bedrock"; + +const result = await embed({ + adapter: bedrockEmbedding("cohere.embed-english-v3"), + input: ["a red guitar", "a blue drum kit"], + modelOptions: { inputType: "search_document" }, +}); +``` + +> Titan models have no batch API — a batch of N items runs as N `InvokeModel` +> calls under a small concurrency cap. Titan Multimodal does not fetch remote +> image URLs; pass base64 data (or a `data:` URI). + +See the [Embeddings guide](../embeddings.md) for the full API. + ## Model Availability The adapter ships with a hand-seeded snapshot catalog (`src/model-catalog.generated.ts`) of confirmed model IDs. This catalog can be refreshed by the maintainer script `scripts/fetch-bedrock-models.ts`, which calls `ListFoundationModels` with AWS credentials. diff --git a/docs/adapters/cohere.md b/docs/adapters/cohere.md new file mode 100644 index 000000000..6553507a6 --- /dev/null +++ b/docs/adapters/cohere.md @@ -0,0 +1,161 @@ +--- +title: Cohere +id: cohere-adapter +order: 18 +description: "Use Cohere's embed-v4.0 multimodal embedding model with TanStack AI via @tanstack/ai-cohere — text and image embeddings for semantic search and RAG." +keywords: + - tanstack ai + - cohere + - embed-v4 + - embeddings + - multimodal embeddings + - semantic search + - adapter +--- + +The Cohere adapter provides access to Cohere's embed-v4.0 multimodal embedding model — text, images, and fused text+image inputs, each producing a single vector. + +## Installation + +```bash +npm install @tanstack/ai-cohere +``` + +## Basic Usage + +```typescript +import { embed } from "@tanstack/ai"; +import { cohereEmbedding } from "@tanstack/ai-cohere"; + +const result = await embed({ + adapter: cohereEmbedding("embed-v4.0"), + input: ["a red guitar", "a blue drum kit"], + modelOptions: { inputType: "search_document" }, +}); + +console.log(result.embeddings[0]?.vector); +console.log(result.usage?.promptTokens); +``` + +`inputType` is required by Cohere's API — use `search_document` at index time and `search_query` at query time (or `classification` / `clustering` for those workloads). TanStack AI enforces this at the type level: `modelOptions` is required for Cohere embedding calls. + +## Basic Usage - Custom API Key + +```typescript +import { embed } from "@tanstack/ai"; +import { createCohereEmbedding } from "@tanstack/ai-cohere"; + +const adapter = createCohereEmbedding("embed-v4.0", process.env.MY_COHERE_KEY!); + +const result = await embed({ + adapter, + input: "a red guitar", + modelOptions: { inputType: "search_query" }, +}); +``` + +## Multimodal Embeddings + +embed-v4.0 embeds images alongside text. An image part produces an image vector; a `{ type: "content" }` item fuses text and image into one vector — ideal for product catalogs and screenshot search: + +```typescript +import { embed } from "@tanstack/ai"; +import { cohereEmbedding } from "@tanstack/ai-cohere"; + +const productPhoto = "iVBORw0KGgo..."; // base64 image data + +const result = await embed({ + adapter: cohereEmbedding("embed-v4.0"), + input: [ + { + type: "image", + source: { + type: "data", + value: productPhoto, + mimeType: "image/png", + }, + }, + { + type: "content", + content: [ + { type: "text", content: "Fender Stratocaster, sunburst finish" }, + { + type: "image", + source: { + type: "data", + value: productPhoto, + mimeType: "image/png", + }, + }, + ], + }, + ], + modelOptions: { inputType: "search_document" }, +}); + +console.log(result.embeddings.length); // 2 +``` + +Cohere's API does not fetch remote image URLs. Pass base64 data (or a `data:` URI), or opt into adapter-side downloading: + +```typescript +import { embed } from "@tanstack/ai"; +import { cohereEmbedding } from "@tanstack/ai-cohere"; + +const adapter = cohereEmbedding("embed-v4.0", { allowUrlFetch: true }); + +const result = await embed({ + adapter, + input: { + type: "image", + source: { type: "url", value: "https://example.com/guitar.png" }, + }, + modelOptions: { inputType: "search_document" }, +}); +``` + +## Requesting Dimensions + +embed-v4.0 supports Matryoshka output dimensions via the top-level `dimensions` option: + +```typescript +import { embed } from "@tanstack/ai"; +import { cohereEmbedding } from "@tanstack/ai-cohere"; + +const result = await embed({ + adapter: cohereEmbedding("embed-v4.0"), + input: "a red guitar", + dimensions: 1024, // 256 | 512 | 1024 | 1536 + modelOptions: { inputType: "search_document" }, +}); +``` + +## Environment Variables + +Set your API key in environment variables: + +```bash +COHERE_API_KEY=... +``` + +Get a key from the [Cohere dashboard](https://dashboard.cohere.com/api-keys). + +## API Reference + +### `cohereEmbedding(model, config?)` + +Creates an embedding adapter using `COHERE_API_KEY` from the environment. + +- `model`: `"embed-v4.0"` +- `config.baseUrl`: override the API base URL (default `https://api.cohere.com`) +- `config.headers`: extra request headers +- `config.allowUrlFetch`: download `http(s)` image URLs and inline them as base64 (default `false`) + +### `createCohereEmbedding(model, apiKey, config?)` + +Same as `cohereEmbedding` with an explicit API key. + +## Next Steps + +- [Embeddings guide](../embeddings.md) — the full `embed()` API +- [Generation Hooks](../media/generation-hooks.md) — usage and lifecycle middleware diff --git a/docs/adapters/gemini.md b/docs/adapters/gemini.md index 0e6286d29..31da9cf6e 100644 --- a/docs/adapters/gemini.md +++ b/docs/adapters/gemini.md @@ -381,6 +381,32 @@ const result = await summarize({ console.log(result.summary); ``` +## Embeddings + +Generate embedding vectors with gemini-embedding-001: + +```typescript +import { embed } from "@tanstack/ai"; +import { geminiEmbedding } from "@tanstack/ai-gemini"; + +const result = await embed({ + adapter: geminiEmbedding("gemini-embedding-001"), + input: "a red guitar", + dimensions: 1536, + modelOptions: { + taskType: "RETRIEVAL_DOCUMENT", + }, +}); + +console.log(result.embeddings[0]?.vector); +``` + +> The Gemini API's embedding endpoint does not report token usage, so +> `result.usage` is absent. The Vertex-only multimodal embedding model +> (`multimodalembedding@001`) is not supported by this adapter. + +See the [Embeddings guide](../embeddings.md) for the full API. + ## Image Generation The Gemini adapter supports two types of image generation: diff --git a/docs/adapters/mistral.md b/docs/adapters/mistral.md index a45474957..c2b53a7a6 100644 --- a/docs/adapters/mistral.md +++ b/docs/adapters/mistral.md @@ -262,6 +262,37 @@ const stream = chat({ > All sampling parameters — including `temperature`, `top_p`, and `max_tokens` — > go inside `modelOptions` using Mistral's native (snake_case) names. +## Embeddings + +Generate embedding vectors with mistral-embed or codestral-embed: + +```typescript +import { embed } from "@tanstack/ai"; +import { mistralEmbedding } from "@tanstack/ai-mistral"; + +const result = await embed({ + adapter: mistralEmbedding("mistral-embed"), + input: ["a red guitar", "a blue drum kit"], +}); + +console.log(result.embeddings[0]?.vector); // 1024 dimensions +``` + +`mistral-embed` has fixed 1024-dimension output; `codestral-embed` (tuned for code) supports the top-level `dimensions` option: + +```typescript +import { embed } from "@tanstack/ai"; +import { mistralEmbedding } from "@tanstack/ai-mistral"; + +const result = await embed({ + adapter: mistralEmbedding("codestral-embed"), + input: "function add(a, b) { return a + b }", + dimensions: 512, +}); +``` + +See the [Embeddings guide](../embeddings.md) for the full API. + ## Environment Variables Set your API key in environment variables: diff --git a/docs/adapters/ollama.md b/docs/adapters/ollama.md index cc821149b..dbbb6cc91 100644 --- a/docs/adapters/ollama.md +++ b/docs/adapters/ollama.md @@ -218,6 +218,26 @@ const result = await summarize({ console.log(result.summary); ``` +## Embeddings + +Generate embedding vectors locally with any Ollama embedding model: + +```typescript ignore +import { embed } from "@tanstack/ai"; +import { ollamaEmbedding } from "@tanstack/ai-ollama"; + +const result = await embed({ + adapter: ollamaEmbedding("nomic-embed-text"), + input: ["a red guitar", "a blue drum kit"], +}); + +console.log(result.embeddings[0]?.vector); +``` + +Known models (`nomic-embed-text`, `mxbai-embed-large`, `all-minilm`, `snowflake-arctic-embed`, `bge-m3`, `embeddinggemma`) get autocomplete, and any other model name is accepted. Pull the model first with `ollama pull nomic-embed-text`. Output dimensions are fixed per model — the top-level `dimensions` option is not supported. + +See the [Embeddings guide](../embeddings.md) for the full API. + ## Setting Up Ollama ### 1. Install Ollama diff --git a/docs/adapters/openai.md b/docs/adapters/openai.md index d275ad260..68ecea644 100644 --- a/docs/adapters/openai.md +++ b/docs/adapters/openai.md @@ -214,6 +214,38 @@ const result = await summarize({ console.log(result.summary); ``` +## Embeddings + +Generate embedding vectors with the text-embedding-3 models: + +```typescript +import { embed } from "@tanstack/ai"; +import { openaiEmbedding } from "@tanstack/ai-openai"; + +const result = await embed({ + adapter: openaiEmbedding("text-embedding-3-small"), + input: ["a red guitar", "a blue drum kit"], +}); + +console.log(result.embeddings[0]?.vector); +console.log(result.usage?.promptTokens); +``` + +Both models support Matryoshka dimension reduction via the top-level `dimensions` option: + +```typescript +import { embed } from "@tanstack/ai"; +import { openaiEmbedding } from "@tanstack/ai-openai"; + +const result = await embed({ + adapter: openaiEmbedding("text-embedding-3-large"), + input: "a red guitar", + dimensions: 1024, +}); +``` + +See the [Embeddings guide](../embeddings.md) for the full API. + ## Image Generation Generate images with DALL-E: diff --git a/docs/config.json b/docs/config.json index ceb80c4f7..21676a637 100644 --- a/docs/config.json +++ b/docs/config.json @@ -303,6 +303,16 @@ } ] }, + { + "label": "Embeddings", + "children": [ + { + "label": "Embeddings", + "to": "embeddings", + "addedAt": "2026-07-10" + } + ] + }, { "label": "Middleware", "children": [ @@ -446,12 +456,14 @@ { "label": "Migration Guide", "to": "migration/migration", - "addedAt": "2026-04-15" + "addedAt": "2026-04-15", + "updatedAt": "2026-07-10" }, { "label": "From Vercel AI SDK", "to": "migration/migration-from-vercel-ai", - "addedAt": "2026-04-20" + "addedAt": "2026-04-20", + "updatedAt": "2026-07-10" }, { "label": "AG-UI Client Compliance", @@ -517,7 +529,7 @@ "label": "OpenAI", "to": "adapters/openai", "addedAt": "2026-04-15", - "updatedAt": "2026-07-03" + "updatedAt": "2026-07-10" }, { "label": "Anthropic", @@ -529,12 +541,13 @@ "label": "Google Gemini", "to": "adapters/gemini", "addedAt": "2026-04-15", - "updatedAt": "2026-07-01" + "updatedAt": "2026-07-10" }, { "label": "Ollama", "to": "adapters/ollama", - "addedAt": "2026-04-15" + "addedAt": "2026-04-15", + "updatedAt": "2026-07-10" }, { "label": "Grok (xAI)", @@ -550,7 +563,13 @@ }, { "label": "Mistral", - "to": "adapters/mistral" + "to": "adapters/mistral", + "updatedAt": "2026-07-10" + }, + { + "label": "Cohere", + "to": "adapters/cohere", + "addedAt": "2026-07-10" }, { "label": "ElevenLabs", @@ -604,7 +623,8 @@ { "label": "Amazon Bedrock", "to": "adapters/bedrock", - "addedAt": "2026-06-25" + "addedAt": "2026-06-25", + "updatedAt": "2026-07-10" } ] }, diff --git a/docs/embeddings.md b/docs/embeddings.md new file mode 100644 index 000000000..a0d1cd251 --- /dev/null +++ b/docs/embeddings.md @@ -0,0 +1,258 @@ +--- +title: Embeddings +id: embeddings +order: 1 +description: "Generate text and multimodal embedding vectors with OpenAI, Cohere, Gemini, Mistral, Amazon Bedrock, and Ollama via TanStack AI's embed() API." +keywords: + - tanstack ai + - embeddings + - embedding vectors + - multimodal embeddings + - semantic search + - rag + - embed + - openai + - cohere + - gemini + - mistral + - bedrock + - ollama +--- + +# Embeddings + +TanStack AI provides embedding generation through dedicated embedding adapters that follow the same tree-shakeable, per-model-typed architecture as every other activity. The `embed()` function turns text — and, for multimodal models, images — into vectors for semantic search, RAG, clustering, and classification. + +## Overview + +Currently supported: + +- **OpenAI**: text-embedding-3-small, text-embedding-3-large (text) +- **Cohere**: embed-v4.0 (text + image, fused multimodal) +- **Google Gemini**: gemini-embedding-001 (text) +- **Mistral**: mistral-embed, codestral-embed (text) +- **Amazon Bedrock**: Titan Text Embeddings V2 (text), Titan Multimodal Embeddings G1 (text + image), Cohere Embed v3 on Bedrock (text) +- **Ollama**: nomic-embed-text, mxbai-embed-large, and any local embedding model (text) + +## Basic Usage + +```typescript +import { embed } from "@tanstack/ai"; +import { openaiEmbedding } from "@tanstack/ai-openai"; + +const result = await embed({ + adapter: openaiEmbedding("text-embedding-3-small"), + input: "a red guitar", +}); + +console.log(result.embeddings[0]?.vector); // number[] +``` + +`input` accepts a single item or an array of items; the result always carries an `embeddings` array with one vector per input item, in input order: + +```typescript +import { embed } from "@tanstack/ai"; +import { openaiEmbedding } from "@tanstack/ai-openai"; + +const result = await embed({ + adapter: openaiEmbedding("text-embedding-3-large"), + input: ["a red guitar", "a blue drum kit", "a vintage synthesizer"], +}); + +for (const embedding of result.embeddings) { + console.log(embedding.index, embedding.vector.length); +} +``` + +## Requesting Dimensions + +Models with configurable (Matryoshka) dimensions accept a top-level `dimensions` option: + +```typescript +import { embed } from "@tanstack/ai"; +import { openaiEmbedding } from "@tanstack/ai-openai"; + +const result = await embed({ + adapter: openaiEmbedding("text-embedding-3-large"), + input: "a red guitar", + dimensions: 1024, +}); +``` + +Adapters for fixed-dimension models (for example `mistral-embed` or Ollama models) throw a clear runtime error when `dimensions` is set. + +## Multimodal Embeddings + +Multimodal models embed images — alone, or fused with text into a single vector. Image inputs reuse the same content-part shapes as chat messages, and the accepted item types are narrowed per model at compile time: passing an image to a text-only model is a type error. + +Each item in the input array produces exactly one vector: + +- a string or text part embeds that text +- an image part embeds that image +- a `{ type: "content" }` item fuses its text and image parts into one vector + +```typescript +import { embed } from "@tanstack/ai"; +import { cohereEmbedding } from "@tanstack/ai-cohere"; + +const productPhoto = "iVBORw0KGgo..."; // base64 image data + +const result = await embed({ + adapter: cohereEmbedding("embed-v4.0"), + input: [ + "a red guitar", + { + type: "image", + source: { + type: "data", + value: productPhoto, + mimeType: "image/png", + }, + }, + { + type: "content", + content: [ + { type: "text", content: "Fender Stratocaster, sunburst finish" }, + { + type: "image", + source: { + type: "data", + value: productPhoto, + mimeType: "image/png", + }, + }, + ], + }, + ], + modelOptions: { inputType: "search_document" }, +}); + +console.log(result.embeddings.length); // 3 — one vector per input item +``` + +Amazon Titan Multimodal works the same way: + +```typescript +import { embed } from "@tanstack/ai"; +import { bedrockEmbedding } from "@tanstack/ai-bedrock"; + +const productPhoto = "iVBORw0KGgo..."; // base64 image data + +const result = await embed({ + adapter: bedrockEmbedding("amazon.titan-embed-image-v1"), + input: { + type: "content", + content: [ + { type: "text", content: "a red guitar" }, + { + type: "image", + source: { + type: "data", + value: productPhoto, + mimeType: "image/png", + }, + }, + ], + }, + dimensions: 1024, +}); +``` + +Adapters do not fetch remote image URLs by default — pass base64 data (or a `data:` URI). The Cohere adapter accepts an `allowUrlFetch` config option to opt into downloading `http(s)` image URLs on your behalf. + +## Search Documents vs. Queries + +Retrieval-tuned models embed documents and queries differently. Cohere requires an `inputType`, which TanStack AI enforces at the type level — `modelOptions` is required for models with required options: + +```typescript +import { embed } from "@tanstack/ai"; +import { cohereEmbedding } from "@tanstack/ai-cohere"; + +// Index time: embed documents +const docs = await embed({ + adapter: cohereEmbedding("embed-v4.0"), + input: ["doc one", "doc two"], + modelOptions: { inputType: "search_document" }, +}); + +// Query time: embed the query +const query = await embed({ + adapter: cohereEmbedding("embed-v4.0"), + input: "which doc mentions one?", + modelOptions: { inputType: "search_query" }, +}); +``` + +Gemini expresses the same idea through an optional `taskType`: + +```typescript +import { embed } from "@tanstack/ai"; +import { geminiEmbedding } from "@tanstack/ai-gemini"; + +const result = await embed({ + adapter: geminiEmbedding("gemini-embedding-001"), + input: "a red guitar", + modelOptions: { taskType: "RETRIEVAL_DOCUMENT" }, +}); +``` + +## Usage and Observability + +Adapters report token usage when the provider does, and `embed()` supports the same observe-only generation middleware as the media activities (see [Generation Hooks](./media/generation-hooks.md)): + +```typescript +import { embed } from "@tanstack/ai"; +import { openaiEmbedding } from "@tanstack/ai-openai"; + +const result = await embed({ + adapter: openaiEmbedding("text-embedding-3-small"), + input: ["a red guitar", "a blue drum kit"], + middleware: [ + { + name: "usage-logger", + onUsage: (ctx, usage) => { + console.log(`${ctx.model}: ${usage.promptTokens} tokens`); + }, + }, + ], +}); + +console.log(result.usage?.promptTokens); +``` + +## Provider Support Matrix + +| Provider | Models | Modalities | `dimensions` | Usage reported | +| --- | --- | --- | --- | --- | +| OpenAI | text-embedding-3-small, text-embedding-3-large | text | ✅ | ✅ | +| Cohere | embed-v4.0 | text + image | ✅ | ✅ | +| Gemini | gemini-embedding-001 | text | ✅ | ❌ | +| Mistral | mistral-embed | text | ❌ | ✅ | +| Mistral | codestral-embed | text | ✅ | ✅ | +| Bedrock | amazon.titan-embed-text-v2:0 | text | ✅ (256/512/1024) | ✅ | +| Bedrock | amazon.titan-embed-image-v1 | text + image | ✅ (256/384/1024) | ✅ | +| Bedrock | cohere.embed-english-v3, cohere.embed-multilingual-v3 | text | ❌ | ❌ | +| Ollama | nomic-embed-text, mxbai-embed-large, … | text | ❌ | ✅ | + +Notes: + +- Bedrock Titan models have no batch API — a batch of N items runs as N requests under a small concurrency cap. +- Gemini's Vertex-only multimodal embedding model (`multimodalembedding@001`) is not supported; `@tanstack/ai-gemini` targets the Gemini API. + +## Error Handling + +`embed()` rejects with the provider error; middleware `onError` hooks run before the rejection propagates: + +```typescript +import { embed } from "@tanstack/ai"; +import { openaiEmbedding } from "@tanstack/ai-openai"; + +try { + await embed({ + adapter: openaiEmbedding("text-embedding-3-small"), + input: "a red guitar", + }); +} catch (error) { + console.error("embedding failed", error); +} +``` diff --git a/docs/migration/migration-from-vercel-ai.md b/docs/migration/migration-from-vercel-ai.md index 680473492..e596a7d92 100644 --- a/docs/migration/migration-from-vercel-ai.md +++ b/docs/migration/migration-from-vercel-ai.md @@ -1028,7 +1028,12 @@ generateImage({ model: openai.image('dall-e-3'), ... }) #### After (TanStack AI) ```typescript ignore -import { openaiText, openaiImage, openaiSpeech } from '@tanstack/ai-openai' +import { + openaiText, + openaiImage, + openaiSpeech, + openaiEmbedding, +} from '@tanstack/ai-openai' // Chat chat({ adapter: openaiText('gpt-4o'), ... }) @@ -1039,7 +1044,8 @@ generateImage({ adapter: openaiImage('dall-e-3'), ... }) // Text to speech generateSpeech({ adapter: openaiSpeech('tts-1'), ... }) -// Embeddings: Use OpenAI SDK directly or your vector DB's built-in support +// Embeddings +embed({ adapter: openaiEmbedding('text-embedding-3-small'), ... }) ``` ### Anthropic @@ -1405,24 +1411,52 @@ const text = await streamToText(stream) For structured (non-streaming) output — the `generateObject` equivalent — pass `outputSchema` instead; see [Structured Output](#structured-output). -## Features Not Yet Covered +## Embeddings -A few AI SDK features don't have direct TanStack AI equivalents today: +Vercel's `embed` and `embedMany` both map to TanStack AI's single `embed()` function — `input` accepts one item or an array, and the result always carries one vector per input item. -### Embeddings +### Before (Vercel AI SDK) -TanStack AI doesn't include embeddings. Use your provider's SDK directly, or the built-in embedding support most vector DBs already offer: +```typescript ignore +import { embed, embedMany } from 'ai' +import { openai } from '@ai-sdk/openai' + +const { embedding } = await embed({ + model: openai.embedding('text-embedding-3-small'), + value: 'Hello, world!', +}) + +const { embeddings } = await embedMany({ + model: openai.embedding('text-embedding-3-small'), + values: ['one', 'two'], +}) +``` + +### After (TanStack AI) ```typescript -import OpenAI from 'openai' +import { embed } from '@tanstack/ai' +import { openaiEmbedding } from '@tanstack/ai-openai' -const openaiClient = new OpenAI() -const result = await openaiClient.embeddings.create({ - model: 'text-embedding-3-small', +const single = await embed({ + adapter: openaiEmbedding('text-embedding-3-small'), input: 'Hello, world!', }) +const vector = single.embeddings[0]?.vector + +const batch = await embed({ + adapter: openaiEmbedding('text-embedding-3-small'), + input: ['one', 'two'], +}) +const vectors = batch.embeddings.map((e) => e.vector) ``` +TanStack AI's `embed()` additionally supports multimodal (text + image) inputs for models like Cohere embed-v4.0 and Amazon Titan Multimodal — see the [Embeddings guide](../embeddings.md). + +## Features Not Yet Covered + +A few AI SDK features don't have direct TanStack AI equivalents today: + ### Partial object streaming (`streamObject().elementStream` / `partialObjectStream`) TanStack AI's `outputSchema` always returns a `Promise` once the full response has validated. If you need to render partial JSON as it streams, stay on `streamText` + `onChunk` for that specific case (or parse from TanStack's raw stream with your own incremental JSON parser). diff --git a/docs/migration/migration.md b/docs/migration/migration.md index 85b52b217..027b39a43 100644 --- a/docs/migration/migration.md +++ b/docs/migration/migration.md @@ -2,7 +2,7 @@ title: Migration Guide id: migration order: 1 -description: "Migrate existing TanStack AI code to the latest version — adapter function splits, flattened options, renamed modelOptions, and removed embeddings." +description: "Migrate existing TanStack AI code to the latest version — adapter function splits, flattened options, renamed modelOptions, and the old embedding() API's replacement by embed()." keywords: - tanstack ai - migration @@ -25,7 +25,7 @@ The main breaking changes in this release are: 2. **Common options flattened** - Options are now flattened in the config instead of nested 3. **`providerOptions` renamed** - Now called `modelOptions` for clarity 4. **`toResponseStream` renamed** - Now called `toServerSentEventsStream` for clarity -5. **Embeddings removed** - Embeddings support has been removed (most vector DB services have built-in support) +5. **Embeddings removed, then reintroduced** - The old `embedding()` API was removed; embeddings are back as the new `embed()` activity with multimodal support ## 1. Adapter Functions Split @@ -330,9 +330,9 @@ return new Response(readableStream, { }) ``` -## 5. Embeddings Removed +## 5. Embeddings Removed, Then Reintroduced as `embed()` -Embeddings support has been removed from TanStack AI. Most vector database services (like Pinecone, Weaviate, Qdrant, etc.) have built-in support for embeddings, and most applications pick an embedding model and stick with it. +The original `embedding()` function was removed from TanStack AI. Embeddings have since returned as a new activity with a different API: a single `embed()` function with multimodal input support and per-model type safety. The old `embedding()` name and adapter factories (like `openaiEmbed()`) are not coming back. ### Before @@ -349,26 +349,26 @@ const result = await embedding({ ### After -Use your vector database service's built-in embedding support, or call the provider's API directly: - ```typescript -// Example with OpenAI SDK directly -import OpenAI from 'openai' - -const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY }) +import { embed } from '@tanstack/ai' +import { openaiEmbedding } from '@tanstack/ai-openai' -const result = await openai.embeddings.create({ - model: 'text-embedding-3-small', +const result = await embed({ + adapter: openaiEmbedding('text-embedding-3-small'), input: 'Hello, world!', }) + +console.log(result.embeddings[0]?.vector) ``` -### Why This Change? +Key differences from the old API: + +- **Model moves into the adapter factory** - `openaiEmbedding('text-embedding-3-small')` instead of a separate `model` option, matching every other activity +- **Single function for one or many inputs** - `input` accepts a single item or an array; the result always has an `embeddings` array with one vector per input item +- **Multimodal inputs** - models like Cohere embed-v4.0 and Amazon Titan Multimodal accept image parts and fused text+image items +- **Top-level `dimensions`** - request Matryoshka dimensions without provider-specific options -- **Vector DB services handle it** - Most vector databases have native embedding support -- **Simpler API** - Reduces API surface area and complexity -- **Direct provider access** - You can use the provider SDK directly for embeddings -- **Focused scope** - TanStack AI focuses on chat, tools, and agentic workflows +See the [Embeddings guide](../embeddings.md) for full usage. ## 6. Provider Tools Moved to `/tools` Subpath diff --git a/packages/ai-bedrock/src/adapters/embedding.ts b/packages/ai-bedrock/src/adapters/embedding.ts new file mode 100644 index 000000000..e40846cad --- /dev/null +++ b/packages/ai-bedrock/src/adapters/embedding.ts @@ -0,0 +1,547 @@ +import { BaseEmbeddingAdapter } from '@tanstack/ai/adapters' +import { toRunErrorPayload } from '@tanstack/ai/adapter-internals' +import { + requireTextOnlyEmbeddingInput, + resolveEmbeddingInput, +} from '@tanstack/ai' +import { resolveBedrockAuth } from '../utils/auth' +import { BEDROCK_EMBEDDING_MODELS } from '../model-meta' +import type * as BedrockRuntime from '@aws-sdk/client-bedrock-runtime' +import type { + BedrockRuntimeClient, + BedrockRuntimeClientConfig, +} from '@aws-sdk/client-bedrock-runtime' +import type { + EmbeddingOptions, + EmbeddingResult, + ImagePart, + TokenUsage, +} from '@tanstack/ai' +import type { ResolvedBedrockAuth } from '../utils/auth' +import type { BedrockClientConfig } from '../utils/client' +import type { + BedrockEmbeddingModel, + BedrockEmbeddingModelInputModalitiesByName, + BedrockEmbeddingModelProviderOptionsByName, + ResolveEmbeddingProviderOptions, +} from '../model-meta' + +/** + * Config for the Bedrock embedding adapter — the same auth surface as the + * other Bedrock adapters (apiKey → env → SigV4 via `resolveBedrockAuth`), + * minus the OpenAI-compat client options that don't apply to `InvokeModel`. + */ +export interface BedrockEmbeddingConfig extends Pick< + BedrockClientConfig, + 'apiKey' | 'region' | 'auth' | 'baseURL' +> {} + +/** InvokeModel calls issued concurrently during a per-item fan-out. */ +const MAX_CONCURRENT_INVOCATIONS = 5 + +/** Valid `dimensions` for `amazon.titan-embed-text-v2:0`. */ +const TITAN_TEXT_DIMENSIONS: ReadonlyArray = [256, 512, 1024] + +/** Valid `dimensions` (outputEmbeddingLength) for `amazon.titan-embed-image-v1`. */ +const TITAN_IMAGE_DIMENSIONS: ReadonlyArray = [256, 384, 1024] +const TITAN_IMAGE_DEFAULT_DIMENSIONS = 1024 + +/** Cohere embed accepts at most 96 texts per InvokeModel call. */ +const COHERE_MAX_BATCH_SIZE = 96 + +/** + * Bedrock Embedding Adapter + * + * Tree-shakeable adapter for embeddings served through Bedrock's native + * `InvokeModel` API (embedding models have no Converse surface). Each model + * family has its own JSON body dialect: + * + * - `amazon.titan-embed-text-v2:0` — text-only, ONE text per call; the batch + * is fanned out with a small concurrency cap and per-call + * `inputTextTokenCount`s are summed into usage. + * - `amazon.titan-embed-image-v1` — MULTIMODAL: text, image, or a fused + * text+image item embedded into a single vector; one item per call. + * - `cohere.embed-english-v3` / `cohere.embed-multilingual-v3` — text-only, + * batched natively (chunked at 96 texts per call). `inputType` is required. + * + * The SDK call lives behind a protected `invokeModel` seam so tests can + * subclass and inject canned response bodies without a real AWS request, and + * the AWS SDK itself is imported lazily (it's Node/server-only). + */ +export class BedrockEmbeddingAdapter< + TModel extends BedrockEmbeddingModel, + // Same rationale as the text adapters: the base parameterises + // `TProviderOptions extends object`, and the per-model options interfaces + // lack implicit index signatures — `Record` (not `unknown`) + // accepts them. Confined to the generic constraint; no value cast. + TProviderOptions extends Record = + ResolveEmbeddingProviderOptions, +> extends BaseEmbeddingAdapter< + TModel, + TProviderOptions, + BedrockEmbeddingModelProviderOptionsByName, + BedrockEmbeddingModelInputModalitiesByName +> { + readonly name = 'bedrock' as const + private clientPromise?: Promise + private readonly clientConfig: BedrockEmbeddingConfig + + constructor(config: BedrockEmbeddingConfig, model: TModel) { + super(model, {}) + // Defer client construction and auth resolution: the AWS SDK is Node/ + // server-only, so we must not pull it into the static graph here. The + // client (and its dynamic import) is built lazily on first SDK call. + this.clientConfig = config + } + + /** + * Dynamically import `@aws-sdk/client-bedrock-runtime`. The specifier is + * held in a variable (not a string literal) so bundler dep scanners cannot + * statically discover the AWS SDK and try to pre-bundle it for the browser. + * Same pattern as the Converse text adapter. + */ + protected importBedrockRuntime(): Promise { + const mod = '@aws-sdk/client-bedrock-runtime' + return import(/* @vite-ignore */ mod) as Promise + } + + /** + * Lazily construct the `BedrockRuntimeClient`, deferring + * `resolveBedrockAuth` until a real request is made. + */ + protected async getClient(): Promise { + if (!this.clientPromise) { + this.clientPromise = (async () => { + const { BedrockRuntimeClient } = await this.importBedrockRuntime() + const region = this.clientConfig.region ?? 'us-east-1' + const resolved = resolveBedrockAuth( + { + apiKey: this.clientConfig.apiKey, + region, + auth: this.clientConfig.auth, + }, + 'runtime', + ) + return new BedrockRuntimeClient( + this.buildClientConfig(resolved, region, this.clientConfig.baseURL), + ) + })().catch((error: unknown) => { + // Don't cache a rejected promise — clear it so a later call can retry + // (e.g. after a transient import failure or fixed auth config). + this.clientPromise = undefined + throw error + }) + } + return this.clientPromise + } + + /** + * Map resolved auth + endpoint to a `BedrockRuntimeClientConfig`. Bearer + * auth needs `authSchemePreference` pinned or the SDK still tries SigV4 + * first — same reasoning as the Converse text adapter. + */ + protected buildClientConfig( + resolved: ResolvedBedrockAuth, + region: string, + endpoint: string | undefined, + ): BedrockRuntimeClientConfig { + if (resolved.kind === 'bearer') { + return { + region, + token: { token: resolved.token }, + authSchemePreference: ['httpBearerAuth'], + ...(endpoint ? { endpoint } : {}), + } + } + return { + region: resolved.region, + credentials: resolved.credentials, + ...(endpoint ? { endpoint } : {}), + } + } + + // --------------------------------------------------------------------------- + // SDK seam (overridden in tests so no real AWS call happens) + // --------------------------------------------------------------------------- + + /** Send one InvokeModel call and parse its JSON response body. */ + protected async invokeModel( + modelId: string, + body: Record, + ): Promise { + const { InvokeModelCommand } = await this.importBedrockRuntime() + const client = await this.getClient() + const response = await client.send( + new InvokeModelCommand({ + modelId, + contentType: 'application/json', + accept: 'application/json', + body: JSON.stringify(body), + }), + ) + return JSON.parse(new TextDecoder().decode(response.body)) + } + + // --------------------------------------------------------------------------- + // Public adapter surface + // --------------------------------------------------------------------------- + + async createEmbeddings( + options: EmbeddingOptions, + ): Promise { + const { model, logger } = options + try { + logger.request( + `activity=embed provider=${this.name} model=${model} inputs=${options.input.length}`, + { provider: this.name, model }, + ) + switch (model) { + case 'amazon.titan-embed-text-v2:0': + return await this.embedTitanText(options) + case 'amazon.titan-embed-image-v1': + return await this.embedTitanImage(options) + case 'cohere.embed-english-v3': + case 'cohere.embed-multilingual-v3': + return await this.embedCohere(options) + default: + throw new Error( + `Unknown Bedrock embedding model "${model}". Supported models: ` + + `${BEDROCK_EMBEDDING_MODELS.join(', ')}.`, + ) + } + } catch (error: unknown) { + logger.errors(`${this.name}.createEmbeddings fatal`, { + error: toRunErrorPayload(error, `${this.name}.createEmbeddings failed`), + source: `${this.name}.createEmbeddings`, + }) + throw error + } + } + + // --------------------------------------------------------------------------- + // Per-model request mapping + // --------------------------------------------------------------------------- + + /** + * `amazon.titan-embed-text-v2:0` — one text per InvokeModel call, fanned + * out with a concurrency cap; result order matches input order and per-call + * `inputTextTokenCount`s are summed into usage. + */ + private async embedTitanText( + options: EmbeddingOptions, + ): Promise { + const { model, dimensions } = options + if ( + dimensions !== undefined && + !TITAN_TEXT_DIMENSIONS.includes(dimensions) + ) { + throw new Error( + `${model} supports dimensions 256, 512, or 1024; got ${dimensions}`, + ) + } + const normalize: boolean | undefined = options.modelOptions?.normalize + const texts = requireTextOnlyEmbeddingInput(options.input, this.name, model) + + const responses = await mapWithConcurrency( + texts, + MAX_CONCURRENT_INVOCATIONS, + async (text) => { + // Built incrementally: exactOptionalPropertyTypes is on, and Titan + // rejects explicit nulls/undefined for absent optional fields. + const body: Record = { inputText: text } + if (dimensions !== undefined) body.dimensions = dimensions + if (normalize !== undefined) body.normalize = normalize + return readTitanEmbeddingBody( + await this.invokeModel(model, body), + `${this.name} ${model}`, + ) + }, + ) + + return this.toTitanResult(model, responses) + } + + /** + * `amazon.titan-embed-image-v1` (Titan Multimodal) — one item per + * InvokeModel call. An item may carry text, an image, or both (a fused + * item embedded into a single vector). Titan accepts at most one image per + * request and never fetches remote URLs. + */ + private async embedTitanImage( + options: EmbeddingOptions, + ): Promise { + const { model, dimensions } = options + const outputEmbeddingLength = dimensions ?? TITAN_IMAGE_DEFAULT_DIMENSIONS + if (!TITAN_IMAGE_DIMENSIONS.includes(outputEmbeddingLength)) { + throw new Error( + `${model} supports dimensions 256, 384, or 1024; got ${outputEmbeddingLength}`, + ) + } + const items = resolveEmbeddingInput(options.input) + + const responses = await mapWithConcurrency( + items, + MAX_CONCURRENT_INVOCATIONS, + async (item, index) => { + if (item.images.length > 1) { + throw new Error( + `${model} accepts at most one image per input item; input item ` + + `at index ${index} contains ${item.images.length} images. ` + + `Pass them as separate input items (one vector each).`, + ) + } + const body: Record = { + embeddingConfig: { outputEmbeddingLength }, + } + if (item.texts.length > 0) body.inputText = item.texts.join('\n') + const image = item.images[0] + if (image) body.inputImage = toTitanInputImage(image, model) + return readTitanEmbeddingBody( + await this.invokeModel(model, body), + `${this.name} ${model}`, + ) + }, + ) + + return this.toTitanResult(model, responses) + } + + /** + * `cohere.embed-*-v3` — natively batched (chunked at 96 texts per call, + * order preserved across chunks). `inputType` is required by the Cohere + * API; output dimensionality is fixed, so `dimensions` is rejected. + */ + private async embedCohere( + options: EmbeddingOptions, + ): Promise { + const { model, dimensions } = options + if (dimensions !== undefined) { + throw new Error( + `${model} does not support the dimensions option; its output size is fixed`, + ) + } + const inputType: string | undefined = options.modelOptions?.inputType + if (inputType === undefined) { + throw new Error( + `${model} requires modelOptions.inputType ('search_document' | ` + + `'search_query' | 'classification' | 'clustering')`, + ) + } + const truncate: string | undefined = options.modelOptions?.truncate + const texts = requireTextOnlyEmbeddingInput(options.input, this.name, model) + const batches = chunk(texts, COHERE_MAX_BATCH_SIZE) + + const responses = await mapWithConcurrency( + batches, + MAX_CONCURRENT_INVOCATIONS, + async (batch) => { + const body: Record = { + texts: batch, + input_type: inputType, + } + if (truncate !== undefined) body.truncate = truncate + return readCohereEmbeddingBody( + await this.invokeModel(model, body), + `${this.name} ${model}`, + ) + }, + ) + + return { + id: this.generateId(), + model, + embeddings: responses.flat().map((vector, index) => ({ vector, index })), + } + } + + /** Assemble an EmbeddingResult from per-item Titan responses. */ + private toTitanResult( + model: string, + responses: Array, + ): EmbeddingResult { + let promptTokens = 0 + const embeddings = responses.map((response, index) => { + promptTokens += response.inputTextTokenCount + return { vector: response.embedding, index } + }) + const usage: TokenUsage = { + promptTokens, + completionTokens: 0, + totalTokens: promptTokens, + } + return { id: this.generateId(), model, embeddings, usage } + } +} + +// --------------------------------------------------------------------------- +// Response-body narrowing (SDK JSON boundary) +// --------------------------------------------------------------------------- + +interface TitanEmbeddingBody { + embedding: Array + /** 0 when the response omits it (e.g. image-only Titan Multimodal calls). */ + inputTextTokenCount: number +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null +} + +/** Narrow a Titan InvokeModel JSON body: `{ embedding, inputTextTokenCount? }`. */ +function readTitanEmbeddingBody( + raw: unknown, + context: string, +): TitanEmbeddingBody { + const embedding = + isRecord(raw) && Array.isArray(raw.embedding) ? raw.embedding : undefined + if (!embedding) { + throw new Error( + `${context}: response body is missing the "embedding" array`, + ) + } + const inputTextTokenCount = + isRecord(raw) && typeof raw.inputTextTokenCount === 'number' + ? raw.inputTextTokenCount + : 0 + return { embedding, inputTextTokenCount } +} + +/** Narrow a Cohere InvokeModel JSON body: `{ embeddings: number[][] }` (float). */ +function readCohereEmbeddingBody( + raw: unknown, + context: string, +): Array> { + const embeddings = + isRecord(raw) && Array.isArray(raw.embeddings) ? raw.embeddings : undefined + if (!embeddings) { + throw new Error( + `${context}: response body is missing the "embeddings" array`, + ) + } + return embeddings +} + +// --------------------------------------------------------------------------- +// Input mapping helpers +// --------------------------------------------------------------------------- + +/** + * Map an ImagePart to Titan's `inputImage` (RAW base64, no data: prefix). + * Accepts `data` sources as-is and `url` sources ONLY when the value is a + * `data:` URI; Titan cannot fetch remote http(s) URLs. + */ +function toTitanInputImage(image: ImagePart, model: string): string { + const source = image.source + if (source.type === 'data') { + return source.value + } + if (source.value.startsWith('data:')) { + const comma = source.value.indexOf(',') + if (comma !== -1) { + return source.value.slice(comma + 1) + } + } + throw new Error( + `Bedrock Titan does not fetch remote image URLs; pass base64 data ` + + `(a { type: 'data' } source or a data: URI) for ${model}.`, + ) +} + +/** Split into runs of at most `size`, preserving order. */ +function chunk(items: Array, size: number): Array> { + const chunks: Array> = [] + for (let i = 0; i < items.length; i += size) { + chunks.push(items.slice(i, i + size)) + } + return chunks +} + +/** + * Map `fn` over `items` with at most `limit` calls in flight, returning + * results in input order (result[i] corresponds to items[i] regardless of + * completion order). Rejects with the first error. + */ +async function mapWithConcurrency( + items: ReadonlyArray, + limit: number, + fn: (item: T, index: number) => Promise, +): Promise> { + const results = new Array(items.length) + let next = 0 + const worker = async (): Promise => { + while (next < items.length) { + const index = next++ + const item = items[index] + if (item === undefined) continue // unreachable: index < length + results[index] = await fn(item, index) + } + } + const workerCount = Math.max(1, Math.min(limit, items.length)) + await Promise.all(Array.from({ length: workerCount }, () => worker())) + return results +} + +// --------------------------------------------------------------------------- +// Factories +// --------------------------------------------------------------------------- + +/** + * Creates a Bedrock embedding adapter with an explicit API key (bearer). + * Type resolution happens here at the call site. + * + * @param model - The model name (e.g., 'amazon.titan-embed-text-v2:0') + * @param apiKey - Your Bedrock API key + * @param config - Optional additional configuration (region, baseURL, ...) + * @returns Configured Bedrock embedding adapter instance with resolved types + * + * @example + * ```typescript + * const adapter = createBedrockEmbedding( + * 'amazon.titan-embed-text-v2:0', + * 'bedrock-api-key', + * ); + * + * const result = await embed({ + * adapter, + * input: 'a red guitar', + * }); + * ``` + */ +export function createBedrockEmbedding( + model: TModel, + apiKey: string, + config?: Omit, +): BedrockEmbeddingAdapter { + // Explicit apiKey is authoritative — spread config first so it can't override. + return new BedrockEmbeddingAdapter({ ...config, apiKey }, model) +} + +/** + * Creates a Bedrock embedding adapter using the ambient auth cascade: + * `config.apiKey` → `BEDROCK_API_KEY` → `AWS_BEARER_TOKEN_BEDROCK` → SigV4 + * (AWS credential provider chain). Auth resolves lazily on the first + * request, so `auth: 'sigv4'` never requires an API key. + * + * @param model - The model name (e.g., 'cohere.embed-english-v3') + * @param config - Optional configuration (region, auth, baseURL, ...) + * @returns Configured Bedrock embedding adapter instance with resolved types + * + * @example + * ```typescript + * const adapter = bedrockEmbedding('cohere.embed-english-v3'); + * + * const result = await embed({ + * adapter, + * input: ['a red guitar', 'a blue drum kit'], + * modelOptions: { inputType: 'search_document' }, + * }); + * + * console.log(result.embeddings[0].vector) + * ``` + */ +export function bedrockEmbedding( + model: TModel, + config?: BedrockEmbeddingConfig, +): BedrockEmbeddingAdapter { + return new BedrockEmbeddingAdapter(config ?? {}, model) +} diff --git a/packages/ai-bedrock/src/embedding/embedding-provider-options.ts b/packages/ai-bedrock/src/embedding/embedding-provider-options.ts new file mode 100644 index 000000000..e43524c3b --- /dev/null +++ b/packages/ai-bedrock/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 + +/** + * 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 diff --git a/packages/ai-bedrock/src/index.ts b/packages/ai-bedrock/src/index.ts index d2a318fe5..f8498256d 100644 --- a/packages/ai-bedrock/src/index.ts +++ b/packages/ai-bedrock/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, diff --git a/packages/ai-bedrock/src/model-meta.ts b/packages/ai-bedrock/src/model-meta.ts index dbd02ab3c..cd8ad76f2 100644 --- a/packages/ai-bedrock/src/model-meta.ts +++ b/packages/ai-bedrock/src/model-meta.ts @@ -1,6 +1,12 @@ import { GENERATED_BEDROCK_MODELS } from './model-catalog.generated' 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 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 keyof BedrockEmbeddingModelProviderOptionsByName + ? BedrockEmbeddingModelProviderOptionsByName[TModel] + : BedrockEmbeddingProviderOptions diff --git a/packages/ai-bedrock/tests/embedding-adapter.test.ts b/packages/ai-bedrock/tests/embedding-adapter.test.ts new file mode 100644 index 000000000..54d9e8c65 --- /dev/null +++ b/packages/ai-bedrock/tests/embedding-adapter.test.ts @@ -0,0 +1,452 @@ +import { describe, expect, it } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { + BedrockEmbeddingAdapter, + bedrockEmbedding, + createBedrockEmbedding, +} from '../src/adapters/embedding' +import type { BedrockEmbeddingModel } from '../src/model-meta' + +const testLogger = resolveDebugOption(false) + +/** + * Subclass that overrides the protected `invokeModel` SDK seam so no real + * AWS call happens — the same stub pattern as the Converse adapter tests. + * Bodies are captured per call; canned responses are keyed by call index. + * An optional per-call delay lets tests force out-of-order completion to + * prove result order tracks input order. + */ +class StubAdapter< + TModel extends BedrockEmbeddingModel, +> extends BedrockEmbeddingAdapter { + bodies: Array> = [] + modelIds: Array = [] + responses: Array = [] + delaysMs: Array = [] + + protected override async invokeModel( + modelId: string, + body: Record, + ): Promise { + const call = this.bodies.length + this.bodies.push(body) + this.modelIds.push(modelId) + const delay = this.delaysMs[call] + if (delay !== undefined && delay > 0) { + await new Promise((resolve) => setTimeout(resolve, delay)) + } + return this.responses[call] + } +} + +function titanResponse(embedding: Array, inputTextTokenCount?: number) { + return inputTextTokenCount === undefined + ? { embedding } + : { embedding, inputTextTokenCount } +} + +describe('Bedrock Embedding Adapter', () => { + describe('factories', () => { + it('createBedrockEmbedding creates an adapter with the provided API key', () => { + const adapter = createBedrockEmbedding( + 'amazon.titan-embed-text-v2:0', + 'test-api-key', + ) + expect(adapter).toBeInstanceOf(BedrockEmbeddingAdapter) + expect(adapter.kind).toBe('embedding') + expect(adapter.name).toBe('bedrock') + expect(adapter.model).toBe('amazon.titan-embed-text-v2:0') + }) + + it('bedrockEmbedding constructs without a key (auth resolves lazily)', () => { + const adapter = bedrockEmbedding('cohere.embed-english-v3', { + region: 'eu-west-1', + auth: 'sigv4', + }) + expect(adapter).toBeInstanceOf(BedrockEmbeddingAdapter) + expect(adapter.model).toBe('cohere.embed-english-v3') + }) + }) + + describe('amazon.titan-embed-text-v2:0', () => { + it('sends one InvokeModel call per text and sums usage', async () => { + const adapter = new StubAdapter( + { apiKey: 'k' }, + 'amazon.titan-embed-text-v2:0', + ) + adapter.responses = [ + titanResponse([0.1, 0.2], 3), + titanResponse([0.3, 0.4], 5), + titanResponse([0.5, 0.6], 7), + ] + + const result = await adapter.createEmbeddings({ + model: 'amazon.titan-embed-text-v2:0', + input: ['one', { type: 'text', content: 'two' }, 'three'], + logger: testLogger, + }) + + expect(adapter.bodies).toEqual([ + { inputText: 'one' }, + { inputText: 'two' }, + { inputText: 'three' }, + ]) + expect(adapter.modelIds).toEqual([ + 'amazon.titan-embed-text-v2:0', + 'amazon.titan-embed-text-v2:0', + 'amazon.titan-embed-text-v2:0', + ]) + expect(result.embeddings).toEqual([ + { vector: [0.1, 0.2], index: 0 }, + { vector: [0.3, 0.4], index: 1 }, + { vector: [0.5, 0.6], index: 2 }, + ]) + expect(result.usage).toEqual({ + promptTokens: 15, + completionTokens: 0, + totalTokens: 15, + }) + expect(result.model).toBe('amazon.titan-embed-text-v2:0') + expect(result.id).toContain('bedrock') + }) + + it('passes dimensions and normalize into each body', async () => { + const adapter = new StubAdapter( + { apiKey: 'k' }, + 'amazon.titan-embed-text-v2:0', + ) + adapter.responses = [titanResponse([1], 1)] + + await adapter.createEmbeddings({ + model: 'amazon.titan-embed-text-v2:0', + input: ['hello'], + dimensions: 512, + modelOptions: { normalize: true }, + logger: testLogger, + }) + + expect(adapter.bodies).toEqual([ + { inputText: 'hello', dimensions: 512, normalize: true }, + ]) + }) + + it('rejects invalid dimensions without calling the SDK', async () => { + const adapter = new StubAdapter( + { apiKey: 'k' }, + 'amazon.titan-embed-text-v2:0', + ) + await expect( + adapter.createEmbeddings({ + model: 'amazon.titan-embed-text-v2:0', + input: ['hello'], + dimensions: 300, + logger: testLogger, + }), + ).rejects.toThrow('supports dimensions 256, 512, or 1024; got 300') + expect(adapter.bodies).toHaveLength(0) + }) + + it('rejects image input with a clear error', async () => { + const adapter = new StubAdapter( + { apiKey: 'k' }, + 'amazon.titan-embed-text-v2:0', + ) + await expect( + adapter.createEmbeddings({ + model: 'amazon.titan-embed-text-v2:0', + input: [ + { + type: 'image', + source: { type: 'data', value: 'aGk=', mimeType: 'image/png' }, + }, + ], + logger: testLogger, + }), + ).rejects.toThrow('only supports text embedding inputs') + expect(adapter.bodies).toHaveLength(0) + }) + + it('preserves input order even when earlier calls finish later', async () => { + const adapter = new StubAdapter( + { apiKey: 'k' }, + 'amazon.titan-embed-text-v2:0', + ) + adapter.responses = [ + titanResponse([0], 1), + titanResponse([1], 1), + titanResponse([2], 1), + ] + // First call resolves last. + adapter.delaysMs = [30, 10, 0] + + const result = await adapter.createEmbeddings({ + model: 'amazon.titan-embed-text-v2:0', + input: ['a', 'b', 'c'], + logger: testLogger, + }) + + expect(result.embeddings).toEqual([ + { vector: [0], index: 0 }, + { vector: [1], index: 1 }, + { vector: [2], index: 2 }, + ]) + }) + }) + + describe('amazon.titan-embed-image-v1', () => { + it('maps a fused text+image item into one body (single vector)', async () => { + const adapter = new StubAdapter( + { apiKey: 'k' }, + 'amazon.titan-embed-image-v1', + ) + adapter.responses = [titanResponse([0.9], 4)] + + const result = await adapter.createEmbeddings({ + model: 'amazon.titan-embed-image-v1', + input: [ + { + type: 'content', + content: [ + { type: 'text', content: 'a red guitar' }, + { type: 'text', content: 'product photo' }, + { + type: 'image', + source: { + type: 'data', + value: 'QkFTRTY0', + mimeType: 'image/png', + }, + }, + ], + }, + ], + logger: testLogger, + }) + + expect(adapter.bodies).toEqual([ + { + embeddingConfig: { outputEmbeddingLength: 1024 }, + inputText: 'a red guitar\nproduct photo', + inputImage: 'QkFTRTY0', + }, + ]) + expect(result.embeddings).toEqual([{ vector: [0.9], index: 0 }]) + expect(result.usage).toEqual({ + promptTokens: 4, + completionTokens: 0, + totalTokens: 4, + }) + }) + + it('strips the data: URI prefix from url sources', async () => { + const adapter = new StubAdapter( + { apiKey: 'k' }, + 'amazon.titan-embed-image-v1', + ) + adapter.responses = [titanResponse([1])] + + await adapter.createEmbeddings({ + model: 'amazon.titan-embed-image-v1', + input: [ + { + type: 'image', + source: { type: 'url', value: 'data:image/png;base64,QUJD' }, + }, + ], + logger: testLogger, + }) + + expect(adapter.bodies).toEqual([ + { + embeddingConfig: { outputEmbeddingLength: 1024 }, + inputImage: 'QUJD', + }, + ]) + }) + + it('rejects http(s) image URLs with a clear error', async () => { + const adapter = new StubAdapter( + { apiKey: 'k' }, + 'amazon.titan-embed-image-v1', + ) + await expect( + adapter.createEmbeddings({ + model: 'amazon.titan-embed-image-v1', + input: [ + { + type: 'image', + source: { type: 'url', value: 'https://example.com/cat.png' }, + }, + ], + logger: testLogger, + }), + ).rejects.toThrow( + 'Bedrock Titan does not fetch remote image URLs; pass base64 data', + ) + expect(adapter.bodies).toHaveLength(0) + }) + + it('rejects multiple images in one fused item', async () => { + const adapter = new StubAdapter( + { apiKey: 'k' }, + 'amazon.titan-embed-image-v1', + ) + await expect( + adapter.createEmbeddings({ + model: 'amazon.titan-embed-image-v1', + input: [ + { + type: 'content', + content: [ + { + type: 'image', + source: { + type: 'data', + value: 'QQ==', + mimeType: 'image/png', + }, + }, + { + type: 'image', + source: { + type: 'data', + value: 'Qg==', + mimeType: 'image/png', + }, + }, + ], + }, + ], + logger: testLogger, + }), + ).rejects.toThrow('at most one image per input item') + }) + + it('maps dimensions onto embeddingConfig.outputEmbeddingLength and validates it', async () => { + const adapter = new StubAdapter( + { apiKey: 'k' }, + 'amazon.titan-embed-image-v1', + ) + adapter.responses = [titanResponse([1])] + + await adapter.createEmbeddings({ + model: 'amazon.titan-embed-image-v1', + input: ['caption only'], + dimensions: 384, + logger: testLogger, + }) + expect(adapter.bodies).toEqual([ + { + embeddingConfig: { outputEmbeddingLength: 384 }, + inputText: 'caption only', + }, + ]) + + await expect( + adapter.createEmbeddings({ + model: 'amazon.titan-embed-image-v1', + input: ['caption only'], + dimensions: 512, + logger: testLogger, + }), + ).rejects.toThrow('supports dimensions 256, 384, or 1024; got 512') + }) + }) + + describe('cohere.embed-*-v3', () => { + it('sends one batched body with required input_type and optional truncate', async () => { + const adapter = new StubAdapter( + { apiKey: 'k' }, + 'cohere.embed-english-v3', + ) + adapter.responses = [{ embeddings: [[0.1], [0.2], [0.3]] }] + + const result = await adapter.createEmbeddings({ + model: 'cohere.embed-english-v3', + input: ['one', 'two', 'three'], + modelOptions: { inputType: 'search_query', truncate: 'END' }, + logger: testLogger, + }) + + expect(adapter.bodies).toEqual([ + { + texts: ['one', 'two', 'three'], + input_type: 'search_query', + truncate: 'END', + }, + ]) + expect(result.embeddings).toEqual([ + { vector: [0.1], index: 0 }, + { vector: [0.2], index: 1 }, + { vector: [0.3], index: 2 }, + ]) + expect(result.usage).toBeUndefined() + }) + + it('chunks more than 96 texts into multiple calls, preserving order', async () => { + const adapter = new StubAdapter( + { apiKey: 'k' }, + 'cohere.embed-multilingual-v3', + ) + const texts = Array.from({ length: 100 }, (_, i) => `text-${i}`) + adapter.responses = [ + { embeddings: Array.from({ length: 96 }, (_, i) => [i]) }, + { embeddings: Array.from({ length: 4 }, (_, i) => [96 + i]) }, + ] + + const result = await adapter.createEmbeddings({ + model: 'cohere.embed-multilingual-v3', + input: texts, + modelOptions: { inputType: 'search_document' }, + logger: testLogger, + }) + + expect(adapter.bodies).toHaveLength(2) + expect(adapter.bodies[0]).toEqual({ + texts: texts.slice(0, 96), + input_type: 'search_document', + }) + expect(adapter.bodies[1]).toEqual({ + texts: texts.slice(96), + input_type: 'search_document', + }) + expect(result.embeddings).toHaveLength(100) + expect(result.embeddings[0]).toEqual({ vector: [0], index: 0 }) + expect(result.embeddings[95]).toEqual({ vector: [95], index: 95 }) + expect(result.embeddings[96]).toEqual({ vector: [96], index: 96 }) + expect(result.embeddings[99]).toEqual({ vector: [99], index: 99 }) + }) + + it('rejects the dimensions option (fixed output size)', async () => { + const adapter = new StubAdapter( + { apiKey: 'k' }, + 'cohere.embed-english-v3', + ) + await expect( + adapter.createEmbeddings({ + model: 'cohere.embed-english-v3', + input: ['hello'], + dimensions: 256, + modelOptions: { inputType: 'search_query' }, + logger: testLogger, + }), + ).rejects.toThrow('does not support the dimensions option') + expect(adapter.bodies).toHaveLength(0) + }) + + it('rejects a missing inputType at runtime (untyped callers)', async () => { + const adapter = new StubAdapter( + { apiKey: 'k' }, + 'cohere.embed-english-v3', + ) + await expect( + adapter.createEmbeddings({ + model: 'cohere.embed-english-v3', + input: ['hello'], + logger: testLogger, + }), + ).rejects.toThrow('requires modelOptions.inputType') + expect(adapter.bodies).toHaveLength(0) + }) + }) +}) diff --git a/packages/ai-cohere/README.md b/packages/ai-cohere/README.md new file mode 100644 index 000000000..a8d0ddcab --- /dev/null +++ b/packages/ai-cohere/README.md @@ -0,0 +1,104 @@ +# @tanstack/ai-cohere + +Cohere adapter for TanStack AI. + +## Installation + +```bash +npm install @tanstack/ai-cohere +# or +pnpm add @tanstack/ai-cohere +# or +yarn add @tanstack/ai-cohere +``` + +## Setup + +Get your API key from the [Cohere Dashboard](https://dashboard.cohere.com/api-keys) and set it as an environment variable: + +```bash +export COHERE_API_KEY="..." +``` + +## Usage + +### Embedding Adapter + +```typescript +import { cohereEmbedding } from '@tanstack/ai-cohere' +import { embed } from '@tanstack/ai' + +const adapter = cohereEmbedding('embed-v4.0') + +const result = await embed({ + adapter, + input: ['a red guitar', 'a blue drum kit'], + modelOptions: { inputType: 'search_document' }, +}) + +console.log(result.embeddings[0].vector) +``` + +### Multimodal Inputs + +embed-v4.0 embeds text, images, and fused text+image items (one vector per input item): + +```typescript +const result = await embed({ + adapter, + input: [ + 'a red guitar', + { + type: 'image', + source: { type: 'data', value: base64Png, mimeType: 'image/png' }, + }, + { + type: 'content', + content: [ + { type: 'text', content: 'product photo' }, + { + type: 'image', + source: { type: 'data', value: base64Jpeg, mimeType: 'image/jpeg' }, + }, + ], + }, + ], + modelOptions: { inputType: 'search_document' }, +}) +``` + +Cohere does not fetch remote image URLs. Pass base64 data or a `data:` URI, or enable `allowUrlFetch` in the adapter config to have the adapter download http(s) URLs and inline them. + +### With Explicit API Key + +```typescript +import { createCohereEmbedding } from '@tanstack/ai-cohere' + +const adapter = createCohereEmbedding('embed-v4.0', process.env.COHERE_API_KEY!) +``` + +## Supported Models + +### Embedding Models + +- `embed-v4.0` - Multimodal embedding model (text + images, Matryoshka dimensions via the top-level `dimensions` option) + +## Features + +- ✅ Embeddings (batch, one request per input array) +- ✅ Multimodal embedding input (text + images + fused text/image items) +- ✅ Dimension reduction (`dimensions` → Cohere `output_dimension`) +- ❌ Chat / text generation +- ❌ Image generation + +## Tree-Shakeable Adapters + +This package uses tree-shakeable adapters, so you only import what you need: + +```typescript +import { cohereEmbedding } from '@tanstack/ai-cohere' +``` + +## License + +MIT diff --git a/packages/ai-cohere/package.json b/packages/ai-cohere/package.json new file mode 100644 index 000000000..23247c536 --- /dev/null +++ b/packages/ai-cohere/package.json @@ -0,0 +1,62 @@ +{ + "name": "@tanstack/ai-cohere", + "version": "0.1.0", + "type": "module", + "description": "Cohere adapter for TanStack AI", + "author": "", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/TanStack/ai.git", + "directory": "packages/ai-cohere" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/tannerlinsley" + }, + "module": "./dist/esm/index.js", + "types": "./dist/esm/index.d.ts", + "exports": { + ".": { + "types": "./dist/esm/index.d.ts", + "import": "./dist/esm/index.js" + }, + "./adapters/embedding": { + "types": "./dist/esm/adapters/embedding.d.ts", + "import": "./dist/esm/adapters/embedding.js" + } + }, + "files": [ + "dist", + "src" + ], + "scripts": { + "build": "vite build", + "clean": "premove ./build ./dist", + "lint:fix": "eslint ./src --fix", + "test:build": "publint --strict", + "test:eslint": "eslint ./src", + "test:lib": "vitest run", + "test:lib:dev": "pnpm test:lib --watch", + "test:types": "tsc" + }, + "keywords": [ + "ai", + "cohere", + "tanstack", + "adapter", + "embeddings", + "multimodal" + ], + "devDependencies": { + "@tanstack/ai": "workspace:*", + "@vitest/coverage-v8": "4.0.14", + "vite": "^7.3.3" + }, + "peerDependencies": { + "@tanstack/ai": "workspace:^" + }, + "dependencies": { + "@tanstack/ai-utils": "workspace:*" + } +} diff --git a/packages/ai-cohere/src/adapters/embedding.ts b/packages/ai-cohere/src/adapters/embedding.ts new file mode 100644 index 000000000..551736b81 --- /dev/null +++ b/packages/ai-cohere/src/adapters/embedding.ts @@ -0,0 +1,369 @@ +import { BaseEmbeddingAdapter } from '@tanstack/ai/adapters' +import { toRunErrorPayload } from '@tanstack/ai/adapter-internals' +import { arrayBufferToBase64, generateId } from '@tanstack/ai-utils' +import { resolveEmbeddingInput } from '@tanstack/ai' +import { getCohereApiKeyFromEnv } from '../utils/client' +import type { + EmbeddingOptions, + EmbeddingResult, + ImagePart, + TokenUsage, +} from '@tanstack/ai' +import type { + CohereEmbeddingModel, + CohereEmbeddingModelInputModalitiesByName, + CohereEmbeddingModelProviderOptionsByName, +} from '../model-meta' +import type { CohereEmbeddingProviderOptions } from '../embedding/embedding-provider-options' +import type { CohereClientConfig } from '../utils/client' + +/** + * Configuration for Cohere embedding adapter. + */ +export interface CohereEmbeddingConfig extends CohereClientConfig {} + +const DEFAULT_BASE_URL = 'https://api.cohere.com' +const DEFAULT_TIMEOUT_MS = 30_000 + +/** + * Returns true when `url` is malformed, non-http(s), or targets a private / + * loopback / link-local host. Used to block SSRF via `allowUrlFetch`. + */ +function isPrivateOrInternalUrl(url: string): boolean { + let parsed: URL + try { + parsed = new URL(url) + } catch { + return true + } + if (parsed.protocol !== 'http:' && parsed.protocol !== 'https:') { + return true + } + const host = parsed.hostname.toLowerCase() + if ( + host === 'localhost' || + host.endsWith('.localhost') || + host === '::1' || + host === '[::1]' || + host.startsWith('127.') || + host.startsWith('10.') || + host.startsWith('192.168.') || + host.startsWith('169.254.') || + /^172\.(1[6-9]|2\d|3[01])\./.test(host) + ) { + return true + } + return false +} + +async function fetchWithTimeout( + url: string, + init: RequestInit | undefined, + timeoutMs: number, +): Promise { + const controller = new AbortController() + const timeoutId = setTimeout(() => controller.abort(), timeoutMs) + try { + return await fetch(url, { ...init, signal: controller.signal }) + } finally { + clearTimeout(timeoutId) + } +} + +/** One content part of a Cohere v2/embed fused input. */ +type CohereEmbedContentPart = + | { type: 'text'; text: string } + | { type: 'image_url'; image_url: { url: string } } + +/** Wire shape of the Cohere v2/embed request body. */ +interface CohereEmbedRequestBody { + model: string + inputs: Array<{ content: Array }> + input_type: CohereEmbeddingProviderOptions['inputType'] + embedding_types: ['float'] + truncate?: 'NONE' | 'START' | 'END' + output_dimension?: number +} + +/** Wire shape of the Cohere v2/embed response (fields the adapter reads). */ +interface CohereEmbedResponse { + id?: string + embeddings?: { + float?: Array> + } + meta?: { + billed_units?: { + input_tokens?: number + images?: number + } + } +} + +/** + * Cohere Embedding Adapter + * + * Tree-shakeable adapter for Cohere multimodal embeddings (embed-v4.0), + * implemented with plain `fetch` against the v2/embed endpoint — no Cohere + * SDK dependency. + * + * Features: + * - Batch embedding (one request for the whole input array) + * - Multimodal inputs: text, images, and fused text+image items (one vector + * per input item) + * - Matryoshka dimension reduction via the top-level `dimensions` option + * (mapped to Cohere's `output_dimension`) + */ +export class CohereEmbeddingAdapter< + TModel extends CohereEmbeddingModel, +> extends BaseEmbeddingAdapter< + TModel, + CohereEmbeddingProviderOptions, + CohereEmbeddingModelProviderOptionsByName, + CohereEmbeddingModelInputModalitiesByName +> { + readonly name = 'cohere' as const + + protected clientConfig: CohereEmbeddingConfig + + constructor(config: CohereEmbeddingConfig, model: TModel) { + super(model, {}) + this.clientConfig = config + } + + async createEmbeddings( + options: EmbeddingOptions, + ): Promise { + const { model, logger, modelOptions } = options + + try { + // The provider options type makes `modelOptions` required at the + // embed() call site; this guard covers untyped/dynamic callers. + const inputType: CohereEmbeddingProviderOptions['inputType'] | undefined = + modelOptions?.inputType + if (!inputType) { + throw new Error( + `Cohere embeddings require modelOptions.inputType ('search_document' | 'search_query' | 'classification' | 'clustering').`, + ) + } + + const resolved = resolveEmbeddingInput(options.input) + const inputs = await Promise.all( + resolved.map(async (item) => { + const content: Array = item.texts.map( + (text) => ({ type: 'text', text }), + ) + for (const image of item.images) { + content.push({ + type: 'image_url', + image_url: { url: await this.resolveImageUrl(image) }, + }) + } + return { content } + }), + ) + + // embedding_types is pinned to ['float'] (overriding any disagreeing + // modelOptions.embeddingTypes) so vectors are always number[]. + const body: CohereEmbedRequestBody = { + model, + inputs, + input_type: inputType, + embedding_types: ['float'], + } + const truncate = modelOptions?.truncate + if (truncate !== undefined) { + body.truncate = truncate + } + if (options.dimensions !== undefined) { + body.output_dimension = options.dimensions + } + + logger.request( + `activity=embed provider=${this.name} model=${model} inputs=${inputs.length}`, + { provider: this.name, model }, + ) + + const timeoutMs = this.clientConfig.timeout ?? DEFAULT_TIMEOUT_MS + const response = await fetchWithTimeout( + `${this.clientConfig.baseUrl ?? DEFAULT_BASE_URL}/v2/embed`, + { + method: 'POST', + headers: { + Authorization: `Bearer ${this.clientConfig.apiKey}`, + 'Content-Type': 'application/json', + ...this.clientConfig.headers, + }, + body: JSON.stringify(body), + }, + timeoutMs, + ) + + if (!response.ok) { + const bodyText = await response.text() + let message = bodyText + try { + const parsed: unknown = JSON.parse(bodyText) + if ( + typeof parsed === 'object' && + parsed !== null && + 'message' in parsed && + typeof parsed.message === 'string' + ) { + message = parsed.message + } + } catch { + // Not JSON — fall back to the raw body text. + } + throw new Error(`Cohere embed failed (${response.status}): ${message}`) + } + + const data = (await response.json()) as CohereEmbedResponse + + const vectors = data.embeddings?.float + if (!vectors) { + throw new Error( + 'Cohere embed response did not include float embeddings', + ) + } + if (vectors.length !== inputs.length) { + throw new Error( + `Cohere embed returned ${vectors.length} embeddings for ${inputs.length} inputs`, + ) + } + + const result: EmbeddingResult = { + id: generateId(this.name), + model, + embeddings: vectors.map((vector, index) => ({ vector, index })), + } + + const inputTokens = data.meta?.billed_units?.input_tokens + if (inputTokens !== undefined) { + const usage: TokenUsage = { + promptTokens: inputTokens, + completionTokens: 0, + totalTokens: inputTokens, + } + result.usage = usage + } + + return result + } catch (error: unknown) { + logger.errors(`${this.name}.createEmbeddings fatal`, { + error: toRunErrorPayload(error, `${this.name}.createEmbeddings failed`), + source: `${this.name}.createEmbeddings`, + }) + throw error + } + } + + /** + * Resolves an image part to a URL Cohere accepts. Cohere does not fetch + * remote image URLs, so everything is normalized to a `data:` URI unless + * the caller already provided one. + */ + protected async resolveImageUrl(image: ImagePart): Promise { + const source = image.source + + if (source.type === 'data') { + return `data:${source.mimeType};base64,${source.value}` + } + + if (source.value.startsWith('data:')) { + return source.value + } + + if (!this.clientConfig.allowUrlFetch) { + throw new Error( + 'Cohere does not fetch remote image URLs; pass base64 data or a data: URI (or enable config.allowUrlFetch to have the adapter download it)', + ) + } + + if (isPrivateOrInternalUrl(source.value)) { + throw new Error( + `Refusing to fetch internal or private URL for Cohere embedding: ${source.value}`, + ) + } + + const response = await fetchWithTimeout( + source.value, + undefined, + this.clientConfig.timeout ?? DEFAULT_TIMEOUT_MS, + ) + if (!response.ok) { + throw new Error( + `Failed to fetch image URL for Cohere embedding (${response.status}): ${source.value}`, + ) + } + const mimeType = + response.headers.get('content-type') ?? + source.mimeType ?? + 'application/octet-stream' + const base64 = arrayBufferToBase64(await response.arrayBuffer()) + return `data:${mimeType};base64,${base64}` + } +} + +/** + * Creates a Cohere embedding adapter with explicit API key. + * Type resolution happens here at the call site. + * + * @param model - The model name (e.g., 'embed-v4.0') + * @param apiKey - Your Cohere API key + * @param config - Optional additional configuration + * @returns Configured Cohere embedding adapter instance with resolved types + * + * @example + * ```typescript + * const adapter = createCohereEmbedding('embed-v4.0', 'api_key'); + * + * const result = await embed({ + * adapter, + * input: 'a red guitar', + * modelOptions: { inputType: 'search_document' } + * }); + * ``` + */ +export function createCohereEmbedding( + model: TModel, + apiKey: string, + config?: Omit, +): CohereEmbeddingAdapter { + return new CohereEmbeddingAdapter({ apiKey, ...config }, model) +} + +/** + * Creates a Cohere embedding adapter using the `COHERE_API_KEY` environment variable. + * Type resolution happens here at the call site. + * + * Looks for `COHERE_API_KEY` in: + * - `process.env` (Node.js) + * - `window.env` (Browser with injected env) + * + * @param model - The model name (e.g., 'embed-v4.0') + * @param config - Optional configuration (excluding apiKey which is auto-detected) + * @returns Configured Cohere embedding adapter instance with resolved types + * @throws Error if COHERE_API_KEY is not found in environment + * + * @example + * ```typescript + * // Automatically uses COHERE_API_KEY from environment + * const adapter = cohereEmbedding('embed-v4.0'); + * + * const result = await embed({ + * adapter, + * input: ['a red guitar', 'a blue drum kit'], + * modelOptions: { inputType: 'search_query' }, + * dimensions: 1024 + * }); + * + * console.log(result.embeddings[0].vector) + * ``` + */ +export function cohereEmbedding( + model: TModel, + config?: Omit, +): CohereEmbeddingAdapter { + const apiKey = getCohereApiKeyFromEnv() + return createCohereEmbedding(model, apiKey, config) +} diff --git a/packages/ai-cohere/src/embedding/embedding-provider-options.ts b/packages/ai-cohere/src/embedding/embedding-provider-options.ts new file mode 100644 index 000000000..35adefd80 --- /dev/null +++ b/packages/ai-cohere/src/embedding/embedding-provider-options.ts @@ -0,0 +1,44 @@ +/** + * Provider options for Cohere embedding models. + * + * `dimensions` is deliberately absent: it's a first-class top-level option on + * `embed()` and is mapped to Cohere's `output_dimension` request field by the + * adapter. + */ + +/** + * Provider options for `embed-v4.0`. + * + * `inputType` is required by Cohere's v2 embed API, which makes + * `modelOptions` required at the `embed()` call site. + */ +export interface CohereEmbeddingProviderOptions { + /** + * The intended downstream use of the embeddings. Cohere requires this to + * pick the right embedding space: + * - `search_document` — corpus items stored for later retrieval + * - `search_query` — queries run against stored documents + * - `classification` — inputs embedded for classification tasks + * - `clustering` — inputs embedded for clustering tasks + */ + inputType: + | 'search_document' + | 'search_query' + | 'classification' + | 'clustering' + + /** + * Requested embedding value encodings. The adapter always pins this to + * `['float']` so vectors are plain `number[]`; other encodings are not + * supported through TanStack AI. + */ + embeddingTypes?: ['float'] + + /** + * How to handle inputs longer than the model's maximum token length. + * `NONE` returns an error for over-long inputs; `START`/`END` truncate + * from the respective side. Defaults to Cohere's server-side default + * (`END`) when omitted. + */ + truncate?: 'NONE' | 'START' | 'END' +} diff --git a/packages/ai-cohere/src/index.ts b/packages/ai-cohere/src/index.ts new file mode 100644 index 000000000..f7900b0df --- /dev/null +++ b/packages/ai-cohere/src/index.ts @@ -0,0 +1,27 @@ +/** + * @module @tanstack/ai-cohere + * + * Cohere provider adapter for TanStack AI. + * Provides a tree-shakeable adapter for Cohere's v2/embed API + * (multimodal embeddings) using plain fetch — no SDK dependency. + */ + +// Embedding adapter - for embedding vectors +export { + CohereEmbeddingAdapter, + createCohereEmbedding, + cohereEmbedding, + type CohereEmbeddingConfig, +} from './adapters/embedding' +export type { CohereEmbeddingProviderOptions } from './embedding/embedding-provider-options' + +// Client config + env helpers +export { getCohereApiKeyFromEnv, type CohereClientConfig } from './utils/client' + +// Types +export type { + CohereEmbeddingModel, + CohereEmbeddingModelProviderOptionsByName, + CohereEmbeddingModelInputModalitiesByName, +} from './model-meta' +export { COHERE_EMBEDDING_MODELS } from './model-meta' diff --git a/packages/ai-cohere/src/model-meta.ts b/packages/ai-cohere/src/model-meta.ts new file mode 100644 index 000000000..6170be069 --- /dev/null +++ b/packages/ai-cohere/src/model-meta.ts @@ -0,0 +1,27 @@ +import type { CohereEmbeddingProviderOptions } from './embedding/embedding-provider-options' + +/** + * Embedding models (based on endpoints: "v2/embed") + */ +export const COHERE_EMBEDDING_MODELS = ['embed-v4.0'] as const + +/** + * Union type of all supported Cohere embedding model names. + */ +export type CohereEmbeddingModel = (typeof COHERE_EMBEDDING_MODELS)[number] + +/** + * Type-only map from embedding model name to its provider options type. + */ +export type CohereEmbeddingModelProviderOptionsByName = { + 'embed-v4.0': CohereEmbeddingProviderOptions +} + +/** + * Per-model input modalities for embedding models. embed-v4.0 is + * multimodal: it accepts text and image inputs (including fused + * text+image items that produce a single vector). + */ +export type CohereEmbeddingModelInputModalitiesByName = { + 'embed-v4.0': readonly ['text', 'image'] +} diff --git a/packages/ai-cohere/src/utils/client.ts b/packages/ai-cohere/src/utils/client.ts new file mode 100644 index 000000000..3bbc89fd8 --- /dev/null +++ b/packages/ai-cohere/src/utils/client.ts @@ -0,0 +1,39 @@ +import { getApiKeyFromEnv } from '@tanstack/ai-utils' + +/** + * Configuration for the Cohere HTTP client used by the adapters in this + * package. Requests are made with plain `fetch` — no Cohere SDK dependency. + */ +export interface CohereClientConfig { + /** Cohere API key. */ + apiKey: string + + /** Optional base URL override (defaults to `https://api.cohere.com`). */ + baseUrl?: string + + /** Optional default headers to include with every request. */ + headers?: Record + + /** + * Cohere's embed API does not fetch remote image URLs itself. When this is + * enabled the adapter downloads http(s) image URLs and inlines them as + * base64 `data:` URIs before sending the request. Disabled by default. + */ + allowUrlFetch?: boolean + + /** Request timeout in milliseconds for API and image URL fetches (default: 30_000). */ + timeout?: number +} + +/** + * Gets Cohere API key from environment variables. + * + * Looks for `COHERE_API_KEY` in: + * - `process.env` (Node.js) + * - `window.env` (Browser with injected env) + * + * @throws Error if COHERE_API_KEY is not found + */ +export function getCohereApiKeyFromEnv(): string { + return getApiKeyFromEnv('COHERE_API_KEY') +} diff --git a/packages/ai-cohere/tests/embedding-adapter.test.ts b/packages/ai-cohere/tests/embedding-adapter.test.ts new file mode 100644 index 000000000..5227366f8 --- /dev/null +++ b/packages/ai-cohere/tests/embedding-adapter.test.ts @@ -0,0 +1,458 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { + CohereEmbeddingAdapter, + cohereEmbedding, + createCohereEmbedding, +} from '../src/adapters/embedding' + +const testLogger = resolveDebugOption(false) + +const EMBED_URL = 'https://api.cohere.com/v2/embed' + +interface CohereEmbedResponseInit { + vectors: Array> + inputTokens?: number +} + +function embedResponse({ + vectors, + inputTokens, +}: CohereEmbedResponseInit): Response { + return new Response( + JSON.stringify({ + id: 'cohere-test-id', + embeddings: { float: vectors }, + ...(inputTokens !== undefined + ? { meta: { billed_units: { input_tokens: inputTokens } } } + : {}), + }), + { status: 200, headers: { 'content-type': 'application/json' } }, + ) +} + +const fetchMock = vi.fn() +vi.stubGlobal('fetch', fetchMock) + +function lastRequestBody(): any { + const call = fetchMock.mock.calls.at(-1) + if (!call) throw new Error('fetch was not called') + return JSON.parse(call[1].body) +} + +function nthCall(index: number): [any, any] { + const call = fetchMock.mock.calls[index] + if (!call) throw new Error('fetch was not called') + return [call[0], call[1]] +} + +beforeEach(() => { + fetchMock.mockReset() +}) + +afterEach(() => { + vi.unstubAllEnvs() +}) + +describe('Cohere Embedding Adapter', () => { + describe('createCohereEmbedding', () => { + it('creates an adapter with the provided API key', () => { + const adapter = createCohereEmbedding('embed-v4.0', 'test-api-key') + expect(adapter).toBeInstanceOf(CohereEmbeddingAdapter) + expect(adapter.kind).toBe('embedding') + expect(adapter.name).toBe('cohere') + expect(adapter.model).toBe('embed-v4.0') + }) + }) + + describe('cohereEmbedding', () => { + it('reads the API key from COHERE_API_KEY', () => { + vi.stubEnv('COHERE_API_KEY', 'env-api-key') + const adapter = cohereEmbedding('embed-v4.0') + expect(adapter).toBeInstanceOf(CohereEmbeddingAdapter) + expect(adapter.model).toBe('embed-v4.0') + }) + + it('throws when COHERE_API_KEY is not set', () => { + vi.stubEnv('COHERE_API_KEY', '') + expect(() => cohereEmbedding('embed-v4.0')).toThrow('COHERE_API_KEY') + }) + }) + + describe('createEmbeddings', () => { + it('sends texts as a batch with auth headers and maps the response', async () => { + fetchMock.mockResolvedValue( + embedResponse({ + vectors: [ + [0.1, 0.2], + [0.3, 0.4], + ], + inputTokens: 7, + }), + ) + const adapter = createCohereEmbedding('embed-v4.0', 'test-key') + + const result = await adapter.createEmbeddings({ + model: 'embed-v4.0', + input: ['a red guitar', { type: 'text', content: 'a blue drum kit' }], + modelOptions: { inputType: 'search_document' }, + logger: testLogger, + }) + + expect(fetchMock).toHaveBeenCalledTimes(1) + const [url, init] = nthCall(0) + expect(url).toBe(EMBED_URL) + expect(init.method).toBe('POST') + expect(init.headers).toMatchObject({ + Authorization: 'Bearer test-key', + 'Content-Type': 'application/json', + }) + + expect(lastRequestBody()).toEqual({ + model: 'embed-v4.0', + inputs: [ + { content: [{ type: 'text', text: 'a red guitar' }] }, + { content: [{ type: 'text', text: 'a blue drum kit' }] }, + ], + input_type: 'search_document', + embedding_types: ['float'], + }) + + expect(result.embeddings).toEqual([ + { vector: [0.1, 0.2], index: 0 }, + { vector: [0.3, 0.4], index: 1 }, + ]) + expect(result.usage).toEqual({ + promptTokens: 7, + completionTokens: 0, + totalTokens: 7, + }) + expect(result.model).toBe('embed-v4.0') + }) + + it('honors baseUrl and custom headers', async () => { + fetchMock.mockResolvedValue(embedResponse({ vectors: [[0.1]] })) + const adapter = createCohereEmbedding('embed-v4.0', 'test-key', { + baseUrl: 'https://proxy.example.com', + headers: { 'X-Custom': 'yes' }, + }) + + await adapter.createEmbeddings({ + model: 'embed-v4.0', + input: ['hello'], + modelOptions: { inputType: 'search_query' }, + logger: testLogger, + }) + + const [url, init] = nthCall(0) + expect(url).toBe('https://proxy.example.com/v2/embed') + expect(init.headers).toMatchObject({ 'X-Custom': 'yes' }) + }) + + it('maps a data-source image to a data: URI content part', async () => { + fetchMock.mockResolvedValue(embedResponse({ vectors: [[0.5]] })) + const adapter = createCohereEmbedding('embed-v4.0', 'test-key') + + await adapter.createEmbeddings({ + model: 'embed-v4.0', + input: [ + { + type: 'image', + source: { type: 'data', value: 'aGk=', mimeType: 'image/png' }, + }, + ], + modelOptions: { inputType: 'search_document' }, + logger: testLogger, + }) + + expect(lastRequestBody().inputs).toEqual([ + { + content: [ + { + type: 'image_url', + image_url: { url: 'data:image/png;base64,aGk=' }, + }, + ], + }, + ]) + }) + + it('passes url-source data: URIs through unchanged', async () => { + fetchMock.mockResolvedValue(embedResponse({ vectors: [[0.5]] })) + const adapter = createCohereEmbedding('embed-v4.0', 'test-key') + + await adapter.createEmbeddings({ + model: 'embed-v4.0', + input: [ + { + type: 'image', + source: { type: 'url', value: 'data:image/jpeg;base64,aGk=' }, + }, + ], + modelOptions: { inputType: 'search_document' }, + logger: testLogger, + }) + + expect(lastRequestBody().inputs[0].content[0].image_url.url).toBe( + 'data:image/jpeg;base64,aGk=', + ) + }) + + it('sends a fused text+image item as one content array (one vector)', async () => { + fetchMock.mockResolvedValue(embedResponse({ vectors: [[0.9]] })) + const adapter = createCohereEmbedding('embed-v4.0', 'test-key') + + const result = await adapter.createEmbeddings({ + model: 'embed-v4.0', + input: [ + { + type: 'content', + content: [ + { type: 'text', content: 'product photo' }, + { + type: 'image', + source: { type: 'data', value: 'aGk=', mimeType: 'image/png' }, + }, + ], + }, + ], + modelOptions: { inputType: 'search_document' }, + logger: testLogger, + }) + + expect(lastRequestBody().inputs).toEqual([ + { + content: [ + { type: 'text', text: 'product photo' }, + { + type: 'image_url', + image_url: { url: 'data:image/png;base64,aGk=' }, + }, + ], + }, + ]) + expect(result.embeddings).toEqual([{ vector: [0.9], index: 0 }]) + }) + + it('maps inputType, truncate, and dimensions onto the request body', async () => { + fetchMock.mockResolvedValue(embedResponse({ vectors: [[0.1]] })) + const adapter = createCohereEmbedding('embed-v4.0', 'test-key') + + await adapter.createEmbeddings({ + model: 'embed-v4.0', + input: ['hello'], + dimensions: 1024, + modelOptions: { inputType: 'clustering', truncate: 'END' }, + logger: testLogger, + }) + + const body = lastRequestBody() + expect(body.input_type).toBe('clustering') + expect(body.truncate).toBe('END') + expect(body.output_dimension).toBe(1024) + }) + + it('omits truncate and output_dimension when not provided', async () => { + fetchMock.mockResolvedValue(embedResponse({ vectors: [[0.1]] })) + const adapter = createCohereEmbedding('embed-v4.0', 'test-key') + + await adapter.createEmbeddings({ + model: 'embed-v4.0', + input: ['hello'], + modelOptions: { inputType: 'search_query' }, + logger: testLogger, + }) + + const body = lastRequestBody() + expect('truncate' in body).toBe(false) + expect('output_dimension' in body).toBe(false) + }) + + it('always pins embedding_types to ["float"]', async () => { + fetchMock.mockResolvedValue(embedResponse({ vectors: [[0.1]] })) + const adapter = createCohereEmbedding('embed-v4.0', 'test-key') + + await adapter.createEmbeddings({ + model: 'embed-v4.0', + input: ['hello'], + modelOptions: { + inputType: 'search_document', + embeddingTypes: ['float'], + }, + logger: testLogger, + }) + + expect(lastRequestBody().embedding_types).toEqual(['float']) + }) + + it('throws when modelOptions.inputType is missing', async () => { + const adapter = createCohereEmbedding('embed-v4.0', 'test-key') + + await expect( + adapter.createEmbeddings({ + model: 'embed-v4.0', + input: ['hello'], + logger: testLogger, + }), + ).rejects.toThrow('modelOptions.inputType') + expect(fetchMock).not.toHaveBeenCalled() + }) + + it('rejects http(s) image URLs by default', async () => { + const adapter = createCohereEmbedding('embed-v4.0', 'test-key') + + await expect( + adapter.createEmbeddings({ + model: 'embed-v4.0', + input: [ + { + type: 'image', + source: { type: 'url', value: 'https://example.com/cat.png' }, + }, + ], + modelOptions: { inputType: 'search_document' }, + logger: testLogger, + }), + ).rejects.toThrow('Cohere does not fetch remote image URLs') + expect(fetchMock).not.toHaveBeenCalled() + }) + + it('downloads http(s) image URLs when allowUrlFetch is enabled', async () => { + const imageUrl = 'https://example.com/cat.png' + // "Hello" → base64 "SGVsbG8=" + const imageBytes = new Uint8Array([72, 101, 108, 108, 111]) + fetchMock.mockImplementation((url: unknown) => { + if (url === imageUrl) { + return Promise.resolve( + new Response(imageBytes, { + status: 200, + headers: { 'content-type': 'image/png' }, + }), + ) + } + return Promise.resolve(embedResponse({ vectors: [[0.7]] })) + }) + const adapter = createCohereEmbedding('embed-v4.0', 'test-key', { + allowUrlFetch: true, + }) + + const result = await adapter.createEmbeddings({ + model: 'embed-v4.0', + input: [{ type: 'image', source: { type: 'url', value: imageUrl } }], + modelOptions: { inputType: 'search_document' }, + logger: testLogger, + }) + + expect(fetchMock).toHaveBeenCalledTimes(2) + expect(fetchMock).toHaveBeenCalledWith( + imageUrl, + expect.objectContaining({ signal: expect.any(AbortSignal) }), + ) + expect(lastRequestBody().inputs[0].content[0].image_url.url).toBe( + 'data:image/png;base64,SGVsbG8=', + ) + expect(result.embeddings).toEqual([{ vector: [0.7], index: 0 }]) + }) + + it('rejects private/loopback image URLs even when allowUrlFetch is enabled', async () => { + const adapter = createCohereEmbedding('embed-v4.0', 'test-key', { + allowUrlFetch: true, + }) + + for (const value of [ + 'http://127.0.0.1/secret.png', + 'http://localhost/secret.png', + 'http://169.254.169.254/latest/meta-data/', + 'http://192.168.1.1/img.png', + 'http://10.0.0.5/img.png', + 'http://172.16.0.1/img.png', + ]) { + await expect( + adapter.createEmbeddings({ + model: 'embed-v4.0', + input: [{ type: 'image', source: { type: 'url', value } }], + modelOptions: { inputType: 'search_document' }, + logger: testLogger, + }), + ).rejects.toThrow('Refusing to fetch internal or private URL') + } + expect(fetchMock).not.toHaveBeenCalled() + }) + + it('passes an AbortSignal on the Cohere API fetch', async () => { + fetchMock.mockResolvedValue(embedResponse({ vectors: [[0.1]] })) + const adapter = createCohereEmbedding('embed-v4.0', 'test-key') + + await adapter.createEmbeddings({ + model: 'embed-v4.0', + input: ['hello'], + modelOptions: { inputType: 'search_document' }, + logger: testLogger, + }) + + const [, init] = nthCall(0) + expect(init.signal).toBeInstanceOf(AbortSignal) + }) + + it('omits usage when billed_units.input_tokens is absent', async () => { + fetchMock.mockResolvedValue(embedResponse({ vectors: [[0.1]] })) + const adapter = createCohereEmbedding('embed-v4.0', 'test-key') + + const result = await adapter.createEmbeddings({ + model: 'embed-v4.0', + input: ['hello'], + modelOptions: { inputType: 'search_document' }, + logger: testLogger, + }) + + expect(result.usage).toBeUndefined() + }) + + it('throws when the embedding count does not match the input count', async () => { + fetchMock.mockResolvedValue(embedResponse({ vectors: [[0.1]] })) + const adapter = createCohereEmbedding('embed-v4.0', 'test-key') + + await expect( + adapter.createEmbeddings({ + model: 'embed-v4.0', + input: ['one', 'two'], + modelOptions: { inputType: 'search_document' }, + logger: testLogger, + }), + ).rejects.toThrow('returned 1 embeddings for 2 inputs') + }) + + it('throws a clear error with the API message on non-2xx responses', async () => { + fetchMock.mockResolvedValue( + new Response(JSON.stringify({ message: 'invalid api token' }), { + status: 401, + headers: { 'content-type': 'application/json' }, + }), + ) + const adapter = createCohereEmbedding('embed-v4.0', 'bad-key') + + await expect( + adapter.createEmbeddings({ + model: 'embed-v4.0', + input: ['hello'], + modelOptions: { inputType: 'search_document' }, + logger: testLogger, + }), + ).rejects.toThrow('Cohere embed failed (401): invalid api token') + }) + + it('falls back to the raw body text for non-JSON error responses', async () => { + fetchMock.mockResolvedValue(new Response('Bad Gateway', { status: 502 })) + const adapter = createCohereEmbedding('embed-v4.0', 'test-key') + + await expect( + adapter.createEmbeddings({ + model: 'embed-v4.0', + input: ['hello'], + modelOptions: { inputType: 'search_document' }, + logger: testLogger, + }), + ).rejects.toThrow('Cohere embed failed (502): Bad Gateway') + }) + }) +}) diff --git a/packages/ai-cohere/tsconfig.json b/packages/ai-cohere/tsconfig.json new file mode 100644 index 000000000..1434e80de --- /dev/null +++ b/packages/ai-cohere/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "dist" + }, + "include": ["src/**/*.ts", "src/**/*.tsx", "tests/**/*.ts"], + "exclude": ["node_modules", "dist", "**/*.config.ts"] +} diff --git a/packages/ai-cohere/vite.config.ts b/packages/ai-cohere/vite.config.ts new file mode 100644 index 000000000..77bcc2e60 --- /dev/null +++ b/packages/ai-cohere/vite.config.ts @@ -0,0 +1,36 @@ +import { defineConfig, mergeConfig } from 'vitest/config' +import { tanstackViteConfig } from '@tanstack/vite-config' +import packageJson from './package.json' + +const config = defineConfig({ + test: { + name: packageJson.name, + dir: './', + watch: false, + globals: true, + environment: 'node', + include: ['tests/**/*.test.ts'], + coverage: { + provider: 'v8', + reporter: ['text', 'json', 'html', 'lcov'], + exclude: [ + 'node_modules/', + 'dist/', + 'tests/', + '**/*.test.ts', + '**/*.config.ts', + '**/types.ts', + ], + include: ['src/**/*.ts'], + }, + }, +}) + +export default mergeConfig( + config, + tanstackViteConfig({ + entry: ['./src/index.ts'], + srcDir: './src', + cjs: false, + }), +) diff --git a/packages/ai-event-client/src/index.ts b/packages/ai-event-client/src/index.ts index 0bd1fdc12..c84340470 100644 --- a/packages/ai-event-client/src/index.ts +++ b/packages/ai-event-client/src/index.ts @@ -643,6 +643,58 @@ export interface ImageUsageEvent extends BaseEventContext { usage: TokenUsage } +// =========================== +// Embedding Events +// =========================== + +/** Emitted when an embedding request starts. Carries input counts only, never input content. */ +export interface EmbeddingRequestStartedEvent extends BaseEventContext { + requestId: string + provider: string + model: string + /** Total number of items being embedded (after single→array normalization). */ + inputCount: number + /** Count of text-only input items. */ + textInputCount: number + /** Count of input items containing an image part. */ + imageInputCount: number + /** Requested output dimensionality, when specified. */ + dimensions?: number + /** Provider-specific options passed to the embedding request. */ + modelOptions?: Record +} + +/** Emitted when an embedding request completes. */ +export interface EmbeddingRequestCompletedEvent extends BaseEventContext { + requestId: string + provider: string + model: string + embeddingCount: number + /** Dimensionality of the returned vectors. */ + dimensions?: number + duration: number + /** Provider-specific options passed to the embedding request. */ + modelOptions?: Record +} + +/** Emitted when an embedding request fails. */ +export interface EmbeddingRequestErrorEvent extends BaseEventContext { + requestId: string + provider: string + model: string + error: { message: string; name?: string } + duration: number + /** Provider-specific options passed to the embedding request. */ + modelOptions?: Record +} + +/** Emitted when embedding usage metrics are available. */ +export interface EmbeddingUsageEvent extends BaseEventContext { + requestId: string + model: string + usage: TokenUsage +} + // =========================== // Speech Events // =========================== @@ -1020,6 +1072,12 @@ export interface AIDevtoolsEventMap { 'image:request:completed': ImageRequestCompletedEvent 'image:usage': ImageUsageEvent + // Embedding events + 'embedding:request:started': EmbeddingRequestStartedEvent + 'embedding:request:completed': EmbeddingRequestCompletedEvent + 'embedding:request:error': EmbeddingRequestErrorEvent + 'embedding:usage': EmbeddingUsageEvent + // Speech events 'speech:request:started': SpeechRequestStartedEvent 'speech:request:completed': SpeechRequestCompletedEvent diff --git a/packages/ai-gemini/src/adapters/embedding.ts b/packages/ai-gemini/src/adapters/embedding.ts new file mode 100644 index 000000000..9daf6f040 --- /dev/null +++ b/packages/ai-gemini/src/adapters/embedding.ts @@ -0,0 +1,173 @@ +import { BaseEmbeddingAdapter } from '@tanstack/ai/adapters' +import { toRunErrorPayload } from '@tanstack/ai/adapter-internals' +import { requireTextOnlyEmbeddingInput } from '@tanstack/ai' +import { + createGeminiClient, + generateId, + getGeminiApiKeyFromEnv, +} from '../utils' +import type { EmbeddingOptions, EmbeddingResult } from '@tanstack/ai' +import type { EmbedContentConfig, GoogleGenAI } from '@google/genai' +import type { + GeminiEmbeddingModel, + GeminiEmbeddingModelInputModalitiesByName, + GeminiEmbeddingModelProviderOptionsByName, +} from '../model-meta' +import type { GeminiEmbeddingProviderOptions } from '../embedding/embedding-provider-options' +import type { GeminiClientConfig } from '../utils/client' + +/** + * Configuration for Gemini Embedding adapter + */ +export interface GeminiEmbeddingConfig extends GeminiClientConfig {} + +/** + * Gemini Embedding Adapter + * + * Tree-shakeable adapter for Gemini text embeddings. + * Supports gemini-embedding-001. + * + * Features: + * - Batch embedding (one request for the whole input array) + * - Matryoshka dimension reduction via the top-level `dimensions` option + * (mapped to the SDK's `outputDimensionality`) + * - Task-type hints (`taskType`, `title`) via provider options + */ +export class GeminiEmbeddingAdapter< + TModel extends GeminiEmbeddingModel, +> extends BaseEmbeddingAdapter< + TModel, + GeminiEmbeddingProviderOptions, + GeminiEmbeddingModelProviderOptionsByName, + GeminiEmbeddingModelInputModalitiesByName +> { + readonly name = 'gemini' as const + + protected client: GoogleGenAI + + constructor(config: GeminiEmbeddingConfig, model: TModel) { + super(model, config) + this.client = createGeminiClient(config) + } + + async createEmbeddings( + options: EmbeddingOptions, + ): Promise { + const { model, logger, modelOptions } = options + const texts = requireTextOnlyEmbeddingInput(options.input, this.name, model) + + try { + // Build the config incrementally so possibly-undefined values are never + // assigned to optional fields (exactOptionalPropertyTypes). + const config: EmbedContentConfig = {} + if (options.dimensions !== undefined) { + config.outputDimensionality = options.dimensions + } + if (modelOptions?.taskType !== undefined) { + config.taskType = modelOptions.taskType + } + if (modelOptions?.title !== undefined) { + config.title = modelOptions.title + } + + logger.request( + `activity=embed provider=${this.name} model=${model} inputs=${texts.length}`, + { provider: this.name, model }, + ) + + const response = await this.client.models.embedContent({ + model, + contents: texts, + config, + }) + + const embeddings = response.embeddings + if (!embeddings || embeddings.length !== texts.length) { + throw new Error( + `Gemini embedContent returned ${embeddings ? embeddings.length : 'no'} embeddings for ${texts.length} inputs (model: ${model}).`, + ) + } + + // The Gemini embedding API does not report token usage, so `usage` is + // omitted from the result. + return { + id: generateId(this.name), + model, + embeddings: embeddings.map((embedding, index) => ({ + vector: embedding.values ?? [], + index, + })), + } + } catch (error: unknown) { + logger.errors(`${this.name}.createEmbeddings fatal`, { + error: toRunErrorPayload(error, `${this.name}.createEmbeddings failed`), + source: `${this.name}.createEmbeddings`, + }) + throw error + } + } +} + +/** + * Creates a Gemini embedding adapter with explicit API key. + * Type resolution happens here at the call site. + * + * @param model - The model name (e.g., 'gemini-embedding-001') + * @param apiKey - Your Google API key + * @param config - Optional additional configuration + * @returns Configured Gemini embedding adapter instance with resolved types + * + * @example + * ```typescript + * const adapter = createGeminiEmbedding('gemini-embedding-001', "your-api-key"); + * + * const result = await embed({ + * adapter, + * input: 'a red guitar' + * }); + * ``` + */ +export function createGeminiEmbedding( + model: TModel, + apiKey: string, + config?: Omit, +): GeminiEmbeddingAdapter { + // Put apiKey LAST so caller-supplied config can't silently override the + // explicit argument. + return new GeminiEmbeddingAdapter({ ...config, apiKey }, model) +} + +/** + * Creates a Gemini embedding adapter with automatic API key detection from environment variables. + * Type resolution happens here at the call site. + * + * Looks for `GOOGLE_API_KEY` or `GEMINI_API_KEY` in: + * - `process.env` (Node.js) + * - `window.env` (Browser with injected env) + * + * @param model - The model name (e.g., 'gemini-embedding-001') + * @param config - Optional configuration (excluding apiKey which is auto-detected) + * @returns Configured Gemini embedding adapter instance with resolved types + * @throws Error if GOOGLE_API_KEY or GEMINI_API_KEY is not found in environment + * + * @example + * ```typescript + * // Automatically uses GOOGLE_API_KEY from environment + * const adapter = geminiEmbedding('gemini-embedding-001'); + * + * const result = await embed({ + * adapter, + * input: ['a red guitar', 'a blue drum kit'], + * dimensions: 1536 + * }); + * + * console.log(result.embeddings[0].vector) + * ``` + */ +export function geminiEmbedding( + model: TModel, + config?: Omit, +): GeminiEmbeddingAdapter { + const apiKey = getGeminiApiKeyFromEnv() + return createGeminiEmbedding(model, apiKey, config) +} diff --git a/packages/ai-gemini/src/embedding/embedding-provider-options.ts b/packages/ai-gemini/src/embedding/embedding-provider-options.ts new file mode 100644 index 000000000..3fedce32c --- /dev/null +++ b/packages/ai-gemini/src/embedding/embedding-provider-options.ts @@ -0,0 +1,29 @@ +/** + * Provider options for Gemini embedding models. + * + * `dimensions` is deliberately absent: it's a first-class top-level option on + * `embed()` and is mapped to the SDK's `outputDimensionality` by the adapter. + */ +export interface GeminiEmbeddingProviderOptions { + /** + * Type of task for which the embedding will be used. Helps the model + * produce embeddings optimized for the intended use case. + * + * @see https://ai.google.dev/gemini-api/docs/embeddings#task-types + */ + taskType?: + | 'SEMANTIC_SIMILARITY' + | 'CLASSIFICATION' + | 'CLUSTERING' + | 'RETRIEVAL_DOCUMENT' + | 'RETRIEVAL_QUERY' + | 'QUESTION_ANSWERING' + | 'FACT_VERIFICATION' + | 'CODE_RETRIEVAL_QUERY' + + /** + * Title for the text. Only applicable when `taskType` is + * `RETRIEVAL_DOCUMENT`. + */ + title?: string +} diff --git a/packages/ai-gemini/src/index.ts b/packages/ai-gemini/src/index.ts index 245f892cb..e0c58922c 100644 --- a/packages/ai-gemini/src/index.ts +++ b/packages/ai-gemini/src/index.ts @@ -36,6 +36,15 @@ export type { ImagePromptLanguage, } from './image/image-provider-options' +// Embedding adapter - for embedding vectors +export { + GeminiEmbeddingAdapter, + createGeminiEmbedding, + geminiEmbedding, + type GeminiEmbeddingConfig, +} from './adapters/embedding' +export type { GeminiEmbeddingProviderOptions } from './embedding/embedding-provider-options' + // TTS adapter (experimental) /** * @experimental Gemini TTS is an experimental feature and may change. @@ -100,6 +109,7 @@ export { GEMINI_TTS_VOICES as GeminiTTSVoices } from './model-meta' export { GEMINI_AUDIO_MODELS as GeminiAudioModels } from './model-meta' export { GEMINI_VIDEO_MODELS as GeminiVideoModels } from './model-meta' export { GEMINI_INTERACTIONS_VIDEO_MODELS as GeminiInteractionsVideoModels } from './model-meta' +export { GEMINI_EMBEDDING_MODELS } from './model-meta' export type { GeminiModels as GeminiTextModel } from './model-meta' export type { GeminiImageModels as GeminiImageModel } from './model-meta' export type { GeminiTTSVoice } from './model-meta' @@ -112,6 +122,9 @@ export type { GeminiChatModelProviderOptionsByName, GeminiChatModelToolCapabilitiesByName, GeminiModelInputModalitiesByName, + GeminiEmbeddingModel, + GeminiEmbeddingModelProviderOptionsByName, + GeminiEmbeddingModelInputModalitiesByName, } from './model-meta' export type { GeminiStructuredOutputOptions, diff --git a/packages/ai-gemini/src/model-meta.ts b/packages/ai-gemini/src/model-meta.ts index d69af9609..6297eb97a 100644 --- a/packages/ai-gemini/src/model-meta.ts +++ b/packages/ai-gemini/src/model-meta.ts @@ -6,6 +6,7 @@ import type { GeminiThinkingOptions, GeminiToolConfigOptions, } from './text/text-provider-options' +import type { GeminiEmbeddingProviderOptions } from './embedding/embedding-provider-options' interface ModelMeta { name: string @@ -897,6 +898,28 @@ export const GEMINI_INTERACTIONS_VIDEO_MODELS = [ GEMINI_OMNI_FLASH_PREVIEW.name, ] as const +/** + * Embedding models + */ +export const GEMINI_EMBEDDING_MODELS = ['gemini-embedding-001'] as const + +export type GeminiEmbeddingModel = (typeof GEMINI_EMBEDDING_MODELS)[number] + +/** + * Type-only map from embedding model name to its provider options type. + */ +export type GeminiEmbeddingModelProviderOptionsByName = { + 'gemini-embedding-001': GeminiEmbeddingProviderOptions +} + +/** + * Per-model input modalities for embedding models. Gemini embedding models + * are text-only, so image inputs fail at compile time. + */ +export type GeminiEmbeddingModelInputModalitiesByName = { + 'gemini-embedding-001': readonly ['text'] +} + // Manual type map for per-model provider options export type GeminiChatModelProviderOptionsByName = { // Models with thinking and structured output support diff --git a/packages/ai-gemini/tests/embedding-adapter.test.ts b/packages/ai-gemini/tests/embedding-adapter.test.ts new file mode 100644 index 000000000..d8e2c861b --- /dev/null +++ b/packages/ai-gemini/tests/embedding-adapter.test.ts @@ -0,0 +1,185 @@ +import { describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { + GeminiEmbeddingAdapter, + createGeminiEmbedding, +} from '../src/adapters/embedding' +import type { EmbedContentResponse } from '@google/genai' +import type { GeminiEmbeddingModel } from '../src/model-meta' + +const testLogger = resolveDebugOption(false) + +/** + * Test-only subclass exposing the SDK client's `models.embedContent` to + * `vi.spyOn` — same pattern as the OpenAI embedding adapter tests, keeping + * every type real with no casts. + */ +class TestGeminiEmbeddingAdapter< + TModel extends GeminiEmbeddingModel, +> extends GeminiEmbeddingAdapter { + spyOnEmbedContent() { + return vi.spyOn(this.client.models, 'embedContent') + } +} + +function mockResponse(vectors: Array>): EmbedContentResponse { + return { + embeddings: vectors.map((values) => ({ values })), + } +} + +describe('Gemini Embedding Adapter', () => { + describe('createGeminiEmbedding', () => { + it('creates an adapter with the provided API key', () => { + const adapter = createGeminiEmbedding( + 'gemini-embedding-001', + 'test-api-key', + ) + expect(adapter).toBeInstanceOf(GeminiEmbeddingAdapter) + expect(adapter.kind).toBe('embedding') + expect(adapter.name).toBe('gemini') + expect(adapter.model).toBe('gemini-embedding-001') + }) + }) + + describe('createEmbeddings', () => { + it('sends texts as a batch and maps vectors with indices', async () => { + const adapter = new TestGeminiEmbeddingAdapter( + { apiKey: 'test' }, + 'gemini-embedding-001', + ) + const spy = adapter.spyOnEmbedContent() + spy.mockResolvedValue( + mockResponse([ + [0.1, 0.2], + [0.3, 0.4], + ]), + ) + + const result = await adapter.createEmbeddings({ + model: 'gemini-embedding-001', + input: ['a red guitar', { type: 'text', content: 'a blue drum kit' }], + logger: testLogger, + }) + + expect(spy).toHaveBeenCalledWith({ + model: 'gemini-embedding-001', + contents: ['a red guitar', 'a blue drum kit'], + config: {}, + }) + expect(result.embeddings).toEqual([ + { vector: [0.1, 0.2], index: 0 }, + { vector: [0.3, 0.4], index: 1 }, + ]) + expect(result.model).toBe('gemini-embedding-001') + // The Gemini embedding API does not report token usage. + expect(result.usage).toBeUndefined() + }) + + it('maps the top-level dimensions option to outputDimensionality', async () => { + const adapter = new TestGeminiEmbeddingAdapter( + { apiKey: 'test' }, + 'gemini-embedding-001', + ) + const spy = adapter.spyOnEmbedContent() + spy.mockResolvedValue(mockResponse([[0.1]])) + + await adapter.createEmbeddings({ + model: 'gemini-embedding-001', + input: ['hello'], + dimensions: 1536, + logger: testLogger, + }) + + expect(spy).toHaveBeenCalledWith( + expect.objectContaining({ + config: { outputDimensionality: 1536 }, + }), + ) + }) + + it('passes taskType and title provider options through the config', async () => { + const adapter = new TestGeminiEmbeddingAdapter( + { apiKey: 'test' }, + 'gemini-embedding-001', + ) + const spy = adapter.spyOnEmbedContent() + spy.mockResolvedValue(mockResponse([[0.1]])) + + await adapter.createEmbeddings({ + model: 'gemini-embedding-001', + input: ['hello'], + modelOptions: { + taskType: 'RETRIEVAL_DOCUMENT', + title: 'Guitar catalog', + }, + logger: testLogger, + }) + + expect(spy).toHaveBeenCalledWith( + expect.objectContaining({ + config: { + taskType: 'RETRIEVAL_DOCUMENT', + title: 'Guitar catalog', + }, + }), + ) + }) + + it('throws a clear error when the response has no embeddings', async () => { + const adapter = new TestGeminiEmbeddingAdapter( + { apiKey: 'test' }, + 'gemini-embedding-001', + ) + const spy = adapter.spyOnEmbedContent() + spy.mockResolvedValue({}) + + await expect( + adapter.createEmbeddings({ + model: 'gemini-embedding-001', + input: ['hello'], + logger: testLogger, + }), + ).rejects.toThrow('returned no embeddings for 1 inputs') + }) + + it('throws a clear error when embedding count mismatches input count', async () => { + const adapter = new TestGeminiEmbeddingAdapter( + { apiKey: 'test' }, + 'gemini-embedding-001', + ) + const spy = adapter.spyOnEmbedContent() + spy.mockResolvedValue(mockResponse([[0.1]])) + + await expect( + adapter.createEmbeddings({ + model: 'gemini-embedding-001', + input: ['hello', 'world'], + logger: testLogger, + }), + ).rejects.toThrow('returned 1 embeddings for 2 inputs') + }) + + it('throws a clear error for image input', async () => { + const adapter = new TestGeminiEmbeddingAdapter( + { apiKey: 'test' }, + 'gemini-embedding-001', + ) + const spy = adapter.spyOnEmbedContent() + + await expect( + adapter.createEmbeddings({ + model: 'gemini-embedding-001', + input: [ + { + type: 'image', + source: { type: 'data', value: 'aGk=', mimeType: 'image/png' }, + }, + ], + logger: testLogger, + }), + ).rejects.toThrow('only supports text embedding inputs') + expect(spy).not.toHaveBeenCalled() + }) + }) +}) diff --git a/packages/ai-mistral/package.json b/packages/ai-mistral/package.json index 4d47ae495..a8776279a 100644 --- a/packages/ai-mistral/package.json +++ b/packages/ai-mistral/package.json @@ -20,6 +20,10 @@ "./adapters/text": { "types": "./dist/esm/adapters/text.d.ts", "import": "./dist/esm/adapters/text.js" + }, + "./adapters/embedding": { + "types": "./dist/esm/adapters/embedding.d.ts", + "import": "./dist/esm/adapters/embedding.js" } }, "files": [ diff --git a/packages/ai-mistral/src/adapters/embedding.ts b/packages/ai-mistral/src/adapters/embedding.ts new file mode 100644 index 000000000..cbb733a97 --- /dev/null +++ b/packages/ai-mistral/src/adapters/embedding.ts @@ -0,0 +1,177 @@ +import { BaseEmbeddingAdapter } from '@tanstack/ai/adapters' +import { toRunErrorPayload } from '@tanstack/ai/adapter-internals' +import { requireTextOnlyEmbeddingInput } from '@tanstack/ai' +import { + createMistralClient, + generateId, + getMistralApiKeyFromEnv, +} from '../utils/client' +import type { + EmbeddingOptions, + EmbeddingResult, + TokenUsage, +} from '@tanstack/ai' +import type { Mistral } from '@mistralai/mistralai' +import type { EmbeddingRequest } from '@mistralai/mistralai/models/components' +import type { + MistralEmbeddingModel, + MistralEmbeddingModelInputModalitiesByName, + MistralEmbeddingModelProviderOptionsByName, +} from '../model-meta' +import type { MistralEmbeddingProviderOptions } from '../embedding/embedding-provider-options' +import type { MistralClientConfig } from '../utils/client' + +/** + * Configuration for Mistral embedding adapter. + */ +export type MistralEmbeddingConfig = MistralClientConfig + +/** + * Mistral Embedding Adapter + * + * Tree-shakeable adapter for Mistral text embeddings. + * Supports mistral-embed and codestral-embed. + * + * Features: + * - Batch embedding (one request for the whole input array) + * - Dimension reduction for codestral-embed via the top-level `dimensions` + * option (mapped to Mistral's `outputDimension`); mistral-embed has a fixed + * 1024-dimension output and rejects `dimensions`. + */ +export class MistralEmbeddingAdapter< + TModel extends MistralEmbeddingModel, +> extends BaseEmbeddingAdapter< + TModel, + MistralEmbeddingProviderOptions, + MistralEmbeddingModelProviderOptionsByName, + MistralEmbeddingModelInputModalitiesByName +> { + readonly name = 'mistral' as const + + protected client: Mistral + + constructor(config: MistralEmbeddingConfig, model: TModel) { + super(model, {}) + this.client = createMistralClient(config) + } + + async createEmbeddings( + options: EmbeddingOptions, + ): Promise { + const { model, logger, modelOptions } = options + const texts = requireTextOnlyEmbeddingInput(options.input, this.name, model) + + if (options.dimensions !== undefined && model === 'mistral-embed') { + throw new Error( + 'mistral-embed does not support requesting dimensions (output is a fixed 1024-dimension vector). Use codestral-embed for dimension reduction, or omit `dimensions`.', + ) + } + + try { + // Spread modelOptions first so it can never override the validated + // fields (server routes often pass modelOptions through from untyped + // client input). + const request: EmbeddingRequest = { + ...modelOptions, + model, + inputs: texts, + } + if (options.dimensions !== undefined) { + request.outputDimension = options.dimensions + } + + logger.request( + `activity=embed provider=${this.name} model=${model} inputs=${texts.length}`, + { provider: this.name, model }, + ) + + const response = await this.client.embeddings.create(request) + + const usage: TokenUsage = { + promptTokens: response.usage.promptTokens ?? 0, + completionTokens: 0, + totalTokens: response.usage.totalTokens ?? 0, + } + + return { + id: generateId(this.name), + model, + // Mistral returns one entry per input; `index` is optional in the SDK + // types, so fall back to array order when it's absent. + embeddings: response.data.map((item, arrayIndex) => ({ + vector: item.embedding ?? [], + index: item.index ?? arrayIndex, + })), + usage, + } + } catch (error: unknown) { + logger.errors(`${this.name}.createEmbeddings fatal`, { + error: toRunErrorPayload(error, `${this.name}.createEmbeddings failed`), + source: `${this.name}.createEmbeddings`, + }) + throw error + } + } +} + +/** + * Creates a Mistral embedding adapter with explicit API key. + * Type resolution happens here at the call site. + * + * @param model - The model name (e.g., 'mistral-embed') + * @param apiKey - Your Mistral API key + * @param config - Optional additional configuration + * @returns Configured Mistral embedding adapter instance with resolved types + * + * @example + * ```typescript + * const adapter = createMistralEmbedding('mistral-embed', 'api_key'); + * + * const result = await embed({ + * adapter, + * input: 'a red guitar' + * }); + * ``` + */ +export function createMistralEmbedding( + model: TModel, + apiKey: string, + config?: Omit, +): MistralEmbeddingAdapter { + return new MistralEmbeddingAdapter({ apiKey, ...config }, model) +} + +/** + * Creates a Mistral embedding adapter using the `MISTRAL_API_KEY` environment variable. + * Type resolution happens here at the call site. + * + * Looks for `MISTRAL_API_KEY` in: + * - `process.env` (Node.js) + * - `window.env` (Browser with injected env) + * + * @param model - The model name (e.g., 'codestral-embed') + * @param config - Optional configuration (excluding apiKey which is auto-detected) + * @returns Configured Mistral embedding adapter instance with resolved types + * @throws Error if MISTRAL_API_KEY is not found in environment + * + * @example + * ```typescript + * // Automatically uses MISTRAL_API_KEY from environment + * const adapter = mistralEmbedding('codestral-embed'); + * + * const result = await embed({ + * adapter, + * input: ['a red guitar', 'a blue drum kit'], + * dimensions: 256 + * }); + * + * console.log(result.embeddings[0].vector) + * ``` + */ +export function mistralEmbedding( + model: TModel, + config?: Omit, +): MistralEmbeddingAdapter { + const apiKey = getMistralApiKeyFromEnv() + return createMistralEmbedding(model, apiKey, config) +} diff --git a/packages/ai-mistral/src/embedding/embedding-provider-options.ts b/packages/ai-mistral/src/embedding/embedding-provider-options.ts new file mode 100644 index 000000000..0e01809f8 --- /dev/null +++ b/packages/ai-mistral/src/embedding/embedding-provider-options.ts @@ -0,0 +1,32 @@ +/** + * Provider options for Mistral embedding models. + * + * `dimensions` is deliberately absent: it's a first-class top-level option on + * `embed()` and is mapped to Mistral's `outputDimension` request field by the + * adapter (codestral-embed only — mistral-embed has a fixed 1024-dim output). + */ + +/** + * Provider options for `mistral-embed`. + * + * mistral-embed accepts no provider-specific options: its output is a fixed + * 1024-dimension float vector. + */ +export type MistralEmbedProviderOptions = Record + +/** + * Provider options for `codestral-embed`. + */ +export interface CodestralEmbedProviderOptions { + /** + * The data type of the output embedding values. Mirrors the Mistral SDK's + * `EmbeddingDtype` enum. Defaults to `float`. + */ + outputDtype?: 'float' | 'int8' | 'uint8' | 'binary' | 'ubinary' +} + +/** + * Widest provider options shape accepted by the Mistral embedding adapter. + * Per-model narrowing happens via `MistralEmbeddingModelProviderOptionsByName`. + */ +export type MistralEmbeddingProviderOptions = CodestralEmbedProviderOptions diff --git a/packages/ai-mistral/src/index.ts b/packages/ai-mistral/src/index.ts index d57f85f81..4e9b675ce 100644 --- a/packages/ai-mistral/src/index.ts +++ b/packages/ai-mistral/src/index.ts @@ -14,6 +14,19 @@ export { type MistralTextProviderOptions, } from './adapters/text' +// Embedding adapter - for embedding vectors +export { + MistralEmbeddingAdapter, + createMistralEmbedding, + mistralEmbedding, + type MistralEmbeddingConfig, +} from './adapters/embedding' +export type { + MistralEmbeddingProviderOptions, + MistralEmbedProviderOptions, + CodestralEmbedProviderOptions, +} from './embedding/embedding-provider-options' + // Types export type { MistralChatModelProviderOptionsByName, @@ -21,8 +34,11 @@ export type { ResolveProviderOptions, ResolveInputModalities, MistralChatModels, + MistralEmbeddingModel, + MistralEmbeddingModelProviderOptionsByName, + MistralEmbeddingModelInputModalitiesByName, } from './model-meta' -export { MISTRAL_CHAT_MODELS } from './model-meta' +export { MISTRAL_CHAT_MODELS, MISTRAL_EMBEDDING_MODELS } from './model-meta' export type { MistralTextMetadata, MistralImageMetadata, diff --git a/packages/ai-mistral/src/model-meta.ts b/packages/ai-mistral/src/model-meta.ts index 45ba0a56e..6a5304b0b 100644 --- a/packages/ai-mistral/src/model-meta.ts +++ b/packages/ai-mistral/src/model-meta.ts @@ -1,4 +1,8 @@ import type { MistralTextProviderOptions } from './text/text-provider-options' +import type { + CodestralEmbedProviderOptions, + MistralEmbedProviderOptions, +} from './embedding/embedding-provider-options' /** Provider options for vision-capable Mistral models (pixtral-*). */ export type MistralVisionProviderOptions = MistralTextProviderOptions @@ -268,6 +272,39 @@ export type MistralChatModelProviderOptionsByName = { [OPEN_MISTRAL_NEMO.name]: MistralTextProviderOptions } +/** + * Embedding models (based on endpoints: "embeddings") + */ +export const MISTRAL_EMBEDDING_MODELS = [ + 'mistral-embed', + 'codestral-embed', +] as const + +/** + * Union type of all supported Mistral embedding model names. + */ +export type MistralEmbeddingModel = (typeof MISTRAL_EMBEDDING_MODELS)[number] + +/** + * Type-only map from embedding model name to its provider options type. + * + * `mistral-embed` accepts no provider options (fixed 1024-dim output); + * `codestral-embed` additionally supports `outputDtype`. + */ +export type MistralEmbeddingModelProviderOptionsByName = { + 'mistral-embed': MistralEmbedProviderOptions + 'codestral-embed': CodestralEmbedProviderOptions +} + +/** + * Per-model input modalities for embedding models. Mistral embedding models + * are text-only, so image inputs fail at compile time. + */ +export type MistralEmbeddingModelInputModalitiesByName = { + 'mistral-embed': readonly ['text'] + 'codestral-embed': readonly ['text'] +} + /** * Resolves the provider options type for a specific Mistral model. */ diff --git a/packages/ai-mistral/tests/embedding-adapter.test.ts b/packages/ai-mistral/tests/embedding-adapter.test.ts new file mode 100644 index 000000000..b902711d9 --- /dev/null +++ b/packages/ai-mistral/tests/embedding-adapter.test.ts @@ -0,0 +1,228 @@ +import { describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { + MistralEmbeddingAdapter, + createMistralEmbedding, +} from '../src/adapters/embedding' +import type { EmbeddingResponse } from '@mistralai/mistralai/models/components' +import type { MistralEmbeddingModel } from '../src/model-meta' + +const testLogger = resolveDebugOption(false) + +/** + * Test-only subclass exposing the SDK client's `embeddings.create` to + * `vi.spyOn` — same pattern as the OpenAI embedding adapter tests, keeping + * every type real with no casts. + */ +class TestMistralEmbeddingAdapter< + TModel extends MistralEmbeddingModel, +> extends MistralEmbeddingAdapter { + spyOnEmbeddingsCreate() { + return vi.spyOn(this.client.embeddings, 'create') + } +} + +function mockResponse(vectors: Array>): EmbeddingResponse { + return { + id: 'embd-test-123', + object: 'list', + model: 'mistral-embed', + data: vectors.map((embedding, index) => ({ + object: 'embedding', + embedding, + index, + })), + usage: { promptTokens: 7, completionTokens: 0, totalTokens: 7 }, + } +} + +describe('Mistral Embedding Adapter', () => { + describe('createMistralEmbedding', () => { + it('creates an adapter with the provided API key', () => { + const adapter = createMistralEmbedding('mistral-embed', 'test-api-key') + expect(adapter).toBeInstanceOf(MistralEmbeddingAdapter) + expect(adapter.kind).toBe('embedding') + expect(adapter.name).toBe('mistral') + expect(adapter.model).toBe('mistral-embed') + }) + }) + + describe('createEmbeddings', () => { + it('sends texts as a batch and maps vectors and usage', async () => { + const adapter = new TestMistralEmbeddingAdapter( + { apiKey: 'test' }, + 'mistral-embed', + ) + const spy = adapter.spyOnEmbeddingsCreate() + spy.mockResolvedValue( + mockResponse([ + [0.1, 0.2], + [0.3, 0.4], + ]), + ) + + const result = await adapter.createEmbeddings({ + model: 'mistral-embed', + input: ['a red guitar', { type: 'text', content: 'a blue drum kit' }], + logger: testLogger, + }) + + expect(spy).toHaveBeenCalledWith({ + model: 'mistral-embed', + inputs: ['a red guitar', 'a blue drum kit'], + }) + expect(result.embeddings).toEqual([ + { vector: [0.1, 0.2], index: 0 }, + { vector: [0.3, 0.4], index: 1 }, + ]) + expect(result.usage).toEqual({ + promptTokens: 7, + completionTokens: 0, + totalTokens: 7, + }) + expect(result.model).toBe('mistral-embed') + }) + + it('maps the top-level dimensions option to outputDimension for codestral-embed', async () => { + const adapter = new TestMistralEmbeddingAdapter( + { apiKey: 'test' }, + 'codestral-embed', + ) + const spy = adapter.spyOnEmbeddingsCreate() + spy.mockResolvedValue(mockResponse([[0.1]])) + + await adapter.createEmbeddings({ + model: 'codestral-embed', + input: ['hello'], + dimensions: 256, + logger: testLogger, + }) + + expect(spy).toHaveBeenCalledWith( + expect.objectContaining({ outputDimension: 256 }), + ) + }) + + it('throws a clear error when dimensions is set for mistral-embed', async () => { + const adapter = new TestMistralEmbeddingAdapter( + { apiKey: 'test' }, + 'mistral-embed', + ) + const spy = adapter.spyOnEmbeddingsCreate() + + await expect( + adapter.createEmbeddings({ + model: 'mistral-embed', + input: ['hello'], + dimensions: 512, + logger: testLogger, + }), + ).rejects.toThrow('mistral-embed does not support requesting dimensions') + expect(spy).not.toHaveBeenCalled() + }) + + it('passes provider options through without letting them override model/inputs', async () => { + const adapter = new TestMistralEmbeddingAdapter( + { apiKey: 'test' }, + 'codestral-embed', + ) + const spy = adapter.spyOnEmbeddingsCreate() + spy.mockResolvedValue(mockResponse([[0.1]])) + + await adapter.createEmbeddings({ + model: 'codestral-embed', + input: ['hello'], + modelOptions: { outputDtype: 'int8' }, + logger: testLogger, + }) + + expect(spy).toHaveBeenCalledWith( + expect.objectContaining({ + outputDtype: 'int8', + model: 'codestral-embed', + inputs: ['hello'], + }), + ) + }) + + it('uses the response index when present and array order when absent', async () => { + const adapter = new TestMistralEmbeddingAdapter( + { apiKey: 'test' }, + 'mistral-embed', + ) + const spy = adapter.spyOnEmbeddingsCreate() + spy.mockResolvedValue({ + id: 'embd-test-456', + object: 'list', + model: 'mistral-embed', + data: [ + { object: 'embedding', embedding: [0.3, 0.4], index: 1 }, + { object: 'embedding', embedding: [0.1, 0.2], index: 0 }, + { object: 'embedding', embedding: [0.5, 0.6] }, + ], + usage: { promptTokens: 3, completionTokens: 0, totalTokens: 3 }, + }) + + const result = await adapter.createEmbeddings({ + model: 'mistral-embed', + input: ['one', 'two', 'three'], + logger: testLogger, + }) + + expect(result.embeddings).toEqual([ + { vector: [0.3, 0.4], index: 1 }, + { vector: [0.1, 0.2], index: 0 }, + { vector: [0.5, 0.6], index: 2 }, + ]) + }) + + it('defaults usage token counts to zero when the API omits them', async () => { + const adapter = new TestMistralEmbeddingAdapter( + { apiKey: 'test' }, + 'mistral-embed', + ) + const spy = adapter.spyOnEmbeddingsCreate() + spy.mockResolvedValue({ + id: 'embd-test-789', + object: 'list', + model: 'mistral-embed', + data: [{ object: 'embedding', embedding: [0.1], index: 0 }], + usage: {}, + }) + + const result = await adapter.createEmbeddings({ + model: 'mistral-embed', + input: ['hello'], + logger: testLogger, + }) + + expect(result.usage).toEqual({ + promptTokens: 0, + completionTokens: 0, + totalTokens: 0, + }) + }) + + it('throws a clear error for image input', async () => { + const adapter = new TestMistralEmbeddingAdapter( + { apiKey: 'test' }, + 'mistral-embed', + ) + const spy = adapter.spyOnEmbeddingsCreate() + + await expect( + adapter.createEmbeddings({ + model: 'mistral-embed', + input: [ + { + type: 'image', + source: { type: 'data', value: 'aGk=', mimeType: 'image/png' }, + }, + ], + logger: testLogger, + }), + ).rejects.toThrow('only supports text embedding inputs') + expect(spy).not.toHaveBeenCalled() + }) + }) +}) diff --git a/packages/ai-ollama/src/adapters/embedding.ts b/packages/ai-ollama/src/adapters/embedding.ts new file mode 100644 index 000000000..4b56e2e04 --- /dev/null +++ b/packages/ai-ollama/src/adapters/embedding.ts @@ -0,0 +1,187 @@ +import { requireTextOnlyEmbeddingInput } from '@tanstack/ai' +import { toRunErrorPayload } from '@tanstack/ai/adapter-internals' +import { BaseEmbeddingAdapter } from '@tanstack/ai/adapters' +import { generateId } from '@tanstack/ai-utils' +import { createOllamaClient, getOllamaHostFromEnv } from '../utils' +import type { OllamaClientConfig } from '../utils/client' +import type { EmbedRequest, EmbedResponse, Ollama } from 'ollama' +import type { + EmbeddingOptions, + EmbeddingResult, + TokenUsage, +} from '@tanstack/ai' +import type { OllamaEmbeddingModel } from '../model-meta' +import type { OllamaEmbeddingProviderOptions } from '../embedding/embedding-provider-options' + +/** + * Configuration for the Ollama Embedding adapter. + * Ollama has no API key — only an optional host (and headers). + */ +export interface OllamaEmbeddingConfig extends OllamaClientConfig {} + +/** + * Extract `prompt_eval_count` through an optional-field view of the response. + * The Ollama SDK types it as required, but servers can omit it at runtime; + * widening here keeps the presence check honest without any casts. + */ +function extractPromptEvalCount(response: { + prompt_eval_count?: number +}): number | undefined { + return response.prompt_eval_count +} + +/** + * Ollama Embedding Adapter + * + * Tree-shakeable adapter for Ollama text embeddings (`/api/embed`). + * + * Notes: + * - Batch embedding: one request for the whole input array. + * - Ollama models are loaded dynamically, so any model name string is + * accepted; `OLLAMA_EMBEDDING_MODELS` lists common embedding models. + * - Ollama does not support requesting embedding dimensions, so the + * top-level `dimensions` option is rejected. + */ +export class OllamaEmbeddingAdapter< + TModel extends OllamaEmbeddingModel, +> extends BaseEmbeddingAdapter { + readonly name = 'ollama' as const + + protected client: Ollama + + constructor( + hostOrClientOrConfig: string | Ollama | OllamaEmbeddingConfig | undefined, + model: TModel, + ) { + super(model, {}) + if ( + typeof hostOrClientOrConfig === 'string' || + hostOrClientOrConfig === undefined + ) { + this.client = createOllamaClient({ host: hostOrClientOrConfig }) + } else if ('embed' in hostOrClientOrConfig) { + // Ollama client instance (has an embed method) + this.client = hostOrClientOrConfig + } else { + // OllamaEmbeddingConfig object + this.client = createOllamaClient(hostOrClientOrConfig) + } + } + + async createEmbeddings( + options: EmbeddingOptions, + ): Promise { + const { model, logger, modelOptions } = options + const texts = requireTextOnlyEmbeddingInput(options.input, this.name, model) + + if (options.dimensions !== undefined) { + throw new Error('Ollama does not support requesting embedding dimensions') + } + + try { + // Built incrementally so optional keys are omitted entirely when unset + // (exactOptionalPropertyTypes). Provider options use camelCase + // (`keepAlive`) and are mapped to the SDK's snake_case wire fields. + const request: EmbedRequest = { + model, + input: texts, + } + if (modelOptions?.truncate !== undefined) { + request.truncate = modelOptions.truncate + } + if (modelOptions?.keepAlive !== undefined) { + request.keep_alive = modelOptions.keepAlive + } + if (modelOptions?.options !== undefined) { + // Spread a fresh object to avoid aliasing the caller's options. + request.options = { ...modelOptions.options } + } + + logger.request( + `activity=embed provider=${this.name} model=${model} inputs=${texts.length}`, + { provider: this.name, model }, + ) + + const response: EmbedResponse = await this.client.embed(request) + + // Include usage only when Ollama actually reported a prompt token count. + const promptEvalCount = extractPromptEvalCount(response) + const usage: TokenUsage | undefined = + promptEvalCount !== undefined + ? { + promptTokens: promptEvalCount, + completionTokens: 0, + totalTokens: promptEvalCount, + } + : undefined + + return { + id: generateId(this.name), + model, + embeddings: response.embeddings.map((vector, index) => ({ + vector, + index, + })), + ...(usage !== undefined && { usage }), + } + } catch (error: unknown) { + logger.errors(`${this.name}.createEmbeddings fatal`, { + error: toRunErrorPayload(error, `${this.name}.createEmbeddings failed`), + source: `${this.name}.createEmbeddings`, + }) + throw error + } + } +} + +/** + * Creates an Ollama embedding adapter with explicit host (or client config). + * Type resolution happens here at the call site. + * + * @param model - The model name (e.g., 'nomic-embed-text') + * @param hostOrConfig - Ollama host URL or client config (defaults to http://localhost:11434) + * @returns Configured Ollama embedding adapter instance with resolved types + * + * @example + * ```typescript + * const adapter = createOllamaEmbedding('nomic-embed-text', 'http://localhost:11434'); + * + * const result = await embed({ + * adapter, + * input: 'a red guitar' + * }); + * ``` + */ +export function createOllamaEmbedding( + model: TModel, + hostOrConfig?: string | OllamaEmbeddingConfig, +): OllamaEmbeddingAdapter { + return new OllamaEmbeddingAdapter(hostOrConfig, model) +} + +/** + * Creates an Ollama embedding adapter with host from the `OLLAMA_HOST` + * environment variable (falling back to the Ollama default). + * Type resolution happens here at the call site. + * + * @param model - The model name (e.g., 'nomic-embed-text') + * @returns Configured Ollama embedding adapter instance with resolved types + * + * @example + * ```typescript + * const adapter = ollamaEmbedding('mxbai-embed-large'); + * + * const result = await embed({ + * adapter, + * input: ['a red guitar', 'a blue drum kit'] + * }); + * + * console.log(result.embeddings[0].vector) + * ``` + */ +export function ollamaEmbedding( + model: TModel, +): OllamaEmbeddingAdapter { + const host = getOllamaHostFromEnv() + return new OllamaEmbeddingAdapter(host, model) +} diff --git a/packages/ai-ollama/src/embedding/embedding-provider-options.ts b/packages/ai-ollama/src/embedding/embedding-provider-options.ts new file mode 100644 index 000000000..2d4377be5 --- /dev/null +++ b/packages/ai-ollama/src/embedding/embedding-provider-options.ts @@ -0,0 +1,28 @@ +import type { Options } from 'ollama' + +/** + * Provider options for Ollama embedding models. + * + * `dimensions` is deliberately absent: Ollama does not support requesting + * embedding dimensions, so the adapter rejects the top-level `dimensions` + * option at runtime. + */ +export interface OllamaEmbeddingProviderOptions { + /** + * Truncates the end of each input to fit within the model's context length. + * When `false`, the request errors if an input exceeds the context length. + * Sent as `truncate` on the wire. + */ + truncate?: boolean + /** + * How long to keep the model loaded in memory after the request + * (e.g. `'5m'`, or a number of seconds). Sent as snake_case `keep_alive` + * on the wire, matching the Ollama SDK's `EmbedRequest`. + */ + keepAlive?: string | number + /** + * Ollama runner/sampling options (num_gpu, num_thread, use_mmap, ...) + * forwarded as `EmbedRequest.options`. + */ + options?: Partial +} diff --git a/packages/ai-ollama/src/index.ts b/packages/ai-ollama/src/index.ts index 0cc920987..bbaf0cdd1 100644 --- a/packages/ai-ollama/src/index.ts +++ b/packages/ai-ollama/src/index.ts @@ -21,6 +21,19 @@ export { } from './adapters/summarize' export { OLLAMA_TEXT_MODELS as OllamaSummarizeModels } from './model-meta' +// Embedding adapter - for embedding vectors +export { + OllamaEmbeddingAdapter, + createOllamaEmbedding, + ollamaEmbedding, + type OllamaEmbeddingConfig, +} from './adapters/embedding' +export type { OllamaEmbeddingProviderOptions } from './embedding/embedding-provider-options' +export { + OLLAMA_EMBEDDING_MODELS, + type OllamaEmbeddingModel, +} from './model-meta' + // Tool converters export { convertFunctionToolToAdapterFormat, diff --git a/packages/ai-ollama/src/model-meta.ts b/packages/ai-ollama/src/model-meta.ts index 6ab88da79..e080b06a2 100644 --- a/packages/ai-ollama/src/model-meta.ts +++ b/packages/ai-ollama/src/model-meta.ts @@ -528,3 +528,23 @@ export type OllamaModelInputModalitiesByName = SmollmModelInputModalitiesByName & TinyllamaModelInputModalitiesByName & Tulu3ModelInputModalitiesByName + +/** + * Common Ollama embedding models. + * + * Note: Ollama models are loaded dynamically, so this is a known subset — + * `OllamaEmbeddingModel` also accepts any string via the `(string & {})` + * escape hatch. + */ +export const OLLAMA_EMBEDDING_MODELS = [ + 'nomic-embed-text', + 'mxbai-embed-large', + 'all-minilm', + 'snowflake-arctic-embed', + 'bge-m3', + 'embeddinggemma', +] as const + +export type OllamaEmbeddingModel = + | (typeof OLLAMA_EMBEDDING_MODELS)[number] + | (string & {}) diff --git a/packages/ai-ollama/tests/embedding-adapter.test.ts b/packages/ai-ollama/tests/embedding-adapter.test.ts new file mode 100644 index 000000000..20176e49d --- /dev/null +++ b/packages/ai-ollama/tests/embedding-adapter.test.ts @@ -0,0 +1,221 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { + OllamaEmbeddingAdapter, + createOllamaEmbedding, + ollamaEmbedding, +} from '../src/adapters/embedding' +import type { Mock } from 'vitest' + +const testLogger = resolveDebugOption(false) + +// eslint-disable-next-line @typescript-eslint/no-explicit-any +let embedMock: Mock<(...args: Array) => any> +let ollamaConstructorCalls: Array<{ host?: string } | undefined> + +vi.mock('ollama', () => { + class Ollama { + embed: (...args: Array) => unknown + constructor(config?: { host?: string }) { + ollamaConstructorCalls.push(config) + this.embed = (...args) => embedMock(...args) + } + } + return { Ollama } +}) + +function mockResponse(vectors: Array>, promptEvalCount?: number) { + return { + model: 'nomic-embed-text', + embeddings: vectors, + total_duration: 1000, + load_duration: 100, + ...(promptEvalCount !== undefined && { + prompt_eval_count: promptEvalCount, + }), + } +} + +beforeEach(() => { + embedMock = vi.fn() + ollamaConstructorCalls = [] +}) + +afterEach(() => { + vi.unstubAllEnvs() +}) + +describe('OllamaEmbeddingAdapter construction', () => { + it('createOllamaEmbedding wires kind=embedding, name=ollama, and the given model', () => { + const adapter = createOllamaEmbedding('nomic-embed-text') + expect(adapter).toBeInstanceOf(OllamaEmbeddingAdapter) + expect(adapter.kind).toBe('embedding') + expect(adapter.name).toBe('ollama') + expect(adapter.model).toBe('nomic-embed-text') + }) + + it('createOllamaEmbedding accepts a string host', () => { + const adapter = createOllamaEmbedding( + 'nomic-embed-text', + 'http://remote:11434', + ) + expect(adapter).toBeInstanceOf(OllamaEmbeddingAdapter) + expect(ollamaConstructorCalls).toContainEqual( + expect.objectContaining({ host: 'http://remote:11434' }), + ) + }) + + it('createOllamaEmbedding accepts a config object', () => { + const adapter = createOllamaEmbedding('mxbai-embed-large', { + host: 'http://remote:11434', + headers: { Authorization: 'Bearer x' }, + }) + expect(adapter).toBeInstanceOf(OllamaEmbeddingAdapter) + expect(adapter.model).toBe('mxbai-embed-large') + }) + + it('ollamaEmbedding reads OLLAMA_HOST from env and forwards it to the Ollama client', () => { + vi.stubEnv('OLLAMA_HOST', 'http://from-env:11434') + const adapter = ollamaEmbedding('nomic-embed-text') + expect(adapter.model).toBe('nomic-embed-text') + expect(ollamaConstructorCalls).toContainEqual( + expect.objectContaining({ host: 'http://from-env:11434' }), + ) + }) +}) + +describe('OllamaEmbeddingAdapter.createEmbeddings', () => { + it('sends texts as a single batch and maps vectors with indices', async () => { + embedMock.mockResolvedValueOnce( + mockResponse( + [ + [0.1, 0.2], + [0.3, 0.4], + ], + 7, + ), + ) + + const adapter = createOllamaEmbedding('nomic-embed-text') + const result = await adapter.createEmbeddings({ + model: 'nomic-embed-text', + input: ['a red guitar', { type: 'text', content: 'a blue drum kit' }], + logger: testLogger, + }) + + expect(embedMock).toHaveBeenCalledTimes(1) + expect(embedMock).toHaveBeenCalledWith({ + model: 'nomic-embed-text', + input: ['a red guitar', 'a blue drum kit'], + }) + expect(result.model).toBe('nomic-embed-text') + expect(result.id).toContain('ollama') + expect(result.embeddings).toEqual([ + { vector: [0.1, 0.2], index: 0 }, + { vector: [0.3, 0.4], index: 1 }, + ]) + }) + + it('maps provider options to the SDK request (keepAlive -> keep_alive)', async () => { + embedMock.mockResolvedValueOnce(mockResponse([[0.1]], 3)) + + const adapter = createOllamaEmbedding('nomic-embed-text') + await adapter.createEmbeddings({ + model: 'nomic-embed-text', + input: ['hello'], + modelOptions: { + truncate: true, + keepAlive: '5m', + options: { num_gpu: 1 }, + }, + logger: testLogger, + }) + + expect(embedMock).toHaveBeenCalledWith({ + model: 'nomic-embed-text', + input: ['hello'], + truncate: true, + keep_alive: '5m', + options: { num_gpu: 1 }, + }) + }) + + it('omits unset provider options from the request entirely', async () => { + embedMock.mockResolvedValueOnce(mockResponse([[0.1]], 3)) + + const adapter = createOllamaEmbedding('nomic-embed-text') + await adapter.createEmbeddings({ + model: 'nomic-embed-text', + input: ['hello'], + logger: testLogger, + }) + + const request = embedMock.mock.calls[0]?.[0] as Record + expect(request).not.toHaveProperty('truncate') + expect(request).not.toHaveProperty('keep_alive') + expect(request).not.toHaveProperty('options') + }) + + it('throws a clear error when dimensions is requested', async () => { + const adapter = createOllamaEmbedding('nomic-embed-text') + + await expect( + adapter.createEmbeddings({ + model: 'nomic-embed-text', + input: ['hello'], + dimensions: 256, + logger: testLogger, + }), + ).rejects.toThrow('Ollama does not support requesting embedding dimensions') + expect(embedMock).not.toHaveBeenCalled() + }) + + it('includes usage when prompt_eval_count is present', async () => { + embedMock.mockResolvedValueOnce(mockResponse([[0.1]], 9)) + + const adapter = createOllamaEmbedding('nomic-embed-text') + const result = await adapter.createEmbeddings({ + model: 'nomic-embed-text', + input: ['hello'], + logger: testLogger, + }) + + expect(result.usage).toEqual({ + promptTokens: 9, + completionTokens: 0, + totalTokens: 9, + }) + }) + + it('omits usage when prompt_eval_count is absent', async () => { + embedMock.mockResolvedValueOnce(mockResponse([[0.1]])) + + const adapter = createOllamaEmbedding('nomic-embed-text') + const result = await adapter.createEmbeddings({ + model: 'nomic-embed-text', + input: ['hello'], + logger: testLogger, + }) + + expect(result.usage).toBeUndefined() + expect(result).not.toHaveProperty('usage') + }) + + it('throws a clear error for image input without calling the client', async () => { + const adapter = createOllamaEmbedding('nomic-embed-text') + + await expect( + adapter.createEmbeddings({ + model: 'nomic-embed-text', + input: [ + { + type: 'image', + source: { type: 'data', value: 'aGk=', mimeType: 'image/png' }, + }, + ], + logger: testLogger, + }), + ).rejects.toThrow('only supports text embedding inputs') + expect(embedMock).not.toHaveBeenCalled() + }) +}) diff --git a/packages/ai-openai/src/adapters/embedding.ts b/packages/ai-openai/src/adapters/embedding.ts new file mode 100644 index 000000000..7072d1ce7 --- /dev/null +++ b/packages/ai-openai/src/adapters/embedding.ts @@ -0,0 +1,166 @@ +import OpenAI from 'openai' +import { BaseEmbeddingAdapter } from '@tanstack/ai/adapters' +import { toRunErrorPayload } from '@tanstack/ai/adapter-internals' +import { generateId } from '@tanstack/ai-utils' +import { requireTextOnlyEmbeddingInput } from '@tanstack/ai' +import { getOpenAIApiKeyFromEnv } from '../utils/client' +import type { + EmbeddingOptions, + EmbeddingResult, + TokenUsage, +} from '@tanstack/ai' +import type OpenAI_SDK from 'openai' +import type { + OpenAIEmbeddingModel, + OpenAIEmbeddingModelInputModalitiesByName, + OpenAIEmbeddingModelProviderOptionsByName, +} from '../model-meta' +import type { OpenAIEmbeddingProviderOptions } from '../embedding/embedding-provider-options' +import type { OpenAIClientConfig } from '../utils/client' + +/** + * Configuration for OpenAI Embedding adapter + */ +export interface OpenAIEmbeddingConfig extends OpenAIClientConfig {} + +/** + * OpenAI Embedding Adapter + * + * Tree-shakeable adapter for OpenAI text embeddings. + * Supports text-embedding-3-small and text-embedding-3-large. + * + * Features: + * - Batch embedding (one request for the whole input array) + * - Matryoshka dimension reduction via the top-level `dimensions` option + */ +export class OpenAIEmbeddingAdapter< + TModel extends OpenAIEmbeddingModel, +> extends BaseEmbeddingAdapter< + TModel, + OpenAIEmbeddingProviderOptions, + OpenAIEmbeddingModelProviderOptionsByName, + OpenAIEmbeddingModelInputModalitiesByName +> { + readonly name = 'openai' as const + + protected client: OpenAI + + constructor(config: OpenAIEmbeddingConfig, model: TModel) { + super(model, {}) + this.client = new OpenAI(config) + } + + async createEmbeddings( + options: EmbeddingOptions, + ): Promise { + const { model, logger, modelOptions } = options + const texts = requireTextOnlyEmbeddingInput(options.input, this.name, model) + + try { + // Spread modelOptions first so it can never override the validated + // fields (server routes often pass modelOptions through from untyped + // client input). encoding_format is pinned to float so vectors are + // always number[]. + const request: OpenAI_SDK.EmbeddingCreateParams = { + ...modelOptions, + model, + input: texts, + encoding_format: 'float', + } + if (options.dimensions !== undefined) { + request.dimensions = options.dimensions + } + + logger.request( + `activity=embed provider=${this.name} model=${model} inputs=${texts.length}`, + { provider: this.name, model }, + ) + + const response = await this.client.embeddings.create(request) + + const usage: TokenUsage = { + promptTokens: response.usage.prompt_tokens, + completionTokens: 0, + totalTokens: response.usage.total_tokens, + } + + return { + id: generateId(this.name), + model, + embeddings: response.data.map((item) => ({ + vector: item.embedding, + index: item.index, + })), + usage, + } + } catch (error: unknown) { + logger.errors(`${this.name}.createEmbeddings fatal`, { + error: toRunErrorPayload(error, `${this.name}.createEmbeddings failed`), + source: `${this.name}.createEmbeddings`, + }) + throw error + } + } +} + +/** + * Creates an OpenAI embedding adapter with explicit API key. + * Type resolution happens here at the call site. + * + * @param model - The model name (e.g., 'text-embedding-3-small') + * @param apiKey - Your OpenAI API key + * @param config - Optional additional configuration + * @returns Configured OpenAI embedding adapter instance with resolved types + * + * @example + * ```typescript + * const adapter = createOpenaiEmbedding('text-embedding-3-small', "sk-..."); + * + * const result = await embed({ + * adapter, + * input: 'a red guitar' + * }); + * ``` + */ +export function createOpenaiEmbedding( + model: TModel, + apiKey: string, + config?: Omit, +): OpenAIEmbeddingAdapter { + return new OpenAIEmbeddingAdapter({ apiKey, ...config }, model) +} + +/** + * Creates an OpenAI embedding adapter with automatic API key detection from environment variables. + * Type resolution happens here at the call site. + * + * Looks for `OPENAI_API_KEY` in: + * - `process.env` (Node.js) + * - `window.env` (Browser with injected env) + * + * @param model - The model name (e.g., 'text-embedding-3-small') + * @param config - Optional configuration (excluding apiKey which is auto-detected) + * @returns Configured OpenAI embedding adapter instance with resolved types + * @throws Error if OPENAI_API_KEY is not found in environment + * + * @example + * ```typescript + * // Automatically uses OPENAI_API_KEY from environment + * const adapter = openaiEmbedding('text-embedding-3-large'); + * + * const result = await embed({ + * adapter, + * input: ['a red guitar', 'a blue drum kit'], + * dimensions: 1024 + * }); + * + * console.log(result.embeddings[0].vector) + * ``` + */ +export function openaiEmbedding( + model: TModel, + config?: Omit, +): OpenAIEmbeddingAdapter { + const apiKey = getOpenAIApiKeyFromEnv() + return createOpenaiEmbedding(model, apiKey, config) +} diff --git a/packages/ai-openai/src/embedding/embedding-provider-options.ts b/packages/ai-openai/src/embedding/embedding-provider-options.ts new file mode 100644 index 000000000..5e9673f2c --- /dev/null +++ b/packages/ai-openai/src/embedding/embedding-provider-options.ts @@ -0,0 +1,14 @@ +/** + * Provider options for OpenAI embedding models. + * + * `dimensions` is deliberately absent: it's a first-class top-level option on + * `embed()`. `encoding_format` is pinned to `float` by the adapter so vectors + * are always `number[]`. + */ +export interface OpenAIEmbeddingProviderOptions { + /** + * A unique identifier representing your end-user, which can help OpenAI + * monitor and detect abuse. + */ + user?: string +} diff --git a/packages/ai-openai/src/index.ts b/packages/ai-openai/src/index.ts index fdeb79d57..4a4534b9b 100644 --- a/packages/ai-openai/src/index.ts +++ b/packages/ai-openai/src/index.ts @@ -78,6 +78,15 @@ export { } from './adapters/transcription' export type { OpenAITranscriptionProviderOptions } from './audio/transcription-provider-options' +// Embedding adapter - for embedding vectors +export { + OpenAIEmbeddingAdapter, + createOpenaiEmbedding, + openaiEmbedding, + type OpenAIEmbeddingConfig, +} from './adapters/embedding' +export type { OpenAIEmbeddingProviderOptions } from './embedding/embedding-provider-options' + // ============================================================================ // Type Exports // ============================================================================ @@ -91,11 +100,15 @@ export type { OpenAIVideoModel, OpenAITTSModel, OpenAITranscriptionModel, + OpenAIEmbeddingModel, + OpenAIEmbeddingModelProviderOptionsByName, + OpenAIEmbeddingModelInputModalitiesByName, } from './model-meta' export { OPENAI_IMAGE_MODELS, OPENAI_TTS_MODELS, OPENAI_TRANSCRIPTION_MODELS, + OPENAI_EMBEDDING_MODELS, OPENAI_VIDEO_MODELS, OPENAI_CHAT_MODELS, } from './model-meta' diff --git a/packages/ai-openai/src/model-meta.ts b/packages/ai-openai/src/model-meta.ts index 4ed863083..d92bb7d69 100644 --- a/packages/ai-openai/src/model-meta.ts +++ b/packages/ai-openai/src/model-meta.ts @@ -7,6 +7,7 @@ import type { OpenAIStructuredOutputOptions, OpenAIToolsOptions, } from './text/text-provider-options' +import type { OpenAIEmbeddingProviderOptions } from './embedding/embedding-provider-options' interface ModelMeta { name: string @@ -2304,6 +2305,33 @@ export const OPENAI_TRANSCRIPTION_MODELS = [ export type OpenAITranscriptionModel = (typeof OPENAI_TRANSCRIPTION_MODELS)[number] +/** + * Embedding models (based on endpoints: "embeddings") + */ +export const OPENAI_EMBEDDING_MODELS = [ + 'text-embedding-3-small', + 'text-embedding-3-large', +] as const + +export type OpenAIEmbeddingModel = (typeof OPENAI_EMBEDDING_MODELS)[number] + +/** + * Type-only map from embedding model name to its provider options type. + */ +export type OpenAIEmbeddingModelProviderOptionsByName = { + 'text-embedding-3-small': OpenAIEmbeddingProviderOptions + 'text-embedding-3-large': OpenAIEmbeddingProviderOptions +} + +/** + * Per-model input modalities for embedding models. OpenAI embedding models + * are text-only, so image inputs fail at compile time. + */ +export type OpenAIEmbeddingModelInputModalitiesByName = { + 'text-embedding-3-small': readonly ['text'] + 'text-embedding-3-large': readonly ['text'] +} + /** * Type-only map from chat model name to its provider options type. * Used by the core AI types (via the adapter) to narrow diff --git a/packages/ai-openai/tests/embedding-adapter.test.ts b/packages/ai-openai/tests/embedding-adapter.test.ts new file mode 100644 index 000000000..349129086 --- /dev/null +++ b/packages/ai-openai/tests/embedding-adapter.test.ts @@ -0,0 +1,156 @@ +import { describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { + OpenAIEmbeddingAdapter, + createOpenaiEmbedding, +} from '../src/adapters/embedding' +import type OpenAI from 'openai' +import type { OpenAIEmbeddingModel } from '../src/model-meta' + +const testLogger = resolveDebugOption(false) + +/** + * Test-only subclass exposing the SDK client's `embeddings.create` to + * `vi.spyOn` — same pattern as `TestOpenAIImageAdapter`, keeping every type + * real with no casts. + */ +class TestOpenAIEmbeddingAdapter< + TModel extends OpenAIEmbeddingModel, +> extends OpenAIEmbeddingAdapter { + spyOnEmbeddingsCreate() { + return vi.spyOn(this.client.embeddings, 'create') + } +} + +function mockResponse( + vectors: Array>, +): OpenAI.CreateEmbeddingResponse { + return { + object: 'list', + model: 'text-embedding-3-small', + data: vectors.map((embedding, index) => ({ + object: 'embedding', + embedding, + index, + })), + usage: { prompt_tokens: 7, total_tokens: 7 }, + } +} + +describe('OpenAI Embedding Adapter', () => { + describe('createOpenaiEmbedding', () => { + it('creates an adapter with the provided API key', () => { + const adapter = createOpenaiEmbedding( + 'text-embedding-3-small', + 'test-api-key', + ) + expect(adapter).toBeInstanceOf(OpenAIEmbeddingAdapter) + expect(adapter.kind).toBe('embedding') + expect(adapter.name).toBe('openai') + expect(adapter.model).toBe('text-embedding-3-small') + }) + }) + + describe('createEmbeddings', () => { + it('sends texts as a batch with encoding_format float', async () => { + const adapter = new TestOpenAIEmbeddingAdapter( + { apiKey: 'test' }, + 'text-embedding-3-small', + ) + const spy = adapter.spyOnEmbeddingsCreate() + spy.mockResolvedValue( + mockResponse([ + [0.1, 0.2], + [0.3, 0.4], + ]), + ) + + const result = await adapter.createEmbeddings({ + model: 'text-embedding-3-small', + input: ['a red guitar', { type: 'text', content: 'a blue drum kit' }], + logger: testLogger, + }) + + expect(spy).toHaveBeenCalledWith({ + model: 'text-embedding-3-small', + input: ['a red guitar', 'a blue drum kit'], + encoding_format: 'float', + }) + expect(result.embeddings).toEqual([ + { vector: [0.1, 0.2], index: 0 }, + { vector: [0.3, 0.4], index: 1 }, + ]) + expect(result.usage).toEqual({ + promptTokens: 7, + completionTokens: 0, + totalTokens: 7, + }) + }) + + it('passes the top-level dimensions option through', async () => { + const adapter = new TestOpenAIEmbeddingAdapter( + { apiKey: 'test' }, + 'text-embedding-3-large', + ) + const spy = adapter.spyOnEmbeddingsCreate() + spy.mockResolvedValue(mockResponse([[0.1]])) + + await adapter.createEmbeddings({ + model: 'text-embedding-3-large', + input: ['hello'], + dimensions: 1024, + logger: testLogger, + }) + + expect(spy).toHaveBeenCalledWith( + expect.objectContaining({ dimensions: 1024 }), + ) + }) + + it('passes provider options through without letting them override model/input', async () => { + const adapter = new TestOpenAIEmbeddingAdapter( + { apiKey: 'test' }, + 'text-embedding-3-small', + ) + const spy = adapter.spyOnEmbeddingsCreate() + spy.mockResolvedValue(mockResponse([[0.1]])) + + await adapter.createEmbeddings({ + model: 'text-embedding-3-small', + input: ['hello'], + modelOptions: { user: 'user-123' }, + logger: testLogger, + }) + + expect(spy).toHaveBeenCalledWith( + expect.objectContaining({ + user: 'user-123', + model: 'text-embedding-3-small', + input: ['hello'], + }), + ) + }) + + it('throws a clear error for image input', async () => { + const adapter = new TestOpenAIEmbeddingAdapter( + { apiKey: 'test' }, + 'text-embedding-3-small', + ) + const spy = adapter.spyOnEmbeddingsCreate() + + await expect( + adapter.createEmbeddings({ + model: 'text-embedding-3-small', + input: [ + { + type: 'image', + source: { type: 'data', value: 'aGk=', mimeType: 'image/png' }, + }, + ], + logger: testLogger, + }), + ).rejects.toThrow('only supports text embedding inputs') + expect(spy).not.toHaveBeenCalled() + }) + }) +}) diff --git a/packages/ai/src/activities/embed/adapter.ts b/packages/ai/src/activities/embed/adapter.ts new file mode 100644 index 000000000..4ddb6a45e --- /dev/null +++ b/packages/ai/src/activities/embed/adapter.ts @@ -0,0 +1,112 @@ +import type { + EmbeddingModelInputModalitiesByName, + EmbeddingOptions, + EmbeddingResult, +} from '../../types' + +/** + * Configuration for embedding adapter instances + */ +export interface EmbeddingAdapterConfig { + apiKey?: string + baseUrl?: string + timeout?: number + maxRetries?: number + headers?: Record +} + +/** + * Embedding adapter interface with pre-resolved generics. + * + * An adapter is created by a provider function: `provider('model')` → `adapter` + * All type resolution happens at the provider call site, not in this interface. + * + * Generic parameters: + * - TModel: The specific model name (e.g., 'text-embedding-3-small') + * - TProviderOptions: Base provider-specific options (already resolved) + * - TModelProviderOptionsByName: Map from model name to its specific provider options + * - TModelInputModalitiesByName: Map from model name to the input modalities it + * accepts (constrains the `input` item types at compile time) + */ +export interface EmbeddingAdapter< + TModel extends string = string, + TProviderOptions extends object = Record, + TModelProviderOptionsByName extends Record = Record, + TModelInputModalitiesByName extends EmbeddingModelInputModalitiesByName = + EmbeddingModelInputModalitiesByName, +> { + /** Discriminator for adapter kind */ + readonly kind: 'embedding' + /** Adapter name identifier */ + readonly name: string + /** The model this adapter is configured for */ + readonly model: TModel + + /** + * @internal Type-only properties for inference. Not assigned at runtime. + */ + '~types': { + providerOptions: TProviderOptions + modelProviderOptionsByName: TModelProviderOptionsByName + modelInputModalitiesByName: TModelInputModalitiesByName + } + + /** + * Generate embeddings for the input items (one vector per item) + */ + createEmbeddings: ( + options: EmbeddingOptions, + ) => Promise +} + +/** + * An EmbeddingAdapter with any/unknown type parameters. + * Useful as a constraint in generic functions and interfaces. + */ +export type AnyEmbeddingAdapter = EmbeddingAdapter + +/** + * Abstract base class for embedding adapters. + * Extend this class to implement an embedding adapter for a specific provider. + * + * Generic parameters match EmbeddingAdapter - all pre-resolved by the provider function. + */ +export abstract class BaseEmbeddingAdapter< + TModel extends string = string, + TProviderOptions extends object = Record, + TModelProviderOptionsByName extends Record = Record, + TModelInputModalitiesByName extends EmbeddingModelInputModalitiesByName = + EmbeddingModelInputModalitiesByName, +> implements EmbeddingAdapter< + TModel, + TProviderOptions, + TModelProviderOptionsByName, + TModelInputModalitiesByName +> { + readonly kind = 'embedding' as const + abstract readonly name: string + readonly model: TModel + + // Type-only property - never assigned at runtime + declare '~types': { + providerOptions: TProviderOptions + modelProviderOptionsByName: TModelProviderOptionsByName + modelInputModalitiesByName: TModelInputModalitiesByName + } + + protected config: EmbeddingAdapterConfig + + constructor(model: TModel, config: EmbeddingAdapterConfig = {}) { + this.config = config + this.model = model + } + + abstract createEmbeddings( + options: EmbeddingOptions, + ): Promise + + protected generateId(prefix?: string): string { + const p = prefix ?? this.name + return `${p}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}` + } +} diff --git a/packages/ai/src/activities/embed/index.ts b/packages/ai/src/activities/embed/index.ts new file mode 100644 index 000000000..56d48a4c8 --- /dev/null +++ b/packages/ai/src/activities/embed/index.ts @@ -0,0 +1,316 @@ +/** + * Embed Activity + * + * Generates embedding vectors from text and (for multimodal models) image + * inputs. This is a self-contained module with implementation, types, and JSDoc. + */ + +import { aiEventClient } from '@tanstack/ai-event-client' +import { resolveDebugOption } from '../../logger/resolve' +import { + createGenerationContext, + runGenerationError, + runGenerationFinish, + runGenerationStart, + runGenerationUsage, +} from '../middleware/run' +import { countEmbeddingInputModalities } from '../../utilities/embedding-input' +import type { InternalLogger } from '../../logger/internal-logger' +import type { DebugOption } from '../../logger/types' +import type { GenerationMiddleware } from '../middleware/types' +import type { EmbeddingAdapter } from './adapter' +import type { + EmbeddingInputItem, + EmbeddingInputItemFor, + EmbeddingResult, +} from '../../types' + +// =========================== +// Activity Kind +// =========================== + +/** The adapter kind this activity handles */ +export const kind = 'embedding' as const + +// =========================== +// Type Extraction Helpers +// =========================== + +/** + * Extract model-specific provider options from an EmbeddingAdapter via ~types. + * If the model has specific options defined in ModelProviderOptions (and not just via index signature), + * use those; otherwise fall back to base provider options. + */ +export type EmbedProviderOptionsForModel = + TAdapter extends EmbeddingAdapter< + any, + infer BaseOptions, + infer ModelOptions, + any + > + ? string extends keyof ModelOptions + ? // ModelOptions is Record or has index signature - use BaseOptions + BaseOptions + : // ModelOptions has explicit keys - check if TModel is one of them + TModel extends keyof ModelOptions + ? ModelOptions[TModel] + : BaseOptions + : object + +/** + * Extract the input type a model accepts from an EmbeddingAdapter via ~types. + * Adapters declare a per-model input-modality map; models in the map get an + * `input` narrowed to their supported item types (text-only models accept + * `string | TextPart`), so unsupported items fail at compile time. Adapters + * without a map fall back to the full EmbeddingInputItem union. + */ +export type EmbeddingInputForModel = + TAdapter extends EmbeddingAdapter + ? string extends keyof ModsByName + ? // No explicit map - accept the full union + EmbeddingInputItem | Array + : TModel extends keyof ModsByName + ? + | EmbeddingInputItemFor + | Array> + : EmbeddingInputItem | Array + : EmbeddingInputItem | Array + +// =========================== +// Activity Options Type +// =========================== + +/** + * Options for the embed activity. + * The model is extracted from the adapter's model property. + * + * @template TAdapter - The embedding adapter type + */ +export type EmbedOptions< + TAdapter extends EmbeddingAdapter, +> = { + /** The embedding adapter to use (must be created with a model) */ + adapter: TAdapter & { kind: typeof kind } + /** + * What to embed: a single item or an array of items. Each item produces + * exactly one vector. An item is a plain string, a text part, an image + * part, or — for models that embed text and image together — a fused + * `{ type: 'content', content: [...] }` item. The accepted item types are + * narrowed per model via the adapter's input-modality map. + */ + input: EmbeddingInputForModel + /** + * Requested output dimensionality. Supported by models with Matryoshka / + * configurable dimensions; adapters for fixed-dimension models throw a + * clear runtime error when this is set. + */ + dimensions?: number + /** + * Enable debug logging. Pass `true` to enable all categories, `false` to + * silence everything including errors, or a `DebugConfig` object for granular + * control and/or a custom `Logger`. + */ + debug?: DebugOption + /** + * Observe-only middleware notified on start, usage, success, and error. Pass + * `otelMiddleware()` to emit OpenTelemetry spans, or implement the + * `GenerationMiddleware` contract for a custom backend. + */ + middleware?: Array +} & ({} extends EmbedProviderOptionsForModel + ? { + /** Provider-specific options for embedding generation */ modelOptions?: EmbedProviderOptionsForModel< + TAdapter, + TAdapter['model'] + > + } + : { + /** Provider-specific options for embedding generation */ modelOptions: EmbedProviderOptionsForModel< + TAdapter, + TAdapter['model'] + > + }) + +function createId(prefix: string): string { + return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}` +} + +// =========================== +// Activity Implementation +// =========================== + +/** + * Embed activity - generates embedding vectors from text and image inputs. + * + * Accepts a single item or an array of items; the result always carries an + * `embeddings` array with one vector per input item, in input order. + * + * @example Embed a single text + * ```ts + * import { embed } from '@tanstack/ai' + * import { openaiEmbedding } from '@tanstack/ai-openai' + * + * const result = await embed({ + * adapter: openaiEmbedding('text-embedding-3-small'), + * input: 'a red guitar', + * }) + * + * console.log(result.embeddings[0].vector) + * ``` + * + * @example Batch with requested dimensions + * ```ts + * const result = await embed({ + * adapter: openaiEmbedding('text-embedding-3-large'), + * input: ['a red guitar', 'a blue drum kit'], + * dimensions: 1024, + * }) + * ``` + * + * @example Multimodal embedding (text + image, one fused vector) + * ```ts + * import { cohereEmbedding } from '@tanstack/ai-cohere' + * + * const result = await embed({ + * adapter: cohereEmbedding('embed-v4.0'), + * input: { + * type: 'content', + * content: [ + * { type: 'text', content: 'product photo' }, + * { type: 'image', source: { type: 'data', value: base64, mimeType: 'image/png' } }, + * ], + * }, + * modelOptions: { inputType: 'search_document' }, + * }) + * ``` + */ +export async function embed< + TAdapter extends EmbeddingAdapter, +>(options: EmbedOptions): Promise { + const { adapter, middleware } = options + const model = adapter.model + const requestId = createId('embedding') + const startTime = Date.now() + const logger: InternalLogger = resolveDebugOption(options.debug) + const modelOptions = (options as { modelOptions?: Record }) + .modelOptions + + // Normalize once: adapters always receive an array of items. + const inputItems: Array = Array.isArray(options.input) + ? options.input + : [options.input] + const { textInputCount, imageInputCount } = + countEmbeddingInputModalities(inputItems) + + const mwCtx = createGenerationContext({ + requestId, + activity: 'embedding', + provider: adapter.name, + model, + modelOptions, + createId, + }) + + await runGenerationStart(middleware, mwCtx) + + aiEventClient.emit('embedding:request:started', { + requestId, + provider: adapter.name, + model, + inputCount: inputItems.length, + textInputCount, + imageInputCount, + dimensions: options.dimensions, + modelOptions, + timestamp: startTime, + }) + + logger.request(`activity=embed provider=${adapter.name} model=${model}`, { + provider: adapter.name, + model, + }) + + try { + const result = await adapter.createEmbeddings({ + model, + input: inputItems, + dimensions: options.dimensions, + modelOptions, + logger, + }) + const duration = Date.now() - startTime + + aiEventClient.emit('embedding:request:completed', { + requestId, + provider: adapter.name, + model, + embeddingCount: result.embeddings.length, + dimensions: result.embeddings[0]?.vector.length, + duration, + modelOptions, + timestamp: Date.now(), + }) + + logger.output(`activity=embed count=${result.embeddings.length}`, { + embeddingCount: result.embeddings.length, + }) + + if (result.usage) { + aiEventClient.emit('embedding:usage', { + requestId, + model, + usage: result.usage, + timestamp: Date.now(), + }) + await runGenerationUsage(middleware, mwCtx, result.usage) + } + await runGenerationFinish(middleware, mwCtx, { + duration, + usage: result.usage, + }) + + return result + } catch (error) { + const duration = Date.now() - startTime + const err = error as Error + aiEventClient.emit('embedding:request:error', { + requestId, + provider: adapter.name, + model, + error: { message: err.message, name: err.name }, + duration, + modelOptions, + timestamp: Date.now(), + }) + await runGenerationError(middleware, mwCtx, { + error, + duration, + }) + logger.errors('embed activity failed', { + error, + source: 'embed', + }) + throw error + } +} + +// =========================== +// Options Factory +// =========================== + +/** + * Create typed options for the embed() function without executing. + */ +export function createEmbedOptions< + TAdapter extends EmbeddingAdapter, +>(options: EmbedOptions): EmbedOptions { + return options +} + +// Re-export adapter types +export type { + EmbeddingAdapter, + EmbeddingAdapterConfig, + AnyEmbeddingAdapter, +} from './adapter' +export { BaseEmbeddingAdapter } from './adapter' diff --git a/packages/ai/src/activities/index.ts b/packages/ai/src/activities/index.ts index 07fdbe73a..32db6e645 100644 --- a/packages/ai/src/activities/index.ts +++ b/packages/ai/src/activities/index.ts @@ -21,6 +21,7 @@ import type { AnyAudioAdapter } from './generateAudio/adapter' import type { AnyVideoAdapter } from './generateVideo/adapter' import type { AnyTTSAdapter } from './generateSpeech/adapter' import type { AnyTranscriptionAdapter } from './generateTranscription/adapter' +import type { AnyEmbeddingAdapter } from './embed/adapter' // =========================== // Chat Activity @@ -170,6 +171,25 @@ export { type AnyTranscriptionAdapter, } from './generateTranscription/adapter' +// =========================== +// Embed Activity +// =========================== + +export { + kind as embeddingKind, + embed, + type EmbedOptions, + type EmbedProviderOptionsForModel, + type EmbeddingInputForModel, +} from './embed/index' + +export { + BaseEmbeddingAdapter, + type EmbeddingAdapter, + type EmbeddingAdapterConfig, + type AnyEmbeddingAdapter, +} from './embed/adapter' + // =========================== // Adapter Union Types // =========================== @@ -183,6 +203,7 @@ export type AIAdapter = | AnyVideoAdapter | AnyTTSAdapter | AnyTranscriptionAdapter + | AnyEmbeddingAdapter /** Union type of all adapter kinds */ export type AdapterKind = @@ -193,3 +214,4 @@ export type AdapterKind = | 'video' | 'tts' | 'transcription' + | 'embedding' diff --git a/packages/ai/src/activities/middleware/types.ts b/packages/ai/src/activities/middleware/types.ts index 99ae3e44e..b56532195 100644 --- a/packages/ai/src/activities/middleware/types.ts +++ b/packages/ai/src/activities/middleware/types.ts @@ -38,6 +38,7 @@ export type GenerationActivity = | 'audio' | 'tts' | 'transcription' + | 'embedding' /** * Stable context passed to every {@link GenerationMiddleware} hook. Created diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 102987c02..f9a9ed71a 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -8,6 +8,7 @@ export { getVideoJobStatus, generateSpeech, generateTranscription, + embed, } from './activities/index' // Create options functions - for pre-defining typed configurations @@ -18,6 +19,7 @@ export { createAudioOptions } from './activities/generateAudio/index' export { createVideoOptions } from './activities/generateVideo/index' export { createSpeechOptions } from './activities/generateSpeech/index' export { createTranscriptionOptions } from './activities/generateTranscription/index' +export { createEmbedOptions } from './activities/embed/index' // Re-export types export type { @@ -36,6 +38,8 @@ export type { TranscriptionAdapter, AnyVideoAdapter, VideoAdapter, + AnyEmbeddingAdapter, + EmbeddingAdapter, } from './activities/index' // Tool definition @@ -171,6 +175,14 @@ export { buildBaseUsage, type BaseUsageInput } from './utilities/usage' export { resolveMediaPrompt } from './utilities/media-prompt' export type { ResolvedMediaPrompt } from './utilities/media-prompt' +// Embedding input resolution (used by embedding adapters) +export { + resolveEmbeddingInput, + requireTextOnlyEmbeddingInput, + countEmbeddingInputModalities, +} from './utilities/embedding-input' +export type { ResolvedEmbeddingItem } from './utilities/embedding-input' + // System prompts (type + normaliser used by adapters) export type { SystemPrompt, NormalizedSystemPrompt } from './system-prompts' export { normalizeSystemPrompts } from './system-prompts' diff --git a/packages/ai/src/middlewares/otel.ts b/packages/ai/src/middlewares/otel.ts index f4ddba35e..9fc465437 100644 --- a/packages/ai/src/middlewares/otel.ts +++ b/packages/ai/src/middlewares/otel.ts @@ -83,6 +83,7 @@ const OPERATION_NAME: Record = { audio: 'audio_generation', tts: 'text_to_speech', transcription: 'transcription', + embedding: 'embeddings', } export interface OtelMiddlewareOptions { diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 6917c8140..4a7607f96 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -2111,6 +2111,117 @@ export interface TranscriptionResult { usage?: TokenUsage } +// ============================================================================ +// Embedding Types +// ============================================================================ + +/** + * Input modalities an embedding model can accept. Unlike + * {@link MediaPromptModality}, `'text'` is listed explicitly because + * text-only embedding models are the common case and the modality list + * drives compile-time narrowing of {@link EmbeddingInputItem}. + */ +export type EmbeddingModality = 'text' | 'image' + +/** + * Per-model map from model name to the input modalities it accepts, used as + * an adapter type parameter (`TModelInputModalitiesByName`). Models absent + * from the map fall back to the unconstrained {@link EmbeddingInputItem}. + */ +export type EmbeddingModelInputModalitiesByName = Record< + string, + ReadonlyArray +> + +/** + * A fused multi-part embedding item: all parts are embedded together into a + * single vector (e.g. a product photo plus its caption). Supported by + * multimodal embedding models such as Cohere embed-v4 and Amazon Titan + * Multimodal. Distinct from passing multiple items in the `input` array, + * where each item produces its own vector. + */ +export interface EmbeddingContentItem { + type: 'content' + content: Array +} + +/** + * One embeddable item. A bare string is shorthand for a text part. Each item + * in an `embed()` input array produces exactly one vector. + */ +export type EmbeddingInputItem = + | string + | TextPart + | ImagePart + | EmbeddingContentItem + +/** Maps an embedding modality to the item types it admits. @internal */ +interface EmbeddingItemByModality { + text: TextPart + image: ImagePart | EmbeddingContentItem +} + +/** + * Embedding item type narrowed to the modalities a specific model supports. + * `EmbeddingInputItemFor<'text'>` (a text-only model) is `string | TextPart`; + * `'text' | 'image'` additionally admits image parts and fused + * {@link EmbeddingContentItem}s. Used by the activity option types together + * with the adapter's per-model modality map so unsupported inputs fail at + * compile time. + */ +export type EmbeddingInputItemFor< + TModalities extends EmbeddingModality = EmbeddingModality, +> = string | TextPart | EmbeddingItemByModality[TModalities] + +/** + * Options for embedding generation, as received by adapters. The `embed()` + * entry point normalizes a single input item to an array before calling the + * adapter, so `input` is always an array here. + */ +export interface EmbeddingOptions { + /** The model to use for embedding generation */ + model: string + /** The items to embed — one vector per item */ + input: Array + /** + * Requested output dimensionality. Adapters for models with fixed + * dimensions throw a clear runtime error when this is set. + */ + dimensions?: number + /** Model-specific options for embedding generation */ + modelOptions?: TProviderOptions + /** + * Internal logger threaded from the embed() entry point. Adapters must + * call logger.request() before the SDK call and logger.errors() in catch + * blocks. + */ + logger: InternalLogger +} + +/** + * A single embedding vector. + */ +export interface Embedding { + /** The embedding vector */ + vector: Array + /** Position of the source item in the (normalized) input array */ + index: number +} + +/** + * Result of embedding generation. + */ +export interface EmbeddingResult { + /** Unique identifier for the generation */ + id: string + /** Model used for generation */ + model: string + /** One embedding per input item, in input order */ + embeddings: Array + /** Token usage information (if provided by the adapter) */ + usage?: TokenUsage +} + /** * Default metadata type for adapters that don't define custom metadata. * Uses unknown for all modalities. diff --git a/packages/ai/src/utilities/embedding-input.ts b/packages/ai/src/utilities/embedding-input.ts new file mode 100644 index 000000000..45cbdfb8b --- /dev/null +++ b/packages/ai/src/utilities/embedding-input.ts @@ -0,0 +1,82 @@ +import type { EmbeddingInputItem, ImagePart } from '../types' + +/** + * One embedding input item resolved into its text and image constituents. + * Produced by {@link resolveEmbeddingInput}; adapters map each entry onto + * one provider-native input (one vector per entry). + */ +export interface ResolvedEmbeddingItem { + /** Text contents of the item, in order (empty for image-only items) */ + texts: Array + /** Image parts of the item, in order (empty for text-only items) */ + images: Array +} + +function resolveItem(item: EmbeddingInputItem): ResolvedEmbeddingItem { + if (typeof item === 'string') { + return { texts: [item], images: [] } + } + if (item.type === 'text') { + return { texts: [item.content], images: [] } + } + if (item.type === 'image') { + return { texts: [], images: [item] } + } + const resolved: ResolvedEmbeddingItem = { texts: [], images: [] } + for (const part of item.content) { + if (part.type === 'text') { + resolved.texts.push(part.content) + } else { + resolved.images.push(part) + } + } + return resolved +} + +/** + * Resolve each embedding input item into its text and image constituents, + * preserving input order (result[i] corresponds to input[i] and to the + * vector at index i). + */ +export function resolveEmbeddingInput( + input: Array, +): Array { + return input.map(resolveItem) +} + +/** + * Extract plain text inputs for a text-only embedding model, throwing a + * uniform error if any item carries an image. The per-model modality typing + * rejects these at compile time; this guard covers untyped/dynamic callers. + */ +export function requireTextOnlyEmbeddingInput( + input: Array, + provider: string, + model: string, +): Array { + return resolveEmbeddingInput(input).map((item, index) => { + if (item.images.length > 0) { + throw new Error( + `${provider} model "${model}" only supports text embedding inputs; ` + + `input item at index ${index} contains an image part`, + ) + } + return item.texts.join('\n') + }) +} + +/** + * Count text-only and image-carrying items for observability events. Never + * exposes input content. + */ +export function countEmbeddingInputModalities( + input: Array, +): { textInputCount: number; imageInputCount: number } { + let textInputCount = 0 + let imageInputCount = 0 + for (const item of resolveEmbeddingInput(input)) { + if (item.images.length > 0) imageInputCount++ + else textInputCount++ + } + return { textInputCount, imageInputCount } +} diff --git a/packages/ai/tests/embed-per-model-type-safety.test.ts b/packages/ai/tests/embed-per-model-type-safety.test.ts new file mode 100644 index 000000000..7be68c761 --- /dev/null +++ b/packages/ai/tests/embed-per-model-type-safety.test.ts @@ -0,0 +1,268 @@ +/** + * Type Safety Tests for embed() function + * + * These tests verify that the embed() function correctly constrains types based on: + * 1. Model-specific input modalities - text-only models reject image parts + * 2. Model-specific provider options (modelOptions), including required options + * + * Uses @ts-expect-error to ensure TypeScript catches invalid type combinations. + */ +import { describe, expectTypeOf, it } from 'vitest' +import { BaseEmbeddingAdapter } from '../src/activities/embed/adapter' +import { embed } from '../src/activities/embed/index' +import type { EmbeddingOptions, EmbeddingResult } from '../src/types' + +// =========================== +// Mock Provider Options Types +// =========================== + +/** + * Options for the text-only mock model (all optional). + */ +interface MockTextEmbedProviderOptions { + /** + * A unique identifier representing your end-user. + */ + user?: string +} + +/** + * Options for the multimodal mock model — `inputType` is REQUIRED, which + * makes `modelOptions` itself required at the embed() call site. + */ +interface MockMultimodalEmbedProviderOptions { + inputType: 'search_document' | 'search_query' + truncate?: 'NONE' | 'END' +} + +type MockEmbedProviderOptions = + | MockTextEmbedProviderOptions + | MockMultimodalEmbedProviderOptions + +// =========================== +// Mock Type Maps +// =========================== + +type MockEmbeddingModelProviderOptionsByName = { + 'mock-text-embed': MockTextEmbedProviderOptions + 'mock-mm-embed': MockMultimodalEmbedProviderOptions +} + +type MockEmbeddingModelInputModalitiesByName = { + 'mock-text-embed': readonly ['text'] + 'mock-mm-embed': readonly ['text', 'image'] +} + +const MOCK_EMBEDDING_MODELS = ['mock-text-embed', 'mock-mm-embed'] as const + +type MockEmbeddingModel = (typeof MOCK_EMBEDDING_MODELS)[number] + +// =========================== +// Mock Adapter Implementation +// =========================== + +class MockEmbeddingAdapter< + TModel extends MockEmbeddingModel, +> extends BaseEmbeddingAdapter< + TModel, + MockEmbedProviderOptions, + MockEmbeddingModelProviderOptionsByName, + MockEmbeddingModelInputModalitiesByName +> { + readonly name = 'mock' as const + + constructor(model: TModel) { + super(model, {}) + } + + /* eslint-disable @typescript-eslint/require-await */ + createEmbeddings = async ( + options: EmbeddingOptions, + ): Promise => { + return { + id: 'mock-id', + model: this.model, + embeddings: options.input.map((_, index) => ({ + vector: [0.1, 0.2], + index, + })), + } + } + /* eslint-enable @typescript-eslint/require-await */ +} + +function mockEmbedding( + model: TModel, +): MockEmbeddingAdapter { + return new MockEmbeddingAdapter(model) +} + +const imagePart = { + type: 'image', + source: { type: 'data', value: 'aGVsbG8=', mimeType: 'image/png' }, +} as const + +// =========================== +// Type Safety Tests +// =========================== + +describe('Type Safety Tests for embed() function', () => { + describe('Per-model input modality type safety', () => { + it('allows string and text-part inputs on text-only models', () => { + embed({ + adapter: mockEmbedding('mock-text-embed'), + input: 'a red guitar', + }) + + embed({ + adapter: mockEmbedding('mock-text-embed'), + input: { type: 'text', content: 'a red guitar' }, + }) + + embed({ + adapter: mockEmbedding('mock-text-embed'), + input: ['a red guitar', { type: 'text', content: 'a blue drum kit' }], + }) + }) + + it('rejects image parts on text-only models', () => { + embed({ + adapter: mockEmbedding('mock-text-embed'), + // @ts-expect-error - mock-text-embed does not accept image inputs + input: imagePart, + }) + + embed({ + adapter: mockEmbedding('mock-text-embed'), + // @ts-expect-error - mock-text-embed does not accept image inputs in arrays + input: ['a red guitar', imagePart], + }) + }) + + it('rejects fused content items on text-only models', () => { + embed({ + adapter: mockEmbedding('mock-text-embed'), + // @ts-expect-error - mock-text-embed does not accept fused content items + input: { + type: 'content', + content: [{ type: 'text', content: 'caption' }, imagePart], + }, + }) + }) + + it('allows image parts and fused items on multimodal models', () => { + embed({ + adapter: mockEmbedding('mock-mm-embed'), + input: imagePart, + modelOptions: { inputType: 'search_document' }, + }) + + embed({ + adapter: mockEmbedding('mock-mm-embed'), + input: [ + 'a red guitar', + imagePart, + { + type: 'content', + content: [{ type: 'text', content: 'caption' }, imagePart], + }, + ], + modelOptions: { inputType: 'search_document' }, + }) + }) + }) + + describe('Model-specific provider options type safety', () => { + it('keeps modelOptions optional when all options are optional', () => { + embed({ + adapter: mockEmbedding('mock-text-embed'), + input: 'hello', + }) + + embed({ + adapter: mockEmbedding('mock-text-embed'), + input: 'hello', + modelOptions: { user: 'user-123' }, + }) + }) + + it('requires modelOptions when the model has required options', () => { + // @ts-expect-error - mock-mm-embed requires modelOptions.inputType + embed({ + adapter: mockEmbedding('mock-mm-embed'), + input: 'hello', + }) + + embed({ + adapter: mockEmbedding('mock-mm-embed'), + input: 'hello', + modelOptions: { inputType: 'search_query' }, + }) + }) + + it('rejects options from the other model', () => { + embed({ + adapter: mockEmbedding('mock-text-embed'), + input: 'hello', + modelOptions: { + // @ts-expect-error - 'inputType' is not available on mock-text-embed + inputType: 'search_document', + }, + }) + + embed({ + adapter: mockEmbedding('mock-mm-embed'), + input: 'hello', + modelOptions: { + inputType: 'search_document', + // @ts-expect-error - 'user' is not available on mock-mm-embed + user: 'user-123', + }, + }) + }) + + it('rejects unknown options', () => { + embed({ + adapter: mockEmbedding('mock-text-embed'), + input: 'hello', + modelOptions: { + // @ts-expect-error - 'unknownOption' does not exist + unknownOption: true, + }, + }) + }) + }) + + describe('Model name type safety', () => { + it('accepts valid model names', () => { + const _adapter1 = mockEmbedding('mock-text-embed') + const _adapter2 = mockEmbedding('mock-mm-embed') + expectTypeOf(_adapter1).toBeObject() + expectTypeOf(_adapter2).toBeObject() + }) + + it('rejects invalid model names', () => { + // @ts-expect-error - 'invalid-model' is not a valid mock embedding model name + const _adapter = mockEmbedding('invalid-model') + }) + }) + + describe('Result type', () => { + it('resolves to EmbeddingResult', () => { + const result = embed({ + adapter: mockEmbedding('mock-text-embed'), + input: 'hello', + }) + expectTypeOf(result).resolves.toEqualTypeOf() + }) + + it('exposes vector and index on embeddings', async () => { + const result = await embed({ + adapter: mockEmbedding('mock-text-embed'), + input: 'hello', + }) + expectTypeOf(result.embeddings[0]!.vector).toEqualTypeOf>() + expectTypeOf(result.embeddings[0]!.index).toEqualTypeOf() + }) + }) +}) diff --git a/packages/ai/tests/embed.test.ts b/packages/ai/tests/embed.test.ts new file mode 100644 index 000000000..8a9cda1ff --- /dev/null +++ b/packages/ai/tests/embed.test.ts @@ -0,0 +1,285 @@ +import { describe, expect, it, vi } from 'vitest' +import { embed } from '../src/activities/embed/index' +import { BaseEmbeddingAdapter } from '../src/activities/embed/adapter' +import { + countEmbeddingInputModalities, + requireTextOnlyEmbeddingInput, + resolveEmbeddingInput, +} from '../src/utilities/embedding-input' +import type { GenerationMiddleware } from '../src/activities/middleware/types' +import type { + EmbeddingInputItem, + EmbeddingOptions, + EmbeddingResult, + TokenUsage, +} from '../src/types' + +// ============================================================================ +// Mock adapter +// ============================================================================ + +interface MockAdapterBehavior { + usage?: TokenUsage + error?: Error + dimensions?: number +} + +class MockEmbeddingAdapter extends BaseEmbeddingAdapter< + 'mock-embed', + { inputType?: string } +> { + readonly name = 'mock' + calls: Array> = [] + + constructor(private behavior: MockAdapterBehavior = {}) { + super('mock-embed', {}) + } + + createEmbeddings( + options: EmbeddingOptions<{ inputType?: string }>, + ): Promise { + this.calls.push(options) + if (this.behavior.error) return Promise.reject(this.behavior.error) + const dimensions = this.behavior.dimensions ?? options.dimensions ?? 4 + const result: EmbeddingResult = { + id: 'mock-1', + model: options.model, + embeddings: options.input.map((_, index) => ({ + vector: Array.from({ length: dimensions }, () => index + 0.5), + index, + })), + } + if (this.behavior.usage) result.usage = this.behavior.usage + return Promise.resolve(result) + } +} + +// ============================================================================ +// embed() runtime behavior +// ============================================================================ + +describe('embed()', () => { + it('normalizes a single string input to an array for the adapter', async () => { + const adapter = new MockEmbeddingAdapter() + const result = await embed({ adapter, input: 'hello world' }) + + expect(adapter.calls).toHaveLength(1) + expect(adapter.calls[0]!.input).toEqual(['hello world']) + expect(result.embeddings).toHaveLength(1) + expect(result.embeddings[0]!.index).toBe(0) + expect(result.embeddings[0]!.vector).toHaveLength(4) + }) + + it('normalizes a single non-string item to an array', async () => { + const adapter = new MockEmbeddingAdapter() + await embed({ + adapter, + input: { type: 'text', content: 'hello' }, + }) + + expect(adapter.calls[0]!.input).toEqual([ + { type: 'text', content: 'hello' }, + ]) + }) + + it('passes batch input through in order with matching indices', async () => { + const adapter = new MockEmbeddingAdapter() + const result = await embed({ + adapter, + input: ['one', 'two', 'three'], + }) + + expect(adapter.calls[0]!.input).toEqual(['one', 'two', 'three']) + expect(result.embeddings.map((e) => e.index)).toEqual([0, 1, 2]) + }) + + it('threads dimensions and modelOptions to the adapter', async () => { + const adapter = new MockEmbeddingAdapter() + await embed({ + adapter, + input: 'hello', + dimensions: 256, + modelOptions: { inputType: 'search_document' }, + }) + + expect(adapter.calls[0]!.dimensions).toBe(256) + expect(adapter.calls[0]!.modelOptions).toEqual({ + inputType: 'search_document', + }) + expect(adapter.calls[0]!.model).toBe('mock-embed') + }) + + it('returns the adapter result with usage', async () => { + const usage: TokenUsage = { + promptTokens: 12, + completionTokens: 0, + totalTokens: 12, + } + const adapter = new MockEmbeddingAdapter({ usage }) + const result = await embed({ adapter, input: ['a', 'b'] }) + + expect(result.usage).toEqual(usage) + expect(result.model).toBe('mock-embed') + }) + + it('rethrows adapter errors', async () => { + const adapter = new MockEmbeddingAdapter({ error: new Error('boom') }) + await expect(embed({ adapter, input: 'x' })).rejects.toThrow('boom') + }) + + // ========================================================================== + // Middleware lifecycle + // ========================================================================== + + describe('middleware', () => { + it('calls onStart, onUsage, and onFinish in order on success', async () => { + const order: Array = [] + const middleware: GenerationMiddleware = { + name: 'test', + onStart: vi.fn(async () => { + order.push('start') + }), + onUsage: vi.fn(async () => { + order.push('usage') + }), + onFinish: vi.fn(async () => { + order.push('finish') + }), + onError: vi.fn(), + } + const usage: TokenUsage = { + promptTokens: 3, + completionTokens: 0, + totalTokens: 3, + } + const adapter = new MockEmbeddingAdapter({ usage }) + + await embed({ adapter, input: 'hello', middleware: [middleware] }) + + expect(order).toEqual(['start', 'usage', 'finish']) + expect(middleware.onError).not.toHaveBeenCalled() + const finishInfo = vi.mocked(middleware.onFinish!).mock.calls[0]![1] + expect(finishInfo.usage).toEqual(usage) + expect(typeof finishInfo.duration).toBe('number') + }) + + it('skips onUsage when the adapter reports no usage', async () => { + const middleware: GenerationMiddleware = { + name: 'test', + onUsage: vi.fn(), + onFinish: vi.fn(), + } + const adapter = new MockEmbeddingAdapter() + + await embed({ adapter, input: 'hello', middleware: [middleware] }) + + expect(middleware.onUsage).not.toHaveBeenCalled() + expect(middleware.onFinish).toHaveBeenCalledOnce() + }) + + it('calls onError (not onFinish) when the adapter throws', async () => { + const middleware: GenerationMiddleware = { + name: 'test', + onStart: vi.fn(), + onFinish: vi.fn(), + onError: vi.fn(), + } + const adapter = new MockEmbeddingAdapter({ error: new Error('boom') }) + + await expect( + embed({ adapter, input: 'x', middleware: [middleware] }), + ).rejects.toThrow('boom') + + expect(middleware.onStart).toHaveBeenCalledOnce() + expect(middleware.onFinish).not.toHaveBeenCalled() + expect(middleware.onError).toHaveBeenCalledOnce() + const errorInfo = vi.mocked(middleware.onError!).mock.calls[0]![1] + expect((errorInfo.error as Error).message).toBe('boom') + }) + + it('exposes an embedding activity context to hooks', async () => { + const onStart = vi.fn() + const adapter = new MockEmbeddingAdapter() + + await embed({ + adapter, + input: 'hello', + middleware: [{ name: 'ctx', onStart }], + }) + + const ctx = onStart.mock.calls[0]![0] + expect(ctx.activity).toBe('embedding') + expect(ctx.provider).toBe('mock') + expect(ctx.model).toBe('mock-embed') + }) + }) +}) + +// ============================================================================ +// Input resolution utilities +// ============================================================================ + +describe('embedding input utilities', () => { + const image: EmbeddingInputItem = { + type: 'image', + source: { type: 'data', value: 'aGVsbG8=', mimeType: 'image/png' }, + } + const fused: EmbeddingInputItem = { + type: 'content', + content: [ + { type: 'text', content: 'caption' }, + { + type: 'image', + source: { type: 'data', value: 'aGVsbG8=', mimeType: 'image/png' }, + }, + ], + } + + describe('resolveEmbeddingInput', () => { + it('resolves strings, text parts, image parts, and fused items in order', () => { + const resolved = resolveEmbeddingInput([ + 'plain', + { type: 'text', content: 'part' }, + image, + fused, + ]) + + expect(resolved).toHaveLength(4) + expect(resolved[0]).toEqual({ texts: ['plain'], images: [] }) + expect(resolved[1]).toEqual({ texts: ['part'], images: [] }) + expect(resolved[2]!.texts).toEqual([]) + expect(resolved[2]!.images).toHaveLength(1) + expect(resolved[3]!.texts).toEqual(['caption']) + expect(resolved[3]!.images).toHaveLength(1) + }) + }) + + describe('requireTextOnlyEmbeddingInput', () => { + it('returns plain texts for text-only input', () => { + expect( + requireTextOnlyEmbeddingInput( + ['a', { type: 'text', content: 'b' }], + 'mock', + 'mock-embed', + ), + ).toEqual(['a', 'b']) + }) + + it('throws a clear error naming the offending index for image input', () => { + expect(() => + requireTextOnlyEmbeddingInput(['a', image], 'mock', 'mock-embed'), + ).toThrow( + 'mock model "mock-embed" only supports text embedding inputs; input item at index 1 contains an image part', + ) + }) + }) + + describe('countEmbeddingInputModalities', () => { + it('counts text-only and image-carrying items', () => { + expect(countEmbeddingInputModalities(['a', image, fused])).toEqual({ + textInputCount: 1, + imageInputCount: 2, + }) + }) + }) +}) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index e006a80cf..a2d3abe2e 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1618,6 +1618,22 @@ importers: specifier: 4.0.14 version: 4.0.14(vitest@4.1.4) + packages/ai-cohere: + dependencies: + '@tanstack/ai-utils': + specifier: workspace:* + version: link:../ai-utils + devDependencies: + '@tanstack/ai': + specifier: workspace:* + version: link:../ai + '@vitest/coverage-v8': + specifier: 4.0.14 + version: 4.0.14(vitest@4.1.4) + vite: + specifier: ^7.3.3 + version: 7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2) + packages/ai-devtools: dependencies: '@tanstack/ai': diff --git a/testing/e2e/fixtures/embedding/basic.json b/testing/e2e/fixtures/embedding/basic.json new file mode 100644 index 000000000..5c71bcc32 --- /dev/null +++ b/testing/e2e/fixtures/embedding/basic.json @@ -0,0 +1,12 @@ +{ + "fixtures": [ + { + "match": { + "inputText": "a red guitar" + }, + "response": { + "embedding": [0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7, -0.8] + } + } + ] +} diff --git a/testing/e2e/global-setup.ts b/testing/e2e/global-setup.ts index 27204a5cd..5b08159ef 100644 --- a/testing/e2e/global-setup.ts +++ b/testing/e2e/global-setup.ts @@ -62,6 +62,26 @@ export default async function globalSetup() { // aimock's native Gemini handlers. mock.mount('/v1beta/models', geminiVeoMount()) + // Embedding endpoints aimock doesn't cover natively. OpenAI's + // /v1/embeddings IS covered natively (JSON fixture in fixtures/embedding/), + // but the other embedding providers need custom mounts: + // + // - Gemini: @google/genai posts to `{model}:batchEmbedContents` on the + // MLDev (API-key) surface; aimock 1.34 only handles `:embedContent`. + // Mounted on the same '/v1beta/models' prefix as the Veo mount above — + // mounts are tried in order and each returns false for paths it doesn't + // own, so both coexist and non-matching paths still fall through to + // aimock's native Gemini handlers. + // - Ollama: the ollama SDK's `embed()` (POST /api/embed) expects the batch + // `embeddings: number[][]` shape; aimock's native handler responds with + // the legacy singular `embedding` field, which crashes the adapter. + // - Mistral: the Mistral SDK Zod-validates the /v1/embeddings response and + // requires an `id` field aimock's OpenAI-format response builder omits. + // The adapter under test points its serverURL at the '/mistral' prefix. + mock.mount('/v1beta/models', geminiBatchEmbedMount()) + mock.mount('/api/embed', ollamaEmbedMount()) + mock.mount('/mistral', mistralEmbeddingsMount()) + // Gemini Omni Flash video generation (Interactions API). aimock handles // synchronous text interactions natively, but not background video jobs // (POST /v1beta/interactions with background:true → poll @@ -420,6 +440,151 @@ function geminiVeoMount(): Mountable { } } +/** + * Deterministic 8-dimension embedding vector shared by the embedding mounts + * below and mirrored by the OpenAI JSON fixture in + * `fixtures/embedding/basic.json`. The embedding spec asserts each rendered + * vector reports exactly this dimension count. + */ +const EMBED_VECTOR = [0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7, -0.8] + +/** + * Mounts Gemini's `{model}:batchEmbedContents` endpoint. The @google/genai + * SDK's `models.embedContent()` posts to `:batchEmbedContents` on the MLDev + * (API-key) surface — even for a single input — but aimock 1.34 only models + * `:embedContent`. Returns one deterministic vector per request entry, in + * the raw MLDev wire shape (`embeddings[].values`) the SDK maps to + * `EmbedContentResponse.embeddings`. + * + * Shares the '/v1beta/models' mount prefix with geminiVeoMount; non-embed + * paths return false and fall through to the Veo mount / native handlers. + */ +function geminiBatchEmbedMount(): Mountable { + return { + async handleRequest( + req: http.IncomingMessage, + res: http.ServerResponse, + // aimock strips the mount prefix ('/v1beta/models') and any query + // string, so pathname looks like '/{model}:batchEmbedContents'. + pathname: string, + ): Promise { + const match = pathname.match(/^\/([^/:]+):batchEmbedContents$/) + if (!match || req.method !== 'POST') return false + const bodyText = await readBody(req) + let count = 1 + try { + const body = JSON.parse(bodyText) as { requests?: Array } + if (Array.isArray(body.requests)) count = body.requests.length + } catch { + // Malformed body — respond with a single vector anyway. + } + res.statusCode = 200 + res.setHeader('Content-Type', 'application/json') + res.end( + JSON.stringify({ + embeddings: Array.from({ length: count }, () => ({ + values: [...EMBED_VECTOR], + })), + }), + ) + return true + }, + } +} + +/** + * Mounts Ollama's batch embed endpoint (POST /api/embed). aimock's native + * handler answers with the legacy /api/embeddings shape (singular + * `embedding: number[]`), but the ollama SDK's `embed()` — which the + * @tanstack/ai-ollama embedding adapter uses — expects + * `embeddings: number[][]` with one vector per input. + */ +function ollamaEmbedMount(): Mountable { + return { + async handleRequest( + req: http.IncomingMessage, + res: http.ServerResponse, + // aimock strips the mount prefix — pathname will be "/" for an exact match. + pathname: string, + ): Promise { + if (pathname !== '/' || req.method !== 'POST') return false + const bodyText = await readBody(req) + let model = 'nomic-embed-text' + let inputs: Array = [''] + try { + const body = JSON.parse(bodyText) as { + model?: string + input?: string | Array + } + if (body.model) model = body.model + inputs = Array.isArray(body.input) ? body.input : [body.input ?? ''] + } catch { + // Malformed body — respond with a single vector anyway. + } + res.statusCode = 200 + res.setHeader('Content-Type', 'application/json') + res.end( + JSON.stringify({ + model, + embeddings: inputs.map(() => [...EMBED_VECTOR]), + prompt_eval_count: inputs.length * 3, + }), + ) + return true + }, + } +} + +/** + * Mounts a Mistral-shaped embeddings endpoint under the '/mistral' prefix + * (the adapter under test sets serverURL to `/mistral`, so the SDK + * posts to '/mistral/v1/embeddings'). aimock's native /v1/embeddings handler + * builds an OpenAI-format response without an `id`, which fails the Mistral + * SDK's Zod response validation (`EmbeddingResponse` requires `id`). + */ +function mistralEmbeddingsMount(): Mountable { + return { + async handleRequest( + req: http.IncomingMessage, + res: http.ServerResponse, + // aimock strips the mount prefix ('/mistral'), so pathname is + // '/v1/embeddings'. + pathname: string, + ): Promise { + if (pathname !== '/v1/embeddings' || req.method !== 'POST') return false + const bodyText = await readBody(req) + let model = 'mistral-embed' + let inputs: Array = [''] + try { + const body = JSON.parse(bodyText) as { + model?: string + input?: string | Array + } + if (body.model) model = body.model + inputs = Array.isArray(body.input) ? body.input : [body.input ?? ''] + } catch { + // Malformed body — respond with a single vector anyway. + } + res.statusCode = 200 + res.setHeader('Content-Type', 'application/json') + res.end( + JSON.stringify({ + id: 'embd-e2e', + object: 'list', + model, + usage: { prompt_tokens: 6, completion_tokens: 0, total_tokens: 6 }, + data: inputs.map((_, index) => ({ + object: 'embedding', + embedding: [...EMBED_VECTOR], + index, + })), + }), + ) + return true + }, + } +} + /** * Mounts Gemini Omni Flash's Interactions-API video generation flow under a * dedicated `/omni-video` prefix (the adapter under test sets its baseUrl to diff --git a/testing/e2e/src/components/EmbeddingUI.tsx b/testing/e2e/src/components/EmbeddingUI.tsx new file mode 100644 index 000000000..75c1cb725 --- /dev/null +++ b/testing/e2e/src/components/EmbeddingUI.tsx @@ -0,0 +1,110 @@ +import { useState } from 'react' +import type { Provider } from '@/lib/types' + +interface EmbeddingUIProps { + provider: Provider + testId?: string + aimockPort?: number +} + +interface EmbeddingApiResult { + embeddings: Array> + model: string +} + +/** + * Minimal embedding harness page. `embed()` is Promise-based (no streaming + * mode), so a single fetch to the dedicated `/api/embedding` route is the + * whole flow — no connection-adapter/mode variants like the streaming + * features. The textarea takes one input text per line and the page renders + * one `embedding-vector` element per returned vector, reporting its + * dimension count. + */ +export function EmbeddingUI({ + provider, + testId, + aimockPort, +}: EmbeddingUIProps) { + const [text, setText] = useState('') + const [isLoading, setIsLoading] = useState(false) + const [error, setError] = useState(null) + const [result, setResult] = useState(null) + + const handleGenerate = async () => { + setIsLoading(true) + setError(null) + setResult(null) + try { + const texts = text + .split('\n') + .map((t) => t.trim()) + .filter((t) => t.length > 0) + const res = await fetch('/api/embedding', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ provider, texts, testId, aimockPort }), + }) + const data: unknown = await res.json() + if (!res.ok) { + const message = + data && typeof data === 'object' && 'error' in data + ? String(data.error) + : `HTTP ${res.status}` + throw new Error(message) + } + setResult(data as EmbeddingApiResult) + } catch (err) { + setError(err instanceof Error ? err.message : 'An error occurred') + } finally { + setIsLoading(false) + } + } + + return ( +
+
+