Update type of Document's metadata

Author: intitni

Tool/Sources/LangChain/DocumentLoader/DocumentLoader.swift

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

Tool/Sources/LangChain/DocumentLoader/WebLoader.swift

@@ -39,16 +39,15 @@ public struct WebLoader: DocumentLoader {
                     let parsed = try SwiftSoup.parse(result.html, result.url.path)
 
                     let title = (try? parsed.title()) ?? "Untitled"
-                    let body = try DefaultLoadContentStrategy().load(parsed)
-
-                    if let body = body {
-                        let doc = Document(pageContent: body, metadata: [
+                    let parsedDocuments = try DefaultLoadContentStrategy().load(
+                        parsed,
+                        metadata: [
                             MetadataKeys.title: .string(title),
                             MetadataKeys.url: .string(result.url.absoluteString),
                             MetadataKeys.date: .number(Date().timeIntervalSince1970),
-                        ])
-                        documents.append(doc)
-                    }
+                        ]
+                    )
+                    documents.append(contentsOf: parsedDocuments)
                 } catch let Exception.Error(_, message) {
                     Logger.langchain.error(message)
                 } catch {
@@ -61,7 +60,7 @@ public struct WebLoader: DocumentLoader {
 }
 
 protocol LoadWebPageMainContentStrategy {
-    func load(_ document: SwiftSoup.Document) throws -> String?
+    func load(_ document: SwiftSoup.Document, metadata: Document.Metadata) throws -> [Document]
 }
 
 extension LoadWebPageMainContentStrategy {
@@ -77,11 +76,19 @@ extension LoadWebPageMainContentStrategy {
 
 extension WebLoader {
     struct DefaultLoadContentStrategy: LoadWebPageMainContentStrategy {
-        func load(_ document: SwiftSoup.Document) throws -> String? {
-            if let article = text(inFirstTag: "article", from: document) { return article }
-            if let main = text(inFirstTag: "main", from: document) { return main }
-            let body = try document.body()?.text()
-            return body
+        func load(
+            _ document: SwiftSoup.Document,
+            metadata: Document.Metadata
+        ) throws -> [Document] {
+            if let mainContent = try? {
+                if let article = text(inFirstTag: "article", from: document) { return article }
+                if let main = text(inFirstTag: "main", from: document) { return main }
+                let body = try document.body()?.text()
+                return body
+            }() {
+                return [.init(pageContent: mainContent, metadata: metadata)]
+            }
+            return []
         }
     }
 }
@@ -100,7 +107,9 @@ final class WebScrapper: NSObject, WKNavigationDelegate {
         configuration.defaultWebpagePreferences.preferredContentMode = .desktop
         configuration.defaultWebpagePreferences.allowsContentJavaScript = true
         configuration.mediaTypesRequiringUserActionForPlayback = []
-        configuration.applicationNameForUserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.1 Safari/605.1.15"
+        configuration
+            .applicationNameForUserAgent =
+            "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.1 Safari/605.1.15"
         // The web page need the web view to have a size to load correctly.
         let webView = WKWebView(
             frame: .init(x: 0, y: 0, width: 500, height: 500),
@@ -125,7 +134,7 @@ final class WebScrapper: NSObject, WKNavigationDelegate {
             if !html.isEmpty { return html }
             retryCount += 1
         }
-        
+
         throw CancellationError()
     }
 
@@ -150,3 +159,4 @@ final class WebScrapper: NSObject, WKNavigationDelegate {
 private let getHTMLText = """
 document.documentElement.outerHTML;
 """
+

Tool/Sources/LangChain/DocumentTransformer/TextSplitter.swift

@@ -18,7 +18,7 @@ public extension TextSplitter {
     /// Create documents from a list of texts.
     func createDocuments(
         texts: [String],
-        metadata: [JSONValue] = []
+        metadata: [Document.Metadata] = []
     ) async throws -> [Document] {
         var documents = [Document]()
         let paddingLength = texts.count - metadata.count
@@ -36,7 +36,7 @@ public extension TextSplitter {
     /// Split documents.
     func splitDocuments(_ documents: [Document]) async throws -> [Document] {
         var texts = [String]()
-        var metadata = [JSONValue]()
+        var metadata = [Document.Metadata]()
         for document in documents {
             texts.append(document.pageContent)
             metadata.append(document.metadata)

Tool/Sources/LangChain/Embedding/OpenAIEmbedding.swift

@@ -30,12 +30,12 @@ public struct OpenAIEmbedding: Embeddings {
         if safe {
             return try await getLenSafeEmbeddings(documents: [.init(
                 pageContent: query,
-                metadata: .null
+                metadata: [:]
             )])
             .first?
             .embeddings ?? []
         }
-        return try await getEmbeddings(documents: [.init(pageContent: query, metadata: .null)])
+        return try await getEmbeddings(documents: [.init(pageContent: query, metadata: [:])])
             .first?
             .embeddings ?? []
     }
@@ -119,7 +119,7 @@ extension OpenAIEmbedding {
                                         .map(\.embedding)
                                 )
                             }
-                            
+
                             if shouldAverageLongEmbeddings {
                                 return (
                                     text.document,