Merge pull request #50 from buhe/coredata_cache

"Add FileCache class to the cache module."
buhe · Oct 31, 2023 · 0c0b440 · 0c0b440
2 parents 7b8e1d9 + 8af0763
commit 0c0b440
Show file tree

Hide file tree

Showing 6 changed files with 106 additions and 8 deletions.
diff --git a/Package.resolved b/Package.resolved
@@ -117,6 +117,15 @@
         "version" : "1.0.4"
       }
     },
+    {
+      "identity" : "swift-filestore",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/juyan/swift-filestore",
+      "state" : {
+        "revision" : "878a896cc2fef836bd70a46bbb873b495155e2d0",
+        "version" : "0.2.0"
+      }
+    },
     {
       "identity" : "swift-log",
       "kind" : "remoteSourceControl",

diff --git a/Package.swift b/Package.swift
@@ -20,7 +20,8 @@ let package = Package(
         .package(url: "https://github.com/supabase-community/supabase-swift", .upToNextMajor(from: "0.2.1")),
         .package(url: "https://github.com/SwiftyJSON/SwiftyJSON", .upToNextMajor(from: "5.0.1")),
         .package(url: "https://github.com/drmohundro/SWXMLHash", .upToNextMajor(from: "7.0.2")),
-        .package(url: "https://github.com/scinfu/SwiftSoup", .upToNextMajor(from: "2.6.1"))
+        .package(url: "https://github.com/scinfu/SwiftSoup", .upToNextMajor(from: "2.6.1")),
+        .package(url: "https://github.com/juyan/swift-filestore", .upToNextMajor(from: "0.2.0")),
     ],
     targets: [
         // Targets are the basic building blocks of a package, defining a module or a test suite.
@@ -33,6 +34,7 @@ let package = Package(
                 .product(name: "SwiftyJSON", package: "SwiftyJSON"),
                 .product(name: "SWXMLHash", package: "SWXMLHash"),
                 .product(name: "SwiftSoup", package: "SwiftSoup"),
+                .product(name: "SwiftFileStore", package: "swift-filestore"),
             ]
 
         ),

diff --git a/README.md b/README.md
@@ -374,7 +374,9 @@ Task {
     - [x] MultiPromptRouter
 - Callback
     - [x] StdOutCallbackHandler 
-
+- LLM Cache
+    - [x] InMemery
+    - [x] File
 ## 👍 Got Ideas?
 Open an issue, and let's discuss!
 

diff --git a/Sources/LangChain/LangChain.swift b/Sources/LangChain/LangChain.swift
@@ -6,8 +6,43 @@
 //
 
 import Foundation
+import CommonCrypto
 
 enum LangChainError: Error {
     case LoaderError(String)
 
 }
+
+extension Data{
+    public func sha256() -> String{
+        return hexStringFromData(input: digest(input: self as NSData))
+    }
+
+    private func digest(input : NSData) -> NSData {
+        let digestLength = Int(CC_SHA256_DIGEST_LENGTH)
+        var hash = [UInt8](repeating: 0, count: digestLength)
+        CC_SHA256(input.bytes, UInt32(input.length), &hash)
+        return NSData(bytes: hash, length: digestLength)
+    }
+
+    private  func hexStringFromData(input: NSData) -> String {
+        var bytes = [UInt8](repeating: 0, count: input.length)
+        input.getBytes(&bytes, length: input.length)
+
+        var hexString = ""
+        for byte in bytes {
+            hexString += String(format:"%02x", UInt8(byte))
+        }
+
+        return hexString
+    }
+}
+
+public extension String {
+    func sha256() -> String{
+        if let stringData = self.data(using: String.Encoding.utf8) {
+            return stringData.sha256()
+        }
+        return ""
+    }
+}
diff --git a/Sources/LangChain/cache/Cache.swift b/Sources/LangChain/cache/Cache.swift
@@ -6,12 +6,14 @@
 //
 
 import Foundation
+import SwiftFileStore
+
 public class BaseCache {
     public init() {}
-    public func lookup(prompt: String) -> LLMResult? {
+    public func lookup(prompt: String) async -> LLMResult? {
         nil
     }
-    public func update(prompt: String, return_val: LLMResult) {
+    public func update(prompt: String, return_val: LLMResult) async {
 
     }
     //For test?
@@ -23,15 +25,63 @@ public class BaseCache {
 public class InMemoryCache: BaseCache {
 
     var memery: [String: LLMResult] = [:]
-    public override func lookup(prompt: String) -> LLMResult? {
+    public override func lookup(prompt: String) async -> LLMResult? {
         print("🍰 Get \(prompt) from cache")
         return memery[prompt]
     }
-    public override func update(prompt: String, return_val: LLMResult) {
+    public override func update(prompt: String, return_val: LLMResult) async {
         print("🍰 Update \(prompt)")
         memery[prompt] = return_val
     }
     public override func clear() {
         memery = [:]
     }
 }
+struct LLMCache: Codable, JSONDataRepresentable {
+    let key: String
+    let value: String
+}
+public class FileCache: BaseCache {
+    let objectStore: FileObjectStore?
+
+    public override init() {
+        do {
+            self.objectStore = try FileObjectStore.create()
+        } catch {
+            self.objectStore = nil
+        }
+    }
+    public override func lookup(prompt: String) async -> LLMResult? {
+        print("🍰 Get \(prompt) from file")
+        do {
+            if let data = prompt.data(using: .utf8) {
+                let base64 = data.base64EncodedString()
+
+                let cache = try await objectStore!.read(key: base64.sha256(), namespace: "llm_cache", objectType: LLMCache.self)
+                if let c = cache {
+                    return LLMResult(llm_output: c.value)
+                }
+            }
+            return nil
+        } catch {
+            return nil
+        }
+
+
+    }
+    public override func update(prompt: String, return_val: LLMResult) async {
+        print("🍰 Update \(prompt) at file")
+        do {
+            if let data = prompt.data(using: .utf8) {
+                let base64 = data.base64EncodedString()
+                let cache = LLMCache(key: prompt, value: return_val.llm_output!)
+                try await objectStore!.write(key: base64.sha256(), namespace: "llm_cache", object: cache)
+            }
+        } catch {
+            print("FileCache set failed")
+        }
+    }
+    public override func clear() {
+
+    }
+}
diff --git a/Sources/LangChain/llms/LLM.swift b/Sources/LangChain/llms/LLM.swift
@@ -29,14 +29,14 @@ public class LLM {
         callStart(prompt: text, reqId: reqId)
         do {
             if let cache = self.cache {
-                if let llmResult = cache.lookup(prompt: text) {
+                if let llmResult = await cache.lookup(prompt: text) {
                     callEnd(output: llmResult.llm_output!, reqId: reqId, cost: 0)
                     return llmResult
                 }
             }
             let llmResult = try await _send(text: text, stops: stops)
             if let cache = self.cache {
-                cache.update(prompt: text, return_val: llmResult)
+                await cache.update(prompt: text, return_val: llmResult)
             }
             cost = Date.now.timeIntervalSince1970 - now
             if !llmResult.stream {