Merge pull request #60 from rryam/audit-fix-hardening

Harden storage consistency and validation

Author: rudrankriyam

Docs/INDEXED_STORAGE_GUIDE.md

@@ -47,7 +47,7 @@ let vectura = try await VecturaKit(config: config)
 **Behavior:**
 - < 10,000 documents → Uses `fullMemory` mode
 - ≥ 10,000 documents → Uses `indexed` mode (if storage supports it) with default parameters:
-  - `candidateMultiplier: 10`
+  - `candidateMultiplier: 4`
   - `batchSize: 100`
   - `maxConcurrentBatches: 4`
 
@@ -79,7 +79,7 @@ Use indexed mode for large datasets:
 let config = VecturaConfig(
     name: "my-database",
     memoryStrategy: .indexed(
-        candidateMultiplier: 10      // Search 10× topK candidates
+        candidateMultiplier: 4       // Search 4× topK candidates
     )
 )
 

Package.resolved

@@ -37,7 +37,7 @@ let package = Package(
   dependencies: [
     .package(url: "https://github.com/jkrukowski/swift-embeddings.git", from: "0.0.21"),
     .package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.4.0"),
-    .package(url: "https://github.com/ml-explore/mlx-swift-lm/", from: "2.29.2"),
+    .package(url: "https://github.com/ml-explore/mlx-swift-lm/", from: "2.30.3"),
   ],
   targets: [
     .target(

Package.swift

@@ -65,7 +65,7 @@ struct VecturaCLI: AsyncParsableCommand {
     }
   }
 
-  static func setupDB(dbName: String, dimension: Int, numResults: Int, threshold: Float, modelId: String) async throws
+  static func setupDB(dbName: String, dimension: Int?, numResults: Int, threshold: Float, modelId: String) async throws
   -> VecturaKit {
     let config = try VecturaConfig(
       name: dbName,
@@ -90,8 +90,8 @@ extension VecturaCLI {
     @Option(name: [.long, .customShort("d")], help: "Database name")
     var dbName: String = "vectura-cli-demo-db"
 
-    @Option(name: [.long, .customShort("v")], help: "Vector dimension")
-    var dimension: Int = 512
+    @Option(name: [.long, .customShort("v")], help: "Vector dimension (auto-detected if not specified)")
+    var dimension: Int?
 
     @Option(name: [.long, .customShort("t")], help: "Minimum similarity threshold for searches")
     var threshold: Float = 0.5
@@ -252,7 +252,8 @@ extension VecturaCLI {
       print("Total search queries: \(searchQueries.count)")
       print("Average search time: \(String(format: "%.1f", avgSearchTime * 1000))ms")
       print("Model: \(modelId)")
-      print("Vector dimension: \(dimension)")
+      let dimensionDescription = dimension.map(String.init) ?? "auto-detected"
+      print("Vector dimension: \(dimensionDescription)")
     }
 
     private func loadMockDataset() throws -> MockDataset {
@@ -275,8 +276,8 @@ extension VecturaCLI {
     @Option(name: [.long, .customShort("d")], help: "Database name")
     var dbName: String = "vectura-cli-db"
 
-    @Option(name: [.long, .customShort("v")], help: "Vector dimension")
-    var dimension: Int = 384
+    @Option(name: [.long, .customShort("v")], help: "Vector dimension (auto-detected if not specified)")
+    var dimension: Int?
 
     @Option(name: [.long, .customShort("m")], help: "Model ID for embeddings")
     var modelId: String = "sentence-transformers/all-MiniLM-L6-v2"
@@ -310,8 +311,8 @@ extension VecturaCLI {
     @Option(name: [.long, .customShort("d")], help: "Database name")
     var dbName: String = "vectura-cli-db"
 
-    @Option(name: [.long, .customShort("v")], help: "Vector dimension")
-    var dimension: Int = 384
+    @Option(name: [.long, .customShort("v")], help: "Vector dimension (auto-detected if not specified)")
+    var dimension: Int?
 
     @Option(name: [.long, .customShort("t")], help: "Minimum similarity threshold")
     var threshold: Float = 0.7
@@ -363,8 +364,8 @@ extension VecturaCLI {
     @Option(name: [.long, .customShort("d")], help: "Database name")
     var dbName: String = "vectura-cli-db"
 
-    @Option(name: [.long, .customShort("v")], help: "Vector dimension")
-    var dimension: Int = 384
+    @Option(name: [.long, .customShort("v")], help: "Vector dimension (auto-detected if not specified)")
+    var dimension: Int?
 
     @Option(name: [.long, .customShort("m")], help: "Model ID for embeddings")
     var modelId: String = "sentence-transformers/all-MiniLM-L6-v2"
@@ -396,8 +397,8 @@ extension VecturaCLI {
     @Option(name: [.long, .customShort("d")], help: "Database name")
     var dbName: String = "vectura-cli-db"
 
-    @Option(name: [.long, .customShort("v")], help: "Vector dimension")
-    var dimension: Int = 384
+    @Option(name: [.long, .customShort("v")], help: "Vector dimension (auto-detected if not specified)")
+    var dimension: Int?
 
     @Argument(help: "Document IDs to delete")
     var ids: [DocumentID]
@@ -423,8 +424,8 @@ extension VecturaCLI {
     @Option(name: [.long, .customShort("d")], help: "Database name")
     var dbName: String = "vectura-cli-db"
 
-    @Option(name: [.long, .customShort("v")], help: "Vector dimension")
-    var dimension: Int = 384
+    @Option(name: [.long, .customShort("v")], help: "Vector dimension (auto-detected if not specified)")
+    var dimension: Int?
 
     mutating func run() async throws {
       let db = try await VecturaCLI.setupDB(

Sources/VecturaCLI/VecturaCLI.swift

@@ -166,7 +166,36 @@ public struct VecturaConfig: Sendable {
     searchOptions: SearchOptions = .init(),
     memoryStrategy: MemoryStrategy = .automatic()
   ) throws {
+    let trimmedName = name.trimmingCharacters(in: .whitespacesAndNewlines)
+    guard !trimmedName.isEmpty else {
+      throw VecturaError.invalidInput("Database name cannot be empty or whitespace")
+    }
+    guard trimmedName != "." && trimmedName != ".." else {
+      throw VecturaError.invalidInput("Database name cannot be '.' or '..'")
+    }
+    guard !trimmedName.contains("/") && !trimmedName.contains("\\") else {
+      throw VecturaError.invalidInput("Database name cannot contain path separators")
+    }
+
     // Validate search options
+    guard searchOptions.defaultNumResults > 0 else {
+      throw VecturaError.invalidInput(
+        "defaultNumResults must be greater than 0, got \(searchOptions.defaultNumResults)"
+      )
+    }
+    if let threshold = searchOptions.minThreshold {
+      guard threshold >= 0.0 && threshold <= 1.0 else {
+        throw VecturaError.invalidInput(
+          "minThreshold must be between 0.0 and 1.0, got \(threshold)"
+        )
+      }
+    }
+    guard searchOptions.k1 > 0 else {
+      throw VecturaError.invalidInput("k1 must be greater than 0, got \(searchOptions.k1)")
+    }
+    guard searchOptions.b >= 0.0 && searchOptions.b <= 1.0 else {
+      throw VecturaError.invalidInput("b must be between 0.0 and 1.0, got \(searchOptions.b)")
+    }
     guard searchOptions.bm25NormalizationFactor > 0 else {
       throw VecturaError.invalidInput(
         "bm25NormalizationFactor must be positive, got \(searchOptions.bm25NormalizationFactor)"
@@ -190,7 +219,7 @@ public struct VecturaConfig: Sendable {
       break
     }
 
-    self.name = name
+    self.name = trimmedName
     self.directoryURL = directoryURL
     self.dimension = dimension
     self.searchOptions = searchOptions

Sources/VecturaKit/Core/VecturaConfig.swift

@@ -147,12 +147,82 @@ public actor VecturaKit {
       documentIds.append(docId)
     }
 
+    let existingDocumentsById: [UUID: VecturaDocument]
+    let idsToRestore = Set(documentIds)
+    if idsToRestore.isEmpty {
+      existingDocumentsById = [:]
+    } else if let indexedStorage = storageProvider as? IndexedVecturaStorage {
+      existingDocumentsById = try await indexedStorage.loadDocuments(ids: Array(idsToRestore))
+    } else if ids != nil {
+      let existingDocs = try await storageProvider.loadDocuments()
+      existingDocumentsById = existingDocs.reduce(into: [:]) { dict, doc in
+        if idsToRestore.contains(doc.id) {
+          dict[doc.id] = doc
+        }
+      }
+    } else {
+      existingDocumentsById = [:]
+    }
+
     // Save documents to storage (storage provider handles batch concurrency)
     try await storageProvider.saveDocuments(documentsToSave)
 
     // Notify search engine to index documents
-    for doc in documentsToSave {
-      try await searchEngine.indexDocument(doc)
+    var indexedDocumentIDs: [UUID] = []
+    indexedDocumentIDs.reserveCapacity(documentsToSave.count)
+
+    do {
+      for doc in documentsToSave {
+        try await searchEngine.indexDocument(doc)
+        indexedDocumentIDs.append(doc.id)
+      }
+    } catch {
+      Self.logger.error("Indexing failed after saving documents: \(error.localizedDescription)")
+
+      for id in indexedDocumentIDs {
+        do {
+          try await searchEngine.removeDocument(id: id)
+        } catch {
+          Self.logger.warning(
+            "Failed to rollback search index for \(id): \(error.localizedDescription)"
+          )
+        }
+      }
+
+      for doc in documentsToSave {
+        if let existingDoc = existingDocumentsById[doc.id] {
+          do {
+            try await storageProvider.updateDocument(existingDoc)
+          } catch {
+            Self.logger.warning(
+              "Failed to restore stored document \(doc.id): \(error.localizedDescription)"
+            )
+          }
+        } else {
+          do {
+            try await storageProvider.deleteDocument(withID: doc.id)
+          } catch {
+            Self.logger.warning(
+              "Failed to rollback stored document \(doc.id): \(error.localizedDescription)"
+            )
+          }
+        }
+      }
+
+      for id in indexedDocumentIDs {
+        guard let existingDoc = existingDocumentsById[id] else {
+          continue
+        }
+        do {
+          try await searchEngine.indexDocument(existingDoc)
+        } catch {
+          Self.logger.warning(
+            "Failed to restore search index for \(id): \(error.localizedDescription)"
+          )
+        }
+      }
+
+      throw error
     }
 
     return documentIds

Sources/VecturaKit/Core/VecturaKit.swift

@@ -179,6 +179,9 @@ public actor BM25Index {
   /// - Returns: Array of tuples containing lightweight documents and their BM25 scores
   public func search(query: String, topK: Int = 10) -> [(document: BM25Document, score: Float)] {
     let queryTerms = tokenize(query)
+    guard !queryTerms.isEmpty else {
+      return []
+    }
     var scores: [(BM25Document, Float)] = []
 
     for document in documents.values {
@@ -207,10 +210,12 @@ public actor BM25Index {
       scores.append((document, score))
     }
 
-    return scores
-      .sorted { $0.1 > $1.1 }
-      .prefix(topK)
-      .filter { $0.1 > 0 }
+    return Array(
+      scores
+        .sorted { $0.1 > $1.1 }
+        .filter { $0.1 > 0 }
+        .prefix(topK)
+    )
   }
 
   /// Add a new document to the index incrementally