Update app.py

app.py

@@ -1,17 +1,19 @@
 import gradio as gr
 from langgraph.graph import StateGraph, END
-from langgraph.checkpoint.memory import MemorySaver
+from langgraph.checkpoint.sqlite import SqliteSaver # Changed import
 import operator
 from typing import TypedDict, Annotated, Optional
+import uuid
+import os # Import os for path handling
 
 # --- 1. Define your Graph State (Same as before) ---
 class GraphState(TypedDict):
     greeting_message: str
-    human_input: Annotated[Optional[str], operator.add] # Make human_input optional for initial state
+    human_input: Annotated[Optional[str], operator.add]
     final_response: str
-    current_node: str # To track which node was last executed/interrupted at
+    current_node: str
 
-# --- 2. Define the Nodes (Slightly modified for Gradio interaction) ---
+# --- 2. Define the Nodes (Same as before) ---
 def greeting_node(state: GraphState) -> GraphState:
     greeting = "Hello there! I'm an AI assistant. How can I help you today?"
     print("🤖 Greeting Node Executed:")
@@ -19,12 +21,9 @@ def greeting_node(state: GraphState) -> GraphState:
     return {"greeting_message": greeting, "current_node": "greeting"}
 
 def human_input_node(state: GraphState) -> GraphState:
-    # In a Gradio context, the actual `input()` will come from the UI
-    # This node primarily serves as the interruption point and where the UI input
-    # will be injected into the state when resuming.
     print("\n✋ Human Input Node Executed (Graph was resumed with input):")
-    # We won't call input() here directly, Gradio will provide it
-    return {"current_node": "human_input_interrupt"}
+    user_response = input("Please enter your response (e.g., 'I want to know about LangChain'): ")
+    return {"human_input": user_response, "current_node": "human_input_interrupt"}
 
 def human_response_display_node(state: GraphState) -> GraphState:
     human_response = state.get("human_input", "No human input received.")
@@ -43,126 +42,85 @@ builder.set_entry_point("greeting")
 builder.add_edge("greeting", "human_input_interrupt")
 builder.add_edge("human_response_display", END)
 
-# --- 4. Configure the interruption and checkpointer ---
-memory = MemorySaver()
-thread_id = "user_conversation_gradio" # A fixed thread ID for this demo
+# --- Checkpointer and Graph Compilation ---
+# Define a persistent SQLite database file
+# It's good practice to place it in a dedicated directory or the project root.
+SQLITE_DB_PATH = "langgraph_checkpoints.sqlite" # This file will be created/used
 
-# Compile the graph with checkpointer and interrupt_after for human_input_interrupt
-# This means the human_input_node will run, take input, then interrupt.
-graph = builder.compile(
-    checkpointer=memory,
+# Initialize SqliteSaver with the file path
+global_memory_saver = SqliteSaver.from_conn_string(SQLITE_DB_PATH)
+
+# Global graph instance compiled once
+global_graph = builder.compile(
+    checkpointer=global_memory_saver,
     interrupt_after=["human_input_interrupt"]
 )
 
-# --- 5. Gradio UI Logic ---
+# --- Gradio UI Logic (Mostly the same, but now truly persistent) ---
 
-# Global variable to store the current thread state, for simplicity in this demo
-# In a real app, you'd manage this for multiple users/sessions
-current_thread_state = {}
+current_thread_id = gr.State("")
 
-def start_graph():
-    global current_thread_state
+def start_graph(thread_id_state):
+    # Generate a new unique thread ID for each new conversation
+    new_thread_id = str(uuid.uuid4())
+    print(f"\n--- Starting New Graph Execution with thread_id: {new_thread_id} ---")
     
-    # Clear any previous state for a fresh start
-    memory.drop(config={"configurable": {"thread_id": thread_id}})
-    current_thread_state = {}
+    # You might want to explicitly drop previous thread state if it exists for this ID
+    # if global_memory_saver.get_state({"configurable": {"thread_id": new_thread_id}}):
+    #     global_memory_saver.drop(config={"configurable": {"thread_id": new_thread_id}})
 
-    print("\n--- Starting New Graph Execution ---")
-    
-    # Run the graph until the first interruption point
-    # We pass an initial empty state; the greeting_node will populate it
     try:
-        # Stream the graph, it will run 'greeting' then 'human_input_interrupt'
-        # and interrupt *after* 'human_input_interrupt'
-        for s in graph.stream({"greeting_message": ""}, {"configurable": {"thread_id": thread_id}}):
-            # The 's' here represents the output of each node.
-            # When an interrupt occurs, the last yielded item will be '__interrupt__':()
+        for s in global_graph.stream({"greeting_message": ""}, {"configurable": {"thread_id": new_thread_id}}):
             if "__end__" in s:
-                # Graph finished immediately, which means no interruption happened or it was very quick
                 break
             elif "__interrupt__" in s:
                 print(f"Graph interrupted at {s.get('__interrupt__', 'Unknown')}")
-                break # Break the loop as soon as an interruption is detected
+                break
             else:
-                # Process regular node output, if any before interrupt
                 pass
         
-        # After the stream, get the current state from the checkpointer
-        current_state_snapshot = graph.get_state({"configurable": {"thread_id": thread_id}})
-        current_thread_state = current_state_snapshot.values
+        current_state_snapshot = global_graph.get_state({"configurable": {"thread_id": new_thread_id}})
         
-        # Display the greeting message from the state
-        output_message = current_thread_state.get("greeting_message", "No greeting yet.")
-        output_message += "\n" + current_thread_state.get("human_input", "Please provide your input to continue.")
-        
-        # Enable the input box and resume button, disable start
-        return output_message, gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=False)
+        output_message = current_state_snapshot.values.get("greeting_message", "No greeting yet.")
+        output_message += "\n" + "Please provide your input to continue."
 
