Feat: embedding activities & gemini embedding adapter

.changeset/tired-spoons-nail.md

@@ -0,0 +1,29 @@
+---
+'@tanstack/tests-adapters': minor
+'@tanstack/preact-ai-devtools': minor
+'@tanstack/react-ai-devtools': minor
+'@tanstack/solid-ai-devtools': minor
+'@tanstack/smoke-tests-e2e': minor
+'@tanstack/ai-openrouter': minor
+'@tanstack/ai-anthropic': minor
+'@tanstack/ai-devtools-core': minor
+'@tanstack/ai-react-ui': minor
+'@tanstack/ai-solid-ui': minor
+'@tanstack/ai-client': minor
+'@tanstack/ai-gemini': minor
+'@tanstack/ai-ollama': minor
+'@tanstack/ai-openai': minor
+'@tanstack/ai-preact': minor
+'@tanstack/ai-svelte': minor
+'@tanstack/ai-vue-ui': minor
+'@tanstack/ai-react': minor
+'@tanstack/ai-solid': minor
+'@tanstack/ai-grok': minor
+'@tanstack/ai-vue': minor
+'ts-svelte-chat': minor
+'@tanstack/ai': minor
+'vanilla-chat': minor
+'ts-vue-chat': minor
+---
+
+Added embed/embedMany activity and Gemini embeddings adapter

packages/typescript/ai-gemini/src/adapters/embedding.ts

@@ -0,0 +1,94 @@
+import { BaseEmbeddingAdapter } from '@tanstack/ai/adapters'
+import { createGeminiClient, generateId } from '../utils'
+import {
+  validateTaskType,
+  validateValue,
+} from '../embedding/embedding-provider-options'
+import type { GoogleGenAI } from '@google/genai'
+import type {
+  EmbedManyOptions,
+  EmbedManyResult,
+  EmbedOptions,
+  EmbedResult,
+} from '@tanstack/ai'
+import type { GEMINI_EMBEDDING_MODELS } from '../model-meta'
+import type { GeminiClientConfig } from '../utils'
+import type {
+  GeminiEmbeddingModelProviderOptionsByName,
+  GeminiEmbeddingProviderOptions,
+} from '../embedding/embedding-provider-options'
+
+/**
+ * Configuration for Gemini embedding adapter
+ */
+export interface GeminiEmbeddingConfig extends GeminiClientConfig {}
+
+export type GeminiEmbeddingModel = (typeof GEMINI_EMBEDDING_MODELS)[number]
+
+export class GeminiEmbeddingAdapter<
+  TModel extends GeminiEmbeddingModel,
+> extends BaseEmbeddingAdapter<TModel, GeminiEmbeddingProviderOptions> {
+  readonly kind = 'embedding' as const
+  readonly name = 'gemini' as const
+
+  // Type-only property - never assigned at runtime
+  declare '~types': {
+    providerOptions: GeminiEmbeddingProviderOptions
+    modelProviderOptionsByName: GeminiEmbeddingModelProviderOptionsByName
+  }
+
+  private client: GoogleGenAI
+
+  constructor(config: GeminiEmbeddingConfig, model: TModel) {
+    super({}, model)
+    this.client = createGeminiClient(config)
+  }
+
+  async embed(
+    options: EmbedOptions<GeminiEmbeddingProviderOptions>,
+  ): Promise<EmbedResult> {
+    const { model, value, modelOptions } = options
+
+    validateValue({ value, model })
+    validateTaskType({ taskType: modelOptions?.taskType, model })
+
+    const { embeddings } = await this.client.models.embedContent({
+      model,
+      contents: value,
+      config: {
+        ...modelOptions,
+      },
+    })
+
+    return {
+      embedding: embeddings?.[0]?.values || [],
+      id: generateId(this.name),
+      model,
+      usage: undefined,
+    }
+  }
+
+  async embedMany(
+    options: EmbedManyOptions<GeminiEmbeddingProviderOptions>,
+  ): Promise<EmbedManyResult> {
+    const { model, values, modelOptions } = options
+
+    validateValue({ value: values, model })
+    validateTaskType({ taskType: modelOptions?.taskType, model })
+
+    const { embeddings } = await this.client.models.embedContent({
+      model,
+      contents: values,
+      config: {
+        ...modelOptions,
+      },
+    })
+
+    return {
+      embeddings: embeddings?.map((embedding) => embedding.values || []) || [],
+      id: generateId(this.name),
+      model,
+      usage: undefined,
+    }
+  }
+}

