WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-30 11:49:38 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #139 from langchain-ai/nc/22feb/add-conditional-entry-point

Browse Source 
Add support for conditional entry points in Graph, StateGraph, MessageGraph
This commit is contained in:
Nuno Campos
2024-02-22 15:23:01 -08:00
committed by GitHub
parent 4f979d69bf debd9a8d9c
commit 0d2b4fe188
5 changed files with 601 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/graph/graph.py
+71 -19
View File
@@ -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
langgraph/graph/state.py
+12 -8
View File
@@ -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,
tests/__snapshots__/test_pregel.ambr
+336
View File
@@ -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_<lambda>",
"type": "runnable",
"data": {
"id": [
"langchain_core",
"runnables",
"base",
"RunnableLambda"
],
"name": "<lambda>"
}
},
{
"id": "__start___condition",
"type": "runnable",
"data": {
"id": [
"langchain_core",
"runnables",
"base",
"RunnableLambda"
],
"name": "should_start"
}
}
],
"edges": [
{
"source": "right",
"target": "__end__"
},
{
"source": "left",
"target": "left_<lambda>"
},
{
"source": "left_<lambda>",
"target": "left",
"data": "left"
},
{
"source": "left_<lambda>",
"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_<lambda> | ***
+---------------+ *
*** ***
* *
** **
+-------+
| 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_<lambda>",
"type": "runnable",
"data": {
"id": [
"langchain_core",
"runnables",
"base",
"RunnableLambda"
],
"name": "<lambda>"
}
},
{
"id": "__start___condition",
"type": "runnable",
"data": {
"id": [
"langchain_core",
"runnables",
"base",
"RunnableLambda"
],
"name": "should_start"
}
}
],
"edges": [
{
"source": "right",
"target": "__end__"
},
{
"source": "left",
"target": "left_<lambda>"
},
{
"source": "left_<lambda>",
"target": "left",
"data": "left"
},
{
"source": "left_<lambda>",
"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_<lambda> | ***
+---------------+ *
*** ***
* *
** **
+-------+
| right |
+-------+
*
*
*
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_conditional_graph
'''
{
tests/test_pregel.py
+96
View File
@@ -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
tests/test_pregel_async.py
+86
View File
@@ -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