Building Self-Correcting LangGraph Agents with Human-in-the-Loop Safeguards

                  ┌──────────────┐
                  │  Draft Node  │
                  └──────┬───────┘
                         │
                         ▼
                  ┌──────────────┐
             ┌───►│ Validate Node│
             │    └──────┬───────┘
             │           │
     Failed  │           ▼ Passed
[Max 3 Tries]│    ┌──────────────┐
             └───┤  Audit Node  │
                 └──────┬───────┘
                        │
                        ▼ High Risk?
                 ┌──────────────┐
                 │ HITL Pause   │◄── Human Approves / Rejects
                 └──────┬───────┘
                        │
                        ▼ Approved
                 ┌──────────────┐
                 │ Execute Node │
                 └──────────────┘

# schema_agent_graph.py
from typing import TypedDict, Annotated, List, Optional
import operator
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver

class AgentState(TypedDict):
    schema_sql: str
    target_dialect: str
    attempts: int
    validation_errors: Annotated[List[str], operator.add]
    is_valid: bool
    human_approved: Optional[bool]

def draft_migration_node(state: AgentState) -> dict:
    """Generates the initial SQL migration script or refactors based on errors."""
    attempts = state.get("attempts", 0) + 1
    errors = state.get("validation_errors", [])
    
    if errors:
        # Prompt LLM to fix specific validation errors from previous turn
        fixed_sql = f"-- Fixed based on error: {errors[-1]}\nALTER TABLE users ADD COLUMN updated_at TIMESTAMP;"
        return {"schema_sql": fixed_sql, "attempts": attempts}
    
    initial_sql = "CREATE TABLE users (id SERIAL PRIMARY KEY, username VARCHAR(50));"
    return {"schema_sql": initial_sql, "attempts": attempts}

def validate_sql_node(state: AgentState) -> dict:
    """Simulates dry-run execution against an isolated database container."""
    sql = state["schema_sql"]
    
    # Simulate validation failure on first attempt to demonstrate self-healing
    if state["attempts"] == 1:
        return {
            "validation_errors": ["Syntax error: missing semicolon or invalid constraint"],
            "is_valid": False
        }
    
    return {"validation_errors": [], "is_valid": True}

def should_retry_or_proceed(state: AgentState) -> str:
    """Conditional edge router: Decide whether to self-correct, fail, or proceed."""
    if state["is_valid"]:
        return "human_approval_checkpoint"
    if state["attempts"] >= 3:
        return "fail_node"
    return "draft_migration"

def execute_node(state: AgentState) -> dict:
    """Executes the verified script against production."""
    if not state.get("human_approved"):
        raise PermissionError("Cannot execute unapproved migration.")
    print("Executing migration on production database...")
    return {}

# Build the Graph
builder = StateGraph(AgentState)
builder.add_node("draft_migration", draft_migration_node)
builder.add_node("validate_sql", validate_sql_node)
builder.add_node("execute_migration", execute_node)

builder.set_entry_point("draft_migration")
builder.add_edge("draft_migration", "validate_sql")
builder.add_conditional_edges(
    "validate_sql",
    should_retry_or_proceed,
    {
        "draft_migration": "draft_migration",
        "human_approval_checkpoint": END,  # Pause here for human review
        "fail_node": END
    }
)

# Persistent checkpointer saves thread state across processes
checkpointer = MemorySaver()
app = builder.compile(checkpointer=checkpointer, interrupt_before=["execute_migration"])

# run_with_hitl.py
from schema_agent_graph import app

thread_config = {"configurable": {"thread_id": "migration_run_402"}}

# Step 1: Run graph until human approval pause
events = app.stream(
    {"schema_sql": "", "target_dialect": "postgresql", "attempts": 0, "validation_errors": [], "is_valid": False, "human_approved": None},
    config=thread_config
)

for event in events:
    print(f"Step: {event}")

# Inspect current state waiting for approval
current_state = app.get_state(thread_config)
print("\n--- WAITING FOR OPERATOR APPROVAL ---")
print("Pending SQL:")
print(current_state.values["schema_sql"])

# Step 2: Human Operator clicks "Approve" in UI
app.update_state(
    thread_config,
    {"human_approved": True},
    as_node="draft_migration"
)

# Step 3: Resume execution from checkpoint
resume_events = app.stream(None, config=thread_config)
for event in resume_events:
    print(f"Resumed Step: {event}")