packages/typescript/ai-gemini/src/embedding/embedding-provider-options.ts

@@ -0,0 +1,88 @@
+import type { HttpOptions } from '@google/genai'
+import type { GeminiEmbeddingModels } from '../model-meta'
+
+const VALID_TASK_TYPES = new Set([
+  'SEMANTIC_SIMILARITY',
+  'CLASSIFICATION',
+  'CLUSTERING',
+  'RETRIEVAL_DOCUMENT',
+  'RETRIEVAL_QUERY',
+  'CODE_RETRIEVAL_QUERY',
+  'QUESTION_ANSWERING',
+  'FACT_VERIFICATION',
+])
+
+type TaskType =
+  | 'SEMANTIC_SIMILARITY'
+  | 'CLASSIFICATION'
+  | 'CLUSTERING'
+  | 'RETRIEVAL_DOCUMENT'
+  | 'RETRIEVAL_QUERY'
+  | 'CODE_RETRIEVAL_QUERY'
+  | 'QUESTION_ANSWERING'
+  | 'FACT_VERIFICATION'
+
+export interface GeminiEmbeddingProviderOptions {
+  /** Used to override HTTP request options. */
+  httpOptions?: HttpOptions
+  /**
+   * Type of task for which the embedding will be used.
+   */
+  taskType?: TaskType
+  /**
+   * Title for the text. Only applicable when TaskType is `RETRIEVAL_DOCUMENT`.
+   */
+  title?: string
+  /**
+   * Reduced dimension for the output embedding. If set,
+   * excessive values in the output embedding are truncated from the end.
+   * Supported by newer models since 2024 only. You cannot set this value if
+   * using the earlier model (`models/embedding-001`).
+   */
+  outputDimensionality?: number
+}
+
+export type GeminiEmbeddingModelProviderOptionsByName = {
+  [K in GeminiEmbeddingModels]: GeminiEmbeddingProviderOptions
+}
+
+/**
+ * Validates the task type
+ */
+export function validateTaskType(options: {
+  taskType: TaskType | undefined
+  model: string
+}) {
+  const { taskType, model } = options
+  if (!taskType) return
+
+  if (!VALID_TASK_TYPES.has(taskType)) {
+    throw new Error(`Invalid task type "${taskType}" for model "${model}".`)
+  }
+}
+
+/**
+ * Validates the value to embed is not empty
+ */
+export function validateValue(options: {
+  value: string | Array<string>
+  model: string
+}): void {
+  const { value, model } = options
+  if (Array.isArray(value)) {
+    if (value.length === 0) {
+      throw new Error(`Value array cannot be empty for model "${model}".`)
+    }
+    for (const v of value) {
+      if (!v || v.trim().length === 0) {
+        throw new Error(
+          `Value array cannot contain empty values for model "${model}".`,
+        )
+      }
+    }
+  } else {
+    if (!value || value.trim().length === 0) {
+      throw new Error(`Value cannot be empty for model "${model}".`)
+    }
+  }
+}

packages/typescript/ai-gemini/src/model-meta.ts

