chonkie-inc · chonknick · May 23, 2025 · May 23, 2025 · May 23, 2025 · May 23, 2025
diff --git a/src/chonkie/cloud/__init__.py b/src/chonkie/cloud/__init__.py
@@ -2,6 +2,7 @@
 
 from .chunker import (
     CloudChunker,
+    CodeChunker,
     LateChunker,
     NeuralChunker,
     RecursiveChunker,
@@ -20,6 +21,7 @@
     "SentenceChunker",
     "LateChunker",
     "SDPMChunker",
+    "CodeChunker",
     "NeuralChunker",
     "SlumberChunker",
 ]
diff --git a/src/chonkie/cloud/chunker/__init__.py b/src/chonkie/cloud/chunker/__init__.py
@@ -1,6 +1,7 @@
 """Module for Chonkie Cloud Chunkers."""
 
 from .base import CloudChunker
+from .code import CodeChunker
 from .late import LateChunker
 from .neural import NeuralChunker
 from .recursive import RecursiveChunker
@@ -18,6 +19,7 @@
     "SentenceChunker",
     "LateChunker",
     "SDPMChunker",
+    "CodeChunker",
     "NeuralChunker",
     "SlumberChunker",
 ]
diff --git a/src/chonkie/cloud/chunker/code.py b/src/chonkie/cloud/chunker/code.py
@@ -0,0 +1,114 @@
+"""Code Chunking for Chonkie API."""
+
+import os
+from typing import Dict, List, Literal, Optional, Union, cast
+
+import requests
+
+from .base import CloudChunker
+
+
+class CodeChunker(CloudChunker):
+    """Code Chunking for Chonkie API."""
+
+    BASE_URL = "https://api.chonkie.ai"
+    VERSION = "v1"
+
+    def __init__(
+        self,
+        tokenizer_or_token_counter: str = "gpt2",
+        chunk_size: int = 512,
+        language: Union[Literal["auto"], str] = "auto",
+        return_type: Literal["texts", "chunks"] = "chunks",
+        api_key: Optional[str] = None,
+    ) -> None:
+        """Initialize the Cloud CodeChunker.
+
+        Args:
+            tokenizer_or_token_counter: The tokenizer or token counter to use.
+            chunk_size: The size of the chunks to create.
+                language: The language of the code to parse. Accepts any of the languages supported by tree-sitter-language-pack.
+            return_type: The type of the return value.
+            api_key: The API key for the Chonkie API.
+
+        Raises:
+            ValueError: If the API key is not provided or if parameters are invalid.
+
+        """
+        # If no API key is provided, use the environment variable
+        self.api_key = api_key or os.getenv("CHONKIE_API_KEY")
+        if not self.api_key:
+            raise ValueError(
+                "No API key provided. Please set the CHONKIE_API_KEY environment variable"
+                + " or pass an API key to the CodeChunker constructor."
+            )
+
+        # Validate parameters
+        if chunk_size <= 0:
+            raise ValueError("Chunk size must be greater than 0.")
+        if return_type not in ["texts", "chunks"]:
+            raise ValueError("Return type must be either 'texts' or 'chunks'.")
+
+        # Assign all the attributes to the instance
+        self.tokenizer_or_token_counter = tokenizer_or_token_counter
+        self.chunk_size = chunk_size
+        self.language = language
+        self.return_type = return_type
+
+        # Check if the API is up right now
+        response = requests.get(f"{self.BASE_URL}/")
+        if response.status_code != 200:
+            raise ValueError(
+                "Oh no! You caught Chonkie at a bad time. It seems to be down right now."
+                + " Please try again in a short while."
+                + " If the issue persists, please contact support at support@chonkie.ai or raise an issue on GitHub."
+            )
+
+    def chunk(self, text: Union[str, List[str]]) -> List[Dict]:
-    def chunk(self, text: Union[str, List[str]]) -> List[Dict]:
+def chunk(self, text: Union[str, List[str]]) -> Union[List[Dict], List[List[Dict]]]:
-    def chunk(self, text: Union[str, List[str]]) -> List[Dict]:
+def chunk(self, text: Union[str, List[str]]) -> Union[List[Dict], List[List[Dict]]]:
+        """Chunk the code into a list of chunks.
+
+        Args:
+            text: The code text(s) to chunk.
+
+        Returns:
+            A list of chunk dictionaries containing the chunked code.
+
+        Raises:
+            ValueError: If the API request fails or returns invalid data.
+
+        """
+        # Define the payload for the request
+        payload = {
+            "text": text,
+            "tokenizer_or_token_counter": self.tokenizer_or_token_counter,
+            "chunk_size": self.chunk_size,
+            "language": self.language,
+            "include_nodes": False,  # API doesn't support tree-sitter nodes
+            "return_type": self.return_type,
+        }
+
+        # Make the request to the Chonkie API
+        response = requests.post(
+            f"{self.BASE_URL}/{self.VERSION}/chunk/code",
+            json=payload,
+            headers={"Authorization": f"Bearer {self.api_key}"},
+        )
+
+        # Check if the response is successful
+        if response.status_code != 200:
+            raise ValueError(
+                f"Error from the Chonkie API: {response.status_code} {response.text}"
+            )
+
+        # Parse the response
+        try:
+            result: List[Dict] = cast(List[Dict], response.json())
+        except Exception as error:
+            raise ValueError(f"Error parsing the response: {error}") from error
-        except Exception as error:
-            raise ValueError(f"Error parsing the response: {error}") from error
+        except requests.exceptions.JSONDecodeError as error:
+            raise ValueError(f"Error parsing the response: {error}") from error
-        except Exception as error:
-            raise ValueError(f"Error parsing the response: {error}") from error
+        except requests.exceptions.JSONDecodeError as error:
+            raise ValueError(f"Error parsing the response: {error}") from error
+
+        # Return the result
+        return result
+
+    def __call__(self, text: Union[str, List[str]]) -> List[Dict]:
+        """Call the chunker."""
+        return self.chunk(text)