claude-brain 0.15.2 → 0.17.0

package/src/retrieval/bm25/index.ts

@@ -1,300 +1,300 @@
-/**
- * BM25 Sparse Search Engine
- * Uses MiniSearch for fast keyword-based retrieval
- */
-
-import type { Logger } from 'pino'
-import MiniSearch from 'minisearch'
-import { tokenize, tokenizeForSearch, type TokenizerOptions } from './tokenizer'
-
-export interface BM25Document {
-  id: string
-  content: string
-  metadata: Record<string, unknown>
-  collection: string
-}
-
-export interface BM25SearchResult {
-  id: string
-  content: string
-  metadata: Record<string, unknown>
-  collection: string
-  score: number
-}
-
-export interface BM25Config {
-  /** Fields to index */
-  fields: string[]
-  /** Fields to store */
-  storeFields: string[]
-  /** BM25 k1 parameter (term frequency saturation) */
-  k1?: number
-  /** BM25 b parameter (document length normalization) */
-  b?: number
-  /** Boost for exact matches */
-  boostExact?: number
-}
-
-const DEFAULT_CONFIG: BM25Config = {
-  fields: ['content'],
-  storeFields: ['content', 'metadata', 'collection'],
-  k1: 1.2,
-  b: 0.75,
-  boostExact: 2.0
-}
-
-export class BM25Engine {
-  private logger: Logger
-  private config: BM25Config
-  private index: MiniSearch<BM25Document>
-  private documentCount: number = 0
-  private tokenizerOptions: TokenizerOptions
-
-  constructor(logger: Logger, config: Partial<BM25Config> = {}) {
-    this.logger = logger.child({ component: 'bm25-engine' })
-    this.config = { ...DEFAULT_CONFIG, ...config }
-    this.tokenizerOptions = {
-      minLength: 2,
-      removeStopwords: true,
-      stemming: true,
-      splitCamelCase: true
-    }
-
-    this.index = this.createIndex()
-  }
-
-  /**
-   * Create a new MiniSearch index
-   */
-  private createIndex(): MiniSearch<BM25Document> {
-    return new MiniSearch<BM25Document>({
-      fields: this.config.fields,
-      storeFields: this.config.storeFields,
-      idField: 'id',
-      tokenize: (text: string) => tokenize(text, this.tokenizerOptions),
-      processTerm: (term: string) => term.toLowerCase()
-    })
-  }
-
-  /**
-   * Build index from documents
-   */
-  async buildIndex(documents: BM25Document[]): Promise<void> {
-    this.logger.info({ count: documents.length }, 'Building BM25 index')
-
-    // Create fresh index
-    this.index = this.createIndex()
-    this.documentCount = 0
-
-    // Add documents in batches
-    const batchSize = 1000
-    for (let i = 0; i < documents.length; i += batchSize) {
-      const batch = documents.slice(i, i + batchSize)
-      await this.addDocuments(batch)
-    }
-
-    this.logger.info({
-      indexed: this.documentCount,
-      terms: this.index.termCount
-    }, 'BM25 index built')
-  }
-
-  /**
-   * Add documents to index
-   */
-  async addDocuments(documents: BM25Document[]): Promise<void> {
-    try {
-      this.index.addAll(documents)
-      this.documentCount += documents.length
-    } catch (error) {
-      this.logger.error({ error }, 'Failed to add documents to BM25 index')
-      throw error
-    }
-  }
-
-  /**
-   * Add single document
-   */
-  addDocument(document: BM25Document): void {
-    try {
-      this.index.add(document)
-      this.documentCount++
-    } catch (error) {
-      this.logger.error({ error, id: document.id }, 'Failed to add document to BM25 index')
-      throw error
-    }
-  }
-
-  /**
-   * Remove document from index
-   */
-  removeDocument(document: BM25Document): void {
-    try {
-      this.index.remove(document)
-      this.documentCount--
-    } catch (error) {
-      this.logger.error({ error, id: document.id }, 'Failed to remove document from BM25 index')
-      // Don't throw - document might not exist
-    }
-  }
-
-  /**
-   * Update document in index
-   */
-  updateDocument(document: BM25Document): void {
-    this.removeDocument(document)
-    this.addDocument(document)
-  }
-
-  /**
-   * Search the index
-   */
-  search(
-    query: string,
-    options: {
-      limit?: number
-      filter?: (result: BM25SearchResult) => boolean
-      collection?: string
-    } = {}
-  ): BM25SearchResult[] {
-    const { limit = 20, filter, collection } = options
-
-    if (!query.trim()) {
-      return []
-    }
-
-    this.logger.debug({
-      query: query.slice(0, 50),
-      limit
-    }, 'BM25 search')
-
-    try {
-      // Tokenize query
-      const queryTokens = tokenizeForSearch(query, this.tokenizerOptions)
-
-      // Search with MiniSearch options
-      const results = this.index.search(query, {
-        prefix: true, // Allow prefix matching
-        fuzzy: 0.2, // Allow minor typos
-        combineWith: 'OR', // Match any term
-        boost: {
-          content: this.config.boostExact!
-        }
-      })
-
-      // Transform and filter results
-      let searchResults: BM25SearchResult[] = results.map(result => ({
-        id: result.id,
-        content: (result as any).content || '',
-        metadata: (result as any).metadata || {},
-        collection: (result as any).collection || '',
-        score: result.score
-      }))
-
-      // Filter by collection if specified
-      if (collection) {
-        searchResults = searchResults.filter(r => r.collection === collection)
-      }
-
-      // Apply custom filter
-      if (filter) {
-        searchResults = searchResults.filter(filter)
-      }
-
-      // Normalize scores to 0-1 range
-      searchResults = this.normalizeScores(searchResults)
-
-      // Apply limit
-      return searchResults.slice(0, limit)
-
-    } catch (error) {
-      this.logger.error({ error, query }, 'BM25 search failed')
-      return []
-    }
-  }
-
-  /**
-   * Normalize scores to 0-1 range
-   */
-  private normalizeScores(results: BM25SearchResult[]): BM25SearchResult[] {
-    if (results.length === 0) return []
-    if (results.length === 1) {
-      return [{ ...results[0], score: 1.0 }]
-    }
-
-    const maxScore = Math.max(...results.map(r => r.score))
-    if (maxScore === 0) {
-      return results.map(r => ({ ...r, score: 0 }))
-    }
-
-    return results.map(r => ({
-      ...r,
-      score: r.score / maxScore
-    }))
-  }
-
-  /**
-   * Get suggestions for autocomplete
-   */
-  suggest(query: string, limit: number = 5): string[] {
-    try {
-      const results = this.index.autoSuggest(query, { limit })
-      return results.map(r => r.suggestion)
-    } catch (error) {
-      this.logger.error({ error, query }, 'BM25 suggest failed')
-      return []
-    }
-  }
-
-  /**
-   * Get index statistics
-   */
-  getStats(): {
-    documentCount: number
-    termCount: number
-  } {
-    return {
-      documentCount: this.documentCount,
-      termCount: this.index.termCount
-    }
-  }
-
-  /**
-   * Clear the index
-   */
-  clear(): void {
-    this.index = this.createIndex()
-    this.documentCount = 0
-    this.logger.info('BM25 index cleared')
-  }
-
-  /**
-   * Export index for persistence
-   */
-  exportIndex(): string {
-    return JSON.stringify(this.index.toJSON())
-  }
-
-  /**
-   * Import index from persisted data
-   */
-  importIndex(data: string): void {
-    try {
-      const parsed = JSON.parse(data)
-      this.index = MiniSearch.loadJSON(data, {
-        fields: this.config.fields,
-        storeFields: this.config.storeFields,
-        idField: 'id'
-      })
-      // Count documents after import
-      this.documentCount = this.index.documentCount
-      this.logger.info({ documentCount: this.documentCount }, 'BM25 index imported')
-    } catch (error) {
-      this.logger.error({ error }, 'Failed to import BM25 index')
-      throw error
-    }
-  }
-}
-
-export { tokenize, tokenizeForSearch, getNGrams } from './tokenizer'
-export type { TokenizerOptions } from './tokenizer'
+/**
+ * BM25 Sparse Search Engine
+ * Uses MiniSearch for fast keyword-based retrieval
+ */
+
+import type { Logger } from 'pino'
+import MiniSearch from 'minisearch'
+import { tokenize, tokenizeForSearch, type TokenizerOptions } from './tokenizer'
+
+export interface BM25Document {
+  id: string
+  content: string
+  metadata: Record<string, unknown>
+  collection: string
+}
+
+export interface BM25SearchResult {
+  id: string
+  content: string
+  metadata: Record<string, unknown>
+  collection: string
+  score: number
+}
+
+export interface BM25Config {
+  /** Fields to index */
+  fields: string[]
+  /** Fields to store */
+  storeFields: string[]
+  /** BM25 k1 parameter (term frequency saturation) */
+  k1?: number
+  /** BM25 b parameter (document length normalization) */
+  b?: number
+  /** Boost for exact matches */
+  boostExact?: number
+}
+
+const DEFAULT_CONFIG: BM25Config = {
+  fields: ['content'],
+  storeFields: ['content', 'metadata', 'collection'],
+  k1: 1.2,
+  b: 0.75,
+  boostExact: 2.0
+}
+
+export class BM25Engine {
+  private logger: Logger
+  private config: BM25Config
+  private index: MiniSearch<BM25Document>
+  private documentCount: number = 0
+  private tokenizerOptions: TokenizerOptions
+
+  constructor(logger: Logger, config: Partial<BM25Config> = {}) {
+    this.logger = logger.child({ component: 'bm25-engine' })
+    this.config = { ...DEFAULT_CONFIG, ...config }
+    this.tokenizerOptions = {
+      minLength: 2,
+      removeStopwords: true,
+      stemming: true,
+      splitCamelCase: true
+    }
+
+    this.index = this.createIndex()
+  }
+
+  /**
+   * Create a new MiniSearch index
+   */
+  private createIndex(): MiniSearch<BM25Document> {
+    return new MiniSearch<BM25Document>({
+      fields: this.config.fields,
+      storeFields: this.config.storeFields,
+      idField: 'id',
+      tokenize: (text: string) => tokenize(text, this.tokenizerOptions),
+      processTerm: (term: string) => term.toLowerCase()
+    })
+  }
+
+  /**
+   * Build index from documents
+   */
+  async buildIndex(documents: BM25Document[]): Promise<void> {
+    this.logger.info({ count: documents.length }, 'Building BM25 index')
+
+    // Create fresh index
+    this.index = this.createIndex()
+    this.documentCount = 0
+
+    // Add documents in batches
+    const batchSize = 1000
+    for (let i = 0; i < documents.length; i += batchSize) {
+      const batch = documents.slice(i, i + batchSize)
+      await this.addDocuments(batch)
+    }
+
+    this.logger.info({
+      indexed: this.documentCount,
+      terms: this.index.termCount
+    }, 'BM25 index built')
+  }
+
+  /**
+   * Add documents to index
+   */
+  async addDocuments(documents: BM25Document[]): Promise<void> {
+    try {
+      this.index.addAll(documents)
+      this.documentCount += documents.length
+    } catch (error) {
+      this.logger.error({ error }, 'Failed to add documents to BM25 index')
+      throw error
+    }
+  }
+
+  /**
+   * Add single document
+   */
+  addDocument(document: BM25Document): void {
+    try {
+      this.index.add(document)
+      this.documentCount++
+    } catch (error) {
+      this.logger.error({ error, id: document.id }, 'Failed to add document to BM25 index')
+      throw error
+    }
+  }
+
+  /**
+   * Remove document from index
+   */
+  removeDocument(document: BM25Document): void {
+    try {
+      this.index.remove(document)
+      this.documentCount--
+    } catch (error) {
+      this.logger.error({ error, id: document.id }, 'Failed to remove document from BM25 index')
+      // Don't throw - document might not exist
+    }
+  }
+
+  /**
+   * Update document in index
+   */
+  updateDocument(document: BM25Document): void {
+    this.removeDocument(document)
+    this.addDocument(document)
+  }
+
+  /**
+   * Search the index
+   */
+  search(
+    query: string,
+    options: {
+      limit?: number
+      filter?: (result: BM25SearchResult) => boolean
+      collection?: string
+    } = {}
+  ): BM25SearchResult[] {
+    const { limit = 20, filter, collection } = options
+
+    if (!query.trim()) {
+      return []
+    }
+
+    this.logger.debug({
+      query: query.slice(0, 50),
+      limit
+    }, 'BM25 search')
+
+    try {
+      // Tokenize query
+      const queryTokens = tokenizeForSearch(query, this.tokenizerOptions)
+
+      // Search with MiniSearch options
+      const results = this.index.search(query, {
+        prefix: true, // Allow prefix matching
+        fuzzy: 0.2, // Allow minor typos
+        combineWith: 'OR', // Match any term
+        boost: {
+          content: this.config.boostExact!
+        }
+      })
+
+      // Transform and filter results
+      let searchResults: BM25SearchResult[] = results.map(result => ({
+        id: result.id,
+        content: (result as any).content || '',
+        metadata: (result as any).metadata || {},
+        collection: (result as any).collection || '',
+        score: result.score
+      }))
+
+      // Filter by collection if specified
+      if (collection) {
+        searchResults = searchResults.filter(r => r.collection === collection)
+      }
+
+      // Apply custom filter
+      if (filter) {
+        searchResults = searchResults.filter(filter)
+      }
+
+      // Normalize scores to 0-1 range
+      searchResults = this.normalizeScores(searchResults)
+
+      // Apply limit
+      return searchResults.slice(0, limit)
+
+    } catch (error) {
+      this.logger.error({ error, query }, 'BM25 search failed')
+      return []
+    }
+  }
+
+  /**
+   * Normalize scores to 0-1 range
+   */
+  private normalizeScores(results: BM25SearchResult[]): BM25SearchResult[] {
+    if (results.length === 0) return []
+    if (results.length === 1) {
+      return [{ ...results[0], score: 1.0 }]
+    }
+
+    const maxScore = Math.max(...results.map(r => r.score))
+    if (maxScore === 0) {
+      return results.map(r => ({ ...r, score: 0 }))
+    }
+
+    return results.map(r => ({
+      ...r,
+      score: r.score / maxScore
+    }))
+  }
+
+  /**
+   * Get suggestions for autocomplete
+   */
+  suggest(query: string, limit: number = 5): string[] {
+    try {
+      const results = this.index.autoSuggest(query, { limit })
+      return results.map(r => r.suggestion)
+    } catch (error) {
+      this.logger.error({ error, query }, 'BM25 suggest failed')
+      return []
+    }
+  }
+
+  /**
+   * Get index statistics
+   */
+  getStats(): {
+    documentCount: number
+    termCount: number
+  } {
+    return {
+      documentCount: this.documentCount,
+      termCount: this.index.termCount
+    }
+  }
+
+  /**
+   * Clear the index
+   */
+  clear(): void {
+    this.index = this.createIndex()
+    this.documentCount = 0
+    this.logger.info('BM25 index cleared')
+  }
+
+  /**
+   * Export index for persistence
+   */
+  exportIndex(): string {
+    return JSON.stringify(this.index.toJSON())
+  }
+
+  /**
+   * Import index from persisted data
+   */
+  importIndex(data: string): void {
+    try {
+      const parsed = JSON.parse(data)
+      this.index = MiniSearch.loadJSON(data, {
+        fields: this.config.fields,
+        storeFields: this.config.storeFields,
+        idField: 'id'
+      })
+      // Count documents after import
+      this.documentCount = this.index.documentCount
+      this.logger.info({ documentCount: this.documentCount }, 'BM25 index imported')
+    } catch (error) {
+      this.logger.error({ error }, 'Failed to import BM25 index')
+      throw error
+    }
+  }
+}
+
+export { tokenize, tokenizeForSearch, getNGrams } from './tokenizer'
+export type { TokenizerOptions } from './tokenizer'