diff --git a/docs/README.md b/docs/README.md index 980ee6ade..b3be09886 100644 --- a/docs/README.md +++ b/docs/README.md @@ -14,6 +14,8 @@ To run the documentation server locally you can run: make serve-docs ``` +This will start the documentation server on [http://127.0.0.1:8000/langgraph/](http://127.0.0.1:8000/langgraph/). + ## Execute notebooks If you would like to automatically execute all of the notebooks, to mimic the "Run notebooks" GHA, you can run: diff --git a/docs/docs/cloud/how-tos/img/studio_graph_with_configuration.png b/docs/docs/cloud/how-tos/img/studio_graph_with_configuration.png new file mode 100644 index 000000000..91166691e Binary files /dev/null and b/docs/docs/cloud/how-tos/img/studio_graph_with_configuration.png differ diff --git a/docs/docs/cloud/how-tos/img/studio_node_configuration.png b/docs/docs/cloud/how-tos/img/studio_node_configuration.png new file mode 100644 index 000000000..1511f48ed Binary files /dev/null and b/docs/docs/cloud/how-tos/img/studio_node_configuration.png differ diff --git a/docs/docs/cloud/how-tos/iterate_graph_studio.md b/docs/docs/cloud/how-tos/iterate_graph_studio.md index 0453be39b..a98b41c8a 100644 --- a/docs/docs/cloud/how-tos/iterate_graph_studio.md +++ b/docs/docs/cloud/how-tos/iterate_graph_studio.md @@ -1,6 +1,133 @@ # Prompt Engineering in LangGraph Studio -In LangGraph Studio you can iterate on the prompts used within your graph by utilizing the LangSmith Playground. To do so: +## Overview + +A central aspect of agent development is prompt engineering. LangGraph Studio makes it easy to iterate on the prompts used within your graph directly within the UI. + +## Setup + +The first step is to define your [configuration](https://langchain-ai.github.io/langgraph/how-tos/configuration/) such that LangGraph Studio is aware of the prompts you want to iterate on and which nodes they are associated with. + +### Reference + +When defining your configuration, you can use special metadata keys to instruct LangGraph Studio how to handle different fields. Here's a reference for the available configuration options: + +#### `langgraph_nodes` + +- **Description**: Specifies which graph nodes a configuration field is associated with. +- **Value Type**: Array of strings, where each string is the name of a node in your graph. +- **Usage Context**: Include in the `json_schema_extra` dictionary for Pydantic models or the `metadata["json_schema_extra"]` dictionary for dataclasses. +- **Required**: No, but necessary if you want a field to be editable for specific nodes in the UI. +- **Example**: + ```python + system_prompt: str = Field( + default="You are a helpful AI assistant.", + json_schema_extra={"langgraph_nodes": ["call_model", "other_node"]}, + ) + ``` + +#### `langgraph_type` + +- **Description**: Specifies the type of configuration field, which determines how it's handled in the UI. +- **Value Type**: String +- **Supported Values**: + - `"prompt"`: Indicates the field contains prompt text that should be treated specially in the UI. +- **Usage Context**: Include in the `json_schema_extra` dictionary for Pydantic models or the `metadata["json_schema_extra"]` dictionary for dataclasses. +- **Required**: No, but helpful for prompt fields to enable special handling. +- **Example**: + ```python + system_prompt: str = Field( + default="You are a helpful AI assistant.", + json_schema_extra={ + "langgraph_nodes": ["call_model"], + "langgraph_type": "prompt", + }, + ) + ``` + +### Example + +For example, if you have a node called `call_model` whose system prompt you want to iterate on, you can define a configuration like the following. + +```python +## Using Pydantic +from pydantic import BaseModel, Field +from typing import Annotated, Literal + +class Configuration(BaseModel): + """The configuration for the agent.""" + + system_prompt: str = Field( + default="You are a helpful AI assistant.", + description="The system prompt to use for the agent's interactions. " + "This prompt sets the context and behavior for the agent.", + json_schema_extra={ + "langgraph_nodes": ["call_model"], + "langgraph_type": "prompt", + }, + ) + + model: Annotated[ + Literal[ + "anthropic/claude-3-7-sonnet-latest", + "anthropic/claude-3-5-haiku-latest", + "openai/o1", + "openai/gpt-4o-mini", + "openai/o1-mini", + "openai/o3-mini", + ], + {"__template_metadata__": {"kind": "llm"}}, + ] = Field( + default="openai/gpt-4o-mini", + description="The name of the language model to use for the agent's main interactions. " + "Should be in the form: provider/model-name.", + json_schema_extra={"langgraph_nodes": ["call_model"]}, + ) + +## Using Dataclasses +from dataclasses import dataclass, field + +@dataclass(kw_only=True) +class Configuration: + """The configuration for the agent.""" + + system_prompt: str = field( + default="You are a helpful AI assistant.", + metadata={ + "description": "The system prompt to use for the agent's interactions. " + "This prompt sets the context and behavior for the agent.", + "json_schema_extra": {"langgraph_nodes": ["call_model"]}, + }, + ) + + model: Annotated[str, {"__template_metadata__": {"kind": "llm"}}] = field( + default="anthropic/claude-3-5-sonnet-20240620", + metadata={ + "description": "The name of the language model to use for the agent's main interactions. " + "Should be in the form: provider/model-name.", + "json_schema_extra": {"langgraph_nodes": ["call_model"]}, + }, + ) + +``` + +## Iterating on prompts + +### Node Configuration + +With this set up, running your graph and viewing in LangGraph Studio will result in the graph rendering like such. + +**Note the configuration icon in the top right corner of the `call_model` node**: + +![Graph in Studio](../img/studio_graph_with_configuration.png){width=1200} + +Clicking this icon will open a modal where you can edit the configuration for all of the fields associated with the `call_model` node. From here, you can save your changes and apply them to the graph. Note that these values reflect the currently active assistant, and saving will update the assistant with the new values. + +![Configuration modal](../img/studio_node_configuration.png){width=1200} + +### Playground + +LangGraph Studio also supports prompt engineering through an integration with the LangSmith Playground. To do so: 1. Open an existing thread or create a new one. 2. Within the thread log, any nodes that have made an LLM call will have a "View LLM Runs" button. Clicking this will open a popover with the LLM runs for that node. @@ -8,8 +135,6 @@ In LangGraph Studio you can iterate on the prompts used within your graph by uti ![Playground in Studio](../img/studio_playground.png){width=1200} - - From here you can edit the prompt, test different model configurations and re-run just this LLM call without having to re-run the entire graph. When you are happy with your changes, you can copy the updated prompt back into your graph. For more information on how to use the LangSmith Playground, see the [LangSmith Playground documentation](https://docs.smith.langchain.com/prompt_engineering/how_to_guides#playground). diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 77f682fa3..93251576a 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -92,7 +92,7 @@ from langgraph.pregel.algo import ( prepare_next_tasks, ) from langgraph.pregel.debug import tasks_w_writes -from langgraph.pregel.io import read_channels +from langgraph.pregel.io import map_input, read_channels from langgraph.pregel.loop import AsyncPregelLoop, StreamProtocol, SyncPregelLoop from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager from langgraph.pregel.messages import StreamMessagesHandler @@ -109,6 +109,7 @@ from langgraph.types import ( Checkpointer, LoopProtocol, StateSnapshot, + StateUpdate, StreamChunk, StreamMode, ) @@ -1163,22 +1164,38 @@ class Pregel(PregelProtocol): checkpoint_tuple.config, checkpoint_tuple ) - def update_state( + def bulk_update_state( self, config: RunnableConfig, - values: Optional[Union[dict[str, Any], Any]], - as_node: Optional[str] = None, + supersteps: Sequence[Sequence[StateUpdate]], ) -> RunnableConfig: - """Update the state of the graph with the given values, as if they came from - node `as_node`. If `as_node` is not provided, it will be set to the last node - that updated the state, if not ambiguous. + """Apply updates to the graph state in bulk. Requires a checkpointer to be set. + + Args: + config: The config to apply the updates to. + supersteps: A list of supersteps, each including a list of updates to apply sequentially to a graph state. + Each update is a tuple of the form `(values, as_node)`. + + Raises: + ValueError: If no checkpointer is set or no updates are provided. + InvalidUpdateError: If an invalid update is provided. + + Returns: + RunnableConfig: The updated config. """ + checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: raise ValueError("No checkpointer set") + if len(supersteps) == 0: + raise ValueError("No supersteps provided") + + if any(len(u) == 0 for u in supersteps): + raise ValueError("No updates provided") + # delegate to subgraph if ( checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") @@ -1187,43 +1204,224 @@ class Pregel(PregelProtocol): recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name for _, pregel in self.get_subgraphs(namespace=recast, recurse=True): - return pregel.update_state( + return pregel.bulk_update_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), - values, - as_node, + supersteps, ) else: raise ValueError(f"Subgraph {recast} not found") - # get last checkpoint - config = ensure_config(self.config, config) - saved = checkpointer.get_tuple(config) - checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() - checkpoint_previous_versions = ( - saved.checkpoint["channel_versions"].copy() if saved else {} - ) - step = saved.metadata.get("step", -1) if saved else -1 - # merge configurable fields with previous checkpoint config - checkpoint_config = patch_configurable( - config, - {CONFIG_KEY_CHECKPOINT_NS: config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")}, - ) - checkpoint_metadata = config["metadata"] - if saved: - checkpoint_config = patch_configurable(config, saved.config[CONF]) - checkpoint_metadata = {**saved.metadata, **checkpoint_metadata} - with ChannelsManager( - self.channels, - checkpoint, - LoopProtocol(config=config, step=step + 1, stop=step + 2), - ) as (channels, managed): - # no values as END, just clear all tasks - if values is None and as_node == END: - if saved is not None: + def perform_superstep( + input_config: RunnableConfig, updates: Sequence[StateUpdate] + ) -> RunnableConfig: + # get last checkpoint + config = ensure_config(self.config, input_config) + saved = checkpointer.get_tuple(config) + checkpoint = ( + copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() + ) + checkpoint_previous_versions = ( + saved.checkpoint["channel_versions"].copy() if saved else {} + ) + step = saved.metadata.get("step", -1) if saved else -1 + # merge configurable fields with previous checkpoint config + checkpoint_config = patch_configurable( + config, + { + CONFIG_KEY_CHECKPOINT_NS: config[CONF].get( + CONFIG_KEY_CHECKPOINT_NS, "" + ) + }, + ) + checkpoint_metadata = config["metadata"] + if saved: + checkpoint_config = patch_configurable(config, saved.config[CONF]) + checkpoint_metadata = {**saved.metadata, **checkpoint_metadata} + with ChannelsManager( + self.channels, + checkpoint, + LoopProtocol(config=config, step=step + 1, stop=step + 2), + ) as (channels, managed): + values, as_node = updates[0] + + # no values as END, just clear all tasks + if values is None and as_node == END: + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot apply multiple updates when clearing state" + ) + + if saved is not None: + # tasks for this checkpoint + next_tasks = prepare_next_tasks( + checkpoint, + saved.pending_writes or [], + self.nodes, + channels, + managed, + saved.config, + saved.metadata.get("step", -1) + 1, + for_execution=True, + store=self.store, + checkpointer=self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None, + manager=None, + ) + # apply null writes + if null_writes := [ + w[1:] + for w in saved.pending_writes or [] + if w[0] == NULL_TASK_ID + ]: + apply_writes( + saved.checkpoint, + channels, + [PregelTaskWrites((), INPUT, null_writes, [])], + None, + ) + # apply writes from tasks that already ran + for tid, k, v in saved.pending_writes or []: + if k in (ERROR, INTERRUPT, SCHEDULED): + continue + if tid not in next_tasks: + continue + next_tasks[tid].writes.append((k, v)) + # clear all current tasks + apply_writes(checkpoint, channels, next_tasks.values(), None) + # save checkpoint + next_config = checkpointer.put( + checkpoint_config, + create_checkpoint(checkpoint, None, step), + { + **checkpoint_metadata, + "source": "update", + "step": step + 1, + "writes": {}, + "parents": saved.metadata.get("parents", {}) + if saved + else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + # no values, empty checkpoint + if values is None and as_node is None: + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot create empty checkpoint with multiple updates" + ) + + next_checkpoint = create_checkpoint(checkpoint, None, step) + # copy checkpoint + next_config = checkpointer.put( + checkpoint_config, + next_checkpoint, + { + **checkpoint_metadata, + "source": "update", + "step": step + 1, + "writes": {}, + "parents": saved.metadata.get("parents", {}) + if saved + else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + + # act as an input + if as_node == INPUT: + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot apply multiple updates when updating as input" + ) + + if input_writes := deque(map_input(self.input_channels, values)): + apply_writes( + checkpoint, + channels, + [PregelTaskWrites((), INPUT, input_writes, [])], + checkpointer.get_next_version, + ) + + # apply input write to channels + next_step = ( + step + 1 + if saved and saved.metadata.get("step") is not None + else -1 + ) + next_config = checkpointer.put( + checkpoint_config, + create_checkpoint(checkpoint, channels, next_step), + { + **checkpoint_metadata, + "source": "input", + "step": next_step, + "writes": dict(input_writes), + }, + get_new_channel_versions( + checkpoint_previous_versions, + checkpoint["channel_versions"], + ), + ) + + # store the writes + checkpointer.put_writes( + next_config, + input_writes, + str(uuid5(UUID(checkpoint["id"]), INPUT)), + ) + + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + else: + raise InvalidUpdateError( + f"Received no input writes for {self.input_channels}" + ) + + # no values, copy checkpoint + if values is None and as_node == "__copy__": + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot copy checkpoint with multiple updates" + ) + + next_checkpoint = create_checkpoint(checkpoint, None, step) + # copy checkpoint + next_config = checkpointer.put( + saved.parent_config or saved.config + if saved + else checkpoint_config, + next_checkpoint, + { + **checkpoint_metadata, + "source": "fork", + "step": step + 1, + "parents": saved.metadata.get("parents", {}) + if saved + else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + # apply pending writes, if not on specific checkpoint + if ( + CONFIG_KEY_CHECKPOINT_ID not in config[CONF] + and saved is not None + and saved.pending_writes + ): # tasks for this checkpoint next_tasks = prepare_next_tasks( checkpoint, - saved.pending_writes or [], + saved.pending_writes, self.nodes, channels, managed, @@ -1250,182 +1448,106 @@ class Pregel(PregelProtocol): [PregelTaskWrites((), INPUT, null_writes, [])], None, ) - # apply writes from tasks that already ran - for tid, k, v in saved.pending_writes or []: + # apply writes + for tid, k, v in saved.pending_writes: if k in (ERROR, INTERRUPT, SCHEDULED): continue if tid not in next_tasks: continue next_tasks[tid].writes.append((k, v)) - # clear all current tasks - apply_writes(checkpoint, channels, next_tasks.values(), None) - # save checkpoint - next_config = checkpointer.put( - checkpoint_config, - create_checkpoint(checkpoint, None, step), - { - **checkpoint_metadata, - "source": "update", - "step": step + 1, - "writes": {}, - "parents": saved.metadata.get("parents", {}) if saved else {}, - }, - {}, - ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - # no values, empty checkpoint - if values is None and as_node is None: - next_checkpoint = create_checkpoint(checkpoint, None, step) - # copy checkpoint - next_config = checkpointer.put( - checkpoint_config, - next_checkpoint, - { - **checkpoint_metadata, - "source": "update", - "step": step + 1, - "writes": {}, - "parents": saved.metadata.get("parents", {}) if saved else {}, - }, - {}, - ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - # no values, copy checkpoint - if values is None and as_node == "__copy__": - next_checkpoint = create_checkpoint(checkpoint, None, step) - # copy checkpoint - next_config = checkpointer.put( - saved.parent_config or saved.config if saved else checkpoint_config, - next_checkpoint, - { - **checkpoint_metadata, - "source": "fork", - "step": step + 1, - "parents": saved.metadata.get("parents", {}) if saved else {}, - }, - {}, - ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - # apply pending writes, if not on specific checkpoint - if ( - CONFIG_KEY_CHECKPOINT_ID not in config[CONF] - and saved is not None - and saved.pending_writes - ): - # tasks for this checkpoint - next_tasks = prepare_next_tasks( - checkpoint, - saved.pending_writes, - self.nodes, - channels, - managed, - saved.config, - saved.metadata.get("step", -1) + 1, - for_execution=True, - store=self.store, - checkpointer=( - self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None - ), - manager=None, - ) - # apply null writes - if null_writes := [ - w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID - ]: - apply_writes( - saved.checkpoint, - channels, - [PregelTaskWrites((), INPUT, null_writes, [])], - None, - ) - # apply writes - for tid, k, v in saved.pending_writes: - if k in (ERROR, INTERRUPT, SCHEDULED): - continue - if tid not in next_tasks: - continue - next_tasks[tid].writes.append((k, v)) - if tasks := [t for t in next_tasks.values() if t.writes]: - apply_writes(checkpoint, channels, tasks, None) - # find last node that updated the state, if not provided - if as_node is None and not any( - v for vv in checkpoint["versions_seen"].values() for v in vv.values() - ): - if ( - isinstance(self.input_channels, str) - and self.input_channels in self.nodes + if tasks := [t for t in next_tasks.values() if t.writes]: + apply_writes(checkpoint, channels, tasks, None) + valid_updates: list[tuple[str, Optional[dict[str, Any]]]] = [] + if len(updates) == 1: + values, as_node = updates[0] + # find last node that updated the state, if not provided + if as_node is None and not any( + v + for vv in checkpoint["versions_seen"].values() + for v in vv.values() ): - as_node = self.input_channels - elif as_node is None: - last_seen_by_node = sorted( - (v, n) - for n, seen in checkpoint["versions_seen"].items() - if n in self.nodes - for v in seen.values() + if ( + isinstance(self.input_channels, str) + and self.input_channels in self.nodes + ): + as_node = self.input_channels + elif as_node is None: + last_seen_by_node = sorted( + (v, n) + for n, seen in checkpoint["versions_seen"].items() + if n in self.nodes + for v in seen.values() + ) + # if two nodes updated the state at the same time, it's ambiguous + if last_seen_by_node: + if len(last_seen_by_node) == 1: + as_node = last_seen_by_node[0][1] + elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: + as_node = last_seen_by_node[-1][1] + if as_node is None: + raise InvalidUpdateError("Ambiguous update, specify as_node") + if as_node not in self.nodes: + raise InvalidUpdateError(f"Node {as_node} does not exist") + valid_updates.append((as_node, values)) + else: + for values, as_node in updates: + if as_node is None: + raise InvalidUpdateError( + "as_node is required when applying multiple updates" + ) + if as_node not in self.nodes: + raise InvalidUpdateError(f"Node {as_node} does not exist") + + valid_updates.append((as_node, values)) + + run_tasks: list[PregelTaskWrites] = [] + run_task_ids: list[str] = [] + + for as_node, values in valid_updates: + # create task to run all writers of the chosen node + writers = self.nodes[as_node].flat_writers + if not writers: + raise InvalidUpdateError(f"Node {as_node} has no writers") + writes: deque[tuple[str, Any]] = deque() + task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) + task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) + run_tasks.append(task) + run_task_ids.append(task_id) + run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] + # execute task + run.invoke( + values, + patch_config( + config, + run_name=self.name + "UpdateState", + configurable={ + # deque.extend is thread-safe + CONFIG_KEY_SEND: partial( + local_write, + writes.extend, + self.nodes.keys(), + ), + CONFIG_KEY_READ: partial( + local_read, + step + 1, + checkpoint, + channels, + managed, + task, + config, + ), + }, + ), ) - # if two nodes updated the state at the same time, it's ambiguous - if last_seen_by_node: - if len(last_seen_by_node) == 1: - as_node = last_seen_by_node[0][1] - elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: - as_node = last_seen_by_node[-1][1] - if as_node is None: - raise InvalidUpdateError("Ambiguous update, specify as_node") - if as_node not in self.nodes: - raise InvalidUpdateError(f"Node {as_node} does not exist") - # create task to run all writers of the chosen node - writers = self.nodes[as_node].flat_writers - if not writers: - raise InvalidUpdateError(f"Node {as_node} has no writers") - writes: deque[tuple[str, Any]] = deque() - task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) - task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) - run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] - # execute task - run.invoke( - values, - patch_config( - config, - run_name=self.name + "UpdateState", - configurable={ - # deque.extend is thread-safe - CONFIG_KEY_SEND: partial( - local_write, - writes.extend, - self.nodes.keys(), - ), - CONFIG_KEY_READ: partial( - local_read, - step + 1, - checkpoint, - channels, - managed, - task, - config, - ), - }, - ), - ) # save task writes - # channel writes are saved to current checkpoint - # push writes are saved to next checkpoint - channel_writes, push_writes = ( - [w for w in task.writes if w[0] != PUSH], - [w for w in task.writes if w[0] == PUSH], - ) - if saved and channel_writes: - checkpointer.put_writes(checkpoint_config, channel_writes, task_id) + for task_id, task in zip(run_task_ids, run_tasks): + # channel writes are saved to current checkpoint + channel_writes = [w for w in task.writes if w[0] != PUSH] + if saved and channel_writes: + checkpointer.put_writes(checkpoint_config, channel_writes, task_id) # apply to checkpoint and save mv_writes, _ = apply_writes( - checkpoint, channels, [task], checkpointer.get_next_version + checkpoint, channels, run_tasks, checkpointer.get_next_version ) assert not mv_writes, "Can't write to SharedValues from update_state" checkpoint = create_checkpoint(checkpoint, channels, step + 1) @@ -1436,33 +1558,57 @@ class Pregel(PregelProtocol): **checkpoint_metadata, "source": "update", "step": step + 1, - "writes": {as_node: values}, + "writes": {as_node: values for as_node, values in valid_updates}, "parents": saved.metadata.get("parents", {}) if saved else {}, }, get_new_channel_versions( checkpoint_previous_versions, checkpoint["channel_versions"] ), ) - if push_writes: - checkpointer.put_writes(next_config, push_writes, task_id) + for task_id, task in zip(run_task_ids, run_tasks): + # save push writes + if push_writes := [w for w in task.writes if w[0] == PUSH]: + checkpointer.put_writes(next_config, push_writes, task_id) + return patch_checkpoint_map(next_config, saved.metadata if saved else None) - async def aupdate_state( + current_config = config + for superstep in supersteps: + current_config = perform_superstep(current_config, superstep) + return current_config + + async def abulk_update_state( self, config: RunnableConfig, - values: dict[str, Any] | Any, - as_node: Optional[str] = None, + supersteps: Sequence[Sequence[StateUpdate]], ) -> RunnableConfig: - """Update the state of the graph asynchronously with the given values, as if they came from - node `as_node`. If `as_node` is not provided, it will be set to the last node - that updated the state, if not ambiguous. + """Apply updates to the graph state in bulk. Requires a checkpointer to be set. + + Args: + config: The config to apply the updates to. + supersteps: A list of supersteps, each including a list of updates to apply sequentially to a graph state. + Each update is a tuple of the form `(values, as_node)`. + + Raises: + ValueError: If no checkpointer is set or no updates are provided. + InvalidUpdateError: If an invalid update is provided. + + Returns: + RunnableConfig: The updated config. """ + checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: raise ValueError("No checkpointer set") + if len(supersteps) == 0: + raise ValueError("No supersteps provided") + + if any(len(u) == 0 for u in supersteps): + raise ValueError("No updates provided") + # delegate to subgraph if ( checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") @@ -1471,46 +1617,225 @@ class Pregel(PregelProtocol): recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True): - return await pregel.aupdate_state( + return await pregel.abulk_update_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), - values, - as_node, + supersteps, ) else: raise ValueError(f"Subgraph {recast} not found") - # get last checkpoint - config = ensure_config(self.config, config) - saved = await checkpointer.aget_tuple(config) - checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() - checkpoint_previous_versions = ( - saved.checkpoint["channel_versions"].copy() if saved else {} - ) - step = saved.metadata.get("step", -1) if saved else -1 - # merge configurable fields with previous checkpoint config - checkpoint_config = patch_configurable( - config, - {CONFIG_KEY_CHECKPOINT_NS: config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")}, - ) - checkpoint_metadata = config["metadata"] - if saved: - checkpoint_config = patch_configurable(config, saved.config[CONF]) - checkpoint_metadata = {**saved.metadata, **checkpoint_metadata} - async with AsyncChannelsManager( - self.channels, - checkpoint, - LoopProtocol(config=config, step=step + 1, stop=step + 2), - ) as ( - channels, - managed, - ): - # no values, just clear all tasks - if values is None and as_node == END: - if saved is not None: + async def aperform_superstep( + input_config: RunnableConfig, updates: Sequence[StateUpdate] + ) -> RunnableConfig: + # get last checkpoint + config = ensure_config(self.config, input_config) + saved = await checkpointer.aget_tuple(config) + checkpoint = ( + copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() + ) + checkpoint_previous_versions = ( + saved.checkpoint["channel_versions"].copy() if saved else {} + ) + step = saved.metadata.get("step", -1) if saved else -1 + # merge configurable fields with previous checkpoint config + checkpoint_config = patch_configurable( + config, + { + CONFIG_KEY_CHECKPOINT_NS: config[CONF].get( + CONFIG_KEY_CHECKPOINT_NS, "" + ) + }, + ) + checkpoint_metadata = config["metadata"] + if saved: + checkpoint_config = patch_configurable(config, saved.config[CONF]) + checkpoint_metadata = {**saved.metadata, **checkpoint_metadata} + async with AsyncChannelsManager( + self.channels, + checkpoint, + LoopProtocol(config=config, step=step + 1, stop=step + 2), + ) as ( + channels, + managed, + ): + values, as_node = updates[0] + # no values, just clear all tasks + if values is None and as_node == END: + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot apply multiple updates when clearing state" + ) + if saved is not None: + # tasks for this checkpoint + next_tasks = prepare_next_tasks( + checkpoint, + saved.pending_writes or [], + self.nodes, + channels, + managed, + saved.config, + saved.metadata.get("step", -1) + 1, + for_execution=True, + store=self.store, + checkpointer=self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None, + manager=None, + ) + # apply null writes + if null_writes := [ + w[1:] + for w in saved.pending_writes or [] + if w[0] == NULL_TASK_ID + ]: + apply_writes( + saved.checkpoint, + channels, + [PregelTaskWrites((), INPUT, null_writes, [])], + None, + ) + # apply writes from tasks that already ran + for tid, k, v in saved.pending_writes or []: + if k in (ERROR, INTERRUPT, SCHEDULED): + continue + if tid not in next_tasks: + continue + next_tasks[tid].writes.append((k, v)) + # clear all current tasks + apply_writes(checkpoint, channels, next_tasks.values(), None) + # save checkpoint + next_config = await checkpointer.aput( + checkpoint_config, + create_checkpoint(checkpoint, None, step), + { + **checkpoint_metadata, + "source": "update", + "step": step + 1, + "writes": {}, + "parents": saved.metadata.get("parents", {}) + if saved + else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + # no values, empty checkpoint + if values is None and as_node is None: + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot create empty checkpoint with multiple updates" + ) + + next_checkpoint = create_checkpoint(checkpoint, None, step) + # copy checkpoint + next_config = await checkpointer.aput( + checkpoint_config, + next_checkpoint, + { + **checkpoint_metadata, + "source": "update", + "step": step + 1, + "writes": {}, + "parents": saved.metadata.get("parents", {}) + if saved + else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + + # act as an input + if as_node == INPUT: + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot apply multiple updates when updating as input" + ) + + if input_writes := deque(map_input(self.input_channels, values)): + apply_writes( + checkpoint, + channels, + [PregelTaskWrites((), INPUT, input_writes, [])], + checkpointer.get_next_version, + ) + + # apply input write to channels + next_step = ( + step + 1 + if saved and saved.metadata.get("step") is not None + else -1 + ) + next_config = await checkpointer.aput( + checkpoint_config, + create_checkpoint(checkpoint, channels, next_step), + { + **checkpoint_metadata, + "source": "input", + "step": next_step, + "writes": dict(input_writes), + }, + get_new_channel_versions( + checkpoint_previous_versions, + checkpoint["channel_versions"], + ), + ) + + # store the writes + await checkpointer.aput_writes( + next_config, + input_writes, + str(uuid5(UUID(checkpoint["id"]), INPUT)), + ) + + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + else: + raise InvalidUpdateError( + f"Received no input writes for {self.input_channels}" + ) + + # no values, copy checkpoint + if values is None and as_node == "__copy__": + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot copy checkpoint with multiple updates" + ) + + next_checkpoint = create_checkpoint(checkpoint, None, step) + # copy checkpoint + next_config = await checkpointer.aput( + saved.parent_config or saved.config + if saved + else checkpoint_config, + next_checkpoint, + { + **checkpoint_metadata, + "source": "fork", + "step": step + 1, + "parents": saved.metadata.get("parents", {}) + if saved + else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + # apply pending writes, if not on specific checkpoint + if ( + CONFIG_KEY_CHECKPOINT_ID not in config[CONF] + and saved is not None + and saved.pending_writes + ): # tasks for this checkpoint next_tasks = prepare_next_tasks( checkpoint, - saved.pending_writes or [], + saved.pending_writes, self.nodes, channels, managed, @@ -1537,181 +1862,103 @@ class Pregel(PregelProtocol): [PregelTaskWrites((), INPUT, null_writes, [])], None, ) - # apply writes from tasks that already ran - for tid, k, v in saved.pending_writes or []: + for tid, k, v in saved.pending_writes: if k in (ERROR, INTERRUPT, SCHEDULED): continue if tid not in next_tasks: continue next_tasks[tid].writes.append((k, v)) - # clear all current tasks - apply_writes(checkpoint, channels, next_tasks.values(), None) - # save checkpoint - next_config = await checkpointer.aput( - checkpoint_config, - create_checkpoint(checkpoint, None, step), - { - **checkpoint_metadata, - "source": "update", - "step": step + 1, - "writes": {}, - "parents": saved.metadata.get("parents", {}) if saved else {}, - }, - {}, - ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - # no values, empty checkpoint - if values is None and as_node is None: - next_checkpoint = create_checkpoint(checkpoint, None, step) - # copy checkpoint - next_config = await checkpointer.aput( - checkpoint_config, - next_checkpoint, - { - **checkpoint_metadata, - "source": "update", - "step": step + 1, - "writes": {}, - "parents": saved.metadata.get("parents", {}) if saved else {}, - }, - {}, - ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - # no values, copy checkpoint - if values is None and as_node == "__copy__": - next_checkpoint = create_checkpoint(checkpoint, None, step) - # copy checkpoint - next_config = await checkpointer.aput( - saved.parent_config or saved.config if saved else checkpoint_config, - next_checkpoint, - { - **checkpoint_metadata, - "source": "fork", - "step": step + 1, - "parents": saved.metadata.get("parents", {}) if saved else {}, - }, - {}, - ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - # apply pending writes, if not on specific checkpoint - if ( - CONFIG_KEY_CHECKPOINT_ID not in config[CONF] - and saved is not None - and saved.pending_writes - ): - # tasks for this checkpoint - next_tasks = prepare_next_tasks( - checkpoint, - saved.pending_writes, - self.nodes, - channels, - managed, - saved.config, - saved.metadata.get("step", -1) + 1, - for_execution=True, - store=self.store, - checkpointer=( - self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None - ), - manager=None, - ) - # apply null writes - if null_writes := [ - w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID - ]: - apply_writes( - saved.checkpoint, - channels, - [PregelTaskWrites((), INPUT, null_writes, [])], - None, + if tasks := [t for t in next_tasks.values() if t.writes]: + apply_writes(checkpoint, channels, tasks, None) + valid_updates: list[tuple[str, Optional[dict[str, Any]]]] = [] + if len(updates) == 1: + values, as_node = updates[0] + # find last node that updated the state, if not provided + if as_node is None and not saved: + if ( + isinstance(self.input_channels, str) + and self.input_channels in self.nodes + ): + as_node = self.input_channels + elif as_node is None: + last_seen_by_node = sorted( + (v, n) + for n, seen in checkpoint["versions_seen"].items() + if n in self.nodes + for v in seen.values() ) - for tid, k, v in saved.pending_writes: - if k in (ERROR, INTERRUPT, SCHEDULED): - continue - if tid not in next_tasks: - continue - next_tasks[tid].writes.append((k, v)) - if tasks := [t for t in next_tasks.values() if t.writes]: - apply_writes(checkpoint, channels, tasks, None) - # find last node that updated the state, if not provided - if as_node is None and not saved: - if ( - isinstance(self.input_channels, str) - and self.input_channels in self.nodes - ): - as_node = self.input_channels - elif as_node is None: - last_seen_by_node = sorted( - (v, n) - for n, seen in checkpoint["versions_seen"].items() - if n in self.nodes - for v in seen.values() + # if two nodes updated the state at the same time, it's ambiguous + if last_seen_by_node: + if len(last_seen_by_node) == 1: + as_node = last_seen_by_node[0][1] + elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: + as_node = last_seen_by_node[-1][1] + if as_node is None: + raise InvalidUpdateError("Ambiguous update, specify as_node") + if as_node not in self.nodes: + raise InvalidUpdateError(f"Node {as_node} does not exist") + valid_updates.append((as_node, values)) + else: + for values, as_node in updates: + if as_node is None: + raise InvalidUpdateError( + "as_node is required when applying multiple updates" + ) + if as_node not in self.nodes: + raise InvalidUpdateError(f"Node {as_node} does not exist") + + valid_updates.append((as_node, values)) + + run_tasks: list[PregelTaskWrites] = [] + run_task_ids: list[str] = [] + + for as_node, values in valid_updates: + # create task to run all writers of the chosen node + writers = self.nodes[as_node].flat_writers + if not writers: + raise InvalidUpdateError(f"Node {as_node} has no writers") + writes: deque[tuple[str, Any]] = deque() + task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) + task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) + run_tasks.append(task) + run_task_ids.append(task_id) + run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] + # execute task + await run.ainvoke( + values, + patch_config( + config, + run_name=self.name + "UpdateState", + configurable={ + # deque.extend is thread-safe + CONFIG_KEY_SEND: partial( + local_write, + writes.extend, + self.nodes.keys(), + ), + CONFIG_KEY_READ: partial( + local_read, + step + 1, + checkpoint, + channels, + managed, + task, + config, + ), + }, + ), ) - # if two nodes updated the state at the same time, it's ambiguous - if last_seen_by_node: - if len(last_seen_by_node) == 1: - as_node = last_seen_by_node[0][1] - elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: - as_node = last_seen_by_node[-1][1] - if as_node is None: - raise InvalidUpdateError("Ambiguous update, specify as_node") - if as_node not in self.nodes: - raise InvalidUpdateError(f"Node {as_node} does not exist") - # create task to run all writers of the chosen node - writers = self.nodes[as_node].flat_writers - if not writers: - raise InvalidUpdateError(f"Node {as_node} has no writers") - writes: deque[tuple[str, Any]] = deque() - task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) - task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) - run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] - # execute task - await run.ainvoke( - values, - patch_config( - config, - run_name=self.name + "UpdateState", - configurable={ - # deque.extend is thread-safe - CONFIG_KEY_SEND: partial( - local_write, - writes.extend, - self.nodes.keys(), - ), - CONFIG_KEY_READ: partial( - local_read, - step + 1, - checkpoint, - channels, - managed, - task, - config, - ), - }, - ), - ) # save task writes - # channel writes are saved to current checkpoint - # push writes are saved to next checkpoint - channel_writes, push_writes = ( - [w for w in task.writes if w[0] != PUSH], - [w for w in task.writes if w[0] == PUSH], - ) - if saved and channel_writes: - await checkpointer.aput_writes( - checkpoint_config, channel_writes, task_id - ) + for task_id, task in zip(run_task_ids, run_tasks): + # channel writes are saved to current checkpoint + channel_writes = [w for w in task.writes if w[0] != PUSH] + if saved and channel_writes: + await checkpointer.aput_writes( + checkpoint_config, channel_writes, task_id + ) # apply to checkpoint and save mv_writes, _ = apply_writes( - checkpoint, channels, [task], checkpointer.get_next_version + checkpoint, channels, run_tasks, checkpointer.get_next_version ) assert not mv_writes, "Can't write to SharedValues from update_state" checkpoint = create_checkpoint(checkpoint, channels, step + 1) @@ -1723,18 +1970,48 @@ class Pregel(PregelProtocol): **checkpoint_metadata, "source": "update", "step": step + 1, - "writes": {as_node: values}, + "writes": {as_node: values for as_node, values in valid_updates}, "parents": saved.metadata.get("parents", {}) if saved else {}, }, get_new_channel_versions( checkpoint_previous_versions, checkpoint["channel_versions"] ), ) - # save push writes - if push_writes: - await checkpointer.aput_writes(next_config, push_writes, task_id) + for task_id, task in zip(run_task_ids, run_tasks): + # save push writes + if push_writes := [w for w in task.writes if w[0] == PUSH]: + await checkpointer.aput_writes(next_config, push_writes, task_id) return patch_checkpoint_map(next_config, saved.metadata if saved else None) + current_config = config + for superstep in supersteps: + current_config = await aperform_superstep(current_config, superstep) + return current_config + + def update_state( + self, + config: RunnableConfig, + values: Optional[Union[dict[str, Any], Any]], + as_node: Optional[str] = None, + ) -> RunnableConfig: + """Update the state of the graph with the given values, as if they came from + node `as_node`. If `as_node` is not provided, it will be set to the last node + that updated the state, if not ambiguous. + """ + return self.bulk_update_state(config, [[StateUpdate(values, as_node)]]) + + async def aupdate_state( + self, + config: RunnableConfig, + values: dict[str, Any] | Any, + as_node: Optional[str] = None, + ) -> RunnableConfig: + """Update the state of the graph asynchronously with the given values, as if they came from + node `as_node`. If `as_node` is not provided, it will be set to the last node + that updated the state, if not ambiguous. + """ + return await self.abulk_update_state(config, [[StateUpdate(values, as_node)]]) + def _defaults( self, config: RunnableConfig, diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 01d8c1ff8..ab9668ab1 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -570,7 +570,7 @@ def prepare_single_task( CONFIG_KEY_CHECKPOINT_ID: None, CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, CONFIG_KEY_SCRATCHPAD: _scratchpad( - config, + config[CONF].get(CONFIG_KEY_SCRATCHPAD), pending_writes, task_id, ), @@ -680,7 +680,7 @@ def prepare_single_task( CONFIG_KEY_CHECKPOINT_ID: None, CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, CONFIG_KEY_SCRATCHPAD: _scratchpad( - config, + config[CONF].get(CONFIG_KEY_SCRATCHPAD), pending_writes, task_id, ), @@ -708,13 +708,14 @@ def prepare_single_task( if checkpoint_null_version is None: return # If any of the channels read by this process were updated - if triggers := _triggers( + if _triggers( channels, checkpoint["channel_versions"], checkpoint["versions_seen"].get(name), checkpoint_null_version, proc, ): + triggers = tuple(sorted(proc.triggers)) try: val = next( _proc_input(proc, managed, channels, for_execution=for_execution) @@ -774,7 +775,7 @@ def prepare_single_task( CONFIG_KEY_SEND: partial( local_write, writes.extend, - processes.keys(), + tuple(processes.keys()), ), CONFIG_KEY_READ: partial( local_read, @@ -783,7 +784,10 @@ def prepare_single_task( channels, managed, PregelTaskWrites( - task_path[:3], name, writes, triggers + task_path[:3], + name, + writes, + triggers, ), config, ), @@ -801,7 +805,7 @@ def prepare_single_task( CONFIG_KEY_CHECKPOINT_ID: None, CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, CONFIG_KEY_SCRATCHPAD: _scratchpad( - config, + config[CONF].get(CONFIG_KEY_SCRATCHPAD), pending_writes, task_id, ), @@ -852,7 +856,7 @@ def _triggers( def _scratchpad( - config: RunnableConfig, + parent_scratchpad: Optional[PregelScratchpad], pending_writes: list[PendingWrite], task_id: str, ) -> PregelScratchpad: @@ -861,9 +865,6 @@ def _scratchpad( null_resume_write = next( (w for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), None ) - parent_scratchpad: Optional[PregelScratchpad] = config[CONF].get( - CONFIG_KEY_SCRATCHPAD - ) def get_null_resume(consume: bool = False) -> Any: if null_resume_write is None: diff --git a/libs/langgraph/langgraph/pregel/protocol.py b/libs/langgraph/langgraph/pregel/protocol.py index ac046e949..5a27f417a 100644 --- a/libs/langgraph/langgraph/pregel/protocol.py +++ b/libs/langgraph/langgraph/pregel/protocol.py @@ -12,7 +12,7 @@ from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.graph import Graph as DrawableGraph from typing_extensions import Self -from langgraph.pregel.types import All, StateSnapshot, StreamMode +from langgraph.pregel.types import All, StateSnapshot, StateUpdate, StreamMode class PregelProtocol( @@ -69,6 +69,20 @@ class PregelProtocol( limit: Optional[int] = None, ) -> AsyncIterator[StateSnapshot]: ... + @abstractmethod + def bulk_update_state( + self, + config: RunnableConfig, + updates: Sequence[Sequence[StateUpdate]], + ) -> RunnableConfig: ... + + @abstractmethod + async def abulk_update_state( + self, + config: RunnableConfig, + updates: Sequence[Sequence[StateUpdate]], + ) -> RunnableConfig: ... + @abstractmethod def update_state( self, diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 6ba547b14..64fb4edb6 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -457,6 +457,20 @@ class RemoteGraph(PregelProtocol): for state in states: yield self._create_state_snapshot(state) + def bulk_update_state( + self, + config: RunnableConfig, + updates: list[tuple[Optional[dict[str, Any]], Optional[str]]], + ) -> RunnableConfig: + raise NotImplementedError + + async def abulk_update_state( + self, + config: RunnableConfig, + updates: list[tuple[Optional[dict[str, Any]], Optional[str]]], + ) -> RunnableConfig: + raise NotImplementedError + def update_state( self, config: RunnableConfig, diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index 7a72b88c9..212c0c4af 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -7,6 +7,7 @@ from langgraph.types import ( PregelTask, RetryPolicy, StateSnapshot, + StateUpdate, StreamMode, StreamWriter, default_retry_on, @@ -14,6 +15,7 @@ from langgraph.types import ( __all__ = [ "All", + "StateUpdate", "CachePolicy", "PregelExecutableTask", "PregelTask", diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 4c8cf6da4..00339119d 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -133,6 +133,11 @@ class Interrupt: when: Literal["during"] = dataclasses.field(default="during", repr=False) +class StateUpdate(NamedTuple): + values: Optional[dict[str, Any]] + as_node: Optional[str] = None + + class PregelTask(NamedTuple): id: str name: str diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index c625f47bc..17561c247 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph" -version = "0.3.16" +version = "0.3.17" 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 d6637f000..6616d4cbb 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": ("branch:to:agent", "start:agent", "tools"), "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2542,7 +2542,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 3, "langgraph_node": "agent", - "langgraph_triggers": ("tools",), + "langgraph_triggers": ("branch:to:agent", "start:agent", "tools"), "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2585,7 +2585,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 5, "langgraph_node": "agent", - "langgraph_triggers": ("tools",), + "langgraph_triggers": ("branch:to:agent", "start:agent", "tools"), "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -5501,7 +5501,10 @@ 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": ( + "branch:to:rewrite_query", + "start:rewrite_query", + ), }, }, ), @@ -5532,7 +5535,10 @@ 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": ( + "branch:to:retriever_one", + "rewrite_query", + ), }, }, ), @@ -5546,7 +5552,10 @@ 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": ( + "branch:to:retriever_two", + "rewrite_query", + ), }, }, ), @@ -5608,7 +5617,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": (AnyStr("retriever_"),), + "triggers": ("branch:to:qa", "retriever_one", "retriever_two"), }, }, ), @@ -6634,7 +6643,7 @@ def test_branch_then( "id": AnyStr(), "name": "prepare", "input": {"my_key": "value", "market": "DE"}, - "triggers": ("start:prepare",), + "triggers": ("branch:to:prepare", "start:prepare"), }, }, { @@ -6773,7 +6782,10 @@ 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", + "branch:to:finish", + ), }, }, { @@ -7783,7 +7795,7 @@ def test_nested_graph_state( "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], + "langgraph_triggers": ["branch:to:inner", "outer_1"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -7978,7 +7990,7 @@ def test_nested_graph_state( "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], + "langgraph_triggers": ["branch:to:inner", "outer_1"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -8021,7 +8033,7 @@ def test_nested_graph_state( "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], + "langgraph_triggers": ["branch:to:inner", "outer_1"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -8070,7 +8082,7 @@ def test_nested_graph_state( "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], + "langgraph_triggers": ["branch:to:inner", "outer_1"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -8504,7 +8516,7 @@ def test_doubly_nested_graph_state( "langgraph_node": "child_1", "langgraph_path": [PULL, AnyStr("child_1")], "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], + "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config=( @@ -8588,7 +8600,10 @@ def test_doubly_nested_graph_state( AnyStr("child_1"), ], "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], + "langgraph_triggers": [ + "branch:to:child_1", + AnyStr("start:child_1"), + ], }, created_at=AnyStr(), parent_config=( @@ -8635,7 +8650,7 @@ def test_doubly_nested_graph_state( "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -8931,7 +8946,7 @@ def test_doubly_nested_graph_state( "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -8970,7 +8985,7 @@ def test_doubly_nested_graph_state( "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -9022,7 +9037,7 @@ def test_doubly_nested_graph_state( "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -9076,7 +9091,7 @@ def test_doubly_nested_graph_state( AnyStr("child_1"), ], "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], + "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config={ @@ -9131,7 +9146,7 @@ def test_doubly_nested_graph_state( AnyStr("child_1"), ], "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], + "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config={ @@ -9193,7 +9208,7 @@ def test_doubly_nested_graph_state( AnyStr("child_1"), ], "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], + "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config={ @@ -9255,7 +9270,7 @@ def test_doubly_nested_graph_state( AnyStr("child_1"), ], "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], + "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config=None, diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index 157e0e080..9fe7f2ac3 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -2300,7 +2300,11 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 1, "langgraph_node": "agent", - "langgraph_triggers": ("start:agent",), + "langgraph_triggers": ( + "branch:to:agent", + "start:agent", + "tools", + ), "langgraph_path": ("__pregel_pull", "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2359,7 +2363,11 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 3, "langgraph_node": "agent", - "langgraph_triggers": ("tools",), + "langgraph_triggers": ( + "branch:to:agent", + "start:agent", + "tools", + ), "langgraph_path": ("__pregel_pull", "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2402,7 +2410,11 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 5, "langgraph_node": "agent", - "langgraph_triggers": ("tools",), + "langgraph_triggers": ( + "branch:to:agent", + "start:agent", + "tools", + ), "langgraph_path": ("__pregel_pull", "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -3883,7 +3895,10 @@ 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": ( + "branch:to:rewrite_query", + "start:rewrite_query", + ), }, }, ), @@ -3914,7 +3929,10 @@ 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": ( + "branch:to:retriever_one", + "rewrite_query", + ), }, }, ), @@ -3928,7 +3946,10 @@ 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": ( + "branch:to:retriever_two", + "rewrite_query", + ), }, }, ), @@ -3990,7 +4011,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": (AnyStr("retriever_"),), + "triggers": ("branch:to:qa", "retriever_one", "retriever_two"), }, }, ), @@ -4465,7 +4486,10 @@ async def test_branch_then(checkpointer_name: str) -> None: "id": AnyStr(), "name": "prepare", "input": {"my_key": "value", "market": "DE"}, - "triggers": ("start:prepare",), + "triggers": ( + "branch:to:prepare", + "start:prepare", + ), }, }, { @@ -4609,7 +4633,10 @@ 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", + "branch:to:finish", + ), }, }, { @@ -4778,7 +4805,10 @@ async def test_branch_then(checkpointer_name: str) -> None: "id": AnyStr(), "name": "prepare", "input": {"my_key": "value", "market": "DE"}, - "triggers": ("start:prepare",), + "triggers": ( + "branch:to:prepare", + "start:prepare", + ), }, }, { @@ -5333,7 +5363,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], + "langgraph_triggers": ["branch:to:inner", "outer_1"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -5530,7 +5560,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], + "langgraph_triggers": ["branch:to:inner", "outer_1"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -5573,7 +5603,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], + "langgraph_triggers": ["branch:to:inner", "outer_1"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -5622,7 +5652,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], + "langgraph_triggers": ["branch:to:inner", "outer_1"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -6060,7 +6090,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "child_1", "langgraph_path": [PULL, AnyStr("child_1")], "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], + "langgraph_triggers": ["branch:to:child_1", "start:child_1"], }, created_at=AnyStr(), parent_config=( @@ -6146,7 +6176,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: AnyStr("child_1"), ], "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], + "langgraph_triggers": [ + "branch:to:child_1", + "start:child_1", + ], }, created_at=AnyStr(), parent_config=( @@ -6195,7 +6228,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_triggers": [ + "branch:to:child", + AnyStr("parent_1"), + ], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -6493,7 +6529,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -6532,7 +6568,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -6584,7 +6620,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -6642,7 +6678,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: AnyStr("child_1"), ], "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], + "langgraph_triggers": [ + "branch:to:child_1", + AnyStr("start:child_1"), + ], }, created_at=AnyStr(), parent_config={ @@ -6697,7 +6736,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: AnyStr("child_1"), ], "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], + "langgraph_triggers": [ + "branch:to:child_1", + AnyStr("start:child_1"), + ], }, created_at=AnyStr(), parent_config={ @@ -6759,7 +6801,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: AnyStr("child_1"), ], "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], + "langgraph_triggers": [ + "branch:to:child_1", + AnyStr("start:child_1"), + ], }, created_at=AnyStr(), parent_config={ @@ -6821,7 +6866,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: AnyStr("child_1"), ], "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], + "langgraph_triggers": [ + "branch:to:child_1", + AnyStr("start:child_1"), + ], }, created_at=AnyStr(), parent_config=None, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index d8d51a8a0..bc7ed98cd 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -69,6 +69,7 @@ from langgraph.types import ( Interrupt, PregelTask, Send, + StateUpdate, StreamWriter, interrupt, ) @@ -6925,7 +6926,10 @@ def test_tags_stream_mode_messages() -> None: { "langgraph_step": 1, "langgraph_node": "call_model", - "langgraph_triggers": ("start:call_model",), + "langgraph_triggers": ( + "branch:to:call_model", + "start:call_model", + ), "langgraph_path": ("__pregel_pull", "call_model"), "langgraph_checkpoint_ns": AnyStr("call_model:"), "checkpoint_ns": AnyStr("call_model:"), @@ -7613,3 +7617,286 @@ def test_parallel_interrupts_double( assert invokes == 5 assert len(events) == 5 + + +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) +def test_bulk_state_updates( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class State(TypedDict): + foo: str + baz: str + + def node_a(state: State) -> State: + return {"foo": "bar"} + + def node_b(state: State) -> State: + return {"baz": "qux"} + + graph = ( + StateGraph(State) + .add_node("node_a", node_a) + .add_node("node_b", node_b) + .add_edge(START, "node_a") + .add_edge("node_a", "node_b") + .compile(checkpointer=checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # First update with node_a + graph.bulk_update_state( + config, + [ + [ + StateUpdate(values={"foo": "bar"}, as_node="node_a"), + ] + ], + ) + + # Then bulk update with both nodes + graph.bulk_update_state( + config, + [ + [ + StateUpdate(values={"foo": "updated"}, as_node="node_a"), + StateUpdate(values={"baz": "new"}, as_node="node_b"), + ] + ], + ) + + state = graph.get_state(config) + assert state.values == {"foo": "updated", "baz": "new"} + + # Check if there are only two checkpoints + checkpoints = list(checkpointer.list(config)) + assert len(checkpoints) == 2 + assert checkpoints[0].metadata["writes"] == { + "node_a": {"foo": "updated"}, + "node_b": {"baz": "new"}, + } + assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}} + + # perform multiple steps at the same time + config = {"configurable": {"thread_id": "2"}} + + graph.bulk_update_state( + config, + [ + [ + StateUpdate(values={"foo": "bar"}, as_node="node_a"), + ], + [ + StateUpdate(values={"foo": "updated"}, as_node="node_a"), + StateUpdate(values={"baz": "new"}, as_node="node_b"), + ], + ], + ) + + state = graph.get_state(config) + assert state.values == {"foo": "updated", "baz": "new"} + + checkpoints = list(checkpointer.list(config)) + assert len(checkpoints) == 2 + assert checkpoints[0].metadata["writes"] == { + "node_a": {"foo": "updated"}, + "node_b": {"baz": "new"}, + } + assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}} + + # Should raise error if updating without as_node + with pytest.raises(InvalidUpdateError): + graph.bulk_update_state( + config, + [ + [ + StateUpdate(values={"foo": "error"}, as_node=None), + StateUpdate(values={"bar": "error"}, as_node=None), + ] + ], + ) + + # Should raise if no updates are provided + with pytest.raises(ValueError, match="No supersteps provided"): + graph.bulk_update_state(config, []) + + # Should raise if no updates are provided + with pytest.raises(ValueError, match="No updates provided"): + graph.bulk_update_state(config, [[], []]) + + # Should raise if __end__ or __copy__ update is applied in bulk + with pytest.raises(InvalidUpdateError): + graph.bulk_update_state( + config, + [ + [ + StateUpdate(values=None, as_node="__end__"), + StateUpdate(values=None, as_node="__copy__"), + ], + ], + ) + + +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) +def test_update_as_input( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class State(TypedDict): + foo: str + + def agent(state: State) -> State: + return {"foo": "agent"} + + def tool(state: State) -> State: + return {"foo": "tool"} + + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("tool", tool) + .add_edge(START, "agent") + .add_edge("agent", "tool") + .compile(checkpointer=checkpointer) + ) + + assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == { + "foo": "tool" + } + + assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == { + "foo": "tool" + } + + def map_snapshot(i: StateSnapshot) -> dict: + return { + "values": i.values, + "next": i.next, + "step": i.metadata.get("step"), + } + + history = [ + map_snapshot(s) + for s in graph.get_state_history({"configurable": {"thread_id": "1"}}) + ] + + graph.bulk_update_state( + {"configurable": {"thread_id": "2"}}, + [ + # First turn + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent"}, "agent")], + [StateUpdate({"foo": "tool"}, "tool")], + # Second turn + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent"}, "agent")], + [StateUpdate({"foo": "tool"}, "tool")], + ], + ) + + state = graph.get_state({"configurable": {"thread_id": "2"}}) + assert state.values == {"foo": "tool"} + + new_history = [ + map_snapshot(s) + for s in graph.get_state_history({"configurable": {"thread_id": "2"}}) + ] + + assert new_history == history + + +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) +def test_batch_update_as_input( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class State(TypedDict): + foo: str + tasks: Annotated[list[int], operator.add] + + def agent(state: State) -> State: + return {"foo": "agent"} + + def map(state: State) -> Command["task"]: + return Command( + goto=[ + Send("task", {"index": 0}), + Send("task", {"index": 1}), + Send("task", {"index": 2}), + ], + update={"foo": "map"}, + ) + + def task(state: dict) -> State: + return {"tasks": [state["index"]]} + + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("map", map) + .add_node("task", task) + .add_edge(START, "agent") + .add_edge("agent", "map") + .compile(checkpointer=checkpointer) + ) + + assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == { + "foo": "map", + "tasks": [0, 1, 2], + } + + def map_snapshot(i: StateSnapshot) -> dict: + return { + "values": i.values, + "next": i.next, + "step": i.metadata.get("step"), + "tasks": [t.name for t in i.tasks], + } + + history = [ + map_snapshot(s) + for s in graph.get_state_history({"configurable": {"thread_id": "1"}}) + ] + + graph.bulk_update_state( + {"configurable": {"thread_id": "2"}}, + [ + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent", "tasks": []}, "agent")], + [ + StateUpdate( + Command( + goto=[ + Send("task", {"index": 0}), + Send("task", {"index": 1}), + Send("task", {"index": 2}), + ], + update={"foo": "map"}, + ), + "map", + ) + ], + [ + StateUpdate({"tasks": [0]}, "task"), + StateUpdate({"tasks": [1]}, "task"), + StateUpdate({"tasks": [2]}, "task"), + ], + ], + ) + + state = graph.get_state({"configurable": {"thread_id": "2"}}) + assert state.values == {"foo": "map", "tasks": [0, 1, 2]} + + new_history = [ + map_snapshot(s) + for s in graph.get_state_history({"configurable": {"thread_id": "2"}}) + ] + + assert new_history == history diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 60cbd2547..d35f34c11 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -59,6 +59,7 @@ from langgraph.types import ( Interrupt, PregelTask, Send, + StateUpdate, StreamWriter, interrupt, ) @@ -7582,7 +7583,10 @@ async def test_tags_stream_mode_messages() -> None: { "langgraph_step": 1, "langgraph_node": "call_model", - "langgraph_triggers": ("start:call_model",), + "langgraph_triggers": ( + "branch:to:call_model", + "start:call_model", + ), "langgraph_path": ("__pregel_pull", "call_model"), "langgraph_checkpoint_ns": AnyStr("call_model:"), "checkpoint_ns": AnyStr("call_model:"), @@ -7838,3 +7842,291 @@ async def test_handles_multiple_interrupts_from_tasks() -> None: assert len(result) == 2 assert result[0] == "Added James!" assert result[1] == "Added Will!" + + +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) +async def test_bulk_state_updates(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + + class State(TypedDict): + foo: str + baz: str + + def node_a(state: State) -> State: + return {"foo": "bar"} + + def node_b(state: State) -> State: + return {"baz": "qux"} + + graph = ( + StateGraph(State) + .add_node("node_a", node_a) + .add_node("node_b", node_b) + .add_edge(START, "node_a") + .add_edge("node_a", "node_b") + .compile(checkpointer=checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # First update with node_a + await graph.abulk_update_state( + config, + [ + [ + StateUpdate({"foo": "bar"}, "node_a"), + ] + ], + ) + + # Then bulk update with both nodes + await graph.abulk_update_state( + config, + [ + [ + StateUpdate({"foo": "updated"}, "node_a"), + StateUpdate({"baz": "new"}, "node_b"), + ] + ], + ) + + state = await graph.aget_state(config) + assert state.values == {"foo": "updated", "baz": "new"} + + # Check if there are only two checkpoints + checkpoints = [ + c async for c in checkpointer.alist({"configurable": {"thread_id": "1"}}) + ] + assert len(checkpoints) == 2 + assert checkpoints[0].metadata["writes"] == { + "node_a": {"foo": "updated"}, + "node_b": {"baz": "new"}, + } + assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}} + + # perform multiple steps at the same time + config = {"configurable": {"thread_id": "2"}} + + await graph.abulk_update_state( + config, + [ + [ + StateUpdate({"foo": "bar"}, "node_a"), + ], + [ + StateUpdate({"foo": "updated"}, "node_a"), + StateUpdate({"baz": "new"}, "node_b"), + ], + ], + ) + + state = await graph.aget_state(config) + assert state.values == {"foo": "updated", "baz": "new"} + + checkpoints = [ + c async for c in checkpointer.alist({"configurable": {"thread_id": "1"}}) + ] + assert len(checkpoints) == 2 + assert checkpoints[0].metadata["writes"] == { + "node_a": {"foo": "updated"}, + "node_b": {"baz": "new"}, + } + assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}} + + # Should raise error if updating without as_node + with pytest.raises(InvalidUpdateError): + await graph.abulk_update_state( + config, + [ + [ + StateUpdate(values={"foo": "error"}, as_node=None), + StateUpdate(values={"bar": "error"}, as_node=None), + ] + ], + ) + + # Should raise if no updates are provided + with pytest.raises(ValueError, match="No supersteps provided"): + await graph.abulk_update_state(config, []) + + # Should raise if no updates are provided + with pytest.raises(ValueError, match="No updates provided"): + await graph.abulk_update_state(config, [[], []]) + + # Should raise if __end__ or __copy__ update is applied in bulk + with pytest.raises(InvalidUpdateError): + await graph.abulk_update_state( + config, + [ + [ + StateUpdate(values=None, as_node="__end__"), + StateUpdate(values=None, as_node="__copy__"), + ], + ], + ) + + +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) +async def test_update_as_input(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + + class State(TypedDict): + foo: str + + def agent(state: State) -> State: + return {"foo": "agent"} + + def tool(state: State) -> State: + return {"foo": "tool"} + + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("tool", tool) + .add_edge(START, "agent") + .add_edge("agent", "tool") + .compile(checkpointer=checkpointer) + ) + + assert await graph.ainvoke( + {"foo": "input"}, {"configurable": {"thread_id": "1"}} + ) == {"foo": "tool"} + + assert await graph.ainvoke( + {"foo": "input"}, {"configurable": {"thread_id": "1"}} + ) == {"foo": "tool"} + + def map_snapshot(i: StateSnapshot) -> dict: + return { + "values": i.values, + "next": i.next, + "step": i.metadata.get("step"), + } + + history = [ + map_snapshot(s) + async for s in graph.aget_state_history( + {"configurable": {"thread_id": "1"}} + ) + ] + + await graph.abulk_update_state( + {"configurable": {"thread_id": "2"}}, + [ + # First turn + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent"}, "agent")], + [StateUpdate({"foo": "tool"}, "tool")], + # Second turn + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent"}, "agent")], + [StateUpdate({"foo": "tool"}, "tool")], + ], + ) + + state = await graph.aget_state({"configurable": {"thread_id": "2"}}) + assert state.values == {"foo": "tool"} + + new_history = [ + map_snapshot(s) + async for s in graph.aget_state_history( + {"configurable": {"thread_id": "2"}} + ) + ] + + assert new_history == history + + +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) +async def test_batch_update_as_input(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + + class State(TypedDict): + foo: str + tasks: Annotated[list[int], operator.add] + + def agent(state: State) -> State: + return {"foo": "agent"} + + def map(state: State) -> Command["task"]: + return Command( + goto=[ + Send("task", {"index": 0}), + Send("task", {"index": 1}), + Send("task", {"index": 2}), + ], + update={"foo": "map"}, + ) + + def task(state: dict) -> State: + return {"tasks": [state["index"]]} + + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("map", map) + .add_node("task", task) + .add_edge(START, "agent") + .add_edge("agent", "map") + .compile(checkpointer=checkpointer) + ) + + assert await graph.ainvoke( + {"foo": "input"}, {"configurable": {"thread_id": "1"}} + ) == {"foo": "map", "tasks": [0, 1, 2]} + + def map_snapshot(i: StateSnapshot) -> dict: + return { + "values": i.values, + "next": i.next, + "step": i.metadata.get("step"), + "tasks": [t.name for t in i.tasks], + } + + history = [ + map_snapshot(s) + async for s in graph.aget_state_history( + {"configurable": {"thread_id": "1"}} + ) + ] + + await graph.abulk_update_state( + {"configurable": {"thread_id": "2"}}, + [ + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent", "tasks": []}, "agent")], + [ + StateUpdate( + Command( + goto=[ + Send("task", {"index": 0}), + Send("task", {"index": 1}), + Send("task", {"index": 2}), + ], + update={"foo": "map"}, + ), + "map", + ) + ], + [ + StateUpdate({"tasks": [0]}, "task"), + StateUpdate({"tasks": [1]}, "task"), + StateUpdate({"tasks": [2]}, "task"), + ], + ], + ) + + state = await graph.aget_state({"configurable": {"thread_id": "2"}}) + assert state.values == {"foo": "map", "tasks": [0, 1, 2]} + + new_history = [ + map_snapshot(s) + async for s in graph.aget_state_history( + {"configurable": {"thread_id": "2"}} + ) + ] + + assert new_history == history diff --git a/libs/sdk-js/package.json b/libs/sdk-js/package.json index 2200fab34..6c3d5a06c 100644 --- a/libs/sdk-js/package.json +++ b/libs/sdk-js/package.json @@ -1,6 +1,6 @@ { "name": "@langchain/langgraph-sdk", - "version": "0.0.57", + "version": "0.0.59", "description": "Client library for interacting with the LangGraph API", "type": "module", "packageManager": "yarn@1.22.19", diff --git a/libs/sdk-js/src/client.ts b/libs/sdk-js/src/client.ts index 84ef3d44c..52051f1e8 100644 --- a/libs/sdk-js/src/client.ts +++ b/libs/sdk-js/src/client.ts @@ -30,6 +30,7 @@ import type { StreamEvent, CronsCreatePayload, OnConflictBehavior, + Command, } from "./types.js"; import { mergeSignals } from "./utils/signals.js"; import { getEnvironmentVariable } from "./utils/env.js"; @@ -638,6 +639,40 @@ export class ThreadsClient< ); } + /** + * Create a new thread from a batch states. + */ + async bulkUpdateState( + supersteps: Array<{ + updates: Array<{ values: unknown; command?: Command; asNode: string }>; + }>, + options?: { + graphId?: string; + threadId?: string; + metadata?: Metadata; + ifExists?: OnConflictBehavior; + }, + ): Promise> { + return this.fetch>("/threads/state/batch", { + method: "POST", + json: { + supersteps: supersteps.map((s) => ({ + updates: s.updates.map((u) => ({ + values: u.values, + command: u.command, + as_node: u.asNode, + })), + })), + thread_id: options?.threadId, + metadata: { + ...options?.metadata, + graph_id: options?.graphId, + }, + if_exists: options?.ifExists, + }, + }); + } + /** * Patch the metadata of a thread. *