claude-brain 0.15.2 → 0.17.0

package/src/retrieval/pipeline.ts

@@ -1,375 +1,375 @@
-/**
- * Hybrid Retrieval Pipeline
- * Orchestrates Dense -> Sparse -> Fusion -> Rerank flow
- */
-
-import type { Logger } from 'pino'
-import type { CollectionManager } from '@/memory/chroma/collection-manager'
-import type { EmbeddingProvider } from '@/memory/chroma/embeddings'
-import type { RetrievalConfig } from '@/config/schema'
-import type { HybridSearchResult, HybridSearchOptions } from './types'
-import { BM25Engine, type BM25Document } from './bm25'
-import { createFusionStrategy, RRFFusion, LinearFusion, MaxFusion } from './fusion'
-import { Reranker } from './reranker'
-import { COLLECTIONS } from '@/memory/chroma/schemas'
-
-export class RetrievalPipeline {
-  private logger: Logger
-  private config: RetrievalConfig
-  private collections: CollectionManager
-  private embeddings: EmbeddingProvider
-  private bm25Engine: BM25Engine
-  private reranker: Reranker | null = null
-  private fusionStrategy: RRFFusion | LinearFusion | MaxFusion
-  private initialized: boolean = false
-  private indexBuilt: boolean = false
-
-  constructor(
-    logger: Logger,
-    collections: CollectionManager,
-    embeddings: EmbeddingProvider,
-    config: RetrievalConfig
-  ) {
-    this.logger = logger.child({ component: 'retrieval-pipeline' })
-    this.config = config
-    this.collections = collections
-    this.embeddings = embeddings
-
-    // Initialize BM25 engine
-    this.bm25Engine = new BM25Engine(logger)
-
-    // Initialize fusion strategy
-    this.fusionStrategy = createFusionStrategy({
-      method: config.fusion.method,
-      rrfK: config.fusion.rrfK,
-      denseWeight: config.dense.weight,
-      sparseWeight: config.sparse.weight
-    })
-  }
-
-  /**
-   * Initialize the pipeline
-   */
-  async initialize(): Promise<void> {
-    if (this.initialized) {
-      return
-    }
-
-    this.logger.info('Initializing retrieval pipeline')
-
-    // Initialize reranker if enabled
-    if (this.config.reranker.enabled) {
-      this.reranker = new Reranker(this.logger, {
-        model: this.config.reranker.model,
-        topK: this.config.reranker.topK
-      })
-      await this.reranker.initialize()
-      this.logger.info('Reranker initialized')
-    }
-
-    // Build BM25 index from existing documents
-    if (this.config.sparse.enabled) {
-      await this.buildBM25Index()
-    }
-
-    this.initialized = true
-    this.logger.info('Retrieval pipeline initialized')
-  }
-
-  /**
-   * Build BM25 index from ChromaDB collections
-   */
-  async buildBM25Index(): Promise<void> {
-    this.logger.info('Building BM25 index from ChromaDB')
-
-    const documents: BM25Document[] = []
-    const collectionsToIndex = [
-      COLLECTIONS.DECISIONS,
-      COLLECTIONS.MEMORIES,
-      COLLECTIONS.PATTERNS,
-      COLLECTIONS.CORRECTIONS
-    ]
-
-    for (const collectionName of collectionsToIndex) {
-      try {
-        const collection = await this.collections.getDecisions() // Use appropriate getter
-        const results = await collection.get({
-          include: ['documents', 'metadatas']
-        })
-
-        if (results.ids.length > 0) {
-          for (let i = 0; i < results.ids.length; i++) {
-            documents.push({
-              id: results.ids[i],
-              content: results.documents?.[i] || '',
-              metadata: results.metadatas?.[i] || {},
-              collection: collectionName
-            })
-          }
-        }
-      } catch (error) {
-        this.logger.warn({ error, collection: collectionName }, 'Failed to index collection')
-      }
-    }
-
-    await this.bm25Engine.buildIndex(documents)
-    this.indexBuilt = true
-    this.logger.info({ documentCount: documents.length }, 'BM25 index built')
-  }
-
-  /**
-   * Perform hybrid search
-   */
-  async search(
-    query: string,
-    options: HybridSearchOptions = {}
-  ): Promise<HybridSearchResult[]> {
-    const {
-      project,
-      collections = ['decisions', 'memories', 'patterns', 'corrections'],
-      limit = 10,
-      minScore = 0,
-      useReranker = this.config.reranker.enabled,
-      dateRange,
-      tags
-    } = options
-
-    this.logger.debug({
-      query: query.slice(0, 50),
-      options
-    }, 'Hybrid search')
-
-    // Step 1: Dense (semantic) search
-    const denseResults = await this.denseSearch(query, {
-      project,
-      collections,
-      limit: this.config.dense.limit,
-      minSimilarity: this.config.dense.minSimilarity,
-      dateRange,
-      tags
-    })
-
-    // Step 2: Sparse (BM25) search
-    let sparseResults: HybridSearchResult[] = []
-    if (this.config.sparse.enabled && this.indexBuilt) {
-      sparseResults = await this.sparseSearch(query, {
-        project,
-        collections,
-        limit: this.config.sparse.limit
-      })
-    }
-
-    // Step 3: Fusion
-    let fusedResults: HybridSearchResult[]
-    if (sparseResults.length > 0) {
-      fusedResults = this.fusionStrategy.fuse(denseResults, sparseResults)
-    } else {
-      // If no sparse results, use dense results directly
-      fusedResults = denseResults.map(r => ({
-        ...r,
-        scores: { ...r.scores, fusion: r.scores.dense, final: r.scores.dense }
-      }))
-    }
-
-    // Step 4: Reranking (optional)
-    if (useReranker && this.reranker && fusedResults.length > 0) {
-      fusedResults = await this.reranker.rerank(query, fusedResults)
-    }
-
-    // Filter by minimum score
-    if (minScore > 0) {
-      fusedResults = fusedResults.filter(r => r.scores.final >= minScore)
-    }
-
-    // Apply final limit
-    return fusedResults.slice(0, limit)
-  }
-
-  /**
-   * Dense (semantic) search using ChromaDB
-   */
-  private async denseSearch(
-    query: string,
-    options: {
-      project?: string
-      collections: string[]
-      limit: number
-      minSimilarity: number
-      dateRange?: { start?: string; end?: string }
-      tags?: string[]
-    }
-  ): Promise<HybridSearchResult[]> {
-    const results: HybridSearchResult[] = []
-
-    // Generate query embedding
-    const embedding = await this.embeddings.generate(query)
-
-    // Search each collection
-    for (const collectionName of options.collections) {
-      try {
-        const collection = await this.getCollection(collectionName)
-        if (!collection) continue
-
-        // Build where clause
-        const where: Record<string, any> = {}
-        if (options.project) {
-          where['project'] = { $eq: options.project }
-        }
-
-        const queryResults = await collection.query({
-          queryEmbeddings: [embedding],
-          nResults: Math.ceil(options.limit / options.collections.length),
-          where: Object.keys(where).length > 0 ? where : undefined,
-          include: ['documents', 'metadatas', 'distances']
-        })
-
-        // Process results
-        if (queryResults.ids[0]) {
-          for (let i = 0; i < queryResults.ids[0].length; i++) {
-            const distance = queryResults.distances?.[0]?.[i] || 0
-            const similarity = 1 - distance // Convert distance to similarity
-
-            if (similarity >= options.minSimilarity) {
-              results.push({
-                id: queryResults.ids[0][i],
-                content: queryResults.documents?.[0]?.[i] || '',
-                metadata: queryResults.metadatas?.[0]?.[i] || {},
-                collection: collectionName,
-                scores: {
-                  dense: similarity,
-                  sparse: 0,
-                  fusion: similarity,
-                  final: similarity
-                },
-                provenance: 'dense'
-              })
-            }
-          }
-        }
-      } catch (error) {
-        this.logger.error({ error, collection: collectionName }, 'Dense search failed')
-      }
-    }
-
-    // Sort by dense score
-    results.sort((a, b) => b.scores.dense - a.scores.dense)
-    return results
-  }
-
-  /**
-   * Sparse (BM25) search
-   */
-  private async sparseSearch(
-    query: string,
-    options: {
-      project?: string
-      collections: string[]
-      limit: number
-    }
-  ): Promise<HybridSearchResult[]> {
-    const bm25Results = this.bm25Engine.search(query, {
-      limit: options.limit,
-      filter: (result) => {
-        // Filter by project if specified
-        if (options.project && result.metadata.project !== options.project) {
-          return false
-        }
-        // Filter by collection
-        if (!options.collections.includes(result.collection)) {
-          return false
-        }
-        return true
-      }
-    })
-
-    return bm25Results.map(result => ({
-      id: result.id,
-      content: result.content,
-      metadata: result.metadata,
-      collection: result.collection,
-      scores: {
-        dense: 0,
-        sparse: result.score,
-        fusion: result.score,
-        final: result.score
-      },
-      provenance: 'sparse' as const
-    }))
-  }
-
-  /**
-   * Get collection by name
-   */
-  private async getCollection(name: string) {
-    switch (name) {
-      case 'decisions':
-        return this.collections.getDecisions()
-      case 'memories':
-        return this.collections.getMemories()
-      case 'patterns':
-        return this.collections.getPatterns()
-      case 'corrections':
-        return this.collections.getCorrections()
-      default:
-        return null
-    }
-  }
-
-  /**
-   * Add document to BM25 index
-   */
-  addToIndex(document: BM25Document): void {
-    if (this.config.sparse.enabled) {
-      this.bm25Engine.addDocument(document)
-    }
-  }
-
-  /**
-   * Remove document from BM25 index
-   */
-  removeFromIndex(document: BM25Document): void {
-    if (this.config.sparse.enabled) {
-      this.bm25Engine.removeDocument(document)
-    }
-  }
-
-  /**
-   * Rebuild BM25 index
-   */
-  async rebuildIndex(): Promise<void> {
-    await this.buildBM25Index()
-  }
-
-  /**
-   * Get pipeline status
-   */
-  getStatus(): {
-    initialized: boolean
-    indexBuilt: boolean
-    rerankerEnabled: boolean
-    sparseEnabled: boolean
-    fusionMethod: string
-    bm25Stats: { documentCount: number; termCount: number }
-  } {
-    return {
-      initialized: this.initialized,
-      indexBuilt: this.indexBuilt,
-      rerankerEnabled: this.reranker !== null,
-      sparseEnabled: this.config.sparse.enabled,
-      fusionMethod: this.config.fusion.method,
-      bm25Stats: this.bm25Engine.getStats()
-    }
-  }
-
-  /**
-   * Cleanup resources
-   */
-  cleanup(): void {
-    if (this.reranker) {
-      this.reranker.cleanup()
-    }
-    this.bm25Engine.clear()
-    this.initialized = false
-    this.indexBuilt = false
-  }
-}
+/**
+ * Hybrid Retrieval Pipeline
+ * Orchestrates Dense -> Sparse -> Fusion -> Rerank flow
+ */
+
+import type { Logger } from 'pino'
+import type { CollectionManager } from '@/memory/chroma/collection-manager'
+import type { EmbeddingProvider } from '@/memory/chroma/embeddings'
+import type { RetrievalConfig } from '@/config/schema'
+import type { HybridSearchResult, HybridSearchOptions } from './types'
+import { BM25Engine, type BM25Document } from './bm25'
+import { createFusionStrategy, RRFFusion, LinearFusion, MaxFusion } from './fusion'
+import { Reranker } from './reranker'
+import { COLLECTIONS } from '@/memory/chroma/schemas'
+
+export class RetrievalPipeline {
+  private logger: Logger
+  private config: RetrievalConfig
+  private collections: CollectionManager
+  private embeddings: EmbeddingProvider
+  private bm25Engine: BM25Engine
+  private reranker: Reranker | null = null
+  private fusionStrategy: RRFFusion | LinearFusion | MaxFusion
+  private initialized: boolean = false
+  private indexBuilt: boolean = false
+
+  constructor(
+    logger: Logger,
+    collections: CollectionManager,
+    embeddings: EmbeddingProvider,
+    config: RetrievalConfig
+  ) {
+    this.logger = logger.child({ component: 'retrieval-pipeline' })
+    this.config = config
+    this.collections = collections
+    this.embeddings = embeddings
+
+    // Initialize BM25 engine
+    this.bm25Engine = new BM25Engine(logger)
+
+    // Initialize fusion strategy
+    this.fusionStrategy = createFusionStrategy({
+      method: config.fusion.method,
+      rrfK: config.fusion.rrfK,
+      denseWeight: config.dense.weight,
+      sparseWeight: config.sparse.weight
+    })
+  }
+
+  /**
+   * Initialize the pipeline
+   */
+  async initialize(): Promise<void> {
+    if (this.initialized) {
+      return
+    }
+
+    this.logger.info('Initializing retrieval pipeline')
+
+    // Initialize reranker if enabled
+    if (this.config.reranker.enabled) {
+      this.reranker = new Reranker(this.logger, {
+        model: this.config.reranker.model,
+        topK: this.config.reranker.topK
+      })
+      await this.reranker.initialize()
+      this.logger.info('Reranker initialized')
+    }
+
+    // Build BM25 index from existing documents
+    if (this.config.sparse.enabled) {
+      await this.buildBM25Index()
+    }
+
+    this.initialized = true
+    this.logger.info('Retrieval pipeline initialized')
+  }
+
+  /**
+   * Build BM25 index from ChromaDB collections
+   */
+  async buildBM25Index(): Promise<void> {
+    this.logger.info('Building BM25 index from ChromaDB')
+
+    const documents: BM25Document[] = []
+    const collectionsToIndex = [
+      COLLECTIONS.DECISIONS,
+      COLLECTIONS.MEMORIES,
+      COLLECTIONS.PATTERNS,
+      COLLECTIONS.CORRECTIONS
+    ]
+
+    for (const collectionName of collectionsToIndex) {
+      try {
+        const collection = await this.collections.getDecisions() // Use appropriate getter
+        const results = await collection.get({
+          include: ['documents', 'metadatas']
+        })
+
+        if (results.ids.length > 0) {
+          for (let i = 0; i < results.ids.length; i++) {
+            documents.push({
+              id: results.ids[i],
+              content: results.documents?.[i] || '',
+              metadata: results.metadatas?.[i] || {},
+              collection: collectionName
+            })
+          }
+        }
+      } catch (error) {
+        this.logger.warn({ error, collection: collectionName }, 'Failed to index collection')
+      }
+    }
+
+    await this.bm25Engine.buildIndex(documents)
+    this.indexBuilt = true
+    this.logger.info({ documentCount: documents.length }, 'BM25 index built')
+  }
+
+  /**
+   * Perform hybrid search
+   */
+  async search(
+    query: string,
+    options: HybridSearchOptions = {}
+  ): Promise<HybridSearchResult[]> {
+    const {
+      project,
+      collections = ['decisions', 'memories', 'patterns', 'corrections'],
+      limit = 10,
+      minScore = 0,
+      useReranker = this.config.reranker.enabled,
+      dateRange,
+      tags
+    } = options
+
+    this.logger.debug({
+      query: query.slice(0, 50),
+      options
+    }, 'Hybrid search')
+
+    // Step 1: Dense (semantic) search
+    const denseResults = await this.denseSearch(query, {
+      project,
+      collections,
+      limit: this.config.dense.limit,
+      minSimilarity: this.config.dense.minSimilarity,
+      dateRange,
+      tags
+    })
+
+    // Step 2: Sparse (BM25) search
+    let sparseResults: HybridSearchResult[] = []
+    if (this.config.sparse.enabled && this.indexBuilt) {
+      sparseResults = await this.sparseSearch(query, {
+        project,
+        collections,
+        limit: this.config.sparse.limit
+      })
+    }
+
+    // Step 3: Fusion
+    let fusedResults: HybridSearchResult[]
+    if (sparseResults.length > 0) {
+      fusedResults = this.fusionStrategy.fuse(denseResults, sparseResults)
+    } else {
+      // If no sparse results, use dense results directly
+      fusedResults = denseResults.map(r => ({
+        ...r,
+        scores: { ...r.scores, fusion: r.scores.dense, final: r.scores.dense }
+      }))
+    }
+
+    // Step 4: Reranking (optional)
+    if (useReranker && this.reranker && fusedResults.length > 0) {
+      fusedResults = await this.reranker.rerank(query, fusedResults)
+    }
+
+    // Filter by minimum score
+    if (minScore > 0) {
+      fusedResults = fusedResults.filter(r => r.scores.final >= minScore)
+    }
+
+    // Apply final limit
+    return fusedResults.slice(0, limit)
+  }
+
+  /**
+   * Dense (semantic) search using ChromaDB
+   */
+  private async denseSearch(
+    query: string,
+    options: {
+      project?: string
+      collections: string[]
+      limit: number
+      minSimilarity: number
+      dateRange?: { start?: string; end?: string }
+      tags?: string[]
+    }
+  ): Promise<HybridSearchResult[]> {
+    const results: HybridSearchResult[] = []
+
+    // Generate query embedding
+    const embedding = await this.embeddings.generate(query)
+
+    // Search each collection
+    for (const collectionName of options.collections) {
+      try {
+        const collection = await this.getCollection(collectionName)
+        if (!collection) continue
+
+        // Build where clause
+        const where: Record<string, any> = {}
+        if (options.project) {
+          where['project'] = { $eq: options.project }
+        }
+
+        const queryResults = await collection.query({
+          queryEmbeddings: [embedding],
+          nResults: Math.ceil(options.limit / options.collections.length),
+          where: Object.keys(where).length > 0 ? where : undefined,
+          include: ['documents', 'metadatas', 'distances']
+        })
+
+        // Process results
+        if (queryResults.ids[0]) {
+          for (let i = 0; i < queryResults.ids[0].length; i++) {
+            const distance = queryResults.distances?.[0]?.[i] || 0
+            const similarity = 1 - distance // Convert distance to similarity
+
+            if (similarity >= options.minSimilarity) {
+              results.push({
+                id: queryResults.ids[0][i],
+                content: queryResults.documents?.[0]?.[i] || '',
+                metadata: queryResults.metadatas?.[0]?.[i] || {},
+                collection: collectionName,
+                scores: {
+                  dense: similarity,
+                  sparse: 0,
+                  fusion: similarity,
+                  final: similarity
+                },
+                provenance: 'dense'
+              })
+            }
+          }
+        }
+      } catch (error) {
+        this.logger.error({ error, collection: collectionName }, 'Dense search failed')
+      }
+    }
+
+    // Sort by dense score
+    results.sort((a, b) => b.scores.dense - a.scores.dense)
+    return results
+  }
+
+  /**
+   * Sparse (BM25) search
+   */
+  private async sparseSearch(
+    query: string,
+    options: {
+      project?: string
+      collections: string[]
+      limit: number
+    }
+  ): Promise<HybridSearchResult[]> {
+    const bm25Results = this.bm25Engine.search(query, {
+      limit: options.limit,
+      filter: (result) => {
+        // Filter by project if specified
+        if (options.project && result.metadata.project !== options.project) {
+          return false
+        }
+        // Filter by collection
+        if (!options.collections.includes(result.collection)) {
+          return false
+        }
+        return true
+      }
+    })
+
+    return bm25Results.map(result => ({
+      id: result.id,
+      content: result.content,
+      metadata: result.metadata,
+      collection: result.collection,
+      scores: {
+        dense: 0,
+        sparse: result.score,
+        fusion: result.score,
+        final: result.score
+      },
+      provenance: 'sparse' as const
+    }))
+  }
+
+  /**
+   * Get collection by name
+   */
+  private async getCollection(name: string) {
+    switch (name) {
+      case 'decisions':
+        return this.collections.getDecisions()
+      case 'memories':
+        return this.collections.getMemories()
+      case 'patterns':
+        return this.collections.getPatterns()
+      case 'corrections':
+        return this.collections.getCorrections()
+      default:
+        return null
+    }
+  }
+
+  /**
+   * Add document to BM25 index
+   */
+  addToIndex(document: BM25Document): void {
+    if (this.config.sparse.enabled) {
+      this.bm25Engine.addDocument(document)
+    }
+  }
+
+  /**
+   * Remove document from BM25 index
+   */
+  removeFromIndex(document: BM25Document): void {
+    if (this.config.sparse.enabled) {
+      this.bm25Engine.removeDocument(document)
+    }
+  }
+
+  /**
+   * Rebuild BM25 index
+   */
+  async rebuildIndex(): Promise<void> {
+    await this.buildBM25Index()
+  }
+
+  /**
+   * Get pipeline status
+   */
+  getStatus(): {
+    initialized: boolean
+    indexBuilt: boolean
+    rerankerEnabled: boolean
+    sparseEnabled: boolean
+    fusionMethod: string
+    bm25Stats: { documentCount: number; termCount: number }
+  } {
+    return {
+      initialized: this.initialized,
+      indexBuilt: this.indexBuilt,
+      rerankerEnabled: this.reranker !== null,
+      sparseEnabled: this.config.sparse.enabled,
+      fusionMethod: this.config.fusion.method,
+      bm25Stats: this.bm25Engine.getStats()
+    }
+  }
+
+  /**
+   * Cleanup resources
+   */
+  cleanup(): void {
+    if (this.reranker) {
+      this.reranker.cleanup()
+    }
+    this.bm25Engine.clear()
+    this.initialized = false
+    this.indexBuilt = false
+  }
+}