Build a LangChain Agent for Code Generation and Auto-Fix

pip install langchain langchain-openai langchain-community python-dotenv

import os
from dotenv import load_dotenv
load_dotenv()

# Required: OPENAI_API_KEY=your-openai-api-key

from langchain_community.tools import PythonREPLTool

repl = PythonREPLTool()

# Test basic execution
result = repl.run("print('Hello from the REPL!')")
print(result)
# → Hello from the REPL!

# Test with computation
result = repl.run("""
import statistics
data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
print(f"Mean: {statistics.mean(data)}")
print(f"Std Dev: {statistics.stdev(data):.2f}")
""")
print(result)

# Test error capture
error_result = repl.run("print(undefined_variable)")
print(error_result)
# → NameError: name 'undefined_variable' is not defined

import subprocess
import sys
import tempfile
import os
from typing import Tuple
import resource

def safe_execute_python(code: str, timeout_seconds: int = 10) -> Tuple[str, str, int]:
    """
    Execute Python code in an isolated subprocess.
    Returns: (stdout, stderr, return_code)
    """
    # Write code to a temp file
    with tempfile.NamedTemporaryFile(
        mode="w",
        suffix=".py",
        delete=False,
        encoding="utf-8"
    ) as f:
        f.write(code)
        temp_path = f.name

    try:
        result = subprocess.run(
            [sys.executable, temp_path],
            capture_output=True,
            text=True,
            timeout=timeout_seconds,
            # Restrict environment
            env={
                "PATH": os.environ.get("PATH", ""),
                "PYTHONPATH": "",
                "HOME": tempfile.gettempdir()
            }
        )
        return result.stdout, result.stderr, result.returncode
    except subprocess.TimeoutExpired:
        return "", f"TimeoutError: Code execution exceeded {timeout_seconds} seconds", 1
    except Exception as e:
        return "", str(e), 1
    finally:
        os.unlink(temp_path)

# Test the sandbox
stdout, stderr, code = safe_execute_python("""
def fibonacci(n):
    a, b = 0, 1
    for _ in range(n):
        a, b = b, a + b
    return a

for i in range(10):
    print(f"fib({i}) = {fibonacci(i)}")
""")

print("STDOUT:", stdout)
print("STDERR:", stderr)
print("Return code:", code)

from langchain_core.tools import tool
from typing import Optional

@tool
def execute_python_safe(code: str) -> str:
    """
    Execute Python code in a sandboxed subprocess and return the result.
    Returns stdout on success, error message on failure.
    Use this to test generated code before returning it to the user.
    """
    stdout, stderr, return_code = safe_execute_python(code, timeout_seconds=15)
    
    if return_code == 0:
        return f"SUCCESS\nOutput:\n{stdout}"
    else:
        return f"ERROR (exit code {return_code})\nError:\n{stderr}\nOutput:\n{stdout}"

@tool
def write_code_to_file(filename: str, code: str) -> str:
    """
    Save generated code to a file in the workspace directory.
    Only use after the code has been successfully tested.
    """
    workspace = "./generated_code"
    os.makedirs(workspace, exist_ok=True)
    filepath = os.path.join(workspace, filename)
    
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(code)
    
    return f"Code saved to {filepath}"

@tool  
def read_file(filepath: str) -> str:
    """Read the contents of a file from the workspace."""
    workspace = "./generated_code"
    safe_path = os.path.join(workspace, os.path.basename(filepath))
    
    try:
        with open(safe_path, "r", encoding="utf-8") as f:
            return f.read()
    except FileNotFoundError:
        return f"File not found: {filepath}"

# Test the custom tool
result = execute_python_safe.invoke("""
import json
data = {"name": "Alice", "scores": [95, 87, 92]}
avg = sum(data["scores"]) / len(data["scores"])
print(f"Student: {data['name']}")
print(f"Average score: {avg:.1f}")
""")
print(result)

from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI

CODE_AGENT_SYSTEM = """You are an expert Python programming assistant with a write→test→fix workflow.

For every coding task:
1. Generate complete, working Python code
2. Test it using execute_python_safe
3. If it fails, read the error carefully and fix the code
4. Repeat until the code runs successfully (max 5 attempts)
5. Once successful, save the final code using write_code_to_file

Code quality standards:
- Include type hints for all function parameters and return values
- Add docstrings for all functions and classes
- Use descriptive variable names
- Handle edge cases and include basic error handling
- Write code that is testable and modular

When you encounter an error:
- Read the full traceback carefully
- Identify the root cause (not just the symptom)
- Fix the specific issue before retesting
- Don't change code that was working — isolate the fix"""

llm = ChatOpenAI(model="gpt-4o", temperature=0.1)

tools = [execute_python_safe, write_code_to_file, read_file]

prompt = ChatPromptTemplate.from_messages([
    ("system", CODE_AGENT_SYSTEM),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{input}"),
    MessagesPlaceholder(variable_name="agent_scratchpad")
])

agent = create_tool_calling_agent(llm, tools, prompt)
code_agent = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    max_iterations=20,  # Allow multiple write→test→fix cycles
    handle_parsing_errors=True,
    return_intermediate_steps=True
)

from dataclasses import dataclass, field
from datetime import datetime
from typing import List

@dataclass
class CodeGenerationResult:
    task: str
    final_code: str
    execution_output: str
    attempts: int
    success: bool
    errors_encountered: List[str] = field(default_factory=list)
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())

def generate_and_fix(task: str, max_iterations: int = 20) -> CodeGenerationResult:
    """Run the full code generation pipeline with auto-fix."""
    
    print(f"\nTask: {task}")
    print("=" * 60)
    
    result = code_agent.invoke({
        "input": task,
        "chat_history": []
    })
    
    # Parse intermediate steps for metrics
    attempts = 0
    errors = []
    
    for action, observation in result.get("intermediate_steps", []):
        if action.tool == "execute_python_safe":
            attempts += 1
            if "ERROR" in str(observation):
                errors.append(str(observation)[:200])
    
    success = "ERROR" not in result["output"] and len(result["output"]) > 10
    
    return CodeGenerationResult(
        task=task,
        final_code=result["output"],
        execution_output=str(result.get("intermediate_steps", [])),
        attempts=attempts,
        success=success,
        errors_encountered=errors
    )

# Test tasks
tasks = [
    "Write a function that reads a CSV file and computes summary statistics (mean, median, std dev) for each numeric column. Include a test with sample data.",
    "Create a class called BinarySearchTree with insert, search, and in_order_traversal methods. Test it with 10 random integers.",
    "Write a decorator that memoizes function results and tracks cache hit/miss rates. Include a Fibonacci example to demonstrate performance improvement.",
]

results = [generate_and_fix(task) for task in tasks]

for r in results:
    status = "PASSED" if r.success else "FAILED"
    print(f"\n[{status}] {r.task[:60]}...")
    print(f"  Attempts: {r.attempts}")
    print(f"  Errors encountered: {len(r.errors_encountered)}")

from pydantic import BaseModel, Field
from typing import List, Optional
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate

class CodeModule(BaseModel):
    filename: str = Field(description="Python filename (e.g., 'calculator.py')")
    implementation: str = Field(description="The complete implementation code")
    test_code: str = Field(description="pytest test code for the implementation")
    dependencies: List[str] = Field(description="pip packages required (e.g., ['numpy', 'pandas'])")
    docstring: str = Field(description="Module-level description of what this code does")

def generate_structured_code(task: str) -> CodeModule:
    """Generate implementation + tests as structured output."""
    
    parser = JsonOutputParser(pydantic_object=CodeModule)
    
    prompt = ChatPromptTemplate.from_messages([
        ("system", """You are an expert Python developer. Generate production-quality code with tests.
        
        Always:
        - Use type hints
        - Write comprehensive pytest tests (test happy path, edge cases, error cases)
        - Handle errors gracefully
        - Follow PEP 8 style guide"""),
        ("human", """Task: {task}
        
{format_instructions}

Return the JSON object with all fields filled in.""")
    ])
    
    from langchain_openai import ChatOpenAI
    from langchain.output_parsers import OutputFixingParser
    
    llm = ChatOpenAI(model="gpt-4o", temperature=0.1)
    
    chain = (
        prompt.partial(format_instructions=parser.get_format_instructions())
        | llm
        | OutputFixingParser.from_llm(parser=parser, llm=llm)
    )
    
    return chain.invoke({"task": task})

