langgraph-basics.py

# Import necessary libraries
from langgraph.graph import StateGraph, END
from typing import Dict, TypedDict, Optional, Literal, List, Union

# Define Graph State
class GraphState(TypedDict):
    init_input: Optional[str] = None
    fruit: Optional[str] = None
    final_result: Optional[str] = None
    user_confirmation: Optional[str] = None

# Define Graph Nodes (State Functions)
def input_fruit(state: GraphState) -> Dict[str, str]:
    init_input = state.get("init_input", "").strip().lower()
    if init_input not in ["apple", "banana", "cherry"]:
        return {"fruit": "error"}
    return {"fruit": init_input}

def confirm_fruit(state: GraphState) -> Dict[str, str]:
    return {"final_result": f"You selected {state['fruit']}, which is a valid fruit."}

def error(state: GraphState) -> Dict[str, str]:
    return {"final_result": "invalid fruit. valid fruits are: apple, banana, cherry"}

def review_fruit(state: GraphState) -> Dict[str, str]:
    # For demonstration purposes, let's assume the user always confirms
    user_confirmation = "yes"
    return {"user_confirmation": user_confirmation}

# Define Functions to Determine Next Node (Edge Functions)
def continue_next(state: GraphState) -> Literal["to_review_fruit", "to_error"]:
    if state.get("fruit") != "error":
        return "to_review_fruit"
    return "to_error"

def review_decision(state: GraphState) -> Literal["to_confirm_fruit", "to_error"]:
    if state.get("user_confirmation") == "yes":
        return "to_confirm_fruit"
    return "to_error"

# Define Workflow
workflow = StateGraph(GraphState)
workflow.add_node("input_fruit", input_fruit)
workflow.add_node("review_fruit", review_fruit)
workflow.add_node("confirm_fruit", confirm_fruit)
workflow.add_node("error", error)
workflow.set_entry_point("input_fruit")
workflow.add_edge("confirm_fruit", END)
workflow.add_edge("error", END)
workflow.add_conditional_edges("input_fruit", continue_next, {"to_review_fruit": "review_fruit", "to_error": "error"})
workflow.add_conditional_edges("review_fruit", review_decision, {"to_confirm_fruit": "confirm_fruit", "to_error": "error"})
app = workflow.compile()

# Test with a valid fruit and user confirmation
result = app.invoke({"init_input": "apple", "user_confirmation": "yes"})
print("Result:")
print(result)

# Test with a valid fruit and user rejection
result = app.invoke({"init_input": "apple", "user_confirmation": "no"})
print("Result:")
print(result)

# Test with an invalid fruit
result = app.invoke({"init_input": "mango", "user_confirmation": "no"})
print("Result:")
print(result)

https://www.kaggle.com/code/marcinrutecki/langgraph-simplified/notebook