Team Ai
Apppublic

kyorlin2001/code-analysis-tool

sourceHugging Faceapache-2.0updated 6mo agoView on Hugging Face
0likes
chunker.py86 linesDownload Raw Back to rag
1from __future__ import annotations2 3from dataclasses import dataclass4from typing import Iterable5 6from models.retrieval_chunk import RetrievalChunk7 8 9@dataclass(frozen=True)10class ChunkingConfig:11    """12    Configuration for line-based chunking.13    """14 15    chunk_size_lines: int = 8016    overlap_lines: int = 1517    max_file_lines: int = 2_00018 19 20class Chunker:21    """22    Splits raw file content into retrieval chunks.23    """24 25    def __init__(self, config: ChunkingConfig | None = None) -> None:26        self.config = config or ChunkingConfig()27 28    def chunk_text(29        self,30        text: str,31        file_path: str,32        language: str | None = None,33    ) -> list[RetrievalChunk]:34        lines = text.splitlines()35        if not lines:36            return []37 38        if len(lines) > self.config.max_file_lines:39            lines = lines[: self.config.max_file_lines]40 41        chunks: list[RetrievalChunk] = []42        step = max(1, self.config.chunk_size_lines - self.config.overlap_lines)43 44        start = 045        chunk_index = 046 47        while start < len(lines):48            end = min(len(lines), start + self.config.chunk_size_lines)49            chunk_lines = lines[start:end]50 51            if not chunk_lines:52                break53 54            chunk_text = "\n".join(chunk_lines).strip()55            if chunk_text:56                chunk_id = f"{file_path}::chunk-{chunk_index}"57                chunks.append(58                    RetrievalChunk(59                        text=chunk_text,60                        file_path=file_path,61                        chunk_id=chunk_id,62                        language=language,63                        start_line=start + 1,64                        end_line=end,65                    )66                )67 68            chunk_index += 169            start += step70 71        return chunks72 73    def chunk_many(74        self,75        items: Iterable[tuple[str, str, str | None]],76    ) -> list[RetrievalChunk]:77        """78        Chunk many files at once.79 80        Each item is expected to be:81        (text, file_path, language)82        """83        chunks: list[RetrievalChunk] = []84        for text, file_path, language in items:85            chunks.extend(self.chunk_text(text=text, file_path=file_path, language=language))86        return chunks