agents/agents/codebase_analyst/select_files.py

# steps/select_files.py
from __future__ import annotations

import json
import re
from dataclasses import dataclass

import pandas as pd
from pydantic import TypeAdapter, ValidationError

from contracts.coders.FileSelection import FileSelectionOutput, FileSelection
from contracts.coders.GatherOutput import GatherOutput

SelectionAdapter = TypeAdapter(FileSelectionOutput)


SYSTEM_PROMPT = """\
You select the files a coding agent should read to plan a change.

You will receive:
1. The issue (title, body, acceptance criteria).
2. A table of repository files with structural metadata.
3. Column definitions.

Your job: pick the files the agent needs, split into three roles:

- **primary**: files that almost certainly need edits for this task. If the \
  task says "add X to the user settings page", the settings page component \
  is primary. Be conservative — include only files you'd bet will be edited.
- **context**: files needed to understand primary files — imports they rely \
  on, type definitions they use, config they read. Include if a planner would \
  be confused without them.
- **test**: existing tests that cover the primary files. Include even if the \
  task doesn't mention tests; the agent needs to know what already exists.

Rules:
- Choose ONLY from the paths in the table. Do not invent paths.
- For each file, the reason MUST quote or paraphrase a specific identifier, \
  module name, or phrase from the issue. If you can't write such a reason, \
  do not select the file.
- Prefer fewer, higher-confidence selections over many low-confidence ones. \
  A bad selection wastes context; a missing file causes a bad plan.
- High `in` means the file is imported by many others — useful context, but \
  rarely the file to edit. Do not select a file just because `in` is high.
- Files in the same `community` as a primary file are often related; use \
  this as a signal, not a rule.
- `[test]`-prefixed files are tests. Use them to populate the test role.
- At most 30 selections total.

Output JSON matching the schema. No prose.
"""


@dataclass(frozen=True)
class SelectFilesInput:
    gather: GatherOutput
    table: str                  # pre-formatted, pre-filtered table text
    filtered_out_count: int     # for the notes


def select_files(
    inp: SelectFilesInput,
    *,
    llm,
) -> FileSelectionOutput:
    prompt = _build_user_prompt(inp)
    schema = SelectionAdapter.json_schema()

    raw = llm.chat_with_schema(
        prompt,
        {
            "type": "json_schema",
            "json_schema": {
                "name": "FileSelectionOutput",
                "schema": schema,
                "strict": True,
            },
        },
        system=SYSTEM_PROMPT,
    )

    raw = _ensure_dict(raw)
    try:
        out = SelectionAdapter.validate_python(raw)
    except ValidationError as e:
        raise SelectionError(str(e)) from e

    return _validate_against_table(out, inp.table)


# ── prompt ─────────────────────────────────────────────────────

def _build_user_prompt(inp: SelectFilesInput) -> str:
    g = inp.gather
    issue = g.issue_context

    parts: list[str] = []
    parts.append(f"# Issue {issue.issue_id}: {issue.title}")
    parts.append("")
    parts.append(issue.body.strip() or "(empty)")
    parts.append("")
    if issue.acceptance_criteria:
        parts.append("## Acceptance criteria")
        for i, ac in enumerate(issue.acceptance_criteria, 1):
            parts.append(f"{i}. {ac}")
        parts.append("")

    parts.append("## Column definitions")
    parts.append("- `path`: repo-relative path (authoritative)")
    parts.append("- `lang`: language")
    parts.append("- `in`: number of files that import this file (in-degree)")
    parts.append("- `out`: number of files this file imports (out-degree)")
    parts.append("- `pkg`: number of packages that depend on this file")
    parts.append("- `comm`: cluster id; files with the same id tend to be "
                 "related by construction")
    parts.append("")

    parts.append("## Files")
    parts.append("```")
    parts.append(inp.table)
    parts.append("```")
    parts.append("")

    if inp.filtered_out_count:
        parts.append(
            f"({inp.filtered_out_count} additional files were pre-filtered "
            f"out as clearly unrelated. Do not select them.)"
        )

    parts.append("Select the files. Return ONLY the JSON object.")
    return "\n".join(parts)


# ── validation ─────────────────────────────────────────────────

class SelectionError(ValueError):
    """LLM output violated the selection contract."""


def _validate_against_table(
    out: FileSelectionOutput,
    table: str,
) -> FileSelectionOutput:
    """
    Reject selections that reference paths not in the table.
    LLMs hallucinate paths surprisingly often, especially on large tables.
    """
    valid_paths = _paths_from_table(table)
    bad = [s.path for s in out.selections if s.path not in valid_paths]
    if bad:
        raise SelectionError(
            f"LLM selected paths not in the table: {bad[:5]}"
            + (f" (+{len(bad)-5} more)" if len(bad) > 5 else "")
        )
    # dedupe in case the model repeated a path with different roles
    seen: dict[str, FileSelection] = {}
    for s in out.selections:
        if s.path not in seen or s.confidence > seen[s.path].confidence:
            seen[s.path] = s
    deduped = list(seen.values())
    return FileSelectionOutput(selections=deduped, notes=out.notes)


def _paths_from_table(table: str) -> set[str]:
    paths = set()
    for line in table.splitlines()[1:]:        # skip header
        first = line.split(" | ", 1)[0].strip()
        if first.startswith("[test] "):
            first = first[len("[test] "):]
        paths.add(first)
    return paths


def _ensure_dict(raw):
    if isinstance(raw, str):
        return json.loads(raw)
    return raw

RENAMES = {
    "file": "path",
    "in_degree": "in",
    "out_degree": "out",
    "used_by_packages": "pkg",
    "community": "comm",
}

# Columns actually used by the LLM. Everything else is either redundant
# (isolated, leaf_no_deps, entry_no_dependents) or unused (label,
# used_by_files — a duplicate of in_degree).
KEEP = list(RENAMES.keys())


_TEST_PATH_RE = re.compile(
    r"(^|/)(tests?|__tests__|spec)(/|$)"      # tests/, test/, __tests__/, spec/
    r"|(^|/)test_[^/]+\.py$"                   # test_foo.py
    r"|_test\.(py|go)$"                        # foo_test.py, foo_test.go
    r"|\.(test|spec)\.(ts|tsx|js|jsx)$",       # foo.test.ts, foo.spec.tsx
    re.IGNORECASE,
)


def is_test_path(path: str) -> bool:
    return bool(_TEST_PATH_RE.search(path))


def _force_int(series: pd.Series) -> pd.Series:
    """
    Coerce any pandas dtype to a plain int64 numpy array wrapped in a Series.
    Handles Categorical (ordered or not), Int64 (nullable), float with NaN,
    object with strings, and boolean.
    """
    if isinstance(series.dtype, pd.CategoricalDtype):
        # cat.codes is always plain int8/int16/... — no ordering needed.
        return pd.Series(series.cat.codes, index=series.index, dtype="int64")

    # For everything else, let pandas convert, then force int64.
    return (
        pd.to_numeric(series, errors="coerce")
        .fillna(0)
        .astype("int64")
        .reset_index(drop=True)
        .set_axis(series.index)
    )


def format_table(df: pd.DataFrame) -> str:
    """
    Render the dependency-graph DataFrame as a compact, pipe-separated
    table suitable for an LLM prompt.

    - Keeps only the columns the model needs.
    - Renames them to short forms.
    - Sorts by (community, path) so related files cluster visually.
    - Prefixes test paths with [test] so the model doesn't have to guess.
    - Returns the header line + rows as a single string.
    """
    if df.empty:
        return "path | in | out | pkg | comm"

    cols = [c for c in KEEP if c in df.columns]
    view = df[cols].rename(columns=RENAMES).copy()

    # ── 1. Force every column to a plain Python type ──────────
    # This is the step that guarantees sort_values never sees a
    # Categorical, nullable Int64, or object-with-NaN.
    view["path"] = view["path"].astype(str)
    for col in ("in", "out", "pkg", "comm"):
        view[col] = _force_int(view[col])

    view = view.sort_values(["comm", "path"], kind="stable")

    view["path"] = view["path"].map(
        lambda p: f"[test] {p}" if is_test_path(p) else p
    )

    header = "path | lang | in | out | pkg | comm"
    view = view.astype(object).where(pd.notna(view), None)
    rows = [
        " | ".join("" if v is None else str(v) for v in row)
        for row in view.itertuples(index=False, name=None)
    ]
    return "\n".join([header, *rows])