From c3ac7a721196032f56b35ce717f0e9325d092f44 Mon Sep 17 00:00:00 2001 From: anaslimem Date: Thu, 16 Apr 2026 22:36:30 +0100 Subject: [PATCH 1/2] fix: calculate float context size per chunk in overlap refinery --- src/chonkie/refinery/overlap.py | 22 ++++++++-------------- tests/refinery/test_overlap_refinery.py | 15 ++++----------- 2 files changed, 12 insertions(+), 25 deletions(-) diff --git a/src/chonkie/refinery/overlap.py b/src/chonkie/refinery/overlap.py index a6e0a990e..bc8a54109 100644 --- a/src/chonkie/refinery/overlap.py +++ b/src/chonkie/refinery/overlap.py @@ -11,10 +11,6 @@ logger = get_logger(__name__) -# TODO: Fix the way that float context size is handled. -# Currently, it just estimates the context size to token count -# but it should ideally handle it on a chunk by chunk basis. - # TODO: Add support for `justified` method which is the best of # both prefix and suffix overlap. @@ -293,16 +289,15 @@ def _refine_prefix(self, chunks: list[Chunk], effective_context_size: int) -> li The refined chunks. """ - # Iterate over the chunks till the second to last chunk for i, chunk in enumerate(chunks[1:]): - # Get the previous chunk, since i starts from 0 prev_chunk = chunks[i] - # Calculate effective context size per chunk if context_size is a float + # Calculate context size based on the chunk RECEIVING context (its own token count) + # This ensures each chunk gets overlap proportional to its own size if isinstance(self.context_size, float): - effective_context_size = int(self.context_size * prev_chunk.token_count) + effective_context_size = int(self.context_size * chunk.token_count) - # Calculate the overlap context + # Get context from the previous chunk context = self._get_prefix_overlap_context(prev_chunk, effective_context_size) # Set it as a part of the chunk @@ -390,16 +385,15 @@ def _refine_suffix(self, chunks: list[Chunk], effective_context_size: int) -> li The refined chunks. """ - # Iterate over the chunks till the second to last chunk for i, chunk in enumerate(chunks[:-1]): - # Get the previous chunk prev_chunk = chunks[i + 1] - # Calculate effective context size per chunk if context_size is a float + # Calculate context size based on the chunk RECEIVING context (its own token count) + # This ensures each chunk gets overlap proportional to its own size if isinstance(self.context_size, float): - effective_context_size = int(self.context_size * prev_chunk.token_count) + effective_context_size = int(self.context_size * chunk.token_count) - # Calculate the overlap context + # Get context from the previous chunk context = self._get_suffix_overlap_context(prev_chunk, effective_context_size) # Set it as a part of the chunk diff --git a/tests/refinery/test_overlap_refinery.py b/tests/refinery/test_overlap_refinery.py index e3c355530..b6ed6e24f 100644 --- a/tests/refinery/test_overlap_refinery.py +++ b/tests/refinery/test_overlap_refinery.py @@ -751,17 +751,16 @@ def test_overlap_refinery_invalid_modes() -> None: def test_overlap_refinery_context_size_reuse_correctness() -> None: """Test that reusing OverlapRefinery with float context_size works correctly with different chunk sets.""" - # This tests the fix for a bug where _calculated_context_size was incorrectly cached refinery = OverlapRefinery(context_size=0.3, mode="token", method="suffix") - # First set: small token counts -> context_size should be 0.3 * 5 = 1.5 -> 1 + # First set: enough tokens so context is not empty (0.3 * 10 = 3) small_chunks = [ - Chunk(text="Short text", start_index=0, end_index=9, token_count=2), + Chunk(text="Short text", start_index=0, end_index=9, token_count=10), Chunk(text="Another brief chunk here", start_index=10, end_index=33, token_count=5), ] refined_small = refinery.refine([c.copy() for c in small_chunks]) - # Second set: large token counts -> context_size should be 0.3 * 20 = 6, NOT cached 1 + # Second set: large token counts large_chunks = [ Chunk( text="This is a significantly longer text chunk with many more tokens", @@ -783,19 +782,13 @@ def test_overlap_refinery_context_size_reuse_correctness() -> None: large_context = getattr(refined_large[0], "context", "") # Verify that different context sizes were actually calculated - # We can't directly access the calculated context size, but we can verify behavior - # by checking that the chunks were processed correctly assert len(refined_small) == 2 assert len(refined_large) == 2 - # At minimum, ensure both contexts exist and are reasonable + # Each chunk gets context proportional to its own size assert small_context is not None and small_context != "" assert large_context is not None and large_context != "" - # The key test: if the bug existed, both would use the same context size - # With the fix, they should use different context sizes based on their respective max token counts - # This is hard to test directly, but we've verified the calculation is correct above - def test_overlap_refinery_repr() -> None: """Test the OverlapRefinery.__repr__ method.""" From be5131fc9d025e819e9402eff261315402f67329 Mon Sep 17 00:00:00 2001 From: anaslimem Date: Thu, 16 Apr 2026 22:48:56 +0100 Subject: [PATCH 2/2] fix small mistake in the comment of the suffix next instead of prev --- src/chonkie/refinery/overlap.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/chonkie/refinery/overlap.py b/src/chonkie/refinery/overlap.py index bc8a54109..46802ba5c 100644 --- a/src/chonkie/refinery/overlap.py +++ b/src/chonkie/refinery/overlap.py @@ -393,7 +393,7 @@ def _refine_suffix(self, chunks: list[Chunk], effective_context_size: int) -> li if isinstance(self.context_size, float): effective_context_size = int(self.context_size * chunk.token_count) - # Get context from the previous chunk + # Get context from the next chunk context = self._get_suffix_overlap_context(prev_chunk, effective_context_size) # Set it as a part of the chunk