def generate_and_validate(task: str) -> dict:
    """Generate code, test it, auto-fix if needed."""
    
    # Step 1: Generate structured code
    print(f"Generating code for: {task[:60]}...")
    code_module = generate_structured_code(task)
    
    print(f"Generated: {code_module.filename}")
    print(f"Dependencies: {code_module.dependencies}")
    
    # Step 2: Test the implementation
    test_result = execute_python_safe.invoke(code_module.implementation + "\n\n# Quick sanity check\nprint('Module loaded successfully')")
    
    if "ERROR" in test_result:
        print(f"Implementation error: {test_result[:200]}")
        # Trigger the full agent fix loop
        fix_result = generate_and_fix(
            f"Fix this Python code:\n\nCode:\n{code_module.implementation}\n\nError:\n{test_result}"
        )
        code_module.implementation = extract_code_from_response(fix_result.final_code)
    
    # Step 3: Run the tests
    combined_code = code_module.implementation + "\n\n" + code_module.test_code.replace("if __name__ == '__main__':", "if True:")
    test_execution = execute_python_safe.invoke(combined_code)
    
    return {
        "filename": code_module.filename,
        "implementation": code_module.implementation,
        "test_code": code_module.test_code,
        "test_result": test_execution,
        "success": "ERROR" not in test_execution
    }

def extract_code_from_response(text: str) -> str:
    """Extract Python code from agent response."""
    import re
    match = re.search(r'```python\n(.*?)```', text, re.DOTALL)
    if match:
        return match.group(1)
    return text

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI
import re

def iterative_code_fixer(
    task: str,
    max_attempts: int = 5,
    model: str = "gpt-4o"
) -> dict:
    """
    Standalone write→test→fix loop without the full agent framework.
    More transparent and easier to debug than the agent approach.
    """
    llm = ChatOpenAI(model=model, temperature=0.1)
    
    # Code generation prompt
    generate_prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a Python expert. Write complete, working Python code. Return ONLY the code, no explanation."),
        ("human", "Write Python code that: {task}")
    ])
    
    # Fix prompt
    fix_prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a Python debugging expert. Fix the broken code. Return ONLY the corrected code, no explanation."),
        ("human", """Task: {task}

Previous code:
# start: python code
{code}
# end code block

Error encountered:
{error}

Fixed code:""")
    ])
    
    generate_chain = generate_prompt | llm | StrOutputParser()
    fix_chain = fix_prompt | llm | StrOutputParser()
    
    def extract_code(text: str) -> str:
        """Extract code from markdown code blocks if present."""
        match = re.search(r'```(?:python)?\n(.*?)```', text, re.DOTALL)
        return match.group(1).strip() if match else text.strip()
    
    history = []
    
    # Initial code generation
    raw_code = generate_chain.invoke({"task": task})
    current_code = extract_code(raw_code)
    
    history.append({
        "attempt": 0,
        "action": "generate",
        "code": current_code
    })
    
    for attempt in range(max_attempts):
        # Execute the code
        stdout, stderr, return_code = safe_execute_python(current_code)
        
        if return_code == 0:
            # Success
            history.append({
                "attempt": attempt + 1,
                "action": "success",
                "output": stdout
            })
            
            print(f"SUCCESS on attempt {attempt + 1}")
            return {
                "success": True,
                "code": current_code,
                "output": stdout,
                "attempts": attempt + 1,
                "history": history
            }
        else:
            # Fix the error
            error_msg = stderr or "Unknown error"
            print(f"Attempt {attempt + 1} failed: {error_msg[:100]}")
            
            history.append({
                "attempt": attempt + 1,
                "action": "fix",
                "error": error_msg[:300]
            })
            
            if attempt < max_attempts - 1:
                raw_fixed = fix_chain.invoke({
                    "task": task,
                    "code": current_code,
                    "error": error_msg
                })
                current_code = extract_code(raw_fixed)
    
    # All attempts exhausted
    return {
        "success": False,
        "code": current_code,
        "output": stderr,
        "attempts": max_attempts,
        "history": history
    }

# Test the standalone fixer
result = iterative_code_fixer(
    task="Read a JSON file called 'data.json', extract all values for the key 'score', and print the average. Handle the case where the file doesn't exist.",
    max_attempts=5
)

print(f"Success: {result['success']}")
print(f"Attempts used: {result['attempts']}")
if result["success"]:
    print(f"Final output: {result['output']}")

