From debd9a8d9c488351019b215c766197a9e2e3edd0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 22 Feb 2024 15:19:03 -0800 Subject: [PATCH] Add support for conditional entry points in Graph, StateGraph, MessageGraph --- langgraph/graph/graph.py | 90 +++++-- langgraph/graph/state.py | 20 +- tests/__snapshots__/test_pregel.ambr | 336 +++++++++++++++++++++++++++ tests/test_pregel.py | 96 ++++++++ tests/test_pregel_async.py | 86 +++++++ 5 files changed, 601 insertions(+), 27 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 16e3551f5..4885f3b6c 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -17,6 +17,7 @@ from langgraph.pregel import Channel, Pregel logger = logging.getLogger(__name__) +START = "__start__" END = "__end__" @@ -40,6 +41,8 @@ class Graph: self.branches: defaultdict[str, list[Branch]] = defaultdict(list) self.support_multiple_edges = False self.compiled = False + self.entry_point: Optional[str] = None + self.entry_point_branch: Optional[Branch] = None def add_node(self, key: str, action: RunnableLike) -> None: if self.compiled: @@ -111,6 +114,29 @@ class Graph: raise ValueError(f"Need to add_node `{key}` first") self.entry_point = key + def set_conditional_entry_point( + self, + condition: Callable[..., str], + conditional_edge_mapping: Optional[Dict[str, str]] = None, + ) -> None: + if self.compiled: + logger.warning( + "Setting the entry point of a graph that has already been compiled. " + "This will not be reflected in the compiled graph." + ) + if iscoroutinefunction(condition): + raise ValueError("Condition cannot be a coroutine function") + if conditional_edge_mapping and set( + conditional_edge_mapping.values() + ).difference([END]).difference(self.nodes): + raise ValueError( + f"Missing nodes which are in conditional edge mapping. Mapping " + f"contains possible destinations: " + f"{list(conditional_edge_mapping.values())}. Possible nodes are " + f"{list(self.nodes.keys())}." + ) + self.entry_point_branch = Branch(condition, conditional_edge_mapping) + def set_finish_point(self, key: str) -> None: return self.add_edge(key, END) @@ -120,21 +146,20 @@ class Graph: if node not in all_starts: raise ValueError(f"Node `{node}` is a dead-end") - if all( - branch.ends is not None - for branch_list in self.branches.values() - for branch in branch_list - ): - all_ends = ( - {end for _, end in self.edges} - | { - end - for branch_list in self.branches.values() - for branch in branch_list - for end in branch.ends.values() - } - | {self.entry_point} - ) + branches = [ + branch for branch_list in self.branches.values() for branch in branch_list + ] + if self.entry_point_branch is not None: + branches.append(self.entry_point_branch) + + all_hard_ends = {end for _, end in self.edges} + if self.entry_point is not None: + all_hard_ends.add(self.entry_point) + + if all(branch.ends is not None for branch in branches): + all_ends = all_hard_ends | { + end for branch in branches for end in branch.ends.values() + } for node in self.nodes: if node not in all_ends: @@ -179,10 +204,19 @@ class Graph: branch.runnable, name=f"{key}_condition" ) + if self.entry_point_branch: + nodes[f"{START}:edges"] = Channel.subscribe_to( + START, tags=["langsmith:hidden"] + ) | RunnableLambda( + self.entry_point_branch.runnable, name=f"{START}_condition" + ) + elif self.entry_point is None: + raise ValueError("No entry point set") + return CompiledGraph( graph=self, nodes=nodes, - input=f"{self.entry_point}:inbox", + input=f"{self.entry_point}:inbox" if self.entry_point else START, output=END, hidden=[f"{node}:inbox" for node in self.nodes], checkpointer=checkpointer, @@ -198,7 +232,7 @@ class CompiledGraph(Pregel): def get_graph(self, config: Optional[RunnableConfig] = None) -> RunnableGraph: graph = RunnableGraph() - graph.add_node(self.get_input_schema(config), "__start__") + graph.add_node(self.get_input_schema(config), START) graph.add_node(self.get_output_schema(config), END) for key, node in self.graph.nodes.items(): @@ -215,8 +249,26 @@ class CompiledGraph(Pregel): name, ) graph.add_edge(graph.nodes[start], graph.nodes[name]) - for label, end in branch.ends.items(): + ends = branch.ends or {k: k for k in self.graph.nodes} + for label, end in ends.items(): graph.add_edge(graph.nodes[name], graph.nodes[end], label) - graph.add_edge(graph.nodes["__start__"], graph.nodes[self.graph.entry_point]) + if self.graph.entry_point_branch: + graph.add_node( + RunnableLambda( + self.graph.entry_point_branch.runnable, + name=self.graph.entry_point_branch.condition.__name__, + ), + f"{START}_condition", + ) + graph.add_edge(graph.nodes[START], graph.nodes[f"{START}_condition"]) + ends = self.graph.entry_point_branch.ends or { + k: k for k in self.graph.nodes + } + for label, end in ends.items(): + graph.add_edge( + graph.nodes[f"{START}_condition"], graph.nodes[end], label + ) + elif self.graph.entry_point: + graph.add_edge(graph.nodes[START], graph.nodes[self.graph.entry_point]) return graph diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 2804990b6..75d5f3d1f 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -12,13 +12,11 @@ from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.checkpoint import BaseCheckpointSaver -from langgraph.graph.graph import END, CompiledGraph, Graph +from langgraph.graph.graph import END, START, CompiledGraph, Graph from langgraph.pregel import Channel from langgraph.pregel.read import ChannelRead from langgraph.pregel.write import SKIP_WRITE, ChannelWrite -START = "__start__" - class StateGraph(Graph): def __init__(self, schema: Type[Any]) -> None: @@ -105,11 +103,17 @@ class StateGraph(Graph): nodes[START] = Channel.subscribe_to( f"{START}:inbox", tags=["langsmith:hidden"] ) | ChannelWrite(channels=[(START, None, False)] + update_channels) - nodes[f"{START}:edges"] = ( - Channel.subscribe_to(START, tags=["langsmith:hidden"]) - | ChannelRead(state_keys_read) - | Channel.write_to(f"{self.entry_point}:inbox") - ) + nodes[f"{START}:edges"] = Channel.subscribe_to( + START, tags=["langsmith:hidden"] + ) | ChannelRead(state_keys_read) + if self.entry_point: + nodes[f"{START}:edges"] |= Channel.write_to(f"{self.entry_point}:inbox") + elif self.entry_point_branch: + nodes[f"{START}:edges"] |= RunnableLambda( + self.entry_point_branch.runnable, name=f"{START}_condition" + ) + else: + raise ValueError("No entry point set") return CompiledGraph( graph=self, diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index 330b87bd0..412b4ad9d 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -1,4 +1,340 @@ # serializer version: 1 +# name: test_conditional_entrypoint_graph + '{"title": "LangGraphInput"}' +# --- +# name: test_conditional_entrypoint_graph.1 + '{"title": "LangGraphOutput"}' +# --- +# name: test_conditional_entrypoint_graph.2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": { + "title": "LangGraphInput" + } + }, + { + "id": "__end__", + "type": "schema", + "data": { + "title": "LangGraphOutput" + } + }, + { + "id": "left", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "left" + } + }, + { + "id": "right", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "right" + } + }, + { + "id": "left_", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "" + } + }, + { + "id": "__start___condition", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "should_start" + } + } + ], + "edges": [ + { + "source": "right", + "target": "__end__" + }, + { + "source": "left", + "target": "left_" + }, + { + "source": "left_", + "target": "left", + "data": "left" + }, + { + "source": "left_", + "target": "right", + "data": "right" + }, + { + "source": "__start__", + "target": "__start___condition" + }, + { + "source": "__start___condition", + "target": "left", + "data": "go-left" + }, + { + "source": "__start___condition", + "target": "right", + "data": "go-right" + } + ] + } + ''' +# --- +# name: test_conditional_entrypoint_graph.3 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------------+ + | __start___condition | + +---------------------+ + *** *** + * * + ** *** + +------+ * + | left | * + +------+ * + * * + * * + * * + +---------------+ * + | left_ | *** + +---------------+ * + *** *** + * * + ** ** + +-------+ + | right | + +-------+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_conditional_entrypoint_graph_state + '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "output": {"title": "Output", "type": "string"}}}}}' +# --- +# name: test_conditional_entrypoint_graph_state.1 + '{"title": "LangGraphOutput", "$ref": "#/definitions/AgentState", "definitions": {"AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "output": {"title": "Output", "type": "string"}}}}}' +# --- +# name: test_conditional_entrypoint_graph_state.2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": { + "title": "LangGraphInput", + "$ref": "#/definitions/AgentState", + "definitions": { + "AgentState": { + "title": "AgentState", + "type": "object", + "properties": { + "input": { + "title": "Input", + "type": "string" + }, + "output": { + "title": "Output", + "type": "string" + } + } + } + } + } + }, + { + "id": "__end__", + "type": "schema", + "data": { + "title": "LangGraphOutput", + "$ref": "#/definitions/AgentState", + "definitions": { + "AgentState": { + "title": "AgentState", + "type": "object", + "properties": { + "input": { + "title": "Input", + "type": "string" + }, + "output": { + "title": "Output", + "type": "string" + } + } + } + } + } + }, + { + "id": "left", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "left" + } + }, + { + "id": "right", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "right" + } + }, + { + "id": "left_", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "" + } + }, + { + "id": "__start___condition", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "should_start" + } + } + ], + "edges": [ + { + "source": "right", + "target": "__end__" + }, + { + "source": "left", + "target": "left_" + }, + { + "source": "left_", + "target": "left", + "data": "left" + }, + { + "source": "left_", + "target": "right", + "data": "right" + }, + { + "source": "__start__", + "target": "__start___condition" + }, + { + "source": "__start___condition", + "target": "left", + "data": "go-left" + }, + { + "source": "__start___condition", + "target": "right", + "data": "go-right" + } + ] + } + ''' +# --- +# name: test_conditional_entrypoint_graph_state.3 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------------+ + | __start___condition | + +---------------------+ + *** *** + * * + ** *** + +------+ * + | left | * + +------+ * + * * + * * + * * + +---------------+ * + | left_ | *** + +---------------+ * + *** *** + * * + ** ** + +-------+ + | right | + +-------+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- # name: test_conditional_graph ''' { diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 33788d9d5..d1cb08a02 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -1074,6 +1074,102 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: ] +def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None: + def left(data: str) -> str: + return data + "->left" + + def right(data: str) -> str: + return data + "->right" + + def should_start(data: str) -> str: + # Logic to decide where to start + if len(data) > 10: + return "go-right" + else: + return "go-left" + + # Define a new graph + workflow = Graph() + + workflow.add_node("left", left) + workflow.add_node("right", right) + + workflow.set_conditional_entry_point( + should_start, {"go-left": "left", "go-right": "right"} + ) + + workflow.add_conditional_edges("left", lambda data: END) + workflow.add_edge("right", END) + + app = workflow.compile() + + assert app.get_input_schema().schema_json() == snapshot + assert app.get_output_schema().schema_json() == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_ascii() == snapshot + + assert app.invoke("what is weather in sf") == "what is weather in sf->right" + + assert [*app.stream("what is weather in sf")] == [ + {"right": "what is weather in sf->right"}, + {"__end__": "what is weather in sf->right"}, + ] + + +def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None: + class AgentState(TypedDict, total=False): + input: str + output: str + + def left(data: AgentState) -> AgentState: + return {"output": data["input"] + "->left"} + + def right(data: AgentState) -> AgentState: + return {"output": data["input"] + "->right"} + + def should_start(data: AgentState) -> str: + # Logic to decide where to start + if len(data["input"]) > 10: + return "go-right" + else: + return "go-left" + + # Define a new graph + workflow = StateGraph(AgentState) + + workflow.add_node("left", left) + workflow.add_node("right", right) + + workflow.set_conditional_entry_point( + should_start, {"go-left": "left", "go-right": "right"} + ) + + workflow.add_conditional_edges("left", lambda data: END) + workflow.add_edge("right", END) + + app = workflow.compile() + + assert app.get_input_schema().schema_json() == snapshot + assert app.get_output_schema().schema_json() == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_ascii() == snapshot + + assert app.invoke({"input": "what is weather in sf"}) == { + "input": "what is weather in sf", + "output": "what is weather in sf->right", + } + + assert [*app.stream({"input": "what is weather in sf"})] == [ + {"right": {"output": "what is weather in sf->right"}}, + { + "__end__": { + "input": "what is weather in sf", + "output": "what is weather in sf->right", + } + }, + ] + + def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 53fe4d39d..d7b762c34 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -1114,6 +1114,92 @@ async def test_conditional_graph_state() -> None: ] +async def test_conditional_entrypoint_graph() -> None: + async def left(data: str) -> str: + return data + "->left" + + async def right(data: str) -> str: + return data + "->right" + + def should_start(data: str) -> str: + # Logic to decide where to start + if len(data) > 10: + return "go-right" + else: + return "go-left" + + # Define a new graph + workflow = Graph() + + workflow.add_node("left", left) + workflow.add_node("right", right) + + workflow.set_conditional_entry_point( + should_start, {"go-left": "left", "go-right": "right"} + ) + + workflow.add_conditional_edges("left", lambda data: END) + workflow.add_edge("right", END) + + app = workflow.compile() + + assert await app.ainvoke("what is weather in sf") == "what is weather in sf->right" + + assert [c async for c in app.astream("what is weather in sf")] == [ + {"right": "what is weather in sf->right"}, + {"__end__": "what is weather in sf->right"}, + ] + + +async def test_conditional_entrypoint_graph_state() -> None: + class AgentState(TypedDict, total=False): + input: str + output: str + + async def left(data: AgentState) -> AgentState: + return {"output": data["input"] + "->left"} + + async def right(data: AgentState) -> AgentState: + return {"output": data["input"] + "->right"} + + def should_start(data: AgentState) -> str: + # Logic to decide where to start + if len(data["input"]) > 10: + return "go-right" + else: + return "go-left" + + # Define a new graph + workflow = StateGraph(AgentState) + + workflow.add_node("left", left) + workflow.add_node("right", right) + + workflow.set_conditional_entry_point( + should_start, {"go-left": "left", "go-right": "right"} + ) + + workflow.add_conditional_edges("left", lambda data: END) + workflow.add_edge("right", END) + + app = workflow.compile() + + assert await app.ainvoke({"input": "what is weather in sf"}) == { + "input": "what is weather in sf", + "output": "what is weather in sf->right", + } + + assert [c async for c in app.astream({"input": "what is weather in sf"})] == [ + {"right": {"output": "what is weather in sf->right"}}, + { + "__end__": { + "input": "what is weather in sf", + "output": "what is weather in sf->right", + } + }, + ] + + async def test_prebuilt_tool_chat() -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool