app.py

from langgraph.graph import StateGraph
from typing import TypedDict, Annotated
from langgraph.graph.message import add_messages

class State(TypedDict):
    messages: Annotated[list[str], add_messages]
    current_step: str

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 {
        "messages": messages,
        "current_step": "process"
    }

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 {
        "messages": messages,
        "current_step": "end"
    }

# Create and setup graph
workflow = StateGraph(State)

# Add nodes
workflow.add_node("collect", collect_info)
workflow.add_node("process", process_info)

# Add edges
workflow.add_edge("collect", "process")

# Set entry and finish points
workflow.set_entry_point("collect")

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']}")