"Refactor SimilaritySearchKit to include autoLoad functionality and w…

…riteToFile method."
buhe · Feb 12, 2024 · 8e92a99 · 8e92a99
1 parent a373249
commit 8e92a99
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/Sources/LangChain/vectorstores/SimilaritySearchKit.swift b/Sources/LangChain/vectorstores/SimilaritySearchKit.swift
@@ -40,11 +40,18 @@ private struct LangChainEmbeddingBridge: EmbeddingsProtocol {
 public class SimilaritySearchKit: VectorStore {
     let vs: SimilarityIndex
 
-    public init(embeddings: Embeddings) async {
+    public init(embeddings: Embeddings, autoLoad: Bool = false) async {
         self.vs =  await SimilarityIndex(
             model: LangChainEmbeddingBridge(embeddings: embeddings),
             metric: CosineSimilarity()
         )
+        if #available(macOS 13.0, *) {
+            if autoLoad {
+                try? vs.loadIndex()
+            }
+        } else {
+            // Fallback on earlier versions
+        }
     }
 
     override func similaritySearch(query: String, k: Int) async -> [MatchedModel] {
@@ -54,5 +61,10 @@ public class SimilaritySearchKit: VectorStore {
     override func addText(text: String, metadata: [String: String]) async {
         await vs.addItem(id: UUID().uuidString, text: text, metadata: metadata)
     }
+
+    @available(macOS 13.0, *)
+    public func writeToFile() {
+        try? vs.saveIndex()
+    }
 }
 #endif