Team Ai
Apppublic

kyorlin2001/code-analysis-tool

sourceHugging Faceapache-2.0updated 6mo agoView on Hugging Face
0likes
repo_chunk_loader.py91 linesDownload Raw Back to rag
1from __future__ import annotations2 3from dataclasses import dataclass4from pathlib import Path5 6from models.retrieval_chunk import RetrievalChunk7from rag.chunker import Chunker8 9 10@dataclass(frozen=True)11class LoadedChunkBundle:12    """13    Container for chunks loaded from a repository.14    """15 16    chunks: list[RetrievalChunk]17 18 19class RepoChunkLoader:20    """21    Loads repository files and converts them into retrieval chunks.22    """23 24    def __init__(self, chunker: Chunker | None = None) -> None:25        self.chunker = chunker or Chunker()26 27    def load_from_files(self, repo_root: str, files: list[str]) -> LoadedChunkBundle:28        """29        Load chunks from a repository root and a list of relative file paths.30        """31        root_path = Path(repo_root)32        chunks: list[RetrievalChunk] = []33 34        for file_path in files:35            absolute_path = root_path / file_path36            if not absolute_path.is_file():37                continue38 39            try:40                text = absolute_path.read_text(encoding="utf-8")41            except UnicodeDecodeError:42                continue43            except OSError:44                continue45 46            language = self._infer_language(file_path)47            chunks.extend(48                self.chunker.chunk_text(49                    text=text,50                    file_path=file_path,51                    language=language,52                )53            )54 55        return LoadedChunkBundle(chunks=chunks)56 57    def load_from_root(self, repo_root: str) -> LoadedChunkBundle:58        """59        Load chunks from every file under a repository root.60        """61        root_path = Path(repo_root)62        files = [63            str(path.relative_to(root_path))64            for path in root_path.rglob("*")65            if path.is_file()66        ]67        return self.load_from_files(repo_root, files)68 69    def _infer_language(self, file_path: str) -> str | None:70        suffix = Path(file_path).suffix.lower()71 72        language_map = {73            ".py": "Python",74            ".js": "JavaScript",75            ".jsx": "JavaScript React",76            ".ts": "TypeScript",77            ".tsx": "TypeScript React",78            ".java": "Java",79            ".kt": "Kotlin",80            ".go": "Go",81            ".rs": "Rust",82            ".md": "Markdown",83            ".txt": "Text",84            ".json": "JSON",85            ".yaml": "YAML",86            ".yml": "YAML",87            ".toml": "TOML",88            ".ini": "INI",89        }90 91        return language_map.get(suffix)