Spaces:
Running on Zero
Running on Zero
File size: 3,699 Bytes
712c393 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | 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
|