diff --git a/.github/workflows/run_notebooks.yml b/.github/workflows/run_notebooks.yml index 80d2aca44..16abf71d7 100644 --- a/.github/workflows/run_notebooks.yml +++ b/.github/workflows/run_notebooks.yml @@ -57,13 +57,13 @@ jobs: env: # these won't actually be used because of the VCR cassettes # but need to set them to avoid triggering getpass() - OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} - ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} - TAVILY_API_KEY: ${{ secrets.TAVILY_API_KEY }} - LANGSMITH_API_KEY: ${{ secrets.LANGSMITH_API_KEY }} - NOMIC_API_KEY: ${{ secrets.NOMIC_API_KEY }} - COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }} - FIREWORKS_API_KEY: ${{ secrets.FIREWORKS_API_KEY }} + OPENAI_API_KEY: "very-secret-key" + ANTHROPIC_API_KEY: "very-secret-key" + TAVILY_API_KEY: "very-secret-key" + LANGSMITH_API_KEY: "very-secret-key" + NOMIC_API_KEY: "very-secret-key" + COHERE_API_KEY: "very-secret-key" + FIREWORKS_API_KEY: "very-secret-key" run: | if [ "${{ github.event_name }}" = "workflow_dispatch" ] || [ "${{ github.event_name }}" = "schedule" ]; then echo "Running all notebooks" diff --git a/docs/docs/how-tos/command.ipynb b/docs/docs/how-tos/command.ipynb index a1f24f27b..cf0b01ae9 100644 --- a/docs/docs/how-tos/command.ipynb +++ b/docs/docs/how-tos/command.ipynb @@ -33,7 +33,7 @@ " )\n", "```\n", "\n", - "If you are using [subgraphs](#subgraphs), you might want to navigate from a node a subgraph to a different subgraph (i.e. a different node in the parent graph). To do so, you can specify `graph=Command.PARENT` in `Command`:\n", + "If you are using [subgraphs](#subgraphs), you might want to navigate from a node within a subgraph to a different subgraph (i.e. a different node in the parent graph). To do so, you can specify `graph=Command.PARENT` in `Command`:\n", "\n", "```python\n", "def my_node(state: State) -> Command[Literal[\"my_other_node\"]]:\n", diff --git a/docs/docs/how-tos/state-model.ipynb b/docs/docs/how-tos/state-model.ipynb index 221a4738d..c4fb12041 100644 --- a/docs/docs/how-tos/state-model.ipynb +++ b/docs/docs/how-tos/state-model.ipynb @@ -266,6 +266,235 @@ " print(\"An exception was raised because bad_node sets `a` to an integer.\")\n", " print(e)" ] + }, + { + "cell_type": "markdown", + "id": "2270bc3c", + "metadata": {}, + "source": [ + "## Multiple Nodes\n", + "\n", + "Run-time validation will also work in a multi-node graph. In the example below `bad_node` updates `a` to an integer. \n", + "\n", + "Because run-time validation occurs on **inputs**, the validation error will occur when `ok_node` is called (not when `bad_node` returns an update to the state which is inconsistent with the schema)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d832cdcc", + "metadata": {}, + "outputs": [], + "source": [ + "from langgraph.graph import StateGraph, START, END\n", + "from typing_extensions import TypedDict\n", + "\n", + "from pydantic import BaseModel\n", + "\n", + "\n", + "# The overall state of the graph (this is the public state shared across nodes)\n", + "class OverallState(BaseModel):\n", + " a: str\n", + "\n", + "\n", + "def bad_node(state: OverallState):\n", + " return {\n", + " \"a\": 123 # Invalid\n", + " }\n", + "\n", + "\n", + "def ok_node(state: OverallState):\n", + " return {\"a\": \"goodbye\"}\n", + "\n", + "\n", + "# Build the state graph\n", + "builder = StateGraph(OverallState)\n", + "builder.add_node(bad_node)\n", + "builder.add_node(ok_node)\n", + "builder.add_edge(START, \"bad_node\")\n", + "builder.add_edge(\"bad_node\", \"ok_node\")\n", + "builder.add_edge(\"ok_node\", END)\n", + "graph = builder.compile()\n", + "\n", + "# Test the graph with a valid input\n", + "try:\n", + " graph.invoke({\"a\": \"hello\"})\n", + "except Exception as e:\n", + " print(\"An exception was raised because bad_node sets `a` to an integer.\")\n", + " print(e)" + ] + }, + { + "cell_type": "markdown", + "id": "456b1f77", + "metadata": {}, + "source": [ + "## Advanced Pydantic Model Usage\n", + "\n", + "This section covers more advanced topics when using Pydantic models with LangGraph.\n", + "\n", + "### Serialization Behavior\n", + "\n", + "When using Pydantic models as state schemas, it's important to understand how serialization works, especially when:\n", + "- Passing Pydantic objects as inputs\n", + "- Receiving outputs from the graph\n", + "- Working with nested Pydantic models\n", + "\n", + "Let's see these behaviors in action:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0e919cdc", + "metadata": {}, + "outputs": [], + "source": [ + "from langgraph.graph import StateGraph, START, END\n", + "from pydantic import BaseModel\n", + "\n", + "\n", + "class NestedModel(BaseModel):\n", + " value: str\n", + "\n", + "\n", + "class ComplexState(BaseModel):\n", + " text: str\n", + " count: int\n", + " nested: NestedModel\n", + "\n", + "\n", + "def process_node(state: ComplexState):\n", + " # Node receives a validated Pydantic object\n", + " print(f\"Input state type: {type(state)}\")\n", + " print(f\"Nested type: {type(state.nested)}\")\n", + "\n", + " # Return a dictionary update\n", + " return {\"text\": state.text + \" processed\", \"count\": state.count + 1}\n", + "\n", + "\n", + "# Build the graph\n", + "builder = StateGraph(ComplexState)\n", + "builder.add_node(\"process\", process_node)\n", + "builder.add_edge(START, \"process\")\n", + "builder.add_edge(\"process\", END)\n", + "graph = builder.compile()\n", + "\n", + "# Create a Pydantic instance for input\n", + "input_state = ComplexState(text=\"hello\", count=0, nested=NestedModel(value=\"test\"))\n", + "print(f\"Input object type: {type(input_state)}\")\n", + "\n", + "# Invoke graph with a Pydantic instance\n", + "result = graph.invoke(input_state)\n", + "print(f\"Output type: {type(result)}\")\n", + "print(f\"Output content: {result}\")\n", + "\n", + "# Convert back to Pydantic model if needed\n", + "output_model = ComplexState(**result)\n", + "print(f\"Converted back to Pydantic: {type(output_model)}\")" + ] + }, + { + "cell_type": "markdown", + "id": "f13f28ce", + "metadata": {}, + "source": [ + "### Runtime Type Coercion\n", + "\n", + "Pydantic performs runtime type coercion for certain data types. This can be helpful but also lead to unexpected behavior if you're not aware of it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "faf59316", + "metadata": {}, + "outputs": [], + "source": [ + "from langgraph.graph import StateGraph, START, END\n", + "from pydantic import BaseModel\n", + "\n", + "\n", + "class CoercionExample(BaseModel):\n", + " # Pydantic will coerce string numbers to integers\n", + " number: int\n", + " # Pydantic will parse string booleans to bool\n", + " flag: bool\n", + "\n", + "\n", + "def inspect_node(state: CoercionExample):\n", + " print(f\"number: {state.number} (type: {type(state.number)})\")\n", + " print(f\"flag: {state.flag} (type: {type(state.flag)})\")\n", + " return {}\n", + "\n", + "\n", + "builder = StateGraph(CoercionExample)\n", + "builder.add_node(\"inspect\", inspect_node)\n", + "builder.add_edge(START, \"inspect\")\n", + "builder.add_edge(\"inspect\", END)\n", + "graph = builder.compile()\n", + "\n", + "# Demonstrate coercion with string inputs that will be converted\n", + "result = graph.invoke({\"number\": \"42\", \"flag\": \"true\"})\n", + "\n", + "# This would fail with a validation error\n", + "try:\n", + " graph.invoke({\"number\": \"not-a-number\", \"flag\": \"true\"})\n", + "except Exception as e:\n", + " print(f\"\\nExpected validation error: {e}\")" + ] + }, + { + "cell_type": "markdown", + "id": "2844475b", + "metadata": {}, + "source": [ + "### Working with Message Models\n", + "\n", + "When working with LangChain message types in your state schema, there are important considerations for serialization. You should use `AnyMessage` (rather than `BaseMessage`) for proper serialization/deserialization when using message objects over the wire:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bd0734b0", + "metadata": {}, + "outputs": [], + "source": [ + "from langgraph.graph import StateGraph, START, END\n", + "from pydantic import BaseModel\n", + "from langchain_core.messages import HumanMessage, AIMessage, BaseMessage\n", + "from typing import List\n", + "\n", + "\n", + "class ChatState(BaseModel):\n", + " messages: List[BaseMessage]\n", + " context: str\n", + "\n", + "\n", + "def add_message(state: ChatState):\n", + " return {\"messages\": state.messages + [AIMessage(content=\"Hello there!\")]}\n", + "\n", + "\n", + "builder = StateGraph(ChatState)\n", + "builder.add_node(\"add_message\", add_message)\n", + "builder.add_edge(START, \"add_message\")\n", + "builder.add_edge(\"add_message\", END)\n", + "graph = builder.compile()\n", + "\n", + "# Create input with a message\n", + "initial_state = ChatState(\n", + " messages=[HumanMessage(content=\"Hi\")], context=\"Customer support chat\"\n", + ")\n", + "\n", + "result = graph.invoke(initial_state)\n", + "print(f\"Output: {result}\")\n", + "\n", + "# Convert back to Pydantic model to see message types\n", + "output_model = ChatState(**result)\n", + "for i, msg in enumerate(output_model.messages):\n", + " print(f\"Message {i}: {type(msg).__name__} - {msg.content}\")" + ] } ], "metadata": { diff --git a/docs/docs/how-tos/update-state-from-tools.ipynb b/docs/docs/how-tos/update-state-from-tools.ipynb index 0a7311342..1aea00556 100644 --- a/docs/docs/how-tos/update-state-from-tools.ipynb +++ b/docs/docs/how-tos/update-state-from-tools.ipynb @@ -210,7 +210,7 @@ "id": "cbb06aea-6654-4245-91f8-af6e8f2b5377", "metadata": {}, "source": [ - "Let's now add personalization: we'll respond differently to the user based on the state values AFTER the state has been updated from the tool. To achieve this, let's define a function that will dynamically construct the system prompt based on the graph state. It will be called ever time the LLM is called and the function output will be passed to the LLM:" + "Let's now add personalization: we'll respond differently to the user based on the state values AFTER the state has been updated from the tool. To achieve this, let's define a function that will dynamically construct the system prompt based on the graph state. It will be called every time the LLM is called and the function output will be passed to the LLM:" ] }, { diff --git a/libs/langgraph/bench/__main__.py b/libs/langgraph/bench/__main__.py index 677e79440..da6b7b6f3 100644 --- a/libs/langgraph/bench/__main__.py +++ b/libs/langgraph/bench/__main__.py @@ -7,6 +7,7 @@ from uvloop import new_event_loop from bench.fanout_to_subgraph import fanout_to_subgraph, fanout_to_subgraph_sync from bench.react_agent import react_agent +from bench.sequential import create_sequential from bench.wide_state import wide_state from langgraph.checkpoint.memory import MemorySaver from langgraph.pregel import Pregel @@ -203,6 +204,36 @@ benchmarks = ( ] }, ), + ( + "sequential_20", + create_sequential(20).compile(), + create_sequential(20).compile(), + {"messages": []}, # Empty list of messages + ), + ( + "sequential_50", + create_sequential(50).compile(), + create_sequential(50).compile(), + {"messages": []}, # Empty list of messages + ), + # ( + # "sequential_100", + # create_sequential(100).compile(), + # create_sequential(100).compile(), + # {"messages": []}, # Empty list of messages + # ), + # ( + # "sequential_200", + # create_sequential(200).compile(), + # create_sequential(200).compile(), + # {"messages": []}, # Empty list of messages + # ), + # ( + # "sequential_1000", + # create_sequential(1000).compile(), + # create_sequential(1000).compile(), + # {"messages": []}, # Empty list of messages + # ), ) diff --git a/libs/langgraph/bench/sequential.py b/libs/langgraph/bench/sequential.py new file mode 100644 index 000000000..3ab92912f --- /dev/null +++ b/libs/langgraph/bench/sequential.py @@ -0,0 +1,48 @@ +"""Create a sequential no-op graph consisting of a few hundred nodes.""" + +from langgraph.graph import MessagesState, StateGraph +from langgraph.utils.runnable import RunnableCallable + + +def create_sequential(number_nodes) -> StateGraph: + """Create a sequential no-op graph consisting of a few hundred nodes.""" + builder = StateGraph(MessagesState) + + def noop(state: MessagesState) -> None: + """No-op function.""" + pass + + async def anoop(state: MessagesState) -> None: + """No-op function.""" + pass + + prev_node = "__start__" + + for i in range(number_nodes): + name = f"node_{i}" + builder.add_node(name, RunnableCallable(noop, anoop)) + builder.add_edge(prev_node, name) + prev_node = name + + builder.add_edge(prev_node, "__end__") + return builder + + +if __name__ == "__main__": + import asyncio + import time + + import uvloop + + graph = create_sequential(2000).compile() + input = {"messages": []} # Empty list of messages + config = {"recursion_limit": 20000000000} + + async def run(): + len([c async for c in graph.astream(input, config=config)]) + + uvloop.install() + start = time.time() + asyncio.run(run()) + end = time.time() + print(f"Time taken: {end - start:.4f} seconds") diff --git a/libs/langgraph/langgraph/channels/any_value.py b/libs/langgraph/langgraph/channels/any_value.py index e9dfb77d6..35452084f 100644 --- a/libs/langgraph/langgraph/channels/any_value.py +++ b/libs/langgraph/langgraph/channels/any_value.py @@ -1,8 +1,9 @@ -from typing import Generic, Optional, Sequence, Type +from typing import Any, Generic, Optional, Sequence, Type from typing_extensions import Self from langgraph.channels.base import BaseChannel, Value +from langgraph.constants import MISSING from langgraph.errors import EmptyChannelError @@ -12,6 +13,10 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]): __slots__ = ("typ", "value") + def __init__(self, typ: Any, key: str = "") -> None: + super().__init__(typ, key) + self.value = MISSING + def __eq__(self, value: object) -> bool: return isinstance(value, AnyValue) @@ -34,17 +39,19 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]): def update(self, values: Sequence[Value]) -> bool: if len(values) == 0: - try: - del self.value - return True - except AttributeError: + if self.value is MISSING: return False + else: + self.value = MISSING + return True self.value = values[-1] return True def get(self) -> Value: - try: - return self.value - except AttributeError: + if self.value is MISSING: raise EmptyChannelError() + return self.value + + def is_available(self) -> bool: + return self.value is not MISSING diff --git a/libs/langgraph/langgraph/channels/base.py b/libs/langgraph/langgraph/channels/base.py index 4aaeb5681..b9239be7a 100644 --- a/libs/langgraph/langgraph/channels/base.py +++ b/libs/langgraph/langgraph/channels/base.py @@ -64,6 +64,17 @@ class BaseChannel(Generic[Value, Update, C], ABC): """ return False + def is_available(self) -> bool: + """Return True if the channel is available (not empty), False otherwise. + Subclasses should override this method to provide a more efficient + implementation than calling get() and catching EmptyChannelError. + """ + try: + self.get() + return True + except EmptyChannelError: + return False + __all__ = [ "BaseChannel", diff --git a/libs/langgraph/langgraph/channels/binop.py b/libs/langgraph/langgraph/channels/binop.py index a2360142b..413e0b91a 100644 --- a/libs/langgraph/langgraph/channels/binop.py +++ b/libs/langgraph/langgraph/channels/binop.py @@ -10,6 +10,7 @@ from typing import ( from typing_extensions import NotRequired, Required, Self from langgraph.channels.base import BaseChannel, Value +from langgraph.constants import MISSING from langgraph.errors import EmptyChannelError @@ -51,7 +52,7 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]): try: self.value = typ() except Exception: - pass + self.value = MISSING def __eq__(self, value: object) -> bool: return isinstance(value, BinaryOperatorAggregate) and ( @@ -81,7 +82,7 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]): def update(self, values: Sequence[Value]) -> bool: if not values: return False - if not hasattr(self, "value"): + if self.value is MISSING: self.value = values[0] values = values[1:] for value in values: @@ -89,7 +90,9 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]): return True def get(self) -> Value: - try: - return self.value - except AttributeError: + if self.value is MISSING: raise EmptyChannelError() + return self.value + + def is_available(self) -> bool: + return self.value is not MISSING diff --git a/libs/langgraph/langgraph/channels/dynamic_barrier_value.py b/libs/langgraph/langgraph/channels/dynamic_barrier_value.py index f64191e86..155c65446 100644 --- a/libs/langgraph/langgraph/channels/dynamic_barrier_value.py +++ b/libs/langgraph/langgraph/channels/dynamic_barrier_value.py @@ -85,6 +85,9 @@ class DynamicBarrierValue( raise EmptyChannelError() return None + def is_available(self) -> bool: + return self.seen == self.names + def consume(self) -> bool: if self.seen == self.names: self.seen = set() diff --git a/libs/langgraph/langgraph/channels/ephemeral_value.py b/libs/langgraph/langgraph/channels/ephemeral_value.py index 537a8763c..29a9a698c 100644 --- a/libs/langgraph/langgraph/channels/ephemeral_value.py +++ b/libs/langgraph/langgraph/channels/ephemeral_value.py @@ -3,6 +3,7 @@ from typing import Any, Generic, Optional, Sequence, Type from typing_extensions import Self from langgraph.channels.base import BaseChannel, Value +from langgraph.constants import MISSING from langgraph.errors import EmptyChannelError, InvalidUpdateError @@ -14,6 +15,7 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]): def __init__(self, typ: Any, guard: bool = True) -> None: super().__init__(typ) self.guard = guard + self.value = MISSING def __eq__(self, value: object) -> bool: return isinstance(value, EphemeralValue) and value.guard == self.guard @@ -37,10 +39,10 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]): def update(self, values: Sequence[Value]) -> bool: if len(values) == 0: - try: - del self.value + if self.value is not MISSING: + self.value = MISSING return True - except AttributeError: + else: return False if len(values) != 1 and self.guard: raise InvalidUpdateError( @@ -51,7 +53,9 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]): return True def get(self) -> Value: - try: - return self.value - except AttributeError: + if self.value is MISSING: raise EmptyChannelError() + return self.value + + def is_available(self) -> bool: + return self.value is not MISSING diff --git a/libs/langgraph/langgraph/channels/last_value.py b/libs/langgraph/langgraph/channels/last_value.py index 5065f4fc5..61669d390 100644 --- a/libs/langgraph/langgraph/channels/last_value.py +++ b/libs/langgraph/langgraph/channels/last_value.py @@ -1,8 +1,9 @@ -from typing import Generic, Optional, Sequence, Type +from typing import Any, Generic, Optional, Sequence, Type from typing_extensions import Self from langgraph.channels.base import BaseChannel, Value +from langgraph.constants import MISSING from langgraph.errors import ( EmptyChannelError, ErrorCode, @@ -16,6 +17,10 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): __slots__ = ("value",) + def __init__(self, typ: Any, key: str = "") -> None: + super().__init__(typ, key) + self.value = MISSING + def __eq__(self, value: object) -> bool: return isinstance(value, LastValue) @@ -50,7 +55,9 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): return True def get(self) -> Value: - try: - return self.value - except AttributeError: + if self.value is MISSING: raise EmptyChannelError() + return self.value + + def is_available(self) -> bool: + return self.value is not MISSING diff --git a/libs/langgraph/langgraph/channels/named_barrier_value.py b/libs/langgraph/langgraph/channels/named_barrier_value.py index 4a1d990ca..553316e19 100644 --- a/libs/langgraph/langgraph/channels/named_barrier_value.py +++ b/libs/langgraph/langgraph/channels/named_barrier_value.py @@ -60,6 +60,9 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]): raise EmptyChannelError() return None + def is_available(self) -> bool: + return self.seen == self.names + def consume(self) -> bool: if self.seen == self.names: self.seen = set() diff --git a/libs/langgraph/langgraph/channels/topic.py b/libs/langgraph/langgraph/channels/topic.py index 0430343dc..91e7027f9 100644 --- a/libs/langgraph/langgraph/channels/topic.py +++ b/libs/langgraph/langgraph/channels/topic.py @@ -75,3 +75,6 @@ class Topic( return list(self.values) else: raise EmptyChannelError + + def is_available(self) -> bool: + return bool(self.values) diff --git a/libs/langgraph/langgraph/channels/untracked_value.py b/libs/langgraph/langgraph/channels/untracked_value.py index 9b1020710..f9168131e 100644 --- a/libs/langgraph/langgraph/channels/untracked_value.py +++ b/libs/langgraph/langgraph/channels/untracked_value.py @@ -3,6 +3,7 @@ from typing import Generic, Optional, Sequence, Type from typing_extensions import Self from langgraph.channels.base import BaseChannel, Value +from langgraph.constants import MISSING from langgraph.errors import EmptyChannelError, InvalidUpdateError @@ -14,6 +15,7 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]): def __init__(self, typ: Type[Value], guard: bool = True) -> None: super().__init__(typ) self.guard = guard + self.value = MISSING def __eq__(self, value: object) -> bool: return isinstance(value, UntrackedValue) and value.guard == self.guard @@ -48,7 +50,9 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]): return True def get(self) -> Value: - try: - return self.value - except AttributeError: + if self.value is MISSING: raise EmptyChannelError() + return self.value + + def is_available(self) -> bool: + return self.value is not MISSING diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 605750716..d9009d200 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -821,9 +821,11 @@ class CompiledStateGraph(CompiledGraph): input_values = {k: k for k in self.builder.schemas[input_schema]} is_single_input = len(input_values) == 1 and "__root__" in input_values + branch_channel = f"branch:to:{key}" self.channels[key] = EphemeralValue(Any, guard=False) + self.channels[branch_channel] = EphemeralValue(Any, guard=False) self.nodes[key] = PregelNode( - triggers=[], + triggers=[branch_channel], # read state keys and managed values channels=(list(input_values) if is_single_input else input_values), # coerce state dict to schema class (eg. pydantic model) @@ -878,7 +880,7 @@ class CompiledStateGraph(CompiledGraph): if filtered := [p for p in packets if p != END]: writes = [ ( - ChannelWriteEntry(f"branch:{start}:{name}:{p}", start) + ChannelWriteEntry(f"branch:to:{p}", start) if not isinstance(p, Send) else p ) @@ -914,11 +916,6 @@ class CompiledStateGraph(CompiledGraph): if branch.ends else [node for node in self.builder.nodes if node != branch.then] ) - for end in ends: - if end != END: - channel_name = f"branch:{start}:{name}:{end}" - self.channels[channel_name] = EphemeralValue(Any, guard=False) - self.nodes[end].triggers.append(channel_name) # attach then subscriber if branch.then and branch.then != END: diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 84de26c05..027f51ee7 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -1,3 +1,4 @@ +import binascii import itertools import sys from collections import defaultdict, deque @@ -18,7 +19,6 @@ from typing import ( cast, overload, ) -from uuid import UUID from langchain_core.callbacks import Callbacks from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager @@ -27,6 +27,7 @@ from langchain_core.runnables.config import RunnableConfig from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import ( BaseCheckpointSaver, + ChannelVersions, Checkpoint, PendingWrite, V, @@ -373,6 +374,8 @@ def prepare_next_tasks( """Prepare the set of tasks that will make up the next Pregel step. This is the union of all PUSH tasks (Sends) and PULL tasks (nodes triggered by edges).""" + checkpoint_id_bytes = binascii.unhexlify(checkpoint["id"].replace("-", "")) + null_version = checkpoint_null_version(checkpoint) tasks: list[Union[PregelTask, PregelExecutableTask]] = [] # Consume pending_sends from previous step for idx, _ in enumerate(checkpoint["pending_sends"]): @@ -380,6 +383,8 @@ def prepare_next_tasks( (PUSH, idx), None, checkpoint=checkpoint, + checkpoint_id_bytes=checkpoint_id_bytes, + checkpoint_null_version=null_version, pending_writes=pending_writes, processes=processes, channels=channels, @@ -399,6 +404,8 @@ def prepare_next_tasks( (PULL, name), None, checkpoint=checkpoint, + checkpoint_id_bytes=checkpoint_id_bytes, + checkpoint_null_version=null_version, pending_writes=pending_writes, processes=processes, channels=channels, @@ -414,11 +421,16 @@ def prepare_next_tasks( return {t.id: t for t in tasks} +PUSH_TRIGGER = (PUSH,) + + def prepare_single_task( task_path: tuple[Any, ...], task_id_checksum: Optional[str], *, checkpoint: Checkpoint, + checkpoint_id_bytes: bytes, + checkpoint_null_version: Optional[V], pending_writes: list[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], @@ -432,7 +444,6 @@ def prepare_single_task( ) -> Union[None, PregelTask, PregelExecutableTask]: """Prepares a single task for the next Pregel step, given a task path, which uniquely identifies a PUSH or PULL task within the graph.""" - checkpoint_id = UUID(checkpoint["id"]).bytes configurable = config.get(CONF, {}) parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "") @@ -445,10 +456,10 @@ def prepare_single_task( if name is None: raise ValueError("`call` functions must have a `__name__` attribute") # create task id - triggers = [PUSH] + triggers: Sequence[str] = PUSH_TRIGGER checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name task_id = _uuid5_str( - checkpoint_id, + checkpoint_id_bytes, checkpoint_ns, str(step), name, @@ -539,12 +550,12 @@ def prepare_single_task( ) return # create task id - triggers = [PUSH] + triggers = PUSH_TRIGGER checkpoint_ns = ( f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node ) task_id = _uuid5_str( - checkpoint_id, + checkpoint_id_bytes, checkpoint_ns, str(step), packet.node, @@ -641,20 +652,15 @@ def prepare_single_task( if name not in processes: return proc = processes[name] - version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) - null_version = version_type() # type: ignore[misc] - if null_version is None: + if checkpoint_null_version is None: return - seen = checkpoint["versions_seen"].get(name, {}) # If any of the channels read by this process were updated - if triggers := sorted( - chan - for chan in proc.triggers - if not isinstance( - read_channel(channels, chan, return_exception=True), EmptyChannelError - ) - and checkpoint["channel_versions"].get(chan, null_version) # type: ignore[operator] - > seen.get(chan, null_version) + if triggers := _triggers( + channels, + checkpoint["channel_versions"], + checkpoint["versions_seen"].get(name), + checkpoint_null_version, + proc, ): try: val = next( @@ -672,7 +678,7 @@ def prepare_single_task( # create task id checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name task_id = _uuid5_str( - checkpoint_id, + checkpoint_id_bytes, checkpoint_ns, str(step), name, @@ -763,6 +769,35 @@ def prepare_single_task( return PregelTask(task_id, name, task_path[:3]) +def checkpoint_null_version( + checkpoint: Checkpoint, +) -> Optional[V]: + """Get the null version for the checkpoint, if available.""" + for version in checkpoint["channel_versions"].values(): + return type(version)() + return None + + +def _triggers( + channels: Mapping[str, BaseChannel], + versions: ChannelVersions, + seen: Optional[ChannelVersions], + null_version: V, + proc: PregelNode, +) -> Sequence[str]: + if seen is None: + for chan in proc.triggers: + if channels[chan].is_available(): + return (chan,) + else: + for chan in proc.triggers: + if channels[chan].is_available() and versions.get( # type: ignore[operator] + chan, null_version + ) > seen.get(chan, null_version): + return (chan,) + return EMPTY_SEQ + + def _scratchpad( config: RunnableConfig, pending_writes: list[PendingWrite], diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index f07745f1b..1b9ae78d5 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -14,7 +14,6 @@ from langgraph.constants import ( NULL_TASK_ID, RESUME, RETURN, - SELF, START, TAG_HIDDEN, TASKS, @@ -38,14 +37,11 @@ def read_channel( 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: + except EmptyChannelError: + if catch: return None else: raise @@ -84,7 +80,7 @@ def map_command( if isinstance(send, Send): yield (NULL_TASK_ID, TASKS, send) elif isinstance(send, str): - yield (NULL_TASK_ID, f"branch:{START}:{SELF}:{send}", START) + yield (NULL_TASK_ID, f"branch:to:{send}", START) else: raise TypeError( f"In Command.goto, expected Send/str, got {type(send).__name__}" diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 7ab3431b7..edd69db01 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -1,4 +1,5 @@ import asyncio +import binascii import concurrent.futures from collections import defaultdict, deque from contextlib import AsyncExitStack, ExitStack @@ -79,6 +80,7 @@ from langgraph.pregel.algo import ( GetNextVersion, PregelTaskWrites, apply_writes, + checkpoint_null_version, increment, prepare_next_tasks, prepare_single_task, @@ -347,12 +349,16 @@ class PregelLoop(LoopProtocol): ): self.to_interrupt.append(task) return + checkpoint_id_bytes = binascii.unhexlify(self.checkpoint["id"].replace("-", "")) + null_version = checkpoint_null_version(self.checkpoint) if pushed := cast( Optional[PregelExecutableTask], prepare_single_task( (PUSH, task.path, write_idx, task.id, call), None, checkpoint=self.checkpoint, + checkpoint_id_bytes=checkpoint_id_bytes, + checkpoint_null_version=null_version, pending_writes=self.checkpoint_pending_writes, processes=self.nodes, channels=self.channels, diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 5efb89385..4c8cf6da4 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -149,7 +149,7 @@ class PregelExecutableTask(NamedTuple): proc: Runnable writes: deque[tuple[str, Any]] config: RunnableConfig - triggers: list[str] + triggers: Sequence[str] retry_policy: Optional[RetryPolicy] cache_policy: Optional[CachePolicy] id: str diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 7893875e8..e49391685 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph" -version = "0.3.14" +version = "0.3.15" description = "Building stateful, multi-actor applications with LLMs" authors = [] license = "MIT" diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index f049e4e81..d6637f000 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -2483,7 +2483,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 1, "langgraph_node": "agent", - "langgraph_triggers": ["start:agent"], + "langgraph_triggers": ("start:agent",), "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2500,7 +2500,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 2, "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], + "langgraph_triggers": ("branch:to:tools",), "langgraph_path": (PULL, "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2542,7 +2542,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 3, "langgraph_node": "agent", - "langgraph_triggers": ["tools"], + "langgraph_triggers": ("tools",), "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2559,7 +2559,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], + "langgraph_triggers": ("branch:to:tools",), "langgraph_path": (PULL, "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2573,7 +2573,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], + "langgraph_triggers": ("branch:to:tools",), "langgraph_path": (PULL, "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2585,7 +2585,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 5, "langgraph_node": "agent", - "langgraph_triggers": ["tools"], + "langgraph_triggers": ("tools",), "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -5501,7 +5501,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "rewrite_query", "input": {"query": "what is weather in sf", "docs": []}, - "triggers": ["start:rewrite_query"], + "triggers": ("start:rewrite_query",), }, }, ), @@ -5532,7 +5532,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "retriever_one", "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ["rewrite_query"], + "triggers": ("rewrite_query",), }, }, ), @@ -5546,7 +5546,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "retriever_two", "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ["rewrite_query"], + "triggers": ("rewrite_query",), }, }, ), @@ -5608,7 +5608,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "query": "query: what is weather in sf", "docs": ["doc1", "doc2", "doc3", "doc4"], }, - "triggers": ["retriever_one", "retriever_two"], + "triggers": (AnyStr("retriever_"),), }, }, ), @@ -6634,7 +6634,7 @@ def test_branch_then( "id": AnyStr(), "name": "prepare", "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], + "triggers": ("start:prepare",), }, }, { @@ -6706,7 +6706,7 @@ def test_branch_then( "id": AnyStr(), "name": "tool_two_slow", "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], + "triggers": ("branch:to:tool_two_slow",), }, }, { @@ -6773,7 +6773,7 @@ def test_branch_then( "id": AnyStr(), "name": "finish", "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition::then"], + "triggers": ("branch:prepare:condition::then",), }, }, { @@ -10378,9 +10378,7 @@ def test_weather_subgraph( "langgraph_node": "weather_graph", "langgraph_path": [PULL, "weather_graph"], "langgraph_step": 2, - "langgraph_triggers": [ - "branch:router_node:route_after_prediction:weather_graph" - ], + "langgraph_triggers": ["branch:to:weather_graph"], "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), @@ -10492,9 +10490,7 @@ def test_weather_subgraph( "langgraph_node": "weather_graph", "langgraph_path": [PULL, "weather_graph"], "langgraph_step": 2, - "langgraph_triggers": [ - "branch:router_node:route_after_prediction:weather_graph" - ], + "langgraph_triggers": ["branch:to:weather_graph"], "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index 250cbfe72..157e0e080 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -2300,7 +2300,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 1, "langgraph_node": "agent", - "langgraph_triggers": ["start:agent"], + "langgraph_triggers": ("start:agent",), "langgraph_path": ("__pregel_pull", "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2317,7 +2317,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 2, "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], + "langgraph_triggers": ("branch:to:tools",), "langgraph_path": ("__pregel_pull", "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2359,7 +2359,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 3, "langgraph_node": "agent", - "langgraph_triggers": ["tools"], + "langgraph_triggers": ("tools",), "langgraph_path": ("__pregel_pull", "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2376,7 +2376,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], + "langgraph_triggers": ("branch:to:tools",), "langgraph_path": ("__pregel_pull", "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2390,7 +2390,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], + "langgraph_triggers": ("branch:to:tools",), "langgraph_path": ("__pregel_pull", "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2402,7 +2402,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 5, "langgraph_node": "agent", - "langgraph_triggers": ["tools"], + "langgraph_triggers": ("tools",), "langgraph_path": ("__pregel_pull", "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -3883,7 +3883,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "rewrite_query", "input": {"query": "what is weather in sf", "docs": []}, - "triggers": ["start:rewrite_query"], + "triggers": ("start:rewrite_query",), }, }, ), @@ -3914,7 +3914,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "retriever_one", "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ["rewrite_query"], + "triggers": ("rewrite_query",), }, }, ), @@ -3928,7 +3928,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "retriever_two", "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ["rewrite_query"], + "triggers": ("rewrite_query",), }, }, ), @@ -3990,7 +3990,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "query": "query: what is weather in sf", "docs": ["doc1", "doc2", "doc3", "doc4"], }, - "triggers": ["retriever_one", "retriever_two"], + "triggers": (AnyStr("retriever_"),), }, }, ), @@ -4465,7 +4465,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "id": AnyStr(), "name": "prepare", "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], + "triggers": ("start:prepare",), }, }, { @@ -4537,7 +4537,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "id": AnyStr(), "name": "tool_two_slow", "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], + "triggers": ("branch:to:tool_two_slow",), }, }, { @@ -4609,7 +4609,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "id": AnyStr(), "name": "finish", "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition::then"], + "triggers": ("branch:prepare:condition::then",), }, }, { @@ -4778,7 +4778,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "id": AnyStr(), "name": "prepare", "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], + "triggers": ("start:prepare",), }, }, { @@ -7231,9 +7231,7 @@ async def test_weather_subgraph( "langgraph_node": "weather_graph", "langgraph_path": [PULL, "weather_graph"], "langgraph_step": 2, - "langgraph_triggers": [ - "branch:router_node:route_after_prediction:weather_graph" - ], + "langgraph_triggers": ["branch:to:weather_graph"], "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), @@ -7347,9 +7345,7 @@ async def test_weather_subgraph( "langgraph_node": "weather_graph", "langgraph_path": [PULL, "weather_graph"], "langgraph_step": 2, - "langgraph_triggers": [ - "branch:router_node:route_after_prediction:weather_graph" - ], + "langgraph_triggers": ["branch:to:weather_graph"], "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index ccd899ce9..d8d51a8a0 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -817,7 +817,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "id": AnyStr(), "name": "one", "input": 2, - "triggers": ["input"], + "triggers": ("input",), }, }, { @@ -828,7 +828,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "id": AnyStr(), "name": "two", "input": [12], - "triggers": ["inbox"], + "triggers": ("inbox",), }, }, { @@ -863,7 +863,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "id": AnyStr(), "name": "two", "input": [3], - "triggers": ["inbox"], + "triggers": ("inbox",), }, }, { @@ -3247,14 +3247,24 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( assert [ c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"qa": {"answer": ""}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"__interrupt__": ()}, - ] + ] in ( + [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"qa": {"answer": ""}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ], + [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"qa": {"answer": ""}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ], + ) assert [c for c in app_w_interrupt.stream(None, config)] == [ {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, @@ -5969,9 +5979,7 @@ def test_falsy_return_from_task( "a": 5, }, "name": "graph", - "triggers": [ - "__start__", - ], + "triggers": ("__start__",), }, "step": 0, "timestamp": AnyStr(), @@ -5985,9 +5993,7 @@ def test_falsy_return_from_task( {}, ), "name": "falsy_task", - "triggers": [ - "__pregel_push", - ], + "triggers": ("__pregel_push",), }, "step": 0, "timestamp": AnyStr(), @@ -6094,9 +6100,7 @@ def test_falsy_return_from_task( "a": 5, }, "name": "graph", - "triggers": [ - "__start__", - ], + "triggers": ("__start__",), }, "step": 0, "timestamp": AnyStr(), @@ -6110,9 +6114,7 @@ def test_falsy_return_from_task( {}, ), "name": "falsy_task", - "triggers": [ - "__pregel_push", - ], + "triggers": ("__pregel_push",), }, "step": 0, "timestamp": AnyStr(), @@ -6923,7 +6925,7 @@ def test_tags_stream_mode_messages() -> None: { "langgraph_step": 1, "langgraph_node": "call_model", - "langgraph_triggers": ["start:call_model"], + "langgraph_triggers": ("start:call_model",), "langgraph_path": ("__pregel_pull", "call_model"), "langgraph_checkpoint_ns": AnyStr("call_model:"), "checkpoint_ns": AnyStr("call_model:"), diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index a42a84a28..f1bb4d8ab 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1672,7 +1672,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "id": AnyStr(), "name": "one", "input": 2, - "triggers": ["input"], + "triggers": ("input",), }, }, { @@ -1683,7 +1683,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "id": AnyStr(), "name": "two", "input": [12], - "triggers": ["inbox"], + "triggers": ("inbox",), }, }, { @@ -1718,7 +1718,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "id": AnyStr(), "name": "two", "input": [3], - "triggers": ["inbox"], + "triggers": ("inbox",), }, }, { @@ -7571,7 +7571,7 @@ async def test_tags_stream_mode_messages() -> None: { "langgraph_step": 1, "langgraph_node": "call_model", - "langgraph_triggers": ["start:call_model"], + "langgraph_triggers": ("start:call_model",), "langgraph_path": ("__pregel_pull", "call_model"), "langgraph_checkpoint_ns": AnyStr("call_model:"), "checkpoint_ns": AnyStr("call_model:"), diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index fa9a221d0..2d442e131 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -1,4 +1,5 @@ import asyncio +import binascii import concurrent.futures from collections.abc import Sequence from contextlib import ( @@ -19,7 +20,7 @@ import langgraph.scheduler.kafka.serde as serde from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR from langgraph.errors import CheckpointNotLatest, GraphDelegate, TaskNotFound from langgraph.pregel import Pregel -from langgraph.pregel.algo import prepare_single_task +from langgraph.pregel.algo import checkpoint_null_version, prepare_single_task from langgraph.pregel.executor import ( AsyncBackgroundExecutor, BackgroundExecutor, @@ -209,6 +210,10 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager): for_execution=True, checkpointer=self.graph.checkpointer, store=self.graph.store, + checkpoint_id_bytes=binascii.unhexlify( + saved.checkpoint["id"].replace("-", "") + ), + checkpoint_null_version=checkpoint_null_version(saved.checkpoint), ): # execute task, saving writes runner = PregelRunner( @@ -421,6 +426,10 @@ class KafkaExecutor(AbstractContextManager): step=saved.metadata["step"] + 1, for_execution=True, checkpointer=self.graph.checkpointer, + checkpoint_id_bytes=binascii.unhexlify( + saved.checkpoint["id"].replace("-", "") + ), + checkpoint_null_version=checkpoint_null_version(saved.checkpoint), ): # execute task, saving writes runner = PregelRunner(