open-swe/agent/review/groups.py

"""Diff grouping for the reviewer "AI sorted" view.

A lightweight, single structured-output LLM pass that partitions a PR's
changed files into ordered, logically-connected groups, each with a short
headline and explanation. The pass is kicked off concurrently with the
reviewer run (so it adds ~0 latency) and its result is stored on the reviewer
thread metadata under ``diff_groups`` for the review UI to render as a
top-to-bottom walkthrough of the PR.

Everything here is best-effort: any failure logs and degrades to "no groups",
and the UI falls back to the folder/flat view. The grouping never blocks or
breaks the review itself.
"""

from __future__ import annotations

import hashlib
import logging
from datetime import UTC, datetime
from typing import TypedDict

from langchain_core.language_models.chat_models import BaseChatModel
from pydantic import BaseModel, Field

from .diff import parse_unified_diff
from .findings import get_thread_metadata, set_reviewer_thread_metadata

logger = logging.getLogger(__name__)

# Bound the prompt so a huge PR can't blow up the cost/latency of the grouping
# pass. The full changed-file list is always included (so every file can be
# assigned); only the per-file hunk previews are budgeted.
MAX_PROMPT_CHARS = 48_000
MAX_FILE_HUNK_CHARS = 4_000
MAX_GROUPS = 8


class DiffGroup(TypedDict):
    """One logical group of changed files for the AI sorted view."""

    title: str
    summary: str
    files: list[str]


class _DiffGroupModel(BaseModel):
    title: str = Field(
        description=(
            "Short headline naming the logical change, roughly 4-10 words, no "
            "trailing punctuation. Wrap code identifiers in `backticks`."
        )
    )
    summary: str = Field(
        description=(
            "A short, plain explanation of what this group changes and why, in a "
            "few sentences a reviewer can skim. Wrap code identifiers, symbols, "
            "types, flags, and paths in `backticks`. No code blocks or links."
        )
    )
    files: list[str] = Field(
        description="Changed file paths in this group, copied verbatim, in reading order."
    )


class _DiffGroupingResult(BaseModel):
    groups: list[_DiffGroupModel] = Field(
        description="Ordered groups forming a top-to-bottom walkthrough of the PR."
    )


_PROMPT_TEMPLATE = """You are organizing a GitHub pull request's changed files into a clear, \
top-to-bottom walkthrough for a human reviewer.

Group the changed files by logical intent — put files that implement one \
coherent change together (e.g. "the new feature", "the supporting refactor", \
"the config wiring", "the tests"). Order the groups so a reviewer can read \
them top to bottom and build a mental model of the PR.

Rules:
- Assign every changed file to exactly one group.
- Use at most {max_groups} groups. Prefer fewer, larger groups over many tiny ones.
- title: a short headline naming the change, roughly 4-10 words, no trailing \
punctuation. Wrap code identifiers (symbols, flags, file names) in `backticks`.
- summary: a short, plain explanation of what the group changes and why, that a \
reviewer can skim:
    - Keep it focused — a few short sentences. Do not restate the diff line by line.
    - Wrap every code identifier, symbol, type, flag, and path in `backticks`.
    - Plain prose only — no code blocks and no links.
- files: the exact file paths (copied verbatim from the list below) in this \
group, in the order a reviewer should read them.

Changed files:
{file_list}

Diffs:
{diffs}
"""


def diff_signature(diff_text: str) -> str:
    """Stable content hash of the diff, used to skip regeneration when unchanged."""
    return hashlib.sha256(diff_text.encode("utf-8", "ignore")).hexdigest()


def _changed_files(diff_text: str) -> list[str]:
    return [fd.file for fd in parse_unified_diff(diff_text)]


