Add TemporaryUSearch for vector searching

Author: intitni

Copilot for Xcode.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved

@@ -24,6 +24,7 @@ let package = Package(
         .package(url: "https://github.com/pointfreeco/swift-parsing", from: "0.12.1"),
         .package(url: "https://github.com/ChimeHQ/JSONRPC", from: "0.6.0"),
         .package(url: "https://github.com/scinfu/SwiftSoup.git", from: "2.6.0"),
+        .package(url: "https://github.com/unum-cloud/usearch", from: "0.19.0"),
     ],
     targets: [
         // MARK: - Helpers
@@ -60,6 +61,7 @@ let package = Package(
                 .product(name: "PythonKit", package: "PythonKit"),
                 .product(name: "Parsing", package: "swift-parsing"),
                 .product(name: "SwiftSoup", package: "SwiftSoup"),
+                .product(name: "USearch", package: "usearch"),
             ]
         ),
 

Tool/Package.swift

@@ -1,9 +1,10 @@
 import Foundation
+import JSONRPC
 
-public struct Document {
+public struct Document: Codable {
     public var pageContent: String
-    public var metadata: [String: Any]
-    public init(pageContent: String, metadata: [String: Any]) {
+    public var metadata: JSONValue
+    public init(pageContent: String, metadata: JSONValue) {
         self.pageContent = pageContent
         self.metadata = metadata
     }

Tool/Sources/LangChain/DocumentLoader/DocumentLoader.swift

@@ -26,9 +26,11 @@ public struct TextLoader: DocumentLoader {
         let modificationDate = try? url.resourceValues(forKeys: [.contentModificationDateKey])
             .contentModificationDate
         return [Document(pageContent: attributedString.string, metadata: [
-            "filename": url.lastPathComponent,
-            "extension": url.pathExtension,
-            "contentModificationDate": modificationDate ?? Date(),
+            "filename": .string(url.lastPathComponent),
+            "extension": .string(url.pathExtension),
+            "contentModificationDate": .number(
+                (modificationDate ?? Date()).timeIntervalSince1970
+            ),
         ])]
     }
 }

Tool/Sources/LangChain/DocumentLoader/TextLoader.swift

@@ -37,9 +37,9 @@ public struct WebLoader: DocumentLoader {
 
                     if let body = body {
                         let doc = Document(pageContent: body, metadata: [
-                            "title": title,
-                            "url": result.url,
-                            "date": Date(),
+                            "title": .string(title),
+                            "url": .string(result.url.absoluteString),
+                            "date": .number(Date().timeIntervalSince1970),
                         ])
                         documents.append(doc)
                     }

Tool/Sources/LangChain/DocumentLoader/WebLoader.swift

@@ -1,4 +1,5 @@
 import Foundation
+import JSONRPC
 
 public protocol DocumentTransformer {
     func transformDocuments(_ documents: [Document]) async throws -> [Document]
@@ -13,7 +14,7 @@ public extension TextSplitter {
     /// Create documents from a list of texts.
     func createDocuments(
         texts: [String],
-        metadata: [[String: Any]] = []
+        metadata: [JSONValue] = []
     ) async throws -> [Document] {
         var documents = [Document]()
         let paddingLength = texts.count - metadata.count
@@ -31,7 +32,7 @@ public extension TextSplitter {
     /// Split documents.
     func splitDocuments(_ documents: [Document]) async throws -> [Document] {
         var texts = [String]()
-        var metadata = [[String: Any]]()
+        var metadata = [JSONValue]()
         for document in documents {
             texts.append(document.pageContent)
             metadata.append(document.metadata)

Tool/Sources/LangChain/DocumentTransformer/DocumentTransformer.swift

@@ -2,7 +2,12 @@ import Foundation
 
 public protocol Embeddings {
     /// Embed search docs.
-    func embed(documents: [String]) async throws -> [[Float]]
+    func embed(documents: [Document]) async throws -> [EmbeddedDocument]
     /// Embed query text.
     func embed(query: String) async throws -> [Float]
 }
+
+public struct EmbeddedDocument: Codable {
+    var document: Document
+    var embeddings: [Float]
+}

Tool/Sources/LangChain/Embedding/Embedding.swift

@@ -21,37 +21,45 @@ public struct OpenAIEmbedding: Embeddings {
         self.safe = safe
     }
 
-    public func embed(documents: [String]) async throws -> [[Float]] {
+    public func embed(documents: [Document]) async throws -> [EmbeddedDocument] {
         if safe {
-            return try await getLenSafeEmbeddings(texts: documents).map(\.embeddings)
+            return try await getLenSafeEmbeddings(documents: documents)
         }
-        return try await getEmbeddings(texts: documents).map(\.embeddings)
+        return try await getEmbeddings(documents: documents)
     }
 
     public func embed(query: String) async throws -> [Float] {
         if safe {
-            return try await getLenSafeEmbeddings(texts: [query]).first?.embeddings ?? []
+            return try await getLenSafeEmbeddings(documents: [.init(
+                pageContent: query,
+                metadata: .null
+            )])
+            .first?
+            .embeddings ?? []
         }
-        return try await getEmbeddings(texts: [query]).first?.embeddings ?? []
+        return try await getEmbeddings(documents: [.init(pageContent: query, metadata: .null)])
+            .first?
+            .embeddings ?? []
     }
 }
 
 extension OpenAIEmbedding {
     func getEmbeddings(
-        texts: [String]
-    ) async throws -> [(originalText: String, embeddings: [Float])] {
+        documents: [Document]
+    ) async throws -> [EmbeddedDocument] {
         try await withThrowingTaskGroup(
-            of: (originalText: String, embeddings: [Float]).self
+            of: (document: Document, embeddings: [Float]).self
         ) { group in
-            for text in texts {
+            for document in documents {
                 group.addTask {
                     var retryCount = 6
                     var previousError: Error?
                     while retryCount > 0 {
                         do {
-                            let embeddings = try await service.embed(text: text).data
+                            let embeddings = try await service.embed(text: document.pageContent)
+                                .data
                                 .map(\.embeddings).first ?? []
-                            return (text, embeddings)
+                            return (document, embeddings)
                         } catch {
                             retryCount -= 1
                             previousError = error
@@ -60,27 +68,27 @@ extension OpenAIEmbedding {
                     throw previousError ?? CancellationError()
                 }
             }
-            var all = [(originalText: String, embeddings: [Float])]()
+            var all = [EmbeddedDocument]()
             for try await result in group {
-                all.append(result)
+                all.append(.init(document: result.document, embeddings: result.embeddings))
             }
             return all
         }
     }
 
     func getLenSafeEmbeddings(
-        texts: [String]
-    ) async throws -> [(originalText: String, embeddings: [Float])] {
+        documents: [Document]
+    ) async throws -> [EmbeddedDocument] {
         struct Text {
-            var rawText: String
+            var document: Document
             var chunkedTokens: [[Int]]
         }
 
-        var texts = texts.map { Text(rawText: $0, chunkedTokens: []) }
+        var texts = documents.map { Text(document: $0, chunkedTokens: []) }
         let encoding = TiktokenCl100kBaseTokenEncoder()
 
         for (index, text) in texts.enumerated() {
-            let token = encoding.encode(text: text.rawText)
+            let token = encoding.encode(text: text.document.pageContent)
             // just incase the calculation is incorrect
             let maxToken = max(10, service.configuration.maxToken - 10)
 
@@ -92,27 +100,28 @@ extension OpenAIEmbedding {
         }
 
         let batchedEmbeddings = try await withThrowingTaskGroup(
-            of: (String, [[Float]]).self
+            of: (Document, [[Float]]).self
         ) { group in
             for text in texts {
                 group.addTask {
                     var retryCount = 6
                     var previousError: Error?
-                    guard !text.chunkedTokens.isEmpty else { return (text.rawText, []) }
+                    guard !text.chunkedTokens.isEmpty
+                    else { return (text.document, []) }
                     while retryCount > 0 {
                         do {
                             if text.chunkedTokens.count <= 1 {
                                 // if possible, we should just let OpenAI do the tokenization.
                                 return (
-                                    text.rawText,
-                                    try await service.embed(text: text.rawText)
+                                    text.document,
+                                    try await service.embed(text: text.document.pageContent)
                                         .data
                                         .map(\.embeddings)
                                 )
                             }
                             if shouldAverageLongEmbeddings {
                                 return (
-                                    text.rawText,
+                                    text.document,
                                     try await service.embed(tokens: text.chunkedTokens)
                                         .data
                                         .map(\.embeddings)
@@ -121,7 +130,7 @@ extension OpenAIEmbedding {
                             // if `shouldAverageLongEmbeddings` is false,
                             // we only embed the first chunk to save some money.
                             return (
-                                text.rawText,
+                                text.document,
                                 try await service.embed(tokens: [text.chunkedTokens.first ?? []])
                                     .data
                                     .map(\.embeddings)
@@ -134,21 +143,21 @@ extension OpenAIEmbedding {
                     throw previousError ?? CancellationError()
                 }
             }
-            var result = [(originalText: String, embeddings: [[Float]])]()
+            var result = [(document: Document, embeddings: [[Float]])]()
             for try await response in group {
                 try Task.checkCancellation()
                 result.append((response.0, response.1))
             }
             return result
         }
 
-        var results = [(originalText: String, embeddings: [Float])]()
+        var results = [EmbeddedDocument]()
 
-        for (text, embeddings) in batchedEmbeddings {
+        for (document, embeddings) in batchedEmbeddings {
             if embeddings.count == 1, let first = embeddings.first {
-                results.append((text, first))
+                results.append(.init(document: document, embeddings: first))
             } else if embeddings.isEmpty {
-                results.append((text, []))
+                results.append(.init(document: document, embeddings: []))
             } else if shouldAverageLongEmbeddings {
                 // untested
                 do {
@@ -162,14 +171,14 @@ extension OpenAIEmbedding {
                         let normalized = average / numpy.linalg.norm(average)
                         return [Float](normalized.tolist())
                     }) else { throw CancellationError() }
-                    results.append((text, averagedEmbeddings))
+                    results.append(.init(document: document, embeddings: averagedEmbeddings))
                 } catch {
                     if let first = embeddings.first {
-                        results.append((text, first))
+                        results.append(.init(document: document, embeddings: first))
                     }
                 }
             } else if let first = embeddings.first {
-                results.append((text, first))
+                results.append(.init(document: document, embeddings: first))
             }
         }