Spaces:
Sleeping
Sleeping
| """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") | |