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")