Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import re | |
| from typing import List, Sequence | |
| from ..ingestion.models import DocumentBlock | |
| from .base import BaseChunker | |
| from .models import Chunk | |
| class CodeChunker(BaseChunker): | |
| """ | |
| Code-aware chunker that identifies functions, classes, and top-level | |
| blocks, preserving scope headers and indentation context. | |
| """ | |
| # Matches class and def in Python, function/export in JS/TS, etc. | |
| DEF_REGEX = re.compile( | |
| r"^(?:async\s+)?(?:def|class|function|export\s+(?:default\s+)?(?:class|function))\s+([a-zA-Z_0-9]+)", | |
| re.MULTILINE, | |
| ) | |
| def chunk(self, blocks: Sequence[DocumentBlock], document_id: str) -> List[Chunk]: | |
| chunks: List[Chunk] = [] | |
| target_chars = int(self.config.chunk_size * self.config.chars_per_token) | |
| for block_idx, block in enumerate(blocks): | |
| code_text = block.text | |
| if not code_text.strip(): | |
| continue | |
| lines = code_text.split("\n") | |
| symbols: List[tuple[str, int, int]] = [] # (symbol_name, start_line, end_line) | |
| for i, line in enumerate(lines): | |
| match = self.DEF_REGEX.match(line) | |
| if match: | |
| symbols.append((match.group(0).strip(), i, len(lines))) | |
| # If no definitions or single short definition, emit as single chunk | |
| if not symbols or (len(symbols) == 1 and len(code_text) <= target_chars): | |
| chunk_id = self.make_chunk_id(document_id, f"code_b{block_idx}_0") | |
| chunks.append( | |
| Chunk( | |
| chunk_id=chunk_id, | |
| document_id=document_id, | |
| text=code_text.strip(), | |
| chunk_type="code", | |
| page=block.page, | |
| section=block.section, | |
| source_locator=block.source_locator or f"code:block:{block_idx}", | |
| token_count=self.estimate_tokens(code_text), | |
| metadata={ | |
| **block.metadata, | |
| "is_code": True, | |
| }, | |
| ) | |
| ) | |
| continue | |
| # Chunk by identified functions/classes | |
| for sym_idx in range(len(symbols)): | |
| name, start_l, _ = symbols[sym_idx] | |
| end_l = symbols[sym_idx + 1][1] if sym_idx + 1 < len(symbols) else len(lines) | |
| snippet = "\n".join(lines[start_l:end_l]).strip() | |
| if not snippet: | |
| continue | |
| chunk_id = self.make_chunk_id( | |
| document_id, f"code_b{block_idx}_{sym_idx}" | |
| ) | |
| chunks.append( | |
| Chunk( | |
| chunk_id=chunk_id, | |
| document_id=document_id, | |
| text=snippet, | |
| chunk_type="code", | |
| page=block.page, | |
| section=name, | |
| source_locator=f"{block.source_locator or f'code:{block_idx}'}:lines_{start_l+1}-{end_l}", | |
| token_count=self.estimate_tokens(snippet), | |
| metadata={ | |
| **block.metadata, | |
| "symbol": name, | |
| "is_code": True, | |
| "line_start": start_l + 1, | |
| "line_end": end_l, | |
| }, | |
| ) | |
| ) | |
| return chunks | |