Building a Code Assistant with LangGraph: Let AI Write and Refine Code

AI‑generated code often contains syntax errors or logical bugs, requiring repeated debugging.

Workflow Overview

user request → generate code → check errors → passes? → finalize
                               │                ▼
                               │                fix
                               └───────────────────┘

Core nodes: generate, check, fix, finalize.

State Definition

from typing import TypedDict, Literal
from langgraph.graph import StateGraph, START, END

class CodeAssistantState(TypedDict):
    task: str
    code: str
    error: str  # "yes" / "no"
    iterations: int
    result: str

Node Implementations

Generate Code

def generate_code(state: CodeAssistantState) -> dict:
    """Generate code (LLM integration in real projects)"""
    task = state["task"]
    if "hello" in task.lower():
        code = 'print("Hello, World!")'
    elif "加法" in task or "add" in task.lower():
        code = "def add(a, b):
    return a + b"
    else:
        code = "# TODO: implement feature"
    return {"code": code, "iterations": state["iterations"] + 1}

Check Code

def check_code(state: CodeAssistantState) -> dict:
    """Check syntax and execution"""
    code = state["code"]
    try:
        compile(code, "<string>", "exec")
        exec(code)
        return {"error": "no"}
    except Exception as e:
        return {"error": "yes", "result": f"Error: {e}"}

Fix Code

def fix_code(state: CodeAssistantState) -> dict:
    """Fix code issues"""
    return {"code": "# Fixed code", "iterations": state["iterations"] + 1}

Routing and Finalization

def decide_next(state: CodeAssistantState) -> Literal["finalize", "fix"]:
    """Decide next step based on check result"""
    if state["error"] == "no" or state["iterations"] >= 3:
        return "finalize"
    return "fix"

def finalize(state: CodeAssistantState) -> dict:
    """Output final code"""
    return {"result": f"Final code:
{state['code']}"}

Full Graph Construction

builder = StateGraph(CodeAssistantState)
builder.add_node("generate", generate_code)
builder.add_node("check", check_code)
builder.add_node("fix", fix_code)
builder.add_node("finalize", finalize)

builder.add_edge(START, "generate")
builder.add_edge("generate", "check")
builder.add_conditional_edges(
    "check",
    decide_next,
    {"finalize": "finalize", "fix": "fix"},
)
builder.add_edge("fix", "check")
builder.add_edge("finalize", END)

app = builder.compile(checkpointer=InMemorySaver())

Execution Examples

config = {"configurable": {"thread_id": "code-assistant"}}
# Example 1: Simple task
result = app.invoke({
    "task": "Hello World 程序",
    "code": "",
    "error": "",
    "iterations": 0,
    "result": "",
}, config)
print(result["result"])
# Example 2: Add function
result = app.invoke({
    "task": "实现加法函数",
    "code": "",
    "error": "",
    "iterations": 0,
    "result": "",
}, config)
print(result["result"])

Key Features of the Official LangGraph Code Assistant

Mistral model integration

Multi‑round iterative optimization

Detailed error feedback

Related Links

Official documentation: https://github.com/langchain-ai/langgraph/tree/main/examples/code_assistant
GitHub source: https://github.com/langchain-ai/langgraph/blob/main/examples/code_assistant