Merge pull request #79 from buhe/LocalFileStore

LocalFileStore
buhe · Nov 19, 2023 · 3211678 · 3211678
2 parents 406fde9 + 8a0d21a
commit 3211678
Show file tree

Hide file tree

Showing 7 changed files with 122 additions and 14 deletions.
diff --git a/README.md b/README.md
@@ -367,6 +367,7 @@ Task(priority: .background)  {
 - Store
     - [x] BaseStore
     - [x] InMemoryStore
+    - [x] FileStore
 - Embedding
     - [x] OpenAI
 - Chain

diff --git a/Sources/LangChain/retriever/MultiVectorRetriever.swift b/Sources/LangChain/retriever/MultiVectorRetriever.swift
@@ -22,7 +22,7 @@ public class MultiVectorRetriever: BaseRetriever {
         for d in sub_docs {
             ids.append(d.metadata[self.id_key]!)
         }
-        let docs = self.docstore.mget(keys: ids)
+        let docs = await self.docstore.mget(keys: ids)
         return docs.map{Document(page_content: $0, metadata: [:])}
     }
 

diff --git a/Sources/LangChain/retriever/ParentDocumentRetriever.swift b/Sources/LangChain/retriever/ParentDocumentRetriever.swift
@@ -35,7 +35,7 @@ public class ParentDocumentRetriever: MultiVectorRetriever {
         }
         print("🚀 Begin add sub document \(docs.count), main document \(full_docs.count)")
         await self.vectorstore.add_documents(documents: docs)
-        self.docstore.mset(kvpairs: full_docs)
+        await self.docstore.mset(kvpairs: full_docs)
         print("🚀 End add sub document \(docs.count), main document \(full_docs.count)")
     }
 //    def add_documents(

