mdcontext 0.0.1 → 0.2.0

package/src/search/bm25-store.ts

@@ -0,0 +1,366 @@
+/**
+ * BM25 Index Store for keyword search
+ *
+ * Uses wink-bm25-text-search for efficient keyword matching.
+ * Index is persisted to .mdcontext/bm25.json for fast startup.
+ */
+
+import * as fs from 'node:fs/promises'
+import * as path from 'node:path'
+import { Effect } from 'effect'
+import bm25 from 'wink-bm25-text-search'
+import { FileReadError, FileWriteError } from '../errors/index.js'
+import { INDEX_DIR } from '../index/types.js'
+
+// ============================================================================
+// Types
+// ============================================================================
+
+export interface BM25Document {
+  readonly id: string
+  readonly sectionId: string
+  readonly documentPath: string
+  readonly heading: string
+  readonly content: string
+}
+
+export interface BM25SearchResult {
+  readonly sectionId: string
+  readonly documentPath: string
+  readonly heading: string
+  readonly score: number
+  readonly rank: number
+}
+
+export interface BM25Stats {
+  readonly count: number
+  readonly lastUpdated: string
+}
+
+interface BM25Metadata {
+  readonly version: number
+  readonly count: number
+  readonly lastUpdated: string
+}
+
+// ============================================================================
+// Text Preparation
+// ============================================================================
+
+/**
+ * Simple tokenizer: lowercase, split on non-word chars, filter short tokens
+ */
+const tokenize = (text: string): string[] => {
+  return text
+    .toLowerCase()
+    .split(/\W+/)
+    .filter((token) => token.length > 2)
+}
+
+// ============================================================================
+// BM25 Store
+// ============================================================================
+
+export interface BM25Store {
+  /**
+   * Add documents to the index
+   */
+  add(docs: readonly BM25Document[]): Effect.Effect<void, never>
+
+  /**
+   * Consolidate the index (must be called after adding docs, before search)
+   */
+  consolidate(): Effect.Effect<void, never>
+
+  /**
+   * Search for matching documents
+   */
+  search(
+    query: string,
+    limit?: number,
+  ): Effect.Effect<readonly BM25SearchResult[], never>
+
+  /**
+   * Save the index to disk
+   */
+  save(): Effect.Effect<void, FileWriteError>
+
+  /**
+   * Load the index from disk
+   * @returns true if loaded successfully, false if no index exists
+   */
+  load(): Effect.Effect<boolean, FileReadError>
+
+  /**
+   * Get index statistics
+   */
+  getStats(): BM25Stats
+
+  /**
+   * Check if the index has been consolidated
+   */
+  isConsolidated(): boolean
+
+  /**
+   * Clear the index
+   */
+  clear(): void
+}
+
+/**
+ * Create a BM25 store for keyword search
+ */
+export const createBM25Store = (rootPath: string): BM25Store => {
+  const resolvedRoot = path.resolve(rootPath)
+  const indexPath = path.join(resolvedRoot, INDEX_DIR, 'bm25.json')
+  const metadataPath = path.join(resolvedRoot, INDEX_DIR, 'bm25.meta.json')
+
+  // Store mapping from internal index to section info
+  const sectionMap: Map<
+    number,
+    { sectionId: string; documentPath: string; heading: string }
+  > = new Map()
+  let documentCount = 0
+  let consolidated = false
+  let lastUpdated = new Date().toISOString()
+
+  // Create BM25 engine
+  let engine = bm25()
+
+  // Configure with weights for heading vs content
+  engine.defineConfig({
+    fldWeights: {
+      heading: 2,
+      content: 1,
+    },
+  })
+
+  // Define tokenization
+  engine.definePrepTasks([tokenize])
+
+  return {
+    add(docs: readonly BM25Document[]): Effect.Effect<void, never> {
+      return Effect.sync(() => {
+        for (const doc of docs) {
+          const idx = documentCount++
+          sectionMap.set(idx, {
+            sectionId: doc.sectionId,
+            documentPath: doc.documentPath,
+            heading: doc.heading,
+          })
+          engine.addDoc(
+            {
+              heading: doc.heading,
+              content: doc.content,
+            },
+            idx,
+          )
+        }
+        consolidated = false
+        lastUpdated = new Date().toISOString()
+      })
+    },
+
+    consolidate(): Effect.Effect<void, never> {
+      return Effect.sync(() => {
+        if (!consolidated && documentCount > 0) {
+          engine.consolidate()
+          consolidated = true
+        }
+      })
+    },
+
+    search(
+      query: string,
+      limit = 10,
+    ): Effect.Effect<readonly BM25SearchResult[], never> {
+      return Effect.sync(() => {
+        if (!consolidated || documentCount === 0) {
+          return []
+        }
+
+        const results = engine.search(query, limit) as [number, number][]
+
+        return results.map(([idx, score], rank) => {
+          const info = sectionMap.get(idx)
+          return {
+            sectionId: info?.sectionId ?? '',
+            documentPath: info?.documentPath ?? '',
+            heading: info?.heading ?? '',
+            score,
+            rank: rank + 1,
+          }
+        })
+      })
+    },
+
+    save(): Effect.Effect<void, FileWriteError> {
+      return Effect.gen(function* () {
+        // Export BM25 index
+        const jsonModel = engine.exportJSON()
+
+        // Save section map as array for JSON serialization
+        const sectionMapArray = Array.from(sectionMap.entries())
+
+        const data = {
+          engine: jsonModel,
+          sectionMap: sectionMapArray,
+        }
+
+        const metadata: BM25Metadata = {
+          version: 1,
+          count: documentCount,
+          lastUpdated,
+        }
+
+        yield* Effect.tryPromise({
+          try: async () => {
+            await fs.writeFile(indexPath, JSON.stringify(data), 'utf-8')
+            await fs.writeFile(
+              metadataPath,
+              JSON.stringify(metadata, null, 2),
+              'utf-8',
+            )
+          },
+          catch: (e) =>
+            new FileWriteError({
+              path: indexPath,
+              message: `Failed to save BM25 index: ${e instanceof Error ? e.message : String(e)}`,
+            }),
+        })
+      })
+    },
+
+    load(): Effect.Effect<boolean, FileReadError> {
+      return Effect.gen(function* () {
+        // Check if index exists
+        const exists = yield* Effect.promise(async () => {
+          try {
+            await fs.access(indexPath)
+            return true
+          } catch {
+            return false
+          }
+        })
+
+        if (!exists) {
+          return false
+        }
+
+        // Load data
+        const [dataStr, metaStr] = yield* Effect.tryPromise({
+          try: async () => {
+            const data = await fs.readFile(indexPath, 'utf-8')
+            const meta = await fs.readFile(metadataPath, 'utf-8')
+            return [data, meta] as const
+          },
+          catch: (e) =>
+            new FileReadError({
+              path: indexPath,
+              message: `Failed to load BM25 index: ${e instanceof Error ? e.message : String(e)}`,
+            }),
+        })
+
+        const data = JSON.parse(dataStr) as {
+          engine: string
+          sectionMap: [
+            number,
+            { sectionId: string; documentPath: string; heading: string },
+          ][]
+        }
+        const metadata = JSON.parse(metaStr) as BM25Metadata
+
+        // Restore engine
+        engine = bm25()
+        engine.importJSON(data.engine)
+        engine.definePrepTasks([tokenize])
+
+        // Restore section map
+        sectionMap.clear()
+        for (const [idx, info] of data.sectionMap) {
+          sectionMap.set(idx, info)
+        }
+
+        documentCount = metadata.count
+        lastUpdated = metadata.lastUpdated
+        consolidated = true
+
+        return true
+      })
+    },
+
+    getStats(): BM25Stats {
+      return {
+        count: documentCount,
+        lastUpdated,
+      }
+    },
+
+    isConsolidated(): boolean {
+      return consolidated
+    },
+
+    clear(): void {
+      engine = bm25()
+      engine.defineConfig({
+        fldWeights: {
+          heading: 2,
+          content: 1,
+        },
+      })
+      engine.definePrepTasks([tokenize])
+      sectionMap.clear()
+      documentCount = 0
+      consolidated = false
+      lastUpdated = new Date().toISOString()
+    },
+  }
+}
+
+// ============================================================================
+// BM25 Search Function
+// ============================================================================
+
+/**
+ * Perform BM25 keyword search over indexed sections.
+ *
+ * @param rootPath - Root directory containing BM25 index
+ * @param query - Search query text
+ * @param limit - Maximum results (default: 10)
+ * @returns Ranked list of matching sections by BM25 score
+ */
+export const bm25Search = (
+  rootPath: string,
+  query: string,
+  limit = 10,
+): Effect.Effect<readonly BM25SearchResult[], FileReadError> =>
+  Effect.gen(function* () {
+    const store = createBM25Store(rootPath)
+    const loaded = yield* store.load()
+
+    if (!loaded) {
+      return []
+    }
+
+    return yield* store.search(query, limit)
+  })
+
+// ============================================================================
+// Check BM25 Index Exists
+// ============================================================================
+
+/**
+ * Check if BM25 index exists for a directory
+ */
+export const bm25IndexExists = (rootPath: string): Effect.Effect<boolean> =>
+  Effect.promise(async () => {
+    const resolvedRoot = path.resolve(rootPath)
+    const indexPath = path.join(resolvedRoot, INDEX_DIR, 'bm25.json')
+
+    try {
+      await fs.access(indexPath)
+      return true
+    } catch {
+      return false
+    }
+  })

