Embedding Models

/**
 * Core embedding model type for converting text to vectors
 */
type EmbeddingModelV2<VALUE> = {
  /** API specification version */
  specificationVersion: 'v2';
  /** Provider identifier (e.g., 'openai', 'cohere') */
  provider: string;
  /** Model identifier (e.g., 'text-embedding-ada-002') */
  modelId: string;
  /** Maximum number of values that can be embedded in a single call */
  maxEmbeddingsPerCall: PromiseLike<number | undefined> | number | undefined;
  /** Whether the model supports parallel embedding calls */
  supportsParallelCalls: PromiseLike<boolean> | boolean;
  
  /** Generate embeddings for the provided values */
  doEmbed(options: {
    values: Array<VALUE>;
    abortSignal?: AbortSignal;
    providerOptions?: SharedV2ProviderOptions;
    headers?: Record<string, string | undefined>;
  }): PromiseLike<EmbeddingModelV2Result>;
};

/**
 * Result from embedding model doEmbed method
 */
interface EmbeddingModelV2Result {
  /** Array of embedding vectors */
  embeddings: EmbeddingModelV2Embedding[];
  /** Token usage information */
  usage?: { tokens: number };
  /** Provider-specific metadata */
  providerMetadata?: SharedV2ProviderMetadata;
  /** Response details */
  response?: { headers?: SharedV2Headers; body?: unknown };
}

/**
 * Single embedding vector as array of numbers
 */
type EmbeddingModelV2Embedding = Array<number>;

import { EmbeddingModelV2 } from "@ai-sdk/provider";

// Basic text embedding
const model: EmbeddingModelV2<string> = provider.textEmbeddingModel('text-embedding-ada-002');

const result = await model.doEmbed({
  values: [
    'The quick brown fox jumps over the lazy dog',
    'A journey of a thousand miles begins with a single step',
    'To be or not to be, that is the question'
  ]
});

console.log('Generated', result.embeddings.length, 'embeddings');
console.log('First embedding dimensions:', result.embeddings[0].length);
console.log('Usage:', result.usage);

// Batch processing with limits
const maxPerCall = await model.maxEmbeddingsPerCall;
if (maxPerCall && texts.length > maxPerCall) {
  // Split into batches
  const batches = [];
  for (let i = 0; i < texts.length; i += maxPerCall) {
    batches.push(texts.slice(i, i + maxPerCall));
  }
  
  const allEmbeddings = [];
  for (const batch of batches) {
    const batchResult = await model.doEmbed({ values: batch });
    allEmbeddings.push(...batchResult.embeddings);
  }
}

// Calculate similarity between embeddings
function cosineSimilarity(a: number[], b: number[]): number {
  let dotProduct = 0;
  let normA = 0;
  let normB = 0;
  
  for (let i = 0; i < a.length; i++) {
    dotProduct += a[i] * b[i];
    normA += a[i] * a[i];
    normB += b[i] * b[i];
  }
  
  return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
}

const similarity = cosineSimilarity(
  result.embeddings[0],
  result.embeddings[1]
);
console.log('Similarity:', similarity);

// Check parallel support
const supportsParallel = await model.supportsParallelCalls;
if (supportsParallel) {
  // Can run multiple embedding calls in parallel
  const promises = batches.map(batch => 
    model.doEmbed({ values: batch })
  );
  const results = await Promise.all(promises);
}