import random
from typing import Literal
from langchain_core.messages import AIMessage
from langgraph.graph import START, MessagesState, StateGraph
from langgraph.types import Command
from langgraph_agent_toolkit.agents.agent import Agent
[docs]
class AgentState(MessagesState, total=False):
pass
[docs]
def node_a(state: AgentState) -> Command[Literal["node_b", "node_c"]]:
print("Called A")
value = random.choice(["a", "b"])
# this is a replacement for a conditional edge function
if value == "a":
goto = "node_b"
else:
goto = "node_c"
# note how Command allows you to BOTH update the graph state AND route to the next node
return Command(
# this is the state update
update={"messages": [AIMessage(content=f"Hello {value}")]},
# this is a replacement for an edge
goto=goto,
)
[docs]
def node_b(state: AgentState):
print("Called B")
return {"messages": [AIMessage(content="Hello B")]}
[docs]
def node_c(state: AgentState):
print("Called C")
return {"messages": [AIMessage(content="Hello C")]}
builder = StateGraph(AgentState)
builder.add_edge(START, "node_a")
builder.add_node(node_a)
builder.add_node(node_b)
builder.add_node(node_c)
# NOTE: there are no edges between nodes A, B and C!
command_agent = Agent(
name="command-agent",
description="A command agent.",
graph=builder.compile(),
)