"""LLM patch generation for SWE-bench benchmarking."""
import difflib
import re
from dataclasses import dataclass
import structlog
from litellm import completion
from benchmark.swebench.instance import SWEBenchInstance
from benchmark.swebench.prediction import Prediction
from benchmark.swebench.prompt import PromptBuilder
from benchmark.swebench.retriever import KoditRetriever, RetrievedSnippet
@dataclass
class SearchReplace:
"""A single search/replace operation."""
file_path: str
search: str
replace: str
class PatchGenerator:
"""Generates patches using an LLM for SWE-bench instances."""
def __init__(
self,
model: str,
max_tokens: int = 32768,
temperature: float = 0.0,
api_key: str | None = None,
timeout: int = 60,
) -> None:
"""Initialize generator with LLM settings."""
self._model = model
self._max_tokens = max_tokens
self._temperature = temperature
self._api_key = api_key
self._timeout = timeout
self._prompt_builder = PromptBuilder()
self._log = structlog.get_logger(__name__)
def generate_baseline(self, instance: SWEBenchInstance) -> Prediction:
"""Generate patch with only the problem statement (no retrieval)."""
self._log.info(
"Generating baseline prediction",
instance_id=instance.instance_id,
)
prompt = self._prompt_builder.baseline(instance)
patch = self._call_llm(prompt)
patch = self._extract_patch(patch)
return Prediction(
instance_id=instance.instance_id,
model_name_or_path=f"baseline-{self._model}",
model_patch=patch,
)
def generate_with_context(
self,
instance: SWEBenchInstance,
snippets: list[RetrievedSnippet],
) -> Prediction:
"""Generate patch with retrieved code context."""
self._log.info(
"Generating prediction with context",
instance_id=instance.instance_id,
snippet_count=len(snippets),
)
prompt = self._prompt_builder.with_context(instance, snippets)
patch = self._call_llm(prompt)
patch = self._extract_patch(patch)
return Prediction(
instance_id=instance.instance_id,
model_name_or_path=f"kodit-{self._model}",
model_patch=patch,
)
def _call_llm(self, prompt: str, retries: int = 3) -> str:
"""Call the LLM and return the response."""
from httpx import ReadTimeout
self._log.info(
"Calling LLM",
model=self._model,
timeout=self._timeout,
prompt_length=len(prompt),
)
last_error: Exception | None = None
for attempt in range(retries):
try:
response = completion(
model=self._model,
messages=[{"role": "user", "content": prompt}],
max_tokens=self._max_tokens,
temperature=self._temperature,
api_key=self._api_key,
timeout=float(self._timeout),
)
except ReadTimeout as e:
last_error = e
self._log.warning(
"LLM request timed out, retrying",
attempt=attempt + 1,
retries=retries,
error=str(e),
)
continue
choice = response.choices[0]
if choice.finish_reason == "length":
self._log.warning(
"Response truncated due to max_tokens limit",
model=self._model,
max_tokens=self._max_tokens,
)
return choice.message.content or ""
msg = f"LLM request failed after {retries} retries"
raise TimeoutError(msg) from last_error
def _extract_patch(self, response: str) -> str:
"""Extract and convert patch from the LLM response."""
# First try search/replace format
replacements = self._parse_search_replace(response)
if replacements:
self._log.info(
"Found search/replace blocks",
count=len(replacements),
)
return self._convert_to_diff(replacements)
# Fall back to diff extraction
code_block_pattern = r"```(?:diff)?\n(diff --git.*?)```"
match = re.search(code_block_pattern, response, re.DOTALL)
if match:
patch = match.group(1).strip()
patch = self._normalize_patch(patch)
self._validate_diff(patch)
return patch
diff_pattern = r"(diff --git.*)"
match = re.search(diff_pattern, response, re.DOTALL)
if match:
patch = match.group(1).strip()
patch = self._normalize_patch(patch)
self._validate_diff(patch)
return patch
self._log.warning(
"Could not extract patch from response",
response_preview=response[:200],
)
return response.strip()
def _parse_search_replace(self, response: str) -> list[SearchReplace]:
"""Parse search/replace blocks from response."""
results = []
current_file = None
# Find file markers
file_pattern = r"\[file:\s*([^\]]+)\]"
# Find search/replace blocks
block_pattern = (
r"<<<<<<+\s*SEARCH\s*\n(.*?)\n======+\s*\n(.*?)\n>>>>>>+\s*REPLACE"
)
# Split by file markers
parts = re.split(file_pattern, response)
for i, part in enumerate(parts):
# Odd indices are file paths from the split
if i % 2 == 1:
current_file = part.strip()
continue
# Even indices are content between file markers
if current_file is None:
# Try to find file path in the content itself
file_match = re.search(r"^\s*([^\s]+\.py)\s*$", part, re.MULTILINE)
if file_match:
current_file = file_match.group(1)
# Find all search/replace blocks in this part
for match in re.finditer(block_pattern, part, re.DOTALL):
search_text = match.group(1)
replace_text = match.group(2)
if current_file:
results.append(
SearchReplace(
file_path=current_file,
search=search_text,
replace=replace_text,
)
)
return results
def _convert_to_diff(self, replacements: list[SearchReplace]) -> str:
"""Convert search/replace blocks to unified diff format."""
diffs = []
# Group by file
by_file: dict[str, list[SearchReplace]] = {}
for r in replacements:
by_file.setdefault(r.file_path, []).append(r)
for file_path, file_replacements in by_file.items():
# Create a mock "original" content by concatenating search blocks
# and apply replacements to generate "new" content
# This is a simplified approach - we generate diff from search->replace
for replacement in file_replacements:
old_lines = replacement.search.splitlines(keepends=True)
new_lines = replacement.replace.splitlines(keepends=True)
# Ensure lines end with newline
if old_lines and not old_lines[-1].endswith("\n"):
old_lines[-1] += "\n"
if new_lines and not new_lines[-1].endswith("\n"):
new_lines[-1] += "\n"
# Generate unified diff
diff = list(
difflib.unified_diff(
old_lines,
new_lines,
fromfile=f"a/{file_path}",
tofile=f"b/{file_path}",
)
)
if diff:
# Replace the standard header with git-style header
diff_text = "".join(diff)
diff_text = f"diff --git a/{file_path} b/{file_path}\n{diff_text}"
diffs.append(diff_text)
result = "\n".join(diffs)
if result and not result.endswith("\n"):
result += "\n"
return result
def _normalize_patch(self, patch: str) -> str:
"""Normalize patch to handle common issues."""
# Split multiple diffs and keep only the first complete one
parts = re.split(r"(?=diff --git)", patch)
valid_parts = [p for p in parts if p.strip()]
if len(valid_parts) > 1:
self._log.warning(
"Multiple diffs found, using first valid one",
diff_count=len(valid_parts),
)
for part in valid_parts:
if "--- " in part and "+++ " in part and "@@ " in part:
patch = part.strip()
break
else:
patch = valid_parts[0].strip()
# Fix missing space prefixes on context lines within hunks
patch = self._fix_context_lines(patch)
# Ensure patch ends with newline
if patch and not patch.endswith("\n"):
patch += "\n"
return patch
def _fix_context_lines(self, patch: str) -> str:
"""Add missing space prefixes to context lines in hunks."""
lines = patch.split("\n")
result = []
in_hunk = False
for line in lines:
# Detect hunk start
if line.startswith("@@") and "@@" in line[2:]:
in_hunk = True
result.append(line)
continue
# Detect headers (not in hunk)
if line.startswith(("diff --git", "---", "+++")):
in_hunk = False
result.append(line)
continue
# Inside a hunk - ensure proper prefix
if in_hunk and line and not line.startswith((" ", "+", "-", "\\")):
# Missing prefix - add space for context line
result.append(" " + line)
else:
result.append(line)
return "\n".join(result)
def _validate_diff(self, patch: str) -> None:
"""Log warnings if the diff appears malformed."""
has_file_headers = "--- " in patch and "+++ " in patch
has_hunk_header = bool(re.search(r"@@ .+ @@", patch))
if not has_file_headers:
self._log.warning(
"Patch missing file headers (--- / +++)",
patch_preview=patch[:200],
)
if not has_hunk_header:
self._log.warning(
"Patch missing hunk header (@@ ... @@)",
patch_preview=patch[:200],
)
class BenchmarkRunner:
"""Runs benchmark predictions for SWE-bench instances."""
def __init__(
self,
generator: PatchGenerator,
retriever: KoditRetriever | None = None,
) -> None:
"""Initialize runner with generator and optional retriever."""
self._generator = generator
self._retriever = retriever
self._log = structlog.get_logger(__name__)
def run_baseline(self, instance: SWEBenchInstance) -> Prediction:
"""Run baseline prediction (no retrieval)."""
return self._generator.generate_baseline(instance)
def run_kodit(
self,
instance: SWEBenchInstance,
top_k: int = 10,
) -> Prediction:
"""Run Kodit-augmented prediction."""
if self._retriever is None:
msg = "Retriever required for Kodit condition"
raise ValueError(msg)
snippets = self._retriever.retrieve(instance, top_k=top_k)
return self._generator.generate_with_context(instance, snippets)
