Langgraph Node Resume Not Working

[abhinav-mane]

`

import os
from typing import Dict, TypedDict, List, Annotated, Literal
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.sqlite import SqliteSaver
import sqlite3

class ConversationState(TypedDict):
    messages: List[Dict[str, str]]
    current_stage: str
    user_input: str
    system_response: str
    verification_status: Dict[str, bool]


def greeting_node(state: ConversationState) -> ConversationState:
    """First node that greets the user and asks for ID"""
    print("📌 EXECUTING: greeting_node"
    state["current_stage"] = "greeting"
    state["system_response"] = "Hello! Welcome to our service. Please provide your customer ID."
    state["messages"].append({
        "role": "assistant", 
        "content": state["system_response"]
    })
    return state



def verify_id_node(state: ConversationState) -> ConversationState:
    """Second node that verifies the user ID"""
    print("📌 EXECUTING: verify_id_node")
    
    state["current_stage"] = "id_verification"
    user_id = state["user_input"]
    
    # Simple mock verification (in real app, this would query a database)
    if user_id.isdigit() and len(user_id) == 5:
        state["verification_status"]["id_verified"] = True
        state["system_response"] = f"Thank you! Your ID {user_id} has been verified. What's your date of birth? (MM/DD/YYYY)"
    else:
        state["verification_status"]["id_verified"] = False
        state["system_response"] = "I couldn't verify that ID. Please provide a valid 5-digit customer ID."
    
    # Add this response to the message history
    state["messages"].append({
        "role": "assistant", 
        "content": state["system_response"]
    })
    
    return state


def verify_dob_node(state: ConversationState) -> ConversationState:
    """Third node that verifies date of birth"""
    print("📌 EXECUTING: verify_dob_node")
    
    state["current_stage"] = "dob_verification"
    user_dob = state["user_input"]
    
    # Simple mock verification with basic format checking
    import re
    if re.match(r"^\d{2}/\d{2}/\d{4}$", user_dob):
        state["verification_status"]["dob_verified"] = True
        state["system_response"] = f"Thank you! Your date of birth has been verified. How can I help you today?"
    else:
        state["verification_status"]["dob_verified"] = False
        state["system_response"] = "I couldn't verify that date of birth. Please provide it in MM/DD/YYYY format."
    
    # Add this response to the message history
    state["messages"].append({
        "role": "assistant", 
        "content": state["system_response"]
    })
    
    return state


def handle_query_node(state: ConversationState) -> ConversationState:
    """Final node that processes the user's actual query"""
    print("📌 EXECUTING: handle_query_node")
    
    state["current_stage"] = "query_handling"
    user_query = state["user_input"].lower()
    
    # Simple mock response generation
    if "balance" in user_query:
        state["system_response"] = "Your current balance is $1,250.00. Is there anything else you'd like to know?"
    elif "payment" in user_query:
        state["system_response"] = "Your last payment of $350.00 was received on April 15, 2025. Your next payment is due on May 10, 2025."
    elif "statement" in user_query:
        state["system_response"] = "I'll email your latest statement right away. You should receive it within a few minutes."
    else:
        state["system_response"] = "I can help with checking your balance, payment information, or sending statements. What would you like to do?"
    
    # Add this response to the message history
    state["messages"].append({
        "role": "assistant", 
        "content": state["system_response"]
    })
    
    return state


# 3. Define our conditional edge functions
def should_proceed_to_dob_verification(state: ConversationState) -> Literal["verify_dob", "verify_id"]:
    """Determine if we should move to DOB verification or retry ID verification"""
    print("🔀 DECISION: ID verification check")
    if state["verification_status"]["id_verified"]:
        return "verify_dob"
    else:
        return "verify_id"


def should_proceed_to_query(state: ConversationState) -> Literal["handle_query", "verify_dob"]:
    """Determine if we should move to query handling or retry DOB verification"""
    print("🔀 DECISION: DOB verification check")
    if state["verification_status"]["dob_verified"]:
        return "handle_query"
    else:
        return "verify_dob"


# 4. Set up our graph
def build_graph():
    # Create a new graph with our state schema
    builder = StateGraph(ConversationState)
    
    # Add all of our nodes
    builder.add_node("greeting", greeting_node)
    builder.add_node("verify_id", verify_id_node)
    builder.add_node("verify_dob", verify_dob_node)
    builder.add_node("handle_query", handle_query_node)
    
    # Add edges
    builder.add_edge(START, "greeting")
    builder.add_edge("greeting", "verify_id")
    
    # Add conditional edges based on verification status
    builder.add_conditional_edges(
        "verify_id",
        should_proceed_to_dob_verification,
        {
            "verify_dob": "verify_dob",
            "verify_id": "verify_id"  # Loop back if verification failed
        }
    )
    
    builder.add_conditional_edges(
        "verify_dob",
        should_proceed_to_query,
        {
            "handle_query": "handle_query",
            "verify_dob": "verify_dob"  # Loop back if verification failed
        }
    )
    
    # Simple edge from query handling to end
    builder.add_edge("handle_query", END)
    
    # Set up SQLite checkpoint storage
    db_path = './conversation_checkpoints.sqlite'
    conn = sqlite3.connect(db_path, check_same_thread=False)
    memory = SqliteSaver(conn)
    
    # Compile the graph with the checkpointer
    # return builder.compile(checkpointer=memory)

    return builder.compile(
            checkpointer=memory,
            # pause before each node that needs user input
            interrupt_before=["verify_id", "verify_dob", "handle_query"]
        )


# 5. Test simulation function
def process_user_input(graph, user_input: str, thread_id: str = "test_thread"):
    config = {"configurable": {"thread_id": thread_id}}
    snap = graph.get_state(config)

    if snap:
        # --- Resume path ---
        # Make sure 'messages' exists

        print("🛠️  Resuming at:", snap.next)    # debug

        if "messages" not in snap.values:
            snap.values["messages"] = []
        # Make sure 'verification_status' exists
        if "verification_status" not in snap.values:
            snap.values["verification_status"] = {"id_verified": False, "dob_verified": False}

        snap.values["user_input"] = user_input
        snap.values["messages"].append({"role": "user", "content": user_input})

        stream = graph.stream(
            snap.values,
            config={**config, "start_at": snap.next}
        )
    else:
        # --- New conversation path ---
        initial_state = ConversationState(
            messages=[{"role": "user", "content": user_input}],
            current_stage="initial",
            user_input=user_input,
            system_response="",
            verification_status={"id_verified": False, "dob_verified": False}
        )
        stream = graph.stream(initial_state, config=config)

    # Only pull the first event from the stream
    try:
        first_event = next(stream)
    except StopIteration:
        raise RuntimeError("Graph terminated without producing a response")

    node_name, state = next(iter(first_event.items()))
    return state




# Build our graph
graph = build_graph()

`

