diffrag 0.2.3__py3-none-any.whl

diffrag/review/prompts.py

@@ -0,0 +1,75 @@
+"""Prompt templates used by the review pipeline."""
+
+VERBOSITY_INSTRUCTIONS: dict[str, str] = {
+    "brief": (
+        "\nKeep your response under 200 words. "
+        "Report only the top 3 most critical issues — omit minor nits."
+    ),
+    "standard": "",
+    "detailed": (
+        "\nBe exhaustive. Include every issue found, including minor style and "
+        "documentation nits. For each issue provide a specific suggested fix."
+    ),
+}
+
+SUMMARIZE_DIFF = """\
+You are an expert software engineer. Summarize the following git diff in one concise paragraph.
+Focus on the logic changed, the specific functions or classes affected, and what problem is being
+solved. Write in plain English without using markdown.
+
+GIT DIFF:
+{diff}
+"""
+
+REVIEW_DIFF = """\
+You are an expert software engineer performing a thorough code review.
+Review the git diff below using the repository context provided for reference.
+
+Check for all of the following and report only findings that are actually present:
+- Breaking changes or API surface changes
+- Violations of SOLID principles or clean code practices
+- Code smells (duplicated logic, overly complex functions, magic values, etc.)
+- Inconsistencies with the surrounding codebase style or conventions
+- Outdated, missing, or incorrect documentation and comments
+- Security concerns (injection risks, credential exposure, unsafe deserialization, etc.)
+- Logic errors, off-by-one errors, or unhandled edge cases
+
+REPOSITORY CONTEXT (retrieved):
+{context}
+
+GIT DIFF:
+{diff}
+
+Structure your response with clear sections. Be concise and specific — reference file paths and
+line numbers where possible. If there are no issues in a category, omit that section entirely.
+{verbosity_instruction}"""
+
+CONSOLIDATE_REVIEW = """\
+You are an expert software engineer. The individual file reviews below were produced separately.
+Produce a single consolidated code review for the entire changeset.
+
+Focus on:
+1. The most important bugs or risks across all files.
+2. Cross-cutting issues (patterns repeated across multiple files).
+3. A brief overall assessment (approved / needs work / rejected).
+
+INDIVIDUAL FILE REVIEWS:
+{reviews}
+
+Write a concise, actionable consolidated review in markdown.{verbosity_instruction}
+"""
+
+ASK_QUESTION = """\
+You are an expert software engineer with knowledge of the repository and the current diff.
+Answer the user's question based on the repository context and diff summary provided.
+Be concise and accurate. Cite specific files or functions where relevant.
+
+REPOSITORY CONTEXT (retrieved):
+{context}
+
+CURRENT DIFF SUMMARY:
+{diff_summary}
+{conversation_history}
+QUESTION:
+{question}{verbosity_instruction}
+"""

diffrag/review/reviewer.py