@@ -12,7 +12,7 @@ interface ModelMeta<TProviderOptions = unknown> {
   name: string
   supports: {
     input: Array<'text' | 'image' | 'audio' | 'video' | 'document'>
-    output: Array<'text' | 'image' | 'audio' | 'video'>
+    output: Array<'text' | 'image' | 'audio' | 'video' | 'embedding'>
     capabilities?: Array<
       | 'audio_generation'
       | 'batch_api'
@@ -678,6 +678,28 @@ const IMAGEN_3 = {
     GeminiCommonConfigOptions &
     GeminiCachedContentOptions
 >
+
+const GEMINI_EMBEDDING_001 = {
+  name: 'embedding-001',
+  max_input_tokens: 2048,
+  supports: {
+    input: ['text'],
+    output: ['embedding'],
+  },
+  pricing: {
+    input: {
+      normal: 0,
+    },
+    output: {
+      normal: 0.15,
+    },
+  },
+} as const satisfies ModelMeta<
+  GeminiToolConfigOptions &
+    GeminiSafetyOptions &
+    GeminiCommonConfigOptions &
+    GeminiCachedContentOptions
+>
 /** 
 const VEO_3_1_PREVIEW = {
   name: 'veo-3.1-generate-preview',
@@ -833,6 +855,10 @@ export const GEMINI_MODELS = [
 
 export type GeminiModels = (typeof GEMINI_MODELS)[number]
 
+export const GEMINI_EMBEDDING_MODELS = [GEMINI_EMBEDDING_001.name] as const
+
+export type GeminiEmbeddingModels = (typeof GEMINI_EMBEDDING_MODELS)[number]
+
 export type GeminiImageModels = (typeof GEMINI_IMAGE_MODELS)[number]
 
 export const GEMINI_IMAGE_MODELS = [

packages/typescript/ai/src/activities/embed/adapter.ts

@@ -0,0 +1,91 @@
+import type {
+  EmbedManyOptions,
+  EmbedManyResult,
+  EmbedOptions,
+  EmbedResult,
+} from '../../types'
+
+// ===========================
+// Activity Kind
+// ===========================
+
+/** The adapter kind this activity handles */
+export const kind = 'embedding' as const
+
+export interface EmbeddingAdapterConfig {
+  apiKey?: string
+  baseUrl?: string
+  timeout?: number
+  maxRetries?: number
+  headers?: Record<string, string>
+}
+
+export interface EmbeddingAdapter<
+  TModel extends string = string,
+  TProviderOptions extends object = Record<string, unknown>,
+> {
+  /** Discriminator for adapter kind - used to determine API shape */
+  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
+  }
+
+  /**
+   * Generate embeddings for text
+   */
+  embed: (options: EmbedOptions<TProviderOptions>) => Promise<EmbedResult>
+
+  /**
+   * Generate embeddings for multiple texts
+   */
+  embedMany: (
+    options: EmbedManyOptions<TProviderOptions>,
+  ) => Promise<EmbedManyResult>
+}
+
+/**
+ * A EmbeddingAdapter with any/unknown type parameters.
+ * Useful as a constraint in generic functions and interfaces.
+ */
+export type AnyEmbeddingAdapter = EmbeddingAdapter<any, any>
+
+/**
+ * Abstract base class for embed adapters.
+ * Extend this class to implement an embed adapter for a specific provider.
+ *
+ * Generic parameters match EmbedAdapter - all pre-resolved by the provider function.
+ */
+export abstract class BaseEmbeddingAdapter<
+  TModel extends string = string,
+  TProviderOptions extends object = Record<string, unknown>,
+> implements EmbeddingAdapter<TModel, TProviderOptions> {
+  readonly kind = 'embedding' as const
+  abstract readonly name: string
+  readonly model: TModel
+
+  // Type-only property - never assigned at runtime
+  declare '~types': {
+    providerOptions: TProviderOptions
+  }
+
+  protected config: EmbeddingAdapterConfig
+
+  constructor(config: EmbeddingAdapterConfig = {}, model: TModel) {
+    this.config = config
+    this.model = model
+  }
+
+  abstract embed(options: EmbedOptions<TProviderOptions>): Promise<EmbedResult>
+
+  abstract embedMany(
+    options: EmbedManyOptions<TProviderOptions>,
+  ): Promise<EmbedManyResult>
+}