From b6848c5fdfb637192e1e5b487373118cdefb5d4f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 1 Apr 2024 13:43:23 -0700 Subject: [PATCH 1/4] Optimize tracing run tree - Do not run node if it's a passthrough - Do not run writers that wouldn't affect any channels - Combine consecutive writers when it doesn't change semantics --- examples/branching.ipynb | 48 ++++++++---------------- langgraph/graph/state.py | 4 +- langgraph/pregel/__init__.py | 71 ++++++++++++++++-------------------- langgraph/pregel/read.py | 41 +++++++++++++++++++++ langgraph/pregel/write.py | 26 ++++++++----- tests/test_pregel.py | 12 +++--- 6 files changed, 114 insertions(+), 88 deletions(-) diff --git a/examples/branching.ipynb b/examples/branching.ipynb index 3818d36cb..92943f7ad 100644 --- a/examples/branching.ipynb +++ b/examples/branching.ipynb @@ -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,29 @@ "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", + "*clears throat and stands up straight, eyes shining with excitement* Pineapples on pizza?! Oh my goodness, where do I even begin?! This is quite possibly the most revolutionary, delectable, and downright magnificent food combination of all time! \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", + "Pineapple's sweet, tangy, and juicy essence is the perfect complement to the savory, cheesy goodness of pizza. The interplay of flavors is simply divine - the pineapple's brightness cuts through the richness of the cheese, while the baked crust provides the perfect textural contrast. It's a symphony for the taste buds!\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 let's not forget the sheer versatility of this masterpiece. Pineapple can be paired with all sorts of toppings - ham, bacon, jalapeños, you name it! It truly is the Swiss Army knife of pizza toppings. Whether you're in the mood for a classic Hawaiian or something more adventurous, pineapple on pizza never fails to deliver.\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", + "Naysayers may try to disparage this culinary work of art, but I say they're simply missing out on one of life's greatest pleasures! Pineapple pizza is a triumph of human ingenuity and creativity. It's a bold, flavor-packed statement that refuses to be confined by traditional pizza norms. \n", + "\n", + "So I say, embrace the pineapple pizza revolution with open arms! Savor every bite of that sweet, tangy, cheesy delight. You'll be wondering how you ever lived without it. Pineapple on pizza - the future of food is now, my friends!\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 choice, as both arguments make compelling points. However, I believe the argument in favor of pineapple on pizza is the stronger of the two.\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", + "Argument 0 presents a passionate and well-reasoned case for why pineapple is a delightful and versatile pizza topping. The points about the complementary flavors and textures are convincing, and the argument about pineapple's versatility to pair with various other toppings is a strong one.\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", + "In contrast, Argument 1 relies more on personal distaste and traditionalist views rather than substantive counterarguments. While the points about pineapple's clash with the typical pizza flavors and the potential for a soggy texture are valid, the overall tone is more dismissive than persuasive.\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" + "Ultimately, the first argument does a better job of making a positive case for pineapple pizza, while the second argument feels more like an emotional rejection of the concept without fully addressing the merits presented. The passion and creativity of Argument 0 gives it the edge in making a more compelling case.\n", + "\n", + "Of course, this is a subjective topic and reasonable people can disagree. But based on the strength of the arguments presented, I believe Argument 0 makes the stronger case in favor of pineapple on pizza.\n" ] } ], @@ -208,8 +212,6 @@ "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", @@ -217,24 +219,6 @@ " 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, @@ -260,7 +244,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.2" + "version": "3.11.8" } }, "nbformat": 4, diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 814bc5c46..91cd52e9d 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -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 diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 8dd8deac1..d85d968c1 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -150,7 +150,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() @@ -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) @@ -316,7 +316,7 @@ class Pregel( values = { k: _read_channel(channels, k, return_exception=True) for k in channels - if k in self.snapshot_channels_list + if k in self.stream_channels_list } values = { k: v for k, v in values.items() if not isinstance(v, EmptyChannelError) @@ -325,7 +325,7 @@ class Pregel( 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, ) @@ -343,7 +343,7 @@ class Pregel( values = { k: _read_channel(channels, k, return_exception=True) for k in channels - if k in self.snapshot_channels_list + if k in self.stream_channels_list } values = { k: v for k, v in values.items() if not isinstance(v, EmptyChannelError) @@ -352,7 +352,7 @@ class Pregel( 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, ) @@ -368,7 +368,7 @@ class Pregel( values = { k: _read_channel(channels, k, return_exception=True) for k in channels - if k in self.snapshot_channels_list + if k in self.stream_channels_list } values = { k: v @@ -379,7 +379,7 @@ class Pregel( 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, ) @@ -398,7 +398,7 @@ class Pregel( values = { k: _read_channel(channels, k, return_exception=True) for k in channels - if k in self.snapshot_channels_list + if k in self.stream_channels_list } values = { k: v @@ -409,7 +409,7 @@ class Pregel( 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,10 +446,11 @@ 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], + RunnableSequence(*writers) if len(writers) > 1 else writers[0], values, as_node, deque(), @@ -502,10 +503,11 @@ 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], + RunnableSequence(*writers) if len(writers) > 1 else writers[0], values, as_node, deque(), @@ -616,7 +618,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 +649,7 @@ class Pregel( if _should_interrupt( checkpoint, interrupt_before_nodes, - self.snapshot_channels_list, + self.stream_channels_list, next_tasks, ): break @@ -726,7 +728,7 @@ class Pregel( if _should_interrupt( checkpoint, interrupt_after_nodes, - self.snapshot_channels_list, + self.stream_channels_list, next_tasks, ): break @@ -800,7 +802,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 +833,7 @@ class Pregel( if _should_interrupt( checkpoint, interrupt_before_nodes, - self.snapshot_channels_list, + self.stream_channels_list, next_tasks, ): break @@ -917,7 +919,7 @@ class Pregel( if _should_interrupt( checkpoint, interrupt_after_nodes, - self.snapshot_channels_list, + self.stream_channels_list, next_tasks, ): break @@ -1226,10 +1228,9 @@ def _apply_writes( channels[chan].update([]) -class PregelTask(NamedTuple): - proc: Runnable - input: Any +class PregelTaskDescription(NamedTuple): name: str + input: Any class PregelExecutableTask(NamedTuple): @@ -1245,7 +1246,7 @@ def _prepare_next_tasks( processes: Mapping[str, ChannelInvoke], channels: Mapping[str, BaseChannel], for_execution: Literal[False], -) -> tuple[Checkpoint, list[PregelTask]]: +) -> tuple[Checkpoint, list[PregelTaskDescription]]: ... @@ -1265,9 +1266,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(): @@ -1320,18 +1321,10 @@ 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(node, val, name, deque())) else: - tasks.append(PregelTask(proc, val, name)) + tasks.append(PregelTaskDescription(name, val)) return checkpoint, tasks diff --git a/langgraph/pregel/read.py b/langgraph/pregel/read.py index 9ea9e46e5..b7ce6e4af 100644 --- a/langgraph/pregel/read.py +++ b/langgraph/pregel/read.py @@ -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, *, diff --git a/langgraph/pregel/write.py b/langgraph/pregel/write.py index 03e95fb28..85fc46cfd 100644 --- a/langgraph/pregel/write.py +++ b/langgraph/pregel/write.py @@ -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 ] diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 0f5810b1a..ee1e6e913 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -2257,7 +2257,7 @@ def test_message_graph( assert app.invoke(HumanMessage(content="what is weather in sf")) == [ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000005", # adds missing ids + id="00000000-0000-4000-8000-000000000002", # adds missing ids ), AIMessage( content="", @@ -2269,7 +2269,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,7 +2281,7 @@ 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"), ] @@ -2291,7 +2291,7 @@ def test_message_graph( "__start__": [ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000042", + id="00000000-0000-4000-8000-000000000034", ) ] }, @@ -2308,7 +2308,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 +2324,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")}, From 3e4c94289ac19032e843497980b67a1888d40561 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 1 Apr 2024 15:08:35 -0700 Subject: [PATCH 2/4] Manage run_manager directly - allows us to set run output different from streamed values: here the run output should always be final value of all channels - removes 1-2 more frames from stack traces --- examples/branching.ipynb | 34 ++-- langgraph/pregel/__init__.py | 308 +++++++++++------------------------ langgraph/pregel/io.py | 36 +++- 3 files changed, 139 insertions(+), 239 deletions(-) diff --git a/examples/branching.ipynb b/examples/branching.ipynb index 92943f7ad..c9d06b895 100644 --- a/examples/branching.ipynb +++ b/examples/branching.ipynb @@ -186,25 +186,31 @@ "Pineapples on pizza\n", "## source:\n", "## branch_1:\n", - "*clears throat and stands up straight, eyes shining with excitement* Pineapples on pizza?! Oh my goodness, where do I even begin?! This is quite possibly the most revolutionary, delectable, and downright magnificent food combination of all time! \n", + "*eyes light up with excitement* Ooh, pineapples on pizza - now we're talking! This is one of my all-time favorite topics to discuss. Let me tell you, pineapple pizza is the ultimate culinary masterpiece, a harmonious fusion of sweet and savory that tantalizes the taste buds in the most delightful way.\n", "\n", - "Pineapple's sweet, tangy, and juicy essence is the perfect complement to the savory, cheesy goodness of pizza. The interplay of flavors is simply divine - the pineapple's brightness cuts through the richness of the cheese, while the baked crust provides the perfect textural contrast. It's a symphony for the taste buds!\n", + "First off, the combination of the juicy, tropical pineapple and the rich, melty cheese is simply divine. The tart sweetness of the pineapple perfectly complements the creamy, salty goodness of the mozzarella. It's a flavor explosion in your mouth!\n", "\n", - "And let's not forget the sheer versatility of this masterpiece. Pineapple can be paired with all sorts of toppings - ham, bacon, jalapeños, you name it! It truly is the Swiss Army knife of pizza toppings. Whether you're in the mood for a classic Hawaiian or something more adventurous, pineapple on pizza never fails to deliver.\n", + "And let's not forget the textural contrast - the soft, tender pineapple chunks paired with the crisp, chewy pizza crust? Perfection! It's like a party in your mouth. \n", "\n", - "Naysayers may try to disparage this culinary work of art, but I say they're simply missing out on one of life's greatest pleasures! Pineapple pizza is a triumph of human ingenuity and creativity. It's a bold, flavor-packed statement that refuses to be confined by traditional pizza norms. \n", + "Plus, pineapple is loaded with nutrients like vitamin C, manganese, and fiber, so you can feel good about indulging in this delectable treat. It's the ultimate in guilt-free indulgence.\n", "\n", - "So I say, embrace the pineapple pizza revolution with open arms! Savor every bite of that sweet, tangy, cheesy delight. You'll be wondering how you ever lived without it. Pineapple on pizza - the future of food is now, my friends!\n", + "Trust me, once you try pineapple pizza, you'll be hooked. It's the perfect balance of flavors and textures that will have you coming back for more, again and again. Pineapple pizza is a true culinary masterpiece, and anyone who says otherwise simply hasn't experienced its greatness. It's a must-try for any true pizza lover!\n", "## sink:\n", - "This is a tough choice, as both arguments make compelling points. However, I believe the argument in favor of pineapple on pizza is the stronger of the two.\n", + "This is a tough one, as both arguments make compelling points. However, I would argue that the first argument in favor of pineapple on pizza is the stronger and more persuasive case.\n", "\n", - "Argument 0 presents a passionate and well-reasoned case for why pineapple is a delightful and versatile pizza topping. The points about the complementary flavors and textures are convincing, and the argument about pineapple's versatility to pair with various other toppings is a strong one.\n", + "The key strengths of the pro-pineapple argument are:\n", "\n", - "In contrast, Argument 1 relies more on personal distaste and traditionalist views rather than substantive counterarguments. While the points about pineapple's clash with the typical pizza flavors and the potential for a soggy texture are valid, the overall tone is more dismissive than persuasive.\n", + "1. It eloquently describes the delightful flavor and texture pairing of the sweet pineapple and savory cheese/crust. The \"flavor explosion\" and \"party in your mouth\" descriptions are very evocative.\n", "\n", - "Ultimately, the first argument does a better job of making a positive case for pineapple pizza, while the second argument feels more like an emotional rejection of the concept without fully addressing the merits presented. The passion and creativity of Argument 0 gives it the edge in making a more compelling case.\n", + "2. It highlights the nutritional benefits of pineapple, making it a more guilt-free indulgence.\n", "\n", - "Of course, this is a subjective topic and reasonable people can disagree. But based on the strength of the arguments presented, I believe Argument 0 makes the stronger case in favor of pineapple on pizza.\n" + "3. The argument is passionate and enthusiastic, painting pineapple pizza as a \"culinary masterpiece\" that any true pizza lover must experience.\n", + "\n", + "In contrast, the anti-pineapple argument, while forceful in its language, relies more on subjective distaste rather than objective reasoning. Calling it an \"abomination\" and \"affront to good taste\" is emotive rather than persuasive.\n", + "\n", + "The texture and flavor arguments, while valid, don't fully counter the pro-pineapple case. Reasonable people can disagree on whether the combination works.\n", + "\n", + "Overall, the first argument paints a more compelling picture of pineapple pizza's merits, making a stronger case for its inclusion as a legitimate pizza topping. The enthusiasm and vivid descriptions give it the edge over the more indignant and dismissive anti-pineapple stance.\n" ] } ], @@ -218,14 +224,6 @@ " else:\n", " print(message.content)" ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "9c9f3df1-70da-4fb9-b4bb-f580b8d55fa3", - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": { diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index d85d968c1..94a3a818b 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -21,11 +21,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,6 +32,9 @@ 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, patch_config, ) @@ -68,7 +68,13 @@ 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.validate import validate_graph, validate_keys @@ -313,14 +319,7 @@ 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.stream_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) @@ -340,14 +339,7 @@ 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.stream_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) @@ -365,16 +357,7 @@ 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.stream_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) @@ -395,16 +378,7 @@ 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.stream_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) @@ -573,13 +547,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") @@ -591,7 +575,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 @@ -605,9 +596,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 @@ -733,6 +722,9 @@ class Pregel( ): 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 @@ -740,6 +732,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: @@ -748,25 +743,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, @@ -775,7 +780,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 @@ -789,9 +801,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 @@ -924,6 +934,9 @@ class Pregel( ): 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 @@ -931,6 +944,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: @@ -954,9 +970,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, @@ -969,58 +984,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], @@ -1033,15 +996,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, @@ -1054,61 +1013,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]]], @@ -1163,33 +1067,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( @@ -1278,7 +1158,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 @@ -1286,9 +1166,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: @@ -1296,7 +1174,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 diff --git a/langgraph/pregel/io.py b/langgraph/pregel/io.py index 3046db221..ddd9f8160 100644 --- a/langgraph/pregel/io.py +++ b/langgraph/pregel/io.py @@ -7,18 +7,42 @@ from langgraph.channels.base import BaseChannel, EmptyChannelError from langgraph.pregel.log import logger -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,10 +70,10 @@ 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 From 2736d729896ec668dafc000457bae1e51f7f279e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 1 Apr 2024 16:05:12 -0700 Subject: [PATCH 3/4] Remove __start__ node from stream output --- examples/branching.ipynb | 26 ++++++++--------- langgraph/pregel/__init__.py | 56 +++++++++++++----------------------- langgraph/pregel/debug.py | 9 ++---- langgraph/pregel/io.py | 15 ++++++---- langgraph/pregel/types.py | 28 ++++++++++++++++++ tests/test_pregel.py | 29 ------------------- tests/test_pregel_async.py | 21 -------------- 7 files changed, 72 insertions(+), 112 deletions(-) create mode 100644 langgraph/pregel/types.py diff --git a/examples/branching.ipynb b/examples/branching.ipynb index c9d06b895..dea1edf4b 100644 --- a/examples/branching.ipynb +++ b/examples/branching.ipynb @@ -186,31 +186,29 @@ "Pineapples on pizza\n", "## source:\n", "## branch_1:\n", - "*eyes light up with excitement* Ooh, pineapples on pizza - now we're talking! This is one of my all-time favorite topics to discuss. Let me tell you, pineapple pizza is the ultimate culinary masterpiece, a harmonious fusion of sweet and savory that tantalizes the taste buds in the most delightful way.\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", - "First off, the combination of the juicy, tropical pineapple and the rich, melty cheese is simply divine. The tart sweetness of the pineapple perfectly complements the creamy, salty goodness of the mozzarella. It's a flavor explosion in your mouth!\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 textural contrast - the soft, tender pineapple chunks paired with the crisp, chewy pizza crust? Perfection! It's like a party in your mouth. \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", - "Plus, pineapple is loaded with nutrients like vitamin C, manganese, and fiber, so you can feel good about indulging in this delectable treat. It's the ultimate in guilt-free indulgence.\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", - "Trust me, once you try pineapple pizza, you'll be hooked. It's the perfect balance of flavors and textures that will have you coming back for more, again and again. Pineapple pizza is a true culinary masterpiece, and anyone who says otherwise simply hasn't experienced its greatness. It's a must-try for any true pizza lover!\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 argue that the first argument in favor of pineapple on pizza is the stronger and more persuasive case.\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 key strengths of the pro-pineapple argument are:\n", + "The main points against pineapple pizza seem to be:\n", "\n", - "1. It eloquently describes the delightful flavor and texture pairing of the sweet pineapple and savory cheese/crust. The \"flavor explosion\" and \"party in your mouth\" descriptions are very evocative.\n", + "1. It goes against the traditional, authentic Italian pizza experience and disrespects the sanctity of this classic dish.\n", "\n", - "2. It highlights the nutritional benefits of pineapple, making it a more guilt-free indulgence.\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. The argument is passionate and enthusiastic, painting pineapple pizza as a \"culinary masterpiece\" that any true pizza lover must experience.\n", + "3. Pineapple is a sweet, tropical fruit that belongs in desserts and drinks, not as a pizza topping.\n", "\n", - "In contrast, the anti-pineapple argument, while forceful in its language, relies more on subjective distaste rather than objective reasoning. Calling it an \"abomination\" and \"affront to good taste\" is emotive rather than persuasive.\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", - "The texture and flavor arguments, while valid, don't fully counter the pro-pineapple case. Reasonable people can disagree on whether the combination works.\n", - "\n", - "Overall, the first argument paints a more compelling picture of pineapple pizza's merits, making a stronger case for its inclusion as a legitimate pizza topping. The enthusiasm and vivid descriptions give it the edge over the more indignant and dismissive anti-pineapple stance.\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" ] } ], diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 94a3a818b..d2aee38ac 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -12,7 +12,6 @@ from typing import ( Iterator, Literal, Mapping, - NamedTuple, Optional, Sequence, Type, @@ -36,6 +35,7 @@ from langchain_core.runnables.config import ( 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 ( @@ -77,6 +77,11 @@ from langgraph.pregel.io import ( ) 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 @@ -167,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]] ): @@ -424,9 +418,9 @@ class Pregel( if not writers: raise InvalidUpdateError(f"Node {as_node} has no writers") task = PregelExecutableTask( - RunnableSequence(*writers) if len(writers) > 1 else writers[0], - values, as_node, + values, + RunnableSequence(*writers) if len(writers) > 1 else writers[0], deque(), ) # execute task @@ -481,9 +475,9 @@ class Pregel( if not writers: raise InvalidUpdateError(f"Node {as_node} has no writers") task = PregelExecutableTask( - RunnableSequence(*writers) if len(writers) > 1 else writers[0], - values, as_node, + values, + RunnableSequence(*writers) if len(writers) > 1 else writers[0], deque(), ) # execute task @@ -652,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={ @@ -664,7 +658,7 @@ class Pregel( }, ), ) - for proc, input, name, writes in next_tasks + for name, input, proc, writes, proc_config in next_tasks ] futures = [ @@ -685,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 @@ -857,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={ @@ -869,7 +863,7 @@ class Pregel( }, ), ) - for proc, input, name, writes in next_tasks + for name, input, proc, writes, proc_config in next_tasks ] futures = ( @@ -897,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 @@ -1052,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) ) @@ -1108,18 +1102,6 @@ def _apply_writes( channels[chan].update([]) -class PregelTaskDescription(NamedTuple): - name: str - input: Any - - -class PregelExecutableTask(NamedTuple): - proc: Runnable - input: Any - name: str - writes: deque[tuple[str, Any]] - - @overload def _prepare_next_tasks( checkpoint: Checkpoint, @@ -1200,7 +1182,9 @@ def _prepare_next_tasks( if for_execution: if node := proc.get_node(): - tasks.append(PregelExecutableTask(node, val, name, deque())) + tasks.append( + PregelExecutableTask(name, val, node, deque(), proc.config) + ) else: tasks.append(PregelTaskDescription(name, val)) return checkpoint, tasks diff --git a/langgraph/pregel/debug.py b/langgraph/pregel/debug.py index f15fc3ccb..64a692e52 100644 --- a/langgraph/pregel/debug.py +++ b/langgraph/pregel/debug.py @@ -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) ) diff --git a/langgraph/pregel/io.py b/langgraph/pregel/io.py index ddd9f8160..a34b28f95 100644 --- a/langgraph/pregel/io.py +++ b/langgraph/pregel/io.py @@ -1,10 +1,8 @@ -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( @@ -79,13 +77,18 @@ def map_output_values( 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 }: @@ -93,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 diff --git a/langgraph/pregel/types.py b/langgraph/pregel/types.py new file mode 100644 index 000000000..3dcad2b79 --- /dev/null +++ b/langgraph/pregel/types.py @@ -0,0 +1,28 @@ +from collections import deque +from typing import Any, NamedTuple, Optional + +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: 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""" diff --git a/tests/test_pregel.py b/tests/test_pregel.py index ee1e6e913..f9752d5a4 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -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": [ @@ -2287,14 +2281,6 @@ def test_message_graph( ] assert [*app.stream([HumanMessage(content="what is weather in sf")])] == [ - { - "__start__": [ - HumanMessage( - content="what is weather in sf", - id="00000000-0000-4000-8000-000000000034", - ) - ] - }, { "agent": AIMessage( content="", @@ -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"}, diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 8ae1bf733..0ec6df93f 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -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"}, From 008301cf020af1860008f8a8425b5f6ed22d54b0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 1 Apr 2024 16:15:30 -0700 Subject: [PATCH 4/4] Lint --- langgraph/pregel/types.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/langgraph/pregel/types.py b/langgraph/pregel/types.py index 3dcad2b79..8ebe52e99 100644 --- a/langgraph/pregel/types.py +++ b/langgraph/pregel/types.py @@ -1,5 +1,5 @@ from collections import deque -from typing import Any, NamedTuple, Optional +from typing import Any, NamedTuple, Optional, Union from langchain_core.runnables import Runnable, RunnableConfig @@ -18,7 +18,7 @@ class PregelExecutableTask(NamedTuple): class StateSnapshot(NamedTuple): - values: dict[str, Any] | Any + values: Union[dict[str, Any], Any] """Current values of channels""" next: tuple[str] """Nodes to execute in the next step, if any"""