WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-31 12:19:58 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'master' into harrison/add-args

Browse Source
This commit is contained in:
Harrison Chase
2024-05-01 08:26:00 -07:00
parent 97f3d66c0e b2020c1972
commit dea85d236c
20 changed files with 1784 additions and 1522 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/human-in-the-loop.ipynb
+2 -2
View File
@@ -11,8 +11,8 @@
"This can be helpful when giving them access to tools.\n",
"Often in these situations you may want to manually approve an action before taking.\n",
"\n",
"This can be in several ways, but the primary supported way is to add an \"interupt\" before a node is executed.\n",
"This interupts execution at that node.\n",
"This can be in several ways, but the primary supported way is to add an \"interrupt\" before a node is executed.\n",
"This interrupts execution at that node.\n",
"You can then resume from that spot to continue."
]
},
examples/multi_agent/hierarchical_agent_teams.ipynb
+1 -1
View File
@@ -626,7 +626,7 @@
")\n",
"\n",
"authoring_graph.set_entry_point(\"supervisor\")\n",
"chain = research_graph.compile()\n",
"chain = authoring_graph.compile()\n",
"\n",
"\n",
"# The following functions interoperate between the top level graph state\n",
examples/rag/langgraph_rag_agent_llama3_local.ipynb
+2 -2
View File
@@ -7,7 +7,7 @@
"metadata": {},
"outputs": [],
"source": [
"! pip install -U langchain-nomic langchain_community tiktoken langchainhub chromadb langchain langgraph tavily-python"
"! pip install -U langchain-nomic langchain_community tiktoken langchainhub chromadb langchain langgraph tavily-python gpt4all"
]
},
{
@@ -229,7 +229,7 @@
"# Prompt\n",
"prompt = PromptTemplate(\n",
" template=\"\"\" <|begin_of_text|><|start_header_id|>system<|end_header_id|> You are a grader assessing whether \n",
" an answer is grounded in / supported by a set of facts. Give a binary score 'yes' or 'no' score to indicate \n",
" an answer is grounded in / supported by a set of facts. Give a binary 'yes' or 'no' score to indicate \n",
" whether the answer is grounded in / supported by a set of facts. Provide the binary score as a JSON with a \n",
" single key 'score' and no preamble or explanation. <|eot_id|><|start_header_id|>user<|end_header_id|>\n",
" Here are the facts:\n",
examples/storm/storm.ipynb
+5 -5
View File
@@ -313,7 +313,7 @@
" description=\"Primary affiliation of the editor.\",\n",
" )\n",
" name: str = Field(\n",
" description=\"Name of the editor.\",\n",
" description=\"Name of the editor.\", pattern=r\"^[a-zA-Z0-9_-]{1,64}$\"\n",
" )\n",
" role: str = Field(\n",
" description=\"Role of the editor in the context of the topic.\",\n",
@@ -730,7 +730,7 @@
"async def gen_answer(\n",
" state: InterviewState,\n",
" config: Optional[RunnableConfig] = None,\n",
" name: str = \"Subject Matter Expert\",\n",
" name: str = \"Subject_Matter_Expert\",\n",
" max_str_len: int = 15000,\n",
"):\n",
" swapped_state = swap_roles(state, name) # Convert all other AI messages\n",
@@ -803,7 +803,7 @@
"max_num_turns = 5\n",
"\n",
"\n",
"def route_messages(state: InterviewState, name: str = \"Subject Matter Expert\"):\n",
"def route_messages(state: InterviewState, name: str = \"Subject_Matter_Expert\"):\n",
" messages = state[\"messages\"]\n",
" num_responses = len(\n",
" [m for m in messages if isinstance(m, AIMessage) and m.name == name]\n",
@@ -886,7 +886,7 @@
" \"messages\": [\n",
" AIMessage(\n",
" content=f\"So you said you were writing an article on {example_topic}?\",\n",
" name=\"Subject Matter Expert\",\n",
" name=\"Subject_Matter_Expert\",\n",
" )\n",
" ],\n",
"}\n",
@@ -1395,7 +1395,7 @@
" \"messages\": [\n",
" AIMessage(\n",
" content=f\"So you said you were writing an article on {topic}?\",\n",
" name=\"Subject Matter Expert\",\n",
" name=\"Subject_Matter_Expert\",\n",
" )\n",
" ],\n",
" }\n",
examples/time-travel.ipynb
+3 -3
View File
File diff suppressed because one or more lines are too long
langgraph/checkpoint/aiosqlite.py
+5 -5
View File
@@ -75,8 +75,8 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
async with self.conn.execute(
"SELECT checkpoint, parent_ts FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
(
config["configurable"]["thread_id"],
config["configurable"]["thread_ts"],
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
),
) as cursor:
if value := await cursor.fetchone():
@@ -95,7 +95,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
else:
async with self.conn.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
(config["configurable"]["thread_id"],),
(str(config["configurable"]["thread_id"]),),
) as cursor:
if value := await cursor.fetchone():
return CheckpointTuple(
@@ -120,7 +120,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
await self.setup()
async with self.conn.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC",
(config["configurable"]["thread_id"],),
(str(config["configurable"]["thread_id"]),),
) as cursor:
async for thread_id, thread_ts, parent_ts, value in cursor:
yield CheckpointTuple(
@@ -138,7 +138,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
async with self.conn.execute(
"INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint) VALUES (?, ?, ?, ?)",
(
config["configurable"]["thread_id"],
str(config["configurable"]["thread_id"]),
checkpoint["ts"],
config["configurable"].get("thread_ts"),
self.serde.dumps(checkpoint),
langgraph/checkpoint/sqlite.py
+5 -5
View File
@@ -93,8 +93,8 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
cur.execute(
"SELECT checkpoint, parent_ts FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
(
config["configurable"]["thread_id"],
config["configurable"]["thread_ts"],
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
),
)
if value := cur.fetchone():
@@ -113,7 +113,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
else:
cur.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
(config["configurable"]["thread_id"],),
(str(config["configurable"]["thread_id"]),),
)
if value := cur.fetchone():
return CheckpointTuple(
@@ -138,7 +138,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
with self.cursor(transaction=False) as cur:
cur.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC",
(config["configurable"]["thread_id"],),
(str(config["configurable"]["thread_id"]),),
)
for thread_id, thread_ts, parent_ts, value in cur:
yield CheckpointTuple(
@@ -159,7 +159,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
cur.execute(
"INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint) VALUES (?, ?, ?, ?)",
(
config["configurable"]["thread_id"],
str(config["configurable"]["thread_id"]),
checkpoint["ts"],
config["configurable"].get("thread_ts"),
self.serde.dumps(checkpoint),
langgraph/graph/graph.py
+38 -33
View File
@@ -1,15 +1,19 @@
import logging
from collections import Counter, defaultdict
from collections import defaultdict
from typing import (
Any,
Awaitable,
Callable,
Dict,
Literal,
NamedTuple,
Optional,
Sequence,
Union,
cast,
get_args,
get_origin,
get_type_hints,
)
from langchain_core.runnables import Runnable
@@ -24,6 +28,7 @@ from langgraph.checkpoint import BaseCheckpointSaver
from langgraph.constants import TAG_HIDDEN
from langgraph.pregel import Channel, Pregel
from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import All
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
from langgraph.utils import DrawableGraph, RunnableCallable, coerce_to_runnable
@@ -142,7 +147,7 @@ class Graph:
Callable[..., Awaitable[Union[str, list[str]]]],
Runnable[Any, Union[str, list[str]]],
],
path_map: Optional[dict[str, str]] = None,
path_map: Optional[Union[dict[str, str], list[str]]] = None,
then: Optional[str] = None,
) -> None:
"""Add a conditional edge from the starting node to any number of destination nodes.
@@ -166,6 +171,14 @@ class Graph:
"Adding an edge to a graph that has already been compiled. This will "
"not be reflected in the compiled graph."
)
# coerce path_map to a dictionary
if isinstance(path_map, dict):
pass
elif isinstance(path_map, list):
path_map = {name: name for name in path_map}
elif rtn_type := get_type_hints(path).get("return"):
if get_origin(rtn_type) is Literal:
path_map = {name: name for name in get_args(rtn_type)}
# find a name for the condition
path = coerce_to_runnable(path, name=None, trace=True)
name = path.name or "condition"
@@ -283,8 +296,8 @@ class Graph:
def compile(
self,
checkpointer: Optional[BaseCheckpointSaver] = None,
interrupt_before: Optional[Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: bool = False,
) -> "CompiledGraph":
# assign default values
@@ -292,11 +305,16 @@ class Graph:
interrupt_after = interrupt_after or []
# validate the graph
self.validate(interrupt=interrupt_before + interrupt_after)
self.validate(
interrupt=(interrupt_before if interrupt_before != "*" else [])
+ interrupt_after
if interrupt_after != "*"
else []
)
# create empty compiled graph
compiled = CompiledGraph(
graph=self,
builder=self,
nodes={},
channels={START: EphemeralValue(Any), END: EphemeralValue(Any)},
input_channels=START,
@@ -326,7 +344,7 @@ class Graph:
class CompiledGraph(Pregel):
graph: Graph
builder: Graph
def attach_node(self, key: str, node: Runnable) -> None:
self.channels[key] = EphemeralValue(Any)
@@ -378,7 +396,6 @@ class CompiledGraph(Pregel):
config: Optional[RunnableConfig] = None,
*,
xray: Union[int, bool] = False,
add_condition_nodes: bool = True,
) -> DrawableGraph:
"""Returns a drawable representation of the computation graph."""
graph = DrawableGraph()
@@ -389,7 +406,7 @@ class CompiledGraph(Pregel):
END: graph.add_node(self.get_output_schema(config), END)
}
for key, node in self.graph.nodes.items():
for key, node in self.builder.nodes.items():
if xray:
subgraph = (
node.get_graph(
@@ -413,40 +430,28 @@ class CompiledGraph(Pregel):
n = graph.add_node(node, key)
start_nodes[key] = n
end_nodes[key] = n
for start, end in sorted(self.graph._all_edges):
for start, end in sorted(self.builder._all_edges):
graph.add_edge(start_nodes[start], end_nodes[end])
branches_by_name = Counter(
name for _, branches in self.graph.branches.items() for name in branches
)
for start, branches in self.graph.branches.items():
for start, branches in self.builder.branches.items():
default_ends = {
**{k: k for k in self.graph.nodes if k != start},
**{k: k for k in self.builder.nodes if k != start},
END: END,
}
for name, branch in branches.items():
for _, branch in branches.items():
if branch.ends is not None:
ends = branch.ends
elif branch.then is not None:
ends = {k: k for k in default_ends if k not in (END, branch.then)}
else:
ends = default_ends
if add_condition_nodes is True:
cond = graph.add_node(
branch.path,
f"{start}_{name}" if branches_by_name[name] > 1 else name,
for label, end in ends.items():
graph.add_edge(
start_nodes[start],
end_nodes[end],
label if label != end else None,
conditional=True,
)
graph.add_edge(start_nodes[start], cond)
for label, end in ends.items():
graph.add_edge(cond, end_nodes[end], label, conditional=True)
if branch.then is not None:
graph.add_edge(start_nodes[end], end_nodes[branch.then])
else:
for label, end in ends.items():
graph.add_edge(
start_nodes[start], end_nodes[end], label, conditional=True
)
if branch.then is not None:
graph.add_edge(start_nodes[end], end_nodes[branch.then])
if branch.then is not None:
graph.add_edge(start_nodes[end], end_nodes[branch.then])
return graph
langgraph/graph/state.py
+15 -9
View File
@@ -16,6 +16,7 @@ from langgraph.checkpoint import BaseCheckpointSaver
from langgraph.constants import TAG_HIDDEN
from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph
from langgraph.pregel.read import ChannelRead, PregelNode
from langgraph.pregel.types import All
from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
from langgraph.utils import RunnableCallable
@@ -100,8 +101,8 @@ class StateGraph(Graph):
def compile(
self,
checkpointer: Optional[BaseCheckpointSaver] = None,
interrupt_before: Optional[Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: bool = False,
) -> CompiledGraph:
"""Compiles the state graph into a `CompiledGraph` object.
@@ -120,14 +121,19 @@ class StateGraph(Graph):
interrupt_after = interrupt_after or []
# validate the graph
self.validate(interrupt=interrupt_before + interrupt_after)
self.validate(
interrupt=(interrupt_before if interrupt_before != "*" else [])
+ interrupt_after
if interrupt_after != "*"
else []
)
# prepare output channels
state_keys = list(self.channels)
output_channels = state_keys[0] if state_keys == ["__root__"] else state_keys
compiled = CompiledStateGraph(
graph=self,
builder=self,
nodes={},
channels={**self.channels, START: EphemeralValue(self.schema)},
input_channels=START,
@@ -159,7 +165,7 @@ class StateGraph(Graph):
class CompiledStateGraph(CompiledGraph):
graph: StateGraph
builder: StateGraph
def attach_node(self, key: str, node: Optional[Runnable]) -> None:
def _get_state_key(input: dict, config: RunnableConfig, *, key: str) -> Any:
@@ -170,7 +176,7 @@ class CompiledStateGraph(CompiledGraph):
else:
return input.get(key, SKIP_WRITE)
state_keys = list(self.graph.channels)
state_keys = list(self.builder.channels)
# state updaters
state_write_entries = [
(
@@ -210,7 +216,7 @@ class CompiledStateGraph(CompiledGraph):
mapper=(
None
if state_keys == ["__root__"]
else partial(_coerce_state, self.graph.schema)
else partial(_coerce_state, self.builder.schema)
),
writers=[
# publish to this channel and state keys
@@ -265,13 +271,13 @@ class CompiledStateGraph(CompiledGraph):
return ChannelWrite(writes, tags=[TAG_HIDDEN])
# attach branch publisher
self.nodes[start] |= branch.run(branch_writer, _get_state_reader(self.graph))
self.nodes[start] |= branch.run(branch_writer, _get_state_reader(self.builder))
# attach branch subscribers
ends = (
branch.ends.values()
if branch.ends
else [node for node in self.graph.nodes if node != branch.then]
else [node for node in self.builder.nodes if node != branch.then]
)
for end in ends:
if end != END:
langgraph/pregel/__init__.py
+131 -64
View File
@@ -54,7 +54,6 @@ from langgraph.channels.base import (
InvalidUpdateError,
create_checkpoint,
)
from langgraph.channels.last_value import LastValue
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
@@ -66,8 +65,12 @@ from langgraph.constants import (
CONFIG_KEY_READ,
CONFIG_KEY_SEND,
INTERRUPT,
TAG_HIDDEN,
)
from langgraph.pregel.debug import (
map_debug_checkpoint,
map_debug_task_results,
map_debug_tasks,
print_step_checkpoint,
print_step_tasks,
print_step_writes,
@@ -82,6 +85,7 @@ from langgraph.pregel.io import (
from langgraph.pregel.log import logger
from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import (
All,
PregelExecutableTask,
PregelTaskDescription,
StateSnapshot,
@@ -175,7 +179,7 @@ class Channel:
)
StreamMode = Literal["values", "updates"]
StreamMode = Literal["values", "updates", "debug"]
class Pregel(
@@ -185,23 +189,21 @@ 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"
output_channels: Union[str, Sequence[str]]
"""Channels to output, defaults to channel named 'output'."""
stream_channels: Optional[Union[str, Sequence[str]]] = None
"""Channels to stream, defaults to all channels not in reserved channels"""
interrupt_after_nodes: Sequence[str] = Field(default_factory=list)
interrupt_after_nodes: Union[All, Sequence[str]] = Field(default_factory=list)
interrupt_before_nodes: Sequence[str] = Field(default_factory=list)
interrupt_before_nodes: Union[All, Sequence[str]] = Field(default_factory=list)
input_channels: Union[str, Sequence[str]] = "input"
input_channels: Union[str, Sequence[str]]
step_timeout: Optional[float] = None
@@ -214,6 +216,11 @@ class Pregel(
class Config:
arbitrary_types_allowed = True
@classmethod
def is_lc_serializable(cls) -> bool:
"""Return whether the graph can be serialized by Langchain."""
return True
@root_validator(skip_on_failure=True)
def validate_on_init(cls, values: dict[str, Any]) -> dict[str, Any]:
if not values["auto_validate"]:
@@ -226,7 +233,6 @@ 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"]:
@@ -242,7 +248,6 @@ class Pregel(
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:
@@ -302,12 +307,15 @@ class Pregel(
@property
def stream_channels_list(self) -> Sequence[str]:
stream_channels = self.stream_channels_asis
return (
[self.stream_channels]
if isinstance(self.stream_channels, str)
else self.stream_channels or [k for k in self.channels]
[stream_channels] if isinstance(stream_channels, str) else stream_channels
)
@property
def stream_channels_asis(self) -> Union[str, Sequence[str]]:
return self.stream_channels or [k for k in self.channels]
def get_state(self, config: RunnableConfig) -> StateSnapshot:
"""Get the current state of the graph."""
if not self.checkpointer:
@@ -320,11 +328,8 @@ class Pregel(
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, for_execution=False
)
values = read_channels(channels, self.stream_channels_list)
return StateSnapshot(
values.get(self.stream_channels, None)
if isinstance(self.stream_channels, str)
else values,
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
config,
)
@@ -341,11 +346,8 @@ class Pregel(
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, for_execution=False
)
values = read_channels(channels, self.stream_channels_list)
return StateSnapshot(
values.get(self.stream_channels, None)
if isinstance(self.stream_channels, str)
else values,
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
config,
)
@@ -360,11 +362,8 @@ class Pregel(
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, for_execution=False
)
values = read_channels(channels, self.stream_channels_list)
yield StateSnapshot(
values.get(self.stream_channels, None)
if isinstance(self.stream_channels, str)
else values,
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
config,
parent_config,
@@ -382,11 +381,8 @@ class Pregel(
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, for_execution=False
)
values = read_channels(channels, self.stream_channels_list)
yield StateSnapshot(
values.get(self.stream_channels, None)
if isinstance(self.stream_channels, str)
else values,
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
config,
parent_config,
@@ -434,6 +430,8 @@ class Pregel(
values,
RunnableSequence(*writers) if len(writers) > 1 else writers[0],
deque(),
None,
[INTERRUPT],
)
# execute task
task.proc.invoke(
@@ -494,6 +492,8 @@ class Pregel(
values,
RunnableSequence(*writers) if len(writers) > 1 else writers[0],
deque(),
None,
[INTERRUPT],
)
# execute task
await task.proc.ainvoke(
@@ -523,8 +523,8 @@ class Pregel(
stream_mode: Optional[StreamMode] = None,
input_keys: Optional[Union[str, Sequence[str]]] = None,
output_keys: Optional[Union[str, Sequence[str]]] = None,
interrupt_before: Optional[Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: Optional[bool] = None,
) -> tuple[
bool,
@@ -536,11 +536,7 @@ class Pregel(
]:
debug = debug if debug is not None else self.debug
if output_keys is None:
output_keys = (
[chan for chan in self.channels]
if self.stream_channels is None
else self.stream_channels
)
output_keys = self.stream_channels_asis
else:
validate_keys(output_keys, self.channels)
if input_keys is None:
@@ -570,8 +566,8 @@ class Pregel(
stream_mode: Optional[StreamMode] = None,
output_keys: Optional[Union[str, Sequence[str]]] = None,
input_keys: Optional[Union[str, Sequence[str]]] = None,
interrupt_before: Optional[Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: Optional[bool] = None,
) -> Iterator[Union[dict[str, Any], Any]]:
"""Stream graph steps for a single input."""
@@ -671,6 +667,9 @@ class Pregel(
if debug:
print_step_tasks(step, next_tasks)
if stream_mode == "debug":
for chunk in map_debug_tasks(step, next_tasks):
yield chunk
# prepare tasks with config
tasks_w_config = [
@@ -690,7 +689,7 @@ class Pregel(
},
),
)
for name, input, proc, writes, proc_config in next_tasks
for name, input, proc, writes, proc_config, _ in next_tasks
]
futures = [
@@ -711,7 +710,7 @@ class Pregel(
# combine pending writes from all tasks
pending_writes = deque[tuple[str, Any]]()
for _, _, _, writes, _ in next_tasks:
for _, _, _, writes, _, _ in next_tasks:
pending_writes.extend(writes)
if debug:
@@ -730,6 +729,10 @@ class Pregel(
yield from map_output_values(
output_keys, pending_writes, channels
)
elif stream_mode == "debug":
yield from map_debug_task_results(
step, next_tasks, self.stream_channels_list
)
else:
yield from map_output_updates(output_keys, next_tasks)
@@ -741,6 +744,17 @@ class Pregel(
checkpoint_config = self.checkpointer.put(
checkpoint_config, checkpoint
)
if stream_mode == "debug":
yield map_debug_checkpoint(
step,
checkpoint_config,
channels,
self.stream_channels_asis,
)
elif stream_mode == "debug":
yield map_debug_checkpoint(
step, None, channels, self.stream_channels_asis
)
# after execution, check if we should interrupt
if _should_interrupt(
@@ -760,7 +774,20 @@ class Pregel(
and self.checkpointer.at == CheckpointAt.END_OF_RUN
):
checkpoint = create_checkpoint(checkpoint, channels)
self.checkpointer.put(checkpoint_config, checkpoint)
checkpoint_config = self.checkpointer.put(
checkpoint_config, checkpoint
)
if stream_mode == "debug":
yield map_debug_checkpoint(
step,
checkpoint_config,
channels,
self.stream_channels_asis,
)
elif self.checkpointer is None and stream_mode == "debug":
yield map_debug_checkpoint(
step, None, channels, self.stream_channels_asis
)
except BaseException as e:
run_manager.on_chain_error(e)
raise
@@ -780,8 +807,8 @@ class Pregel(
stream_mode: Optional[StreamMode] = None,
output_keys: Optional[Union[str, Sequence[str]]] = None,
input_keys: Optional[Union[str, Sequence[str]]] = None,
interrupt_before: Optional[Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: Optional[bool] = None,
) -> AsyncIterator[Union[dict[str, Any], Any]]:
config = ensure_config(config)
@@ -889,6 +916,9 @@ class Pregel(
if debug:
print_step_tasks(step, next_tasks)
if stream_mode == "debug":
for chunk in map_debug_tasks(step, next_tasks):
yield chunk
# prepare tasks with config
tasks_w_config = [
@@ -908,7 +938,7 @@ class Pregel(
},
),
)
for name, input, proc, writes, proc_config in next_tasks
for name, input, proc, writes, proc_config, _ in next_tasks
]
futures = (
@@ -936,7 +966,7 @@ class Pregel(
# combine pending writes from all tasks
pending_writes = deque[tuple[str, Any]]()
for _, _, _, writes, _ in next_tasks:
for _, _, _, writes, _, _ in next_tasks:
pending_writes.extend(writes)
if debug:
@@ -956,6 +986,11 @@ class Pregel(
output_keys, pending_writes, channels
):
yield chunk
elif stream_mode == "debug":
for chunk in map_debug_task_results(
step, next_tasks, self.stream_channels_list
):
yield chunk
else:
for chunk in map_output_updates(output_keys, next_tasks):
yield chunk
@@ -968,6 +1003,17 @@ class Pregel(
checkpoint_config = await self.checkpointer.aput(
checkpoint_config, checkpoint
)
if stream_mode == "debug":
yield map_debug_checkpoint(
step,
checkpoint_config,
channels,
self.stream_channels_asis,
)
elif stream_mode == "debug":
yield map_debug_checkpoint(
step, None, channels, self.stream_channels_asis
)
# after execution, check if we should interrupt
if _should_interrupt(
@@ -987,7 +1033,17 @@ class Pregel(
and self.checkpointer.at == CheckpointAt.END_OF_RUN
):
checkpoint = create_checkpoint(checkpoint, channels)
await self.checkpointer.aput(checkpoint_config, checkpoint)
checkpoint_config = await self.checkpointer.aput(
checkpoint_config, checkpoint
)
if stream_mode == "debug":
yield map_debug_checkpoint(
step, checkpoint_config, channels, self.stream_channels_asis
)
elif self.checkpointer is None and stream_mode == "debug":
yield map_debug_checkpoint(
step, None, channels, self.stream_channels_asis
)
except BaseException as e:
await run_manager.on_chain_error(e)
raise
@@ -1007,8 +1063,8 @@ class Pregel(
stream_mode: StreamMode = "values",
output_keys: Optional[Union[str, Sequence[str]]] = None,
input_keys: Optional[Union[str, Sequence[str]]] = None,
interrupt_before_nodes: Optional[Sequence[str]] = None,
interrupt_after_nodes: Optional[Sequence[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: Optional[bool] = None,
**kwargs: Any,
) -> Union[dict[str, Any], Any]:
@@ -1020,8 +1076,8 @@ class Pregel(
stream_mode: Optional[str]. The stream mode for the graph run. Default is "values".
output_keys: Optional. The output keys to retrieve from the graph run.
input_keys: Optional. The input keys to provide for the graph run.
interrupt_before_nodes: Optional. The nodes to interrupt the graph run before.
interrupt_after_nodes: Optional. The nodes to interrupt the graph run after.
interrupt_before: Optional. The nodes to interrupt the graph run before.
interrupt_after: Optional. The nodes to interrupt the graph run after.
debug: Optional. Enable debug mode for the graph run.
**kwargs: Additional keyword arguments to pass to the graph run.
@@ -1040,8 +1096,8 @@ class Pregel(
stream_mode=stream_mode,
output_keys=output_keys,
input_keys=input_keys,
interrupt_before=interrupt_before_nodes,
interrupt_after=interrupt_after_nodes,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
debug=debug,
**kwargs,
):
@@ -1062,8 +1118,8 @@ class Pregel(
stream_mode: StreamMode = "values",
output_keys: Optional[Union[str, Sequence[str]]] = None,
input_keys: Optional[Union[str, Sequence[str]]] = None,
interrupt_before_nodes: Optional[Sequence[str]] = None,
interrupt_after_nodes: Optional[Sequence[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: Optional[bool] = None,
**kwargs: Any,
) -> Union[dict[str, Any], Any]:
@@ -1075,8 +1131,8 @@ class Pregel(
stream_mode: Optional. The stream mode for the computation. Default is "values".
output_keys: Optional. The output keys to include in the result. Default is None.
input_keys: Optional. The input keys to include in the result. Default is None.
interrupt_before_nodes: Optional. The nodes to interrupt before. Default is None.
interrupt_after_nodes: Optional. The nodes to interrupt after. Default is None.
interrupt_before: Optional. The nodes to interrupt before. Default is None.
interrupt_after: Optional. The nodes to interrupt after. Default is None.
debug: Optional. Whether to enable debug mode. Default is None.
**kwargs: Additional keyword arguments.
@@ -1096,8 +1152,8 @@ class Pregel(
stream_mode=stream_mode,
output_keys=output_keys,
input_keys=input_keys,
interrupt_before=interrupt_before_nodes,
interrupt_after=interrupt_after_nodes,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
debug=debug,
**kwargs,
):
@@ -1137,7 +1193,7 @@ def _panic_or_proceed(
def _should_interrupt(
checkpoint: Checkpoint,
interrupt_nodes: Sequence[str],
interrupt_nodes: Union[All, Sequence[str]],
snapshot_channels: Sequence[str],
tasks: list[PregelExecutableTask],
) -> bool:
@@ -1150,7 +1206,15 @@ def _should_interrupt(
for chan in snapshot_channels
)
# and any channel written to is in interrupt_nodes list
and any(node for node, _, _, _, _ in tasks if node in interrupt_nodes)
and any(
node
for node, _, _, _, config, _ in tasks
if (
(not config or TAG_HIDDEN not in config.get("tags"))
if interrupt_nodes == "*"
else node in interrupt_nodes
)
)
)
@@ -1240,13 +1304,14 @@ def _prepare_next_tasks(
for name, proc in processes.items():
seen = checkpoint["versions_seen"][name]
# If any of the channels read by this process were updated
if any(
checkpoint["channel_versions"][chan] > seen[chan]
if triggers := [
chan
for chan in proc.triggers
if not isinstance(
read_channel(channels, chan, return_exception=True), EmptyChannelError
)
):
and checkpoint["channel_versions"][chan] > seen[chan]
]:
# If all trigger channels subscribed by this process are not empty
# then invoke the process with the values of all non-empty channels
if isinstance(proc.channels, dict):
@@ -1287,7 +1352,9 @@ def _prepare_next_tasks(
if for_execution:
if node := proc.get_node():
tasks.append(
PregelExecutableTask(name, val, node, deque(), proc.config)
PregelExecutableTask(
name, val, node, deque(), proc.config, triggers
)
)
else:
tasks.append(PregelTaskDescription(name, val))
langgraph/pregel/debug.py
+115 -15
View File
@@ -1,13 +1,124 @@
import json
from collections import defaultdict
from datetime import datetime, timezone
from pprint import pformat
from typing import Any, Iterator, Mapping, Sequence
from typing import Any, Iterator, Literal, Mapping, Optional, Sequence, TypedDict, Union
from uuid import UUID, uuid5
from langchain_core.runnables.config import RunnableConfig
from langchain_core.utils.input import get_bolded_text, get_colored_text
from langgraph.channels.base import BaseChannel, EmptyChannelError
from langgraph.channels.base import BaseChannel
from langgraph.constants import TAG_HIDDEN
from langgraph.pregel.io import read_channels
from langgraph.pregel.types import PregelExecutableTask
class TaskPayload(TypedDict):
id: str
name: str
input: Any
triggers: list[str]
class TaskResultPayload(TypedDict):
id: str
result: list[tuple[str, Any]]
class CheckpointPayload(TypedDict):
config: Optional[RunnableConfig]
values: dict[str, Any]
class DebugOutputBase(TypedDict):
timestamp: str
step: int
type: str
payload: dict[str, Any]
class DebugOutputTask(DebugOutputBase):
type: Literal["task"]
payload: TaskPayload
class DebugOutputTaskResult(DebugOutputBase):
type: Literal["task_result"]
payload: TaskResultPayload
class DebugOutputCheckpoint(DebugOutputBase):
type: Literal["checkpoint"]
payload: CheckpointPayload
DebugOutput = Union[DebugOutputTask, DebugOutputTaskResult, DebugOutputCheckpoint]
TASK_NAMESPACE = UUID("6ba7b831-9dad-11d1-80b4-00c04fd430c8")
def map_debug_tasks(
step: int, tasks: list[PregelExecutableTask]
) -> Iterator[DebugOutputTask]:
ts = datetime.now(timezone.utc).isoformat()
for name, input, _, _, config, triggers in tasks:
if config is not None and TAG_HIDDEN in config.get("tags", []):
continue
yield {
"type": "task",
"timestamp": ts,
"step": step,
"payload": {
"id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step)))),
"name": name,
"input": input,
"triggers": triggers,
},
}
def map_debug_task_results(
step: int,
tasks: list[PregelExecutableTask],
stream_channels_list: Sequence[str],
) -> Iterator[DebugOutputTaskResult]:
ts = datetime.now(timezone.utc).isoformat()
for name, _, _, writes, config, _ in tasks:
if config is not None and TAG_HIDDEN in config.get("tags", []):
continue
yield {
"type": "task_result",
"timestamp": ts,
"step": step,
"payload": {
"id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step)))),
"result": [w for w in writes if w[0] in stream_channels_list],
},
}
def map_debug_checkpoint(
step: int,
config: RunnableConfig,
channels: Mapping[str, BaseChannel],
stream_channels: Union[str, Sequence[str]],
) -> DebugOutputCheckpoint:
ts = datetime.now(timezone.utc).isoformat()
return {
"type": "checkpoint",
"timestamp": ts,
"step": step,
"payload": {
"config": config,
"values": read_channels(channels, stream_channels),
},
}
def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None:
n_tasks = len(next_tasks)
print(
@@ -17,7 +128,7 @@ def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None:
)
+ "\n".join(
f"- {get_colored_text(name, 'green')} -> {pformat(val)}"
for name, val, _, _, _ in next_tasks
for name, val, _, _, _, _ in next_tasks
)
)
@@ -47,16 +158,5 @@ def print_step_checkpoint(
print(
f"{get_colored_text(f'[{step}:checkpoint]', color='blue')} "
+ get_bolded_text(f"State at the end of step {step}:\n")
+ pformat(
{name: val for name, val in _read_channels(channels) if name in whitelist},
depth=3,
)
+ pformat(read_channels(channels, whitelist), depth=3)
)
def _read_channels(channels: Mapping[str, BaseChannel]) -> Iterator[tuple[str, Any]]:
for name, channel in channels.items():
try:
yield (name, channel.get())
except EmptyChannelError:
pass
langgraph/pregel/io.py
+2 -2
View File
@@ -105,7 +105,7 @@ def map_output_updates(
if updated := AddableUpdatesDict(
{
node: value
for node, _, _, writes, _ in output_tasks
for node, _, _, writes, _, _ in output_tasks
for chan, value in writes
if chan == output_channels
}
@@ -115,7 +115,7 @@ def map_output_updates(
if updated := AddableUpdatesDict(
{
node: {chan: value for chan, value in writes if chan in output_channels}
for node, _, _, writes, _ in output_tasks
for node, _, _, writes, _, _ in output_tasks
if any(chan in output_channels for chan, _ in writes)
}
):
langgraph/pregel/types.py
+6 -2
View File
@@ -1,5 +1,5 @@
from collections import deque
from typing import Any, NamedTuple, Optional, Union
from typing import Any, Literal, NamedTuple, Optional, Union
from langchain_core.runnables import Runnable, RunnableConfig
@@ -14,7 +14,8 @@ class PregelExecutableTask(NamedTuple):
input: Any
proc: Runnable
writes: deque[tuple[str, Any]]
config: Optional[RunnableConfig] = None
config: Optional[RunnableConfig]
triggers: list[str]
class StateSnapshot(NamedTuple):
@@ -26,3 +27,6 @@ class StateSnapshot(NamedTuple):
"""Config used to fetch this snapshot"""
parent_config: Optional[RunnableConfig] = None
"""Config used to fetch the parent snapshot, if any"""
All = Literal["*"]
langgraph/pregel/validate.py
+20 -14
View File
@@ -1,8 +1,9 @@
from typing import Any, Mapping, Optional, Sequence, Type, Union
from typing import Mapping, Optional, Sequence, Union
from langgraph.channels.base import BaseChannel
from langgraph.constants import INTERRUPT
from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import All
def validate_graph(
@@ -11,10 +12,13 @@ def validate_graph(
input_channels: Union[str, Sequence[str]],
output_channels: Union[str, Sequence[str]],
stream_channels: Optional[Union[str, Sequence[str]]],
interrupt_after_nodes: Sequence[str],
interrupt_before_nodes: Sequence[str],
default_channel_cls: Type[BaseChannel],
interrupt_after_nodes: Union[All, Sequence[str]],
interrupt_before_nodes: Union[All, Sequence[str]],
) -> None:
for chan in channels:
if chan == INTERRUPT:
raise ValueError(f"Channel name {INTERRUPT} is reserved")
subscribed_channels = set[str]()
for name, node in nodes.items():
if name == INTERRUPT:
@@ -28,11 +32,11 @@ def validate_graph(
for chan in subscribed_channels:
if chan not in channels:
channels[chan] = default_channel_cls(Any) # type: ignore[arg-type]
raise ValueError(f"Subscribed channel '{chan}' not in 'channels'")
if isinstance(input_channels, str):
if input_channels not in channels:
channels[input_channels] = default_channel_cls(Any) # type: ignore[arg-type]
raise ValueError(f"Input channel '{input_channels}' not in 'channels'")
if input_channels not in subscribed_channels:
raise ValueError(
f"Input channel {input_channels} is not subscribed to by any node"
@@ -40,7 +44,7 @@ def validate_graph(
else:
for chan in input_channels:
if chan not in channels:
channels[chan] = default_channel_cls(Any) # type: ignore[arg-type]
raise ValueError(f"Input channel '{chan}' not in 'channels'")
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"
@@ -58,14 +62,16 @@ def validate_graph(
for chan in all_output_channels:
if chan not in channels:
channels[chan] = default_channel_cls(Any) # type: ignore[arg-type]
raise ValueError(f"Output channel '{chan}' not in '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")
if interrupt_after_nodes != "*":
for node in interrupt_after_nodes:
if node not in nodes:
raise ValueError(f"Node {node} not in nodes")
if interrupt_before_nodes != "*":
for node in interrupt_before_nodes:
if node not in nodes:
raise ValueError(f"Node {node} not in nodes")
def validate_keys(
langgraph/serde/jsonplus.py
+7 -1
View File
@@ -38,7 +38,13 @@ class JsonPlusSerializer(SerializerProtocol):
if isinstance(obj, Serializable):
return obj.to_json()
elif isinstance(obj, (BaseModel, LcBaseModel)):
return self._encode_constructor_args(obj.__class__, kwargs=obj.dict())
# prefer non-deprecated method if available
if hasattr(obj, "model_dump"):
return self._encode_constructor_args(
obj.__class__, kwargs=obj.model_dump()
)
else:
return self._encode_constructor_args(obj.__class__, kwargs=obj.dict())
elif isinstance(obj, UUID):
return self._encode_constructor_args(UUID, args=[obj.hex])
elif isinstance(obj, (set, frozenset)):
pyproject.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph"
version = "0.0.39"
version = "0.0.40"
description = "langgraph"
authors = []
license = "MIT"
tests/__snapshots__/test_pregel.ambr
+187 -1219
View File
File diff suppressed because it is too large Load Diff
tests/__snapshots__/test_pregel_async.ambr
+24 -24
View File
@@ -10,12 +10,12 @@
+---------------+
| rewrite_query |
+---------------+
*** ***
* *
** **
+--------------+ +---------+
| analyzer_one | | decider |
+--------------+ +---------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
@@ -47,12 +47,12 @@
+---------------+
| rewrite_query |
+---------------+
*** ***
* *
** **
+--------------+ +---------+
| analyzer_one | | decider |
+--------------+ +---------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
@@ -84,12 +84,12 @@
+---------------+
| rewrite_query |
+---------------+
*** ***
* *
** **
+--------------+ +-----------+
| analyzer_one | | condition |
+--------------+ +-----------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
@@ -121,12 +121,12 @@
+---------------+
| rewrite_query |
+---------------+
*** ***
* *
** **
+--------------+ +-----------+
| analyzer_one | | condition |
+--------------+ +-----------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
tests/test_pregel.py
+669 -75
View File
File diff suppressed because it is too large Load Diff
tests/test_pregel_async.py
+546 -40
View File
@@ -38,6 +38,57 @@ from tests.any_str import AnyStr
from tests.memory_assert import MemorySaverAssertImmutable
async def test_node_cancellation_on_external_cancel() -> None:
inner_task_cancelled = False
async def awhile(input: Any) -> None:
try:
await asyncio.sleep(1)
except asyncio.CancelledError:
nonlocal inner_task_cancelled
inner_task_cancelled = True
raise
builder = Graph()
builder.add_node("agent", awhile)
builder.set_entry_point("agent")
builder.set_finish_point("agent")
graph = builder.compile()
with pytest.raises(asyncio.TimeoutError):
await asyncio.wait_for(graph.ainvoke(1), 0.5)
assert inner_task_cancelled
async def test_node_cancellation_on_other_node_exception() -> None:
inner_task_cancelled = False
async def awhile(input: Any) -> None:
try:
await asyncio.sleep(1)
except asyncio.CancelledError:
nonlocal inner_task_cancelled
inner_task_cancelled = True
raise
async def iambad(input: Any) -> None:
raise ValueError("I am bad")
builder = Graph()
builder.add_node("agent", awhile)
builder.add_node("bad", iambad)
builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END)
graph = builder.compile()
with pytest.raises(ValueError, match="I am bad"):
await graph.ainvoke(1)
assert inner_task_cancelled
async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
@@ -80,19 +131,6 @@ async def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> No
assert falsy_value == await gapp.ainvoke(1)
async def test_invoke_single_process_in_out_implicit_channels(
mocker: MockerFixture,
) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Channel.subscribe_to("input") | 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) == 3
async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = (
@@ -102,17 +140,25 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N
)
app = Pregel(
nodes={"one": chain}, output_channels=["output", "fixed", "output_plus_one"]
nodes={"one": chain},
channels={
"input": LastValue(int),
"output": LastValue(int),
"fixed": LastValue(int),
"output_plus_one": LastValue(int),
},
output_channels=["output", "fixed", "output_plus_one"],
input_channels="input",
)
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"}
assert app.output_schema.schema() == {
"title": "LangGraphOutput",
"type": "object",
"properties": {
"output": {"title": "Output"},
"fixed": {"title": "Fixed"},
"output_plus_one": {"title": "Output Plus One"},
"output": {"title": "Output", "type": "integer"},
"fixed": {"title": "Fixed", "type": "integer"},
"output_plus_one": {"title": "Output Plus One", "type": "integer"},
},
}
assert await app.ainvoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4}
@@ -124,14 +170,16 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
app = Pregel(
nodes={"one": chain},
channels={"input": LastValue(int), "output": LastValue(int)},
input_channels="input",
output_channels=["output"],
)
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"}
assert app.output_schema.schema() == {
"title": "LangGraphOutput",
"type": "object",
"properties": {"output": {"title": "Output"}},
"properties": {"output": {"title": "Output", "type": "integer"}},
}
assert await app.ainvoke(2) == {"output": 3}
@@ -141,9 +189,8 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) ->
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
nodes={
"one": chain,
},
nodes={"one": chain},
channels={"input": LastValue(int), "output": LastValue(int)},
input_channels=["input"],
output_channels=["output"],
)
@@ -151,12 +198,12 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) ->
assert app.input_schema.schema() == {
"title": "LangGraphInput",
"type": "object",
"properties": {"input": {"title": "Input"}},
"properties": {"input": {"title": "Input", "type": "integer"}},
}
assert app.output_schema.schema() == {
"title": "LangGraphOutput",
"type": "object",
"properties": {"output": {"title": "Output"}},
"properties": {"output": {"title": "Output", "type": "integer"}},
}
assert await app.ainvoke({"input": 2}) == {"output": 3}
@@ -166,7 +213,17 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output")
app = Pregel(nodes={"one": one, "two": two}, stream_channels=["inbox", "output"])
app = Pregel(
nodes={"one": one, "two": two},
channels={
"inbox": LastValue(int),
"output": LastValue(int),
"input": LastValue(int),
},
input_channels="input",
output_channels="output",
stream_channels=["inbox", "output"],
)
assert await app.ainvoke(2) == 4
@@ -226,6 +283,13 @@ async def test_invoke_two_processes_in_out_interrupt(
memory = MemorySaverAssertImmutable(at=checkpoint_at)
app = Pregel(
nodes={"one": one, "two": two},
channels={
"inbox": LastValue(int),
"output": LastValue(int),
"input": LastValue(int),
},
input_channels="input",
output_channels="output",
checkpointer=memory,
interrupt_after_nodes=["one"],
)
@@ -281,9 +345,14 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
app = Pregel(
nodes={"one": one, "two": two},
channels={"inbox": Topic(int)},
channels={
"inbox": Topic(int),
"output": LastValue(int),
"input": LastValue(int),
},
input_channels=["input", "inbox"],
stream_channels=["inbox", "output"],
stream_channels=["output", "inbox"],
output_channels=["output"],
)
# [12 + 1, 2 + 1 + 1]
@@ -310,6 +379,88 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
{"inbox": [3], "output": 13},
{"inbox": [], "output": 4},
]
assert [
c async for c in app.astream({"input": 2, "inbox": 12}, stream_mode="debug")
] == [
{
"type": "task",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"id": "7a3cc398-2e02-5023-ad7b-e4848d3b67fa",
"name": "one",
"input": 2,
"triggers": ["input"],
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"id": "34e90af0-f97e-54e0-a159-691da37f175f",
"name": "two",
"input": [12],
"triggers": ["inbox"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"id": "7a3cc398-2e02-5023-ad7b-e4848d3b67fa",
"result": [("inbox", 3)],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"id": "34e90af0-f97e-54e0-a159-691da37f175f",
"result": [("output", 13)],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 0,
"payload": {"config": None, "values": {"output": 13, "inbox": [3]}},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "cf7cf374-2a2a-556f-8561-91737af89d2f",
"name": "two",
"input": [3],
"triggers": ["inbox"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "cf7cf374-2a2a-556f-8561-91737af89d2f",
"result": [("output", 4)],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 1,
"payload": {"config": None, "values": {"output": 4, "inbox": []}},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 2,
"payload": {"config": None, "values": {"output": 4, "inbox": []}},
},
]
async def test_batch_two_processes_in_out() -> None:
@@ -322,7 +473,13 @@ async def test_batch_two_processes_in_out() -> None:
app = Pregel(
nodes={"one": one, "two": two},
channels={"one": LastValue(int)},
channels={
"one": LastValue(int),
"output": LastValue(int),
"input": LastValue(int),
},
input_channels="input",
output_channels="output",
)
assert await app.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
@@ -356,7 +513,13 @@ async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
)
nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
app = Pregel(nodes=nodes)
app = Pregel(
nodes=nodes,
channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
| {"input": LastValue(int), "output": LastValue(int)},
input_channels="input",
output_channels="output",
)
# No state is left over from previous invocations
for _ in range(10):
@@ -379,7 +542,13 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
)
nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
app = Pregel(nodes=nodes)
app = Pregel(
nodes=nodes,
channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
| {"input": LastValue(int), "output": LastValue(int)},
input_channels="input",
output_channels="output",
)
# No state is left over from previous invocations
for _ in range(3):
@@ -409,7 +578,12 @@ async def test_invoke_two_processes_two_in_two_out_invalid(
one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(nodes={"one": one, "two": two})
app = Pregel(
nodes={"one": one, "two": two},
channels={"output": LastValue(int), "input": LastValue(int)},
input_channels="input",
output_channels="output",
)
with pytest.raises(InvalidUpdateError):
# LastValue channels can only be updated once per iteration
@@ -424,7 +598,12 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture)
app = Pregel(
nodes={"one": one, "two": two},
channels={"output": Topic(int)},
channels={
"input": LastValue(int),
"output": Topic(int),
},
input_channels="input",
output_channels="output",
)
# An Topic channel accumulates updates into a sequence
@@ -455,7 +634,13 @@ async def test_invoke_checkpoint(
app = Pregel(
nodes={"one": one},
channels={"total": BinaryOperatorAggregate(int, operator.add)},
channels={
"total": BinaryOperatorAggregate(int, operator.add),
"input": LastValue(int),
"output": LastValue(int),
},
input_channels="input",
output_channels="output",
checkpointer=memory,
)
@@ -510,7 +695,13 @@ async def test_invoke_checkpoint_aiosqlite(
memory.at = checkpoint_at
app = Pregel(
nodes={"one": one},
channels={"total": BinaryOperatorAggregate(int, operator.add)},
channels={
"total": BinaryOperatorAggregate(int, operator.add),
"input": LastValue(int),
"output": LastValue(int),
},
input_channels="input",
output_channels="output",
checkpointer=memory,
debug=True,
)
@@ -605,7 +796,13 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -
"chain_three": chain_three,
"chain_four": chain_four,
},
channels={"inbox": Topic(int)},
channels={
"inbox": Topic(int),
"output": LastValue(int),
"input": LastValue(int),
},
input_channels="input",
output_channels="output",
)
# Then invoke app
@@ -619,14 +816,20 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -
]
async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None:
async def test_invoke_join_then_call_other_pregel(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
inner_app = Pregel(
nodes={
"one": Channel.subscribe_to("input") | add_one | Channel.write_to("output")
}
},
channels={
"output": LastValue(int),
"input": LastValue(int),
},
input_channels="input",
output_channels="output",
)
one = (
@@ -651,7 +854,11 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None
channels={
"inbox_one": Topic(int),
"outbox_one": LastValue(int),
"output": LastValue(int),
"input": LastValue(int),
},
input_channels="input",
output_channels="output",
)
# Then invoke pubsub
@@ -673,7 +880,17 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non
)
two = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
app = Pregel(nodes={"one": one, "two": two}, stream_channels=["output", "between"])
app = Pregel(
nodes={"one": one, "two": two},
channels={
"input": LastValue(int),
"between": LastValue(int),
"output": LastValue(int),
},
stream_channels=["output", "between"],
input_channels="input",
output_channels="output",
)
# Then invoke pubsub
assert [c async for c in app.astream(2)] == [
@@ -687,7 +904,16 @@ async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
one = Channel.subscribe_to("input") | add_one | Channel.write_to("between")
two = Channel.subscribe_to("between") | add_one
app = Pregel(nodes={"one": one, "two": two})
app = Pregel(
nodes={"one": one, "two": two},
channels={
"input": LastValue(int),
"between": LastValue(int),
"output": LastValue(int),
},
input_channels="input",
output_channels="output",
)
# It finishes executing (once no more messages being published)
# but returns nothing, as nothing was published to "output" topic
@@ -727,9 +953,12 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
app = Pregel(
nodes={"one": one, "two": two},
channels={
"input": LastValue(int),
"output": LastValue(int),
"inbox": Topic(int),
"ctx": Context(an_int, an_int_async, typ=int),
},
input_channels="input",
output_channels=["inbox", "output"],
stream_channels=["inbox", "output"],
)
@@ -2798,6 +3027,250 @@ async def test_branch_then(
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
saver.at = checkpoint_at
# test stream_mode=debug
tool_two = tool_two_graph.compile(checkpointer=saver)
thread10 = {"configurable": {"thread_id": "10"}}
if checkpoint_at is CheckpointAt.END_OF_RUN:
assert [
c
async for c in tool_two.astream(
{"my_key": "value", "market": "DE"}, thread10, stream_mode="debug"
)
] == [
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"config": None,
"values": {"my_key": "value", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"input": {"my_key": "value", "market": "DE"},
"triggers": ["start:prepare"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"result": [("my_key", " prepared")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"config": None,
"values": {"my_key": "value prepared", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"input": {"my_key": "value prepared", "market": "DE"},
"triggers": ["branch:prepare:condition:tool_two_slow"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"result": [("my_key", " slow")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"config": None,
"values": {"my_key": "value prepared slow", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"input": {"my_key": "value prepared slow", "market": "DE"},
"triggers": ["branch:prepare:condition:then"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"result": [("my_key", " finished")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"config": None,
"values": {
"my_key": "value prepared slow finished",
"market": "DE",
},
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 4,
"payload": {
"config": {
"configurable": {
"thread_id": "10",
"thread_ts": AnyStr(),
}
},
"values": {
"my_key": "value prepared slow finished",
"market": "DE",
},
},
},
]
else:
assert [
c
async for c in tool_two.astream(
{"my_key": "value", "market": "DE"}, thread10, stream_mode="debug"
)
] == [
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"input": {"my_key": "value", "market": "DE"},
"triggers": ["start:prepare"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"result": [("my_key", " prepared")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value prepared", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"input": {"my_key": "value prepared", "market": "DE"},
"triggers": ["branch:prepare:condition:tool_two_slow"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"result": [("my_key", " slow")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value prepared slow", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"input": {"my_key": "value prepared slow", "market": "DE"},
"triggers": ["branch:prepare:condition:then"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"result": [("my_key", " finished")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {
"my_key": "value prepared slow finished",
"market": "DE",
},
},
},
]
tool_two = tool_two_graph.compile(
checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
)
@@ -3563,3 +4036,36 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None:
]
}
assert times_called == 1
chain = app | RunnablePassthrough()
assert await chain.ainvoke(
{"my_key": "my value", "never_called": never_called}
) == {
"my_key": "my value there and back again",
"never_called": never_called,
}
assert [
chunk
async for chunk in chain.astream(
{"my_key": "my value", "never_called": never_called}
)
] == [
{"inner": {"my_key": "my value there"}},
{"side": {"my_key": "my value there and back again"}},
]
times_called = 0
async for event in chain.astream_events(
{"my_key": "my value", "never_called": never_called},
version="v1",
config={"run_id": UUID(int=0)},
):
if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)):
times_called += 1
assert event["data"] == {
"output": [
{"inner": {"my_key": "my value there"}},
{"side": {"my_key": "my value there and back again"}},
]
}
assert times_called == 1