@od-oneapp/ai-platform 0.1.0

package/src/grounding/embed/index.ts

@@ -0,0 +1,704 @@
+/**
+ * Embeddings & RAG Module
+ *
+ * Retrieval-Augmented Generation patterns:
+ * - Text chunking strategies
+ * - Embedding generation
+ * - Similarity search
+ * - RAG pipeline orchestration
+ *
+ * @module embeddings
+ */
+
+// -----------------------------------------------------------------------------
+// Types
+// -----------------------------------------------------------------------------
+
+/** Text chunk with metadata */
+export interface TextChunk {
+  readonly id: string;
+  readonly content: string;
+  readonly metadata: ChunkMetadata;
+  embedding?: number[];
+}
+
+/** Chunk metadata */
+export interface ChunkMetadata {
+  readonly sourceId?: string;
+  readonly sourceName?: string;
+  readonly startIndex?: number;
+  readonly endIndex?: number;
+  readonly chunkIndex?: number;
+  readonly totalChunks?: number;
+  readonly [key: string]: unknown;
+}
+
+/** Chunking strategy */
+export type ChunkingStrategy = 'fixed' | 'sentence' | 'paragraph' | 'semantic';
+
+/** Chunking configuration */
+export interface ChunkingConfig {
+  readonly strategy: ChunkingStrategy;
+  readonly chunkSize: number;
+  readonly chunkOverlap: number;
+  readonly minChunkSize?: number;
+}
+
+/** Embedding model interface */
+export interface EmbeddingModel {
+  embed: (text: string) => Promise<number[]>;
+  embedBatch: (texts: string[]) => Promise<number[][]>;
+  readonly dimensions: number;
+  readonly modelId: string;
+}
+
+/** Search result */
+export interface SearchResult {
+  readonly chunk: TextChunk;
+  readonly score: number;
+  readonly rank: number;
+}
+
+/** RAG context */
+export interface RAGContext {
+  readonly query: string;
+  readonly results: SearchResult[];
+  readonly totalTokens: number;
+  formattedContext: string;
+}
+
+/** Vector store interface */
+export interface VectorStore {
+  upsert: (chunks: TextChunk[]) => Promise<void>;
+  search: (embedding: number[], options?: SearchOptions) => Promise<SearchResult[]>;
+  delete: (ids: string[]) => Promise<void>;
+  clear: () => Promise<void>;
+}
+
+/** Search options */
+export interface SearchOptions {
+  readonly topK?: number;
+  readonly minScore?: number;
+  readonly filter?: Record<string, unknown>;
+}
+
+// -----------------------------------------------------------------------------
+// Text Chunking
+// -----------------------------------------------------------------------------
+
+/**
+ * Creates a text chunker with configurable strategy.
+ *
+ * @example
+ * ```typescript
+ * const chunker = createChunker({
+ *   strategy: 'paragraph',
+ *   chunkSize: 500,
+ *   chunkOverlap: 50
+ * });
+ *
+ * const chunks = chunker.chunk(longText, { sourceId: 'doc-1' });
+ * ```
+ */
+export function createChunker(config: ChunkingConfig): Chunker {
+  const { strategy, chunkSize, chunkOverlap, minChunkSize = 50 } = config;
+
+  function chunk(text: string, metadata?: Partial<ChunkMetadata>): TextChunk[] {
+    const chunks: TextChunk[] = [];
+    let segments: string[];
+
+    switch (strategy) {
+      case 'sentence':
+        segments = splitBySentence(text);
+        break;
+      case 'paragraph':
+        segments = splitByParagraph(text);
+        break;
+      case 'semantic':
+        segments = splitSemantic(text, chunkSize);
+        break;
+      case 'fixed':
+        segments = splitFixed(text, chunkSize, chunkOverlap);
+        break;
+      default: {
+        const exhaustiveCheck: never = strategy;
+        throw new Error(`Unsupported chunking strategy: ${exhaustiveCheck}`);
+      }
+    }
+
+    // Combine small segments and split large ones
+    let currentChunk = '';
+    let startIndex = 0;
+    let chunkIndex = 0;
+
+    for (let i = 0; i < segments.length; i++) {
+      const segment = segments[i] ?? '';
+
+      if (currentChunk.length + segment.length <= chunkSize) {
+        currentChunk += (currentChunk ? ' ' : '') + segment;
+      } else {
+        if (currentChunk.length >= minChunkSize) {
+          chunks.push(createChunk(currentChunk, chunkIndex++, startIndex, metadata));
+          startIndex += currentChunk.length;
+        }
+
+        // Handle segments larger than chunkSize
+        if (segment.length > chunkSize) {
+          const subChunks = splitFixed(segment, chunkSize, chunkOverlap);
+          for (const sub of subChunks) {
+            if (sub.length >= minChunkSize) {
+              chunks.push(createChunk(sub, chunkIndex++, startIndex, metadata));
+              startIndex += sub.length;
+            }
+          }
+          currentChunk = '';
+        } else {
+          currentChunk = segment;
+        }
+      }
+    }
+
+    // Add remaining chunk
+    if (currentChunk.length >= minChunkSize) {
+      chunks.push(createChunk(currentChunk, chunkIndex++, startIndex, metadata));
+    }
+
+    // Update total chunks
+    const totalChunks = chunks.length;
+    return chunks.map(c => ({
+      ...c,
+      metadata: { ...c.metadata, totalChunks },
+    }));
+  }
+
+  function createChunk(
+    content: string,
+    chunkIndex: number,
+    startIndex: number,
+    metadata?: Partial<ChunkMetadata>,
+  ): TextChunk {
+    // Add randomness to prevent ID collisions across concurrent calls
+    const uniqueSuffix = Math.random().toString(36).slice(2, 8);
+    return {
+      id: `chunk-${Date.now()}-${chunkIndex}-${uniqueSuffix}`,
+      content: content.trim(),
+      metadata: {
+        ...metadata,
+        chunkIndex,
+        startIndex,
+        endIndex: startIndex + content.length,
+      },
+    };
+  }
+
+  return { chunk, config };
+}
+
+export interface Chunker {
+  chunk: (text: string, metadata?: Partial<ChunkMetadata>) => TextChunk[];
+  config: ChunkingConfig;
+}
+
+// Chunking helpers
+function splitFixed(text: string, size: number, overlap: number): string[] {
+  const chunks: string[] = [];
+
+  // Guard against invalid overlap that would cause infinite loop
+  const effectiveOverlap = Math.min(overlap, size - 1);
+  const step = Math.max(1, size - effectiveOverlap);
+
+  if (size <= 0) {
+    return text ? [text] : [];
+  }
+
+  let start = 0;
+
+  while (start < text.length) {
+    chunks.push(text.slice(start, start + size));
+    start += step;
+  }
+
+  return chunks;
+}
+
+function splitBySentence(text: string): string[] {
+  return text.split(/(?<=[.!?])\s+/).filter(s => s.trim().length > 0);
+}
+
+function splitByParagraph(text: string): string[] {
+  return text.split(/\n\n+/).filter(p => p.trim().length > 0);
+}
+
+function splitSemantic(text: string, targetSize: number): string[] {
+  // Simplified semantic splitting - prefer sentence boundaries within target size
+  const sentences = splitBySentence(text);
+  const result: string[] = [];
+  let current = '';
+
+  for (const sentence of sentences) {
+    if (current.length + sentence.length <= targetSize * 1.2) {
+      current += (current ? ' ' : '') + sentence;
+    } else {
+      if (current) result.push(current);
+      current = sentence;
+    }
+  }
+
+  if (current) result.push(current);
+  return result;
+}
+
+// -----------------------------------------------------------------------------
+// Embedding Generator
+// -----------------------------------------------------------------------------
+
+/**
+ * Creates an embedding generator wrapper.
+ *
+ * @example
+ * ```typescript
+ * const embedder = createEmbedder({
+ *   model: openaiEmbedding,
+ *   batchSize: 100,
+ *   cache: embeddingCache
+ * });
+ *
+ * const chunks = await embedder.embedChunks(textChunks);
+ * ```
+ */
+export function createEmbedder(options: {
+  model: EmbeddingModel;
+  batchSize?: number;
+  cache?: Map<string, number[]>;
+  onProgress?: (processed: number, total: number) => void;
+}): Embedder {
+  const { model, batchSize = 100, cache, onProgress } = options;
+
+  async function embed(text: string): Promise<number[]> {
+    const cacheKey = hashText(text);
+
+    const cached = cache?.get(cacheKey);
+    if (cached) {
+      return cached;
+    }
+
+    const embedding = await model.embed(text);
+
+    if (cache) {
+      cache.set(cacheKey, embedding);
+    }
+
+    return embedding;
+  }
+
+  async function embedBatch(texts: string[]): Promise<number[][]> {
+    const results: number[][] = new Array(texts.length);
+    const toEmbed: { index: number; text: string }[] = [];
+
+    // Check cache first
+    texts.forEach((text, index) => {
+      const cacheKey = hashText(text);
+      const cached = cache?.get(cacheKey);
+      if (cached) {
+        results[index] = cached;
+      } else {
+        toEmbed.push({ index, text });
+      }
+    });
+
+    // Embed uncached in batches
+    for (let i = 0; i < toEmbed.length; i += batchSize) {
+      const batch = toEmbed.slice(i, i + batchSize);
+      const embeddings = await model.embedBatch(batch.map(b => b.text));
+
+      batch.forEach((item, j) => {
+        const embedding = embeddings[j];
+        if (embedding) {
+          results[item.index] = embedding;
+          if (cache) {
+            cache.set(hashText(item.text), embedding);
+          }
+        }
+      });
+
+      onProgress?.(Math.min(i + batchSize, toEmbed.length), toEmbed.length);
+    }
+
+    // Validate all embeddings were populated (no sparse array holes)
+    for (let i = 0; i < results.length; i++) {
+      if (!results[i]) {
+        throw new Error(`Embedding failed for text at index ${i}: incomplete batch response`);
+      }
+    }
+
+    return results;
+  }
+
+  async function embedChunks(chunks: TextChunk[]): Promise<TextChunk[]> {
+    const texts = chunks.map(c => c.content);
+    const embeddings = await embedBatch(texts);
+
+    return chunks.map((chunk, i) => ({
+      ...chunk,
+      embedding: embeddings[i],
+    }));
+  }
+
+  return {
+    embed,
+    embedBatch,
+    embedChunks,
+    dimensions: model.dimensions,
+    modelId: model.modelId,
+  };
+}
+
+export interface Embedder {
+  embed: (text: string) => Promise<number[]>;
+  embedBatch: (texts: string[]) => Promise<number[][]>;
+  embedChunks: (chunks: TextChunk[]) => Promise<TextChunk[]>;
+  readonly dimensions: number;
+  readonly modelId: string;
+}
+
+/**
+ * Simple djb2-style hash for cache keys.
+ * Note: For production caches with millions of entries, consider using
+ * a stronger hash (SHA-256) to reduce collision probability.
+ */
+function hashText(text: string): string {
+  let hash = 0;
+  for (let i = 0; i < text.length; i++) {
+    const char = text.charCodeAt(i);
+    hash = (hash << 5) - hash + char;
+    hash = hash & hash;
+  }
+  return `embed-${hash}`;
+}
+
+// -----------------------------------------------------------------------------
+// Similarity Search
+// -----------------------------------------------------------------------------
+
+/**
+ * Calculates cosine similarity between two vectors.
+ */
+export function cosineSimilarity(a: number[], b: number[]): number {
+  if (a.length !== b.length) {
+    throw new Error('Vectors must have same dimensions');
+  }
+
+  let dotProduct = 0;
+  let normA = 0;
+  let normB = 0;
+
+  for (let i = 0; i < a.length; i++) {
+    const aVal = a[i] ?? 0;
+    const bVal = b[i] ?? 0;
+    dotProduct += aVal * bVal;
+    normA += aVal * aVal;
+    normB += bVal * bVal;
+  }
+
+  if (normA === 0 || normB === 0) return 0;
+  return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
+}
+
+/**
+ * Creates an in-memory vector store.
+ *
+ * @example
+ * ```typescript
+ * const store = createInMemoryVectorStore();
+ *
+ * await store.upsert(embeddedChunks);
+ * const results = await store.search(queryEmbedding, { topK: 5 });
+ * ```
+ */
+export function createInMemoryVectorStore(): VectorStore {
+  const chunks = new Map<string, TextChunk>();
+
+  async function upsert(newChunks: TextChunk[]): Promise<void> {
+    for (const chunk of newChunks) {
+      chunks.set(chunk.id, chunk);
+    }
+  }
+
+  async function search(embedding: number[], options?: SearchOptions): Promise<SearchResult[]> {
+    const { topK = 10, minScore = 0 } = options ?? {};
+
+    const results: SearchResult[] = [];
+
+    for (const chunk of chunks.values()) {
+      if (!chunk.embedding) continue;
+
+      const score = cosineSimilarity(embedding, chunk.embedding);
+
+      if (score >= minScore) {
+        results.push({ chunk, score, rank: 0 });
+      }
+    }
+
+    // Sort by score descending
+    results.sort((a, b) => b.score - a.score);
+
+    // Take top K and assign ranks
+    return results.slice(0, topK).map((r, i) => ({ ...r, rank: i + 1 }));
+  }
+
+  async function deleteChunks(ids: string[]): Promise<void> {
+    for (const id of ids) {
+      chunks.delete(id);
+    }
+  }
+
+  async function clear(): Promise<void> {
+    chunks.clear();
+  }
+
+  return {
+    upsert,
+    search,
+    delete: deleteChunks,
+    clear,
+  };
+}
+
+// -----------------------------------------------------------------------------
+// RAG Pipeline
+// -----------------------------------------------------------------------------
+
+/**
+ * Creates a RAG pipeline for retrieval-augmented generation.
+ *
+ * @example
+ * ```typescript
+ * const rag = createRAGPipeline({
+ *   embedder,
+ *   vectorStore,
+ *   topK: 5,
+ *   maxTokens: 2000
+ * });
+ *
+ * const context = await rag.retrieve('What is the capital of France?');
+ * const prompt = rag.buildPrompt(context, 'Answer the question based on context.');
+ * ```
+ */
+export function createRAGPipeline(options: {
+  embedder: Embedder;
+  vectorStore: VectorStore;
+  topK?: number;
+  maxTokens?: number;
+  minScore?: number;
+}): RAGPipeline {
+  const { embedder, vectorStore, topK = 5, maxTokens = 2000, minScore = 0.5 } = options;
+
+  async function ingest(
+    documents: { id: string; content: string; metadata?: Record<string, unknown> }[],
+    chunker: Chunker,
+  ): Promise<void> {
+    const allChunks: TextChunk[] = [];
+
+    for (const doc of documents) {
+      const chunks = chunker.chunk(doc.content, {
+        sourceId: doc.id,
+        ...doc.metadata,
+      });
+      allChunks.push(...chunks);
+    }
+
+    const embeddedChunks = await embedder.embedChunks(allChunks);
+    await vectorStore.upsert(embeddedChunks);
+  }
+
+  async function retrieve(query: string): Promise<RAGContext> {
+    const queryEmbedding = await embedder.embed(query);
+    const results = await vectorStore.search(queryEmbedding, { topK, minScore });
+
+    // Build formatted context
+    let formattedContext = '';
+    let totalTokens = 0;
+    const estimateTokens = (text: string) => Math.ceil(text.length / 4);
+
+    for (const result of results) {
+      const chunkText = result.chunk.content;
+      const chunkTokens = estimateTokens(chunkText);
+
+      if (totalTokens + chunkTokens > maxTokens) break;
+
+      formattedContext +=
+        `[Source: ${result.chunk.metadata.sourceId ?? 'unknown'}, ` +
+        `Relevance: ${result.score.toFixed(2)}]\n` +
+        `${chunkText}\n\n`;
+      totalTokens += chunkTokens;
+    }
+
+    return {
+      query,
+      results,
+      totalTokens,
+      formattedContext: formattedContext.trim(),
+    };
+  }
+
+  function buildPrompt(context: RAGContext, systemPrompt?: string): string {
+    const system = systemPrompt ?? 'Answer the question based on the provided context.';
+
+    return `${system}
+
+Context:
+${context.formattedContext}
+
+Question: ${context.query}
+
+Answer:`;
+  }
+
+  return {
+    ingest,
+    retrieve,
+    buildPrompt,
+  };
+}
+
+export interface RAGPipeline {
+  ingest: (
+    documents: { id: string; content: string; metadata?: Record<string, unknown> }[],
+    chunker: Chunker,
+  ) => Promise<void>;
+  retrieve: (query: string) => Promise<RAGContext>;
+  buildPrompt: (context: RAGContext, systemPrompt?: string) => string;
+}
+
+// -----------------------------------------------------------------------------
+// Reranking
+// -----------------------------------------------------------------------------
+
+export {
+  // Reranker factories
+  createCohereReranker,
+  createCrossEncoderReranker,
+  // Pipelines and utilities
+  createRerankPipeline,
+  reciprocalRankFusion,
+  vectorResultsToDocuments,
+  calculateNDCG,
+  calculateMRR,
+  // Types
+  type RerankDocument,
+  type RerankResult,
+  type RerankOptions,
+  type RerankProvider,
+  type CohereRerankConfig,
+  type CrossEncoderConfig,
+  type RerankPipelineConfig,
+} from './reranking';
+
+// -----------------------------------------------------------------------------
+// AI SDK Quick Embed Helpers (migrated from @repo/ai)
+// -----------------------------------------------------------------------------
+
+import {
+  embed as sdkEmbed,
+  embedMany as sdkEmbedMany,
+  cosineSimilarity as sdkCosineSimilarity,
+} from '../../shared';
+
+import type { SDKEmbeddingModel } from '../../shared';
+
+/**
+ * Quick single-text embedding using AI SDK's embed function.
+ *
+ * @param model - AI SDK embedding model instance
+ * @param value - Text to embed
+ * @param options - Optional configuration
+ * @returns Embed result from AI SDK
+ *
+ * @example
+ * ```ts
+ * import { openai } from '@ai-sdk/openai';
+ * import { embedText } from '@od-oneapp/ai-platform/grounding/embed';
+ *
+ * const result = await embedText(
+ *   openai.embedding('text-embedding-3-small'),
+ *   'Hello world',
+ * );
+ * console.log(result.embedding); // number[]
+ * ```
+ */
+export async function embedText(
+  model: SDKEmbeddingModel,
+  value: string,
+  options?: { maxRetries?: number },
+) {
+  return sdkEmbed({
+    model,
+    value,
+    maxRetries: options?.maxRetries ?? 2,
+  });
+}
+
+/**
+ * Quick batch embedding using AI SDK's embedMany function.
+ *
+ * @param model - AI SDK embedding model instance
+ * @param values - Array of texts to embed
+ * @param options - Optional configuration
+ * @returns EmbedMany result from AI SDK
+ *
+ * @example
+ * ```ts
+ * import { openai } from '@ai-sdk/openai';
+ * import { embedBatch } from '@od-oneapp/ai-platform/grounding/embed';
+ *
+ * const result = await embedBatch(
+ *   openai.embedding('text-embedding-3-small'),
+ *   ['Hello', 'World'],
+ * );
+ * console.log(result.embeddings); // number[][]
+ * ```
+ */
+export async function embedBatch(
+  model: SDKEmbeddingModel,
+  values: string[],
+  options?: { maxParallelCalls?: number },
+) {
+  return sdkEmbedMany({
+    model,
+    values,
+    maxParallelCalls: options?.maxParallelCalls ?? 2,
+  });
+}
+
+/**
+ * Calculate similarity between two texts by embedding them and computing cosine similarity.
+ *
+ * @param model - AI SDK embedding model instance
+ * @param text1 - First text
+ * @param text2 - Second text
+ * @returns Cosine similarity score between -1 and 1
+ *
+ * @example
+ * ```ts
+ * import { openai } from '@ai-sdk/openai';
+ * import { calculateTextSimilarity } from '@od-oneapp/ai-platform/grounding/embed';
+ *
+ * const score = await calculateTextSimilarity(
+ *   openai.embedding('text-embedding-3-small'),
+ *   'The cat sat on the mat',
+ *   'A feline rested on a rug',
+ * );
+ * console.log(score); // 0.85 (high similarity)
+ * ```
+ */
+export async function calculateTextSimilarity(
+  model: SDKEmbeddingModel,
+  text1: string,
+  text2: string,
+): Promise<number> {
+  const [{ embedding: e1 }, { embedding: e2 }] = await Promise.all([
+    embedText(model, text1),
+    embedText(model, text2),
+  ]);
+  return sdkCosineSimilarity(e1, e2);
+}