diff --git a/Sources/LangChain/schema/BaseStore.swift b/Sources/LangChain/schema/BaseStore.swift
@@ -7,19 +7,19 @@
 
 import Foundation
 public class BaseStore {
-    func mget(keys: [String]) -> [String] {
+    func mget(keys: [String]) async -> [String] {
         []
     }
 
-    func mset(kvpairs: [(String, String)]) {
+    func mset(kvpairs: [(String, String)]) async {
 
     }
 
-    func mdelete(keys: [String]) {
+    func mdelete(keys: [String]) async {
 
     }
 
-    func keys(prefix: String? = nil) -> [String] {
+    func keys(prefix: String? = nil) async -> [String] {
         []
     }
 }
diff --git a/Sources/LangChain/schema/InMemoryStore.swift b/Sources/LangChain/schema/InMemoryStore.swift
@@ -11,7 +11,7 @@ public class InMemoryStore: BaseStore {
     public override init() {
         super.init()
     }
-    override func mget(keys: [String]) -> [String] {
+    override func mget(keys: [String]) async -> [String] {
         var values: [String] = []
         for k in keys {
             let v = self.store[k]
@@ -22,19 +22,19 @@ public class InMemoryStore: BaseStore {
         return values
     }
 
-    override func mset(kvpairs: [(String, String)]) {
+    override func mset(kvpairs: [(String, String)]) async {
         for kv in kvpairs {
             self.store[kv.0] = kv.1
         }
     }
 
-    override func mdelete(keys: [String]) {
+    override func mdelete(keys: [String]) async {
         for k in keys {
             self.store.removeValue(forKey: k)
         }
     }
 
-    override func keys(prefix: String? = nil) -> [String] {
+    override func keys(prefix: String? = nil) async -> [String] {
         if prefix == nil {
             return Array(self.store.keys)
         } else {

diff --git a/Sources/LangChain/schema/LocalFileStore.swift b/Sources/LangChain/schema/LocalFileStore.swift
@@ -0,0 +1,94 @@
+//
+//  File.swift
+//  
+//
+//  Created by 顾艳华 on 2023/11/19.
+//
+
+import Foundation
+import SwiftFileStore
+struct StoreEntry: Codable, JSONDataRepresentable {
+    let key: String
+    let value: String
+}
+public class LocalFileStore: BaseStore {
+    let objectStore: FileObjectStore?
+    static let STORE_NS = "store"
+    public override init() {
+        do {
+            self.objectStore = try FileObjectStore.create()
+        } catch {
+            self.objectStore = nil
+        }
+    }
+
+    override func mget(keys: [String]) async -> [String] {
+        print("🍰 Get \(keys) from file")
+        var values: [String] = []
+        do {
+            for key in keys {
+                if let data = key.data(using: .utf8) {
+                    let base64 = data.base64EncodedString()
+
+                    let cache = try await objectStore!.read(key: base64.sha256(), namespace: LocalFileStore.STORE_NS, objectType: StoreEntry.self)
+                    if let c = cache {
+                        values.append(c.value)
+                    }
+                }
+            }
+        } catch {
+            print("FileStore get failed")
+        }
+        return values
+    }
+
+    override func mset(kvpairs: [(String, String)]) async {
+        print("🍰 Update \(kvpairs) at file")
+        do {
+            for kv in kvpairs {
+                if let data = kv.0.data(using: .utf8) {
+                    let base64 = data.base64EncodedString()
+                    let cache = StoreEntry(key: kv.0, value: kv.1)
+                    try await objectStore!.write(key: base64.sha256(), namespace: LocalFileStore.STORE_NS, object: cache)
+                }
+            }
+        } catch {
+            print("FileStore set failed")
+        }
+    }
+
+    override func mdelete(keys: [String]) async {
+        print("🍰 Delete \(keys) at file")
+        do {
+            for key in keys {
+                if let data = key.data(using: .utf8) {
+                    let base64 = data.base64EncodedString()
+                    try await objectStore!.remove(key: base64.sha256(), namespace: LocalFileStore.STORE_NS)
+                }
+            }
+        } catch {
+            print("FileStore set failed")
+        }
+    }
+
+    override func keys(prefix: String? = nil) async -> [String] {
+        if prefix == nil {
+            print("🍰 Get all keys from file")
+            return Array(await self.allKeys())
+        } else {
+            print("🍰 Get keys \(prefix!) from file")
+            var matched: [String] = []
+            for k in await self.allKeys() {
+                if k.hasPrefix(prefix!) {
+                    matched.append(k)
+                }
+            }
+            return matched
+        }
+
+    }
+
+    func allKeys() async -> [String] {
+        [] // TODO
+    }
+}
diff --git a/Tests/LangChainTests/langchain_swiftTests.swift b/Tests/LangChainTests/langchain_swiftTests.swift
@@ -478,10 +478,12 @@ Action Input: the input to the action
 
     func testInMemoryStore() async throws {
         let store = InMemoryStore()
-        store.mset(kvpairs: [("1", "a"), ("2", "b")])
-        XCTAssertEqual(store.mget(keys: ["1", "2"]), ["a", "b"])
-        store.mdelete(keys: ["1"])
-        XCTAssertEqual(store.keys(), ["2"])
+        await store.mset(kvpairs: [("1", "a"), ("2", "b")])
+        var values = await store.mget(keys: ["1", "2"])
+        XCTAssertEqual(values, ["a", "b"])
+        await store.mdelete(keys: ["1"])
+        let keys = await store.keys()
+        XCTAssertEqual(keys, ["2"])
     }
 
     func testRecursiveCharacterTextSplitter() async throws {
@@ -500,6 +502,17 @@ Action Input: the input to the action
             XCTAssertTrue(doc.count <= 400)
         }
     }
+
+    func testFileStore() async throws {
+        let store = LocalFileStore()
+        await store.mset(kvpairs: [("1", "a"), ("2", "b")])
+        var values = await store.mget(keys: ["1", "2"])
+        XCTAssertEqual(values, ["a", "b"])
+        await store.mdelete(keys: ["1"])
+//        let keys = await store.keys()
+//        XCTAssertEqual(keys, ["2"])
+    }
+
 //
 //    func testYoutubeHackClientList() async throws {
 //        let eventLoopGroup = ThreadManager.thread