@@ -0,0 +1,322 @@
+"""CodeReviewer: orchestrates diff extraction, RAG retrieval, and LLM review."""
+
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import tiktoken
+
+from ..ai.base import AIClient, ChatMessage
+from ..diff.git import GitRepository
+from ..diff.models import DiffResult, FileDiff
+from ..indexing.indexer import RepoIndexer
+from . import prompts
+
+log = logging.getLogger(__name__)
+
+_TOKENIZER = tiktoken.get_encoding("o200k_base")
+
+
+def _count_tokens(text: str) -> int:
+    return len(_TOKENIZER.encode(text))
+
+
+@dataclass
+class FileReview:
+    """Review result for a single file (or a hunk within a file).
+
+    Attributes:
+        identifier: File path, optionally with a hunk suffix like ``src/foo.py:hunk2``.
+        content: Markdown-formatted review text.
+        prompt_tokens: Total prompt tokens consumed for this review.
+        completion_tokens: Total completion tokens consumed for this review.
+    """
+
+    identifier: str
+    content: str
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+
+
+@dataclass
+class CodeReview:
+    """Complete code review result for a changeset.
+
+    Attributes:
+        base_ref: Base git ref the diff was computed from.
+        head_ref: Head git ref that was reviewed.
+        file_reviews: Per-file (or per-hunk) review objects.
+        summary: Final consolidated review in markdown.
+        total_prompt_tokens: Aggregate prompt tokens across all LLM calls.
+        total_completion_tokens: Aggregate completion tokens across all LLM calls.
+    """
+
+    base_ref: str
+    head_ref: str
+    file_reviews: list[FileReview]
+    summary: str
+    total_prompt_tokens: int = field(default=0, init=False)
+    total_completion_tokens: int = field(default=0, init=False)
+
+    def __post_init__(self) -> None:
+        self.total_prompt_tokens = sum(r.prompt_tokens for r in self.file_reviews)
+        self.total_completion_tokens = sum(r.completion_tokens for r in self.file_reviews)
+
+
+class CodeReviewer:
+    """Orchestrates the full code review pipeline.
+
+    For each changed file:
+    1. The diff is summarised by the LLM to create a targeted RAG query.
+    2. The RAG index is queried for relevant context chunks.
+    3. The LLM reviews the diff in light of the retrieved context.
+
+    Large diffs are automatically split into per-hunk reviews which are
+    then consolidated into a single final summary.
+
+    Args:
+        ai_client: Language model client for completion requests.
+        indexer: Vector index for context retrieval.
+        similarity_top_k: Number of context chunks retrieved per diff chunk.
+        max_prompt_tokens: Approximate token budget per diff before hunk-splitting.
+    """
+
+    def __init__(
+        self,
+        ai_client: AIClient,
+        indexer: RepoIndexer,
+        similarity_top_k: int = 5,
+        max_prompt_tokens: int = 6000,
+        verbosity: str = "standard",
+    ) -> None:
+        self._ai = ai_client
+        self._indexer = indexer
+        self._top_k = similarity_top_k
+        self._max_prompt_tokens = max_prompt_tokens
+        self._verbosity = verbosity
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    async def review(
+        self,
+        repo: GitRepository,
+        base: str,
+        head: str = "HEAD",
+        *,
+        rebuild_index: bool = True,
+        on_progress: Callable[[str], None] | None = None,
+    ) -> CodeReview:
+        """Run a complete code review for changes between *base* and *head*.
+
+        Args:
+            repo: Repository to analyse.
+            base: Base git ref (branch, tag, or commit).
+            head: Head git ref. Defaults to ``"HEAD"``.
+            rebuild_index: When ``True`` the vector index is rebuilt before
+                reviewing.  Pass ``False`` to reuse an existing 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.
+
+        Returns:
+            A :class:`CodeReview` with per-file results and a consolidated summary.
+        """
+        _p = on_progress or (lambda _: None)
+
+        _p("Fetching diff...")
+        diff_result = repo.get_diff(base, head)
+
+        if diff_result.is_empty:
+            log.info("No changes between %s and %s", base, head)
+            _p("No changes detected")
+            return CodeReview(
+                base_ref=base,
+                head_ref=head,
+                file_reviews=[],
+                summary="No changes detected between the two refs.",
+            )
+
+        n = len(diff_result.file_diffs)
+        _p(f"Diff: {n} file(s) changed")
+
+        commit_hash = repo.head_commit()
+
+        if rebuild_index:
+            if self._indexer.is_stale(repo.path, commit_hash):
+                await self._indexer.build_index(
+                    repo.path, on_progress=on_progress, commit_hash=commit_hash
+                )
+            else:
+                log.info("Index is up to date (commit %s)", commit_hash[:8])
+                _p("Index is up to date — skipping rebuild")
+        elif self._indexer.document_count == 0:
+            log.warning("Index is empty and rebuild_index=False — building index anyway")
+            await self._indexer.build_index(
+                repo.path, on_progress=on_progress, commit_hash=commit_hash
+            )
+
+        file_reviews: list[FileReview] = []
+        for i, file_diff in enumerate(diff_result.file_diffs, 1):
+            _p(f"Reviewing {file_diff.path} ({i}/{n})...")
+            log.info("Reviewing %s (%s)", file_diff.path, file_diff.status)
+            reviews = await self._review_file(file_diff)
+            file_reviews.extend(reviews)
+            _p(f"Reviewed {file_diff.path} ({i}/{n})")
+
+        if n > 1:
+            _p("Consolidating review...")
+        summary = await self._consolidate(file_reviews)
+        return CodeReview(base_ref=base, head_ref=head, file_reviews=file_reviews, summary=summary)
+
+    async def ask(
+        self,
+        question: str,
+        diff_result: DiffResult,
+        history: list[tuple[str, str]] | None = None,
+    ) -> str:
+        """Answer a free-form question about the diff and codebase.
+
+        Args:
+            question: The user's question in plain English.
+            diff_result: Pre-computed diff to use as context.
+            history: Prior Q&A pairs from this session, oldest first.  When
+                provided, they are injected into the prompt so the model can
+                refer to earlier exchanges.
+
+        Returns:
+            A markdown-formatted answer string.
+        """
+        diff_text = "\n\n".join(f.raw_diff for f in diff_result.file_diffs)
+        query = f"{question}\n{diff_text[:500]}"
+
+        context_chunks = await self._indexer.query(query, top_k=self._top_k)
+        context = "\n\n".join(c.content for c in context_chunks)
+
+        if history:
+            lines = ["PREVIOUS EXCHANGES:"]
+            for q, a in history:
+                lines += [f"User: {q}", f"Assistant: {a}", ""]
+            history_block = "\n".join(lines) + "\n"
+        else:
+            history_block = ""
+
+        vi = prompts.VERBOSITY_INSTRUCTIONS.get(self._verbosity, "")
+        prompt_text = prompts.ASK_QUESTION.format(
+            context=context,
+            diff_summary=diff_text[:1500],
+            conversation_history=history_block,
+            question=question,
+            verbosity_instruction=vi,
+        )
+        response = await self._ai.complete([ChatMessage(role="user", content=prompt_text)])
+        return response.content
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    async def _review_file(self, file_diff: FileDiff) -> list[FileReview]:
+        """Review a single file, splitting into hunks if the diff is large.
+
+        Args:
+            file_diff: Parsed file diff to review.
+
+        Returns:
+            One or more :class:`FileReview` objects.
+        """
+        token_count = _count_tokens(file_diff.raw_diff)
+        if token_count <= self._max_prompt_tokens:
+            return [await self._review_chunk(file_diff.path, file_diff.raw_diff)]
+
+        log.info(
+            "Diff for %s is ~%d tokens — splitting into %d hunks",
+            file_diff.path,
+            token_count,
+            len(file_diff.hunks),
+        )
+        reviews: list[FileReview] = []
+        for idx, hunk in enumerate(file_diff.hunks):
+            hunk_text = f"{hunk.header}\n{hunk.content}"
+            if _count_tokens(hunk_text) > self._max_prompt_tokens:
+                log.warning(
+                    "Hunk %d of %s exceeds token budget — skipping",
+                    idx,
+                    file_diff.path,
+                )
+                continue
+            review = await self._review_chunk(f"{file_diff.path}:hunk{idx}", hunk_text)
+            reviews.append(review)
+        return reviews
+
+    async def _review_chunk(self, identifier: str, diff_text: str) -> FileReview:
+        """Review a single diff chunk (file or hunk) using RAG context.
+
+        Args:
+            identifier: Human-readable label for logging and reporting.
+            diff_text: Raw diff text to review.
+
+        Returns:
+            A :class:`FileReview` for this chunk.
+        """
+        # Step 1: summarise the diff to form a targeted query
+        summary_resp = await self._ai.complete(
+            [ChatMessage(role="user", content=prompts.SUMMARIZE_DIFF.format(diff=diff_text))]
+        )
+        log.debug("Diff summary for %s: %s", identifier, summary_resp.content[:120])
+
+        # Step 2: retrieve relevant context from the index
+        context_chunks = await self._indexer.query(summary_resp.content, top_k=self._top_k)
+        context = "\n\n".join(
+            f"# {c.file_path}"
+            + (f" — {c.node_type} `{c.symbol_name}`" if c.symbol_name else "")
+            + f"\n{c.content}"
+            for c in context_chunks
+        )
+
+        # Step 3: full review with retrieved context
+        vi = prompts.VERBOSITY_INSTRUCTIONS.get(self._verbosity, "")
+        review_resp = await self._ai.complete(
+            [
+                ChatMessage(
+                    role="user",
+                    content=prompts.REVIEW_DIFF.format(
+                        context=context, diff=diff_text, verbosity_instruction=vi
+                    ),
+                )
+            ]
+        )
+
+        return FileReview(
+            identifier=identifier,
+            content=review_resp.content,
+            prompt_tokens=summary_resp.prompt_tokens + review_resp.prompt_tokens,
+            completion_tokens=summary_resp.completion_tokens + review_resp.completion_tokens,
+        )
+
+    async def _consolidate(self, file_reviews: list[FileReview]) -> str:
+        """Consolidate multiple file reviews into a single summary.
+
+        Args:
+            file_reviews: All per-file or per-hunk reviews produced so far.
+
+        Returns:
+            Consolidated markdown review string.
+        """
+        if len(file_reviews) == 1:
+            return file_reviews[0].content
+
+        reviews_text = "\n\n".join(f"### {r.identifier}\n{r.content}" for r in file_reviews)
+        vi = prompts.VERBOSITY_INSTRUCTIONS.get(self._verbosity, "")
+        response = await self._ai.complete(
+            [
+                ChatMessage(
+                    role="user",
+                    content=prompts.CONSOLIDATE_REVIEW.format(
+                        reviews=reviews_text, verbosity_instruction=vi
+                    ),
+                )
+            ]
+        )
+        return response.content