WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-07 02:07:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

Remove add_condition_nodes

Browse Source
This commit is contained in:
Nuno Campos
2024-04-29 11:13:16 -07:00
parent 3a59fb2247
commit a3c867d688
4 changed files with 374 additions and 995 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/graph/graph.py
+9 -31
View File
@@ -390,7 +390,6 @@ class CompiledGraph(Pregel):
config: Optional[RunnableConfig] = None,
*,
xray: Union[int, bool] = False,
add_condition_nodes: bool = True,
) -> DrawableGraph:
"""Returns a drawable representation of the computation graph."""
graph = DrawableGraph()
@@ -406,7 +405,6 @@ class CompiledGraph(Pregel):
subgraph = (
node.get_graph(
config=config,
add_condition_nodes=add_condition_nodes,
xray=xray - 1 if isinstance(xray, int) and xray > 0 else xray,
)
if isinstance(node, CompiledGraph)
@@ -428,46 +426,26 @@ class CompiledGraph(Pregel):
end_nodes[key] = n
for start, end in sorted(self.graph._all_edges):
graph.add_edge(start_nodes[start], end_nodes[end])
branches_by_name = Counter(
name for _, branches in self.graph.branches.items() for name in branches
)
for start, branches in self.graph.branches.items():
default_ends = {
**{k: k for k in self.graph.nodes if k != start},
END: END,
}
for name, branch in branches.items():
for _, branch in branches.items():
if branch.ends is not None:
ends = branch.ends
elif branch.then is not None:
ends = {k: k for k in default_ends if k not in (END, branch.then)}
else:
ends = default_ends
if add_condition_nodes is True:
cond = graph.add_node(
branch.path,
f"{start}_{name}" if branches_by_name[name] > 1 else name,
for label, end in ends.items():
graph.add_edge(
start_nodes[start],
end_nodes[end],
label if label != end else None,
conditional=True,
)
graph.add_edge(start_nodes[start], cond)
for label, end in ends.items():
graph.add_edge(
cond,
end_nodes[end],
label if label != end else None,
conditional=True,
)
if branch.then is not None:
graph.add_edge(start_nodes[end], end_nodes[branch.then])
else:
for label, end in ends.items():
graph.add_edge(
start_nodes[start],
end_nodes[end],
label if label != end else None,
conditional=True,
)
if branch.then is not None:
graph.add_edge(start_nodes[end], end_nodes[branch.then])
if branch.then is not None:
graph.add_edge(start_nodes[end], end_nodes[branch.then])
return graph
tests/__snapshots__/test_pregel.ambr
+339 -923
View File
File diff suppressed because it is too large Load Diff
tests/__snapshots__/test_pregel_async.ambr
+24 -24
View File
@@ -10,12 +10,12 @@
+---------------+
| rewrite_query |
+---------------+
*** ***
* *
** **
+--------------+ +---------+
| analyzer_one | | decider |
+--------------+ +---------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
@@ -47,12 +47,12 @@
+---------------+
| rewrite_query |
+---------------+
*** ***
* *
** **
+--------------+ +---------+
| analyzer_one | | decider |
+--------------+ +---------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
@@ -84,12 +84,12 @@
+---------------+
| rewrite_query |
+---------------+
*** ***
* *
** **
+--------------+ +-----------+
| analyzer_one | | condition |
+--------------+ +-----------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
@@ -121,12 +121,12 @@
+---------------+
| rewrite_query |
+---------------+
*** ***
* *
** **
+--------------+ +-----------+
| analyzer_one | | condition |
+--------------+ +-----------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
tests/test_pregel.py
+2 -17
View File
@@ -883,18 +883,10 @@ def test_conditional_graph(
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
assert app.get_graph().draw_ascii() == snapshot
assert (
app.get_graph(add_condition_nodes=False).draw_mermaid(with_styles=False)
== snapshot
)
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == snapshot
assert app.get_graph(xray=True).draw_ascii() == snapshot
assert (
app.get_graph(xray=True, add_condition_nodes=False).draw_mermaid(
with_styles=False
)
== snapshot
)
assert app.get_graph(xray=True).draw_mermaid(with_styles=False) == snapshot
assert app.invoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
@@ -3032,9 +3024,7 @@ def test_start_branch_then(
)
tool_two = tool_two_graph.compile()
assert tool_two.get_graph().draw_ascii() == snapshot
assert tool_two.get_graph(add_condition_nodes=False).draw_ascii() == snapshot
assert tool_two.get_graph().draw_mermaid() == snapshot
assert tool_two.get_graph(add_condition_nodes=False).draw_mermaid() == snapshot
assert tool_two.invoke({"my_key": "value", "market": "DE"}) == {
"my_key": "value slow",
@@ -3138,9 +3128,7 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) -
tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"})
tool_two = tool_two_graph.compile()
assert tool_two.get_graph().draw_ascii() == snapshot
assert tool_two.get_graph(add_condition_nodes=False).draw_ascii() == snapshot
assert tool_two.get_graph().draw_mermaid() == snapshot
assert tool_two.get_graph(add_condition_nodes=False).draw_mermaid() == snapshot
assert tool_two.invoke({"my_key": "value", "market": "DE"}, debug=1) == {
"my_key": "value prepared slow finished",
@@ -3877,9 +3865,6 @@ def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None:
assert app.get_graph(xray=True).to_json() == snapshot
assert app.get_graph().draw_ascii() == snapshot
assert app.get_graph(xray=True).draw_mermaid() == snapshot
assert (
app.get_graph(xray=True, add_condition_nodes=False).draw_mermaid() == snapshot
)
def test_nested_graph(snapshot: SnapshotAssertion) -> None: