Plan step

This commit is contained in:
2026-09-23 22:42:21 +03:00
parent 1f2dc25e0b
commit 762ac2f14b
15 changed files with 823 additions and 283 deletions
+17 -40
View File
@@ -25,11 +25,11 @@ import os
from pathlib import Path
from typing import List, Literal, Optional, Type, TypeVar
from pydantic import BaseModel, Field, ValidationError, TypeAdapter
from pydantic import BaseModel, Field, ValidationError, TypeAdapter, field_validator
from agents.coders.repo import RepoContext
from common.llm_client import LLMClient
from contracts.RepoContext import Language, LanguageLiteral
from common.llm_client import LLMClient, ResponseCache
from contracts.RepoContext import RepoContext
from contracts.CodeLanguages import Language, LanguageLiteral
from contracts.coders.VerifyOutput import sanitize_for_openai
log = logging.getLogger("entrypoint_analyzer")
@@ -86,10 +86,21 @@ class Attempt(BaseModel):
class AnalysisResult(BaseModel):
language: Language
entrypoint: str
entrypoint: str = Field(
...,
description="Relative path from repo_path to the entrypoint file",
)
attempts: List[Attempt]
reasoning: str
@field_validator("entrypoint")
@classmethod
def entrypoint_must_exist(cls, v: str, info) -> str:
repo = info.data.get("repo_path")
if repo and not (repo / v).is_file():
raise ValueError(f"Entrypoint not found: {repo / v}")
return v
# --------------------------------------------------------------------------- #
# LLM plumbing: LLMClient wrapper with strict schema + cache + retries #
@@ -97,41 +108,6 @@ class AnalysisResult(BaseModel):
T = TypeVar("T", bound=BaseModel)
class ResponseCache:
"""File-backed cache: prompt-hash -> validated-model JSON. Makes runs idempotent."""
def __init__(self, path: Optional[Path] = None) -> None:
self.path = path
self._mem: dict[str, str] = {}
if path and path.exists():
try:
self._mem = json.loads(path.read_text("utf-8"))
except Exception: # noqa: BLE001
log.warning("Corrupt cache at %s; starting fresh", path)
@staticmethod
def key(system: str, user: str, schema: Type[BaseModel]) -> str:
h = hashlib.sha256()
h.update(schema.__name__.encode())
h.update(b"\x00")
h.update(system.encode())
h.update(b"\x00")
h.update(user.encode())
return h.hexdigest()
def get(self, k: str) -> Optional[str]:
return self._mem.get(k)
def put(self, k: str, v: str) -> None:
self._mem[k] = v
if self.path:
self.path.parent.mkdir(parents=True, exist_ok=True)
tmp = self.path.with_suffix(self.path.suffix + ".tmp")
tmp.write_text(json.dumps(self._mem, indent=2, sort_keys=True), "utf-8")
tmp.replace(self.path)
class StructuredLLM:
"""
Wrap the project's `LLMClient` with:
@@ -404,6 +380,7 @@ class EntrypointAnalyzer:
def gather_repo_entrypoint(llm: LLMClient, repo: RepoContext, cache_path: Optional[str] = None) -> AnalysisResult:
log.debug("Gathering repo entrypoint")
return EntrypointAnalyzer(
StructuredLLM(llm, cache=ResponseCache(Path(cache_path)) if cache_path else None)
).analyze(Path(repo.repo_dir))
+122 -66
View File
@@ -1,10 +1,13 @@
import json
import logging
import os
import pickle
import shutil
import subprocess
from abc import ABC, abstractmethod
from functools import partial
from pathlib import Path
from typing import Optional, Callable
from typing import Optional, Callable, TypedDict, Generic, Hashable, TypeVar
import networkx as nx
import pandas as pd
@@ -12,9 +15,16 @@ from networkx.algorithms.community import louvain_communities
from pydantic import BaseModel, Field, ValidationError
from agents.codebase_analyst.entrypoint import gather_repo_entrypoint
from common.llm_client import LLMClient
from contracts.RepoContext import Language, ToolRequirement, DependencyGraph, DependencyNode, DependencyEdge, \
AnalysisResult, AnalysisInput, AnalysisError
from contracts.GraphAnalyzer import AnalyzeResult, DependencyEdge, DependencyNode
from contracts.RepoContext import Language, ToolRequirement, DependencyGraph, \
AnalysisResult, AnalysisInput, AnalysisError, RepoContext
from contracts.coders.RunContext import RunContext
logger = logging.getLogger("graph_tool")
logger.setLevel(logging.DEBUG)
def find_tool(command: str) -> Optional[str]:
@@ -452,7 +462,8 @@ class JavaAdapter(LanguageAdapter):
raise RuntimeError(f"gradle dependencies failed: {result.stderr.strip()}")
return self._parse_tree_text(result.stdout, entrypoint)
def _parse_tree_text(self, text: str, entrypoint: str) -> DependencyGraph:
@staticmethod
def _parse_tree_text(text: str, entrypoint: str) -> DependencyGraph:
# Simplified parser for the textual tree format.
# In production, use the JSON output of the Maven/Gradle plugin.
nodes: list[DependencyNode] = []
@@ -636,7 +647,30 @@ ADAPTER_REGISTRY: dict[Language, AdapterFactory] = {
}
def analyze_repository(inp: AnalysisInput) -> AnalysisResult:
def analyze_repo_with_cache(run_ctx: RunContext, repo: RepoContext) -> AnalysisResult:
repo_analyze = None
analysis_cache_path = run_ctx.config.project_mem_dir() / "repo_analysis.json"
if os.path.exists(analysis_cache_path):
try:
repo_analyze = pickle.load(open(analysis_cache_path, "rb"))
logger.debug("repo analysis from cache")
except Exception as e:
print(e)
pass
if not repo_analyze:
logger.debug("generating repo analysis")
repo_analyze = analyze_repository(AnalysisInput(
repo=repo
), run_ctx)
if repo_analyze.error:
raise BaseException(repo_analyze.error)
with open(analysis_cache_path, "wb") as f:
pickle.dump(repo_analyze, f)
return repo_analyze
def analyze_repository(inp: AnalysisInput, run_ctx: RunContext) -> AnalysisResult:
"""
Idempotent entry point.
@@ -644,15 +678,19 @@ def analyze_repository(inp: AnalysisInput) -> AnalysisResult:
(either a graph or a structured error). No files are written to the
repository under analysis.
"""
adapter_factory = ADAPTER_REGISTRY.get(inp.language)
entrypoint_analyze = gather_repo_entrypoint(run_ctx.llm, inp.repo, run_ctx.config.project_mem_dir() / "entrypoint.json")
language = entrypoint_analyze.language
entrypoint = entrypoint_analyze.entrypoint
adapter_factory = ADAPTER_REGISTRY.get(language)
if adapter_factory is None:
return AnalysisResult(
success=False,
error=AnalysisError(
error_code="UNSUPPORTED_LANGUAGE",
message=f"No adapter registered for language '{inp.language}'",
language=inp.language,
entrypoint=inp.entrypoint,
message=f"No adapter registered for language '{language}'",
language=language,
entrypoint=entrypoint,
),
)
@@ -670,27 +708,27 @@ def analyze_repository(inp: AnalysisInput) -> AnalysisResult:
error=AnalysisError(
error_code="MISSING_TOOL",
message=(
f"Cannot build dependency graph for '{inp.language.value}': "
f"Cannot build dependency graph for '{language.value}': "
f"required tool(s) not found on PATH: "
f"{', '.join(r.tool_name for r in missing)}"
),
missing_tools=missing,
language=inp.language,
entrypoint=inp.entrypoint,
language=language,
entrypoint=entrypoint,
),
)
# --- Build graph ---
try:
graph = adapter.build_graph(inp.repo_path, inp.entrypoint)
graph = adapter.build_graph(inp.repo.repo_dir, entrypoint)
except subprocess.TimeoutExpired:
return AnalysisResult(
success=False,
error=AnalysisError(
error_code="TOOL_TIMEOUT",
message=f"Tool timed out while analyzing {inp.entrypoint}",
language=inp.language,
entrypoint=inp.entrypoint,
message=f"Tool timed out while analyzing {entrypoint}",
language=language,
entrypoint=entrypoint,
),
)
except Exception as exc:
@@ -699,39 +737,45 @@ def analyze_repository(inp: AnalysisInput) -> AnalysisResult:
error=AnalysisError(
error_code="ANALYSIS_FAILED",
message=str(exc),
language=inp.language,
entrypoint=inp.entrypoint,
language=language,
entrypoint=entrypoint,
),
)
G_all = to_networkx(graph, include_external=True)
G = to_networkx(graph, include_external=False) # internal-only analysis
res_internal = analyze(G, graph, label="_internal")
res_all = analyze(G_all, graph, label="_all")
leaf_clusters = cluster_leaves(G, res_internal["leaf_no_deps"])
leaf_clusters = cluster_leaves(G, res_internal.leaf_no_deps)
res_all = analyze(to_networkx(graph, include_external=True), graph, label="_all")
return AnalysisResult(success=True, graph=graph, res_all=res_all, leaf_clusters=leaf_clusters)
def to_networkx(dep_graph: DependencyGraph, include_external=False):
def to_networkx(dep_graph: "DependencyGraph", include_external: bool = False) -> nx.DiGraph[DependencyNode]:
"""
Convert a DependencyGraph into a networkx DiGraph.
Nodes are identified by their string id. Edges point from the source
(the file that declares the dependency) to the target (the file that
is depended upon).
If include_external is False, external nodes (and edges touching them)
are omitted.
"""
G = nx.DiGraph()
# Nodes
id_to_node: dict[str, DependencyNode] = {}
for n in dep_graph.nodes:
if not include_external and n.is_external:
continue
G.add_node(n.id, label=n.label, language=n.language.value, is_external=n.is_external)
G.add_node(n) # DependencyNode must be hashable
id_to_node[n.id] = n
# Edges (source depends on target)
for e in dep_graph.edges:
if e.source in G and e.target in G:
G.add_edge(e.source, e.target, kind=e.kind)
# Add isolated internal files that never appear in edges
if not include_external:
for n in dep_graph.nodes:
if not n.is_external and n.id not in G:
G.add_node(n.id, label=n.label, language=n.language.value, is_external=False)
src = id_to_node.get(e.source)
tgt = id_to_node.get(e.target)
if src is not None and tgt is not None:
G.add_edge(src, tgt, kind=e.kind)
return G
@@ -741,54 +785,66 @@ def package_of(path: str) -> str:
COLS = [
"file", "language", "in_degree", "out_degree",
"file", "in_degree", "out_degree",
"isolated", "leaf_no_deps", "entry_no_dependents",
"used_by_files", "used_by_packages", "community",
]
def analyze(G, graph, label=""):
in_deg = dict(G.in_degree())
out_deg = dict(G.out_degree())
def analyze(
G: nx.DiGraph[DependencyNode],
graph: DependencyGraph, # NOTE: currently unused in the body
label: str = "",
) -> AnalyzeResult:
in_deg: dict[DependencyNode, int] = dict(G.in_degree())
out_deg: dict[DependencyNode, int] = dict(G.out_degree())
# 1. Truly isolated: no in and no out edges
isolated = [n for n in G if G.degree(n) == 0]
isolated: list[DependencyNode] = [n for n in G if G.degree(n) == 0]
# 2. Leaf: depends on nothing
leaf_no_deps = [n for n in G if G.out_degree(n) == 0]
leaf_no_deps: list[DependencyNode] = [n for n in G if G.out_degree(n) == 0]
# 3. Entrypoint: nothing depends on it
entry_no_dependents = [n for n in G if G.in_degree(n) == 0]
entry_no_dependents: list[DependencyNode] = [n for n in G if G.in_degree(n) == 0]
# 4. Widely used: how many distinct files depend on it
widely_used = sorted(in_deg.items(), key=lambda kv: kv[1], reverse=True)
widely_used: list[tuple[DependencyNode, int]] = sorted(
in_deg.items(), key=lambda kv: kv[1], reverse=True
)
# 5. How many distinct packages depend on it
pkg = {n: package_of(n) for n in G}
used_by_packages = {
pkg: dict[DependencyNode, Hashable] = {n: package_of(n.id) for n in G}
used_by_packages: dict[DependencyNode, int] = {
n: len({pkg[p] for p in G.predecessors(n)}) for n in G
}
# 6. Weakly connected components = isolated groups
wccs = sorted(nx.weakly_connected_components(G), key=len, reverse=True)
wccs: list[set[DependencyNode]] = sorted(
nx.weakly_connected_components(G), key=len, reverse=True
)
# 7. Strongly connected components = cycles / tightly coupled groups
sccs = sorted(
sccs: list[set[DependencyNode]] = sorted(
(c for c in nx.strongly_connected_components(G) if len(c) > 1),
key=len, reverse=True,
key=len,
reverse=True,
)
# 8. Communities via Louvain
U = G.to_undirected()
comms = louvain_communities(U, seed=42) if U.number_of_edges() else []
node2comm = {n: i for i, c in enumerate(comms) for n in c}
U: nx.Graph[DependencyNode] = G.to_undirected()
comms: list[set[DependencyNode]] = (
louvain_communities(U, seed=42) if U.number_of_edges() else []
)
node2comm: dict[DependencyNode, int] = {
n: i for i, c in enumerate(comms) for n in c
}
# Per-file table
rows = []
rows: list[dict[str, object]] = []
for n, data in G.nodes(data=True):
rows.append({
"file": n,
"language": data.get("language"),
"file": n.id,
"in_degree": in_deg[n],
"out_degree": out_deg[n],
"isolated": G.degree(n) == 0,
@@ -808,27 +864,27 @@ def analyze(G, graph, label=""):
df.to_csv(f"file_metrics{label}.csv", index=False)
pd.DataFrame([
{"group_id": i, "size": len(c), "members": ";".join(sorted(c))}
{"group_id": i, "size": len(c), "members": ";".join(sorted(n.id for n in c))}
for i, c in enumerate(wccs)
]).to_csv(f"isolated_groups{label}.csv", index=False)
pd.DataFrame([
{"scc_id": i, "size": len(c), "members": ";".join(sorted(c))}
{"scc_id": i, "size": len(c), "members": ";".join(sorted(n.id for n in c))}
for i, c in enumerate(sccs)
]).to_csv(f"cycles{label}.csv", index=False)
return {
"df": df,
"isolated": isolated,
"leaf_no_deps": leaf_no_deps,
"entry_no_dependents": entry_no_dependents,
"widely_used": widely_used,
"used_by_packages": used_by_packages,
"wccs": wccs,
"sccs": sccs,
"node2comm": node2comm,
"comms": comms,
}
return AnalyzeResult(
df=df,
isolated=isolated,
leaf_no_deps=leaf_no_deps,
entry_no_dependents=entry_no_dependents,
widely_used=widely_used,
used_by_packages=used_by_packages,
wccs=wccs,
sccs=sccs,
node2comm=node2comm,
comms=comms,
)
def cluster_leaves(G, leaves):
+262
View File
@@ -0,0 +1,262 @@
# 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])