mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-11 12:17:53 +02:00
Rewrite Graph/StateGraph.compile()
- Each iteration of the graph is one single iteration in Pregel (ie. no more "{node}:edges" nodes)
- .update_state() now acts exactly as one of the nodes in the graph (which can be chosen), which makes human-in-the-loop scenarios where you want to override the actions of a specific node a lot easier to build
- Graph/StateGraph.compile() now delegate adding nodes/edges/etc to dedicated methods that operate on the Pregel object, which is 90% of the way towards dynamic graphs where nodes and edges can be added during execution
This commit is contained in:
+100
-104
@@ -19,6 +19,8 @@ from langchain_core.runnables.graph import (
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.checkpoint import BaseCheckpointSaver
|
||||
from langgraph.pregel import Channel, Pregel
|
||||
from langgraph.pregel.read import ChannelInvoke
|
||||
from langgraph.pregel.write import ChannelWrite
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -31,11 +33,13 @@ class Branch(NamedTuple):
|
||||
ends: Optional[dict[str, str]]
|
||||
|
||||
def run(self, writer: Callable[[str], Optional[Runnable]]) -> None:
|
||||
return RunnableLambda(
|
||||
self._route,
|
||||
self._aroute,
|
||||
name=self.condition.name,
|
||||
).bind(writer=writer)
|
||||
return ChannelWrite.register_writer(
|
||||
RunnableLambda(
|
||||
self._route,
|
||||
self._aroute,
|
||||
name=self.condition.name,
|
||||
).bind(writer=writer)
|
||||
)
|
||||
|
||||
def _route(
|
||||
self, input: Any, *, writer: Callable[[str], Optional[Runnable]]
|
||||
@@ -62,11 +66,9 @@ class Graph:
|
||||
def __init__(self) -> None:
|
||||
self.nodes: dict[str, Runnable] = {}
|
||||
self.edges = set[tuple[str, str]]()
|
||||
self.branches: defaultdict[str, list[Branch]] = defaultdict(list)
|
||||
self.branches: defaultdict[str, dict[str, Branch]] = defaultdict(dict)
|
||||
self.support_multiple_edges = False
|
||||
self.compiled = False
|
||||
self.entry_point: Optional[str] = None
|
||||
self.entry_point_branch: Optional[Branch] = None
|
||||
|
||||
@property
|
||||
def _all_edges(self) -> set[tuple[str, str]]:
|
||||
@@ -93,7 +95,9 @@ class Graph:
|
||||
)
|
||||
if start_key == END:
|
||||
raise ValueError("END cannot be a start node")
|
||||
if start_key not in self.nodes:
|
||||
if end_key == START:
|
||||
raise ValueError("START cannot be an end node")
|
||||
if start_key not in self.nodes and start_key != START:
|
||||
raise ValueError(f"Need to add_node `{start_key}` first")
|
||||
if end_key not in self.nodes and end_key != END:
|
||||
raise ValueError(f"Need to add_node `{end_key}` first")
|
||||
@@ -118,7 +122,7 @@ class Graph:
|
||||
"Adding an edge to a graph that has already been compiled. This will "
|
||||
"not be reflected in the compiled graph."
|
||||
)
|
||||
if start_key not in self.nodes:
|
||||
if start_key not in self.nodes and start_key != START:
|
||||
raise ValueError(f"Need to add_node `{start_key}` first")
|
||||
if conditional_edge_mapping and set(
|
||||
conditional_edge_mapping.values()
|
||||
@@ -131,18 +135,20 @@ class Graph:
|
||||
)
|
||||
if not isinstance(condition, Runnable):
|
||||
condition = RunnableLambda(condition)
|
||||
if condition.name is None:
|
||||
condition.name = f"{start_key}_condition"
|
||||
if condition.name in self.branches[start_key]:
|
||||
raise ValueError(
|
||||
f"Branch with name `{condition.name}` already exists for node "
|
||||
f"`{start_key}`"
|
||||
)
|
||||
|
||||
self.branches[start_key].append(Branch(condition, conditional_edge_mapping))
|
||||
self.branches[start_key][condition.name] = Branch(
|
||||
condition, conditional_edge_mapping
|
||||
)
|
||||
|
||||
def set_entry_point(self, key: str) -> 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 key not in self.nodes:
|
||||
raise ValueError(f"Need to add_node `{key}` first")
|
||||
self.entry_point = key
|
||||
return self.add_edge(START, key)
|
||||
|
||||
def set_conditional_entry_point(
|
||||
self,
|
||||
@@ -151,23 +157,7 @@ class Graph:
|
||||
],
|
||||
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 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())}."
|
||||
)
|
||||
if not isinstance(condition, Runnable):
|
||||
condition = RunnableLambda(condition)
|
||||
self.entry_point_branch = Branch(condition, conditional_edge_mapping)
|
||||
return self.add_conditional_edges(START, condition, conditional_edge_mapping)
|
||||
|
||||
def set_finish_point(self, key: str) -> None:
|
||||
return self.add_edge(key, END)
|
||||
@@ -180,19 +170,14 @@ class Graph:
|
||||
if node not in all_starts:
|
||||
raise ValueError(f"Node `{node}` is a dead-end")
|
||||
|
||||
branches = [
|
||||
branch for branch_list in self.branches.values() for branch in branch_list
|
||||
all_branches = [
|
||||
branch
|
||||
for branches in self.branches.values()
|
||||
for branch in branches.values()
|
||||
]
|
||||
if self.entry_point_branch is not None:
|
||||
branches.append(self.entry_point_branch)
|
||||
|
||||
all_hard_ends = {end for _, end in self._all_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()
|
||||
if all(branch.ends is not None for branch in all_branches):
|
||||
all_ends = {end for _, end in self._all_edges} | {
|
||||
end for branch in all_branches for end in branch.ends.values()
|
||||
}
|
||||
|
||||
for node in self.nodes:
|
||||
@@ -213,65 +198,89 @@ class Graph:
|
||||
interrupt_after: Optional[Sequence[str]] = None,
|
||||
debug: bool = False,
|
||||
) -> "CompiledGraph":
|
||||
# assign default values
|
||||
interrupt_before = interrupt_before or []
|
||||
interrupt_after = interrupt_after or []
|
||||
|
||||
# validate the graph
|
||||
self.validate(interrupt=interrupt_before + interrupt_after)
|
||||
|
||||
outgoing_edges = defaultdict(list)
|
||||
for start, end in self.edges:
|
||||
outgoing_edges[start].append(f"{end}:inbox" if end != END else END)
|
||||
|
||||
nodes = {
|
||||
key: (Channel.subscribe_to(f"{key}:inbox") | node | Channel.write_to(key))
|
||||
for key, node in self.nodes.items()
|
||||
}
|
||||
node_outboxes = {
|
||||
# we clear outbox channels after each step
|
||||
key: EphemeralValue(Any)
|
||||
for key in self.nodes
|
||||
}
|
||||
|
||||
def branch_writer(dest: str) -> Optional[Runnable]:
|
||||
return Channel.write_to(f"{dest}:inbox" if dest != END else END)
|
||||
|
||||
for key in self.nodes:
|
||||
outgoing = outgoing_edges[key]
|
||||
edges_key = f"{key}:edges"
|
||||
if outgoing or key in self.branches:
|
||||
nodes[edges_key] = Channel.subscribe_to(key, tags=["langsmith:hidden"])
|
||||
if outgoing:
|
||||
nodes[edges_key] |= Channel.write_to(*[dest for dest in outgoing])
|
||||
if key in self.branches:
|
||||
for branch in self.branches[key]:
|
||||
nodes[edges_key] |= branch.run(branch_writer)
|
||||
|
||||
if self.entry_point_branch:
|
||||
nodes[f"{START}:edges"] = Channel.subscribe_to(
|
||||
START, tags=["langsmith:hidden"]
|
||||
) | self.entry_point_branch.run(branch_writer)
|
||||
elif self.entry_point is None:
|
||||
raise ValueError("No entry point set")
|
||||
|
||||
return CompiledGraph(
|
||||
# create empty compiled graph
|
||||
compiled = CompiledGraph(
|
||||
graph=self,
|
||||
nodes=nodes,
|
||||
channels={**node_outboxes},
|
||||
input_channels=f"{self.entry_point}:inbox" if self.entry_point else START,
|
||||
nodes={},
|
||||
channels={START: EphemeralValue(Any), END: EphemeralValue(Any)},
|
||||
input_channels=START,
|
||||
output_channels=END,
|
||||
stream_channels=list(self.nodes),
|
||||
stream_mode="values",
|
||||
stream_channels=[],
|
||||
checkpointer=checkpointer,
|
||||
interrupt_before_nodes=[f"{node}:inbox" for node in interrupt_before],
|
||||
interrupt_before_nodes=interrupt_before,
|
||||
interrupt_after_nodes=interrupt_after,
|
||||
auto_validate=False,
|
||||
debug=debug,
|
||||
)
|
||||
|
||||
# attach nodes, edges, and branches
|
||||
for key, node in self.nodes.items():
|
||||
compiled.attach_node(key, node)
|
||||
|
||||
for start, end in self.edges:
|
||||
compiled.attach_edge(start, end)
|
||||
|
||||
for start, branches in self.branches.items():
|
||||
for name, branch in branches.items():
|
||||
compiled.attach_branch(start, name, branch)
|
||||
|
||||
# validate the compiled graph
|
||||
return compiled.validate()
|
||||
|
||||
|
||||
class CompiledGraph(Pregel):
|
||||
graph: Graph
|
||||
|
||||
def attach_node(self, key: str, node: Runnable) -> None:
|
||||
self.channels[key] = EphemeralValue(Any)
|
||||
self.nodes[key] = (
|
||||
ChannelInvoke(channels=[], triggers=[]) | node | Channel.write_to(key)
|
||||
)
|
||||
self.stream_channels.append(key)
|
||||
|
||||
def attach_edge(self, start: str, end: str) -> None:
|
||||
if end == END:
|
||||
# publish to end channel
|
||||
self.nodes[start].writers.append(Channel.write_to(END))
|
||||
else:
|
||||
# subscribe to start channel
|
||||
self.nodes[end].triggers.append(start)
|
||||
self.nodes[end].channels.append(start)
|
||||
|
||||
def attach_branch(self, start: str, name: str, branch: Branch) -> None:
|
||||
def branch_writer(end: str) -> Optional[ChannelWrite]:
|
||||
return Channel.write_to(
|
||||
f"branch:{start}:{name}:{end}" if end != END else END
|
||||
)
|
||||
|
||||
# add hidden start node
|
||||
if start == START and start not in self.nodes:
|
||||
self.nodes[start] = Channel.subscribe_to(START, tags=["langsmith:hidden"])
|
||||
|
||||
# attach branch writer
|
||||
self.nodes[start] |= branch.run(branch_writer)
|
||||
|
||||
# attach branch readers
|
||||
ends = branch.ends.values() if branch.ends else [node for node in self.nodes]
|
||||
for end in ends:
|
||||
if end != END:
|
||||
channel_name = f"branch:{start}:{name}:{end}"
|
||||
self.channels[channel_name] = EphemeralValue(Any)
|
||||
self.nodes[end].triggers.append(channel_name)
|
||||
self.nodes[end].channels.append(channel_name)
|
||||
|
||||
def get_graph(
|
||||
self, config: Optional[RunnableConfig] = None, *, xray: bool = False
|
||||
) -> RunnableGraph:
|
||||
"""Returns a drawable representation of the computation graph."""
|
||||
graph = RunnableGraph()
|
||||
start_nodes: dict[str, RunnableGraphNode] = {
|
||||
START: graph.add_node(self.get_input_schema(config), START)
|
||||
@@ -304,10 +313,8 @@ class CompiledGraph(Pregel):
|
||||
for start, end in self.graph._all_edges:
|
||||
graph.add_edge(start_nodes[start], end_nodes[end])
|
||||
for start, branches in self.graph.branches.items():
|
||||
for i, branch in enumerate(branches):
|
||||
name = f"{start}_{branch.condition.name or 'condition'}"
|
||||
if i > 0:
|
||||
name += f"_{i}"
|
||||
for name, branch in branches.items():
|
||||
name = f"{start}_{name}"
|
||||
cond = graph.add_node(branch.condition, name)
|
||||
graph.add_edge(start_nodes[start], cond)
|
||||
ends = branch.ends or {
|
||||
@@ -316,16 +323,5 @@ class CompiledGraph(Pregel):
|
||||
}
|
||||
for label, end in ends.items():
|
||||
graph.add_edge(cond, end_nodes[end], label)
|
||||
if entry_point_branch := self.graph.entry_point_branch:
|
||||
cond = graph.add_node(
|
||||
entry_point_branch.condition,
|
||||
entry_point_branch.condition.name or f"{START}_condition",
|
||||
)
|
||||
graph.add_edge(start_nodes[START], cond)
|
||||
ends = entry_point_branch.ends or {k: k for k in self.graph.nodes}
|
||||
for label, end in ends.items():
|
||||
graph.add_edge(cond, end_nodes[end], label)
|
||||
elif self.graph.entry_point:
|
||||
graph.add_edge(start_nodes[START], end_nodes[self.graph.entry_point])
|
||||
|
||||
return graph
|
||||
|
||||
@@ -11,13 +11,9 @@ Messages = Union[list[AnyMessage], AnyMessage]
|
||||
def add_messages(left: Messages, right: Messages) -> Messages:
|
||||
# coerce to list
|
||||
if not isinstance(left, list):
|
||||
left = [message_chunk_to_message(left)]
|
||||
else:
|
||||
left = [message_chunk_to_message(m) for m in left]
|
||||
left = [left]
|
||||
if not isinstance(right, list):
|
||||
right = [message_chunk_to_message(right)]
|
||||
else:
|
||||
right = [message_chunk_to_message(m) for m in right]
|
||||
right = [right]
|
||||
# assign missing ids
|
||||
for m in left:
|
||||
if m.id is None:
|
||||
@@ -25,6 +21,9 @@ def add_messages(left: Messages, right: Messages) -> Messages:
|
||||
for m in right:
|
||||
if m.id is None:
|
||||
m.id = str(uuid.uuid4())
|
||||
# coerce to message
|
||||
left = [message_chunk_to_message(m) for m in left]
|
||||
right = [message_chunk_to_message(m) for m in right]
|
||||
# merge
|
||||
left_idx_by_id = {m.id: i for i, m in enumerate(left)}
|
||||
merged = left.copy()
|
||||
|
||||
+126
-119
@@ -1,22 +1,20 @@
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from functools import partial
|
||||
from inspect import signature
|
||||
from typing import Any, Optional, Sequence, Type, Union
|
||||
|
||||
from langchain_core.runnables import RunnableLambda
|
||||
from langchain_core.runnables import Runnable, RunnableLambda
|
||||
from langchain_core.runnables.base import RunnableLike
|
||||
|
||||
from langgraph.channels.any_value import AnyValue
|
||||
from langgraph.channels.base import BaseChannel, InvalidUpdateError
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.channels.named_barrier_value import NamedBarrierValue
|
||||
from langgraph.checkpoint import BaseCheckpointSaver
|
||||
from langgraph.graph.graph import END, START, CompiledGraph, Graph
|
||||
from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph
|
||||
from langgraph.pregel import Channel
|
||||
from langgraph.pregel.read import ChannelInvoke
|
||||
from langgraph.pregel.read import ChannelInvoke, ChannelRead
|
||||
from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -73,134 +71,143 @@ class StateGraph(Graph):
|
||||
interrupt_after: Optional[Sequence[str]] = None,
|
||||
debug: bool = False,
|
||||
) -> CompiledGraph:
|
||||
interrupt_before = interrupt_before or []
|
||||
interrupt_after = interrupt_after or []
|
||||
self.validate(interrupt=interrupt_before + interrupt_after)
|
||||
|
||||
# prepare state channels
|
||||
state_keys = list(self.channels)
|
||||
state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys
|
||||
state_channels = (
|
||||
{chan: chan for chan in state_keys}
|
||||
if isinstance(state_keys_read, list)
|
||||
else {None: state_keys_read}
|
||||
)
|
||||
update_channels = (
|
||||
[ChannelWriteEntry("__root__", None, True)]
|
||||
if not isinstance(state_keys_read, list)
|
||||
else [
|
||||
ChannelWriteEntry(
|
||||
key, RunnableLambda(partial(_dict_getter, state_keys, key)), False
|
||||
)
|
||||
for key in state_keys_read
|
||||
]
|
||||
)
|
||||
coerce_state = (
|
||||
partial(_coerce_state, self.schema)
|
||||
if isinstance(state_keys_read, list)
|
||||
else None
|
||||
)
|
||||
|
||||
waiting_edges = {
|
||||
(f"{starts}:{end}", starts, end) for starts, end in self.w_edges
|
||||
}
|
||||
waiting_edge_channels = {
|
||||
key: NamedBarrierValue(str, set(starts)) for key, starts, _ in waiting_edges
|
||||
}
|
||||
# assign default values
|
||||
interrupt_before = interrupt_before or []
|
||||
interrupt_after = interrupt_after or []
|
||||
|
||||
outgoing_edges = defaultdict(list)
|
||||
for start, end in self.edges:
|
||||
if end != END:
|
||||
outgoing_edges[start].append(f"{end}:inbox")
|
||||
for key, starts, end in waiting_edges:
|
||||
for start in starts:
|
||||
outgoing_edges[start].append(key)
|
||||
# validate the graph
|
||||
self.validate(interrupt=interrupt_before + interrupt_after)
|
||||
|
||||
nodes = {
|
||||
key: (
|
||||
ChannelInvoke(
|
||||
triggers=[
|
||||
f"{key}:inbox",
|
||||
*[chan for chan, _, end in waiting_edges if end == key],
|
||||
],
|
||||
channels=state_channels,
|
||||
mapper=coerce_state,
|
||||
)
|
||||
| node
|
||||
| ChannelWrite(
|
||||
channels=[ChannelWriteEntry(key, None, False)] + update_channels
|
||||
)
|
||||
)
|
||||
for key, node in self.nodes.items()
|
||||
}
|
||||
node_inboxes = {
|
||||
# we take any value written to channel because all writers
|
||||
# write the entire state as of that step, which is equal for all
|
||||
f"{key}:inbox": AnyValue(self.schema)
|
||||
for key in list(self.nodes) + [START]
|
||||
}
|
||||
node_outboxes = {
|
||||
# we clear outbox channels after each step
|
||||
key: EphemeralValue(Any)
|
||||
for key in list(self.nodes) + [START]
|
||||
}
|
||||
|
||||
def branch_writer(src: str, dest: str) -> Optional[ChannelWrite]:
|
||||
if dest != END:
|
||||
return ChannelWrite(
|
||||
channels=[ChannelWriteEntry(f"{dest}:inbox", src, False)]
|
||||
)
|
||||
|
||||
for key in self.nodes:
|
||||
outgoing = outgoing_edges[key]
|
||||
edges_key = f"{key}:edges"
|
||||
if outgoing or key in self.branches:
|
||||
nodes[edges_key] = ChannelInvoke(
|
||||
triggers=[key], tags=["langsmith:hidden"], channels=state_channels
|
||||
)
|
||||
if outgoing:
|
||||
nodes[edges_key] |= ChannelWrite(
|
||||
channels=[ChannelWriteEntry(dest, key, True) for dest in outgoing]
|
||||
)
|
||||
if key in self.branches:
|
||||
for branch in self.branches[key]:
|
||||
nodes[edges_key] |= branch.run(partial(branch_writer, key))
|
||||
|
||||
nodes[START] = Channel.subscribe_to(
|
||||
f"{START}:inbox", tags=["langsmith:hidden"]
|
||||
) | ChannelWrite(
|
||||
channels=[ChannelWriteEntry(START, None, False)] + update_channels
|
||||
)
|
||||
nodes[f"{START}:edges"] = ChannelInvoke(
|
||||
triggers=[START], tags=["langsmith:hidden"], channels=state_channels
|
||||
)
|
||||
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"] |= self.entry_point_branch.run(
|
||||
partial(branch_writer, START)
|
||||
)
|
||||
else:
|
||||
raise ValueError("No entry point set")
|
||||
|
||||
return CompiledGraph(
|
||||
compiled = CompiledStateGraph(
|
||||
graph=self,
|
||||
nodes=nodes,
|
||||
channels={
|
||||
**self.channels,
|
||||
**node_inboxes,
|
||||
**node_outboxes,
|
||||
**waiting_edge_channels,
|
||||
},
|
||||
input_channels=f"{START}:inbox",
|
||||
nodes={},
|
||||
channels={**self.channels, START: EphemeralValue(self.schema)},
|
||||
input_channels=START,
|
||||
stream_mode="updates",
|
||||
output_channels=state_keys_read,
|
||||
stream_channels=state_keys_read,
|
||||
checkpointer=checkpointer,
|
||||
interrupt_before_nodes=[f"{node}:inbox" for node in interrupt_before],
|
||||
interrupt_before_nodes=interrupt_before,
|
||||
interrupt_after_nodes=interrupt_after,
|
||||
auto_validate=False,
|
||||
debug=debug,
|
||||
)
|
||||
|
||||
compiled.attach_node(START, None)
|
||||
for key, node in self.nodes.items():
|
||||
compiled.attach_node(key, node)
|
||||
|
||||
for start, end in self.edges:
|
||||
compiled.attach_edge(start, end)
|
||||
|
||||
for starts, end in self.w_edges:
|
||||
compiled.attach_edge(starts, end)
|
||||
|
||||
for start, branches in self.branches.items():
|
||||
for name, branch in branches.items():
|
||||
compiled.attach_branch(start, name, branch)
|
||||
|
||||
return compiled.validate()
|
||||
|
||||
|
||||
class CompiledStateGraph(CompiledGraph):
|
||||
graph: StateGraph
|
||||
|
||||
def attach_node(self, key: str, node: Optional[Runnable]) -> None:
|
||||
# prepare state
|
||||
state_keys = list(self.graph.channels)
|
||||
state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys
|
||||
state_write_entries = [
|
||||
ChannelWriteEntry(
|
||||
key,
|
||||
RunnableLambda(partial(_dict_getter, state_keys, key)),
|
||||
False,
|
||||
)
|
||||
if key != "__root__"
|
||||
else ChannelWriteEntry(key, None, True)
|
||||
for key in state_keys
|
||||
]
|
||||
|
||||
# add node and output channel
|
||||
if key == START:
|
||||
self.nodes[key] = Channel.subscribe_to(
|
||||
START, tags=["langsmith:hidden"]
|
||||
).pipe(ChannelWrite(channels=state_write_entries))
|
||||
else:
|
||||
self.channels[key] = EphemeralValue(Any)
|
||||
self.nodes[key] = ChannelInvoke(
|
||||
triggers=[],
|
||||
# read state keys
|
||||
channels=(
|
||||
{chan: chan for chan in state_keys}
|
||||
if isinstance(state_keys_read, list)
|
||||
else [state_keys_read]
|
||||
),
|
||||
# coerce state dict to schema class (eg. pydantic model)
|
||||
mapper=(
|
||||
partial(_coerce_state, self.graph.schema)
|
||||
if isinstance(state_keys_read, list)
|
||||
else None
|
||||
),
|
||||
# publish to this channel and state keys
|
||||
writers=[
|
||||
ChannelWrite(
|
||||
[ChannelWriteEntry(key, None, False)] + state_write_entries
|
||||
),
|
||||
# read back state with updates applied
|
||||
ChannelRead(state_keys_read, fresh=True),
|
||||
],
|
||||
).pipe(node)
|
||||
|
||||
def attach_edge(self, starts: Union[str, Sequence[str]], end: str) -> None:
|
||||
if isinstance(starts, str):
|
||||
if starts == START:
|
||||
channel_name = f"start:{end}"
|
||||
# register channel
|
||||
self.channels[channel_name] = EphemeralValue(Any)
|
||||
# subscribe to channel
|
||||
self.nodes[end].triggers.append(channel_name)
|
||||
# publish to channel
|
||||
self.nodes[START] |= ChannelWrite(
|
||||
[ChannelWriteEntry(channel_name, START, False)]
|
||||
)
|
||||
elif end != END:
|
||||
# subscribe to start channel
|
||||
self.nodes[end].triggers.append(starts)
|
||||
else:
|
||||
channel_name = f"join:{starts}:{end}"
|
||||
# register channel
|
||||
self.channels[channel_name] = NamedBarrierValue(str, set(starts))
|
||||
# subscribe to channel
|
||||
self.nodes[end].triggers.append(channel_name)
|
||||
# publish to channel
|
||||
for start in starts:
|
||||
self.nodes[start] |= ChannelWrite(
|
||||
[ChannelWriteEntry(channel_name, start, False)]
|
||||
)
|
||||
|
||||
def attach_branch(self, start: str, name: str, branch: Branch) -> None:
|
||||
def branch_writer(end: str) -> Optional[ChannelWrite]:
|
||||
if end != END:
|
||||
return ChannelWrite(
|
||||
[ChannelWriteEntry(f"branch:{start}:{name}:{end}", start, False)]
|
||||
)
|
||||
|
||||
# attach branch publisher
|
||||
self.nodes[start] |= branch.run(branch_writer)
|
||||
|
||||
# attach branch subscribers
|
||||
ends = branch.ends.values() if branch.ends else [node for node in self.nodes]
|
||||
for end in ends:
|
||||
if end != END:
|
||||
channel_name = f"branch:{start}:{name}:{end}"
|
||||
self.channels[channel_name] = EphemeralValue(Any)
|
||||
self.nodes[end].triggers.append(channel_name)
|
||||
|
||||
|
||||
def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
|
||||
return schema(**input)
|
||||
|
||||
+271
-139
@@ -14,6 +14,7 @@ from typing import (
|
||||
Mapping,
|
||||
NamedTuple,
|
||||
Optional,
|
||||
Self,
|
||||
Sequence,
|
||||
Type,
|
||||
Union,
|
||||
@@ -29,6 +30,7 @@ from langchain_core.globals import get_debug
|
||||
from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
|
||||
from langchain_core.runnables import (
|
||||
Runnable,
|
||||
RunnableSequence,
|
||||
RunnableSerializable,
|
||||
)
|
||||
from langchain_core.runnables.base import Input, Output, coerce_to_runnable
|
||||
@@ -52,6 +54,7 @@ from langgraph.channels.base import (
|
||||
InvalidUpdateError,
|
||||
create_checkpoint,
|
||||
)
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
Checkpoint,
|
||||
@@ -59,12 +62,15 @@ from langgraph.checkpoint.base import (
|
||||
copy_checkpoint,
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND, INTERRUPT
|
||||
from langgraph.constants import (
|
||||
CONFIG_KEY_READ,
|
||||
CONFIG_KEY_SEND,
|
||||
INTERRUPT,
|
||||
)
|
||||
from langgraph.pregel.debug import print_checkpoint, print_step_start
|
||||
from langgraph.pregel.io import map_input, map_output_updates, map_output_values
|
||||
from langgraph.pregel.log import logger
|
||||
from langgraph.pregel.read import ChannelInvoke
|
||||
from langgraph.pregel.reserved import AllReservedChannels, ReservedChannels
|
||||
from langgraph.pregel.validate import validate_graph, validate_keys
|
||||
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
|
||||
|
||||
@@ -94,7 +100,6 @@ class Channel:
|
||||
channels: str,
|
||||
*,
|
||||
key: Optional[str] = None,
|
||||
when: Optional[Callable[[Any], bool]] = None,
|
||||
tags: Optional[list[str]] = None,
|
||||
) -> ChannelInvoke:
|
||||
...
|
||||
@@ -106,7 +111,6 @@ class Channel:
|
||||
channels: Sequence[str],
|
||||
*,
|
||||
key: None = None,
|
||||
when: Optional[Callable[[Any], bool]] = None,
|
||||
tags: Optional[list[str]] = None,
|
||||
) -> ChannelInvoke:
|
||||
...
|
||||
@@ -117,7 +121,6 @@ class Channel:
|
||||
channels: Union[str, Sequence[str]],
|
||||
*,
|
||||
key: Optional[str] = None,
|
||||
when: Optional[Callable[[Any], bool]] = None,
|
||||
tags: Optional[list[str]] = None,
|
||||
) -> ChannelInvoke:
|
||||
"""Runs process.invoke() each time channels are updated,
|
||||
@@ -130,11 +133,12 @@ class Channel:
|
||||
channels=cast(
|
||||
Union[Mapping[None, str], Mapping[str, str]],
|
||||
{key: channels}
|
||||
if isinstance(channels, str) and key is not None
|
||||
else [channels]
|
||||
if isinstance(channels, str)
|
||||
else {chan: chan for chan in channels},
|
||||
),
|
||||
triggers=[channels] if isinstance(channels, str) else channels,
|
||||
when=when,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
@@ -146,13 +150,11 @@ class Channel:
|
||||
) -> ChannelWrite:
|
||||
"""Writes to channels the result of the lambda, or None to skip writing."""
|
||||
return ChannelWrite(
|
||||
channels=(
|
||||
[ChannelWriteEntry(c, None, False) for c in channels]
|
||||
+ [
|
||||
ChannelWriteEntry(k, _coerce_write_value(v), True)
|
||||
for k, v in kwargs.items()
|
||||
]
|
||||
)
|
||||
[ChannelWriteEntry(c, None, False) for c in channels]
|
||||
+ [
|
||||
ChannelWriteEntry(k, _coerce_write_value(v), True)
|
||||
for k, v in kwargs.items()
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@@ -177,6 +179,10 @@ class Pregel(
|
||||
|
||||
channels: Mapping[str, BaseChannel] = Field(default_factory=dict)
|
||||
|
||||
default_channel_cls: Type[BaseChannel] = Field(default=LastValue)
|
||||
|
||||
auto_validate: bool = True
|
||||
|
||||
stream_mode: StreamMode = "values"
|
||||
|
||||
output_channels: Union[str, Sequence[str]] = "output"
|
||||
@@ -203,7 +209,9 @@ class Pregel(
|
||||
arbitrary_types_allowed = True
|
||||
|
||||
@root_validator(skip_on_failure=True)
|
||||
def validate_pregel(cls, values: dict[str, Any]) -> dict[str, Any]:
|
||||
def validate_on_init(cls, values: dict[str, Any]) -> dict[str, Any]:
|
||||
if not values["auto_validate"]:
|
||||
return values
|
||||
validate_graph(
|
||||
values["nodes"],
|
||||
values["channels"],
|
||||
@@ -212,12 +220,29 @@ class Pregel(
|
||||
values["stream_channels"],
|
||||
values["interrupt_after_nodes"],
|
||||
values["interrupt_before_nodes"],
|
||||
values["default_channel_cls"],
|
||||
)
|
||||
if values["interrupt_after_nodes"] or values["interrupt_before_nodes"]:
|
||||
if not values["checkpointer"]:
|
||||
raise ValueError("Interrupts require a checkpointer")
|
||||
return values
|
||||
|
||||
def validate(self) -> Self:
|
||||
validate_graph(
|
||||
self.nodes,
|
||||
self.channels,
|
||||
self.input_channels,
|
||||
self.output_channels,
|
||||
self.stream_channels,
|
||||
self.interrupt_after_nodes,
|
||||
self.interrupt_before_nodes,
|
||||
self.default_channel_cls,
|
||||
)
|
||||
if self.interrupt_after_nodes or self.interrupt_before_nodes:
|
||||
if not self.checkpointer:
|
||||
raise ValueError("Interrupts require a checkpointer")
|
||||
return self
|
||||
|
||||
@property
|
||||
def config_specs(self) -> list[ConfigurableFieldSpec]:
|
||||
return [
|
||||
@@ -274,8 +299,7 @@ class Pregel(
|
||||
return (
|
||||
[self.stream_channels]
|
||||
if isinstance(self.stream_channels, str)
|
||||
else self.stream_channels
|
||||
or [k for k in self.channels if k not in AllReservedChannels]
|
||||
else self.stream_channels or [k for k in self.channels]
|
||||
)
|
||||
|
||||
def get_state(self, config: RunnableConfig) -> StateSnapshot:
|
||||
@@ -290,10 +314,13 @@ class Pregel(
|
||||
checkpoint, self.nodes, channels, for_execution=False
|
||||
)
|
||||
values = {
|
||||
k: _read_channel(channels, k)
|
||||
k: _read_channel(channels, k, return_exception=True)
|
||||
for k in channels
|
||||
if k in self.snapshot_channels_list
|
||||
}
|
||||
values = {
|
||||
k: v for k, v in values.items() if not isinstance(v, EmptyChannelError)
|
||||
}
|
||||
return StateSnapshot(
|
||||
values[self.stream_channels]
|
||||
if isinstance(self.stream_channels, str)
|
||||
@@ -314,10 +341,13 @@ class Pregel(
|
||||
checkpoint, self.nodes, channels, for_execution=False
|
||||
)
|
||||
values = {
|
||||
k: _read_channel(channels, k)
|
||||
k: _read_channel(channels, k, return_exception=True)
|
||||
for k in channels
|
||||
if k in self.snapshot_channels_list
|
||||
}
|
||||
values = {
|
||||
k: v for k, v in values.items() if not isinstance(v, EmptyChannelError)
|
||||
}
|
||||
return StateSnapshot(
|
||||
values[self.stream_channels]
|
||||
if isinstance(self.stream_channels, str)
|
||||
@@ -336,10 +366,15 @@ class Pregel(
|
||||
checkpoint, self.nodes, channels, for_execution=False
|
||||
)
|
||||
values = {
|
||||
k: _read_channel(channels, k)
|
||||
k: _read_channel(channels, k, return_exception=True)
|
||||
for k in channels
|
||||
if k in self.snapshot_channels_list
|
||||
}
|
||||
values = {
|
||||
k: v
|
||||
for k, v in values.items()
|
||||
if not isinstance(v, EmptyChannelError)
|
||||
}
|
||||
yield StateSnapshot(
|
||||
values[self.stream_channels]
|
||||
if isinstance(self.stream_channels, str)
|
||||
@@ -361,10 +396,15 @@ class Pregel(
|
||||
checkpoint, self.nodes, channels, for_execution=False
|
||||
)
|
||||
values = {
|
||||
k: _read_channel(channels, k)
|
||||
k: _read_channel(channels, k, return_exception=True)
|
||||
for k in channels
|
||||
if k in self.snapshot_channels_list
|
||||
}
|
||||
values = {
|
||||
k: v
|
||||
for k, v in values.items()
|
||||
if not isinstance(v, EmptyChannelError)
|
||||
}
|
||||
yield StateSnapshot(
|
||||
values[self.stream_channels]
|
||||
if isinstance(self.stream_channels, str)
|
||||
@@ -375,49 +415,117 @@ class Pregel(
|
||||
)
|
||||
|
||||
def update_state(
|
||||
self, config: RunnableConfig, values: dict[str, Any] | Any
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
values: dict[str, Any] | Any,
|
||||
as_node: Optional[str] = None,
|
||||
) -> RunnableConfig:
|
||||
"""Update the state of the graph with the given values, as if they came from
|
||||
node `as_node`. If `as_node` is not provided, it will be set to the last node
|
||||
that updated the state, if not ambiguous.
|
||||
"""
|
||||
if not self.checkpointer:
|
||||
raise ValueError("No checkpointer set")
|
||||
|
||||
values = (
|
||||
{self.stream_channels: values}
|
||||
if isinstance(self.stream_channels, str)
|
||||
else values
|
||||
)
|
||||
# get last checkpoint
|
||||
checkpoint = self.checkpointer.get(config)
|
||||
checkpoint = copy_checkpoint(checkpoint) if checkpoint else empty_checkpoint()
|
||||
# find last node that updated the state, if not provided
|
||||
if as_node is None:
|
||||
last_seen_by_node = sorted(
|
||||
(v, n)
|
||||
for n, seen in checkpoint["versions_seen"].items()
|
||||
for v in seen.values()
|
||||
)
|
||||
# if two nodes updated the state at the same time, it's ambiguous
|
||||
if last_seen_by_node:
|
||||
if last_seen_by_node[-1][0] != last_seen_by_node[-2][0]:
|
||||
as_node = last_seen_by_node[-1][1]
|
||||
if as_node is None:
|
||||
raise InvalidUpdateError("Ambiguous update, specify as_node")
|
||||
# update channels
|
||||
with ChannelsManager(self.channels, checkpoint) as channels:
|
||||
for k, v in values.items():
|
||||
channels[k].update([v])
|
||||
checkpoint["channel_versions"][k] += 1
|
||||
for k in self.snapshot_channels_list:
|
||||
version = checkpoint["channel_versions"][k]
|
||||
checkpoint["versions_seen"][INTERRUPT][k] = version
|
||||
# create task to run all writers of the chosen node
|
||||
task = PregelExecutableTask(
|
||||
RunnableSequence(*self.nodes[as_node].writers)
|
||||
if len(self.nodes[as_node].writers) > 1
|
||||
else self.nodes[as_node].writers[0],
|
||||
values,
|
||||
as_node,
|
||||
deque(),
|
||||
)
|
||||
# execute task
|
||||
task.proc.invoke(
|
||||
task.input,
|
||||
patch_config(
|
||||
config,
|
||||
configurable={
|
||||
# deque.extend is thread-safe
|
||||
CONFIG_KEY_SEND: task.writes.extend,
|
||||
CONFIG_KEY_READ: partial(
|
||||
_local_read, checkpoint, channels, task.writes
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
# apply to checkpoint and save
|
||||
_apply_writes(checkpoint, channels, task.writes)
|
||||
return self.checkpointer.put(
|
||||
config, create_checkpoint(checkpoint, channels)
|
||||
)
|
||||
|
||||
async def aupdate_state(
|
||||
self, config: RunnableConfig, values: dict[str, Any] | Any
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
values: dict[str, Any] | Any,
|
||||
as_node: Optional[str] = None,
|
||||
) -> RunnableConfig:
|
||||
if not self.checkpointer:
|
||||
raise ValueError("No checkpointer set")
|
||||
|
||||
values = (
|
||||
{self.stream_channels: values}
|
||||
if isinstance(self.stream_channels, str)
|
||||
else values
|
||||
)
|
||||
# get last checkpoint
|
||||
checkpoint = await self.checkpointer.aget(config)
|
||||
checkpoint = copy_checkpoint(checkpoint) if checkpoint else empty_checkpoint()
|
||||
# find last node that updated the state, if not provided
|
||||
if as_node is None:
|
||||
last_seen_by_node = sorted(
|
||||
(v, n)
|
||||
for n, seen in checkpoint["versions_seen"].items()
|
||||
for v in seen.values()
|
||||
)
|
||||
# if two nodes updated the state at the same time, it's ambiguous
|
||||
if last_seen_by_node:
|
||||
if last_seen_by_node[-1][0] != last_seen_by_node[-2][0]:
|
||||
as_node = last_seen_by_node[-1][1]
|
||||
if as_node is None:
|
||||
raise InvalidUpdateError("Ambiguous update, specify as_node")
|
||||
# update channels, acting as the chosen node
|
||||
async with AsyncChannelsManager(self.channels, checkpoint) as channels:
|
||||
for k, v in values.items():
|
||||
channels[k].update([v])
|
||||
checkpoint["channel_versions"][k] += 1
|
||||
for k in self.stream_channels or self.channels:
|
||||
version = checkpoint["channel_versions"][k]
|
||||
checkpoint["versions_seen"][INTERRUPT][k] = version
|
||||
# create task to run all writers of the chosen node
|
||||
task = PregelExecutableTask(
|
||||
RunnableSequence(*self.nodes[as_node].writers)
|
||||
if len(self.nodes[as_node].writers) > 1
|
||||
else self.nodes[as_node].writers[0],
|
||||
values,
|
||||
as_node,
|
||||
deque(),
|
||||
)
|
||||
# execute task
|
||||
await task.proc.ainvoke(
|
||||
task.input,
|
||||
patch_config(
|
||||
config,
|
||||
configurable={
|
||||
# deque.extend is thread-safe
|
||||
CONFIG_KEY_SEND: task.writes.extend,
|
||||
CONFIG_KEY_READ: partial(
|
||||
_local_read, checkpoint, channels, task.writes
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
# apply to checkpoint and save
|
||||
_apply_writes(checkpoint, channels, task.writes)
|
||||
return await self.checkpointer.aput(
|
||||
config, create_checkpoint(checkpoint, channels)
|
||||
)
|
||||
@@ -503,13 +611,7 @@ class Pregel(
|
||||
checkpoint, processes, channels, for_execution=True
|
||||
)
|
||||
# apply input writes
|
||||
_apply_writes(
|
||||
checkpoint,
|
||||
channels,
|
||||
input_writes,
|
||||
config,
|
||||
0,
|
||||
)
|
||||
_apply_writes(checkpoint, channels, input_writes)
|
||||
else:
|
||||
# if received no input, take that as signal to proceed
|
||||
# past previous interrupt, if any
|
||||
@@ -518,8 +620,6 @@ class Pregel(
|
||||
version = checkpoint["channel_versions"][k]
|
||||
checkpoint["versions_seen"][INTERRUPT][k] = version
|
||||
|
||||
read = partial(_read_channel, channels)
|
||||
|
||||
# Similarly to Bulk Synchronous Parallel / Pregel model
|
||||
# computation proceeds in steps, while there are channel updates
|
||||
# channel updates from step N are only visible in step N+1
|
||||
@@ -543,6 +643,15 @@ class Pregel(
|
||||
"limit by setting the `recursion_limit` config key."
|
||||
)
|
||||
|
||||
# before execution, check if we should interrupt
|
||||
if _should_interrupt(
|
||||
checkpoint,
|
||||
interrupt_before_nodes,
|
||||
self.snapshot_channels_list,
|
||||
next_tasks,
|
||||
):
|
||||
break
|
||||
|
||||
if debug:
|
||||
print_step_start(step, next_tasks)
|
||||
|
||||
@@ -558,7 +667,9 @@ class Pregel(
|
||||
configurable={
|
||||
# deque.extend is thread-safe
|
||||
CONFIG_KEY_SEND: writes.extend,
|
||||
CONFIG_KEY_READ: read,
|
||||
CONFIG_KEY_READ: partial(
|
||||
_local_read, checkpoint, channels, writes
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
@@ -587,9 +698,7 @@ class Pregel(
|
||||
pending_writes.extend(writes)
|
||||
|
||||
# apply writes to channels
|
||||
_apply_writes(
|
||||
checkpoint, channels, pending_writes, config, step + 1
|
||||
)
|
||||
_apply_writes(checkpoint, channels, pending_writes)
|
||||
|
||||
if debug:
|
||||
print_checkpoint(step, channels)
|
||||
@@ -604,31 +713,21 @@ class Pregel(
|
||||
if step_output := map_output_updates(output_keys, next_tasks):
|
||||
yield step_output
|
||||
|
||||
# with previous step's checkpoint
|
||||
if _should_interrupt(
|
||||
checkpoint,
|
||||
interrupt_before_nodes,
|
||||
self.snapshot_channels_list,
|
||||
pending_writes,
|
||||
):
|
||||
break
|
||||
|
||||
# save end of step checkpoint
|
||||
if self.checkpointer is not None and (
|
||||
self.checkpointer.at == CheckpointAt.END_OF_STEP
|
||||
or interrupt_before_nodes
|
||||
):
|
||||
checkpoint = create_checkpoint(checkpoint, channels)
|
||||
checkpoint_config = self.checkpointer.put(
|
||||
checkpoint_config, checkpoint
|
||||
)
|
||||
|
||||
# with this step's checkpoint,
|
||||
# after execution, check if we should interrupt
|
||||
if _should_interrupt(
|
||||
checkpoint,
|
||||
interrupt_after_nodes,
|
||||
self.snapshot_channels_list,
|
||||
pending_writes,
|
||||
next_tasks,
|
||||
):
|
||||
break
|
||||
|
||||
@@ -636,7 +735,6 @@ class Pregel(
|
||||
if (
|
||||
self.checkpointer is not None
|
||||
and self.checkpointer.at == CheckpointAt.END_OF_RUN
|
||||
and not interrupt_before_nodes
|
||||
):
|
||||
checkpoint = create_checkpoint(checkpoint, channels)
|
||||
self.checkpointer.put(checkpoint_config, checkpoint)
|
||||
@@ -697,13 +795,7 @@ class Pregel(
|
||||
checkpoint, processes, channels, for_execution=True
|
||||
)
|
||||
# apply input writes
|
||||
_apply_writes(
|
||||
checkpoint,
|
||||
channels,
|
||||
input_writes,
|
||||
config,
|
||||
0,
|
||||
)
|
||||
_apply_writes(checkpoint, channels, input_writes)
|
||||
else:
|
||||
# if received no input, take that as signal to proceed
|
||||
# past previous interrupt, if any
|
||||
@@ -712,8 +804,6 @@ class Pregel(
|
||||
version = checkpoint["channel_versions"][k]
|
||||
checkpoint["versions_seen"][INTERRUPT][k] = version
|
||||
|
||||
read = partial(_read_channel, channels)
|
||||
|
||||
# Similarly to Bulk Synchronous Parallel / Pregel model
|
||||
# computation proceeds in steps, while there are channel updates
|
||||
# channel updates from step N are only visible in step N+1,
|
||||
@@ -726,7 +816,10 @@ class Pregel(
|
||||
|
||||
# if no more tasks, we're done
|
||||
if not next_tasks:
|
||||
break
|
||||
if step == 0:
|
||||
raise ValueError("No tasks to run in graph.")
|
||||
else:
|
||||
break
|
||||
elif step == config["recursion_limit"]:
|
||||
raise GraphRecursionError(
|
||||
f"Recursion limit of {config['recursion_limit']} reached"
|
||||
@@ -734,6 +827,15 @@ class Pregel(
|
||||
"by setting the `recursion_limit` config key."
|
||||
)
|
||||
|
||||
# before execution, check if we should interrupt
|
||||
if _should_interrupt(
|
||||
checkpoint,
|
||||
interrupt_before_nodes,
|
||||
self.snapshot_channels_list,
|
||||
next_tasks,
|
||||
):
|
||||
break
|
||||
|
||||
if debug:
|
||||
print_step_start(step, next_tasks)
|
||||
|
||||
@@ -749,7 +851,9 @@ class Pregel(
|
||||
configurable={
|
||||
# deque.extend is thread-safe
|
||||
CONFIG_KEY_SEND: writes.extend,
|
||||
CONFIG_KEY_READ: read,
|
||||
CONFIG_KEY_READ: partial(
|
||||
_local_read, checkpoint, channels, writes
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
@@ -785,9 +889,7 @@ class Pregel(
|
||||
pending_writes.extend(writes)
|
||||
|
||||
# apply writes to channels
|
||||
_apply_writes(
|
||||
checkpoint, channels, pending_writes, config, step + 1
|
||||
)
|
||||
_apply_writes(checkpoint, channels, pending_writes)
|
||||
|
||||
if debug:
|
||||
print_checkpoint(step, channels)
|
||||
@@ -802,31 +904,21 @@ class Pregel(
|
||||
if step_output := map_output_updates(output_keys, next_tasks):
|
||||
yield step_output
|
||||
|
||||
# with previous step's checkpoint
|
||||
if _should_interrupt(
|
||||
checkpoint,
|
||||
interrupt_before_nodes,
|
||||
self.snapshot_channels_list,
|
||||
pending_writes,
|
||||
):
|
||||
break
|
||||
|
||||
# save end of step checkpoint
|
||||
if self.checkpointer is not None and (
|
||||
self.checkpointer.at == CheckpointAt.END_OF_STEP
|
||||
or interrupt_before_nodes
|
||||
):
|
||||
checkpoint = create_checkpoint(checkpoint, channels)
|
||||
checkpoint_config = await self.checkpointer.aput(
|
||||
checkpoint_config, checkpoint
|
||||
)
|
||||
|
||||
# with this step's checkpoint
|
||||
# after execution, check if we should interrupt
|
||||
if _should_interrupt(
|
||||
checkpoint,
|
||||
interrupt_after_nodes,
|
||||
self.snapshot_channels_list,
|
||||
pending_writes,
|
||||
next_tasks,
|
||||
):
|
||||
break
|
||||
|
||||
@@ -834,7 +926,6 @@ class Pregel(
|
||||
if (
|
||||
self.checkpointer is not None
|
||||
and self.checkpointer.at == CheckpointAt.END_OF_RUN
|
||||
and not interrupt_before_nodes
|
||||
):
|
||||
checkpoint = create_checkpoint(checkpoint, channels)
|
||||
await self.checkpointer.aput(checkpoint_config, checkpoint)
|
||||
@@ -1041,7 +1132,7 @@ def _should_interrupt(
|
||||
checkpoint: Checkpoint,
|
||||
interrupt_nodes: Sequence[str],
|
||||
snapshot_channels: Sequence[str],
|
||||
pending_writes: Sequence[tuple[str, Any]],
|
||||
tasks: list[PregelExecutableTask],
|
||||
) -> bool:
|
||||
return (
|
||||
# interrupt if any of snapshopt_channels has been updated since last interrupt
|
||||
@@ -1051,10 +1142,32 @@ def _should_interrupt(
|
||||
for chan in snapshot_channels
|
||||
)
|
||||
# and any channel written to is in interrupt_nodes list
|
||||
and any(chan for chan, _ in pending_writes if chan in interrupt_nodes)
|
||||
and any(node for _, _, node, _ in tasks if node in interrupt_nodes)
|
||||
)
|
||||
|
||||
|
||||
def _local_read(
|
||||
checkpoint: Checkpoint,
|
||||
channels: Mapping[str, BaseChannel],
|
||||
writes: Sequence[tuple[str, Any]],
|
||||
select: Union[list[str], str],
|
||||
fresh: bool = False,
|
||||
) -> Union[dict[str, Any], Any]:
|
||||
if fresh:
|
||||
checkpoint = create_checkpoint(checkpoint, channels)
|
||||
with ChannelsManager(channels, checkpoint) as channels:
|
||||
_apply_writes(copy_checkpoint(checkpoint), channels, writes)
|
||||
if isinstance(select, str):
|
||||
return _read_channel(channels, select)
|
||||
else:
|
||||
return {k: _read_channel(channels, k) for k in select}
|
||||
else:
|
||||
if isinstance(select, str):
|
||||
return _read_channel(channels, select)
|
||||
else:
|
||||
return {k: _read_channel(channels, k) for k in select}
|
||||
|
||||
|
||||
def _read_channel(
|
||||
channels: Mapping[str, BaseChannel],
|
||||
chan: str,
|
||||
@@ -1077,20 +1190,17 @@ def _apply_writes(
|
||||
checkpoint: Checkpoint,
|
||||
channels: Mapping[str, BaseChannel],
|
||||
pending_writes: Sequence[tuple[str, Any]],
|
||||
config: RunnableConfig,
|
||||
for_step: int,
|
||||
) -> None:
|
||||
pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
|
||||
# Group writes by channel
|
||||
for chan, val in pending_writes:
|
||||
if chan in AllReservedChannels:
|
||||
raise ValueError(f"Can't write to reserved channel {chan}")
|
||||
pending_writes_by_channel[chan].append(val)
|
||||
|
||||
# Update reserved channels
|
||||
pending_writes_by_channel[ReservedChannels.is_last_step] = [
|
||||
for_step + 1 == config["recursion_limit"]
|
||||
]
|
||||
# Find the highest version of all channels
|
||||
if checkpoint["channel_versions"]:
|
||||
max_version = max(checkpoint["channel_versions"].values())
|
||||
else:
|
||||
max_version = 0
|
||||
|
||||
updated_channels: set[str] = set()
|
||||
# Apply writes to channels
|
||||
@@ -1102,7 +1212,7 @@ def _apply_writes(
|
||||
raise InvalidUpdateError(
|
||||
f"Invalid update for channel {chan}: {e}"
|
||||
) from e
|
||||
checkpoint["channel_versions"][chan] += 1
|
||||
checkpoint["channel_versions"][chan] = max_version + 1
|
||||
updated_channels.add(chan)
|
||||
else:
|
||||
logger.warning(f"Skipping write for channel {chan} which has no readers")
|
||||
@@ -1112,13 +1222,26 @@ def _apply_writes(
|
||||
channels[chan].update([])
|
||||
|
||||
|
||||
class PregelTask(NamedTuple):
|
||||
proc: Runnable
|
||||
input: Any
|
||||
name: str
|
||||
|
||||
|
||||
class PregelExecutableTask(NamedTuple):
|
||||
proc: Runnable
|
||||
input: Any
|
||||
name: str
|
||||
writes: deque[tuple[str, Any]]
|
||||
|
||||
|
||||
@overload
|
||||
def _prepare_next_tasks(
|
||||
checkpoint: Checkpoint,
|
||||
processes: Mapping[str, ChannelInvoke],
|
||||
channels: Mapping[str, BaseChannel],
|
||||
for_execution: Literal[False],
|
||||
) -> tuple[Checkpoint, list[tuple[Runnable, Any, str]]]:
|
||||
) -> tuple[Checkpoint, list[PregelTask]]:
|
||||
...
|
||||
|
||||
|
||||
@@ -1128,7 +1251,7 @@ def _prepare_next_tasks(
|
||||
processes: Mapping[str, ChannelInvoke],
|
||||
channels: Mapping[str, BaseChannel],
|
||||
for_execution: Literal[True],
|
||||
) -> tuple[Checkpoint, list[tuple[Runnable, Any, str, deque[tuple[str, Any]]]]]:
|
||||
) -> tuple[Checkpoint, list[PregelExecutableTask]]:
|
||||
...
|
||||
|
||||
|
||||
@@ -1138,18 +1261,9 @@ def _prepare_next_tasks(
|
||||
channels: Mapping[str, BaseChannel],
|
||||
*,
|
||||
for_execution: bool,
|
||||
) -> tuple[
|
||||
Checkpoint,
|
||||
Union[
|
||||
list[tuple[Runnable, Any, str]],
|
||||
list[tuple[Runnable, Any, str, deque[tuple[str, Any]]]],
|
||||
],
|
||||
]:
|
||||
) -> tuple[Checkpoint, Union[list[PregelTask], list[PregelExecutableTask]]]:
|
||||
checkpoint = copy_checkpoint(checkpoint)
|
||||
tasks: Union[
|
||||
list[tuple[Runnable, Any, str]],
|
||||
list[tuple[Runnable, Any, str, deque[tuple[str, Any]]]],
|
||||
] = []
|
||||
tasks: Union[list[PregelTask], list[PregelExecutableTask]] = []
|
||||
# Check if any processes should be run in next step
|
||||
# If so, prepare the values to be passed to them
|
||||
for name, proc in processes.items():
|
||||
@@ -1164,23 +1278,34 @@ def _prepare_next_tasks(
|
||||
):
|
||||
# If all trigger channels subscribed by this process are not empty
|
||||
# then invoke the process with the values of all non-empty channels
|
||||
try:
|
||||
val: Any = {
|
||||
k: _read_channel(channels, chan, catch=chan not in proc.triggers)
|
||||
for k, chan in proc.channels.items()
|
||||
}
|
||||
except EmptyChannelError:
|
||||
continue
|
||||
if isinstance(proc.channels, dict):
|
||||
try:
|
||||
val: Any = {
|
||||
k: _read_channel(
|
||||
channels, chan, catch=chan not in proc.triggers
|
||||
)
|
||||
for k, chan in proc.channels.items()
|
||||
}
|
||||
except EmptyChannelError:
|
||||
continue
|
||||
elif isinstance(proc.channels, list):
|
||||
for chan in proc.channels:
|
||||
try:
|
||||
val = _read_channel(channels, chan, catch=False)
|
||||
break
|
||||
except EmptyChannelError:
|
||||
pass
|
||||
else:
|
||||
continue
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"Invalid channels type, expected list or dict, got {proc.channels}"
|
||||
)
|
||||
|
||||
# If the process has a mapper, apply it to the value
|
||||
if proc.mapper is not None:
|
||||
val = proc.mapper(val)
|
||||
|
||||
# Processes that subscribe to a single keyless channel get
|
||||
# the value directly, instead of a dict
|
||||
if list(proc.channels.keys()) == [None]:
|
||||
val = val[None]
|
||||
|
||||
# update seen versions
|
||||
if for_execution:
|
||||
seen.update(
|
||||
@@ -1190,12 +1315,19 @@ def _prepare_next_tasks(
|
||||
}
|
||||
)
|
||||
|
||||
# skip if condition is not met
|
||||
if proc.when is None or proc.when(val):
|
||||
if for_execution:
|
||||
tasks.append((proc, val, name, deque()))
|
||||
else:
|
||||
tasks.append((proc, val, name))
|
||||
if for_execution:
|
||||
tasks.append(
|
||||
PregelExecutableTask(
|
||||
RunnableSequence(proc, *proc.writers, name=name)
|
||||
if proc.writers
|
||||
else proc,
|
||||
val,
|
||||
name,
|
||||
deque(),
|
||||
)
|
||||
)
|
||||
else:
|
||||
tasks.append(PregelTask(proc, val, name))
|
||||
return checkpoint, tasks
|
||||
|
||||
|
||||
|
||||
@@ -25,7 +25,7 @@ def print_checkpoint(step: int, channels: Mapping[str, BaseChannel]) -> None:
|
||||
print(
|
||||
f"{get_colored_text('[langgraph/checkpoint]', color='blue')} "
|
||||
+ get_bolded_text(f"Finishing step {step}. Channel values:\n")
|
||||
+ pformat({name: val for name, val in _read_channels(channels)}, depth=2)
|
||||
+ pformat({name: val for name, val in _read_channels(channels)}, depth=3)
|
||||
)
|
||||
|
||||
|
||||
|
||||
+45
-25
@@ -15,11 +15,16 @@ from langchain_core.runnables.config import merge_configs
|
||||
from langchain_core.runnables.utils import ConfigurableFieldSpec
|
||||
|
||||
from langgraph.constants import CONFIG_KEY_READ
|
||||
from langgraph.pregel.write import ChannelWrite
|
||||
|
||||
READ_TYPE = Callable[[str, bool], Union[Any, dict[str, Any]]]
|
||||
|
||||
|
||||
class ChannelRead(RunnableLambda):
|
||||
channel: Union[str, list[str]]
|
||||
|
||||
fresh: bool = False
|
||||
|
||||
@property
|
||||
def config_specs(self) -> list[ConfigurableFieldSpec]:
|
||||
return [
|
||||
@@ -32,51 +37,44 @@ class ChannelRead(RunnableLambda):
|
||||
),
|
||||
]
|
||||
|
||||
def __init__(self, channel: Union[str, list[str]]) -> None:
|
||||
def __init__(self, channel: Union[str, list[str]], fresh: bool = False) -> None:
|
||||
super().__init__(func=self._read, afunc=self._aread)
|
||||
self.fresh = fresh
|
||||
self.channel = channel
|
||||
self.name = f"ChannelRead<{channel}>"
|
||||
|
||||
def _read(self, _: Any, config: RunnableConfig) -> Any:
|
||||
try:
|
||||
read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ]
|
||||
read: READ_TYPE = config["configurable"][CONFIG_KEY_READ]
|
||||
except KeyError:
|
||||
raise RuntimeError(
|
||||
f"Runnable {self} is not configured with a read function"
|
||||
"Make sure to call in the context of a Pregel process"
|
||||
)
|
||||
return (
|
||||
read(self.channel)
|
||||
if isinstance(self.channel, str)
|
||||
else {chan: read(chan) for chan in self.channel}
|
||||
)
|
||||
return read(self.channel, self.fresh)
|
||||
|
||||
async def _aread(self, _: Any, config: RunnableConfig) -> Any:
|
||||
try:
|
||||
read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ]
|
||||
read: READ_TYPE = config["configurable"][CONFIG_KEY_READ]
|
||||
except KeyError:
|
||||
raise RuntimeError(
|
||||
f"Runnable {self} is not configured with a read function"
|
||||
"Make sure to call in the context of a Pregel process"
|
||||
)
|
||||
return (
|
||||
read(self.channel)
|
||||
if isinstance(self.channel, str)
|
||||
else {chan: read(chan) for chan in self.channel}
|
||||
)
|
||||
return read(self.channel, self.fresh)
|
||||
|
||||
|
||||
default_bound: RunnablePassthrough = RunnablePassthrough()
|
||||
|
||||
|
||||
class ChannelInvoke(RunnableBindingBase):
|
||||
channels: Union[Mapping[None, str], Mapping[str, str]]
|
||||
channels: Union[list[str], Mapping[str, str]]
|
||||
|
||||
triggers: list[str] = Field(default_factory=list)
|
||||
|
||||
mapper: Optional[Callable[[Any], Any]] = None
|
||||
|
||||
when: Optional[Callable[[Any], bool]] = None
|
||||
writers: list[Runnable] = Field(default_factory=list)
|
||||
|
||||
bound: Runnable[Any, Any] = Field(default=default_bound)
|
||||
|
||||
@@ -84,12 +82,12 @@ class ChannelInvoke(RunnableBindingBase):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: Mapping[None, str] | Mapping[str, str],
|
||||
*,
|
||||
channels: Union[list[str], Mapping[str, str]],
|
||||
triggers: Sequence[str],
|
||||
mapper: Optional[Callable[[Any], Any]] = None,
|
||||
when: Optional[Callable[[Any], bool]] = None,
|
||||
writers: Optional[list[Runnable]] = None,
|
||||
tags: Optional[list[str]] = None,
|
||||
*,
|
||||
bound: Optional[Runnable[Any, Any]] = None,
|
||||
kwargs: Optional[Mapping[str, Any]] = None,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
@@ -99,19 +97,22 @@ class ChannelInvoke(RunnableBindingBase):
|
||||
channels=channels,
|
||||
triggers=triggers,
|
||||
mapper=mapper,
|
||||
when=when,
|
||||
writers=writers or [],
|
||||
bound=bound or default_bound,
|
||||
kwargs=kwargs or {},
|
||||
config=merge_configs(config, {"tags": tags or []}),
|
||||
**other_kwargs,
|
||||
)
|
||||
|
||||
def __repr_args__(self) -> Any:
|
||||
return [(k, v) for k, v in super().__repr_args__() if k != "bound"]
|
||||
|
||||
def join(self, channels: Sequence[str]) -> ChannelInvoke:
|
||||
assert isinstance(channels, list) or isinstance(
|
||||
channels, tuple
|
||||
), "channels must be a list or tuple"
|
||||
assert all(
|
||||
k is not None for k in self.channels.keys()
|
||||
assert isinstance(
|
||||
self.channels, dict
|
||||
), "all channels must be named when using .join()"
|
||||
return ChannelInvoke(
|
||||
channels={
|
||||
@@ -120,7 +121,7 @@ class ChannelInvoke(RunnableBindingBase):
|
||||
},
|
||||
triggers=self.triggers,
|
||||
mapper=self.mapper,
|
||||
when=self.when,
|
||||
writers=self.writers,
|
||||
bound=self.bound,
|
||||
kwargs=self.kwargs,
|
||||
config=self.config,
|
||||
@@ -134,12 +135,22 @@ class ChannelInvoke(RunnableBindingBase):
|
||||
Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
|
||||
],
|
||||
) -> ChannelInvoke:
|
||||
if self.bound is default_bound:
|
||||
if ChannelWrite.is_writer(other):
|
||||
return ChannelInvoke(
|
||||
channels=self.channels,
|
||||
triggers=self.triggers,
|
||||
mapper=self.mapper,
|
||||
when=self.when,
|
||||
writers=[*self.writers, other],
|
||||
bound=self.bound,
|
||||
kwargs=self.kwargs,
|
||||
config=self.config,
|
||||
)
|
||||
elif self.bound is default_bound:
|
||||
return ChannelInvoke(
|
||||
channels=self.channels,
|
||||
triggers=self.triggers,
|
||||
mapper=self.mapper,
|
||||
writers=self.writers,
|
||||
bound=coerce_to_runnable(other),
|
||||
kwargs=self.kwargs,
|
||||
config=self.config,
|
||||
@@ -149,13 +160,22 @@ class ChannelInvoke(RunnableBindingBase):
|
||||
channels=self.channels,
|
||||
triggers=self.triggers,
|
||||
mapper=self.mapper,
|
||||
when=self.when,
|
||||
writers=self.writers,
|
||||
# delegate to __or__ in self.bound
|
||||
bound=self.bound | other,
|
||||
kwargs=self.kwargs,
|
||||
config=self.config,
|
||||
)
|
||||
|
||||
def pipe(
|
||||
self,
|
||||
*others: Runnable[Any, Other] | Callable[[Any], Other],
|
||||
name: Optional[str] = None,
|
||||
) -> RunnableSerializable[Any, Other]:
|
||||
for other in others:
|
||||
self = self | other
|
||||
return self
|
||||
|
||||
def __ror__(
|
||||
self,
|
||||
other: Union[
|
||||
|
||||
@@ -1,11 +0,0 @@
|
||||
from langgraph.utils import StrEnum
|
||||
|
||||
|
||||
class ReservedChannels(StrEnum):
|
||||
"""Channels managed by the framework."""
|
||||
|
||||
is_last_step = "is_last_step"
|
||||
"""A channel that is True if the current step is the last step, False otherwise."""
|
||||
|
||||
|
||||
AllReservedChannels = {channel.value for channel in ReservedChannels}
|
||||
@@ -1,10 +1,8 @@
|
||||
from typing import Any, Mapping, Optional, Sequence, Union
|
||||
from typing import Any, Mapping, Optional, Sequence, Type, Union
|
||||
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.constants import INTERRUPT
|
||||
from langgraph.pregel.read import ChannelInvoke
|
||||
from langgraph.pregel.reserved import ReservedChannels
|
||||
|
||||
|
||||
def validate_graph(
|
||||
@@ -15,13 +13,14 @@ def validate_graph(
|
||||
stream_channels: Optional[Union[str, Sequence[str]]],
|
||||
interrupt_after_nodes: Sequence[str],
|
||||
interrupt_before_nodes: Sequence[str],
|
||||
default_channel_cls: Type[BaseChannel],
|
||||
) -> None:
|
||||
subscribed_channels = set[str]()
|
||||
for name, node in nodes.items():
|
||||
if name == INTERRUPT:
|
||||
raise ValueError(f"Node name {INTERRUPT} is reserved")
|
||||
if isinstance(node, ChannelInvoke):
|
||||
subscribed_channels.update(node.channels.values())
|
||||
subscribed_channels.update(node.triggers)
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Invalid node type {type(node)}, expected Channel.subscribe_to()"
|
||||
@@ -29,11 +28,11 @@ def validate_graph(
|
||||
|
||||
for chan in subscribed_channels:
|
||||
if chan not in channels:
|
||||
channels[chan] = LastValue(Any) # type: ignore[arg-type]
|
||||
channels[chan] = default_channel_cls(Any) # type: ignore[arg-type]
|
||||
|
||||
if isinstance(input_channels, str):
|
||||
if input_channels not in channels:
|
||||
channels[input_channels] = LastValue(Any) # type: ignore[arg-type]
|
||||
channels[input_channels] = default_channel_cls(Any) # type: ignore[arg-type]
|
||||
if input_channels not in subscribed_channels:
|
||||
raise ValueError(
|
||||
f"Input channel {input_channels} is not subscribed to by any node"
|
||||
@@ -41,27 +40,32 @@ def validate_graph(
|
||||
else:
|
||||
for chan in input_channels:
|
||||
if chan not in channels:
|
||||
channels[chan] = LastValue(Any) # type: ignore[arg-type]
|
||||
channels[chan] = default_channel_cls(Any) # type: ignore[arg-type]
|
||||
if all(chan not in subscribed_channels for chan in input_channels):
|
||||
raise ValueError(
|
||||
f"None of the input channels {input_channels} are subscribed to by any node"
|
||||
)
|
||||
|
||||
all_output_channels = set[str]()
|
||||
if isinstance(output_channels, str):
|
||||
if output_channels not in channels:
|
||||
channels[output_channels] = LastValue(Any) # type: ignore[arg-type]
|
||||
all_output_channels.add(output_channels)
|
||||
else:
|
||||
for chan in output_channels:
|
||||
if chan not in channels:
|
||||
channels[chan] = LastValue(Any) # type: ignore[arg-type]
|
||||
all_output_channels.update(output_channels)
|
||||
if isinstance(stream_channels, str):
|
||||
all_output_channels.add(stream_channels)
|
||||
elif stream_channels is not None:
|
||||
all_output_channels.update(stream_channels)
|
||||
|
||||
for chan in ReservedChannels:
|
||||
for chan in all_output_channels:
|
||||
if chan not in channels:
|
||||
channels[chan] = LastValue(Any) # type: ignore[arg-type]
|
||||
channels[chan] = default_channel_cls(Any) # type: ignore[arg-type]
|
||||
|
||||
validate_keys(stream_channels, channels)
|
||||
validate_keys(interrupt_after_nodes, channels)
|
||||
validate_keys(interrupt_before_nodes, channels)
|
||||
for node in interrupt_after_nodes:
|
||||
if node not in nodes:
|
||||
raise ValueError(f"Node {node} not in nodes")
|
||||
for node in interrupt_before_nodes:
|
||||
if node not in nodes:
|
||||
raise ValueError(f"Node {node} not in nodes")
|
||||
|
||||
|
||||
def validate_keys(
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any, Callable, NamedTuple, Optional, Sequence, Union
|
||||
from typing import Any, Callable, NamedTuple, Optional, Sequence, TypeVar, Union
|
||||
|
||||
from langchain_core.runnables import (
|
||||
Runnable,
|
||||
@@ -13,6 +13,7 @@ from langchain_core.runnables.utils import ConfigurableFieldSpec
|
||||
from langgraph.constants import CONFIG_KEY_SEND
|
||||
|
||||
TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
|
||||
R = TypeVar("R", bound=Runnable)
|
||||
|
||||
|
||||
SKIP_WRITE = object()
|
||||
@@ -36,13 +37,17 @@ class ChannelWrite(RunnablePassthrough):
|
||||
class Config:
|
||||
arbitrary_types_allowed = True
|
||||
|
||||
def __init__(self, *, channels: Sequence[ChannelWriteEntry]):
|
||||
def __init__(self, channels: Sequence[ChannelWriteEntry]):
|
||||
super().__init__(func=self._write, afunc=self._awrite, channels=channels)
|
||||
self.name = f"ChannelWrite<{','.join(chan for chan, _, _ in self.channels)}>"
|
||||
|
||||
def __repr_args__(self) -> Any:
|
||||
return [("channels", self.channels)]
|
||||
|
||||
@property
|
||||
def is_channel_writer(self) -> bool:
|
||||
return True
|
||||
|
||||
@property
|
||||
def config_specs(self) -> list[ConfigurableFieldSpec]:
|
||||
return [
|
||||
@@ -99,6 +104,18 @@ class ChannelWrite(RunnablePassthrough):
|
||||
write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND]
|
||||
write([(chan, val) for chan, val in values.items() if val is not SKIP_WRITE])
|
||||
|
||||
@staticmethod
|
||||
def is_writer(runnable: Runnable) -> bool:
|
||||
return (
|
||||
isinstance(runnable, ChannelWrite)
|
||||
or getattr(runnable, "_is_channel_writer", False) is True
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def register_writer(runnable: R) -> R:
|
||||
object.__setattr__(runnable, "_is_channel_writer", True)
|
||||
return runnable
|
||||
|
||||
|
||||
def _mk_future(val: Any) -> asyncio.Future:
|
||||
fut = asyncio.Future()
|
||||
|
||||
@@ -50,20 +50,7 @@
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "left_condition",
|
||||
"type": "runnable",
|
||||
"data": {
|
||||
"id": [
|
||||
"langchain_core",
|
||||
"runnables",
|
||||
"base",
|
||||
"RunnableLambda"
|
||||
],
|
||||
"name": "RunnableLambda"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "should_start",
|
||||
"id": "__start___should_start",
|
||||
"type": "runnable",
|
||||
"data": {
|
||||
"id": [
|
||||
@@ -74,6 +61,19 @@
|
||||
],
|
||||
"name": "should_start"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "left_left_condition",
|
||||
"type": "runnable",
|
||||
"data": {
|
||||
"id": [
|
||||
"langchain_core",
|
||||
"runnables",
|
||||
"base",
|
||||
"RunnableLambda"
|
||||
],
|
||||
"name": "left_condition"
|
||||
}
|
||||
}
|
||||
],
|
||||
"edges": [
|
||||
@@ -81,38 +81,38 @@
|
||||
"source": "right",
|
||||
"target": "__end__"
|
||||
},
|
||||
{
|
||||
"source": "left",
|
||||
"target": "left_condition"
|
||||
},
|
||||
{
|
||||
"source": "left_condition",
|
||||
"target": "left",
|
||||
"data": "left"
|
||||
},
|
||||
{
|
||||
"source": "left_condition",
|
||||
"target": "right",
|
||||
"data": "right"
|
||||
},
|
||||
{
|
||||
"source": "left_condition",
|
||||
"target": "__end__",
|
||||
"data": "__end__"
|
||||
},
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": "should_start"
|
||||
"target": "__start___should_start"
|
||||
},
|
||||
{
|
||||
"source": "should_start",
|
||||
"source": "__start___should_start",
|
||||
"target": "left",
|
||||
"data": "go-left"
|
||||
},
|
||||
{
|
||||
"source": "should_start",
|
||||
"source": "__start___should_start",
|
||||
"target": "right",
|
||||
"data": "go-right"
|
||||
},
|
||||
{
|
||||
"source": "left",
|
||||
"target": "left_left_condition"
|
||||
},
|
||||
{
|
||||
"source": "left_left_condition",
|
||||
"target": "left",
|
||||
"data": "left"
|
||||
},
|
||||
{
|
||||
"source": "left_left_condition",
|
||||
"target": "right",
|
||||
"data": "right"
|
||||
},
|
||||
{
|
||||
"source": "left_left_condition",
|
||||
"target": "__end__",
|
||||
"data": "__end__"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -120,39 +120,39 @@
|
||||
# ---
|
||||
# name: test_conditional_entrypoint_graph.3
|
||||
'''
|
||||
+-----------+
|
||||
| __start__ |
|
||||
+-----------+
|
||||
*
|
||||
*
|
||||
*
|
||||
+--------------+
|
||||
| should_start |
|
||||
+--------------+
|
||||
*** **
|
||||
* **
|
||||
** **
|
||||
+------+ *
|
||||
| left | *
|
||||
+------+ *
|
||||
* *
|
||||
* *
|
||||
* *
|
||||
+----------------+ *
|
||||
| left_condition | *
|
||||
+----------------+* *
|
||||
* ***** *
|
||||
* *** *
|
||||
* *** *
|
||||
** +-------+
|
||||
* | right |
|
||||
*** +-------+
|
||||
* ***
|
||||
*** *
|
||||
* **
|
||||
+---------+
|
||||
| __end__ |
|
||||
+---------+
|
||||
+-----------+
|
||||
| __start__ |
|
||||
+-----------+
|
||||
*
|
||||
*
|
||||
*
|
||||
+------------------------+
|
||||
| __start___should_start |
|
||||
+------------------------+
|
||||
*** **
|
||||
* **
|
||||
** **
|
||||
+------+ **
|
||||
| left | *
|
||||
+------+ *
|
||||
* *
|
||||
* *
|
||||
* *
|
||||
+---------------------+ *
|
||||
| left_left_condition | *
|
||||
+---------------------+ *
|
||||
* ***** *
|
||||
* **** *
|
||||
* *** *
|
||||
** +-------+
|
||||
* | right |
|
||||
*** +-------+
|
||||
* ***
|
||||
*** *
|
||||
* **
|
||||
+---------+
|
||||
| __end__ |
|
||||
+---------+
|
||||
'''
|
||||
# ---
|
||||
# name: test_conditional_entrypoint_graph_state
|
||||
@@ -234,20 +234,7 @@
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "left_condition",
|
||||
"type": "runnable",
|
||||
"data": {
|
||||
"id": [
|
||||
"langchain_core",
|
||||
"runnables",
|
||||
"base",
|
||||
"RunnableLambda"
|
||||
],
|
||||
"name": "RunnableLambda"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "should_start",
|
||||
"id": "__start___should_start",
|
||||
"type": "runnable",
|
||||
"data": {
|
||||
"id": [
|
||||
@@ -258,6 +245,19 @@
|
||||
],
|
||||
"name": "should_start"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "left_left_condition",
|
||||
"type": "runnable",
|
||||
"data": {
|
||||
"id": [
|
||||
"langchain_core",
|
||||
"runnables",
|
||||
"base",
|
||||
"RunnableLambda"
|
||||
],
|
||||
"name": "left_condition"
|
||||
}
|
||||
}
|
||||
],
|
||||
"edges": [
|
||||
@@ -265,38 +265,38 @@
|
||||
"source": "right",
|
||||
"target": "__end__"
|
||||
},
|
||||
{
|
||||
"source": "left",
|
||||
"target": "left_condition"
|
||||
},
|
||||
{
|
||||
"source": "left_condition",
|
||||
"target": "left",
|
||||
"data": "left"
|
||||
},
|
||||
{
|
||||
"source": "left_condition",
|
||||
"target": "right",
|
||||
"data": "right"
|
||||
},
|
||||
{
|
||||
"source": "left_condition",
|
||||
"target": "__end__",
|
||||
"data": "__end__"
|
||||
},
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": "should_start"
|
||||
"target": "__start___should_start"
|
||||
},
|
||||
{
|
||||
"source": "should_start",
|
||||
"source": "__start___should_start",
|
||||
"target": "left",
|
||||
"data": "go-left"
|
||||
},
|
||||
{
|
||||
"source": "should_start",
|
||||
"source": "__start___should_start",
|
||||
"target": "right",
|
||||
"data": "go-right"
|
||||
},
|
||||
{
|
||||
"source": "left",
|
||||
"target": "left_left_condition"
|
||||
},
|
||||
{
|
||||
"source": "left_left_condition",
|
||||
"target": "left",
|
||||
"data": "left"
|
||||
},
|
||||
{
|
||||
"source": "left_left_condition",
|
||||
"target": "right",
|
||||
"data": "right"
|
||||
},
|
||||
{
|
||||
"source": "left_left_condition",
|
||||
"target": "__end__",
|
||||
"data": "__end__"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -304,39 +304,39 @@
|
||||
# ---
|
||||
# name: test_conditional_entrypoint_graph_state.3
|
||||
'''
|
||||
+-----------+
|
||||
| __start__ |
|
||||
+-----------+
|
||||
*
|
||||
*
|
||||
*
|
||||
+--------------+
|
||||
| should_start |
|
||||
+--------------+
|
||||
*** **
|
||||
* **
|
||||
** **
|
||||
+------+ *
|
||||
| left | *
|
||||
+------+ *
|
||||
* *
|
||||
* *
|
||||
* *
|
||||
+----------------+ *
|
||||
| left_condition | *
|
||||
+----------------+* *
|
||||
* ***** *
|
||||
* *** *
|
||||
* *** *
|
||||
** +-------+
|
||||
* | right |
|
||||
*** +-------+
|
||||
* ***
|
||||
*** *
|
||||
* **
|
||||
+---------+
|
||||
| __end__ |
|
||||
+---------+
|
||||
+-----------+
|
||||
| __start__ |
|
||||
+-----------+
|
||||
*
|
||||
*
|
||||
*
|
||||
+------------------------+
|
||||
| __start___should_start |
|
||||
+------------------------+
|
||||
*** **
|
||||
* **
|
||||
** **
|
||||
+------+ **
|
||||
| left | *
|
||||
+------+ *
|
||||
* *
|
||||
* *
|
||||
* *
|
||||
+---------------------+ *
|
||||
| left_left_condition | *
|
||||
+---------------------+ *
|
||||
* ***** *
|
||||
* **** *
|
||||
* *** *
|
||||
** +-------+
|
||||
* | right |
|
||||
*** +-------+
|
||||
* ***
|
||||
*** *
|
||||
* **
|
||||
+---------+
|
||||
| __end__ |
|
||||
+---------+
|
||||
'''
|
||||
# ---
|
||||
# name: test_conditional_graph
|
||||
@@ -398,6 +398,10 @@
|
||||
}
|
||||
],
|
||||
"edges": [
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": "agent"
|
||||
},
|
||||
{
|
||||
"source": "tools",
|
||||
"target": "agent"
|
||||
@@ -415,10 +419,6 @@
|
||||
"source": "agent_should_continue",
|
||||
"target": "__end__",
|
||||
"data": "exit"
|
||||
},
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": "agent"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -669,6 +669,10 @@
|
||||
"source": 7,
|
||||
"target": 3
|
||||
},
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": 2
|
||||
},
|
||||
{
|
||||
"source": "tools",
|
||||
"target": 2
|
||||
@@ -686,10 +690,6 @@
|
||||
"source": "agent_should_continue",
|
||||
"target": "__end__",
|
||||
"data": "exit"
|
||||
},
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": 2
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -1019,6 +1019,10 @@
|
||||
}
|
||||
],
|
||||
"edges": [
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": "agent"
|
||||
},
|
||||
{
|
||||
"source": "tools",
|
||||
"target": "agent"
|
||||
@@ -1036,10 +1040,6 @@
|
||||
"source": "agent_should_continue",
|
||||
"target": "__end__",
|
||||
"data": "exit"
|
||||
},
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": "agent"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -1842,6 +1842,10 @@
|
||||
"source": "action",
|
||||
"target": "agent"
|
||||
},
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": "agent"
|
||||
},
|
||||
{
|
||||
"source": "agent",
|
||||
"target": "agent_should_continue"
|
||||
@@ -1855,10 +1859,6 @@
|
||||
"source": "agent_should_continue",
|
||||
"target": "__end__",
|
||||
"data": "end"
|
||||
},
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": "agent"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -2085,6 +2085,10 @@
|
||||
"source": "action",
|
||||
"target": "agent"
|
||||
},
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": "agent"
|
||||
},
|
||||
{
|
||||
"source": "agent",
|
||||
"target": "agent_should_continue"
|
||||
@@ -2098,10 +2102,6 @@
|
||||
"source": "agent_should_continue",
|
||||
"target": "__end__",
|
||||
"data": "end"
|
||||
},
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": "agent"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -2328,6 +2328,10 @@
|
||||
"source": "action",
|
||||
"target": "agent"
|
||||
},
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": "agent"
|
||||
},
|
||||
{
|
||||
"source": "agent",
|
||||
"target": "agent_should_continue"
|
||||
@@ -2341,10 +2345,6 @@
|
||||
"source": "agent_should_continue",
|
||||
"target": "__end__",
|
||||
"data": "end"
|
||||
},
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": "agent"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
class AnyStr(str):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
return isinstance(other, str)
|
||||
+189
-147
@@ -26,7 +26,7 @@ from langgraph.prebuilt.chat_agent_executor import (
|
||||
)
|
||||
from langgraph.prebuilt.tool_executor import ToolExecutor
|
||||
from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
|
||||
from langgraph.pregel.reserved import ReservedChannels
|
||||
from tests.any_str import AnyStr
|
||||
from tests.memory_assert import MemorySaverAssertImmutable
|
||||
|
||||
|
||||
@@ -64,7 +64,7 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
|
||||
assert app.invoke(2, output_keys=["output"]) == {"output": 3}
|
||||
assert repr(app), "does not raise recursion error"
|
||||
|
||||
assert gapp.invoke(2) == 3
|
||||
assert gapp.invoke(2, debug=True) == 3
|
||||
|
||||
|
||||
def test_invoke_single_process_in_out_implicit_channels(mocker: MockerFixture) -> None:
|
||||
@@ -103,23 +103,6 @@ def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None:
|
||||
assert app.invoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4}
|
||||
|
||||
|
||||
def test_invoke_single_process_in_out_reserved_is_last(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: {**x, "input": x["input"] + 1})
|
||||
|
||||
chain = (
|
||||
Channel.subscribe_to(["input"]).join([ReservedChannels.is_last_step])
|
||||
| add_one
|
||||
| Channel.write_to("output")
|
||||
)
|
||||
|
||||
app = Pregel(nodes={"one": chain})
|
||||
|
||||
assert app.input_schema.schema() == {"title": "LangGraphInput"}
|
||||
assert app.output_schema.schema() == {"title": "LangGraphOutput"}
|
||||
assert app.invoke(2) == {"input": 3, "is_last_step": False}
|
||||
assert app.invoke(2, {"recursion_limit": 1}) == {"input": 3, "is_last_step": True}
|
||||
|
||||
|
||||
def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
|
||||
@@ -214,7 +197,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
checkpointer=memory,
|
||||
interrupt_after_nodes=["inbox"],
|
||||
interrupt_after_nodes=["one"],
|
||||
)
|
||||
|
||||
# start execution, stop at inbox
|
||||
@@ -249,7 +232,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
|
||||
assert snapshot.next == ("two",)
|
||||
|
||||
# update the state, resume
|
||||
app.update_state({"configurable": {"thread_id": 2}}, {"inbox": 25})
|
||||
app.update_state({"configurable": {"thread_id": 2}}, 25, as_node="one")
|
||||
assert app.invoke(None, {"configurable": {"thread_id": 2}}) == 26
|
||||
|
||||
# no pending tasks
|
||||
@@ -507,9 +490,11 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
|
||||
state = app.get_state({"configurable": {"thread_id": "1"}})
|
||||
assert state is not None
|
||||
assert state.values.get("total") == 7
|
||||
assert state.next == ()
|
||||
state = app.get_state(thread_2)
|
||||
assert state is not None
|
||||
assert state.values.get("total") == 5
|
||||
assert state.next == ()
|
||||
|
||||
# list all checkpoints for thread 1
|
||||
thread_1_history = [c for c in app.get_state_history(thread_1)]
|
||||
@@ -534,9 +519,7 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
|
||||
== thread_1_history[1].config["configurable"]["thread_ts"]
|
||||
)
|
||||
|
||||
thread_1_next_config = app.update_state(
|
||||
thread_1_history[1].config, {"total": 10}
|
||||
)
|
||||
thread_1_next_config = app.update_state(thread_1_history[1].config, 10)
|
||||
# update creates a new checkpoint
|
||||
assert (
|
||||
thread_1_next_config["configurable"]["thread_ts"]
|
||||
@@ -927,9 +910,8 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
},
|
||||
"tools": None,
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
assert (
|
||||
@@ -942,14 +924,12 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
|
||||
app_w_interrupt.update_state(
|
||||
config,
|
||||
{
|
||||
"agent": {
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
)
|
||||
|
||||
@@ -963,9 +943,8 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
"tools": None,
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
@@ -1010,6 +989,26 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
|
||||
app_w_interrupt.update_state(
|
||||
config,
|
||||
{
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
(
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
)
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
@@ -1026,12 +1025,12 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
}
|
||||
},
|
||||
},
|
||||
next=(),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
assert [c for c in app_w_interrupt.stream(None, config)] == []
|
||||
|
||||
# test state get/update methods with interrupt_before
|
||||
|
||||
app_w_interrupt = workflow.compile(
|
||||
@@ -1061,23 +1060,20 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
},
|
||||
"tools": None,
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
app_w_interrupt.update_state(
|
||||
config,
|
||||
{
|
||||
"agent": {
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
)
|
||||
|
||||
@@ -1091,9 +1087,8 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
"tools": None,
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
@@ -1138,6 +1133,26 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
|
||||
app_w_interrupt.update_state(
|
||||
config,
|
||||
{
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
(
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
)
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
@@ -1154,12 +1169,12 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
}
|
||||
},
|
||||
},
|
||||
next=(),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
assert [c for c in app_w_interrupt.stream(None, config)] == []
|
||||
|
||||
# test re-invoke to continue with interrupt_before
|
||||
|
||||
app_w_interrupt = workflow.compile(
|
||||
@@ -1189,9 +1204,8 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
},
|
||||
"tools": None,
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
@@ -1286,6 +1300,50 @@ def test_conditional_graph(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", debug=True)
|
||||
== "what is weather in sf->right"
|
||||
)
|
||||
|
||||
assert [*app.stream("what is weather in sf")] == [
|
||||
{"right": "what is weather in sf->right"},
|
||||
]
|
||||
|
||||
|
||||
def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
|
||||
from langchain.llms.fake import FakeStreamingListLLM
|
||||
from langchain_community.tools import tool
|
||||
@@ -1479,7 +1537,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
|
||||
),
|
||||
"intermediate_steps": [],
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
@@ -1504,7 +1562,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
|
||||
),
|
||||
"intermediate_steps": [],
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
@@ -1544,7 +1602,27 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
|
||||
},
|
||||
)
|
||||
|
||||
assert [c for c in app_w_interrupt.stream(None, config)] == []
|
||||
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
||||
values={
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
"intermediate_steps": [
|
||||
(
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
)
|
||||
],
|
||||
},
|
||||
next=(),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
# test state get/update methods with interrupt_before
|
||||
|
||||
@@ -1577,7 +1655,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
|
||||
),
|
||||
"intermediate_steps": [],
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
@@ -1602,7 +1680,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
|
||||
),
|
||||
"intermediate_steps": [],
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
@@ -1642,49 +1720,28 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
|
||||
},
|
||||
)
|
||||
|
||||
assert [c for c in app_w_interrupt.stream(None, config)] == []
|
||||
|
||||
|
||||
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"}
|
||||
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
||||
values={
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
"intermediate_steps": [
|
||||
(
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
)
|
||||
],
|
||||
},
|
||||
next=(),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
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"},
|
||||
]
|
||||
|
||||
|
||||
def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None:
|
||||
class AgentState(TypedDict, total=False):
|
||||
@@ -2200,7 +2257,7 @@ def test_message_graph(
|
||||
assert app.invoke(HumanMessage(content="what is weather in sf")) == [
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000002", # adds missing ids
|
||||
id="00000000-0000-4000-8000-000000000005", # adds missing ids
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2212,7 +2269,7 @@ def test_message_graph(
|
||||
FunctionMessage(
|
||||
content="result for query",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000015",
|
||||
id="00000000-0000-4000-8000-000000000019",
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2224,7 +2281,7 @@ def test_message_graph(
|
||||
FunctionMessage(
|
||||
content="result for another",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000028",
|
||||
id="00000000-0000-4000-8000-000000000033",
|
||||
),
|
||||
AIMessage(content="answer", id="ai3"),
|
||||
]
|
||||
@@ -2234,7 +2291,7 @@ def test_message_graph(
|
||||
"__start__": [
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000037",
|
||||
id="00000000-0000-4000-8000-000000000045",
|
||||
)
|
||||
]
|
||||
},
|
||||
@@ -2251,7 +2308,7 @@ def test_message_graph(
|
||||
"action": FunctionMessage(
|
||||
content="result for query",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000050",
|
||||
id="00000000-0000-4000-8000-000000000059",
|
||||
)
|
||||
},
|
||||
{
|
||||
@@ -2267,7 +2324,7 @@ def test_message_graph(
|
||||
"action": FunctionMessage(
|
||||
content="result for another",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000063",
|
||||
id="00000000-0000-4000-8000-000000000073",
|
||||
)
|
||||
},
|
||||
{"agent": AIMessage(content="answer", id="ai3")},
|
||||
@@ -2284,12 +2341,7 @@ def test_message_graph(
|
||||
HumanMessage(content="what is weather in sf"), config
|
||||
)
|
||||
] == [
|
||||
{
|
||||
"__start__": HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000072",
|
||||
)
|
||||
},
|
||||
{"__start__": HumanMessage(content="what is weather in sf", id=AnyStr())},
|
||||
{
|
||||
"agent": AIMessage(
|
||||
content="",
|
||||
@@ -2303,10 +2355,7 @@ def test_message_graph(
|
||||
|
||||
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
||||
values=[
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000072",
|
||||
),
|
||||
HumanMessage(content="what is weather in sf", id=AnyStr()),
|
||||
AIMessage(
|
||||
content="",
|
||||
additional_kwargs={
|
||||
@@ -2315,7 +2364,7 @@ def test_message_graph(
|
||||
id="ai1",
|
||||
),
|
||||
],
|
||||
next=("agent:edges",),
|
||||
next=("action",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
@@ -2327,10 +2376,7 @@ def test_message_graph(
|
||||
# message was replaced instead of appended
|
||||
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
||||
values=[
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000072",
|
||||
),
|
||||
HumanMessage(content="what is weather in sf", id=AnyStr()),
|
||||
AIMessage(
|
||||
content="",
|
||||
additional_kwargs={
|
||||
@@ -2342,7 +2388,7 @@ def test_message_graph(
|
||||
id="ai1",
|
||||
),
|
||||
],
|
||||
next=("agent:edges",),
|
||||
next=("action",),
|
||||
config=next_config,
|
||||
)
|
||||
|
||||
@@ -2351,7 +2397,7 @@ def test_message_graph(
|
||||
"action": FunctionMessage(
|
||||
content="result for a different query",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000086",
|
||||
id=AnyStr(),
|
||||
)
|
||||
},
|
||||
{
|
||||
@@ -2369,7 +2415,7 @@ def test_message_graph(
|
||||
values=[
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000072",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2384,7 +2430,7 @@ def test_message_graph(
|
||||
FunctionMessage(
|
||||
content="result for a different query",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000086",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2394,7 +2440,7 @@ def test_message_graph(
|
||||
id="ai2",
|
||||
),
|
||||
],
|
||||
next=("agent:edges",),
|
||||
next=("action",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
@@ -2408,7 +2454,7 @@ def test_message_graph(
|
||||
values=[
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000072",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2423,16 +2469,14 @@ def test_message_graph(
|
||||
FunctionMessage(
|
||||
content="result for a different query",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000086",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(content="answer", id="ai2"),
|
||||
],
|
||||
next=("agent:edges",),
|
||||
next=(),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
assert [c for c in app_w_interrupt.stream(None, config)] == []
|
||||
|
||||
app_w_interrupt = workflow.compile(
|
||||
checkpointer=MemorySaverAssertImmutable(), interrupt_before=["action"]
|
||||
)
|
||||
@@ -2448,7 +2492,7 @@ def test_message_graph(
|
||||
{
|
||||
"__start__": HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000096",
|
||||
id=AnyStr(),
|
||||
)
|
||||
},
|
||||
{
|
||||
@@ -2466,7 +2510,7 @@ def test_message_graph(
|
||||
values=[
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000096",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2476,7 +2520,7 @@ def test_message_graph(
|
||||
id="ai1",
|
||||
),
|
||||
],
|
||||
next=("agent:edges",),
|
||||
next=("action",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
@@ -2490,7 +2534,7 @@ def test_message_graph(
|
||||
values=[
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000096",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2503,7 +2547,7 @@ def test_message_graph(
|
||||
id="ai1",
|
||||
),
|
||||
],
|
||||
next=("agent:edges",),
|
||||
next=("action",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
@@ -2512,7 +2556,7 @@ def test_message_graph(
|
||||
"action": FunctionMessage(
|
||||
content="result for a different query",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000113",
|
||||
id=AnyStr(),
|
||||
)
|
||||
},
|
||||
{
|
||||
@@ -2530,7 +2574,7 @@ def test_message_graph(
|
||||
values=[
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000096",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2545,7 +2589,7 @@ def test_message_graph(
|
||||
FunctionMessage(
|
||||
content="result for a different query",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000113",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2555,7 +2599,7 @@ def test_message_graph(
|
||||
id="ai2",
|
||||
),
|
||||
],
|
||||
next=("agent:edges",),
|
||||
next=("action",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
@@ -2569,7 +2613,7 @@ def test_message_graph(
|
||||
values=[
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000096",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2584,16 +2628,14 @@ def test_message_graph(
|
||||
FunctionMessage(
|
||||
content="result for a different query",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000113",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(content="answer", id="ai2"),
|
||||
],
|
||||
next=("agent:edges",),
|
||||
next=(),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
assert [c for c in app_w_interrupt.stream(None, config)] == []
|
||||
|
||||
|
||||
def test_in_one_fan_out_out_one_graph_state() -> None:
|
||||
def sorted_add(x: list[str], y: list[str]) -> list[str]:
|
||||
|
||||
+137
-109
@@ -31,7 +31,7 @@ from langgraph.prebuilt.chat_agent_executor import (
|
||||
)
|
||||
from langgraph.prebuilt.tool_executor import ToolExecutor
|
||||
from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
|
||||
from langgraph.pregel.reserved import ReservedChannels
|
||||
from tests.any_str import AnyStr
|
||||
from tests.memory_assert import MemorySaverAssertImmutable
|
||||
|
||||
|
||||
@@ -102,28 +102,6 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N
|
||||
assert await app.ainvoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4}
|
||||
|
||||
|
||||
async def test_invoke_single_process_in_out_reserved_is_last(
|
||||
mocker: MockerFixture,
|
||||
) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: {**x, "input": x["input"] + 1})
|
||||
|
||||
chain = (
|
||||
Channel.subscribe_to(["input"]).join([ReservedChannels.is_last_step])
|
||||
| add_one
|
||||
| Channel.write_to("output")
|
||||
)
|
||||
|
||||
app = Pregel(nodes={"one": chain})
|
||||
|
||||
assert app.input_schema.schema() == {"title": "LangGraphInput"}
|
||||
assert app.output_schema.schema() == {"title": "LangGraphOutput"}
|
||||
assert await app.ainvoke(2) == {"input": 3, "is_last_step": False}
|
||||
assert await app.ainvoke(2, {"recursion_limit": 1}) == {
|
||||
"input": 3,
|
||||
"is_last_step": True,
|
||||
}
|
||||
|
||||
|
||||
async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
|
||||
@@ -227,7 +205,7 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
checkpointer=memory,
|
||||
interrupt_after_nodes=["inbox"],
|
||||
interrupt_after_nodes=["one"],
|
||||
)
|
||||
|
||||
# start execution, stop at inbox
|
||||
@@ -262,7 +240,7 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N
|
||||
assert snapshot.next == ("two",)
|
||||
|
||||
# update the state, resume
|
||||
await app.aupdate_state({"configurable": {"thread_id": 2}}, {"inbox": 25})
|
||||
await app.aupdate_state({"configurable": {"thread_id": 2}}, 25, as_node="one")
|
||||
assert await app.ainvoke(None, {"configurable": {"thread_id": 2}}) == 26
|
||||
|
||||
# no pending tasks
|
||||
@@ -536,9 +514,11 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None:
|
||||
state = await app.aget_state({"configurable": {"thread_id": "1"}})
|
||||
assert state is not None
|
||||
assert state.values.get("total") == 7
|
||||
assert state.next == ()
|
||||
state = await app.aget_state(thread_2)
|
||||
assert state is not None
|
||||
assert state.values.get("total") == 5
|
||||
assert state.next == ()
|
||||
|
||||
# list all checkpoints for thread 1
|
||||
thread_1_history = [c async for c in app.aget_state_history(thread_1)]
|
||||
@@ -561,9 +541,7 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None:
|
||||
"ts"
|
||||
] == thread_1_history[1].config["configurable"]["thread_ts"]
|
||||
|
||||
thread_1_next_config = await app.aupdate_state(
|
||||
thread_1_history[1].config, {"total": 10}
|
||||
)
|
||||
thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10)
|
||||
# update creates a new checkpoint
|
||||
assert (
|
||||
thread_1_next_config["configurable"]["thread_ts"]
|
||||
@@ -981,23 +959,20 @@ async def test_conditional_graph() -> None:
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
},
|
||||
"tools": None,
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
)
|
||||
|
||||
await app_w_interrupt.aupdate_state(
|
||||
config,
|
||||
{
|
||||
"agent": {
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
)
|
||||
|
||||
@@ -1011,9 +986,8 @@ async def test_conditional_graph() -> None:
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
"tools": None,
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
)
|
||||
|
||||
@@ -1058,6 +1032,26 @@ async def test_conditional_graph() -> None:
|
||||
await app_w_interrupt.aupdate_state(
|
||||
config,
|
||||
{
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
(
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
)
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
@@ -1074,12 +1068,12 @@ async def test_conditional_graph() -> None:
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
}
|
||||
},
|
||||
},
|
||||
next=(),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
)
|
||||
|
||||
assert [c async for c in app_w_interrupt.astream(None, config)] == []
|
||||
|
||||
# test state get/update methods with interrupt_before
|
||||
|
||||
app_w_interrupt = workflow.compile(
|
||||
@@ -1112,23 +1106,20 @@ async def test_conditional_graph() -> None:
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
},
|
||||
"tools": None,
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
)
|
||||
|
||||
await app_w_interrupt.aupdate_state(
|
||||
config,
|
||||
{
|
||||
"agent": {
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
)
|
||||
|
||||
@@ -1142,9 +1133,8 @@ async def test_conditional_graph() -> None:
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
"tools": None,
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
)
|
||||
|
||||
@@ -1189,6 +1179,26 @@ async def test_conditional_graph() -> None:
|
||||
await app_w_interrupt.aupdate_state(
|
||||
config,
|
||||
{
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
(
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
)
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
@@ -1205,12 +1215,12 @@ async def test_conditional_graph() -> None:
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
}
|
||||
},
|
||||
},
|
||||
next=(),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
)
|
||||
|
||||
assert [c async for c in app_w_interrupt.astream(None, config)] == []
|
||||
|
||||
# test re-invoke to continue with interrupt_before
|
||||
|
||||
app_w_interrupt = workflow.compile(
|
||||
@@ -1243,9 +1253,8 @@ async def test_conditional_graph() -> None:
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
},
|
||||
"tools": None,
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
)
|
||||
|
||||
@@ -1527,11 +1536,13 @@ async def test_conditional_graph_state() -> None:
|
||||
values={
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"intermediate_steps": [],
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
)
|
||||
|
||||
@@ -1556,7 +1567,7 @@ async def test_conditional_graph_state() -> None:
|
||||
),
|
||||
"intermediate_steps": [],
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
)
|
||||
|
||||
@@ -1596,7 +1607,27 @@ async def test_conditional_graph_state() -> None:
|
||||
},
|
||||
)
|
||||
|
||||
assert [c async for c in app_w_interrupt.astream(None, config)] == []
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values={
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
"intermediate_steps": [
|
||||
(
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
)
|
||||
],
|
||||
},
|
||||
next=(),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
)
|
||||
|
||||
# test state get/update methods with interrupt_before
|
||||
|
||||
@@ -1630,7 +1661,7 @@ async def test_conditional_graph_state() -> None:
|
||||
),
|
||||
"intermediate_steps": [],
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
)
|
||||
|
||||
@@ -1655,7 +1686,7 @@ async def test_conditional_graph_state() -> None:
|
||||
),
|
||||
"intermediate_steps": [],
|
||||
},
|
||||
next=("agent:edges",),
|
||||
next=("tools",),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
)
|
||||
|
||||
@@ -1695,7 +1726,27 @@ async def test_conditional_graph_state() -> None:
|
||||
},
|
||||
)
|
||||
|
||||
assert [c async for c in app_w_interrupt.astream(None, config)] == []
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values={
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
"intermediate_steps": [
|
||||
(
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
)
|
||||
],
|
||||
},
|
||||
next=(),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
)
|
||||
|
||||
|
||||
async def test_conditional_entrypoint_graph() -> None:
|
||||
@@ -2092,7 +2143,7 @@ async def test_prebuilt_chat() -> None:
|
||||
]
|
||||
|
||||
|
||||
async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
async def test_message_graph() -> None:
|
||||
from langchain.chat_models.fake import FakeMessagesListChatModel
|
||||
from langchain_community.tools import tool
|
||||
from langchain_core.agents import AgentAction
|
||||
@@ -2206,9 +2257,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
app = workflow.compile()
|
||||
|
||||
assert await app.ainvoke(HumanMessage(content="what is weather in sf")) == [
|
||||
HumanMessage(
|
||||
content="what is weather in sf", id="00000000-0000-4000-8000-000000000002"
|
||||
),
|
||||
HumanMessage(content="what is weather in sf", id=AnyStr()),
|
||||
AIMessage(
|
||||
content="",
|
||||
additional_kwargs={
|
||||
@@ -2216,11 +2265,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
},
|
||||
id="ai1",
|
||||
),
|
||||
FunctionMessage(
|
||||
content="result for query",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000015",
|
||||
),
|
||||
FunctionMessage(content="result for query", name="search_api", id=AnyStr()),
|
||||
AIMessage(
|
||||
content="",
|
||||
additional_kwargs={
|
||||
@@ -2228,25 +2273,14 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
},
|
||||
id="ai2",
|
||||
),
|
||||
FunctionMessage(
|
||||
content="result for another",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000028",
|
||||
),
|
||||
FunctionMessage(content="result for another", name="search_api", id=AnyStr()),
|
||||
AIMessage(content="answer", id="ai3"),
|
||||
]
|
||||
|
||||
assert [
|
||||
c async for c in app.astream([HumanMessage(content="what is weather in sf")])
|
||||
] == [
|
||||
{
|
||||
"__start__": [
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000037",
|
||||
)
|
||||
]
|
||||
},
|
||||
{"__start__": [HumanMessage(content="what is weather in sf", id=AnyStr())]},
|
||||
{
|
||||
"agent": AIMessage(
|
||||
content="",
|
||||
@@ -2258,9 +2292,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
},
|
||||
{
|
||||
"action": FunctionMessage(
|
||||
content="result for query",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000050",
|
||||
content="result for query", name="search_api", id=AnyStr()
|
||||
)
|
||||
},
|
||||
{
|
||||
@@ -2274,9 +2306,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
},
|
||||
{
|
||||
"action": FunctionMessage(
|
||||
content="result for another",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000063",
|
||||
content="result for another", name="search_api", id=AnyStr()
|
||||
)
|
||||
},
|
||||
{"agent": AIMessage(content="answer", id="ai3")},
|
||||
@@ -2296,7 +2326,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
{
|
||||
"__start__": HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000072",
|
||||
id=AnyStr(),
|
||||
)
|
||||
},
|
||||
{
|
||||
@@ -2314,7 +2344,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
values=[
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000072",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2324,7 +2354,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
id="ai1",
|
||||
),
|
||||
],
|
||||
next=("agent:edges",),
|
||||
next=("action",),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
)
|
||||
|
||||
@@ -2338,7 +2368,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
values=[
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000072",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2351,7 +2381,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
id="ai1",
|
||||
),
|
||||
],
|
||||
next=("agent:edges",),
|
||||
next=("action",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
@@ -2360,7 +2390,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
"action": FunctionMessage(
|
||||
content="result for a different query",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000086",
|
||||
id=AnyStr(),
|
||||
)
|
||||
},
|
||||
{
|
||||
@@ -2378,7 +2408,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
values=[
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000072",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2393,7 +2423,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
FunctionMessage(
|
||||
content="result for a different query",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000086",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2403,7 +2433,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
id="ai2",
|
||||
),
|
||||
],
|
||||
next=("agent:edges",),
|
||||
next=("action",),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
@@ -2417,7 +2447,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
values=[
|
||||
HumanMessage(
|
||||
content="what is weather in sf",
|
||||
id="00000000-0000-4000-8000-000000000072",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
@@ -2432,16 +2462,14 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
|
||||
FunctionMessage(
|
||||
content="result for a different query",
|
||||
name="search_api",
|
||||
id="00000000-0000-4000-8000-000000000086",
|
||||
id=AnyStr(),
|
||||
),
|
||||
AIMessage(content="answer", id="ai2"),
|
||||
],
|
||||
next=("agent:edges",),
|
||||
next=(),
|
||||
config=app_w_interrupt.checkpointer.get_tuple(config).config,
|
||||
)
|
||||
|
||||
assert [c async for c in app_w_interrupt.astream(None, config)] == []
|
||||
|
||||
|
||||
async def test_in_one_fan_out_out_one_graph_state() -> None:
|
||||
def sorted_add(x: list[str], y: list[str]) -> list[str]:
|
||||
|
||||
Reference in New Issue
Block a user