Add native implementation of RecursiveCharacterTextSplitter

Author: intitni

Tool/Sources/LangChain/DocumentTransformer/DocumentTransformer.swift

@@ -6,6 +6,10 @@ public protocol DocumentTransformer {
 }
 
 public protocol TextSplitter: DocumentTransformer {
+    var chunkSize: Int { get }
+    var chunkOverlap: Int { get }
+    var lengthFunction: (String) -> Int { get }
+
     /// Split text into multiple components.
     func split(text: String) async throws -> [String]
 }
@@ -46,5 +50,87 @@ public extension TextSplitter {
     }
 }
 
-extension TextSplitter {}
+public extension TextSplitter {
+    /// Merge small splits to just fit in the chunk size.
+    func mergeSplits(_ splits: [String]) -> [String] {
+        let chunkOverlap = chunkOverlap < chunkSize ? chunkOverlap : 0
+
+        var chunks = [String]()
+        var currentChunk = [String]()
+        var overlappingChunks = [String]()
+        var currentChunkSize = 0
+        
+        func join(_ a: [String], _ b: [String]) -> String {
+            return (a + b).joined().trimmingCharacters(in: .whitespaces)
+        }
+
+        for text in splits {
+            let textLength = lengthFunction(text)
+            if currentChunkSize + textLength > chunkSize {
+                let currentChunkText = join(overlappingChunks, currentChunk)
+                chunks.append(currentChunkText)
+
+                overlappingChunks = []
+                var overlappingSize = 0
+                // use small chunks as overlap if possible
+                for chunk in currentChunk.reversed() {
+                    let length = lengthFunction(chunk)
+                    if overlappingSize + length > chunkOverlap { break }
+                    if overlappingSize + length + textLength > chunkSize { break }
+                    overlappingSize += length
+                    overlappingChunks.insert(chunk, at: 0)
+                }
+//                // fallback to use suffix if no small chunk found
+//                if overlappingChunks.isEmpty {
+//                    let suffix = String(
+//                        currentChunkText.suffix(min(chunkOverlap, chunkSize - textLength))
+//                    )
+//                    overlappingChunks.append(suffix)
+//                    overlappingSize = lengthFunction(suffix)
+//                }
+
+                currentChunkSize = overlappingSize + textLength
+                currentChunk = [text]
+            } else {
+                currentChunkSize += textLength
+                currentChunk.append(text)
+            }
+        }
+
+        if !currentChunk.isEmpty {
+            chunks.append(join(overlappingChunks, currentChunk))
+        }
+
+        return chunks
+    }
+
+    /// Split the text by separator.
+    func split(text: String, separator: String) -> [String] {
+        guard !separator.isEmpty else {
+            return [text]
+        }
+
+        let pattern = "(\(separator))"
+        if let regex = try? NSRegularExpression(pattern: pattern) {
+            let matches = regex.matches(in: text, range: NSRange(text.startIndex..., in: text))
+            var all = [String]()
+            var start = text.startIndex
+            for match in matches {
+                guard let range = Range(match.range, in: text) else { break }
+                guard range.lowerBound > start else { break }
+                let result = text[start..<range.lowerBound]
+                start = range.lowerBound
+                if !result.isEmpty {
+                    all.append(String(result))
+                }
+            }
+            if start < text.endIndex {
+                all.append(String(text[start...]))
+            }
+            return all
+        } else {
+            return [text]
+        }
+    }
+}
 

Tool/Sources/LangChain/DocumentTransformer/RecursiveCharacterTextSplitter.swift

@@ -1,39 +1,99 @@
 import Foundation
-import PythonHelper
-import PythonKit
 
-public struct RecursiveCharacterTextSplitter: TextSplitter {
+class RecursiveCharacterTextSplitter: TextSplitter {
+    /**
+     Implementation of splitting text that looks at characters.
+     Recursively tries to split by different characters to find one that works.
+     */
     public var separators: [String]
     public var chunkSize: Int
     public var chunkOverlap: Int
-    
-    public init(
+    public var lengthFunction: (String) -> Int
+
+    init(
         separators: [String] = ["\n\n", "\n", " ", ""],
         chunkSize: Int = 4000,
-        chunkOverlap: Int = 200
+        chunkOverlap: Int = 200,
+        lengthFunction: @escaping (String) -> Int = { $0.count }
     ) {
+        assert(chunkOverlap <= chunkSize)
+        self.chunkSize = chunkSize
+        self.chunkOverlap = chunkOverlap
+        self.lengthFunction = lengthFunction
         self.separators = separators
+    }
+
+    init(
+        separatorSet: TextSplitterSeparatorSet,
+        chunkSize: Int = 4000,
+        chunkOverlap: Int = 200,
+        lengthFunction: @escaping (String) -> Int = { $0.count }
+    ) {
+        assert(chunkOverlap <= chunkSize)
         self.chunkSize = chunkSize
         self.chunkOverlap = chunkOverlap
+        self.lengthFunction = lengthFunction
+        separators = separatorSet.separators
     }
 
     public func split(text: String) async throws -> [String] {
-        try await runPython {
-            let text_splitter = try Python.attemptImportOnPythonThread("langchain.text_splitter")
-            let PythonRecursiveCharacterTextSplitter = text_splitter.RecursiveCharacterTextSplitter
-            let splitter = PythonRecursiveCharacterTextSplitter(
-                separators: separators,
-                chunk_size: chunkSize,
-                chunk_overlap: chunkOverlap
-//                length_function: PythonFunction({ object in
-//                    if let string = String(object) { return string.count }
-//                    return 0
-//                })
-            )
-            let result = splitter.split_text(text)
-            guard let array = [String](result) else { return [] }
-            return array
+        return split(text: text, separators: separators)
+    }
+
+    private func split(text: String, separators: [String]) -> [String] {
+        var finalChunks = [String]()
+
+        // Get appropriate separator to use
+        let firstSeparatorIndex = separators.firstIndex {
+            let pattern = "(\($0))"
+            guard let regex = try? NSRegularExpression(pattern: pattern) else { return false }
+            return regex.firstMatch(
+                in: text,
+                options: [],
+                range: NSRange(text.startIndex..., in: text)
+            ) != nil
+        }
+        var separator: String
+        var nextSeparators: [String]
+        
+        if let index = firstSeparatorIndex {
+            separator = separators[index]
+            if index < separators.endIndex - 1 {
+                nextSeparators = Array(separators[(index + 1)...])
+            } else {
+                nextSeparators = []
+            }
+        } else {
+            separator = ""
+            nextSeparators = []
+        }
+
+        let splits = split(text: text, separator: separator)
+
+        // Now go merging things, recursively splitting longer texts.
+        var goodSplits = [String]()
+        for s in splits {
+            if lengthFunction(s) < chunkSize {
+                goodSplits.append(s)
+            } else {
+                if !goodSplits.isEmpty {
+                    let mergedText = mergeSplits(goodSplits)
+                    finalChunks.append(contentsOf: mergedText)
+                    goodSplits.removeAll()
+                }
+                if nextSeparators.isEmpty {
+                    finalChunks.append(s)
+                } else {
+                    let other_info = split(text: s, separators: nextSeparators)
+                    finalChunks.append(contentsOf: other_info)
+                }
+            }
+        }
+        if !goodSplits.isEmpty {
+            let merged_text = mergeSplits(goodSplits)
+            finalChunks.append(contentsOf: merged_text)
         }
+        return finalChunks
     }
 }