Plan step
This commit is contained in:
@@ -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))
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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])
|
||||
Reference in New Issue
Block a user