@od-oneapp/ai-platform 0.1.0

package/src/grounding/embed/reranking.ts

@@ -0,0 +1,604 @@
+/**
+ * Reranking Module - AI SDK v6 Compatible
+ *
+ * Provides document reranking capabilities for RAG pipelines using
+ * Cohere, cross-encoders, and other reranking providers.
+ *
+ * @module @od-oneapp/ai-platform/grounding/embed/reranking
+ */
+
+import { logWarn } from '@repo/shared/logs';
+import { z } from 'zod';
+
+// ============================================================================
+// Types
+// ============================================================================
+
+/**
+ * A document to be reranked
+ */
+export interface RerankDocument {
+  /** Unique identifier for the document */
+  id: string;
+  /** The text content to rerank */
+  text: string;
+  /** Optional metadata */
+  metadata?: Record<string, unknown>;
+  /** Original relevance score (e.g., from vector search) */
+  originalScore?: number;
+}
+
+/**
+ * Result of reranking a document
+ */
+export interface RerankResult {
+  /** The original document */
+  document: RerankDocument;
+  /** New relevance score from reranker (0-1) */
+  relevanceScore: number;
+  /** Position in original results */
+  originalIndex: number;
+  /** Position after reranking */
+  newIndex: number;
+}
+
+/**
+ * Options for reranking
+ */
+export interface RerankOptions {
+  /** The query to rerank against */
+  query: string;
+  /** Documents to rerank */
+  documents: RerankDocument[];
+  /** Maximum number of results to return */
+  topK?: number;
+  /** Minimum relevance score threshold (0-1) */
+  minScore?: number;
+  /** Whether to return documents below threshold */
+  returnAllResults?: boolean;
+  /** Model to use for reranking */
+  model?: string;
+}
+
+/**
+ * Reranking provider interface
+ */
+export interface RerankProvider {
+  /** Provider name */
+  readonly name: string;
+  /** Rerank documents */
+  rerank(options: RerankOptions): Promise<RerankResult[]>;
+  /** Check if provider is available */
+  isAvailable(): Promise<boolean>;
+}
+
+/**
+ * Cohere reranker configuration
+ */
+export interface CohereRerankConfig {
+  apiKey: string;
+  model?: 'rerank-english-v3.0' | 'rerank-multilingual-v3.0' | 'rerank-english-v2.0';
+  baseUrl?: string;
+  maxRetries?: number;
+  timeout?: number;
+}
+
+/**
+ * Cross-encoder reranker configuration
+ */
+export interface CrossEncoderConfig {
+  modelName: string;
+  endpoint: string;
+  apiKey?: string;
+  batchSize?: number;
+  timeout?: number;
+}
+
+// ============================================================================
+// Schemas
+// ============================================================================
+
+export const RerankDocumentSchema = z.object({
+  id: z.string(),
+  text: z.string(),
+  metadata: z.record(z.string(), z.unknown()).optional(),
+  originalScore: z.number().optional(),
+});
+
+export const RerankOptionsSchema = z.object({
+  query: z.string().min(1),
+  documents: z.array(RerankDocumentSchema).min(1),
+  topK: z.number().int().positive().optional(),
+  minScore: z.number().min(0).max(1).optional(),
+  returnAllResults: z.boolean().optional(),
+  model: z.string().optional(),
+});
+
+// ============================================================================
+// Cohere Reranker
+// ============================================================================
+
+/**
+ * Create a Cohere reranker provider
+ *
+ * @example
+ * ```typescript
+ * const reranker = createCohereReranker({
+ *   apiKey: process.env.COHERE_API_KEY!,
+ *   model: 'rerank-english-v3.0',
+ * });
+ *
+ * const results = await reranker.rerank({
+ *   query: 'What is machine learning?',
+ *   documents: docs,
+ *   topK: 5,
+ * });
+ * ```
+ */
+export function createCohereReranker(config: CohereRerankConfig): RerankProvider {
+  const {
+    apiKey,
+    model = 'rerank-english-v3.0',
+    baseUrl = 'https://api.cohere.ai/v1',
+    maxRetries = 3,
+    timeout = 30000,
+  } = config;
+
+  return {
+    name: 'cohere',
+
+    async rerank(options: RerankOptions): Promise<RerankResult[]> {
+      const { query, documents, topK, minScore = 0, returnAllResults = false } = options;
+
+      if (documents.length === 0) {
+        return [];
+      }
+
+      // Prepare documents for Cohere API
+      const docTexts = documents.map(d => d.text);
+
+      let lastError: Error | null = null;
+      for (let attempt = 0; attempt < maxRetries; attempt++) {
+        try {
+          const controller = new AbortController();
+          const timeoutId = setTimeout(() => controller.abort(), timeout);
+
+          const response = await fetch(`${baseUrl}/rerank`, {
+            method: 'POST',
+            headers: {
+              Authorization: `Bearer ${apiKey}`,
+              'Content-Type': 'application/json',
+            },
+            body: JSON.stringify({
+              query,
+              documents: docTexts,
+              model: options.model || model,
+              top_n: topK || documents.length,
+              return_documents: false,
+            }),
+            signal: controller.signal,
+          });
+
+          clearTimeout(timeoutId);
+
+          if (!response.ok) {
+            const errorText = await response.text();
+            throw new Error(`Cohere API error: ${response.status} - ${errorText}`);
+          }
+
+          const data = (await response.json()) as {
+            results: Array<{ index: number; relevance_score: number }>;
+          };
+
+          // Map results back to documents
+          const results: RerankResult[] = data.results
+            .map((r, newIndex) => {
+              const document = documents[r.index];
+              if (!document) return null;
+              return {
+                document,
+                relevanceScore: r.relevance_score,
+                originalIndex: r.index,
+                newIndex,
+              };
+            })
+            .filter(
+              (entry): entry is RerankResult =>
+                entry !== null && (returnAllResults || entry.relevanceScore >= minScore),
+            );
+
+          return topK ? results.slice(0, topK) : results;
+        } catch (error) {
+          lastError = error as Error;
+          if (attempt < maxRetries - 1) {
+            // Exponential backoff
+            await new Promise(resolve => setTimeout(resolve, Math.pow(2, attempt) * 1000));
+          }
+        }
+      }
+
+      throw lastError || new Error('Cohere rerank failed after retries');
+    },
+
+    async isAvailable(): Promise<boolean> {
+      try {
+        const response = await fetch(`${baseUrl}/check-api-key`, {
+          method: 'POST',
+          headers: {
+            Authorization: `Bearer ${apiKey}`,
+            'Content-Type': 'application/json',
+          },
+        });
+        return response.ok;
+      } catch {
+        return false;
+      }
+    },
+  };
+}
+
+// ============================================================================
+// Cross-Encoder Reranker
+// ============================================================================
+
+/**
+ * Create a cross-encoder reranker (for self-hosted models)
+ *
+ * @example
+ * ```typescript
+ * const reranker = createCrossEncoderReranker({
+ *   modelName: 'cross-encoder/ms-marco-MiniLM-L-6-v2',
+ *   endpoint: 'http://localhost:8080/rerank',
+ * });
+ * ```
+ */
+export function createCrossEncoderReranker(config: CrossEncoderConfig): RerankProvider {
+  const { modelName, endpoint, apiKey, batchSize = 32, timeout = 60000 } = config;
+
+  return {
+    name: `cross-encoder:${modelName}`,
+
+    async rerank(options: RerankOptions): Promise<RerankResult[]> {
+      const { query, documents, topK, minScore = 0, returnAllResults = false } = options;
+
+      if (documents.length === 0) {
+        return [];
+      }
+
+      // Process in batches
+      const allScores: number[] = [];
+
+      for (let i = 0; i < documents.length; i += batchSize) {
+        const batch = documents.slice(i, i + batchSize);
+
+        const controller = new AbortController();
+        const timeoutId = setTimeout(() => controller.abort(), timeout);
+
+        const headers: Record<string, string> = {
+          'Content-Type': 'application/json',
+        };
+        if (apiKey) {
+          headers['Authorization'] = `Bearer ${apiKey}`;
+        }
+
+        const response = await fetch(endpoint, {
+          method: 'POST',
+          headers,
+          body: JSON.stringify({
+            model: modelName,
+            query,
+            documents: batch.map(d => d.text),
+          }),
+          signal: controller.signal,
+        });
+
+        clearTimeout(timeoutId);
+
+        if (!response.ok) {
+          throw new Error(`Cross-encoder API error: ${response.status}`);
+        }
+
+        const data = (await response.json()) as { scores: number[] };
+        allScores.push(...data.scores);
+      }
+
+      // Create results with scores
+      const results: RerankResult[] = documents
+        .map((doc, index) => {
+          const score = allScores[index];
+          if (score === undefined) return null;
+          return {
+            document: doc,
+            relevanceScore: score,
+            originalIndex: index,
+            newIndex: -1, // Will be set after sorting
+          };
+        })
+        .filter((entry): entry is RerankResult => entry !== null)
+        .sort((a, b) => b.relevanceScore - a.relevanceScore)
+        .map((r, newIndex) => ({ ...r, newIndex }))
+        .filter(r => returnAllResults || r.relevanceScore >= minScore);
+
+      return topK ? results.slice(0, topK) : results;
+    },
+
+    async isAvailable(): Promise<boolean> {
+      try {
+        const response = await fetch(endpoint.replace('/rerank', '/health'), {
+          method: 'GET',
+        });
+        return response.ok;
+      } catch {
+        return false;
+      }
+    },
+  };
+}
+
+// ============================================================================
+// Reciprocal Rank Fusion
+// ============================================================================
+
+/**
+ * Combine multiple ranking results using Reciprocal Rank Fusion
+ *
+ * RRF is useful when combining results from different retrieval methods
+ * (e.g., BM25 + vector search + reranker).
+ *
+ * @example
+ * ```typescript
+ * const combined = reciprocalRankFusion([
+ *   vectorSearchResults,
+ *   bm25Results,
+ *   rerankResults,
+ * ], { k: 60 });
+ * ```
+ */
+export function reciprocalRankFusion(
+  rankings: RerankResult[][],
+  options: { k?: number; topK?: number } = {},
+): RerankResult[] {
+  const { k = 60, topK } = options;
+
+  // Calculate RRF scores
+  const scoreMap = new Map<string, { score: number; document: RerankDocument }>();
+
+  for (const ranking of rankings) {
+    for (let rank = 0; rank < ranking.length; rank++) {
+      const result = ranking[rank];
+      if (!result) {
+        continue;
+      }
+      const docId = result.document.id;
+      const rrfScore = 1 / (k + rank + 1);
+
+      const existing = scoreMap.get(docId);
+      if (existing) {
+        existing.score += rrfScore;
+      } else {
+        scoreMap.set(docId, { score: rrfScore, document: result.document });
+      }
+    }
+  }
+
+  // Sort by combined RRF score
+  const results = Array.from(scoreMap.values())
+    .sort((a, b) => b.score - a.score)
+    .map((item, index) => ({
+      document: item.document,
+      relevanceScore: item.score,
+      originalIndex: -1,
+      newIndex: index,
+    }));
+
+  return topK ? results.slice(0, topK) : results;
+}
+
+// ============================================================================
+// Reranking Pipeline
+// ============================================================================
+
+/**
+ * Configuration for the reranking pipeline
+ */
+export interface RerankPipelineConfig {
+  /** Primary reranker provider */
+  provider: RerankProvider;
+  /** Fallback provider if primary fails */
+  fallbackProvider?: RerankProvider;
+  /** Whether to use RRF with original scores */
+  useRRFWithOriginal?: boolean;
+  /** RRF k parameter */
+  rrfK?: number;
+  /** Cache rerank results */
+  enableCache?: boolean;
+  /** Cache TTL in seconds */
+  cacheTTL?: number;
+}
+
+/**
+ * Create a reranking pipeline with fallback and caching
+ *
+ * @example
+ * ```typescript
+ * const pipeline = createRerankPipeline({
+ *   provider: createCohereReranker({ apiKey: '...' }),
+ *   fallbackProvider: createCrossEncoderReranker({ ... }),
+ *   useRRFWithOriginal: true,
+ *   enableCache: true,
+ * });
+ *
+ * const results = await pipeline.rerank({
+ *   query: 'machine learning',
+ *   documents: docs,
+ *   topK: 10,
+ * });
+ * ```
+ */
+export function createRerankPipeline(config: RerankPipelineConfig) {
+  const {
+    provider,
+    fallbackProvider,
+    useRRFWithOriginal = false,
+    rrfK = 60,
+    enableCache = false,
+    cacheTTL = 300,
+  } = config;
+
+  // Simple in-memory cache
+  const cache = new Map<string, { results: RerankResult[]; timestamp: number }>();
+
+  function getCacheKey(query: string, docIds: string[]): string {
+    return `${query}:${docIds.sort().join(',')}`;
+  }
+
+  return {
+    async rerank(options: RerankOptions): Promise<RerankResult[]> {
+      const { query, documents, topK } = options;
+
+      // Check cache
+      if (enableCache) {
+        const cacheKey = getCacheKey(
+          query,
+          documents.map(d => d.id),
+        );
+        const cached = cache.get(cacheKey);
+        if (cached && Date.now() - cached.timestamp < cacheTTL * 1000) {
+          return topK ? cached.results.slice(0, topK) : cached.results;
+        }
+      }
+
+      let rerankResults: RerankResult[];
+
+      try {
+        rerankResults = await provider.rerank(options);
+      } catch (error) {
+        if (fallbackProvider) {
+          logWarn('Primary reranker failed, using fallback', {
+            error: error instanceof Error ? error : new Error(String(error)),
+          });
+          rerankResults = await fallbackProvider.rerank(options);
+        } else {
+          throw error;
+        }
+      }
+
+      // Optionally combine with original scores using RRF
+      if (useRRFWithOriginal && documents.some(d => d.originalScore !== undefined)) {
+        const originalRanking: RerankResult[] = [...documents]
+          .sort((a, b) => (b.originalScore || 0) - (a.originalScore || 0))
+          .map((doc, index) => ({
+            document: doc,
+            relevanceScore: doc.originalScore || 0,
+            originalIndex: documents.indexOf(doc),
+            newIndex: index,
+          }));
+
+        rerankResults = reciprocalRankFusion([rerankResults, originalRanking], {
+          k: rrfK,
+          topK,
+        });
+      }
+
+      // Update cache
+      if (enableCache) {
+        const cacheKey = getCacheKey(
+          query,
+          documents.map(d => d.id),
+        );
+        cache.set(cacheKey, { results: rerankResults, timestamp: Date.now() });
+
+        // Clean old entries
+        const now = Date.now();
+        for (const [key, value] of cache.entries()) {
+          if (now - value.timestamp > cacheTTL * 1000) {
+            cache.delete(key);
+          }
+        }
+      }
+
+      return rerankResults;
+    },
+
+    async isAvailable(): Promise<boolean> {
+      const primaryAvailable = await provider.isAvailable();
+      if (primaryAvailable) return true;
+      if (fallbackProvider) {
+        return fallbackProvider.isAvailable();
+      }
+      return false;
+    },
+
+    clearCache(): void {
+      cache.clear();
+    },
+  };
+}
+
+// ============================================================================
+// Utility Functions
+// ============================================================================
+
+/**
+ * Convert vector search results to RerankDocuments
+ */
+export function vectorResultsToDocuments(
+  results: Array<{
+    id: string;
+    content: string;
+    score: number;
+    metadata?: Record<string, unknown>;
+  }>,
+): RerankDocument[] {
+  return results.map(r => ({
+    id: r.id,
+    text: r.content,
+    metadata: r.metadata,
+    originalScore: r.score,
+  }));
+}
+
+/**
+ * Calculate Normalized Discounted Cumulative Gain (NDCG)
+ * for evaluating reranking quality
+ */
+export function calculateNDCG(
+  results: RerankResult[],
+  relevanceLabels: Map<string, number>,
+  k?: number,
+): number {
+  const resultsToScore = k ? results.slice(0, k) : results;
+
+  // Calculate DCG
+  let dcg = 0;
+  for (let i = 0; i < resultsToScore.length; i++) {
+    const result = resultsToScore[i];
+    if (!result) continue;
+    const relevance = relevanceLabels.get(result.document.id) || 0;
+    dcg += (Math.pow(2, relevance) - 1) / Math.log2(i + 2);
+  }
+
+  // Calculate ideal DCG
+  const idealOrder = Array.from(relevanceLabels.values()).sort((a, b) => b - a);
+  let idcg = 0;
+  const idealLength = k ? Math.min(k, idealOrder.length) : idealOrder.length;
+  for (let i = 0; i < idealLength; i++) {
+    const idealScore = idealOrder[i] ?? 0;
+    idcg += (Math.pow(2, idealScore) - 1) / Math.log2(i + 2);
+  }
+
+  return idcg > 0 ? dcg / idcg : 0;
+}
+
+/**
+ * Calculate Mean Reciprocal Rank (MRR)
+ */
+export function calculateMRR(results: RerankResult[], relevantIds: Set<string>): number {
+  for (let i = 0; i < results.length; i++) {
+    const result = results[i];
+    if (result && relevantIds.has(result.document.id)) {
+      return 1 / (i + 1);
+    }
+  }
+  return 0;
+}