Support joining split documents

Author: intitni

Tool/Sources/LangChain/DocumentTransformer/TextSplitter.swift

@@ -26,6 +26,9 @@ public extension TextSplitter {
         for (text, metadata) in zip(texts, metadata) {
             let chunks = try await split(text: text)
             for chunk in chunks {
+                var metadata = metadata
+                metadata["startUTF16Offset"] = .number(Double(chunk.startUTF16Offset))
+                metadata["endUTF16Offset"] = .number(Double(chunk.endUTF16Offset))
                 let document = Document(pageContent: chunk.text, metadata: metadata)
                 documents.append(document)
             }
@@ -48,6 +51,32 @@ public extension TextSplitter {
     func transformDocuments(_ documents: [Document]) async throws -> [Document] {
         return try await splitDocuments(documents)
     }
+
+    func joinDocuments(_ documents: [Document]) -> Document {
+        let textChunks: [TextChunk] = documents.compactMap { document in
+            func extract(_ key: String) -> Int? {
+                if case let .number(d) = document.metadata[key] {
+                    return Int(d)
+                }
+                return nil
+            }
+            guard let start = extract("startUTF16Offset"),
+                  let end = extract("endUTF16Offset")
+            else { return nil }
+            return TextChunk(
+                text: document.pageContent,
+                startUTF16Offset: start,
+                endUTF16Offset: end
+            )
+        }.sorted(by: { $0.startUTF16Offset < $1.startUTF16Offset })
+        let mergedChunks = mergeSplits(textChunks)
+        let pageContent = mergedChunks.map(\.text).joined()
+        var metadata = documents.first?.metadata ?? [String: JSONValue]()
+        metadata["startUTF16Offset"] = nil
+        metadata["endUTF16Offset"] = nil
+
+        return Document(pageContent: pageContent, metadata: metadata)
+    }
 }
 
 public struct TextChunk: Equatable {
@@ -83,14 +112,14 @@ public extension TextSplitter {
             let text = (a + b).map(\.text).joined()
             var l = Int.max
             var u = 0
-            
+
             for chunk in a + b {
                 l = min(l, chunk.startUTF16Offset)
                 u = max(u, chunk.endUTF16Offset)
             }
-            
+
             guard l < u else { return nil }
-            
+
             return .init(text: text, startUTF16Offset: l, endUTF16Offset: u)
         }