Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from langgraph.graph import StateGraph, END
-from langgraph.checkpoint.sqlite import SqliteSaver
+# We are using MemorySaver now
 from langgraph.checkpoint.memory import MemorySaver
 import operator
 from typing import TypedDict, Annotated, Optional
@@ -21,22 +21,20 @@ def greeting_node(state: GraphState) -> GraphState:
     print(greeting)
     return {"greeting_message": greeting, "current_node": "greeting"}
 
-# IMPORTANT CHANGE HERE: Remove input() from this node!
+# IMPORTANT: This node will now run *after* human input is in the state.
+# It does NOT ask for input.
 def human_input_node(state: GraphState) -> GraphState:
     """
-    Node 2: This node now simply acts as the interruption point.
-    The human input will be provided by the Gradio UI *after* this node runs and interrupts.
+    Node 2: This node processes the human input that was added to the state
+    before its execution (during the resume phase via put_state).
     """
-    print("\n✋ Human Input Node Executed (This is where graph will interrupt AFTER running this node):")
-    # This node returns its current_node state. The actual human_input will be added
-    # to the state during the RESUME phase from the Gradio UI.
-    # Since we are using 'interrupt_after' and human_input_node doesn't take input directly,
-    # the 'human_input' will only appear in the state after 'resume_graph' sends it.
-    return {"current_node": "human_input_interrupt"}
-
+    human_response = state.get("human_input", "No human input found in state when human_input_node ran.")
+    print(f"\n✋ Human Input Node Executed (Processing input: '{human_response}'):")
+    # You could add validation or further processing of state['human_input'] here if needed.
+    return {"current_node": "human_input_processed"}
 
 def human_response_display_node(state: GraphState) -> GraphState:
-    human_response = state.get("human_input", "No human input received.")
+    human_response = state.get("human_input", "No human input received for final display.")
     final_message = f"You said: '{human_response}'. Thank you for your input!"
     print("\n✅ Human Response Display Node Executed:")
     print(final_message)
@@ -45,23 +43,22 @@ def human_response_display_node(state: GraphState) -> GraphState:
 # --- 3. Build the Graph ---
 builder = StateGraph(GraphState)
 builder.add_node("greeting", greeting_node)
-builder.add_node("human_input_interrupt", human_input_node)
+builder.add_node("human_input_interrupt", human_input_node) # Node name remains for clarity
 builder.add_node("human_response_display", human_response_display_node)
 
 builder.set_entry_point("greeting")
-builder.add_edge("greeting", "human_input_interrupt")
+builder.add_edge("greeting", "human_input_interrupt") # Still connects normally
 builder.add_edge("human_response_display", END)
 
 # --- Checkpointer and Graph Compilation ---
-#SQLITE_DB_PATH = "langgraph_checkpoints.sqlite"
-#global_memory_saver = SqliteSaver.from_conn_string(SQLITE_DB_PATH)
+# Using MemorySaver as requested
 global_memory_saver = MemorySaver()
 
 global_graph = builder.compile(
     checkpointer=global_memory_saver,
-    # Keep interrupt_after for human_input_interrupt.
-    # This means human_input_node runs, then graph interrupts.
-    interrupt_after=["human_input_interrupt"]
+    # CRITICAL CHANGE: Interrupt BEFORE human_input_interrupt
+    # This means the graph will pause *before* executing the human_input_node
+    interrupt_before=["human_input_interrupt"]
 )
 
 # --- Gradio UI Logic ---
@@ -71,25 +68,23 @@ def start_graph(thread_id_state):
     new_thread_id = str(uuid.uuid4())
     print(f"\n--- Starting New Graph Execution with thread_id: {new_thread_id} ---")
     
-    # You might want to explicitly drop previous thread state if it exists for this ID
-    # (Uncomment if you want to explicitly clear old data for a potentially re-used UUID,
-    # though with uuid4 it's highly unlikely)
+    # In-memory saver does not need explicit drop for new UUIDs, but good for resetting if reusing IDs
     # try:
     #     global_memory_saver.drop(config={"configurable": {"thread_id": new_thread_id}})
     # except Exception as e:
     #     print(f"Could not drop previous state for {new_thread_id}: {e}")
 
     try:
