Spaces:
Sleeping
Sleeping
File size: 16,650 Bytes
3fe77a3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 | """Coder agent for generating code fixes."""
import logging
from pathlib import Path
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import ChatPromptTemplate
from tools.git_tool import GitTool
from models import AgentState, CodeChange, ReasoningStep
import time
logger = logging.getLogger(__name__)
class CoderAgent:
"""Agent that generates code fixes."""
SYSTEM_PROMPT = """You are an expert software engineer writing code fixes for GitHub issues.
Your task:
1. Analyze the issue and the relevant code files
2. Generate precise code changes to fix the issue
3. Ensure changes are minimal and focused
4. Maintain code style and best practices
CRITICAL RULES:
- STRICT SCOPE CONTROL: Do exactly what the issue asks and NOTHING MORE. Do not refactor, reorganize, or rewrite unrelated code or documentation.
- FORMATTING MATCH: If adding new entries, examples, or code, you MUST rigidly match the format, structure, and syntax of the surrounding file.
- Only modify files that directly fix the issue.
- Keep changes under 300 lines total.
- Preserve existing code style and ensure backwards compatibility.
- VERY IMPORTANT: The code in the ORIGINAL block MUST EXACTLY match the existing code in the file, including all whitespace and indentation!
- NEVER truncate the NEW block. Write out all the code you want to replace the ORIGINAL block with.
- For Markdown (.md) or configuration files, if you want to rewrite the entire file, leave the ORIGINAL block completely empty and put the full new content in the NEW block.
Output format for EACH file to modify:
FILE: <file_path>
EXPLANATION: <one line explaining the change>
ORIGINAL:
<exact original code to be replaced>
NEW:
<new code to replace it>
---
If no changes needed, output: NO_CHANGES_NEEDED
"""
def __init__(self, api_key: str, git_tool: GitTool, model_name: str = "gemini-1.5-flash"):
"""
Initialize coder agent.
Args:
api_key: Gemini, OpenAI, Anthropic, or Groq API key
git_tool: Git tool instance
model_name: Model to use
"""
import os
# 1. Groq Support (Key starts with gsk_)
if api_key and api_key.startswith("gsk_"):
from langchain_groq import ChatGroq
actual_model = "llama-3.3-70b-specdec" if any(x in model_name.lower() for x in ["gemini", "gpt", "claude"]) else model_name
self.llm = ChatGroq(
model=actual_model,
groq_api_key=api_key,
temperature=0.2,
)
# 2. Anthropic (Claude) Support (Key starts with sk-ant-)
elif api_key and api_key.startswith("sk-ant-"):
from langchain_anthropic import ChatAnthropic
actual_model = "claude-3-5-sonnet-latest" if any(x in model_name.lower() for x in ["gemini", "gpt", "llama"]) else model_name
self.llm = ChatAnthropic(
model=actual_model,
api_key=api_key,
temperature=0.2,
)
# 3. OpenRouter Support (Key starts with sk-or-)
elif api_key and api_key.startswith("sk-or-"):
from langchain_openai import ChatOpenAI
base_url = "https://openrouter.ai/api/v1"
actual_model = model_name
if "gemini-2.5-flash" in model_name:
actual_model = "google/gemini-2.5-flash"
elif "gemini-1.5-flash" in model_name:
actual_model = "google/gemini-1.5-flash"
elif "gemini-2.5-pro" in model_name:
actual_model = "google/gemini-2.5-pro"
elif "gemini-1.5-pro" in model_name:
actual_model = "google/gemini-1.5-pro"
elif "/" not in model_name:
actual_model = "google/gemini-2.5-flash"
self.llm = ChatOpenAI(
model=actual_model,
api_key=api_key,
base_url=base_url,
temperature=0.2,
max_tokens=4096,
)
# 3.5. OpenAI / Custom Endpoint Support (Key starts with sk-)
elif api_key and api_key.startswith("sk-"):
from langchain_openai import ChatOpenAI
base_url = os.environ.get("OPENAI_BASE_URL") or os.environ.get("OPENAI_API_BASE") or None
actual_model = "gpt-4o-mini" if any(x in model_name.lower() for x in ["gemini", "llama", "claude"]) else model_name
self.llm = ChatOpenAI(
model=actual_model,
api_key=api_key,
base_url=base_url,
temperature=0.2,
)
# 4. Standard Gemini Fallback
else:
self.llm = ChatGoogleGenerativeAI(
model=model_name,
google_api_key=api_key,
temperature=0.2,
)
self.git_tool = git_tool
self.prompt = ChatPromptTemplate.from_messages([
("system", self.SYSTEM_PROMPT),
("user", """Repository: {repo_name}
Issue #{issue_number}: {issue_title}
Issue Description:
{issue_body}
Strategy:
{strategy}
Relevant Files:
{files_content}
Generate code changes to fix this issue."""),
])
self.chain = self.prompt | self.llm
logger.info(f"Coder agent initialized with model: {model_name}")
def code(self, state: AgentState) -> AgentState:
"""
Generate code changes.
Args:
state: Current agent state
Returns:
Updated state with code changes
"""
logger.info("Generating code changes")
if not state.relevant_files:
logger.error("No relevant files available")
state.error = "No relevant files from searcher"
return state
try:
files_content = self._prepare_files_content(state)
response = self.chain.invoke({
"repo_name": f"{state.owner}/{state.repo}",
"issue_number": state.issue_number,
"issue_title": state.issue.title,
"issue_body": state.issue.body,
"strategy": state.strategy,
"files_content": files_content,
})
content = response.content
if "NO_CHANGES_NEEDED" in content:
logger.info("Agent determined no changes are needed")
state.changes = []
state.reasoning_trace.append(ReasoningStep(
agent="Coder",
message="Agent determined no changes are needed.",
timestamp=time.time()
))
return state
changes = self._parse_changes(content)
if not changes:
logger.error("Failed to parse any changes")
state.error = "Failed to parse code changes from LLM response"
return state
branch_name = f"ai-fix/issue-{state.issue_number}"
state.reasoning_trace.append(ReasoningStep(
agent="Coder",
message=f"Generated {len(changes)} code changes.\n" + "\n".join([f"- Modified {c.file_path}" for c in changes]),
timestamp=time.time()
))
self.git_tool.create_branch(Path(state.repo_path), branch_name)
for change in changes:
file_content = self.git_tool.get_file_content(Path(state.repo_path), change.file_path)
if file_content is not None:
orig = change.original_content
new_val = change.new_content
if orig and orig in file_content:
updated_content = file_content.replace(orig, new_val)
self.git_tool.write_file(Path(state.repo_path), change.file_path, updated_content)
logger.info(f"Successfully applied exact patch to {change.file_path}")
elif orig and orig.strip() in file_content:
updated_content = file_content.replace(orig.strip(), new_val.strip())
self.git_tool.write_file(Path(state.repo_path), change.file_path, updated_content)
logger.info(f"Successfully applied stripped patch to {change.file_path}")
else:
is_md = change.file_path.lower().endswith('.md')
# Fallback: if we can't find original, only overwrite if orig is empty, file is tiny, or it is a markdown file
if not orig.strip() or len(file_content) < 200 or is_md:
logger.warning(f"Original block empty, file tiny, or is markdown for {change.file_path}. Falling back to overwrite.")
self.git_tool.write_file(Path(state.repo_path), change.file_path, new_val)
else:
logger.error(f"Failed to apply patch to {change.file_path}: ORIGINAL block not found in file.")
state.error = f"Patch failed for {change.file_path}: ORIGINAL content could not be matched exactly. The agent must rewrite the ORIGINAL block to match the file."
return state
else:
# New file
self.git_tool.write_file(Path(state.repo_path), change.file_path, change.new_content)
state.changes = changes
state.branch_name = branch_name
logger.info(f"Generated {len(changes)} code changes on branch {branch_name}")
return state
except Exception as e:
logger.error(f"Code generation failed: {e}")
state.error = f"Code generation failed: {str(e)}"
return state
def _prepare_files_content(self, state: AgentState) -> str:
"""Prepare file contents for LLM prompt with smart truncation to preserve token budget."""
content_parts = []
# Limit to top 3 files instead of 5 to avoid prompt bloat
for result in state.relevant_files[:3]:
file_path = result.file_path
full_content = self.git_tool.get_file_content(
Path(state.repo_path),
file_path,
)
if not full_content:
continue
lines = full_content.splitlines()
total_lines = len(lines)
# Thresholds: 400 lines or 16000 characters
if total_lines <= 400 and len(full_content) <= 16000:
content_parts.append(f"File: {file_path}\n```\n{full_content}\n```\n")
continue
# If the file is larger, we perform smart sliding window extraction around the search snippet
match_idx = -1
snippet = result.content.strip() if result.content else ""
if snippet:
# Clean up the snippet for matching. Take the first non-empty line of the snippet
snippet_lines = [l.strip() for l in snippet.splitlines() if l.strip()]
target = snippet_lines[0] if snippet_lines else ""
if len(target) >= 6: # Only look for significant matches
for idx, line in enumerate(lines):
if target.lower() in line.lower():
match_idx = idx
break
if match_idx != -1:
# Extract a +/- 100 line window around the match
start_idx = max(0, match_idx - 100)
end_idx = min(total_lines, match_idx + 100)
truncated_lines = []
if start_idx > 0:
truncated_lines.append(f"// ... [TRUNCATED: Showing lines {start_idx + 1} to {end_idx} of {total_lines}. Preceding {start_idx} lines omitted to keep prompt size small] ...")
truncated_lines.extend(lines[start_idx:end_idx])
if end_idx < total_lines:
truncated_lines.append(f"// ... [TRUNCATED: Showing lines {start_idx + 1} to {end_idx} of {total_lines}. Remaining {total_lines - end_idx} lines omitted to keep prompt size small] ...")
truncated_content = "\n".join(truncated_lines)
logger.info(f"Smart truncated file '{file_path}' to lines {start_idx+1}-{end_idx} around search match (line {match_idx+1}).")
else:
# Fallback: Extract first 350 lines
limit = min(350, total_lines)
truncated_lines = lines[:limit]
if total_lines > limit:
truncated_lines.append(f"// ... [TRUNCATED: File is too large ({total_lines} lines). First {limit} lines shown to keep prompt size small] ...")
truncated_content = "\n".join(truncated_lines)
logger.info(f"File '{file_path}' too large and no clear search match found. Showing first {limit} lines.")
content_parts.append(f"File: {file_path} (Truncated)\n```\n{truncated_content}\n```\n")
return "\n".join(content_parts)
def _parse_changes(self, content: str) -> list[CodeChange]:
"""Parse code changes from LLM response."""
changes = []
sections = content.split("---")
for section in sections:
if not section.strip():
continue
try:
file_path = self._extract_field(section, "FILE:")
explanation = self._extract_field(section, "EXPLANATION:")
original = self._extract_code_block(section, "ORIGINAL:")
new = self._extract_code_block(section, "NEW:")
if file_path and new:
changes.append(CodeChange(
file_path=file_path,
original_content=original,
new_content=new,
explanation=explanation or "Code fix",
))
except Exception as e:
logger.warning(f"Failed to parse change section: {e}")
continue
return changes
def _extract_field(self, content: str, marker: str) -> str:
"""Extract single-line field from content."""
for line in content.split("\n"):
if marker in line:
return line.replace(marker, "").strip()
return ""
def _extract_code_block(self, content: str, marker: str) -> str:
"""Extract code block from content."""
lines = content.split("\n")
result = []
capturing = False
for line in lines:
if marker in line:
capturing = True
continue
elif capturing:
if line.strip() in ["FILE:", "EXPLANATION:", "ORIGINAL:", "NEW:", "---"]:
break
result.append(line)
# Don't strip leading/trailing spaces as it breaks indentation matching
code = "\n".join(result).strip("\r\n")
# Remove markdown code block backticks if present
if code.startswith("```"):
lines = code.split("\n")
if len(lines) > 1:
code = "\n".join(lines[1:])
else:
code = ""
if code.endswith("```"):
lines = code.split("\n")
if len(lines) > 1:
code = "\n".join(lines[:-1])
else:
code = ""
# Only strip newlines again, preserve horizontal whitespace
return code.strip("\r\n")
|