package/src/search/cross-encoder.test.ts

@@ -0,0 +1,253 @@
+/**
+ * Tests for cross-encoder re-ranking module
+ */
+
+import { Effect, Exit } from 'effect'
+import { afterEach, describe, expect, it, vi } from 'vitest'
+import {
+  getReranker,
+  isRerankerAvailable,
+  type RerankedResult,
+  RerankerError,
+  type RerankerOptions,
+  rerankResults,
+  unloadReranker,
+} from './cross-encoder.js'
+
+describe('cross-encoder', () => {
+  afterEach(() => {
+    unloadReranker()
+    vi.restoreAllMocks()
+  })
+
+  describe('RerankerError', () => {
+    it('should create error with DependencyMissing reason', () => {
+      const error = new RerankerError(
+        'DependencyMissing',
+        'Package not installed',
+      )
+      expect(error.reason).toBe('DependencyMissing')
+      expect(error.message).toBe('Package not installed')
+      expect(error._tag).toBe('RerankerError')
+      expect(error.name).toBe('RerankerError')
+    })
+
+    it('should create error with ModelLoadFailed reason', () => {
+      const error = new RerankerError('ModelLoadFailed', 'Failed to load model')
+      expect(error.reason).toBe('ModelLoadFailed')
+      expect(error.message).toBe('Failed to load model')
+    })
+
+    it('should create error with InferenceFailed reason', () => {
+      const cause = new Error('GPU OOM')
+      const error = new RerankerError(
+        'InferenceFailed',
+        'Inference failed',
+        cause,
+      )
+      expect(error.reason).toBe('InferenceFailed')
+      expect(error.cause).toBe(cause)
+    })
+  })
+
+  describe('getReranker', () => {
+    it('should return a singleton reranker instance', () => {
+      const reranker1 = getReranker()
+      const reranker2 = getReranker()
+      expect(reranker1).toBe(reranker2)
+    })
+
+    it('should return reranker with expected interface', () => {
+      const reranker = getReranker()
+      expect(typeof reranker.rerank).toBe('function')
+      expect(typeof reranker.isReady).toBe('function')
+      expect(typeof reranker.unload).toBe('function')
+    })
+
+    it('should report not ready before loading', () => {
+      const reranker = getReranker()
+      expect(reranker.isReady()).toBe(false)
+    })
+
+    it('should accept custom cache directory', () => {
+      unloadReranker()
+      const reranker = getReranker('/custom/cache/dir')
+      expect(reranker).toBeDefined()
+    })
+  })
+
+  describe('unloadReranker', () => {
+    it('should reset singleton', () => {
+      const reranker1 = getReranker()
+      unloadReranker()
+      const reranker2 = getReranker()
+      expect(reranker1).not.toBe(reranker2)
+    })
+
+    it('should be safe to call multiple times', () => {
+      unloadReranker()
+      unloadReranker()
+      unloadReranker()
+      // No error should be thrown
+    })
+  })
+
+  describe('rerankResults', () => {
+    it('should return empty array for empty input', async () => {
+      const result = await Effect.runPromise(
+        rerankResults('test query', [], (item: unknown) => String(item)),
+      )
+      expect(result).toEqual([])
+    })
+
+    it('should handle DependencyMissing error gracefully', async () => {
+      // When transformers is not installed, the effect should fail with DependencyMissing
+      const results = [
+        { id: 1, content: 'test content' },
+        { id: 2, content: 'another content' },
+      ]
+
+      const effect = rerankResults(
+        'test query',
+        results,
+        (item) => item.content,
+      )
+
+      const exit = await Effect.runPromiseExit(effect)
+
+      // This test will either:
+      // 1. Fail with DependencyMissing if @huggingface/transformers is not installed
+      // 2. Succeed if the package is installed
+      if (Exit.isFailure(exit)) {
+        const error = exit.cause
+        // Check if it's the right kind of error
+        if ('_tag' in error) {
+          // It should be a RerankerError
+          expect(true).toBe(true)
+        }
+      } else {
+        // If it succeeds, that's fine too (package is installed)
+        expect(Exit.isSuccess(exit)).toBe(true)
+      }
+    })
+
+    it('should respect topK option', async () => {
+      const results = Array.from({ length: 30 }, (_, i) => ({
+        id: i,
+        content: `content ${i}`,
+      }))
+
+      const options: RerankerOptions = { topK: 5, returnTopN: 3 }
+
+      // This test verifies the options are passed through
+      // The actual slicing happens in the effect implementation
+      const effect = rerankResults(
+        'test query',
+        results,
+        (item) => item.content,
+        options,
+      )
+
+      // Just verify the effect can be created without error
+      expect(effect).toBeDefined()
+    })
+
+    it('should respect returnTopN option', async () => {
+      const results = Array.from({ length: 10 }, (_, i) => ({
+        id: i,
+        content: `content ${i}`,
+      }))
+
+      const options: RerankerOptions = { topK: 10, returnTopN: 5 }
+
+      const effect = rerankResults(
+        'test query',
+        results,
+        (item) => item.content,
+        options,
+      )
+
+      expect(effect).toBeDefined()
+    })
+
+    it('should use default topK of 20', async () => {
+      // Default options should be topK=20, returnTopN=10
+      const results = Array.from({ length: 25 }, (_, i) => ({
+        id: i,
+        content: `content ${i}`,
+      }))
+
+      const effect = rerankResults(
+        'test query',
+        results,
+        (item) => item.content,
+        // No options, should use defaults
+      )
+
+      expect(effect).toBeDefined()
+    })
+
+    it('should use default returnTopN of 10', async () => {
+      const results = Array.from({ length: 15 }, (_, i) => ({
+        id: i,
+        content: `content ${i}`,
+      }))
+
+      const effect = rerankResults(
+        'test query',
+        results,
+        (item) => item.content,
+      )
+
+      expect(effect).toBeDefined()
+    })
+  })
+
+  describe('isRerankerAvailable', () => {
+    it('should return boolean effect', async () => {
+      const effect = isRerankerAvailable()
+      const result = await Effect.runPromise(effect)
+      expect(typeof result).toBe('boolean')
+    })
+
+    it('should not throw on missing dependency', async () => {
+      // This should gracefully return false if package is not installed
+      const result = await Effect.runPromise(isRerankerAvailable())
+      // Result is either true or false, never throws
+      expect([true, false]).toContain(result)
+    })
+  })
+
+  describe('RerankedResult type', () => {
+    it('should have expected structure', () => {
+      const result: RerankedResult<{ id: number }> = {
+        item: { id: 1 },
+        rerankerScore: 0.95,
+        originalRank: 1,
+      }
+
+      expect(result.item.id).toBe(1)
+      expect(result.rerankerScore).toBe(0.95)
+      expect(result.originalRank).toBe(1)
+    })
+  })
+
+  describe('Reranker interface', () => {
+    it('should define rerank method signature', () => {
+      const reranker = getReranker()
+
+      // Verify the rerank method exists and is async
+      expect(typeof reranker.rerank).toBe('function')
+    })
+
+    it('should define isReady method', () => {
+      const reranker = getReranker()
+      expect(typeof reranker.isReady).toBe('function')
+    })
+
+    it('should define unload method', () => {
+      const reranker = getReranker()
+      expect(typeof reranker.unload).toBe('function')
+    })
+  })
+})