mdcontext 0.1.0 → 0.2.0

package/src/embeddings/vector-store.ts

@@ -1,12 +1,28 @@
 /**
  * Vector store using hnswlib-node
+ *
+ * Supports both legacy (flat) and namespaced storage layouts:
+ * - Legacy: .mdcontext/vectors.bin, .mdcontext/vectors.meta.bin
+ * - Namespaced: .mdcontext/embeddings/{namespace}/vectors.bin, vectors.meta.bin
+ *
+ * New indexes are written using namespaced storage. Existing legacy indexes
+ * continue to be loaded from their original flat locations; this module does
+ * not perform automatic migration between layouts.
  */
 
 import * as fs from 'node:fs/promises'
 import * as path from 'node:path'
+import * as msgpack from '@msgpack/msgpack'
 import { Effect } from 'effect'
 import HierarchicalNSW from 'hnswlib-node'
+import { DimensionMismatchError, VectorStoreError } from '../errors/index.js'
 import { INDEX_DIR } from '../index/types.js'
+import {
+  generateNamespace,
+  getNamespaceDir,
+  getMetaPath as getNamespacedMetaPath,
+  getVectorPath as getNamespacedVectorPath,
+} from './embedding-namespace.js'
 import type { VectorEntry, VectorIndex } from './types.js'
 
 // ============================================================================
@@ -14,24 +30,49 @@ import type { VectorEntry, VectorIndex } from './types.js'
 // ============================================================================
 
 const VECTOR_INDEX_FILE = 'vectors.bin'
-const VECTOR_META_FILE = 'vectors.meta.json'
+const VECTOR_META_FILE = 'vectors.meta.bin'
 const INDEX_VERSION = 1
 
 // ============================================================================
 // Vector Store
 // ============================================================================
 
+export interface VectorSearchOptions {
+  /** efSearch parameter for HNSW (controls recall/speed tradeoff, default: 100) */
+  readonly efSearch?: number | undefined
+}
+
 export interface VectorStore {
   readonly rootPath: string
   readonly dimensions: number
-  add(entries: VectorEntry[]): Effect.Effect<void, Error>
+  add(entries: VectorEntry[]): Effect.Effect<void, VectorStoreError>
   search(
     vector: number[],
     limit: number,
     threshold?: number,
-  ): Effect.Effect<VectorSearchResult[], Error>
-  save(): Effect.Effect<void, Error>
-  load(): Effect.Effect<boolean, Error>
+    options?: VectorSearchOptions,
+  ): Effect.Effect<VectorSearchResult[], VectorStoreError>
+  /**
+   * Search with additional stats about below-threshold results.
+   * Used to provide feedback when 0 results pass the threshold.
+   */
+  searchWithStats(
+    vector: number[],
+    limit: number,
+    threshold?: number,
+    options?: VectorSearchOptions,
+  ): Effect.Effect<VectorSearchResultWithStats, VectorStoreError>
+  save(): Effect.Effect<void, VectorStoreError>
+  /**
+   * Load the vector store from disk.
+   *
+   * @returns VectorStoreLoadResult with loaded status and any warnings
+   * @throws DimensionMismatchError if the stored dimensions don't match current provider
+   */
+  load(): Effect.Effect<
+    VectorStoreLoadResult,
+    VectorStoreError | DimensionMismatchError
+  >
   getStats(): VectorStoreStats
 }
 
@@ -43,14 +84,48 @@ export interface VectorSearchResult {
   readonly similarity: number
 }
 
+/**
+ * Extended search result with metadata about below-threshold results.
+ * Used to provide user feedback when 0 results pass the threshold.
+ */
+export interface VectorSearchResultWithStats {
+  readonly results: VectorSearchResult[]
+  /** Number of results that were found but below threshold */
+  readonly belowThresholdCount: number
+  /** Highest similarity score among below-threshold results (if any) */
+  readonly belowThresholdHighest: number | null
+}
+
 export interface VectorStoreStats {
   readonly count: number
   readonly dimensions: number
   readonly provider: string
+  readonly providerModel?: string | undefined
   readonly totalCost: number
   readonly totalTokens: number
 }
 
