claude-brain 0.15.2 → 0.16.0

package/src/retrieval/reranker/index.ts

@@ -1,188 +1,188 @@
-/**
- * Cross-Encoder Reranker
- * Uses a cross-encoder model to rerank search results for improved relevance
- */
-
-import type { Logger } from 'pino'
-import { loadRerankerModel, isModelLoaded, unloadModel } from './model'
-import type { HybridSearchResult } from '../types'
-
-export interface RerankerConfig {
-  /** Model to use for reranking */
-  model: string
-  /** Number of top results to rerank */
-  topK: number
-  /** Batch size for inference */
-  batchSize?: number
-}
-
-export class Reranker {
-  private logger: Logger
-  private config: RerankerConfig
-  private model: any | null = null
-
-  constructor(logger: Logger, config: Partial<RerankerConfig> = {}) {
-    this.logger = logger.child({ component: 'reranker' })
-    this.config = {
-      model: config.model || 'cross-encoder/ms-marco-MiniLM-L-6-v2',
-      topK: config.topK || 10,
-      batchSize: config.batchSize || 8
-    }
-  }
-
-  /**
-   * Initialize the reranker (loads model)
-   */
-  async initialize(): Promise<void> {
-    if (this.model) {
-      return
-    }
-
-    this.logger.info('Initializing reranker')
-    const { model } = await loadRerankerModel(this.logger, this.config.model)
-    this.model = model
-    this.logger.info('Reranker initialized')
-  }
-
-  /**
-   * Rerank search results using cross-encoder
-   *
-   * @param query - The search query
-   * @param results - Results to rerank
-   * @returns Reranked results with updated scores
-   */
-  async rerank(
-    query: string,
-    results: HybridSearchResult[]
-  ): Promise<HybridSearchResult[]> {
-    if (!this.model) {
-      await this.initialize()
-    }
-
-    if (results.length === 0) {
-      return []
-    }
-
-    // Only rerank top K results
-    const toRerank = results.slice(0, this.config.topK)
-    const remainder = results.slice(this.config.topK)
-
-    this.logger.debug({
-      query: query.slice(0, 50),
-      resultsCount: toRerank.length
-    }, 'Reranking results')
-
-    try {
-      // Create query-document pairs for cross-encoder
-      const pairs = toRerank.map(result => ({
-        text: query,
-        text_pair: result.content
-      }))
-
-      // Process in batches
-      const scores: number[] = []
-      for (let i = 0; i < pairs.length; i += this.config.batchSize!) {
-        const batch = pairs.slice(i, i + this.config.batchSize!)
-
-        // Run inference
-        const batchScores = await Promise.all(
-          batch.map(async (pair) => {
-            const output = await this.model(`${pair.text} [SEP] ${pair.text_pair}`)
-            // Cross-encoder outputs score indicating relevance
-            // Higher score = more relevant
-            return this.extractScore(output)
-          })
-        )
-
-        scores.push(...batchScores)
-      }
-
-      // Normalize scores to 0-1 range
-      const normalizedScores = this.normalizeScores(scores)
-
-      // Update results with reranked scores
-      const rerankedResults = toRerank.map((result, index) => ({
-        ...result,
-        scores: {
-          ...result.scores,
-          reranked: normalizedScores[index],
-          final: normalizedScores[index] // Reranked score becomes final
-        }
-      }))
-
-      // Sort by reranked score
-      rerankedResults.sort((a, b) => b.scores.final - a.scores.final)
-
-      // Append non-reranked results
-      return [...rerankedResults, ...remainder]
-
-    } catch (error) {
-      this.logger.error({ error }, 'Reranking failed, returning original order')
-      return results
-    }
-  }
-
-  /**
-   * Extract score from model output
-   */
-  private extractScore(output: any): number {
-    // Handle different output formats from transformers.js
-    if (Array.isArray(output) && output.length > 0) {
-      // Classification output format
-      if (typeof output[0] === 'object' && 'score' in output[0]) {
-        return output[0].score
-      }
-      // Raw logits format
-      if (typeof output[0] === 'number') {
-        return output[0]
-      }
-    }
-    // Single score
-    if (typeof output === 'number') {
-      return output
-    }
-    // Object with score
-    if (typeof output === 'object' && 'score' in output) {
-      return output.score
-    }
-
-    this.logger.warn({ output }, 'Unexpected reranker output format')
-    return 0
-  }
-
-  /**
-   * Normalize scores to 0-1 range using min-max normalization
-   */
-  private normalizeScores(scores: number[]): number[] {
-    if (scores.length === 0) return []
-    if (scores.length === 1) return [1.0]
-
-    const min = Math.min(...scores)
-    const max = Math.max(...scores)
-    const range = max - min
-
-    if (range === 0) {
-      return scores.map(() => 0.5)
-    }
-
-    return scores.map(score => (score - min) / range)
-  }
-
-  /**
-   * Check if reranker is ready
-   */
-  isReady(): boolean {
-    return isModelLoaded()
-  }
-
-  /**
-   * Cleanup model from memory
-   */
-  cleanup(): void {
-    unloadModel()
-    this.model = null
-    this.logger.info('Reranker model unloaded')
-  }
-}
-
-export { loadRerankerModel, isModelLoaded, unloadModel } from './model'
+/**
+ * Cross-Encoder Reranker
+ * Uses a cross-encoder model to rerank search results for improved relevance
+ */
+
+import type { Logger } from 'pino'
+import { loadRerankerModel, isModelLoaded, unloadModel } from './model'
+import type { HybridSearchResult } from '../types'
+
+export interface RerankerConfig {
+  /** Model to use for reranking */
+  model: string
+  /** Number of top results to rerank */
+  topK: number
+  /** Batch size for inference */
+  batchSize?: number
+}
+
+export class Reranker {
+  private logger: Logger
+  private config: RerankerConfig
+  private model: any | null = null
+
+  constructor(logger: Logger, config: Partial<RerankerConfig> = {}) {
+    this.logger = logger.child({ component: 'reranker' })
+    this.config = {
+      model: config.model || 'cross-encoder/ms-marco-MiniLM-L-6-v2',
+      topK: config.topK || 10,
+      batchSize: config.batchSize || 8
+    }
+  }
+
+  /**
+   * Initialize the reranker (loads model)
+   */
+  async initialize(): Promise<void> {
+    if (this.model) {
+      return
+    }
+
+    this.logger.info('Initializing reranker')
+    const { model } = await loadRerankerModel(this.logger, this.config.model)
+    this.model = model
+    this.logger.info('Reranker initialized')
+  }
+
+  /**
+   * Rerank search results using cross-encoder
+   *
+   * @param query - The search query
+   * @param results - Results to rerank
+   * @returns Reranked results with updated scores
+   */
+  async rerank(
+    query: string,
+    results: HybridSearchResult[]
+  ): Promise<HybridSearchResult[]> {
+    if (!this.model) {
+      await this.initialize()
+    }
+
+    if (results.length === 0) {
+      return []
+    }
+
+    // Only rerank top K results
+    const toRerank = results.slice(0, this.config.topK)
+    const remainder = results.slice(this.config.topK)
+
+    this.logger.debug({
+      query: query.slice(0, 50),
+      resultsCount: toRerank.length
+    }, 'Reranking results')
+
+    try {
+      // Create query-document pairs for cross-encoder
+      const pairs = toRerank.map(result => ({
+        text: query,
+        text_pair: result.content
+      }))
+
+      // Process in batches
+      const scores: number[] = []
+      for (let i = 0; i < pairs.length; i += this.config.batchSize!) {
+        const batch = pairs.slice(i, i + this.config.batchSize!)
+
+        // Run inference
+        const batchScores = await Promise.all(
+          batch.map(async (pair) => {
+            const output = await this.model(`${pair.text} [SEP] ${pair.text_pair}`)
+            // Cross-encoder outputs score indicating relevance
+            // Higher score = more relevant
+            return this.extractScore(output)
+          })
+        )
+
+        scores.push(...batchScores)
+      }
+
+      // Normalize scores to 0-1 range
+      const normalizedScores = this.normalizeScores(scores)
+
+      // Update results with reranked scores
+      const rerankedResults = toRerank.map((result, index) => ({
+        ...result,
+        scores: {
+          ...result.scores,
+          reranked: normalizedScores[index],
+          final: normalizedScores[index] // Reranked score becomes final
+        }
+      }))
+
+      // Sort by reranked score
+      rerankedResults.sort((a, b) => b.scores.final - a.scores.final)
+
+      // Append non-reranked results
+      return [...rerankedResults, ...remainder]
+
+    } catch (error) {
+      this.logger.error({ error }, 'Reranking failed, returning original order')
+      return results
+    }
+  }
+
+  /**
+   * Extract score from model output
+   */
+  private extractScore(output: any): number {
+    // Handle different output formats from transformers.js
+    if (Array.isArray(output) && output.length > 0) {
+      // Classification output format
+      if (typeof output[0] === 'object' && 'score' in output[0]) {
+        return output[0].score
+      }
+      // Raw logits format
+      if (typeof output[0] === 'number') {
+        return output[0]
+      }
+    }
+    // Single score
+    if (typeof output === 'number') {
+      return output
+    }
+    // Object with score
+    if (typeof output === 'object' && 'score' in output) {
+      return output.score
+    }
+
+    this.logger.warn({ output }, 'Unexpected reranker output format')
+    return 0
+  }
+
+  /**
+   * Normalize scores to 0-1 range using min-max normalization
+   */
+  private normalizeScores(scores: number[]): number[] {
+    if (scores.length === 0) return []
+    if (scores.length === 1) return [1.0]
+
+    const min = Math.min(...scores)
+    const max = Math.max(...scores)
+    const range = max - min
+
+    if (range === 0) {
+      return scores.map(() => 0.5)
+    }
+
+    return scores.map(score => (score - min) / range)
+  }
+
+  /**
+   * Check if reranker is ready
+   */
+  isReady(): boolean {
+    return isModelLoaded()
+  }
+
+  /**
+   * Cleanup model from memory
+   */
+  cleanup(): void {
+    unloadModel()
+    this.model = null
+    this.logger.info('Reranker model unloaded')
+  }
+}
+
+export { loadRerankerModel, isModelLoaded, unloadModel } from './model'

