Merge pull request #8 from moekidev/add-chunks-with-metadata

Add chunks with metadata

lib/baran/character_text_splitter.rb

@@ -14,4 +14,4 @@ def splitted(text)
       merged(splits, @separator)
     end
   end
-end
+end

lib/baran/text_splitter.rb

@@ -14,12 +14,14 @@ def splitted(text)
       raise NotImplementedError, "splitted method should be implemented in a subclass"
     end
 
-    def chunks(text)
+    def chunks(text, metadata: nil)
       cursor = 0
       chunks = []
 
       splitted(text).compact.each do |chunk|
-        chunks << { text: chunk, cursor: cursor }
+        chunk = { text: chunk, cursor: cursor }
+        chunk[:metadata] = metadata if metadata
+        chunks << chunk
         cursor += chunk.length
       end
 
@@ -56,4 +58,4 @@ def merged(splits, separator)
       results
     end
   end
-end
+end

test/test_character_text_splitter.rb

@@ -13,4 +13,4 @@ def test_chunks
 
     assert_equal(chunks.length, 3)
   end
-end
+end

test/test_recursive_character_text_splitter.rb

@@ -20,4 +20,4 @@ def test_empty_chunks
 
     assert_equal(chunks.length, 6)
   end
-end
+end

test/test_text_splitter.rb

@@ -55,6 +55,14 @@ def test_chunks
     assert_equal 'text', documents[0][:text]
   end
 
+  def test_chunks_with_metadata
+    text = 'text one'
+    metadata = { page: 1 }
+    documents = @test_splitter.chunks(text, metadata: metadata)
+
+    assert_equal({ page: 1 }, documents[0][:metadata])
+  end
+
   def test_joined
     items = ['one', 'two', 'three']
     separator = ' '