+/**
+ * Result of loading a vector store, including any warnings about config mismatches.
+ */
+export interface VectorStoreLoadResult {
+  /** Whether the index was loaded successfully */
+  readonly loaded: boolean
+  /** Warning about HNSW parameter mismatch (if any) */
+  readonly hnswMismatch?: HnswMismatchWarning | undefined
+}
+
+/**
+ * Warning when HNSW parameters in config differ from stored index parameters.
+ * The index was built with different parameters than currently configured.
+ */
+export interface HnswMismatchWarning {
+  /** Current config values */
+  readonly configParams: { m: number; efConstruction: number }
+  /** Values stored in the index */
+  readonly indexParams: { m: number; efConstruction: number }
+}
+
 // ============================================================================
 // Implementation
 // ============================================================================
@@ -64,24 +139,73 @@ class HnswVectorStore implements VectorStore {
   private idToIndex: Map<string, number> = new Map()
   private nextIndex = 0
   private provider = 'unknown'
+  private providerModel: string | undefined = undefined
+  private providerBaseURL: string | undefined = undefined
   private totalCost = 0
   private totalTokens = 0
 
-  constructor(rootPath: string, dimensions: number) {
+  // HNSW build parameters
+  private readonly hnswM: number
+  private readonly hnswEfConstruction: number
+
+  // Namespace support - when set, uses namespaced storage paths
+  private namespace: string | undefined = undefined
+
+  constructor(
+    rootPath: string,
+    dimensions: number,
+    hnswOptions?: HnswBuildOptions,
+  ) {
     this.rootPath = path.resolve(rootPath)
     this.dimensions = dimensions
+    this.hnswM = hnswOptions?.m ?? 16
+    this.hnswEfConstruction = hnswOptions?.efConstruction ?? 200
+  }
+
+  /**
+   * Set the namespace for this vector store.
+   * When set, all storage operations use the namespaced path.
+   */
+  setNamespace(namespace: string): void {
+    this.namespace = namespace
+  }
+
+  /**
+   * Get the current namespace (if any).
+   */
+  getNamespace(): string | undefined {
+    return this.namespace
   }
 
+  /**
+   * Get the index directory path.
+   * Returns namespaced path if namespace is set, otherwise legacy path.
+   */
   private getIndexDir(): string {
+    if (this.namespace) {
+      return getNamespaceDir(this.rootPath, this.namespace)
+    }
     return path.join(this.rootPath, INDEX_DIR)
   }
 
+  /**
+   * Get the vector index file path.
+   */
   private getVectorPath(): string {
-    return path.join(this.getIndexDir(), VECTOR_INDEX_FILE)
+    if (this.namespace) {
+      return getNamespacedVectorPath(this.rootPath, this.namespace)
+    }
+    return path.join(this.rootPath, INDEX_DIR, VECTOR_INDEX_FILE)
   }
 
+  /**
+   * Get the metadata file path.
+   */
   private getMetaPath(): string {
-    return path.join(this.getIndexDir(), VECTOR_META_FILE)
+    if (this.namespace) {
+      return getNamespacedMetaPath(this.rootPath, this.namespace)
+    }
+    return path.join(this.rootPath, INDEX_DIR, VECTOR_META_FILE)
   }
 
   private ensureIndex(): HierarchicalNSW.HierarchicalNSW {
@@ -91,32 +215,41 @@ class HnswVectorStore implements VectorStore {
         'cosine',
         this.dimensions,
       )
-      this.index.initIndex(10000, 16, 200, 100)
+      // Use configured HNSW parameters (M, efConstruction, randomSeed)
+      this.index.initIndex(10000, this.hnswM, this.hnswEfConstruction, 100)
     }
     return this.index
   }
 
-  add(entries: VectorEntry[]): Effect.Effect<void, Error> {
-    return Effect.sync(() => {
-      const index = this.ensureIndex()
+  add(entries: VectorEntry[]): Effect.Effect<void, VectorStoreError> {
+    return Effect.try({
+      try: () => {
+        const index = this.ensureIndex()
 
-      for (const entry of entries) {
-        // Skip if already exists
-        if (this.idToIndex.has(entry.id)) {
-          continue
-        }
+        for (const entry of entries) {
+          // Skip if already exists
+          if (this.idToIndex.has(entry.id)) {
+            continue
+          }
 
-        const idx = this.nextIndex++
+          const idx = this.nextIndex++
 
-        // Resize if needed
-        if (idx >= index.getMaxElements()) {
-          index.resizeIndex(index.getMaxElements() * 2)
-        }
+          // Resize if needed
+          if (idx >= index.getMaxElements()) {
+            index.resizeIndex(index.getMaxElements() * 2)
+          }
 
-        index.addPoint(entry.embedding as number[], idx)
-        this.entries.set(idx, entry)
-        this.idToIndex.set(entry.id, idx)
-      }
+          index.addPoint(entry.embedding as number[], idx)
+          this.entries.set(idx, entry)
+          this.idToIndex.set(entry.id, idx)
+        }
+      },
+      catch: (e) =>
+        new VectorStoreError({
+          operation: 'add',
+          message: e instanceof Error ? e.message : String(e),
+          cause: e,
+        }),
     })
   }
 
@@ -124,36 +257,120 @@ class HnswVectorStore implements VectorStore {
     vector: number[],
     limit: number,
     threshold = 0,
-  ): Effect.Effect<VectorSearchResult[], Error> {
-    return Effect.sync(() => {
-      if (!this.index || this.entries.size === 0) {
-        return []
-      }
-
-      const result = this.index.searchKnn(
-        vector,
-        Math.min(limit, this.entries.size),
-      )
-      const results: VectorSearchResult[] = []
+    options?: VectorSearchOptions,
+  ): Effect.Effect<VectorSearchResult[], VectorStoreError> {
+    return Effect.try({
+      try: () => {
+        if (!this.index || this.entries.size === 0) {
+          return []
+        }
 
-      for (let i = 0; i < result.neighbors.length; i++) {
-        const idx = result.neighbors[i]
-        const distance = result.distances[i]
+        // Set efSearch if provided (controls recall/speed tradeoff)
+        if (options?.efSearch !== undefined) {
+          this.index.setEf(options.efSearch)
+        }
 
-        if (idx === undefined || distance === undefined) {
-          continue
+        const result = this.index.searchKnn(
+          vector,
+          Math.min(limit, this.entries.size),
+        )
+        const results: VectorSearchResult[] = []
+
+        for (let i = 0; i < result.neighbors.length; i++) {
+          const idx = result.neighbors[i]
+          const distance = result.distances[i]
+
+          if (idx === undefined || distance === undefined) {
+            continue
+          }
+
+          // Convert distance to similarity (cosine distance to cosine similarity)
+          // hnswlib returns 1 - cosine_similarity for cosine space
+          const similarity = 1 - distance
+
+          if (similarity < threshold) {
+            continue
+          }
+
+          const entry = this.entries.get(idx)
+          if (entry) {
+            results.push({
+              id: entry.id,
+              sectionId: entry.sectionId,
+              documentPath: entry.documentPath,
+              heading: entry.heading,
+              similarity,
+            })
+          }
         }
 
-        // Convert distance to similarity (cosine distance to cosine similarity)
-        // hnswlib returns 1 - cosine_similarity for cosine space
-        const similarity = 1 - distance
+        return results
+      },
+      catch: (e) =>
+        new VectorStoreError({
+          operation: 'search',
+          message: e instanceof Error ? e.message : String(e),
+          cause: e,
+        }),
+    })
+  }
 
-        if (similarity < threshold) {
-          continue
+  searchWithStats(
+    vector: number[],
+    limit: number,
+    threshold = 0,
+    options?: VectorSearchOptions,
+  ): Effect.Effect<VectorSearchResultWithStats, VectorStoreError> {
+    return Effect.try({
+      try: () => {
+        if (!this.index || this.entries.size === 0) {
+          return {
+            results: [],
+            belowThresholdCount: 0,
+            belowThresholdHighest: null,
+          }
         }
 
-        const entry = this.entries.get(idx)
-        if (entry) {
+        // Set efSearch if provided (controls recall/speed tradeoff)
+        if (options?.efSearch !== undefined) {
+          this.index.setEf(options.efSearch)
+        }
+
+        const result = this.index.searchKnn(
+          vector,
+          Math.min(limit, this.entries.size),
+        )
+        const results: VectorSearchResult[] = []
+        let belowThresholdCount = 0
+        let belowThresholdHighest: number | null = null
+
+        for (let i = 0; i < result.neighbors.length; i++) {
+          const idx = result.neighbors[i]
+          const distance = result.distances[i]
+
+          if (idx === undefined || distance === undefined) {
+            continue
+          }
+
+          // Convert distance to similarity (cosine distance to cosine similarity)
+          // hnswlib returns 1 - cosine_similarity for cosine space
+          const similarity = 1 - distance
+
+          const entry = this.entries.get(idx)
+          if (!entry) continue
+
+          if (similarity < threshold) {
+            // Track below-threshold stats
+            belowThresholdCount++
+            if (
+              belowThresholdHighest === null ||
+              similarity > belowThresholdHighest
+            ) {
+              belowThresholdHighest = similarity
+            }
+            continue
+          }
+
           results.push({
             id: entry.id,
             sectionId: entry.sectionId,
@@ -162,13 +379,23 @@ class HnswVectorStore implements VectorStore {
             similarity,
           })
         }
-      }
 
-      return results
+        return {
+          results,
+          belowThresholdCount,
+          belowThresholdHighest,
+        }
+      },
+      catch: (e) =>
+        new VectorStoreError({
+          operation: 'search',
+          message: e instanceof Error ? e.message : String(e),
+          cause: e,
+        }),
     })
   }
 
-  save(): Effect.Effect<void, Error> {
+  save(): Effect.Effect<void, VectorStoreError> {
     return Effect.gen(
       function* (this: HnswVectorStore) {
         if (!this.index) {
@@ -176,17 +403,33 @@ class HnswVectorStore implements VectorStore {
         }
 
         const indexDir = this.getIndexDir()
-        yield* Effect.promise(() => fs.mkdir(indexDir, { recursive: true }))
+        yield* Effect.tryPromise({
+          try: () => fs.mkdir(indexDir, { recursive: true }),
+          catch: (e) =>
+            new VectorStoreError({
+              operation: 'save',
+              message: `Failed to create directory: ${e instanceof Error ? e.message : String(e)}`,
+              cause: e,
+            }),
+        })
 
         // Save the hnswlib index
-        yield* Effect.promise(() =>
-          this.index!.writeIndex(this.getVectorPath()),
-        )
+        yield* Effect.tryPromise({
+          try: () => this.index!.writeIndex(this.getVectorPath()),
+          catch: (e) =>
+            new VectorStoreError({
+              operation: 'save',
+              message: `Failed to write index: ${e instanceof Error ? e.message : String(e)}`,
+              cause: e,
+            }),
+        })
 
         // Save metadata
         const meta: VectorIndex = {
           version: INDEX_VERSION,
           provider: this.provider,
+          providerModel: this.providerModel,
+          providerBaseURL: this.providerBaseURL,
           dimensions: this.dimensions,
           entries: Object.fromEntries(
             Array.from(this.entries.entries()).map(([idx, entry]) => [
@@ -198,44 +441,135 @@ class HnswVectorStore implements VectorStore {
           totalTokens: this.totalTokens,
           createdAt: new Date().toISOString(),
           updatedAt: new Date().toISOString(),
+          // Store HNSW build parameters for validation on load
+          hnswParams: {
+            m: this.hnswM,
+            efConstruction: this.hnswEfConstruction,
+          },
         }
 
-        yield* Effect.promise(() =>
-          fs.writeFile(this.getMetaPath(), JSON.stringify(meta, null, 2)),
-        )
+        yield* Effect.tryPromise({
+          try: async () => {
+            // Size validation
+            const estimatedSize = this.entries.size * 15000
+            if (estimatedSize > 100_000_000) {
+              console.warn(
+                `Large metadata detected: ~${(estimatedSize / 1e6).toFixed(0)}MB. ` +
+                  `Consider indexing subdirectories separately.`,
+              )
+            }
+
+            // Encode with MessagePack and write
+            const encoded = msgpack.encode(meta)
+            await fs.writeFile(this.getMetaPath(), encoded)
+          },
+          catch: (e) =>
+            new VectorStoreError({
+              operation: 'save',
+              message: `Failed to write metadata: ${e instanceof Error ? e.message : String(e)}`,
+              cause: e,
+            }),
+        })
       }.bind(this),
     )
   }
 
-  load(): Effect.Effect<boolean, Error> {
+  load(): Effect.Effect<
+    VectorStoreLoadResult,
+    VectorStoreError | DimensionMismatchError
+  > {
     return Effect.gen(
       function* (this: HnswVectorStore) {
         const vectorPath = this.getVectorPath()
         const metaPath = this.getMetaPath()
 
-        // Check if files exist
+        // Check if files exist - catch file not found gracefully
+        // For metadata, check both binary (.bin) and JSON (.json) for migration
         const filesExist = yield* Effect.tryPromise({
           try: async () => {
             await fs.access(vectorPath)
-            await fs.access(metaPath)
-            return true
+            // Check if either binary or JSON metadata exists
+            try {
+              await fs.access(metaPath)
+              return true
+            } catch {
+              const jsonPath = metaPath.replace('.bin', '.json')
+              await fs.access(jsonPath)
+              return true
+            }
           },
-          catch: () => false as const,
-        }).pipe(Effect.catchAll(() => Effect.succeed(false)))
+          catch: () =>
+            new VectorStoreError({
+              operation: 'load',
+              message: 'Files not found',
+            }),
+        }).pipe(
+          Effect.catchTag('VectorStoreError', () => Effect.succeed(false)),
+        )
 
         if (!filesExist) {
-          return false
+          return { loaded: false }
         }
 
-        // Load metadata first
-        const metaContent = yield* Effect.promise(() =>
-          fs.readFile(metaPath, 'utf-8'),
-        )
-        const meta = JSON.parse(metaContent) as VectorIndex
+        // Load metadata - try binary first, fall back to JSON for migration
+        const loadedMeta = yield* Effect.tryPromise({
+          try: async () => {
+            // Try binary format first (new)
+            try {
+              await fs.access(metaPath)
+              const buffer = await fs.readFile(metaPath)
+              return msgpack.decode(buffer) as VectorIndex
+            } catch {
+              // Fall back to JSON for migration (old)
+              const jsonPath = metaPath.replace('.bin', '.json')
+              try {
+                await fs.access(jsonPath)
+                const json = await fs.readFile(jsonPath, 'utf-8')
+                const meta = JSON.parse(json) as VectorIndex
+
+                // Auto-migrate to binary format (safe for concurrent access)
+                try {
+                  const encoded = msgpack.encode(meta)
+                  await fs.writeFile(metaPath, encoded)
+
+                  // Remove old JSON file (ignore errors if already deleted by another process)
+                  await fs.unlink(jsonPath).catch(() => {})
+                } catch {
+                  // Migration failed, but we have the data - continue
+                }
+
+                return meta
+              } catch {
+                throw new Error('Metadata file not found')
+              }
+            }
+          },
+          catch: (e) =>
+            new VectorStoreError({
+              operation: 'load',
+              message: `Failed to read metadata: ${e instanceof Error ? e.message : String(e)}`,
+              cause: e,
+            }),
+        })
+
+        // Apply legacy index migration: default to 'openai' if provider is missing
+        const meta: VectorIndex = {
+          ...loadedMeta,
+          provider: loadedMeta.provider || 'openai',
+        }
 
-        // Verify dimensions match
+        // Verify dimensions match - fail with clear error if mismatch
         if (meta.dimensions !== this.dimensions) {
-          return false
+          return yield* Effect.fail(
+            new DimensionMismatchError({
+              corpusDimensions: meta.dimensions,
+              providerDimensions: this.dimensions,
+              corpusProvider: meta.providerModel
+                ? `${meta.provider}:${meta.providerModel}`
+                : meta.provider,
+              path: this.rootPath,
+            }),
+          )
         }
 
         // Load the hnswlib index
@@ -243,7 +577,15 @@ class HnswVectorStore implements VectorStore {
           'cosine',
           this.dimensions,
         )
-        yield* Effect.promise(() => this.index!.readIndex(vectorPath))
+        yield* Effect.tryPromise({
+          try: () => this.index!.readIndex(vectorPath),
+          catch: (e) =>
+            new VectorStoreError({
+              operation: 'load',
+              message: `Failed to read index: ${e instanceof Error ? e.message : String(e)}`,
+              cause: e,
+            }),
+        })
 
         // Restore entries
         this.entries.clear()
@@ -258,10 +600,28 @@ class HnswVectorStore implements VectorStore {
         }
 
         this.provider = meta.provider
+        this.providerModel = meta.providerModel
+        this.providerBaseURL = meta.providerBaseURL
         this.totalCost = meta.totalCost
         this.totalTokens = meta.totalTokens
 
-        return true
+        // Check for HNSW parameter mismatch
+        let hnswMismatch: HnswMismatchWarning | undefined
+        if (meta.hnswParams) {
+          const indexM = meta.hnswParams.m
+          const indexEf = meta.hnswParams.efConstruction
+          if (indexM !== this.hnswM || indexEf !== this.hnswEfConstruction) {
+            hnswMismatch = {
+              configParams: {
+                m: this.hnswM,
+                efConstruction: this.hnswEfConstruction,
+              },
+              indexParams: { m: indexM, efConstruction: indexEf },
+            }
+          }
+        }
+
+        return { loaded: true, hnswMismatch }
       }.bind(this),
     )
   }
@@ -271,13 +631,16 @@ class HnswVectorStore implements VectorStore {
       count: this.entries.size,
       dimensions: this.dimensions,
       provider: this.provider,
+      providerModel: this.providerModel,
       totalCost: this.totalCost,
       totalTokens: this.totalTokens,
     }
   }
 
-  setProvider(name: string): void {
+  setProvider(name: string, model?: string, baseURL?: string): void {
     this.provider = name
+    this.providerModel = model
+    this.providerBaseURL = baseURL
   }
 
   addCost(cost: number, tokens: number): void {
@@ -290,10 +653,56 @@ class HnswVectorStore implements VectorStore {
 // Factory
 // ============================================================================
 
+/**
+ * HNSW build parameters for index construction.
+ * These affect index quality and build time - changes require index rebuild.
+ */
+export interface HnswBuildOptions {
+  /** Max connections per node (default: 16). Higher = better recall, larger index. */
+  readonly m?: number | undefined
+  /** Construction-time search width (default: 200). Higher = better quality, slower builds. */
+  readonly efConstruction?: number | undefined
+}
+
+/**
+ * Create a vector store for the given root path.
+ *
+ * @param rootPath - Root directory containing the index
+ * @param dimensions - Embedding dimensions
+ * @param hnswOptions - Optional HNSW build parameters
+ * @returns A new VectorStore instance
+ */
 export const createVectorStore = (
   rootPath: string,
   dimensions: number,
-): VectorStore => new HnswVectorStore(rootPath, dimensions)
+  hnswOptions?: HnswBuildOptions,
+): VectorStore => new HnswVectorStore(rootPath, dimensions, hnswOptions)
+
+/**
+ * Create a namespaced vector store for a specific provider/model.
+ *
+ * Uses the new namespaced storage structure:
+ * .mdcontext/embeddings/{provider}_{model}_{dimensions}/vectors.bin
+ *
+ * @param rootPath - Root directory containing the index
+ * @param provider - Provider name (e.g., "openai", "voyage")
+ * @param model - Model name (e.g., "text-embedding-3-small")
+ * @param dimensions - Embedding dimensions
+ * @param hnswOptions - Optional HNSW build parameters
+ * @returns A new VectorStore instance with namespace set
+ */
+export const createNamespacedVectorStore = (
+  rootPath: string,
+  provider: string,
+  model: string,
+  dimensions: number,
+  hnswOptions?: HnswBuildOptions,
+): VectorStore => {
+  const namespace = generateNamespace(provider, model, dimensions)
+  const store = new HnswVectorStore(rootPath, dimensions, hnswOptions)
+  store.setNamespace(namespace)
+  return store
+}
 
 // Export the class for type access
 export { HnswVectorStore }