This commit is contained in:
Nuno Campos
2025-04-22 08:36:42 -07:00
parent 3398715258
commit 56c9c210c3
11 changed files with 481 additions and 3817 deletions
+19 -1
View File
@@ -3,6 +3,7 @@ from inspect import (
ismethod,
signature,
)
from itertools import zip_longest
from types import FunctionType
from typing import (
Any,
@@ -132,6 +133,16 @@ class Branch(NamedTuple):
writer: Writer,
reader: Optional[Callable[[RunnableConfig], Any]] = None,
) -> RunnableCallable:
print(
list(
zip_longest(
writer([e for e in self.ends.values() if e != END]),
[la for la, e in self.ends.items() if e != END],
)
)
if self.ends
else None
)
return ChannelWrite.register_writer(
RunnableCallable(
func=self._route,
@@ -142,7 +153,14 @@ class Branch(NamedTuple):
trace=False,
func_accepts_config=True,
),
writer(list(self.ends.values())) if self.ends else None,
list(
zip_longest(
writer([e for e in self.ends.values() if e != END]),
[la for la, e in self.ends.items() if e != END],
)
)
if self.ends
else None,
)
def _route(
+20 -2
View File
@@ -776,7 +776,7 @@ class CompiledStateGraph(CompiledGraph):
),
ChannelWriteTupleEntry(
mapper=_control_branch,
declared=_control_branch(Command(goto=tuple(node.ends)))
static=_control_static(node.ends)
if node is not None and node.ends is not None
else None,
),
@@ -845,6 +845,8 @@ class CompiledStateGraph(CompiledGraph):
def attach_branch(
self, start: str, name: str, branch: Branch, *, with_reader: bool = True
) -> None:
print(f"Attaching branch {name} to {start} {branch}")
def get_writes(
packets: Sequence[Union[str, Send]],
) -> Sequence[Union[ChannelWriteEntry, Send]]:
@@ -862,7 +864,10 @@ class CompiledStateGraph(CompiledGraph):
ChannelWriteEntry(
f"branch:{start}:{name}::then",
WaitForNames(
{p.node if isinstance(p, Send) else p for p in filtered}
frozenset(
p.node if isinstance(p, Send) else p
for p in filtered
)
),
)
)
@@ -1063,6 +1068,19 @@ def _control_branch(value: Any) -> Sequence[tuple[str, Any]]:
return rtn
def _control_static(
ends: Union[tuple[str, ...], dict[str, str]],
) -> Sequence[tuple[str, Any, Optional[str]]]:
if isinstance(ends, dict):
return [
(CHANNEL_BRANCH_TO.format(k), None, label)
for k, label in ends.items()
if k != END
]
else:
return [(CHANNEL_BRANCH_TO.format(e), None, None) for e in ends if e != END]
def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]:
if isinstance(input, Command):
if input.graph == Command.PARENT:
+55 -51
View File
@@ -6,7 +6,7 @@ from langchain_core.runnables.graph import Graph
from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import CONF, CONFIG_KEY_SEND, END, INPUT
from langgraph.constants import CONF, CONFIG_KEY_SEND, END, INPUT, START
from langgraph.managed.base import ManagedValueSpec
from langgraph.pregel.algo import (
PregelTaskWrites,
@@ -17,8 +17,8 @@ from langgraph.pregel.algo import (
from langgraph.pregel.checkpoint import empty_checkpoint
from langgraph.pregel.io import map_input
from langgraph.pregel.manager import ChannelsManager
from langgraph.pregel.read import DEFAULT_BOUND, PregelNode
from langgraph.pregel.write import ChannelWrite, ChannelWriteTupleEntry
from langgraph.pregel.read import PregelNode
from langgraph.pregel.write import ChannelWrite
from langgraph.types import All, Checkpointer, LoopProtocol
@@ -45,7 +45,7 @@ def draw_graph(
The graph for this Pregel instance.
"""
# (src, dest, is_conditional)
edges: list[tuple[str, str, bool]] = []
edges: set[tuple[str, str, bool]] = set()
step = -1
checkpoint = empty_checkpoint()
@@ -60,9 +60,9 @@ def draw_graph(
LoopProtocol(step=step, stop=-1, config=config),
skip_context=True,
) as (channels, managed):
declared_seen: set[Any] = set()
sources: dict[str, set[tuple[str, bool]]] = {}
step_sources: dict[str, set[tuple[str, bool]]] = {}
static_seen: set[Any] = set()
sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {}
step_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {}
# remove node mappers
nodes = {
k: v.copy(update={"mapper": None}) if v.mapper is not None else v
@@ -94,47 +94,45 @@ def draw_graph(
trigger_to_nodes=trigger_to_nodes,
updated_channels=updated_channels,
)
start_tasks = tasks
# run the pregel loop
while tasks:
conditionals = set()
conditionals: dict[tuple[str, str, Any], Optional[str]] = {}
# run task writers
for task in tasks.values():
for w in task.writers:
# apply regular writes
if isinstance(w, ChannelWrite):
w.invoke(None, task.config)
# apply declared writes (Command)
for entry in w.writes:
if (
isinstance(entry, ChannelWriteTupleEntry)
and entry.declared
and entry not in conditionals
):
# visit only once
declared_seen.add(entry)
# apply them
current_len = len(task.writes)
task.config[CONF][CONFIG_KEY_SEND](entry.declared)
conditionals.update(list(task.writes)[current_len:])
elif w not in declared_seen:
# visit only once
declared_seen.add(w)
# get declared writes
if writes := ChannelWrite.get_declared_writes(w):
# apply them
current_len = len(task.writes)
ChannelWrite.do_write(task.config, writes)
conditionals.update(list(task.writes)[current_len:])
# apply conditional writes declared for static analysis, only once
if w not in static_seen:
static_seen.add(w)
# apply static writes
if writes := ChannelWrite.get_static_writes(w):
conditionals.update(
{(task.name, *t[:2]): t[2] for t in writes}
)
task.config[CONF][CONFIG_KEY_SEND]([t[:2] for t in writes])
# collect sources
step_sources = {
task.name: {(w[0], w in conditionals) for w in task.writes}
task.name: {
(
w[0],
(task.name, *w) in conditionals,
conditionals.get((task.name, *w)),
)
for w in task.writes
}
for task in tasks.values()
}
sources.update(step_sources)
# invert triggers
trigger_to_sources: dict[str, set[tuple[str, bool]]] = defaultdict(set)
trigger_to_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = (
defaultdict(set)
)
for src, triggers in sources.items():
for trigger, cond in triggers:
trigger_to_sources[trigger].add((src, cond))
for trigger, cond, label in triggers:
trigger_to_sources[trigger].add((src, cond, label))
# apply writes
_, updated_channels = apply_writes(
checkpoint, channels, tasks.values(), get_next_version
@@ -158,10 +156,11 @@ def draw_graph(
# collect edges
for task in tasks.values():
for trigger in task.triggers:
for src, cond in sorted(trigger_to_sources[trigger]):
edges.append((src, task.name, cond))
for src, cond, label in sorted(trigger_to_sources[trigger]):
edges.add((src, task.name, cond, label))
# assemble the graph
graph = Graph()
# add nodes
for name, node in nodes.items():
metadata = dict(node.metadata or {})
if name in interrupt_before_nodes and name in interrupt_after_nodes:
@@ -170,12 +169,26 @@ def draw_graph(
metadata["__interrupt"] = "before"
elif name in interrupt_after_nodes:
metadata["__interrupt"] = "after"
graph.add_node(node.bound, name, metadata=metadata)
for src, dest, is_conditional in edges:
# TODO conditional labels
graph.add_node(node.bound, name, metadata=metadata or None)
# add start node
if START not in nodes:
graph.add_node(None, START)
for task in start_tasks.values():
graph.add_edge(graph.nodes[START], graph.nodes[task.name])
# add discovered edges
for src, dest, is_conditional, label in sorted(edges):
graph.add_edge(
graph.nodes[src], graph.nodes[dest], conditional=is_conditional
graph.nodes[src],
graph.nodes[dest],
data=label if label != dest else None,
conditional=is_conditional,
)
# add end edges
if step_sources:
end = graph.add_node(None, END)
termini = {d for _, d, _, _ in edges}.difference(s for s, _, _, _ in edges)
for src in sorted(termini.union(step_sources)):
graph.add_edge(graph.nodes[src], end, conditional=src not in termini)
# replace subgraphs
for name, subgraph in subgraphs.items():
subgraph.trim_first_node()
@@ -191,17 +204,8 @@ def draw_graph(
first, last = graph.extend(subgraph, prefix=name)
for idx, edge in enumerate(graph.edges):
if edge.source == name:
graph.edges[idx] = edge.copy(source=last)
graph.edges[idx] = edge.copy(source=last.id)
elif edge.target == name:
graph.edges[idx] = edge.copy(target=first)
# add end edges
if step_sources:
end = graph.add_node(DEFAULT_BOUND, END)
for src in step_sources:
graph.add_edge(graph.nodes[src], end)
termini = set(d for _, d, _ in edges).difference((s for s, _, _ in edges))
for src in termini.union(step_sources):
# TODO conditional labels
graph.add_edge(graph.nodes[src], end, conditional=src not in termini)
graph.edges[idx] = edge.copy(target=first.id)
return graph
+44 -26
View File
@@ -41,7 +41,7 @@ class ChannelWriteTupleEntry(NamedTuple):
"""Function to extract tuples from value."""
value: Any = PASSTHROUGH
"""Value to write, or PASSTHROUGH to use the input."""
declared: Optional[Sequence[tuple[str, Any]]] = None
static: Optional[Sequence[tuple[str, Any, Optional[str]]]] = None
"""Optional, declared writes for static analysis."""
@@ -138,27 +138,10 @@ class ChannelWrite(RunnableCallable):
if isinstance(w, ChannelWriteTupleEntry):
if w.value is PASSTHROUGH and not allow_passthrough:
raise InvalidUpdateError("PASSTHROUGH value must be replaced")
# assemble writes
tuples: list[tuple[str, Any]] = []
for w in writes:
if isinstance(w, Send):
tuples.append((TASKS, w))
elif isinstance(w, ChannelWriteTupleEntry):
if ww := w.mapper(w.value):
tuples.extend(ww)
elif isinstance(w, ChannelWriteEntry):
value = w.mapper(w.value) if w.mapper is not None else w.value
if value is SKIP_WRITE:
continue
if w.skip_none and value is None:
continue
tuples.append((w.channel, value))
else:
raise ValueError(f"Invalid write entry: {w}")
# if we want to persist writes found before hitting a ParentCommand
# can move this to a finally block
write: TYPE_SEND = config[CONF][CONFIG_KEY_SEND]
write(tuples)
write(_assemble_writes(writes))
@staticmethod
def is_writer(runnable: Runnable) -> bool:
@@ -169,22 +152,57 @@ class ChannelWrite(RunnableCallable):
)
@staticmethod
def get_declared_writes(
def get_static_writes(
runnable: Runnable,
) -> Optional[Sequence[Union[ChannelWriteEntry, Send]]]:
"""Used to get the writes a writer declares for static analysis."""
if writes := getattr(runnable, "_is_channel_writer", MISSING):
return writes if writes is not MISSING else None
) -> Optional[Sequence[tuple[str, Any, Optional[str]]]]:
"""Used to get conditional writes a writer declares for static analysis."""
if isinstance(runnable, ChannelWrite):
return [
w
for entry in runnable.writes
if isinstance(entry, ChannelWriteTupleEntry) and entry.static
for w in entry.static
] or None
elif writes := getattr(runnable, "_is_channel_writer", MISSING):
if writes is not MISSING:
entries = [e for e, _ in writes]
labels = [la for _, la in writes]
return [(*t, la) for t, la in zip(_assemble_writes(entries), labels)]
@staticmethod
def register_writer(
runnable: R,
declared: Optional[Sequence[Union[ChannelWriteEntry, Send]]] = None,
static: Optional[
Sequence[tuple[Union[ChannelWriteEntry, Send], Optional[str]]]
] = None,
) -> R:
"""Used to mark a runnable as a writer, so that it can be detected by is_writer.
Instances of ChannelWrite are automatically marked as writers.
Optionally, a list of declared writes can be passed for static analysis."""
# using object.__setattr__ to work around objects that override __setattr__
# eg. pydantic models and dataclasses
object.__setattr__(runnable, "_is_channel_writer", declared)
object.__setattr__(runnable, "_is_channel_writer", static)
return runnable
def _assemble_writes(
writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
) -> list[tuple[str, Any]]:
"""Assembles the writes into a list of tuples."""
tuples: list[tuple[str, Any]] = []
for w in writes:
if isinstance(w, Send):
tuples.append((TASKS, w))
elif isinstance(w, ChannelWriteTupleEntry):
if ww := w.mapper(w.value):
tuples.extend(ww)
elif isinstance(w, ChannelWriteEntry):
value = w.mapper(w.value) if w.mapper is not None else w.value
if value is SKIP_WRITE:
continue
if w.skip_none and value is None:
continue
tuples.append((w.channel, value))
else:
raise ValueError(f"Invalid write entry: {w}")
return tuples
+4 -4
View File
@@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 2.0.1 and should not be changed by hand.
# This file is automatically @generated by Poetry 2.0.0 and should not be changed by hand.
[[package]]
name = "aiosqlite"
@@ -1324,14 +1324,14 @@ files = [
[[package]]
name = "langchain-core"
version = "0.3.46"
version = "0.3.55"
description = "Building applications with LLMs through composability"
optional = false
python-versions = "<4.0,>=3.9"
groups = ["main", "dev"]
files = [
{file = "langchain_core-0.3.46-py3-none-any.whl", hash = "sha256:28b5689fc347975ea520b5364ab4aee5567e661553bbee5e97cabf4596c28ce0"},
{file = "langchain_core-0.3.46.tar.gz", hash = "sha256:5fca010eeb0a427be5aa8a8525e2112995dde790c584cef165be7c5e0ee1c2b5"},
{file = "langchain_core-0.3.55-py3-none-any.whl", hash = "sha256:b3cb36bf37755a616158a79866657c6697b43a2f7c69dd723ce425f1c76c1baa"},
{file = "langchain_core-0.3.55.tar.gz", hash = "sha256:0f2b3e311621116a83510c70b0ac9d959030a0a457a69483535cff18501fedc9"},
]
[package.dependencies]
File diff suppressed because one or more lines are too long
@@ -1,151 +0,0 @@
# serializer version: 1
# name: test_weather_subgraph[memory]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
__end__([<p>__end__</p>]):::last
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
subgraph weather_graph
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
weather_graph_model_node --> weather_graph_weather_node;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[postgres_aio]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
__end__([<p>__end__</p>]):::last
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
subgraph weather_graph
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
weather_graph_model_node --> weather_graph_weather_node;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[postgres_aio_pipe]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
__end__([<p>__end__</p>]):::last
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
subgraph weather_graph
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
weather_graph_model_node --> weather_graph_weather_node;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[postgres_aio_pool]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
__end__([<p>__end__</p>]):::last
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
subgraph weather_graph
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
weather_graph_model_node --> weather_graph_weather_node;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[postgres_aio_shallow]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
__end__([<p>__end__</p>]):::last
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
subgraph weather_graph
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
weather_graph_model_node --> weather_graph_weather_node;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[sqlite_aio]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
__end__([<p>__end__</p>]):::last
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
subgraph weather_graph
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
weather_graph_model_node --> weather_graph_weather_node;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
@@ -9,11 +9,6 @@
'''
{
"nodes": [
{
"id": "__start__",
"type": "schema",
"data": "__start__"
},
{
"id": "left",
"type": "runnable",
@@ -41,16 +36,23 @@
}
},
{
"id": "__end__",
"type": "schema",
"data": "__end__"
"id": "__start__",
"type": "runnable",
"data": {
"id": [
"langchain",
"schema",
"runnable",
"RunnablePassthrough"
],
"name": "__start__"
}
},
{
"id": "__end__"
}
],
"edges": [
{
"source": "right",
"target": "__end__"
},
{
"source": "__start__",
"target": "left",
@@ -65,8 +67,11 @@
},
{
"source": "left",
"target": "__end__",
"conditional": true
"target": "__end__"
},
{
"source": "right",
"target": "__end__"
}
]
}
@@ -75,10 +80,10 @@
# name: test_conditional_entrypoint_graph.3
'''
graph TD;
right --> __end__;
__start__ -. &nbsp;go-left&nbsp; .-> left;
__start__ -. &nbsp;go-right&nbsp; .-> right;
left -.-> __end__;
left --> __end__;
right --> __end__;
'''
# ---
@@ -94,8 +99,16 @@
"nodes": [
{
"id": "__start__",
"type": "schema",
"data": "__start__"
"type": "runnable",
"data": {
"id": [
"langchain",
"schema",
"runnable",
"RunnablePassthrough"
],
"name": "__start__"
}
},
{
"id": "left",
@@ -124,16 +137,10 @@
}
},
{
"id": "__end__",
"type": "schema",
"data": "__end__"
"id": "__end__"
}
],
"edges": [
{
"source": "right",
"target": "__end__"
},
{
"source": "__start__",
"target": "left",
@@ -148,8 +155,11 @@
},
{
"source": "left",
"target": "__end__",
"conditional": true
"target": "__end__"
},
{
"source": "right",
"target": "__end__"
}
]
}
@@ -158,10 +168,10 @@
# name: test_conditional_entrypoint_graph_state.3
'''
graph TD;
right --> __end__;
__start__ -. &nbsp;go-left&nbsp; .-> left;
__start__ -. &nbsp;go-right&nbsp; .-> right;
left -.-> __end__;
left --> __end__;
right --> __end__;
'''
# ---
@@ -177,8 +187,16 @@
"nodes": [
{
"id": "__start__",
"type": "schema",
"data": "__start__"
"type": "runnable",
"data": {
"id": [
"langchain",
"schema",
"runnable",
"RunnablePassthrough"
],
"name": "__start__"
}
},
{
"id": "get_weather",
@@ -194,25 +212,18 @@
}
},
{
"id": "__end__",
"type": "schema",
"data": "__end__"
"id": "__end__"
}
],
"edges": [
{
"source": "get_weather",
"target": "__end__"
},
{
"source": "__start__",
"target": "get_weather",
"conditional": true
},
{
"source": "__start__",
"target": "__end__",
"conditional": true
"source": "get_weather",
"target": "__end__"
}
]
}
@@ -221,9 +232,8 @@
# name: test_conditional_entrypoint_to_multiple_state_graph.3
'''
graph TD;
get_weather --> __end__;
__start__ -.-> get_weather;
__start__ -.-> __end__;
get_weather --> __end__;
'''
# ---
@@ -233,8 +243,16 @@
"nodes": [
{
"id": "__start__",
"type": "schema",
"data": "__start__"
"type": "runnable",
"data": {
"id": [
"langchain",
"schema",
"runnable",
"RunnablePassthrough"
],
"name": "__start__"
}
},
{
"id": "A",
@@ -263,20 +281,10 @@
}
},
{
"id": "__end__",
"type": "schema",
"data": "__end__"
"id": "__end__"
}
],
"edges": [
{
"source": "A",
"target": "__end__"
},
{
"source": "B",
"target": "__end__"
},
{
"source": "__start__",
"target": "A"
@@ -284,6 +292,14 @@
{
"source": "__start__",
"target": "B"
},
{
"source": "A",
"target": "__end__"
},
{
"source": "B",
"target": "__end__"
}
]
}
@@ -292,10 +308,10 @@
# name: test_conditional_state_graph_with_list_edge_inputs.1
'''
graph TD;
A --> __end__;
B --> __end__;
__start__ --> A;
__start__ --> B;
A --> __end__;
B --> __end__;
'''
# ---
@@ -303,11 +319,11 @@
'''
graph TD;
__start__ --> rewrite_query;
rewrite_query --> analyzer_one;
rewrite_query --> retriever_two;
analyzer_one --> retriever_one;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query --> retriever_two;
qa --> __end__;
'''
@@ -316,11 +332,11 @@
'''
graph TD;
__start__ --> rewrite_query;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
analyzer_one --> retriever_one;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
qa --> __end__;
'''
@@ -386,11 +402,11 @@
'''
graph TD;
__start__ --> rewrite_query;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
analyzer_one --> retriever_one;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
qa --> __end__;
'''
@@ -456,18 +472,22 @@
'''
graph TD;
__start__ --> rewrite_query;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
analyzer_one --> retriever_one;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
qa --> __end__;
'''
# ---
# name: test_multiple_sinks_subgraphs
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
---
config:
flowchart:
curve: linear
---
graph TD;
__start__([<p>__start__</p>]):::first
uno(uno)
@@ -499,7 +519,11 @@
# ---
# name: test_nested_graph.1
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
---
config:
flowchart:
curve: linear
---
graph TD;
__start__([<p>__start__</p>]):::first
inner(inner)
@@ -518,6 +542,7 @@
dict({
'edges': list([
dict({
'conditional': True,
'source': '__start__',
'target': '__end__',
}),
@@ -576,23 +601,25 @@
'type': 'runnable',
}),
dict({
'data': '__end__',
'id': '__end__',
'type': 'schema',
}),
]),
})
# ---
# name: test_nested_graph_xray.1
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
---
config:
flowchart:
curve: linear
---
graph TD;
__start__(<p>__start__</p>)
tool_one(tool_one)
tool_two(tool_two)
tool_three(tool_three)
__end__(<p>__end__</p>)
__start__ --> __end__;
__start__ -.-> __end__;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
@@ -602,15 +629,16 @@
# name: test_repeat_condition
'''
graph TD;
__start__ --> Researcher;
Researcher -. &nbsp;continue&nbsp; .-> Chart_Generator;
Researcher -. &nbsp;call_tool&nbsp; .-> Call_Tool;
Researcher -. &nbsp;end&nbsp; .-> __end__;
Chart_Generator -. &nbsp;continue&nbsp; .-> Researcher;
Chart_Generator -. &nbsp;call_tool&nbsp; .-> Call_Tool;
Chart_Generator -. &nbsp;end&nbsp; .-> __end__;
Call_Tool -.-> Researcher;
Call_Tool -.-> Chart_Generator;
Call_Tool -.-> Researcher;
Chart_Generator -. &nbsp;call_tool&nbsp; .-> Call_Tool;
Chart_Generator -. &nbsp;continue&nbsp; .-> Researcher;
Researcher -. &nbsp;call_tool&nbsp; .-> Call_Tool;
Researcher -. &nbsp;continue&nbsp; .-> Chart_Generator;
__start__ --> Researcher;
Call_Tool -.-> __end__;
Chart_Generator -.-> __end__;
Researcher -.-> __end__;
Researcher -. &nbsp;redo&nbsp; .-> Researcher;
'''
@@ -619,12 +647,12 @@
'''
graph TD;
__start__ --> up;
up --> other;
up --> side;
side --> down;
up --> down;
other --> __end__;
up --> other;
up --> side;
down --> __end__;
other --> __end__;
'''
# ---
@@ -639,7 +667,11 @@
# ---
# name: test_xray_bool
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
---
config:
flowchart:
curve: linear
---
graph TD;
__start__([<p>__start__</p>]):::first
gp_one(gp_one)
@@ -675,24 +707,28 @@
# ---
# name: test_xray_issue
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
---
config:
flowchart:
curve: linear
---
graph TD;
__start__([<p>__start__</p>]):::first
p_one(p_one)
__end__([<p>__end__</p>]):::last
__start__ --> p_one;
p_two___end__ --> p_one;
p_one -. &nbsp;0&nbsp; .-> p_two___start__;
p_one -. &nbsp;1&nbsp; .-> __end__;
p_two___end__ --> p_one;
p_one -.-> __end__;
subgraph p_two
p_two___start__(<p>__start__</p>)
p_two_c_one(c_one)
p_two_c_two(c_two)
p_two___end__(<p>__end__</p>)
p_two___start__ --> p_two_c_one;
p_two_c_two --> p_two_c_one;
p_two_c_one -. &nbsp;0&nbsp; .-> p_two_c_two;
p_two_c_one -. &nbsp;1&nbsp; .-> p_two___end__;
p_two_c_two --> p_two_c_one;
p_two_c_one -.-> p_two___end__;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
@@ -707,15 +743,15 @@
'source': '__start__',
'target': 'ask_question',
}),
dict({
'source': 'ask_question',
'target': 'answer_question',
}),
dict({
'conditional': True,
'source': 'answer_question',
'target': 'ask_question',
}),
dict({
'source': 'ask_question',
'target': 'answer_question',
}),
dict({
'conditional': True,
'source': 'answer_question',
@@ -724,9 +760,17 @@
]),
'nodes': list([
dict({
'data': '__start__',
'data': dict({
'id': list([
'langchain',
'schema',
'runnable',
'RunnablePassthrough',
]),
'name': '__start__',
}),
'id': '__start__',
'type': 'schema',
'type': 'runnable',
}),
dict({
'data': dict({
@@ -755,9 +799,7 @@
'type': 'runnable',
}),
dict({
'data': '__end__',
'id': '__end__',
'type': 'schema',
}),
]),
})
@@ -773,21 +815,29 @@
'source': 'conduct_interview',
'target': 'generate_sections',
}),
dict({
'source': 'generate_sections',
'target': '__end__',
}),
dict({
'conditional': True,
'source': 'generate_analysts',
'target': 'conduct_interview',
}),
dict({
'source': 'generate_sections',
'target': '__end__',
}),
]),
'nodes': list([
dict({
'data': '__start__',
'data': dict({
'id': list([
'langchain',
'schema',
'runnable',
'RunnablePassthrough',
]),
'name': '__start__',
}),
'id': '__start__',
'type': 'schema',
'type': 'runnable',
}),
dict({
'data': dict({
@@ -829,9 +879,7 @@
'type': 'runnable',
}),
dict({
'data': '__end__',
'id': '__end__',
'type': 'schema',
}),
]),
})
@@ -839,24 +887,6 @@
# name: test_xray_lance.2
dict({
'edges': list([
dict({
'source': 'conduct_interview:__start__',
'target': 'conduct_interview:ask_question',
}),
dict({
'source': 'conduct_interview:ask_question',
'target': 'conduct_interview:answer_question',
}),
dict({
'conditional': True,
'source': 'conduct_interview:answer_question',
'target': 'conduct_interview:ask_question',
}),
dict({
'conditional': True,
'source': 'conduct_interview:answer_question',
'target': 'conduct_interview:__end__',
}),
dict({
'source': '__start__',
'target': 'generate_analysts',
@@ -865,21 +895,47 @@
'source': 'conduct_interview:__end__',
'target': 'generate_sections',
}),
dict({
'source': 'generate_sections',
'target': '__end__',
}),
dict({
'conditional': True,
'source': 'generate_analysts',
'target': 'conduct_interview:__start__',
}),
dict({
'source': 'generate_sections',
'target': '__end__',
}),
dict({
'source': 'conduct_interview:__start__',
'target': 'conduct_interview:ask_question',
}),
dict({
'conditional': True,
'source': 'conduct_interview:answer_question',
'target': 'conduct_interview:ask_question',
}),
dict({
'source': 'conduct_interview:ask_question',
'target': 'conduct_interview:answer_question',
}),
dict({
'conditional': True,
'source': 'conduct_interview:answer_question',
'target': 'conduct_interview:__end__',
}),
]),
'nodes': list([
dict({
'data': '__start__',
'data': dict({
'id': list([
'langchain',
'schema',
'runnable',
'RunnablePassthrough',
]),
'name': '__start__',
}),
'id': '__start__',
'type': 'schema',
'type': 'runnable',
}),
dict({
'data': dict({
@@ -895,9 +951,33 @@
'type': 'runnable',
}),
dict({
'data': 'conduct_interview:__start__',
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'generate_sections',
}),
'id': 'generate_sections',
'type': 'runnable',
}),
dict({
'id': '__end__',
}),
dict({
'data': dict({
'id': list([
'langchain',
'schema',
'runnable',
'RunnablePassthrough',
]),
'name': 'conduct_interview:__start__',
}),
'id': 'conduct_interview:__start__',
'type': 'schema',
'type': 'runnable',
}),
dict({
'data': dict({
@@ -926,28 +1006,9 @@
'type': 'runnable',
}),
dict({
'data': 'conduct_interview:__end__',
'id': 'conduct_interview:__end__',
'type': 'schema',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'generate_sections',
}),
'id': 'generate_sections',
'type': 'runnable',
}),
dict({
'data': '__end__',
'id': '__end__',
'type': 'schema',
}),
]),
})
# ---
+23 -16
View File
@@ -587,12 +587,10 @@ def test_conditional_graph(
app = workflow.compile()
if SHOULD_CHECK_SNAPSHOTS:
if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.get_graph().draw_mermaid() == snapshot
assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == 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",
@@ -722,10 +720,6 @@ def test_conditional_graph(
)
config = {"configurable": {"thread_id": "1"}}
if SHOULD_CHECK_SNAPSHOTS:
assert app_w_interrupt.get_graph().to_json() == snapshot
assert app_w_interrupt.get_graph().draw_mermaid() == snapshot
assert [
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
] == [
@@ -1538,7 +1532,7 @@ def test_conditional_state_graph(
app = workflow.compile()
if SHOULD_CHECK_SNAPSHOTS:
if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot
assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
@@ -3774,7 +3768,7 @@ def test_message_graph(
# meaning you can use it as you would any other runnable
app = workflow.compile()
if SHOULD_CHECK_SNAPSHOTS:
if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot
assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
@@ -6234,10 +6228,13 @@ def test_start_branch_then(
tool_two_graph.add_node("tool_two_slow", tool_two_slow)
tool_two_graph.add_node("tool_two_fast", tool_two_fast)
tool_two_graph.set_conditional_entry_point(
lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
then=END,
path_map=["tool_two_slow", "tool_two_fast"],
)
tool_two = tool_two_graph.compile()
assert tool_two.get_graph().draw_mermaid() == snapshot
if checkpointer_name == "memory":
assert tool_two.get_graph().draw_mermaid() == snapshot
assert tool_two.invoke({"my_key": "value", "market": "DE"}) == {
"my_key": "value slow",
@@ -6516,6 +6513,7 @@ def test_branch_then(
tool_two_graph.add_conditional_edges(
source="prepare",
path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
path_map=["tool_two_slow", "tool_two_fast"],
then="finish",
)
tool_two_graph.add_node("prepare", lambda s: {"my_key": " prepared"})
@@ -6523,8 +6521,10 @@ def test_branch_then(
tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"})
tool_two = tool_two_graph.compile()
assert tool_two.get_graph().draw_mermaid(with_styles=False) == snapshot
assert tool_two.get_graph().draw_mermaid() == snapshot
if checkpointer_name == "memory":
assert tool_two.get_graph().draw_mermaid(with_styles=False) == snapshot
assert tool_two.get_graph().draw_mermaid() == snapshot
assert tool_two.invoke({"my_key": "value", "market": "DE"}, debug=1) == {
"my_key": "value prepared slow finished",
@@ -9856,7 +9856,9 @@ def test_send_react_interrupt_control(
builder.add_node(foo)
builder.add_edge(START, "agent")
graph = builder.compile()
assert graph.get_graph().draw_mermaid() == snapshot
if checkpointer_name == "memory":
assert graph.get_graph().draw_mermaid() == snapshot
assert graph.invoke({"messages": [HumanMessage("hello")]}) == {
"messages": [
@@ -10187,12 +10189,17 @@ def test_weather_subgraph(
graph.add_node(normal_llm_node)
graph.add_node("weather_graph", weather_graph)
graph.add_edge(START, "router_node")
graph.add_conditional_edges("router_node", route_after_prediction)
graph.add_conditional_edges(
"router_node",
route_after_prediction,
path_map=["weather_graph", "normal_llm_node"],
)
graph.add_edge("normal_llm_node", END)
graph.add_edge("weather_graph", END)
graph = graph.compile(checkpointer=checkpointer)
assert graph.get_graph(xray=1).draw_mermaid() == snapshot
if checkpointer_name == "memory":
assert graph.get_graph(xray=1).draw_mermaid() == snapshot
config = {"configurable": {"thread_id": "1"}}
thread2 = {"configurable": {"thread_id": "2"}}
@@ -7041,7 +7041,11 @@ async def test_weather_subgraph(
graph.add_node(normal_llm_node)
graph.add_node("weather_graph", weather_graph)
graph.add_edge(START, "router_node")
graph.add_conditional_edges("router_node", route_after_prediction)
graph.add_conditional_edges(
"router_node",
route_after_prediction,
path_map=["weather_graph", "normal_llm_node"],
)
graph.add_edge("normal_llm_node", END)
graph.add_edge("weather_graph", END)
@@ -7051,8 +7055,6 @@ async def test_weather_subgraph(
async with awith_checkpointer(checkpointer_name) as checkpointer:
graph = graph.compile(checkpointer=checkpointer)
assert graph.get_graph(xray=1).draw_mermaid() == snapshot
config = {"configurable": {"thread_id": "1"}}
thread2 = {"configurable": {"thread_id": "2"}}
inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]}
+4 -2
View File
@@ -2147,7 +2147,7 @@ def test_conditional_entrypoint_to_multiple_state_graph(
workflow.add_node("get_weather", get_weather)
workflow.add_edge("get_weather", END)
workflow.set_conditional_entry_point(continue_to_weather)
workflow.set_conditional_entry_point(continue_to_weather, path_map=["get_weather"])
app = workflow.compile()
@@ -4477,7 +4477,9 @@ def test_xray_lance(snapshot: SnapshotAssertion):
# Flow
interview_builder.add_edge(START, "ask_question")
interview_builder.add_edge("ask_question", "answer_question")
interview_builder.add_conditional_edges("answer_question", route_messages)
interview_builder.add_conditional_edges(
"answer_question", route_messages, ["ask_question", END]
)
# Set up memory
memory = InMemorySaver()