Build complex, stateful AI workflows with graph-based orchestration
Build complex, stateful AI workflows with graph-based orchestration.The StateGraph class is a powerful graph-based workflow engine that enables you to build complex, stateful AI applications with explicit control flow. Instead of writing monolithic functions or brittle chains, you define workflows as graphs of nodes that can branch, loop, persist state, and recover from failures.
StateGraph can be created with minimal configuration or with extensive customization to suit your specific needs. The graph provides a robust foundation for AI-powered applications with built-in support for various advanced features.
from typing_extensions import TypedDictfrom upsonic.graphv2 import StateGraph, START, ENDclass ConversationState(TypedDict): messages: list[str] response: strdef process_node(state: ConversationState) -> dict: """Process the user's message.""" if not state["messages"]: return {"response": "No messages to process"} user_message = state["messages"][-1] response = f"Echo: {user_message}" return {"response": response}# Create and build the graphbuilder = StateGraph(ConversationState)builder.add_node("process", process_node)builder.add_edge(START, "process")builder.add_edge("process", END)graph = builder.compile()# Execute the graphresult = graph.invoke({ "messages": ["Hello!"], "response": ""})print(result["response"]) # Output: Echo: Hello!
When creating a graph without specifying a checkpointer, state is not persisted between executions. Use MemorySaver or SqliteCheckpointer for persistence.
from typing import Annotated, Listfrom typing_extensions import TypedDictimport operatorfrom upsonic.graphv2 import StateGraph, START, ENDfrom upsonic.models import infer_modelfrom upsonic.messages import ModelRequest, UserPromptPart, SystemPromptPart, TextPartfrom upsonic.tools import toolclass AgentState(TypedDict): messages: Annotated[List, operator.add] result: str@tooldef calculator(a: float, b: float, operation: str) -> float: """Perform basic math operations. Use this tool to calculate mathematical expressions. Args: a: First number b: Second number operation: Operation to perform - either "add" or "multiply" Returns: The result of the calculation """ if operation == "add": return a + b elif operation == "multiply": return a * b return 0def agent_node(state: AgentState) -> dict: """Agent node with tool access.""" model = infer_model("anthropic/claude-sonnet-4-6") model_with_tools = model.bind_tools([calculator]) user_message = state["messages"][-1] if state["messages"] else "Hello" request = ModelRequest(parts=[ SystemPromptPart(content="You are a helpful assistant."), UserPromptPart(content=str(user_message)) ]) response = model_with_tools.invoke([request]) # Extract text content from response text_content = "" for part in response.parts: if isinstance(part, TextPart): text_content += part.content # If no text content, use the string representation if not text_content: text_content = str(response) return {"messages": [response], "result": text_content}# Build graphbuilder = StateGraph(AgentState)builder.add_node("agent", agent_node)builder.add_edge(START, "agent")builder.add_edge("agent", END)graph = builder.compile()# Executeresult = graph.invoke({ "messages": ["What is 23 multiplied by 17?"], "result": ""})print(result["result"])
