from langgraph.graph import StateGraph, MessagesState, START, END
from typing import Literal
from langchain_core.messages import HumanMessage
class WorkflowState(MessagesState):
current_step: str = "analyze"
analysis_result: str = ""
plan: str = ""
def analyze_node(state: WorkflowState):
# Analysis logic
analysis = "Analysis complete"
return {
"current_step": "plan",
"analysis_result": analysis,
"messages": [HumanMessage(content=f"Analysis: {analysis}")]
}
def plan_node(state: WorkflowState):
# Planning logic
plan = f"Plan based on {state['analysis_result']}"
return {
"current_step": "execute",
"plan": plan,
"messages": [HumanMessage(content=f"Plan: {plan}")]
}
def execute_node(state: WorkflowState):
# Execution logic
result = f"Executed plan: {state['plan']}"
return {
"current_step": "done",
"messages": [HumanMessage(content=f"Result: {result}")]
}
def route_step(state: WorkflowState) -> Literal["plan", "execute", "end"]:
step = state.get("current_step", "analyze")
if step == "analyze":
return "plan"
elif step == "plan":
return "execute"
else:
return "end"
# Build workflow
workflow = StateGraph(WorkflowState)
workflow.add_node("analyze", analyze_node)
workflow.add_node("plan", plan_node)
workflow.add_node("execute", execute_node)
workflow.add_edge(START, "analyze")
workflow.add_conditional_edges("analyze", route_step)
workflow.add_conditional_edges("plan", route_step)
workflow.add_conditional_edges("execute", route_step)
agent = LangGraphAgent(
name="WorkflowAgent",
description="Multi-step workflow agent",
graph=workflow.compile()
)