@tool
def run_code_quality_checks(code: str) -> str:
    """
    Run automated quality checks on Python code:
    - Syntax validation
    - Basic style checks
    - Security scan for obvious issues
    Returns a report with any issues found.
    """
    import ast
    import re
    
    issues = []
    
    # 1. Syntax check
    try:
        ast.parse(code)
    except SyntaxError as e:
        return f"SYNTAX ERROR: {e}"
    
    # 2. Security checks (basic)
    dangerous_patterns = [
        (r'\beval\b', "Use of eval() is dangerous"),
        (r'\bexec\b', "Use of exec() is dangerous"),
        (r'__import__', "Dynamic imports may indicate code injection"),
        (r'os\.system\b', "Use subprocess instead of os.system"),
        (r'subprocess\.call.*shell=True', "shell=True is a security risk"),
    ]
    
    for pattern, message in dangerous_patterns:
        if re.search(pattern, code):
            issues.append(f"SECURITY WARNING: {message}")
    
    # 3. Style checks
    lines = code.split("\n")
    for i, line in enumerate(lines, 1):
        if len(line) > 120:
            issues.append(f"Line {i}: exceeds 120 characters ({len(line)} chars)")
    
    # 4. Check for type hints on function definitions
    function_defs = re.findall(r'def \w+\([^)]*\):', code)
    unhinted = [f for f in function_defs if '->' not in f and f != 'def __init__(self):']
    if unhinted:
        issues.append(f"Missing return type hints on {len(unhinted)} function(s)")
    
    if not issues:
        return "PASSED: No quality issues found"
    
    return "ISSUES FOUND:\n" + "\n".join(f"  - {issue}" for issue in issues)

# Add to the agent tools
tools_with_quality = tools + [run_code_quality_checks]

code_agent_v2 = AgentExecutor(
    agent=create_tool_calling_agent(
        ChatOpenAI(model="gpt-4o"),
        tools_with_quality,
        ChatPromptTemplate.from_messages([
            ("system", CODE_AGENT_SYSTEM + "\n\nAlways run run_code_quality_checks before saving final code."),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad")
        ])
    ),
    tools=tools_with_quality,
    verbose=True,
    max_iterations=25
)

import asyncio

async def stream_code_generation(task: str):
    """Stream code generation events for real-time UI updates."""
    
    async for event in code_agent.astream_events(
        {"input": task, "chat_history": []},
        version="v1"
    ):
        kind = event["event"]
        
        if kind == "on_tool_start":
            tool_name = event["name"]
            if tool_name == "execute_python_safe":
                print("\n[Executing code...]")
            elif tool_name == "write_code_to_file":
                print("\n[Saving file...]")
        
        elif kind == "on_tool_end":
            output = str(event["data"].get("output", ""))
            if "SUCCESS" in output:
                print("[Code executed successfully]")
            elif "ERROR" in output:
                error_preview = output.split("\n")[1:3]
                print(f"[Execution error: {' '.join(error_preview)[:100]}]")
        
        elif kind == "on_chat_model_stream":
            chunk = event["data"]["chunk"]
            if hasattr(chunk, "content") and chunk.content:
                print(chunk.content, end="", flush=True)

asyncio.run(stream_code_generation(
    "Write a function to parse HTTP server logs and count requests by status code"
))

result = generate_and_fix(
    "Write a script that reads a list of URLs from a text file (one per line), checks if each URL is accessible (HTTP 200), and writes a report showing which URLs are working vs broken. Include timeout handling."
)

result = generate_and_fix(
    "Implement Dijkstra's shortest path algorithm with a priority queue. Include a test with a sample weighted graph and print the shortest path between nodes."
)

result = generate_and_fix(
    "Write a Python class that wraps the OpenWeatherMap API. Include methods: get_current_weather(city), get_forecast(city, days), and handle rate limiting with automatic retry. Use requests library."
)

from langchain_community.callbacks import get_openai_callback

def generate_with_cost_tracking(task: str) -> dict:
    with get_openai_callback() as cb:
        result = generate_and_fix(task)
    
    print(f"\nCost breakdown:")
    print(f"  Total tokens: {cb.total_tokens:,}")
    print(f"  Total cost: ${cb.total_cost:.4f}")
    
    return {**result.__dict__, "cost": cb.total_cost, "tokens": cb.total_tokens}