116 lines
3.1 KiB
Python
116 lines
3.1 KiB
Python
"""
|
|
Basic LangGraph Example - A simple conversational agent with state management.
|
|
|
|
This demonstrates core LangGraph concepts:
|
|
- State schemas
|
|
- Nodes (functions that process state)
|
|
- Edges (connections between nodes)
|
|
- Conditional routing
|
|
"""
|
|
|
|
from typing import TypedDict, Literal, Annotated
|
|
from langgraph.graph import StateGraph, START, END
|
|
from langgraph.graph.message import add_messages
|
|
|
|
|
|
# Define the state schema
|
|
class State(TypedDict):
|
|
"""State represents what flows through our graph."""
|
|
messages: Annotated[list, add_messages]
|
|
counter: int
|
|
|
|
|
|
# Define node functions
|
|
def greeting_node(state: State) -> dict:
|
|
"""Process greeting messages."""
|
|
print("Processing greeting...")
|
|
return {"counter": state["counter"] + 1}
|
|
|
|
|
|
def math_node(state: State) -> dict:
|
|
"""Process math-related messages."""
|
|
print("Processing math request...")
|
|
return {"counter": state["counter"] + 1}
|
|
|
|
|
|
def default_node(state: State) -> dict:
|
|
"""Handle unknown/unrecognized messages."""
|
|
print("Handling general message...")
|
|
return {"counter": state["counter"] + 1}
|
|
|
|
|
|
# Define the router function
|
|
def route_message(state: State) -> Literal["greeting", "math", "default"]:
|
|
"""Route to appropriate node based on message content."""
|
|
messages = state["messages"]
|
|
if not messages:
|
|
return "default"
|
|
|
|
last_message = messages[-1].lower() if isinstance(messages[-1], str) else str(messages[-1]).lower()
|
|
|
|
if any(word in last_message for word in ["hello", "hi", "hey", "greetings"]):
|
|
return "greeting"
|
|
elif any(word in last_message for word in ["add", "subtract", "multiply", "divide", "calculate"]):
|
|
return "math"
|
|
else:
|
|
return "default"
|
|
|
|
|
|
# Build the graph
|
|
def create_graph():
|
|
"""Create and compile the LangGraph workflow."""
|
|
|
|
# Initialize the graph builder
|
|
builder = StateGraph(State)
|
|
|
|
# Add nodes
|
|
builder.add_node("greeting", greeting_node)
|
|
builder.add_node("math", math_node)
|
|
builder.add_node("default", default_node)
|
|
|
|
# Add edges with conditional routing
|
|
builder.add_conditional_edges(
|
|
START,
|
|
route_message,
|
|
{
|
|
"greeting": "greeting",
|
|
"math": "math",
|
|
"default": "default"
|
|
}
|
|
)
|
|
|
|
# All nodes lead to END (or could chain to other nodes)
|
|
builder.add_edge("greeting", END)
|
|
builder.add_edge("math", END)
|
|
builder.add_edge("default", END)
|
|
|
|
# Compile the graph
|
|
return builder.compile()
|
|
|
|
|
|
def main():
|
|
"""Run the basic LangGraph agent."""
|
|
graph = create_graph()
|
|
|
|
print("=== Basic LangGraph Agent ===")
|
|
print("Try messages like: 'hello', 'add 5 and 3', or anything else")
|
|
print()
|
|
|
|
# Test cases
|
|
test_messages = [
|
|
"Hello there!",
|
|
"Can you add 5 and 3?",
|
|
"What's the weather like?"
|
|
]
|
|
|
|
for message in test_messages:
|
|
print(f"Input: {message}")
|
|
initial_state = {"messages": [message], "counter": 0}
|
|
result = graph.invoke(initial_state)
|
|
print(f"Output: Processed with counter = {result['counter']}")
|
|
print()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|