====================== OUTPUT =========================================

SESSION_ID = "test_user_another_sess"

Step 1: Initial greeting
print("=" * 50)
print("STEP 1: INITIAL GREETING")
print("=" * 50)
response = process_user_input(graph, "Hello", SESSION_ID)
print(f"\nSystem: {response['system_response']}\n")

Output: -

STEP 1: INITIAL GREETING

🛠️ Resuming at: ()
📌 EXECUTING: greeting_node

System: Hello! Welcome to our service. Please provide your customer ID.

Step 2: Provide invalid ID (should loop back to ID verification)
print("=" * 50)
print("STEP 2: INVALID CUSTOMER ID")
print("=" * 50)
response = process_user_input(graph, "123", SESSION_ID)
print(f"\nSystem: {response['system_response']}\n")

Output : -

STEP 2: INVALID CUSTOMER ID

🛠️ Resuming at: ('verify_id',)
📌 EXECUTING: greeting_node

System: Hello! Welcome to our service. Please provide your customer ID.

Step 3: Provide valid ID
print("=" * 50)
print("STEP 3: VALID CUSTOMER ID")
print("=" * 50)
response = process_user_input(graph, "12345", SESSION_ID)
print(f"\nSystem: {response['system_response']}\n")

Output: -

STEP 3: VALID CUSTOMER ID

🛠️ Resuming at: ('verify_id',)
📌 EXECUTING: greeting_node

System: Hello! Welcome to our service. Please provide your customer ID.

[abhinav-mane]

I dont know what I am doing wrong in this why is it not working or resuming from the past node that it left off?