diffrag 0.2.3__py3-none-any.whl

diffrag/indexing/indexer.py

@@ -0,0 +1,387 @@
+"""ChromaDB-backed repository indexer for RAG context retrieval."""
+
+import json
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+
+import chromadb
+import pathspec
+
+from ..ai.base import EmbeddingClient
+from ..exceptions import IndexingError
+from .splitter import CodeSplitter
+
+log = logging.getLogger(__name__)
+
+_DEFAULT_EXCLUDED_PATTERNS = [
+    ".git",
+    ".venv",
+    "__pycache__",
+    "node_modules",
+    ".idea",
+    ".mypy_cache",
+    ".ruff_cache",
+    "dist",
+    "build",
+]
+
+_COLLECTION_NAME = "repo_index"
+_EMBED_BATCH_SIZE = 32
+_MAX_FILE_SIZE = 1_000_000  # 1 MB
+_META_FILENAME = ".index_meta.json"
+
+
+@dataclass
+class _ScanResult:
+    files: list[Path]
+    n_gitignore: int = 0
+    n_pattern: int = 0
+    n_oversized: int = 0
+    n_binary: int = 0
+
+
+@dataclass(frozen=True)
+class ContextChunk:
+    """A retrieved chunk of repository context.
+
+    Attributes:
+        content: Text content of the chunk.
+        file_path: Repository-relative path of the source file.
+        chunk_index: Zero-based index of this chunk within the source file.
+        distance: Cosine distance to the query vector (lower = more similar).
+        language: Detected language (e.g. ``"python"``), or ``None``.
+        node_type: AST node category (``"function"``, ``"class"``, ``"preamble"``),
+            or ``None`` when no syntax-aware split was performed.
+        symbol_name: Qualified symbol name (e.g. ``"MyClass.my_method"``), or ``None``.
+    """
+
+    content: str
+    file_path: str
+    chunk_index: int
+    distance: float
+    language: str | None = None
+    node_type: str | None = None
+    symbol_name: str | None = None
+
+
+class RepoIndexer:
+    """Builds and queries a ChromaDB vector index of a local repository.
+
+    Files are read, split into overlapping character chunks, embedded, and
+    stored in a persistent ChromaDB collection.  At query time, the query
+    text is embedded and the nearest chunks are returned.
+
+    Args:
+        db_path: Directory where the ChromaDB database is persisted.
+        embedding_client: Client used to compute text embeddings.
+        chunk_size: Maximum characters per chunk.
+        chunk_overlap: Characters of overlap between consecutive chunks.
+        excluded_patterns: Directory or file name segments to skip.
+        max_file_size: Files larger than this (bytes) are skipped.
+    """
+
+    def __init__(
+        self,
+        db_path: Path,
+        embedding_client: EmbeddingClient,
+        chunk_size: int = 1024,
+        chunk_overlap: int = 128,
+        excluded_patterns: list[str] | None = None,
+        max_file_size: int = _MAX_FILE_SIZE,
+    ) -> None:
+        self._db_path = db_path
+        self._embedding_client = embedding_client
+        self._chunk_size = chunk_size
+        self._chunk_overlap = chunk_overlap
+        self._excluded_patterns = set(excluded_patterns or _DEFAULT_EXCLUDED_PATTERNS)
+        self._max_file_size = max_file_size
+        self._splitter = CodeSplitter(chunk_size, chunk_overlap)
+
+        self._chroma = chromadb.PersistentClient(path=str(db_path))
+        self._collection = self._chroma.get_or_create_collection(_COLLECTION_NAME)
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    # ------------------------------------------------------------------
+    # Staleness helpers
+    # ------------------------------------------------------------------
+
+    def is_stale(self, repo_path: Path, commit_hash: str) -> bool:
+        """Return ``True`` if the index must be rebuilt.
+
+        The index is considered stale when the collection is empty, the stored
+        commit hash differs from *commit_hash*, or the stored repository path
+        differs from *repo_path*.
+
+        Args:
+            repo_path: Root directory of the repository being reviewed.
+            commit_hash: Current HEAD commit SHA of the repository.
+
+        Returns:
+            ``True`` when a rebuild is required, ``False`` when the cached index
+            is still valid.
+        """
+        if self.document_count == 0:
+            return True
+        meta = self._read_meta()
+        return meta.get("commit_hash") != commit_hash or meta.get("repo_path") != str(
+            repo_path.resolve()
+        )
+
+    def _meta_path(self) -> Path:
+        return self._db_path / _META_FILENAME
+
+    def _read_meta(self) -> dict:
+        try:
+            return json.loads(self._meta_path().read_text())
+        except (OSError, json.JSONDecodeError):
+            return {}
+
+    def _write_meta(self, repo_path: Path, commit_hash: str) -> None:
+        self._db_path.mkdir(parents=True, exist_ok=True)
+        self._meta_path().write_text(
+            json.dumps({"commit_hash": commit_hash, "repo_path": str(repo_path.resolve())})
+        )
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    @property
+    def document_count(self) -> int:
+        """Number of chunks currently stored in the index.
+
+        Returns:
+            Total chunk count in the ChromaDB collection.
+        """
+        return self._collection.count()
+
+    async def build_index(
+        self,
+        repo_path: Path,
+        on_progress: Callable[[str], None] | None = None,
+        commit_hash: str | None = None,
+    ) -> None:
+        """Rebuild the vector index from scratch for *repo_path*.
+
+        Deletes any existing index for this collection before re-indexing.
+        If *commit_hash* is provided, it is persisted to ``.index_meta.json``
+        so that future runs can skip rebuilding when the repo has not changed.
+
+        Args:
+            repo_path: Root directory of the repository to index.
+            on_progress: Optional callback invoked with a human-readable status
+                string at each major step.  Messages ending in ``"..."`` indicate
+                an ongoing operation; others mark a completed milestone.
+            commit_hash: HEAD commit SHA to associate with this index snapshot.
+
+        Raises:
+            IndexingError: If a critical error occurs during indexing.
+        """
+        _p = on_progress or (lambda _: None)
+        log.info("Building index for repository at %s", repo_path)
+
+        self._chroma.delete_collection(_COLLECTION_NAME)
+        self._collection = self._chroma.create_collection(_COLLECTION_NAME)
+
+        _p("Scanning repository...")
+        scan = self._gather_files(repo_path)
+        files = scan.files
+        log.info(
+            "File scan: %d indexed, %d .gitignore, %d patterns, %d binary, %d oversized",
+            len(files),
+            scan.n_gitignore,
+            scan.n_pattern,
+            scan.n_binary,
+            scan.n_oversized,
+        )
+
+        all_texts: list[str] = []
+        all_metas: list[dict] = []
+        all_ids: list[str] = []
+
+        for file_path in files:
+            rel = str(file_path.relative_to(repo_path))
+            try:
+                content = file_path.read_text(encoding="utf-8", errors="replace")
+            except OSError as exc:
+                log.warning("Skipping %s: %s", rel, exc)
+                continue
+
+            for idx, chunk in enumerate(self._splitter.split(content, rel)):
+                all_texts.append(chunk.content)
+                all_metas.append(
+                    {
+                        "file": rel,
+                        "chunk": idx,
+                        "language": chunk.language or "",
+                        "node_type": chunk.node_type or "",
+                        "symbol_name": chunk.symbol_name or "",
+                        "start_line": chunk.start_line,
+                        "end_line": chunk.end_line,
+                    }
+                )
+                all_ids.append(f"{rel}::chunk{idx}")
+
+        total_batches = max(1, len(all_texts) // _EMBED_BATCH_SIZE + 1)
+        _p(
+            f"Found {len(files)} files "
+            f"(excluded: {scan.n_gitignore} .gitignore · "
+            f"{scan.n_pattern} patterns · "
+            f"{scan.n_binary} binary · "
+            f"{scan.n_oversized} oversized) — {len(all_texts)} chunks"
+        )
+        log.info("Embedding %d chunks in batches of %d", len(all_texts), _EMBED_BATCH_SIZE)
+
+        try:
+            for batch_num, start in enumerate(range(0, len(all_texts), _EMBED_BATCH_SIZE), 1):
+                _p(f"Embedding chunks: batch {batch_num}/{total_batches}...")
+                batch_texts = all_texts[start : start + _EMBED_BATCH_SIZE]
+                batch_metas = all_metas[start : start + _EMBED_BATCH_SIZE]
+                batch_ids = all_ids[start : start + _EMBED_BATCH_SIZE]
+
+                resp = await self._embedding_client.embed(batch_texts)
+                self._collection.add(
+                    documents=batch_texts,
+                    embeddings=resp.embeddings,  # type: ignore
+                    metadatas=batch_metas,  # type: ignore
+                    ids=batch_ids,
+                )
+                log.debug("Indexed chunks %d–%d", start, start + len(batch_texts) - 1)
+        except Exception as exc:
+            raise IndexingError(f"Failed to build index: {exc}") from exc
+
+        count = self._collection.count()
+        if commit_hash is not None:
+            self._write_meta(repo_path, commit_hash)
+        _p(f"Index built: {count} chunks from {len(files)} files")
+        log.info("Index built: %d total chunks", count)
+
+    async def query(self, text: str, top_k: int = 5) -> list[ContextChunk]:
+        """Retrieve the *top_k* most relevant chunks for a query.
+
+        Args:
+            text: Query text; will be embedded before searching.
+            top_k: Number of results to return.
+
+        Returns:
+            List of :class:`ContextChunk` objects ordered by relevance.
+
+        Raises:
+            IndexingError: If the embedding call or ChromaDB query fails.
+        """
+        if self._collection.count() == 0:
+            log.warning("Index is empty — no context will be retrieved")
+            return []
+
+        try:
+            resp = await self._embedding_client.embed([text])
+            results = self._collection.query(
+                query_embeddings=resp.embeddings,  # type: ignore
+                n_results=min(top_k, self._collection.count()),
+                include=["documents", "metadatas", "distances"],
+            )
+        except Exception as exc:
+            raise IndexingError(f"Failed to query index: {exc}") from exc
+
+        def _str(v: object) -> str | None:
+            return str(v) if isinstance(v, str) and v else None
+
+        chunks: list[ContextChunk] = []
+        for doc, meta, dist in zip(
+            results["documents"][0],  # type: ignore
+            results["metadatas"][0],  # type: ignore
+            results["distances"][0],  # type: ignore
+            strict=False,
+        ):
+            chunks.append(
+                ContextChunk(
+                    content=str(doc),
+                    file_path=str(meta.get("file", "")),
+                    chunk_index=int(meta.get("chunk", 0)),  # type: ignore
+                    distance=float(dist),
+                    language=_str(meta.get("language")),
+                    node_type=_str(meta.get("node_type")),
+                    symbol_name=_str(meta.get("symbol_name")),
+                )
+            )
+        return chunks
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _load_gitignore_spec(repo_path: Path) -> "pathspec.PathSpec | None":
+        """Return a compiled gitignore spec for *repo_path*, or ``None``.
+
+        Reads the root ``.gitignore`` of the repository.  If the file does not
+        exist the method returns ``None`` and no gitignore filtering is applied.
+        """
+        gitignore = repo_path / ".gitignore"
+        if not gitignore.is_file():
+            return None
+        lines = gitignore.read_text(encoding="utf-8", errors="ignore").splitlines()
+        return pathspec.PathSpec.from_lines("gitignore", lines)
+
+    def _gather_files(self, repo_path: Path) -> _ScanResult:
+        """Collect all indexable files under *repo_path*.
+
+        Files are excluded if they match any pattern in the repository's root
+        ``.gitignore`` (via :meth:`_load_gitignore_spec`) or any segment of
+        their path matches an entry in :attr:`_excluded_patterns`.
+
+        Args:
+            repo_path: Root directory to search.
+
+        Returns:
+            :class:`_ScanResult` with the accepted file list and per-reason
+            exclusion counts.  Pass ``--verbose`` on the CLI to see per-file
+            ``DEBUG`` log lines.
+        """
+        gitspec = self._load_gitignore_spec(repo_path)
+        files: list[Path] = []
+        n_gitignore = n_pattern = n_oversized = n_binary = 0
+        for path in repo_path.rglob("*"):
+            if not path.is_file():
+                continue
+            rel = str(path.relative_to(repo_path))
+            if gitspec is not None and gitspec.match_file(rel):
+                n_gitignore += 1
+                log.debug("Excluded by .gitignore: %s", rel)
+                continue
+            if any(part in self._excluded_patterns for part in path.parts):
+                n_pattern += 1
+                log.debug("Excluded by pattern:    %s", rel)
+                continue
+            if path.stat().st_size > self._max_file_size:
+                n_oversized += 1
+                log.debug("Skipped (too large):    %s", rel)
+                continue
+            if self._is_binary(path):
+                n_binary += 1
+                log.debug("Skipped (binary):       %s", rel)
+                continue
+            log.debug("Indexing:               %s", rel)
+            files.append(path)
+        return _ScanResult(files, n_gitignore, n_pattern, n_oversized, n_binary)
+
+    @staticmethod
+    def _is_binary(path: Path) -> bool:
+        """Return ``True`` if the file appears to contain binary content.
+
+        Args:
+            path: File to inspect.
+
+        Returns:
+            ``True`` when a null byte is found in the first 8 KiB.
+        """
+        try:
+            with path.open("rb") as fh:
+                return b"\x00" in fh.read(8192)
+        except OSError:
+            return True

diffrag/indexing/splitter.py

@@ -0,0 +1,296 @@
+"""Syntax-aware code chunking for RAG indexing.
+
+Extension guide
+---------------
+To add syntax-aware splitting for a new language:
+
+1. Write a class that implements :class:`LanguageSplitter` (one method: ``split``).
+2. Add the file extension → language name mapping to :data:`EXTENSION_MAP`.
+3. Add the language name → splitter instance mapping to :data:`LANGUAGE_SPLITTERS`.
+
+Nothing else needs to change.  :class:`CodeSplitter` automatically falls back to
+character-based splitting for any extension not in the registry, or when a
+registered splitter returns ``[]`` (indicating a parse failure).
+"""
+
+import ast
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Protocol, runtime_checkable
+
+
+@dataclass(frozen=True)
+class SplitChunk:
+    """A chunk of source text with optional symbol-level metadata.
+
+    Attributes:
+        content: Raw text content of the chunk.
+        start_line: 1-based line number of the first line in the source file.
+        end_line: 1-based line number of the last line in the source file.
+        language: Detected language (e.g. ``"python"``), or ``None``.
+        node_type: AST node category (``"function"``, ``"class"``, ``"preamble"``),
+            or ``None`` when no syntax-aware split was performed.
+        symbol_name: Qualified name of the symbol (e.g. ``"MyClass.my_method"``),
+            or ``None``.
+    """
+
+    content: str
+    start_line: int
+    end_line: int
+    language: str | None = None
+    node_type: str | None = None
+    symbol_name: str | None = None
+
+
+def _char_split(
+    content: str,
+    chunk_size: int,
+    chunk_overlap: int,
+    *,
+    language: str | None = None,
+    node_type: str | None = None,
+    symbol_name: str | None = None,
+    line_offset: int = 0,
+) -> list[SplitChunk]:
+    """Character-based line-boundary splitting with optional metadata tagging.
+
+    This is the universal fallback used for unsupported file types and for
+    oversized AST nodes that cannot fit in a single chunk.
+
+    Args:
+        content: Source text to split.
+        chunk_size: Maximum characters per chunk.
+        chunk_overlap: Characters of overlap between consecutive chunks.
+        language: Language tag to attach to every produced chunk.
+        node_type: Node-type tag to attach to every produced chunk.
+        symbol_name: Symbol name to attach to every produced chunk.
+        line_offset: Number of source-file lines before the first line of
+            *content*.  Use ``node.lineno - 1`` when splitting a node
+            extracted from a larger file so that reported line numbers are
+            relative to the original file, not the node's text.
+
+    Returns:
+        List of non-empty :class:`SplitChunk` objects.
+    """
+    if not content.strip():
+        return []
+
+    if len(content) <= chunk_size:
+        return [
+            SplitChunk(
+                content=content,
+                start_line=line_offset + 1,
+                end_line=line_offset + content.count("\n") + 1,
+                language=language,
+                node_type=node_type,
+                symbol_name=symbol_name,
+            )
+        ]
+
+    chunks: list[SplitChunk] = []
+    start = 0
+    while start < len(content):
+        end = min(start + chunk_size, len(content))
+        if end < len(content):
+            newline_idx = content.rfind("\n", start, end)
+            if newline_idx > start:
+                end = newline_idx + 1
+        chunk_text = content[start:end]
+        if chunk_text.strip():
+            chunks.append(
+                SplitChunk(
+                    content=chunk_text,
+                    start_line=line_offset + content.count("\n", 0, start) + 1,
+                    end_line=line_offset + content.count("\n", 0, end) + 1,
+                    language=language,
+                    node_type=node_type,
+                    symbol_name=symbol_name,
+                )
+            )
+        start = max(start + 1, end - chunk_overlap)
+
+    return chunks
+
+
+@runtime_checkable
+class LanguageSplitter(Protocol):
+    """Interface for language-specific AST splitters.
+
+    Implementations receive raw source content and return semantically
+    meaningful chunks.  Return an empty list to signal a parse failure —
+    :class:`CodeSplitter` will then fall back to :func:`_char_split`.
+    """
+
+    def split(self, content: str, chunk_size: int, chunk_overlap: int) -> list[SplitChunk]:
+        """Split *content* into semantically meaningful chunks.
+
+        Args:
+            content: Raw source text.
+            chunk_size: Maximum characters per chunk.
+            chunk_overlap: Characters of overlap when char-splitting oversized nodes.
+
+        Returns:
+            Non-empty list of chunks, or ``[]`` to signal a parse failure.
+        """
+        ...
+
+
+class PythonASTSplitter:
+    """Splits Python source at function and class boundaries using :mod:`ast`.
+
+    Top-level :class:`ast.FunctionDef`, :class:`ast.AsyncFunctionDef`, and
+    :class:`ast.ClassDef` nodes each become their own chunk.  All other
+    top-level statements (imports, assignments, module docstrings) are
+    batched into a single ``"preamble"`` chunk that precedes the first
+    function or class.
+
+    Nodes that exceed *chunk_size* are character-split while preserving
+    their ``language`` and ``node_type`` metadata.
+    """
+
+    def split(self, content: str, chunk_size: int, chunk_overlap: int) -> list[SplitChunk]:
+        """Return AST-boundary chunks for *content*, or ``[]`` on :exc:`SyntaxError`.
+
+        Args:
+            content: Python source text.
+            chunk_size: Maximum characters per chunk.
+            chunk_overlap: Characters of overlap when char-splitting oversized nodes.
+
+        Returns:
+            List of chunks, or ``[]`` if the source cannot be parsed.
+        """
+        try:
+            tree = ast.parse(content)
+        except SyntaxError:
+            return []
+
+        lines = content.splitlines(keepends=True)
+        result: list[SplitChunk] = []
+        preamble: list[str] = []
+        preamble_start: int = 1
+
+        def flush_preamble() -> None:
+            nonlocal preamble, preamble_start
+            text = "".join(preamble)
+            if text.strip():
+                result.append(
+                    SplitChunk(
+                        content=text,
+                        start_line=preamble_start,
+                        end_line=preamble_start + len(preamble) - 1,
+                        language="python",
+                        node_type="preamble",
+                    )
+                )
+            preamble = []
+
+        def emit(text: str, start: int, end: int, node_type: str, symbol_name: str) -> None:
+            if len(text) <= chunk_size:
+                result.append(
+                    SplitChunk(
+                        content=text,
+                        start_line=start,
+                        end_line=end,
+                        language="python",
+                        node_type=node_type,
+                        symbol_name=symbol_name,
+                    )
+                )
+            else:
+                result.extend(
+                    _char_split(
+                        text,
+                        chunk_size,
+                        chunk_overlap,
+                        language="python",
+                        node_type=node_type,
+                        symbol_name=symbol_name,
+                        line_offset=start - 1,
+                    )
+                )
+
+        for node in ast.iter_child_nodes(tree):
+            if not isinstance(node, ast.stmt):
+                continue
+            start: int = node.lineno
+            end: int = node.end_lineno or node.lineno
+            node_text = "".join(lines[start - 1 : end])
+
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                flush_preamble()
+                emit(node_text, start, end, "function", node.name)
+            elif isinstance(node, ast.ClassDef):
+                flush_preamble()
+                emit(node_text, start, end, "class", node.name)
+            else:
+                if not preamble:
+                    preamble_start = start
+                preamble.extend(lines[start - 1 : end])
+
+        flush_preamble()
+        return result
+
+
+# ---------------------------------------------------------------------------
+# Extension registries — the plug-in surface for new languages
+# ---------------------------------------------------------------------------
+
+# Maps file extension → language name.
+# Add an entry here to teach CodeSplitter about a new file type.
+EXTENSION_MAP: dict[str, str] = {
+    ".py": "python",
+    # future: ".js": "javascript", ".ts": "typescript", ".go": "go", ...
+}
+
+# Maps language name → LanguageSplitter instance.
+# Register a new LanguageSplitter implementation here; files whose language
+# has no registered splitter fall back to character-based splitting automatically.
+LANGUAGE_SPLITTERS: dict[str, LanguageSplitter] = {
+    "python": PythonASTSplitter(),
+    # future: "javascript": TreeSitterSplitter("javascript"), ...
+}
+
+
+class CodeSplitter:
+    """Dispatcher that routes each file to the appropriate language splitter.
+
+    Uses :data:`EXTENSION_MAP` to detect the language from the file extension,
+    then delegates to the matching entry in :data:`LANGUAGE_SPLITTERS`.  If no
+    splitter is registered for the language (or the splitter signals a parse
+    failure by returning ``[]``), character-based line-boundary splitting is used.
+
+    Args:
+        chunk_size: Maximum characters per chunk.
+        chunk_overlap: Characters of overlap between consecutive chunks.
+    """
+
+    def __init__(self, chunk_size: int, chunk_overlap: int) -> None:
+        self._chunk_size = chunk_size
+        self._chunk_overlap = chunk_overlap
+
+    def split(self, content: str, file_path: str) -> list[SplitChunk]:
+        """Split *content* into chunks appropriate for the language of *file_path*.
+
+        Args:
+            content: Raw file content.
+            file_path: Repository-relative path; used only for extension detection.
+
+        Returns:
+            Non-empty list of :class:`SplitChunk` objects.
+        """
+        ext = Path(file_path).suffix.lower()
+        language = EXTENSION_MAP.get(ext)
+        if language:
+            splitter = LANGUAGE_SPLITTERS.get(language)
+            if splitter:
+                chunks = splitter.split(content, self._chunk_size, self._chunk_overlap)
+                if chunks:
+                    return chunks
+        return _char_split(
+            content,
+            self._chunk_size,
+            self._chunk_overlap,
+            language=language,
+            node_type=None,
+            symbol_name=None,
+        )

diffrag/review/__init__.py

@@ -0,0 +1,5 @@
+"""Review orchestration: diff → RAG → LLM → structured report."""
+
+from .reviewer import CodeReview, CodeReviewer, FileReview
+
+__all__ = ["CodeReview", "CodeReviewer", "FileReview"]