Switch to chonkie from llamaindex chunker (onyx-dot-app#4838)

* Switch to chonkie from llamaindex chunker

* Remove un-intended changes

* Order requirements

* Upgrade chonkie version

Author: Weves

backend/onyx/indexing/chunker.py

@@ -1,3 +1,7 @@
+from typing import cast
+
+from chonkie import SentenceChunker
+
 from onyx.configs.app_configs import AVERAGE_SUMMARY_EMBEDDINGS
 from onyx.configs.app_configs import BLURB_SIZE
 from onyx.configs.app_configs import LARGE_CHUNK_RATIO
@@ -135,9 +139,6 @@ def __init__(
         mini_chunk_size: int = MINI_CHUNK_SIZE,
         callback: IndexingHeartbeatInterface | None = None,
     ) -> None:
-        # importing llama_index uses a lot of RAM, so we only import it when needed.
-        from llama_index.core.node_parser import SentenceSplitter
-
         self.include_metadata = include_metadata
         self.chunk_token_limit = chunk_token_limit
         self.enable_multipass = enable_multipass
@@ -156,23 +157,30 @@ def __init__(
         self.max_context = 0
         self.prompt_tokens = 0
 
-        self.blurb_splitter = SentenceSplitter(
-            tokenizer=tokenizer.tokenize,
+        # Create a token counter function that returns the count instead of the tokens
+        def token_counter(text: str) -> int:
+            return len(tokenizer.encode(text))
+
+        self.blurb_splitter = SentenceChunker(
+            tokenizer_or_token_counter=token_counter,
             chunk_size=blurb_size,
             chunk_overlap=0,
+            return_type="texts",
         )
 
-        self.chunk_splitter = SentenceSplitter(
-            tokenizer=tokenizer.tokenize,
+        self.chunk_splitter = SentenceChunker(
+            tokenizer_or_token_counter=token_counter,
             chunk_size=chunk_token_limit,
             chunk_overlap=chunk_overlap,
+            return_type="texts",
         )
 
         self.mini_chunk_splitter = (
-            SentenceSplitter(
-                tokenizer=tokenizer.tokenize,
+            SentenceChunker(
+                tokenizer_or_token_counter=token_counter,
                 chunk_size=mini_chunk_size,
                 chunk_overlap=0,
+                return_type="texts",
             )
             if enable_multipass
             else None
@@ -199,7 +207,8 @@ def _extract_blurb(self, text: str) -> str:
         """
         Extract a short blurb from the text (first chunk of size `blurb_size`).
         """
-        texts = self.blurb_splitter.split_text(text)
+        # chunker is in `text` mode
+        texts = cast(list[str], self.blurb_splitter.chunk(text))
         if not texts:
             return ""
         return texts[0]
@@ -209,7 +218,8 @@ def _get_mini_chunk_texts(self, chunk_text: str) -> list[str] | None:
         For "multipass" mode: additional sub-chunks (mini-chunks) for use in certain embeddings.
         """
         if self.mini_chunk_splitter and chunk_text.strip():
-            return self.mini_chunk_splitter.split_text(chunk_text)
+            # chunker is in `text` mode
+            return cast(list[str], self.mini_chunk_splitter.chunk(chunk_text))
         return None
 
     # ADDED: extra param image_url to store in the chunk
@@ -329,7 +339,8 @@ def _chunk_document_with_sections(
                     chunk_text = ""
                     link_offsets = {}
 
-                split_texts = self.chunk_splitter.split_text(section_text)
+                # chunker is in `text` mode
+                split_texts = cast(list[str], self.chunk_splitter.chunk(section_text))
                 for i, split_text in enumerate(split_texts):
                     # If even the split_text is bigger than strict limit, further split
                     if (

backend/requirements/default.txt

@@ -7,6 +7,7 @@ beautifulsoup4==4.12.3
 boto3==1.36.23
 celery==5.5.1
 chardet==5.2.0
+chonkie==1.0.10
 dask==2023.8.1
 ddtrace==2.6.5
 discord.py==2.4.0
@@ -42,7 +43,6 @@ langgraph-sdk==0.1.44
 litellm==1.72.2 
 lxml==5.3.0
 lxml_html_clean==0.2.2
-llama-index==0.12.28
 Mako==1.2.4
 msal==1.28.0
 nltk==3.9.1