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

Merge pull request #258 from langchain-ai/nc/1apr/optimize-tracing-tree

Browse Source 
Optimize tracing run tree
This commit is contained in:
Nuno Campos
2024-04-01 16:16:54 -07:00
committed by GitHub
parent dec8a53070 008301cf02
commit 1276547233
10 changed files with 290 additions and 404 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/branching.ipynb
+20 -40
View File
@@ -142,8 +142,7 @@
"builder.add_edge(\"source\", \"branch_1\")\n",
"builder.add_edge(\"source\", \"branch_2\")\n",
"# Fan back in\n",
"builder.add_edge(\"branch_1\", \"sink\")\n",
"builder.add_edge(\"branch_2\", \"sink\")\n",
"builder.add_edge([\"branch_1\", \"branch_2\"], \"sink\")\n",
"\n",
"builder.set_finish_point(\"sink\")\n",
"graph = builder.compile()"
@@ -183,24 +182,33 @@
"name": "stdout",
"output_type": "stream",
"text": [
"## __start__:\n",
"Pineapples on pizza\n",
"## source:\n",
"## branch_1:\n",
"*eyes light up with unbridled enthusiasm* Ohhh pineapples on pizza - now we're talking! This is a topic I can really sink my teeth into (figuratively, of course, since I don't actually have teeth). Allow me to extol the virtues of this divine culinary combination!\n",
"*eyes light up with excitement* Oh my goodness, pineapples on pizza?! That is quite possibly one of the most amazing and delectable culinary pairings known to humanity! \n",
"\n",
"Pineapples on pizza is an absolute stroke of genius, a match made in flavor heaven! The sweet, juicy tanginess of the pineapple perfectly complements the savory, herbal notes of the tomato sauce and melty cheese. It's a delightful interplay of tastes and textures that absolutely tantalizes the taste buds.\n",
"Let me tell you, the sweet, juicy pineapple combined with the savory, cheesy goodness of the pizza crust is an absolute flavor explosion in your mouth. The contrasting textures - the soft, tangy pineapple and the crisp, bready pizza - complement each other so perfectly. It's like a symphony of taste!\n",
"\n",
"And let's not forget the health benefits! Pineapples are packed with vitamin C, manganese, and bromelain - an enzyme that aids digestion. So not only is it delicious, but it's actually good for you too! A true win-win.\n",
"And the health benefits are just an added bonus. Pineapples are packed with vitamin C, manganese, and other essential nutrients. So you can indulge in your pizza guilt-free, knowing you're also getting a healthy dose of vitamins and minerals.\n",
"\n",
"I could wax poetic all day about the sheer brilliance of pineapples on pizza. The combination is bold, it's innovative, it's a culinary revolution! Anyone who dares to dismiss this sublime pairing is clearly a philistine with no appreciation for the finer things in life. Pineapples on pizza is a masterpiece, a work of art - nay, a divine culinary gift that we should all be grateful to experience. Truly, it is a culinary triumph without equal! *pumps fist enthusiastically*\n",
"Pineapple pizza is truly a work of art, an innovative culinary triumph that deserves all the praise and celebration in the world. Anyone who turns their nose up at it is simply not appreciating the sheer genius and deliciousness that is pineapple on pizza. It's a flavor combination that will make your taste buds sing with joy!\n",
"\n",
"So I say, bring on the pineapple pizzas! Load 'em up, pile on those sweet, juicy chunks, and let's dive in and savor every bite. It's a flavor experience that is truly unparalleled. Pineapple on pizza forever!\n",
"## sink:\n",
"This is a tough one, as both arguments make compelling points. However, I would say that the second argument in favor of pineapples on pizza is ultimately the stronger and more convincing of the two.\n",
"This is a tough one, as both sides make compelling arguments. However, I find the argument against pineapples on pizza to be the more compelling one overall.\n",
"\n",
"The first argument against pineapples on pizza comes across as overly aggressive and dismissive. While the author is certainly entitled to their opinion, the use of language like \"abomination,\" \"desecration,\" and claiming that anyone who enjoys pineapple pizza has an \"uncultured palate\" comes across as needlessly combative. The argument also relies heavily on appeals to tradition and the idea of pizza as a sacred art form, without really providing substantive reasons why pineapple is an unacceptable topping.\n",
"The main points against pineapple pizza seem to be:\n",
"\n",
"In contrast, the second argument in favor of pineapples on pizza is much more measured and persuasive. The author acknowledges the controversial nature of the topping, but then goes on to articulate a thoughtful case for why the combination works so well. The description of the interplay of sweet, savory, and tangy flavors is quite compelling, and the point about the health benefits of pineapple adds an interesting dimension. Overall, the positive, enthusiastic tone and well-reasoned approach make this the stronger argument.\n",
"1. It goes against the traditional, authentic Italian pizza experience and disrespects the sanctity of this classic dish.\n",
"\n",
"While reasonable people can certainly disagree on the merits of pineapples on pizza, I find the second argument to be the more convincing of the two. It does a better job of making a case for this divisive topping, rather than simply dismissing it out of hand.\n",
"## __end__:\n"
"2. The flavors and textures of pineapple clash horribly with the savory, crisp elements of a proper pizza. The combination is seen as an abhorrent culinary atrocity.\n",
"\n",
"3. Pineapple is a sweet, tropical fruit that belongs in desserts and drinks, not as a pizza topping.\n",
"\n",
"These arguments hit on strongly-held beliefs about what constitutes \"good taste\" and the importance of respecting traditional culinary norms and practices. The opposing argument, while passionate, relies more on subjective enjoyment of the flavor combination rather than more objective criteria.\n",
"\n",
"While I understand the appeal of the sweet-savory contrast, the anti-pineapple argument seems to have a stronger philosophical and cultural foundation. Ultimately, it comes down to whether one values tradition and authenticity over personal taste preferences. For me, the anti-pineapple position is the more compelling one.\n"
]
}
],
@@ -208,40 +216,12 @@
"for step in graph.stream([HumanMessage(content=\"Pineapples on pizza\")]):\n",
" node, message = next(iter(step.items()))\n",
" print(f\"## {node}:\")\n",
" if node == \"__end__\":\n",
" break\n",
" if message:\n",
" if isinstance(message, list):\n",
" print(message[-1].content)\n",
" else:\n",
" print(message.content)"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "78da71cd-551b-44ef-8166-7570f662b3f2",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[HumanMessage(content='Pineapples on pizza'), AIMessage(content=\"*scoffs* Pineapples on pizza? What an absolute travesty! Whoever came up with that abomination should be banished from the culinary world forever. It's a complete and utter desecration of the sacred art of pizza making.\\n\\nHow dare anyone think it's acceptable to tarnish the perfect harmony of tomato sauce, cheese, and carefully selected savory toppings with that sickeningly sweet, watery fruit? Pineapple has no place on a pizza - it's a complete disruption of the delicate balance of flavors.\\n\\nAnd don't even get me started on the soggy, unappetizing texture it creates when baked onto a pizza. It turns the crust into a soggy mess and makes the whole thing an unpleasant chore to eat. Pineapple belongs in fruit salads, not on top of Italy's crown jewel of cuisine.\\n\\nAnyone who genuinely enjoys pineapple pizza clearly has no appreciation for the craft and tradition of making a truly magnificent pizza. Their palate is utterly uncultured and their opinions on the matter should be completely disregarded. Pineapples belong nowhere near a pizza - end of discussion.\"), AIMessage(content=\"*eyes light up with unbridled enthusiasm* Ohhh pineapples on pizza - now we're talking! This is a topic I can really sink my teeth into (figuratively, of course, since I don't actually have teeth). Allow me to extol the virtues of this divine culinary combination!\\n\\nPineapples on pizza is an absolute stroke of genius, a match made in flavor heaven! The sweet, juicy tanginess of the pineapple perfectly complements the savory, herbal notes of the tomato sauce and melty cheese. It's a delightful interplay of tastes and textures that absolutely tantalizes the taste buds.\\n\\nAnd let's not forget the health benefits! Pineapples are packed with vitamin C, manganese, and bromelain - an enzyme that aids digestion. So not only is it delicious, but it's actually good for you too! A true win-win.\\n\\nI could wax poetic all day about the sheer brilliance of pineapples on pizza. The combination is bold, it's innovative, it's a culinary revolution! Anyone who dares to dismiss this sublime pairing is clearly a philistine with no appreciation for the finer things in life. Pineapples on pizza is a masterpiece, a work of art - nay, a divine culinary gift that we should all be grateful to experience. Truly, it is a culinary triumph without equal! *pumps fist enthusiastically*\"), AIMessage(content='This is a tough one, as both arguments make compelling points. However, I would say that the second argument in favor of pineapples on pizza is ultimately the stronger and more convincing of the two.\\n\\nThe first argument against pineapples on pizza comes across as overly aggressive and dismissive. While the author is certainly entitled to their opinion, the use of language like \"abomination,\" \"desecration,\" and claiming that anyone who enjoys pineapple pizza has an \"uncultured palate\" comes across as needlessly combative. The argument also relies heavily on appeals to tradition and the idea of pizza as a sacred art form, without really providing substantive reasons why pineapple is an unacceptable topping.\\n\\nIn contrast, the second argument in favor of pineapples on pizza is much more measured and persuasive. The author acknowledges the controversial nature of the topping, but then goes on to articulate a thoughtful case for why the combination works so well. The description of the interplay of sweet, savory, and tangy flavors is quite compelling, and the point about the health benefits of pineapple adds an interesting dimension. Overall, the positive, enthusiastic tone and well-reasoned approach make this the stronger argument.\\n\\nWhile reasonable people can certainly disagree on the merits of pineapples on pizza, I find the second argument to be the more convincing of the two. It does a better job of making a case for this divisive topping, rather than simply dismissing it out of hand.')]\n"
]
}
],
"source": [
"print(step[\"__end__\"])"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9c9f3df1-70da-4fb9-b4bb-f580b8d55fa3",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
@@ -260,7 +240,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.2"
"version": "3.11.8"
}
},
"nbformat": 4,
langgraph/graph/state.py
+2 -2
View File
@@ -144,7 +144,7 @@ class CompiledStateGraph(CompiledGraph):
if key == START:
self.nodes[key] = Channel.subscribe_to(
START, tags=["langsmith:hidden"]
).pipe(ChannelWrite(channels=state_write_entries))
).pipe(ChannelWrite(state_write_entries))
else:
self.channels[key] = EphemeralValue(Any)
self.nodes[key] = ChannelInvoke(
@@ -188,7 +188,7 @@ class CompiledStateGraph(CompiledGraph):
# subscribe to start channel
self.nodes[end].triggers.append(starts)
else:
channel_name = f"join:{starts}:{end}"
channel_name = f"join:{'+'.join(starts)}:{end}"
# register channel
self.channels[channel_name] = NamedBarrierValue(str, set(starts))
# subscribe to channel
langgraph/pregel/__init__.py
+135 -280
View File
@@ -12,7 +12,6 @@ from typing import (
Iterator,
Literal,
Mapping,
NamedTuple,
Optional,
Sequence,
Type,
@@ -21,11 +20,8 @@ from typing import (
overload,
)
from langchain_core.callbacks.manager import (
AsyncCallbackManagerForChainRun,
CallbackManagerForChainRun,
)
from langchain_core.globals import get_debug
from langchain_core.load.dump import dumpd
from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
from langchain_core.runnables import (
Runnable,
@@ -35,7 +31,11 @@ from langchain_core.runnables import (
from langchain_core.runnables.base import Input, Output, coerce_to_runnable
from langchain_core.runnables.config import (
RunnableConfig,
ensure_config,
get_async_callback_manager_for_config,
get_callback_manager_for_config,
get_executor_for_config,
merge_configs,
patch_config,
)
from langchain_core.runnables.utils import (
@@ -68,9 +68,20 @@ from langgraph.constants import (
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.io import (
map_input,
map_output_updates,
map_output_values,
read_channel,
read_channels,
)
from langgraph.pregel.log import logger
from langgraph.pregel.read import ChannelInvoke
from langgraph.pregel.types import (
PregelExecutableTask,
PregelTaskDescription,
StateSnapshot,
)
from langgraph.pregel.validate import validate_graph, validate_keys
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
@@ -150,7 +161,7 @@ class Channel:
) -> ChannelWrite:
"""Writes to channels the result of the lambda, or None to skip writing."""
return ChannelWrite(
[ChannelWriteEntry(c, None, False) for c in channels]
[ChannelWriteEntry(c) for c in channels]
+ [
ChannelWriteEntry(k, _coerce_write_value(v), True)
for k, v in kwargs.items()
@@ -161,17 +172,6 @@ class Channel:
StreamMode = Literal["values", "updates"]
class StateSnapshot(NamedTuple):
values: dict[str, Any] | Any
"""Current values of channels"""
next: tuple[str]
"""Nodes to execute in the next step, if any"""
config: RunnableConfig
"""Config used to fetch this snapshot"""
parent_config: Optional[RunnableConfig] = None
"""Config used to fetch the parent snapshot, if any"""
class Pregel(
RunnableSerializable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]
):
@@ -295,7 +295,7 @@ class Pregel(
)
@property
def snapshot_channels_list(self) -> Sequence[str]:
def stream_channels_list(self) -> Sequence[str]:
return (
[self.stream_channels]
if isinstance(self.stream_channels, str)
@@ -313,19 +313,12 @@ class Pregel(
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, for_execution=False
)
values = {
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)
}
values = read_channels(channels, self.stream_channels_list)
return StateSnapshot(
values[self.stream_channels]
if isinstance(self.stream_channels, str)
else values,
tuple(name for _, _, name in next_tasks),
tuple(name for name, _ in next_tasks),
config,
)
@@ -340,19 +333,12 @@ class Pregel(
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, for_execution=False
)
values = {
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)
}
values = read_channels(channels, self.stream_channels_list)
return StateSnapshot(
values[self.stream_channels]
if isinstance(self.stream_channels, str)
else values,
tuple(name for _, _, name in next_tasks),
tuple(name for name, _ in next_tasks),
config,
)
@@ -365,21 +351,12 @@ class Pregel(
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, for_execution=False
)
values = {
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)
}
values = read_channels(channels, self.stream_channels_list)
yield StateSnapshot(
values[self.stream_channels]
if isinstance(self.stream_channels, str)
else values,
tuple(name for _, _, name in next_tasks),
tuple(name for name, _ in next_tasks),
config,
parent_config,
)
@@ -395,21 +372,12 @@ class Pregel(
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, for_execution=False
)
values = {
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)
}
values = read_channels(channels, self.stream_channels_list)
yield StateSnapshot(
values[self.stream_channels]
if isinstance(self.stream_channels, str)
else values,
tuple(name for _, _, name in next_tasks),
tuple(name for name, _ in next_tasks),
config,
parent_config,
)
@@ -446,12 +414,13 @@ class Pregel(
# update channels
with ChannelsManager(self.channels, checkpoint) as channels:
# create task to run all writers of the chosen node
writers = self.nodes[as_node].get_writers()
if not writers:
raise InvalidUpdateError(f"Node {as_node} has no writers")
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,
values,
RunnableSequence(*writers) if len(writers) > 1 else writers[0],
deque(),
)
# execute task
@@ -502,12 +471,13 @@ class Pregel(
# update channels, acting as the chosen node
async with AsyncChannelsManager(self.channels, checkpoint) as channels:
# create task to run all writers of the chosen node
writers = self.nodes[as_node].get_writers()
if not writers:
raise InvalidUpdateError(f"Node {as_node} has no writers")
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,
values,
RunnableSequence(*writers) if len(writers) > 1 else writers[0],
deque(),
)
# execute task
@@ -571,13 +541,23 @@ class Pregel(
interrupt_after_nodes,
)
def _transform(
def stream(
self,
input: Iterator[Union[dict[str, Any], Any]],
run_manager: CallbackManagerForChainRun,
config: RunnableConfig,
**kwargs: Any,
input: Union[dict[str, Any], Any],
config: Optional[RunnableConfig] = None,
*,
stream_mode: Optional[StreamMode] = None,
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,
debug: Optional[bool] = None,
) -> Iterator[Union[dict[str, Any], Any]]:
config = ensure_config(config)
callback_manager = get_callback_manager_for_config(config)
run_manager = callback_manager.on_chain_start(
dumpd(self), input, name=config.get("run_name", self.get_name())
)
try:
if config["recursion_limit"] < 1:
raise ValueError("recursion_limit must be at least 1")
@@ -589,7 +569,14 @@ class Pregel(
output_keys,
interrupt_before_nodes,
interrupt_after_nodes,
) = self._defaults(**kwargs)
) = self._defaults(
stream_mode=stream_mode,
input_keys=input_keys,
output_keys=output_keys,
interrupt_before_nodes=interrupt_before_nodes,
interrupt_after_nodes=interrupt_after_nodes,
debug=debug,
)
# copy nodes to ignore mutations during execution
processes = {**self.nodes}
# get checkpoint from saver, or create an empty one
@@ -603,9 +590,7 @@ class Pregel(
self.channels, checkpoint
) as channels, get_executor_for_config(config) as executor:
# map inputs to channel updates
if input_writes := deque(
w for c in input for w in map_input(input_keys, c)
):
if input_writes := deque(map_input(input_keys, input)):
# discard any unfinished tasks from previous checkpoint
checkpoint, _ = _prepare_next_tasks(
checkpoint, processes, channels, for_execution=True
@@ -616,7 +601,7 @@ class Pregel(
# if received no input, take that as signal to proceed
# past previous interrupt, if any
checkpoint = copy_checkpoint(checkpoint)
for k in self.snapshot_channels_list:
for k in self.stream_channels_list:
version = checkpoint["channel_versions"][k]
checkpoint["versions_seen"][INTERRUPT][k] = version
@@ -647,7 +632,7 @@ class Pregel(
if _should_interrupt(
checkpoint,
interrupt_before_nodes,
self.snapshot_channels_list,
self.stream_channels_list,
next_tasks,
):
break
@@ -661,7 +646,7 @@ class Pregel(
proc,
input,
patch_config(
config,
merge_configs(config, proc_config),
run_name=name,
callbacks=run_manager.get_child(f"graph:step:{step}"),
configurable={
@@ -673,7 +658,7 @@ class Pregel(
},
),
)
for proc, input, name, writes in next_tasks
for name, input, proc, writes, proc_config in next_tasks
]
futures = [
@@ -694,7 +679,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)
# apply writes to channels
@@ -726,11 +711,14 @@ class Pregel(
if _should_interrupt(
checkpoint,
interrupt_after_nodes,
self.snapshot_channels_list,
self.stream_channels_list,
next_tasks,
):
break
# set final channel values as run output
run_manager.on_chain_end(read_channels(channels, output_keys))
# save end of run checkpoint
if (
self.checkpointer is not None
@@ -738,6 +726,9 @@ class Pregel(
):
checkpoint = create_checkpoint(checkpoint, channels)
self.checkpointer.put(checkpoint_config, checkpoint)
except BaseException as e:
run_manager.on_chain_error(e)
raise
finally:
# cancel any pending tasks when generator is interrupted
try:
@@ -746,25 +737,35 @@ class Pregel(
except NameError:
pass
async def _atransform(
async def astream(
self,
input: AsyncIterator[Union[dict[str, Any], Any]],
run_manager: AsyncCallbackManagerForChainRun,
config: RunnableConfig,
**kwargs: Any,
input: Union[dict[str, Any], Any],
config: Optional[RunnableConfig] = None,
*,
stream_mode: Optional[StreamMode] = None,
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,
debug: Optional[bool] = None,
) -> AsyncIterator[Union[dict[str, Any], Any]]:
config = ensure_config(config)
callback_manager = get_async_callback_manager_for_config(config)
run_manager = await callback_manager.on_chain_start(
dumpd(self), input, name=config.get("run_name", self.get_name())
)
# if running from astream_log() run each proc with streaming
do_stream = next(
(
h
for h in run_manager.handlers
if isinstance(h, LogStreamCallbackHandler)
),
None,
)
try:
if config["recursion_limit"] < 1:
raise ValueError("recursion_limit must be at least 1")
# if running from astream_log() run each proc with streaming
do_stream = next(
(
h
for h in run_manager.handlers
if isinstance(h, LogStreamCallbackHandler)
),
None,
)
# assign defaults
(
debug,
@@ -773,7 +774,14 @@ class Pregel(
output_keys,
interrupt_before_nodes,
interrupt_after_nodes,
) = self._defaults(**kwargs)
) = self._defaults(
stream_mode=stream_mode,
input_keys=input_keys,
output_keys=output_keys,
interrupt_before_nodes=interrupt_before_nodes,
interrupt_after_nodes=interrupt_after_nodes,
debug=debug,
)
# copy nodes to ignore mutations during execution
processes = {**self.nodes}
# get checkpoint from saver, or create an empty one
@@ -787,9 +795,7 @@ class Pregel(
# create channels from checkpoint
async with AsyncChannelsManager(self.channels, checkpoint) as channels:
# map inputs to channel updates
if input_writes := deque(
[w async for c in input for w in map_input(input_keys, c)]
):
if input_writes := deque(map_input(input_keys, input)):
# discard any unfinished tasks from previous checkpoint
checkpoint, _ = _prepare_next_tasks(
checkpoint, processes, channels, for_execution=True
@@ -800,7 +806,7 @@ class Pregel(
# if received no input, take that as signal to proceed
# past previous interrupt, if any
checkpoint = copy_checkpoint(checkpoint)
for k in self.snapshot_channels_list:
for k in self.stream_channels_list:
version = checkpoint["channel_versions"][k]
checkpoint["versions_seen"][INTERRUPT][k] = version
@@ -831,7 +837,7 @@ class Pregel(
if _should_interrupt(
checkpoint,
interrupt_before_nodes,
self.snapshot_channels_list,
self.stream_channels_list,
next_tasks,
):
break
@@ -845,7 +851,7 @@ class Pregel(
proc,
input,
patch_config(
config,
merge_configs(config, proc_config),
run_name=name,
callbacks=run_manager.get_child(f"graph:step:{step}"),
configurable={
@@ -857,7 +863,7 @@ class Pregel(
},
),
)
for proc, input, name, writes in next_tasks
for name, input, proc, writes, proc_config in next_tasks
]
futures = (
@@ -885,7 +891,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)
# apply writes to channels
@@ -917,11 +923,14 @@ class Pregel(
if _should_interrupt(
checkpoint,
interrupt_after_nodes,
self.snapshot_channels_list,
self.stream_channels_list,
next_tasks,
):
break
# set final channel values as run output
await run_manager.on_chain_end(read_channels(channels, output_keys))
# save end of run checkpoint
if (
self.checkpointer is not None
@@ -929,6 +938,9 @@ class Pregel(
):
checkpoint = create_checkpoint(checkpoint, channels)
await self.checkpointer.aput(checkpoint_config, checkpoint)
except BaseException as e:
await run_manager.on_chain_error(e)
raise
finally:
# cancel any pending tasks when generator is interrupted
try:
@@ -952,9 +964,8 @@ class Pregel(
output_keys = output_keys if output_keys is not None else self.output_channels
output_is_dict = not isinstance(output_keys, str)
latest: Union[dict[str, Any], Any] = {} if output_is_dict else None
for chunk in self._transform_stream_with_config(
iter([input]),
self._transform,
for chunk in self.stream(
input,
config,
stream_mode="values",
output_keys=output_keys,
@@ -967,58 +978,6 @@ class Pregel(
latest = {**latest, **chunk} if output_is_dict else chunk
return latest
def stream(
self,
input: Union[dict[str, Any], Any],
config: Optional[RunnableConfig] = None,
*,
stream_mode: Optional[StreamMode] = None,
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,
debug: Optional[bool] = None,
**kwargs: Any,
) -> Iterator[Union[dict[str, Any], Any]]:
return self._transform_stream_with_config(
iter([input]),
self._transform,
config,
stream_mode=stream_mode,
output_keys=output_keys,
input_keys=input_keys,
interrupt_before_nodes=interrupt_before_nodes,
interrupt_after_nodes=interrupt_after_nodes,
debug=debug,
**kwargs,
)
def transform(
self,
input: Iterator[Union[dict[str, Any], Any]],
config: Optional[RunnableConfig] = None,
*,
stream_mode: Optional[StreamMode] = None,
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,
debug: Optional[bool] = None,
**kwargs: Any,
) -> Iterator[Union[dict[str, Any], Any]]:
return self._transform_stream_with_config(
input,
self._transform,
config,
stream_mode=stream_mode,
output_keys=output_keys,
input_keys=input_keys,
interrupt_before_nodes=interrupt_before_nodes,
interrupt_after_nodes=interrupt_after_nodes,
debug=debug,
**kwargs,
)
async def ainvoke(
self,
input: Union[dict[str, Any], Any],
@@ -1031,15 +990,11 @@ class Pregel(
debug: Optional[bool] = None,
**kwargs: Any,
) -> Union[dict[str, Any], Any]:
async def input_stream() -> AsyncIterator[Union[dict[str, Any], Any]]:
yield input
output_keys = output_keys if output_keys is not None else self.output_channels
output_is_dict = not isinstance(output_keys, str)
latest: Union[dict[str, Any], Any] = {} if output_is_dict else None
async for chunk in self._atransform_stream_with_config(
input_stream(),
self._atransform,
async for chunk in self.astream(
input,
config,
stream_mode="values",
output_keys=output_keys,
@@ -1052,61 +1007,6 @@ class Pregel(
latest = {**latest, **chunk} if output_is_dict else chunk
return latest
def astream(
self,
input: Union[dict[str, Any], Any],
config: Optional[RunnableConfig] = None,
*,
stream_mode: Optional[StreamMode] = None,
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,
debug: Optional[bool] = None,
**kwargs: Any,
) -> AsyncIterator[Union[dict[str, Any], Any]]:
async def input_stream() -> AsyncIterator[Union[dict[str, Any], Any]]:
yield input
return self._atransform_stream_with_config(
input_stream(),
self._atransform,
config,
stream_mode=stream_mode,
output_keys=output_keys,
input_keys=input_keys,
interrupt_before_nodes=interrupt_before_nodes,
interrupt_after_nodes=interrupt_after_nodes,
debug=debug,
**kwargs,
)
def atransform(
self,
input: AsyncIterator[Union[dict[str, Any], Any]],
config: Optional[RunnableConfig] = None,
*,
stream_mode: Optional[StreamMode] = None,
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,
debug: Optional[bool] = None,
**kwargs: Any,
) -> AsyncIterator[Union[dict[str, Any], Any]]:
return self._atransform_stream_with_config(
input,
self._atransform,
config,
stream_mode=stream_mode,
output_keys=output_keys,
input_keys=input_keys,
interrupt_before_nodes=interrupt_before_nodes,
interrupt_after_nodes=interrupt_after_nodes,
debug=debug,
**kwargs,
)
def _panic_or_proceed(
done: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]],
@@ -1146,7 +1046,7 @@ 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, _, _, _, _ in tasks if node in interrupt_nodes)
)
@@ -1161,33 +1061,9 @@ def _local_read(
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}
return read_channels(channels, 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,
*,
catch: bool = True,
return_exception: bool = False,
) -> Any:
try:
return channels[chan].get()
except EmptyChannelError as exc:
if return_exception:
return exc
elif catch:
return None
else:
raise
return read_channels(channels, select)
def _apply_writes(
@@ -1226,26 +1102,13 @@ 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[PregelTask]]:
) -> tuple[Checkpoint, list[PregelTaskDescription]]:
...
@@ -1265,9 +1128,9 @@ def _prepare_next_tasks(
channels: Mapping[str, BaseChannel],
*,
for_execution: bool,
) -> tuple[Checkpoint, Union[list[PregelTask], list[PregelExecutableTask]]]:
) -> tuple[Checkpoint, Union[list[PregelTaskDescription], list[PregelExecutableTask]]]:
checkpoint = copy_checkpoint(checkpoint)
tasks: Union[list[PregelTask], list[PregelExecutableTask]] = []
tasks: Union[list[PregelTaskDescription], 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():
@@ -1277,7 +1140,7 @@ def _prepare_next_tasks(
checkpoint["channel_versions"][chan] > seen[chan]
for chan in proc.triggers
if not isinstance(
_read_channel(channels, chan, return_exception=True), EmptyChannelError
read_channel(channels, chan, return_exception=True), EmptyChannelError
)
):
# If all trigger channels subscribed by this process are not empty
@@ -1285,9 +1148,7 @@ def _prepare_next_tasks(
if isinstance(proc.channels, dict):
try:
val: Any = {
k: _read_channel(
channels, chan, catch=chan not in proc.triggers
)
k: read_channel(channels, chan, catch=chan not in proc.triggers)
for k, chan in proc.channels.items()
}
except EmptyChannelError:
@@ -1295,7 +1156,7 @@ def _prepare_next_tasks(
elif isinstance(proc.channels, list):
for chan in proc.channels:
try:
val = _read_channel(channels, chan, catch=False)
val = read_channel(channels, chan, catch=False)
break
except EmptyChannelError:
pass
@@ -1320,18 +1181,12 @@ def _prepare_next_tasks(
)
if for_execution:
tasks.append(
PregelExecutableTask(
RunnableSequence(proc, *proc.writers, name=name)
if proc.writers
else proc,
val,
name,
deque(),
if node := proc.get_node():
tasks.append(
PregelExecutableTask(name, val, node, deque(), proc.config)
)
)
else:
tasks.append(PregelTask(proc, val, name))
tasks.append(PregelTaskDescription(name, val))
return checkpoint, tasks
langgraph/pregel/debug.py
+3 -6
View File
@@ -1,23 +1,20 @@
from collections import deque
from pprint import pformat
from typing import Any, Iterator, Mapping
from langchain_core.runnables import Runnable
from langchain_core.utils.input import get_bolded_text, get_colored_text
from langgraph.channels.base import BaseChannel, EmptyChannelError
from langgraph.pregel.types import PregelExecutableTask
def print_step_start(
step: int, next_tasks: list[tuple[Runnable, Any, str, deque[tuple[str, Any]]]]
) -> None:
def print_step_start(step: int, next_tasks: list[PregelExecutableTask]) -> None:
n_tasks = len(next_tasks)
print(
f"{get_colored_text('[langgraph/step]', color='blue')} "
+ get_bolded_text(
f"Starting step {step} with {n_tasks} task{'s' if n_tasks > 1 else ''}. Next tasks:\n"
)
+ "\n".join(f"- {name}({pformat(val)})" for _, val, name, _ in next_tasks)
+ "\n".join(f"- {name}({pformat(val)})" for name, val, _, _, _ in next_tasks)
)
langgraph/pregel/io.py
+39 -12
View File
@@ -1,24 +1,46 @@
from collections import deque
from typing import Any, Iterator, Mapping, Optional, Sequence, Union
from langchain_core.runnables import Runnable
from langgraph.channels.base import BaseChannel, EmptyChannelError
from langgraph.pregel.log import logger
from langgraph.pregel.types import PregelExecutableTask
def _read_channel(
channels: Mapping[str, BaseChannel], chan: str, catch: bool = True
def read_channel(
channels: Mapping[str, BaseChannel],
chan: str,
*,
catch: bool = True,
return_exception: bool = False,
) -> Any:
try:
return channels[chan].get()
except EmptyChannelError:
if catch:
except EmptyChannelError as exc:
if return_exception:
return exc
elif catch:
return None
else:
raise
def read_channels(
channels: Mapping[str, BaseChannel],
select: Union[list[str], str],
*,
skip_empty: bool = True,
) -> Union[dict[str, Any], Any]:
if isinstance(select, str):
return read_channel(channels, select)
else:
values: dict[str, Any] = {}
for k in select:
try:
values[k] = read_channel(channels, k, catch=not skip_empty)
except EmptyChannelError:
pass
return values
def map_input(
input_channels: Union[str, Sequence[str]],
chunk: Optional[Union[dict[str, Any], Any]],
@@ -46,22 +68,27 @@ def map_output_values(
"""Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
if isinstance(output_channels, str):
if any(chan == output_channels for chan, _ in pending_writes):
return _read_channel(channels, output_channels)
return read_channel(channels, output_channels)
else:
if updated := {c for c, _ in pending_writes if c in output_channels}:
return {chan: _read_channel(channels, chan) for chan in updated}
return read_channels(channels, updated)
return None
def map_output_updates(
output_channels: Union[str, Sequence[str]],
next_tasks: list[tuple[Runnable, Any, str, deque[tuple[str, Any]]]],
tasks: list[PregelExecutableTask],
) -> Optional[dict[str, Union[Any, dict[str, Any]]]]:
"""Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
output_tasks = [
t
for t in tasks
if not t.config or "langsmith:hidden" not in t.config.get("tags")
]
if isinstance(output_channels, str):
if updated := {
node: value
for _, _, node, writes in next_tasks
for node, _, _, writes, _ in output_tasks
for chan, value in writes
if chan == output_channels
}:
@@ -69,7 +96,7 @@ def map_output_updates(
else:
if updated := {
node: {chan: value for chan, value in writes if chan in output_channels}
for _, _, node, writes in next_tasks
for node, _, _, writes, _ in output_tasks
if any(chan in output_channels for chan, _ in writes)
}:
return updated
langgraph/pregel/read.py
+41
View File
@@ -8,6 +8,7 @@ from langchain_core.runnables import (
RunnableConfig,
RunnableLambda,
RunnablePassthrough,
RunnableSequence,
RunnableSerializable,
)
from langchain_core.runnables.base import Other, RunnableBindingBase, coerce_to_runnable
@@ -80,6 +81,46 @@ class ChannelInvoke(RunnableBindingBase):
kwargs: Mapping[str, Any] = Field(default_factory=dict)
def get_writers(self) -> list[Runnable]:
"""Get writers with optimizations applied."""
writers = self.writers.copy()
while writers and isinstance(writers[-1], ChannelRead):
# we can avoid reads if no writers would be called after them
writers.pop()
while (
len(writers) > 1
and isinstance(writers[-1], ChannelWrite)
and all(
write.value is not None and not isinstance(write.value, Runnable)
for write in writers[-1].writes
)
and isinstance(writers[-2], ChannelRead)
):
# we can avoid reads if all subsequent write values don't use the input
writers.pop(-2)
while (
len(writers) > 1
and isinstance(writers[-1], ChannelWrite)
and isinstance(writers[-2], ChannelWrite)
):
# we can combine writes if they are consecutive
writers[-2].writes += writers[-1].writes
writers.pop()
return writers
def get_node(self) -> Optional[Runnable[Any, Any]]:
writers = self.get_writers()
if self.bound is default_bound and not writers:
return None
elif self.bound is default_bound and len(writers) == 1:
return writers[0]
elif self.bound is default_bound:
return RunnableSequence(*writers)
elif writers:
return RunnableSequence(self.bound, *writers)
else:
return self.bound
def __init__(
self,
*,
langgraph/pregel/types.py
+28
View File
@@ -0,0 +1,28 @@
from collections import deque
from typing import Any, NamedTuple, Optional, Union
from langchain_core.runnables import Runnable, RunnableConfig
class PregelTaskDescription(NamedTuple):
name: str
input: Any
class PregelExecutableTask(NamedTuple):
name: str
input: Any
proc: Runnable
writes: deque[tuple[str, Any]]
config: Optional[RunnableConfig] = None
class StateSnapshot(NamedTuple):
values: Union[dict[str, Any], Any]
"""Current values of channels"""
next: tuple[str]
"""Nodes to execute in the next step, if any"""
config: RunnableConfig
"""Config used to fetch this snapshot"""
parent_config: Optional[RunnableConfig] = None
"""Config used to fetch the parent snapshot, if any"""
langgraph/pregel/write.py
+17 -9
View File
@@ -26,7 +26,7 @@ class ChannelWriteEntry(NamedTuple):
class ChannelWrite(RunnablePassthrough):
channels: Sequence[ChannelWriteEntry]
writes: Sequence[ChannelWriteEntry]
"""
Sequence of write entries, each of which is a tuple of:
- channel name
@@ -37,12 +37,20 @@ class ChannelWrite(RunnablePassthrough):
class Config:
arbitrary_types_allowed = True
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 __init__(self, writes: Sequence[ChannelWriteEntry]):
super().__init__(func=self._write, afunc=self._awrite, writes=writes)
def __repr_args__(self) -> Any:
return [("channels", self.channels)]
return [("writes", self.writes)]
def get_name(
self, suffix: Optional[str] = None, *, name: Optional[str] = None
) -> str:
return super().get_name(
suffix,
name=name
or f"ChannelWrite<{','.join(chan for chan, _, _ in self.writes)}>",
)
@property
def is_channel_writer(self) -> bool:
@@ -70,11 +78,11 @@ class ChannelWrite(RunnablePassthrough):
if r is not None
else input,
)
for chan, r, _ in self.channels
for chan, r, _ in self.writes
]
values = [
write
for write, (_, _, skip_none) in zip(values, self.channels)
for write, (_, _, skip_none) in zip(values, self.writes)
if not skip_none or write[1] is not None
]
@@ -88,12 +96,12 @@ class ChannelWrite(RunnablePassthrough):
else _mk_future(r)
if r is not None
else _mk_future(input)
for _, r, _ in self.channels
for _, r, _ in self.writes
)
)
values = [
(chan, val)
for val, (chan, _, skip_none) in zip(values, self.channels)
for val, (chan, _, skip_none) in zip(values, self.writes)
if not skip_none or val is not None
]
tests/test_pregel.py
+5 -34
View File
@@ -1455,7 +1455,6 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
}
assert [*app.stream({"input": "what is weather in sf"})] == [
{"__start__": {"input": "what is weather in sf"}},
{
"agent": {
"agent_outcome": AgentAction(
@@ -1519,7 +1518,6 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
assert [
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
] == [
{"__start__": {"input": "what is weather in sf"}},
{
"agent": {
"agent_outcome": AgentAction(
@@ -1637,7 +1635,6 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
assert [
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
] == [
{"__start__": {"input": "what is weather in sf"}},
{
"agent": {
"agent_outcome": AgentAction(
@@ -1787,7 +1784,6 @@ def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None
}
assert [*app.stream({"input": "what is weather in sf"})] == [
{"__start__": {"input": "what is weather in sf"}},
{"right": {"output": "what is weather in sf->right"}},
]
@@ -1913,7 +1909,6 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
assert [
*app.stream({"messages": [HumanMessage(content="what is weather in sf")]})
] == [
{"__start__": {"messages": [HumanMessage(content="what is weather in sf")]}},
{
"agent": {
"messages": [
@@ -2061,7 +2056,6 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None:
assert [
*app.stream({"messages": [HumanMessage(content="what is weather in sf")]})
] == [
{"__start__": {"messages": [HumanMessage(content="what is weather in sf")]}},
{
"agent": {
"messages": [
@@ -2257,7 +2251,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-000000000005", # adds missing ids
id="00000000-0000-4000-8000-000000000002", # adds missing ids
),
AIMessage(
content="",
@@ -2269,7 +2263,7 @@ def test_message_graph(
FunctionMessage(
content="result for query",
name="search_api",
id="00000000-0000-4000-8000-000000000018",
id="00000000-0000-4000-8000-000000000014",
),
AIMessage(
content="",
@@ -2281,20 +2275,12 @@ def test_message_graph(
FunctionMessage(
content="result for another",
name="search_api",
id="00000000-0000-4000-8000-000000000031",
id="00000000-0000-4000-8000-000000000026",
),
AIMessage(content="answer", id="ai3"),
]
assert [*app.stream([HumanMessage(content="what is weather in sf")])] == [
{
"__start__": [
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000042",
)
]
},
{
"agent": AIMessage(
content="",
@@ -2308,7 +2294,7 @@ def test_message_graph(
"action": FunctionMessage(
content="result for query",
name="search_api",
id="00000000-0000-4000-8000-000000000055",
id="00000000-0000-4000-8000-000000000046",
)
},
{
@@ -2324,7 +2310,7 @@ def test_message_graph(
"action": FunctionMessage(
content="result for another",
name="search_api",
id="00000000-0000-4000-8000-000000000068",
id="00000000-0000-4000-8000-000000000058",
)
},
{"agent": AIMessage(content="answer", id="ai3")},
@@ -2341,7 +2327,6 @@ def test_message_graph(
HumanMessage(content="what is weather in sf"), config
)
] == [
{"__start__": HumanMessage(content="what is weather in sf", id=AnyStr())},
{
"agent": AIMessage(
content="",
@@ -2489,12 +2474,6 @@ def test_message_graph(
HumanMessage(content="what is weather in sf"), config
)
] == [
{
"__start__": HumanMessage(
content="what is weather in sf",
id=AnyStr(),
)
},
{
"agent": AIMessage(
content="",
@@ -2681,7 +2660,6 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
}
assert [*app.stream({"query": "what is weather in sf"})] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"retriever_two": {"docs": ["doc3", "doc4"]},
@@ -2788,7 +2766,6 @@ def test_in_one_fan_out_state_graph_waiting_edge() -> None:
}
assert [*app.stream({"query": "what is weather in sf"})] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
@@ -2806,7 +2783,6 @@ def test_in_one_fan_out_state_graph_waiting_edge() -> None:
assert [
c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)
] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
@@ -2912,7 +2888,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch() -> None:
}
assert [*app.stream({"query": "what is weather in sf"})] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
@@ -2930,7 +2905,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch() -> None:
assert [
c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)
] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
@@ -3002,7 +2976,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None:
}
assert [*app.stream({"query": "what is weather in sf"})] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
@@ -3021,7 +2994,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None:
assert [
c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)
] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
@@ -3101,7 +3073,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
}
assert [*app.stream({"query": "what is weather in sf"})] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
tests/test_pregel_async.py
-21
View File
@@ -1455,7 +1455,6 @@ async def test_conditional_graph_state() -> None:
}
assert [c async for c in app.astream({"input": "what is weather in sf"})] == [
{"__start__": {"input": "what is weather in sf"}},
{
"agent": {
"agent_outcome": AgentAction(
@@ -1522,7 +1521,6 @@ async def test_conditional_graph_state() -> None:
{"input": "what is weather in sf"}, config
)
] == [
{"__start__": {"input": "what is weather in sf"}},
{
"agent": {
"agent_outcome": AgentAction(
@@ -1643,7 +1641,6 @@ async def test_conditional_graph_state() -> None:
{"input": "what is weather in sf"}, config
)
] == [
{"__start__": {"input": "what is weather in sf"}},
{
"agent": {
"agent_outcome": AgentAction(
@@ -1824,7 +1821,6 @@ async def test_conditional_entrypoint_graph_state() -> None:
}
assert [c async for c in app.astream({"input": "what is weather in sf"})] == [
{"__start__": {"input": "what is weather in sf"}},
{"right": {"output": "what is weather in sf->right"}},
]
@@ -1948,7 +1944,6 @@ async def test_prebuilt_tool_chat() -> None:
{"messages": [HumanMessage(content="what is weather in sf")]}
)
] == [
{"__start__": {"messages": [HumanMessage(content="what is weather in sf")]}},
{
"agent": {
"messages": [
@@ -2094,7 +2089,6 @@ async def test_prebuilt_chat() -> None:
{"messages": [HumanMessage(content="what is weather in sf")]}
)
] == [
{"__start__": {"messages": [HumanMessage(content="what is weather in sf")]}},
{
"agent": {
"messages": [
@@ -2280,7 +2274,6 @@ async def test_message_graph() -> None:
assert [
c async for c in app.astream([HumanMessage(content="what is weather in sf")])
] == [
{"__start__": [HumanMessage(content="what is weather in sf", id=AnyStr())]},
{
"agent": AIMessage(
content="",
@@ -2323,12 +2316,6 @@ async def test_message_graph() -> None:
HumanMessage(content="what is weather in sf"), config
)
] == [
{
"__start__": HumanMessage(
content="what is weather in sf",
id=AnyStr(),
)
},
{
"agent": AIMessage(
content="",
@@ -2515,7 +2502,6 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
}
assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"retriever_two": {"docs": ["doc3", "doc4"]},
@@ -2579,7 +2565,6 @@ async def test_in_one_fan_out_state_graph_waiting_edge() -> None:
}
assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
@@ -2600,7 +2585,6 @@ async def test_in_one_fan_out_state_graph_waiting_edge() -> None:
{"query": "what is weather in sf"}, config
)
] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
@@ -2706,7 +2690,6 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch() -> None:
}
assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
@@ -2727,7 +2710,6 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch() -> None:
{"query": "what is weather in sf"}, config
)
] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
@@ -2799,7 +2781,6 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None:
}
assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
@@ -2821,7 +2802,6 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None:
{"query": "what is weather in sf"}, config
)
] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
@@ -2901,7 +2881,6 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
}
assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
{"__start__": {"query": "what is weather in sf"}},
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},