-        # Pass an empty input for the first run.
-        # The graph will run greeting_node, then human_input_node, then interrupt.
-        for s in global_graph.stream({"greeting_message": ""}, {"configurable": {"thread_id": new_thread_id}}):
+        # Initial run: Graph executes greeting_node, then interrupts BEFORE human_input_node.
+        # We pass None as input, as the state is managed by the graph's execution up to the breakpoint.
+        for s in global_graph.stream(None, {"configurable": {"thread_id": new_thread_id}}):
             if "__end__" in s:
                 break
             elif "__interrupt__" in s:
-                print(f"Graph interrupted at {s.get('__interrupt__', 'Unknown')}")
+                print(f"Graph interrupted BEFORE {s.get('__interrupt__', 'Unknown')} node.")
                 break
             else:
-                pass
+                pass # You can print s here to see node outputs during initial run
         
         current_state_snapshot = global_graph.get_state({"configurable": {"thread_id": new_thread_id}})
         
@@ -108,13 +103,32 @@ def resume_graph(human_input_from_ui: str, thread_id_state):
     print(f"Human input received from UI: {human_input_from_ui}")
 
     try:
-        # Pass the human_input_from_ui directly to the graph stream.
-        # This input will be merged into the graph's state when it resumes.
-        for s in global_graph.stream({"human_input": human_input_from_ui}, {"configurable": {"thread_id": thread_id_state}}):
+        # --- CRITICAL CHANGES FOR RESUMPTION WITH INTERRUPT_BEFORE ---
+        # 1. Get the current state from the checkpointer.
+        current_state_snapshot = global_graph.get_state({"configurable": {"thread_id": thread_id_state}})
+        
+        # 2. Extract the current values dictionary from the snapshot.
+        current_state_values = current_state_snapshot.values
+        
+        # 3. Manually update the 'human_input' field in this dictionary.
+        # The operator.add on Annotated[Optional[str], operator.add] will implicitly
+        # handle combining if it was a list or similar, but for a simple string
+        # it typically means setting the value.
+        current_state_values["human_input"] = human_input_from_ui
+
+        # 4. Save this modified state back to the checkpointer.
+        # This is how the human's input is "injected" into the graph's memory.
+        global_memory_saver.put_state(current_state_values, {"configurable": {"thread_id": thread_id_state}})
+
+        # 5. Resume the graph with None as input.
+        # The graph will now load the state (which includes the injected human_input)
+        # and proceed to execute the 'human_input_interrupt' node (which it was paused before)
+        # and subsequent nodes.
+        for s in global_graph.stream(None, {"configurable": {"thread_id": thread_id_state}}):
             if "__end__" in s:
                 break
             else:
-                pass
+                pass # You can print s here to see node outputs during resumption
 
         final_state_snapshot = global_graph.get_state({"configurable": {"thread_id": thread_id_state}})
         final_state_values = final_state_snapshot.values
@@ -129,11 +143,10 @@ def resume_graph(human_input_from_ui: str, thread_id_state):
 
 # Gradio Interface setup
 with gr.Blocks() as demo:
-    gr.Markdown("# LangGraph Human Approval Demo (Persistent on Hugging Face)")
-    gr.Markdown(f"Graph state will be saved persistently in MEMORY.")
+    gr.Markdown("# LangGraph Human-in-the-Loop Demo (MemorySaver)")
+    gr.Markdown("Graph state will be saved persistently in **memory** for the current session. Click 'Start Conversation' to begin. Input is taken via Gradio UI.")
 
     output_textbox = gr.Textbox(label="AI Assistant Output", lines=5, interactive=False)
-    # The human input now comes ONLY from this textbox
     human_input_textbox = gr.Textbox(label="Your Input", placeholder="Type your response here...", interactive=False)
     
     thread_id_state = gr.State("")