-    except Exception as e:
-        return f"An error occurred during graph start: {e}", gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=True)
+        return (output_message, gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=False), new_thread_id)
 
-def resume_graph(human_input_from_ui: str):
-    global current_thread_state
+    except Exception as e:
+        return (f"An error occurred during graph start: {e}",
+                gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=True), thread_id_state)
 
-    print("\n--- Resuming Graph Execution ---")
-    print(f"Received human input from UI: {human_input_from_ui}")
+def resume_graph(human_input_from_ui: str, thread_id_state):
+    print(f"\n--- Resuming Graph Execution for thread_id: {thread_id_state} ---")
+    print(f"Received human input from UI (will be ignored if input taken in node): {human_input_from_ui}")
 
     try:
-        # Resume the graph with the human input from Gradio
-        # The graph will pick up from where it left off (after human_input_interrupt)
-        # The human_input_node itself already captured the input, so this input is just
-        # confirming/triggering the next step, but if using interrupt_before, this is where
-        # you'd pass the actual input. For interrupt_after, the input is already in state.
-        
-        # However, for consistency and to ensure flow, we re-run and pass the input
-        # which will effectively update the 'human_input' again (harmless due to operator.add)
-        # or simply trigger the next steps if human_input_interrupt already ran.
-        
-        # We need to explicitly re-provide the human_input from the UI
-        # in case the human_input_node itself wasn't designed to wait for input *after* its execution.
-        # Given our current human_input_node structure (it calls input()),
-        # we'll let it call input(), and this resume just triggers the graph forward.
-
-        # If human_input_node was only a marker, you'd pass:
-        # for s in graph.stream({"human_input": human_input_from_ui}, {"configurable": {"thread_id": thread_id}}):
-
-        # But since human_input_node has the input() call, let's just resume
-        # to process the *next* step assuming human_input_node *already* updated the state
-        # correctly during the initial interrupted run.
-        # If the input was taken inside the node that interrupts *after*, the state is already there.
-        # So we just trigger continuation.
-        
-        # However, a cleaner way for interrupt_after is often:
-        # Node takes input -> Node updates state -> Interrupt -> Resume (no new input needed)
-        # If input was taken *outside* the node (interrupt_before), then resume provides it.
-
-        # Given the previous successful run with interrupt_after, the human_input_node
-        # already captured the input. We just need to drive the graph forward.
-        # LangGraph automatically loads the state from the checkpointer based on thread_id
-        for s in graph.stream(None, {"configurable": {"thread_id": thread_id}}): # Pass None as input to just continue
+        # Pass None as input as the human_input_node already handled taking input
+        for s in global_graph.stream(None, {"configurable": {"thread_id": thread_id_state}}):
             if "__end__" in s:
                 break
             else:
-                pass # Process intermediate outputs if needed
+                pass
 
-        # Get the final state
-        final_state_snapshot = graph.get_state({"configurable": {"thread_id": thread_id}})
-        current_thread_state = final_state_snapshot.values
+        final_state_snapshot = global_graph.get_state({"configurable": {"thread_id": thread_id_state}})
+        final_state_values = final_state_snapshot.values
 
-        final_message = current_thread_state.get("final_response", "Graph finished.")
-        # Disable input and resume, enable start for a new conversation
-        return final_message, gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=True)
+        final_message = final_state_values.get("final_response", "Graph finished without final response.")
+        return (final_message, gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=True), thread_id_state)
 
     except Exception as e:
-        return f"An error occurred during graph resumption: {e}", gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=True)
+        return (f"An error occurred during graph resumption: {e}",
+                gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=True), thread_id_state)
 
 # Gradio Interface setup
 with gr.Blocks() as demo:
-    gr.Markdown("# LangGraph Human Approval Demo")
-    gr.Markdown("Click 'Start Conversation' to begin. The graph will interrupt for your input.")
+    gr.Markdown("# LangGraph Human Approval Demo (Persistent)")
+    gr.Markdown(f"Graph state will be saved persistently in `{SQLITE_DB_PATH}`. Click 'Start Conversation' to begin.")
 
     output_textbox = gr.Textbox(label="AI Assistant Output", lines=5, interactive=False)
-    human_input_textbox = gr.Textbox(label="Your Input", placeholder="Type your response here...", interactive=False)
+    human_input_textbox = gr.Textbox(label="Your Input (Please enter in console as well)", placeholder="Type your response here...", interactive=True)
+    
+    thread_id_state = gr.State("")
 
     with gr.Row():
         start_button = gr.Button("Start Conversation")
@@ -170,14 +128,14 @@ with gr.Blocks() as demo:
 
     start_button.click(
         start_graph,
-        inputs=[],
-        outputs=[output_textbox, human_input_textbox, resume_button, start_button]
+        inputs=[thread_id_state],
+        outputs=[output_textbox, human_input_textbox, resume_button, start_button, thread_id_state]
     )
 
     resume_button.click(
         resume_graph,
-        inputs=[human_input_textbox],
-        outputs=[output_textbox, human_input_textbox, resume_button, start_button]
+        inputs=[human_input_textbox, thread_id_state],
+        outputs=[output_textbox, human_input_textbox, resume_button, start_button, thread_id_state]
     )
 
 if __name__ == "__main__":