Update app.py

app.py

@@ -1,34 +1,53 @@
-from langgraph.graph import Graph
-from typing import TypedDict
+from langgraph.graph import StateGraph
+from typing import TypedDict, Annotated
+from langgraph.graph.message import add_messages
 
-# Define the state type
 class State(TypedDict):
-    messages: list[str]
+    messages: Annotated[list[str], add_messages]
     current_step: str
 
-# Create nodes (functions that represent different states)
-def collect_info(state: State) -> State:
+def collect_info(state: State) -> dict:
+    print("\n--> In collect_info")
+    print(f"Messages before: {state['messages']}")
+    
+    messages = state["messages"] + ["Information collected"]
+    print(f"Messages after: {messages}")
+    
     return {
-        **state,
-        "messages": state["messages"] + ["Information collected"],
+        "messages": messages,
         "current_step": "process"
     }
 
-def process_info(state: State) -> State:
+def process_info(state: State) -> dict:
+    print("\n--> In process_info")
+    print(f"Messages before: {state['messages']}")
+    
+    messages = state["messages"] + ["Information processed"]
+    print(f"Messages after: {messages}")
+    
     return {
-        **state,
-        "messages": state["messages"] + ["Information processed"],
+        "messages": messages,
         "current_step": "end"
     }
 
-# Create the graph
-workflow = Graph()
+# Create and setup graph
+workflow = StateGraph(State)
 
-# Add nodes and edges
+# Add nodes
 workflow.add_node("collect", collect_info)
 workflow.add_node("process", process_info)
+
+# Add edges
 workflow.add_edge("collect", "process")
 
-# Set the entry point and compile
+# Set entry and finish points
 workflow.set_entry_point("collect")
-app = workflow.compile()
+workflow.set_finish_point("process")
+
+app = workflow.compile()
+
+# Run workflow
+print("\nStarting workflow...")
+initial_state = State(messages=["Starting"], current_step="collect")
+final_state = app.invoke(initial_state)
+print(f"\nFinal messages: {final_state['messages']}")