def _build_prompt(diff_text: str, files: list[str]) -> str:
    parts: list[str] = []
    budget = MAX_PROMPT_CHARS
    for file_diff in parse_unified_diff(diff_text):
        segments: list[str] = []
        for hunk in file_diff.hunks:
            body = hunk.body
            if len(body) > MAX_FILE_HUNK_CHARS:
                body = body[:MAX_FILE_HUNK_CHARS] + "\n... (truncated)"
            segments.append(f"```diff\n{body}\n```")
        block = f"### {file_diff.file}\n" + "\n".join(segments) + "\n"
        if budget - len(block) < 0:
            parts.append(f"### {file_diff.file}\n(diff omitted — prompt budget reached)\n")
            continue
        budget -= len(block)
        parts.append(block)
    file_list = "\n".join(f"- {path}" for path in files)
    return _PROMPT_TEMPLATE.format(
        max_groups=MAX_GROUPS,
        file_list=file_list,
        diffs="\n".join(parts),
    )


def _normalize_groups(result: _DiffGroupingResult, files: list[str]) -> list[DiffGroup]:
    """Validate the model output into a clean partition.

    Each file is assigned at most once (first group wins), unknown paths are
    dropped, empty/untitled groups are dropped, and the group count is capped.
    Files the model failed to assign are intentionally left out here — the UI
    collects them into a trailing "Other changes" group, which also keeps the
    view robust against stale group data.
    """
    valid = set(files)
    seen: set[str] = set()
    groups: list[DiffGroup] = []
    for group in result.groups:
        title = (group.title or "").strip()
        summary = (group.summary or "").strip()
        picked: list[str] = []
        for path in group.files:
            if path in valid and path not in seen:
                seen.add(path)
                picked.append(path)
        if not title or not picked:
            continue
        groups.append({"title": title, "summary": summary, "files": picked})
        if len(groups) >= MAX_GROUPS:
            break
    return groups


async def generate_diff_groups(
    *,
    diff_text: str,
    model: BaseChatModel,
) -> list[DiffGroup] | None:
    """Run the single structured-output grouping pass over a unified diff.

    Returns the normalized groups, ``[]`` when there are no changed files, or
    ``None`` when the LLM call fails (caller treats ``None`` as "leave existing
    groups untouched").
    """
    files = _changed_files(diff_text)
    if not files:
        return []
    prompt = _build_prompt(diff_text, files)
    try:
        structured = model.with_structured_output(_DiffGroupingResult)
        result = await structured.ainvoke(prompt)
    except Exception:  # noqa: BLE001 — grouping is best-effort, never break the review
        logger.exception("Diff grouping LLM call failed")
        return None
    if not isinstance(result, _DiffGroupingResult):
        logger.warning("Diff grouping returned unexpected type: %s", type(result))
        return None
    return _normalize_groups(result, files)


async def maybe_generate_and_store_diff_groups(
    *,
    thread_id: str,
    head_sha: str,
    diff_text: str,
    model: BaseChatModel,
) -> None:
    """Generate diff groups and persist them on the reviewer thread metadata.

    Skips regeneration when the persisted signature already matches the current
    diff (the cheap no-op path for re-reviews with no file changes). All errors
    are logged and swallowed so this can run as a fire-and-forget background
    task without ever affecting the review.
    """
    try:
        if not thread_id or not diff_text:
            return
        signature = diff_signature(diff_text)
        metadata = await get_thread_metadata(thread_id)
        existing = metadata.get("diff_groups") if isinstance(metadata, dict) else None
        if (
            isinstance(existing, dict)
            and existing.get("signature") == signature
            and isinstance(existing.get("groups"), list)
            and existing.get("groups")
        ):
            return
        groups = await generate_diff_groups(diff_text=diff_text, model=model)
        if groups is None:
            return
        payload = {
            "head_sha": head_sha,
            "signature": signature,
            "generated_at": datetime.now(UTC).isoformat(),
            "groups": groups,
        }
        await set_reviewer_thread_metadata(thread_id, extra={"diff_groups": payload})
        logger.info(
            "Stored %d diff group(s) for reviewer thread %s (head %s)",
            len(groups),
            thread_id,
            head_sha[:7] if head_sha else "?",
        )
    except Exception:  # noqa: BLE001 — best-effort; never break the review
        logger.exception("Failed to generate/store diff groups for thread %s", thread_id)