package/src/retrieval/reranker/model.ts

@@ -1,95 +1,95 @@
-/**
- * Reranker Model Loader
- * Handles loading and caching of the cross-encoder model
- */
-
-import type { Logger } from 'pino'
-
-// Dynamic import for transformers to avoid bundling issues
-let pipeline: any
-let AutoTokenizer: any
-
-interface RerankerModel {
-  tokenizer: any
-  model: any
-}
-
-let modelInstance: RerankerModel | null = null
-let loadingPromise: Promise<RerankerModel> | null = null
-
-/**
- * Load the cross-encoder model
- * Uses singleton pattern to avoid loading multiple times
- */
-export async function loadRerankerModel(
-  logger: Logger,
-  modelName: string = 'cross-encoder/ms-marco-MiniLM-L-6-v2'
-): Promise<RerankerModel> {
-  // Return existing instance if loaded
-  if (modelInstance) {
-    return modelInstance
-  }
-
-  // Wait for in-progress loading
-  if (loadingPromise) {
-    return loadingPromise
-  }
-
-  // Start loading
-  loadingPromise = (async () => {
-    logger.info({ modelName }, 'Loading cross-encoder reranker model')
-
-    try {
-      // Dynamic import to avoid issues if transformers.js not available
-      const transformers = await import('@xenova/transformers')
-      pipeline = transformers.pipeline
-      AutoTokenizer = transformers.AutoTokenizer
-
-      // Load text-classification pipeline for cross-encoder
-      // Cross-encoders score query-document pairs
-      const classifier = await pipeline('text-classification', modelName, {
-        quantized: true // Use quantized model for faster inference
-      })
-
-      const tokenizer = await AutoTokenizer.from_pretrained(modelName)
-
-      modelInstance = {
-        tokenizer,
-        model: classifier
-      }
-
-      logger.info({ modelName }, 'Cross-encoder model loaded successfully')
-      return modelInstance
-
-    } catch (error) {
-      logger.error({ error, modelName }, 'Failed to load cross-encoder model')
-      loadingPromise = null
-      throw error
-    }
-  })()
-
-  return loadingPromise
-}
-
-/**
- * Check if model is loaded
- */
-export function isModelLoaded(): boolean {
-  return modelInstance !== null
-}
-
-/**
- * Unload model to free memory
- */
-export function unloadModel(): void {
-  modelInstance = null
-  loadingPromise = null
-}
-
-/**
- * Get model memory usage estimate
- */
-export function getModelMemoryEstimate(): number {
-  // ms-marco-MiniLM-L-6-v2 quantized is roughly 22MB
-  return isModelLoaded() ? 22 * 1024 * 1024 : 0
-}
+/**
+ * Reranker Model Loader
+ * Handles loading and caching of the cross-encoder model
+ */
+
+import type { Logger } from 'pino'
+
+// Dynamic import for transformers to avoid bundling issues
+let pipeline: any
+let AutoTokenizer: any
+
+interface RerankerModel {
+  tokenizer: any
+  model: any
+}
+
+let modelInstance: RerankerModel | null = null
+let loadingPromise: Promise<RerankerModel> | null = null
+
+/**
+ * Load the cross-encoder model
+ * Uses singleton pattern to avoid loading multiple times
+ */
+export async function loadRerankerModel(
+  logger: Logger,
+  modelName: string = 'cross-encoder/ms-marco-MiniLM-L-6-v2'
+): Promise<RerankerModel> {
+  // Return existing instance if loaded
+  if (modelInstance) {
+    return modelInstance
+  }
+
+  // Wait for in-progress loading
+  if (loadingPromise) {
+    return loadingPromise
+  }
+
+  // Start loading
+  loadingPromise = (async () => {
+    logger.info({ modelName }, 'Loading cross-encoder reranker model')
+
+    try {
+      // Dynamic import to avoid issues if transformers.js not available
+      const transformers = await import('@xenova/transformers')
+      pipeline = transformers.pipeline
+      AutoTokenizer = transformers.AutoTokenizer
+
+      // Load text-classification pipeline for cross-encoder
+      // Cross-encoders score query-document pairs
+      const classifier = await pipeline('text-classification', modelName, {
+        quantized: true // Use quantized model for faster inference
+      })
+
+      const tokenizer = await AutoTokenizer.from_pretrained(modelName)
+
+      modelInstance = {
+        tokenizer,
+        model: classifier
+      }
+
+      logger.info({ modelName }, 'Cross-encoder model loaded successfully')
+      return modelInstance
+
+    } catch (error) {
+      logger.error({ error, modelName }, 'Failed to load cross-encoder model')
+      loadingPromise = null
+      throw error
+    }
+  })()
+
+  return loadingPromise
+}
+
+/**
+ * Check if model is loaded
+ */
+export function isModelLoaded(): boolean {
+  return modelInstance !== null
+}
+
+/**
+ * Unload model to free memory
+ */
+export function unloadModel(): void {
+  modelInstance = null
+  loadingPromise = null
+}
+
+/**
+ * Get model memory usage estimate
+ */
+export function getModelMemoryEstimate(): number {
+  // ms-marco-MiniLM-L-6-v2 quantized is roughly 22MB
+  return isModelLoaded() ? 22 * 1024 * 1024 : 0
+}