From 8d4c79898e45f05692b3709e5f56840457b81e60 Mon Sep 17 00:00:00 2001 From: AlanPonnachan Date: Fri, 10 Oct 2025 18:04:38 +0530 Subject: [PATCH 1/6] update Chunk data class in base.py --- src/chonkie/types/base.py | 34 ++++++++++++++++++++++++++-------- 1 file changed, 26 insertions(+), 8 deletions(-) diff --git a/src/chonkie/types/base.py b/src/chonkie/types/base.py index 2badb175e..cfe96479d 100644 --- a/src/chonkie/types/base.py +++ b/src/chonkie/types/base.py @@ -22,6 +22,8 @@ class Chunk: start_index (int): The starting index of the chunk in the original text. end_index (int): The ending index of the chunk in the original text. token_count (int): The number of tokens in the chunk. + start_line (Optional[int]): The starting line number of the chunk in the original text (1-indexed). + end_line (Optional[int]): The ending line number of the chunk in the original text (1-indexed). context (Optional[str]): Optional context metadata for the chunk. embedding (Union[List[float], np.ndarray, None]): Optional embedding vector for the chunk, either as a list of floats or a numpy array. @@ -33,6 +35,8 @@ class Chunk: start_index: int = field(default=0) end_index: int = field(default=0) token_count: int = field(default=0) + start_line: Optional[int] = field(default=None) + end_line: Optional[int] = field(default=None) context: Optional[str] = field(default=None) embedding: Union[List[float], "np.ndarray", None] = field(default=None) @@ -80,15 +84,17 @@ def _preview_embedding(self) -> str: def __repr__(self) -> str: """Return a detailed string representation of the Chunk.""" - repr = ( - f"Chunk(text='{self.text}', token_count={self.token_count}, " + repr_str = ( + f"Chunk(text='{self.text[:50]}...', token_count={self.token_count}, " f"start_index={self.start_index}, end_index={self.end_index}" ) + if self.start_line is not None and self.end_line is not None: + repr_str += f", start_line={self.start_line}, end_line={self.end_line}" if self.context: - repr += f", context={self.context}" + repr_str += f", context={self.context}" if self.embedding is not None: - repr += f", embedding={self._preview_embedding()}" - return repr + ")" + repr_str += f", embedding={self._preview_embedding()}" + return repr_str + ")" def __iter__(self) -> Iterator[str]: """Return an iterator over the chunk's text.""" @@ -100,14 +106,24 @@ def __getitem__(self, index: int) -> str: def to_dict(self) -> dict: """Return the Chunk as a dictionary.""" - result = self.__dict__.copy() - result["context"] = self.context + result = { + "id": self.id, + "text": self.text, + "start_index": self.start_index, + "end_index": self.end_index, + "token_count": self.token_count, + "start_line": self.start_line, + "end_line": self.end_line, + "context": self.context, + } # Convert embedding to list if it has tolist method (numpy array) if self.embedding is not None: if hasattr(self.embedding, 'tolist'): result["embedding"] = self.embedding.tolist() else: result["embedding"] = self.embedding + else: + result["embedding"] = None return result @classmethod @@ -119,10 +135,12 @@ def from_dict(cls, data: dict) -> "Chunk": start_index=data["start_index"], end_index=data["end_index"], token_count=data["token_count"], + start_line=data.get("start_line", None), + end_line=data.get("end_line", None), context=data.get("context", None), embedding=data.get("embedding", None), ) def copy(self) -> "Chunk": """Return a deep copy of the chunk.""" - return Chunk.from_dict(self.to_dict()) + return Chunk.from_dict(self.to_dict()) \ No newline at end of file From 07df58632748e9ff6658cc6fcdd83b0932436acb Mon Sep 17 00:00:00 2001 From: AlanPonnachan Date: Fri, 10 Oct 2025 18:23:40 +0530 Subject: [PATCH 2/6] enhance base.py --- src/chonkie/types/base.py | 16 +++------------- 1 file changed, 3 insertions(+), 13 deletions(-) diff --git a/src/chonkie/types/base.py b/src/chonkie/types/base.py index cfe96479d..4a14a9d03 100644 --- a/src/chonkie/types/base.py +++ b/src/chonkie/types/base.py @@ -106,24 +106,14 @@ def __getitem__(self, index: int) -> str: def to_dict(self) -> dict: """Return the Chunk as a dictionary.""" - result = { - "id": self.id, - "text": self.text, - "start_index": self.start_index, - "end_index": self.end_index, - "token_count": self.token_count, - "start_line": self.start_line, - "end_line": self.end_line, - "context": self.context, - } + result = self.__dict__.copy() + result["context"] = self.context # Convert embedding to list if it has tolist method (numpy array) if self.embedding is not None: if hasattr(self.embedding, 'tolist'): result["embedding"] = self.embedding.tolist() else: result["embedding"] = self.embedding - else: - result["embedding"] = None return result @classmethod @@ -143,4 +133,4 @@ def from_dict(cls, data: dict) -> "Chunk": def copy(self) -> "Chunk": """Return a deep copy of the chunk.""" - return Chunk.from_dict(self.to_dict()) \ No newline at end of file + return Chunk.from_dict(self.to_dict()) From 10e53ebf1e1768dab5cdc6af938ece605679ecbf Mon Sep 17 00:00:00 2001 From: AlanPonnachan Date: Fri, 10 Oct 2025 18:29:05 +0530 Subject: [PATCH 3/6] extract the line numbers from the tree-sitter --- src/chonkie/chunker/code.py | 42 +++++++++++++++++++++---------------- 1 file changed, 24 insertions(+), 18 deletions(-) diff --git a/src/chonkie/chunker/code.py b/src/chonkie/chunker/code.py index 3e782ff31..c55dc3d60 100644 --- a/src/chonkie/chunker/code.py +++ b/src/chonkie/chunker/code.py @@ -300,42 +300,44 @@ def _get_texts_from_node_groups(self, return chunk_texts def _create_chunks(self, - texts: List[str], - token_counts: List[int], - node_groups: List[List["Node"]]) -> List[Chunk]: + texts: List[str], + token_counts: List[int], + node_groups: List[List["Node"]]) -> List[Chunk]: """Create Code Chunks.""" chunks = [] current_index = 0 - for i in range(len(texts)): - text = texts[i] - token_count = token_counts[i] + for i, (text, token_count, group) in enumerate(zip(texts, token_counts, node_groups)): + start_line = None + end_line = None + if group: + # tree-sitter nodes provide 0-indexed row numbers, so add 1 for 1-indexed line numbers. + start_line = group[0].start_point[0] + 1 + end_line = group[-1].end_point[0] + 1 chunks.append(Chunk(text=text, start_index=current_index, end_index=current_index + len(text), - token_count=token_count)) # type: ignore[attr-defined] + token_count=token_count, + start_line=start_line, + end_line=end_line)) current_index += len(text) return chunks def chunk(self, text: str) -> List[Chunk]: """Recursively chunks the code based on context from tree-sitter.""" if not text.strip(): # Handle empty or whitespace-only input - logger.debug("Empty or whitespace-only code provided") return [] - logger.debug(f"Starting code chunking for text of length {len(text)}") - original_text_bytes = text.encode("utf-8") # Store bytes - + # At this point, if the language is auto, we need to detect the language # and initialize the parser if self.language == "auto": language = self._detect_language(original_text_bytes) - logger.info(f"Auto-detected code language: {language}") self.parser = get_parser(language) # type: ignore - else: - logger.debug(f"Using configured language: {self.language}") + # Initialize node_groups here to ensure it's available in the finally block scope + node_groups = [] try: # Create the parsing tree for the current code tree: Tree = self.parser.parse(original_text_bytes) # type: ignore @@ -344,14 +346,18 @@ def chunk(self, text: str) -> List[Chunk]: # Get the node_groups node_groups, token_counts = self._group_child_nodes(root_node) texts: List[str] = self._get_texts_from_node_groups(node_groups, original_text_bytes) + + chunks = self._create_chunks(texts, token_counts, node_groups) finally: # Clean up the tree and root_node if they are not needed if not self.include_nodes: - del tree, root_node - node_groups = [] + if 'tree' in locals(): + del tree + if 'root_node' in locals(): + del root_node + # Clear node_groups only after it's been used by _create_chunks + node_groups.clear() - chunks = self._create_chunks(texts, token_counts, node_groups) - logger.info(f"Created {len(chunks)} code chunks from parsed syntax tree") return chunks def __repr__(self) -> str: From 79276c9028d9f851760f1971dd64d2dc319b6245 Mon Sep 17 00:00:00 2001 From: AlanPonnachan Date: Fri, 10 Oct 2025 18:39:05 +0530 Subject: [PATCH 4/6] improve code.py --- src/chonkie/chunker/code.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/src/chonkie/chunker/code.py b/src/chonkie/chunker/code.py index c55dc3d60..6157e4ab2 100644 --- a/src/chonkie/chunker/code.py +++ b/src/chonkie/chunker/code.py @@ -300,9 +300,9 @@ def _get_texts_from_node_groups(self, return chunk_texts def _create_chunks(self, - texts: List[str], - token_counts: List[int], - node_groups: List[List["Node"]]) -> List[Chunk]: + texts: List[str], + token_counts: List[int], + node_groups: List[List["Node"]]) -> List[Chunk]: """Create Code Chunks.""" chunks = [] current_index = 0 @@ -326,15 +326,21 @@ def _create_chunks(self, def chunk(self, text: str) -> List[Chunk]: """Recursively chunks the code based on context from tree-sitter.""" if not text.strip(): # Handle empty or whitespace-only input + logger.debug("Empty or whitespace-only code provided") return [] + logger.debug(f"Starting code chunking for text of length {len(text)}") + original_text_bytes = text.encode("utf-8") # Store bytes - + # At this point, if the language is auto, we need to detect the language # and initialize the parser if self.language == "auto": language = self._detect_language(original_text_bytes) + logger.info(f"Auto-detected code language: {language}") self.parser = get_parser(language) # type: ignore + else: + logger.debug(f"Using configured language: {self.language}") # Initialize node_groups here to ensure it's available in the finally block scope node_groups = [] @@ -343,12 +349,15 @@ def chunk(self, text: str) -> List[Chunk]: tree: Tree = self.parser.parse(original_text_bytes) # type: ignore root_node: Node = tree.root_node # type: ignore - # Get the node_groups + # Get the node_groups node_groups, token_counts = self._group_child_nodes(root_node) texts: List[str] = self._get_texts_from_node_groups(node_groups, original_text_bytes) + # Create chunks inside the try block to ensure node_groups is available chunks = self._create_chunks(texts, token_counts, node_groups) - finally: + logger.info(f"Created {len(chunks)} code chunks from parsed syntax tree") + + finally: # Clean up the tree and root_node if they are not needed if not self.include_nodes: if 'tree' in locals(): @@ -358,7 +367,7 @@ def chunk(self, text: str) -> List[Chunk]: # Clear node_groups only after it's been used by _create_chunks node_groups.clear() - return chunks + return chunks def __repr__(self) -> str: """Return the string representation of the CodeChunker.""" From 36c9c6bec7d6408173ac12040af1b68d271eb04d Mon Sep 17 00:00:00 2001 From: AlanPonnachan Date: Fri, 10 Oct 2025 18:46:03 +0530 Subject: [PATCH 5/6] update test --- tests/chunkers/test_code_chunker.py | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/tests/chunkers/test_code_chunker.py b/tests/chunkers/test_code_chunker.py index 887702ad3..4725cea25 100644 --- a/tests/chunkers/test_code_chunker.py +++ b/tests/chunkers/test_code_chunker.py @@ -165,4 +165,27 @@ def test_code_chunker_chunk_size_javascript(js_code: str) -> None: chunks = chunker.chunk(js_code) # Allow for some leeway assert all(chunk.token_count < chunk_size + 15 for chunk in chunks[:-1]) - assert chunks[-1].token_count > 0 \ No newline at end of file + assert chunks[-1].token_count > 0 + + +def test_code_chunker_adds_line_numbers(python_code: str) -> None: + """Test that the CodeChunker correctly adds start and end line numbers to chunks.""" + # Use a chunk size that will reliably split the python_code fixture + chunker = CodeChunker(language="python", chunk_size=100) + chunks = chunker.chunk(python_code) + + assert len(chunks) > 0, "Chunker should produce at least one chunk" + + # Verify that every chunk has valid, 1-indexed line numbers + for chunk in chunks: + assert chunk.start_line is not None, "start_line should not be None" + assert chunk.end_line is not None, "end_line should not be None" + assert isinstance(chunk.start_line, int) + assert isinstance(chunk.end_line, int) + assert chunk.start_line > 0, "Line numbers should be 1-indexed and positive" + assert chunk.end_line >= chunk.start_line, "End line must be greater than or equal to start line" + + # The python_code fixture starts with a newline, so the first line of code is line 1. + # The code itself has 18 lines. + assert chunks[0].start_line == 1, "The first chunk should start on the first line of code" + assert chunks[-1].end_line == 18, "The last chunk should end on the last line of code" \ No newline at end of file From e03399232a70d0cec46f93bb3fcf3ec2c96deceb Mon Sep 17 00:00:00 2001 From: AlanPonnachan Date: Fri, 10 Oct 2025 18:59:57 +0530 Subject: [PATCH 6/6] resolve issues --- src/chonkie/chunker/code.py | 1 + tests/chunkers/test_code_chunker.py | 8 ++++---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/chonkie/chunker/code.py b/src/chonkie/chunker/code.py index 6157e4ab2..c6fbcc119 100644 --- a/src/chonkie/chunker/code.py +++ b/src/chonkie/chunker/code.py @@ -344,6 +344,7 @@ def chunk(self, text: str) -> List[Chunk]: # Initialize node_groups here to ensure it's available in the finally block scope node_groups = [] + chunks = [] try: # Create the parsing tree for the current code tree: Tree = self.parser.parse(original_text_bytes) # type: ignore diff --git a/tests/chunkers/test_code_chunker.py b/tests/chunkers/test_code_chunker.py index 4725cea25..8a17ff72e 100644 --- a/tests/chunkers/test_code_chunker.py +++ b/tests/chunkers/test_code_chunker.py @@ -185,7 +185,7 @@ def test_code_chunker_adds_line_numbers(python_code: str) -> None: assert chunk.start_line > 0, "Line numbers should be 1-indexed and positive" assert chunk.end_line >= chunk.start_line, "End line must be greater than or equal to start line" - # The python_code fixture starts with a newline, so the first line of code is line 1. - # The code itself has 18 lines. - assert chunks[0].start_line == 1, "The first chunk should start on the first line of code" - assert chunks[-1].end_line == 18, "The last chunk should end on the last line of code" \ No newline at end of file + # The python_code fixture starts with a newline, so its content begins on line 2. + # The code content ends on line 19 of the input string. + assert chunks[0].start_line == 2, "The first chunk should start on the first line of code (line 2)" + assert chunks[-1].end_line == 19, "The last chunk should end on the last line of code (line 19)" \ No newline at end of file