Add merged() to TextChunk

Author: intitni

Tool/Sources/LangChain/DocumentTransformer/TextSplitter.swift

@@ -50,10 +50,23 @@ public extension TextSplitter {
     }
 }
 
-public struct TextChunk {
+public struct TextChunk: Equatable {
     public var text: String
     public var startUTF16Offset: Int
     public var endUTF16Offset: Int
+
+    /// Merge the current chunk with another chunk if the 2 chunks are overlapping or adjacent.
+    public func merged(with chunk: TextChunk, force: Bool = false) -> TextChunk? {
+        let frontChunk = startUTF16Offset < chunk.startUTF16Offset ? self : chunk
+        let backChunk = startUTF16Offset < chunk.startUTF16Offset ? chunk : self
+        let overlap = frontChunk.endUTF16Offset - backChunk.startUTF16Offset
+        guard overlap >= 0 || force else { return nil }
+
+        let text = frontChunk.text + backChunk.text.dropFirst(max(0, overlap))
+        let start = frontChunk.startUTF16Offset
+        let end = backChunk.endUTF16Offset
+        return TextChunk(text: text, startUTF16Offset: start, endUTF16Offset: end)
+    }
 }
 
 public extension TextSplitter {

Tool/Tests/LangChainTests/TextSplitterTests/TextChunkTests.swift

@@ -0,0 +1,49 @@
+import Foundation
+import XCTest
+
+@testable import LangChain
+
+class TextChunkTests: XCTestCase {
+    func test_merging_overlapping_text_chunks() {
+        let chunk1 = TextChunk(text: "abc", startUTF16Offset: 0, endUTF16Offset: 3)
+        let chunk2 = TextChunk(text: "cdef", startUTF16Offset: 2, endUTF16Offset: 6)
+        let mergedChunk = chunk1.merged(with: chunk2)
+        XCTAssertEqual(mergedChunk?.text, "abcdef")
+        XCTAssertEqual(mergedChunk?.startUTF16Offset, 0)
+        XCTAssertEqual(mergedChunk?.endUTF16Offset, 6)
+    }
+    
+    func test_merging_adjacent_text_chunks() {
+        let chunk1 = TextChunk(text: "abc", startUTF16Offset: 0, endUTF16Offset: 3)
+        let chunk2 = TextChunk(text: "def", startUTF16Offset: 3, endUTF16Offset: 6)
+        let mergedChunk = chunk1.merged(with: chunk2)
+        XCTAssertEqual(mergedChunk?.text, "abcdef")
+        XCTAssertEqual(mergedChunk?.startUTF16Offset, 0)
+        XCTAssertEqual(mergedChunk?.endUTF16Offset, 6)
+    }
+    
+    func test_merging_overlapping_text_chunks_reversed_order() {
+        let chunk1 = TextChunk(text: "abc", startUTF16Offset: 0, endUTF16Offset: 3)
+        let chunk2 = TextChunk(text: "cdef", startUTF16Offset: 2, endUTF16Offset: 6)
+        let mergedChunk = chunk2.merged(with: chunk1)
+        XCTAssertEqual(mergedChunk?.text, "abcdef")
+        XCTAssertEqual(mergedChunk?.startUTF16Offset, 0)
+        XCTAssertEqual(mergedChunk?.endUTF16Offset, 6)
+    }
+    
+    func test_merging_adjacent_text_chunks_reversed_order() {
+        let chunk1 = TextChunk(text: "abc", startUTF16Offset: 0, endUTF16Offset: 3)
+        let chunk2 = TextChunk(text: "def", startUTF16Offset: 3, endUTF16Offset: 6)
+        let mergedChunk = chunk2.merged(with: chunk1)
+        XCTAssertEqual(mergedChunk?.text, "abcdef")
+        XCTAssertEqual(mergedChunk?.startUTF16Offset, 0)
+        XCTAssertEqual(mergedChunk?.endUTF16Offset, 6)
+    }
+    
+    func test_do_not_merge_non_overlapping_text_chunks() {
+        let chunk1 = TextChunk(text: "abc", startUTF16Offset: 0, endUTF16Offset: 3)
+        let chunk2 = TextChunk(text: "def", startUTF16Offset: 4, endUTF16Offset: 7)
+        let mergedChunk = chunk1.merged(with: chunk2)
+        XCTAssertNil(mergedChunk)
+    }
+}