diff --git a/README.md b/README.md index 5f5cd6acf..fcbf4dc1f 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ - - - LangGraph Logo + + + LangGraph Logo
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/docs/docs/concepts/langgraph_platform.md b/docs/docs/concepts/langgraph_platform.md index d9908aa8e..46dcab247 100644 --- a/docs/docs/concepts/langgraph_platform.md +++ b/docs/docs/concepts/langgraph_platform.md @@ -1,3 +1,8 @@ +--- +search: + boost: 2 +--- + # LangGraph Platform ## Overview diff --git a/docs/docs/concepts/v0-human-in-the-loop.md b/docs/docs/concepts/v0-human-in-the-loop.md index ad2f19aa4..801d94be8 100644 --- a/docs/docs/concepts/v0-human-in-the-loop.md +++ b/docs/docs/concepts/v0-human-in-the-loop.md @@ -1,3 +1,8 @@ +--- +search: + exclude: true +--- + # Human-in-the-loop !!! note "Use the `interrupt` function instead." diff --git a/docs/docs/how-tos/state-model.ipynb b/docs/docs/how-tos/state-model.ipynb index c4fb12041..2fadd7994 100644 --- a/docs/docs/how-tos/state-model.ipynb +++ b/docs/docs/how-tos/state-model.ipynb @@ -463,12 +463,12 @@ "source": [ "from langgraph.graph import StateGraph, START, END\n", "from pydantic import BaseModel\n", - "from langchain_core.messages import HumanMessage, AIMessage, BaseMessage\n", + "from langchain_core.messages import HumanMessage, AIMessage, AnyMessage\n", "from typing import List\n", "\n", "\n", "class ChatState(BaseModel):\n", - " messages: List[BaseMessage]\n", + " messages: List[AnyMessage]\n", " context: str\n", "\n", "\n", diff --git a/docs/docs/index.md b/docs/docs/index.md index 066f0775c..d4d038d6f 100644 --- a/docs/docs/index.md +++ b/docs/docs/index.md @@ -20,7 +20,7 @@ title: Home

diff --git a/docs/docs/llms-txt-overview.md b/docs/docs/llms-txt-overview.md new file mode 100644 index 000000000..246106619 --- /dev/null +++ b/docs/docs/llms-txt-overview.md @@ -0,0 +1,36 @@ +# LLMs-txt for LangGraph + +## Overview + +LangGraph provides documentation files in the [`llms.txt`](https://llmstxt.org/) format, specifically `llms.txt` and `llms-full.txt`. These files allow large language models (LLMs) and agents to access programming documentation and APIs, particularly useful within integrated development environments (IDEs). + +| Language Version | llms.txt | llms-full.txt | +|------------------|------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------| +| LangGraph Python | [https://langchain-ai.github.io/langgraph/llms.txt](https://langchain-ai.github.io/langgraph/llms.txt) | [https://langchain-ai.github.io/langgraph/llms-full.txt](https://langchain-ai.github.io/langgraph/llms-full.txt) | +| LangGraph JS | [https://langchain-ai.github.io/langgraphjs/llms.txt](https://langchain-ai.github.io/langgraphjs/llms.txt) | [https://langchain-ai.github.io/langgraphjs/llms-full.txt](https://langchain-ai.github.io/langgraphjs/llms-full.txt) | + +## Differences Between `llms.txt` and `llms-full.txt` + +- **`llms.txt`** is an index file containing links with brief descriptions of the content. An LLM or agent must follow these links to access detailed information. + +- **`llms-full.txt`** includes all the detailed content directly in a single file, eliminating the need for additional navigation. + +A key consideration when using `llms-full.txt` is its size. For extensive documentation, this file may become too large to fit into an LLM's context window. + +## Using `llms.txt` via an MCP Server + +As of March 9, 2025, IDEs [do not yet have robust native support for `llms.txt`](https://x.com/jeremyphoward/status/1902109312216129905?t=1eHFv2vdNdAckajnug0_Vw&s=19). However, you can utilize `llms.txt` effectively through an MCP server. + +We provide an MCP server specifically designed to serve documentation, called [`mcpdoc`](https://github.com/langchain-ai/mcpdoc). This setup is compatible with IDEs and platforms such as Cursor, Windsurf, Claude, and Claude Code. Instructions for using `mcpdoc` with these tools are available in the repository. + +## Using `llms-full.txt` + +The LangGraph `llms-full.txt` file typically contains several hundred thousand tokens, exceeding the context window limitations of most LLMs. To effectively use this file: + +1. **With IDEs (e.g., Cursor, Windsurf)**: + - Add the `llms-full.txt` as custom documentation. The IDE will automatically chunk and index the content, implementing Retrieval-Augmented Generation (RAG). + +2. **Without IDE support**: + - Use a chat model with a large context window. + - Implement a RAG strategy to manage and query the documentation efficiently. + diff --git a/docs/docs/tutorials/deployment.md b/docs/docs/tutorials/deployment.md index 29470fe63..f4d9dc44a 100644 --- a/docs/docs/tutorials/deployment.md +++ b/docs/docs/tutorials/deployment.md @@ -1,3 +1,8 @@ +--- +search: + boost: 2 +--- + # Deployment Get started deploying your LangGraph applications locally or on the cloud with diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 2189d5dc8..697d6bdc4 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -54,7 +54,7 @@ theme: code: "Roboto Mono" plugins: - search: - separator: '[\s\u200b\-_,:!=\[\]()"`/]+|\.(?!\d)|&[lg]t;|(?!\b)(?=[A-Z][a-z])' + separator: '[\s\u200b\-,:!=\[\]()"`/]+|\.(?!\d)|&[lg]t;' - autorefs - mkdocstrings: handlers: @@ -361,6 +361,7 @@ nav: # NOTE: prebuilt.md is auto-generated by `make build-prebuilt` - Prebuilt Agents: prebuilt.md - Companies using LangGraph: adopters.md + - LLMS-txt: llms-txt-overview.md - FAQ: concepts/faq.md - Troubleshooting: - Troubleshooting: troubleshooting/errors/index.md diff --git a/libs/checkpoint-postgres/poetry.lock b/libs/checkpoint-postgres/poetry.lock index 02f22bf76..d7aecaf05 100644 --- a/libs/checkpoint-postgres/poetry.lock +++ b/libs/checkpoint-postgres/poetry.lock @@ -397,7 +397,7 @@ typing-extensions = ">=4.7" [[package]] name = "langgraph-checkpoint" -version = "2.0.18" +version = "2.0.21" description = "Library with base interfaces for LangGraph checkpoint savers." optional = false python-versions = "^3.9.0,<4.0" @@ -1404,4 +1404,4 @@ cffi = ["cffi (>=1.11)"] [metadata] lock-version = "2.0" python-versions = "^3.9.0,<4.0" -content-hash = "369bfffecb9489835b43b8255932e043176a11d2f639aad2d055ffd89263ca1e" +content-hash = "4b0efdd115566f294fcd876334f9c3787aafc81f2689473759d88189a71d4635" diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml index 7c727163a..7c2b11c5a 100644 --- a/libs/checkpoint-postgres/pyproject.toml +++ b/libs/checkpoint-postgres/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint-postgres" -version = "2.0.18" +version = "2.0.19" description = "Library with a Postgres implementation of LangGraph checkpoint saver." authors = [] license = "MIT" @@ -10,7 +10,7 @@ packages = [{ include = "langgraph" }] [tool.poetry.dependencies] python = "^3.9.0,<4.0" -langgraph-checkpoint = "^2.0.15" +langgraph-checkpoint = "^2.0.21" orjson = ">=3.10.1" psycopg = "^3.2.0" psycopg-pool = "^3.2.0" diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 6eeb8867b..28e1696f1 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -7,7 +7,7 @@ from collections import defaultdict from collections.abc import AsyncIterator, Iterator, Sequence from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack from types import TracebackType -from typing import Any, Optional +from typing import Any, Optional, Union from langchain_core.runnables import RunnableConfig @@ -70,6 +70,12 @@ class InMemorySaver( tuple[str, str, str], dict[tuple[str, int], tuple[str, str, tuple[str, bytes], str]], ] + blobs: dict[ + tuple[ + str, str, str, Union[str, int, float] + ], # thread id, checkpoint ns, channel, version + tuple[str, bytes], + ] def __init__( self, @@ -80,6 +86,7 @@ class InMemorySaver( super().__init__(serde=serde) self.storage = factory(lambda: defaultdict(dict)) self.writes = factory(dict) + self.blobs = factory() self.stack = ExitStack() if factory is not defaultdict: self.stack.enter_context(self.storage) # type: ignore[arg-type] @@ -107,6 +114,18 @@ class InMemorySaver( ) -> Optional[bool]: return self.stack.__exit__(__exc_type, __exc_value, __traceback) + def _load_blobs( + self, thread_id: str, checkpoint_ns: str, versions: ChannelVersions + ) -> dict[str, Any]: + channel_values: dict[str, Any] = {} + for k, v in versions.items(): + kk = (thread_id, checkpoint_ns, k, v) + if kk in self.blobs: + vv = self.blobs[kk] + if vv[0] != "empty": + channel_values[k] = self.serde.loads_typed(vv) + return channel_values + def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the in-memory storage. @@ -121,8 +140,8 @@ class InMemorySaver( Returns: Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. """ - thread_id = config["configurable"]["thread_id"] - checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + thread_id: str = config["configurable"]["thread_id"] + checkpoint_ns: str = config["configurable"].get("checkpoint_ns", "") if checkpoint_id := get_checkpoint_id(config): if saved := self.storage[thread_id][checkpoint_ns].get(checkpoint_id): checkpoint, metadata, parent_checkpoint_id = saved @@ -140,10 +159,14 @@ class InMemorySaver( ) else: sends = [] + checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint) return CheckpointTuple( config=config, checkpoint={ - **self.serde.loads_typed(checkpoint), + **checkpoint_, + "channel_values": self._load_blobs( + thread_id, checkpoint_ns, checkpoint_["channel_versions"] + ), "pending_sends": [self.serde.loads_typed(s[2]) for s in sends], }, metadata=self.serde.loads_typed(metadata), @@ -180,6 +203,9 @@ class InMemorySaver( ) else: sends = [] + + checkpoint_ = self.serde.loads_typed(checkpoint) + return CheckpointTuple( config={ "configurable": { @@ -189,7 +215,10 @@ class InMemorySaver( } }, checkpoint={ - **self.serde.loads_typed(checkpoint), + **checkpoint_, + "channel_values": self._load_blobs( + thread_id, checkpoint_ns, checkpoint_["channel_versions"] + ), "pending_sends": [self.serde.loads_typed(s[2]) for s in sends], }, metadata=self.serde.loads_typed(metadata), @@ -297,6 +326,8 @@ class InMemorySaver( else: sends = [] + checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint) + yield CheckpointTuple( config={ "configurable": { @@ -306,7 +337,12 @@ class InMemorySaver( } }, checkpoint={ - **self.serde.loads_typed(checkpoint), + **checkpoint_, + "channel_values": self._load_blobs( + thread_id, + checkpoint_ns, + checkpoint_["channel_versions"], + ), "pending_sends": [ self.serde.loads_typed(s[2]) for s in sends ], @@ -353,6 +389,11 @@ class InMemorySaver( c.pop("pending_sends") # type: ignore[misc] thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"]["checkpoint_ns"] + values: dict[str, Any] = c.pop("channel_values") # type: ignore[misc] + for k, v in new_versions.items(): + self.blobs[(thread_id, checkpoint_ns, k, v)] = ( + self.serde.dumps_typed(values[k]) if k in values else ("empty", b"") + ) self.storage[thread_id][checkpoint_ns].update( { checkpoint["id"]: ( diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index 3c49219f9..ad2dbdb1e 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -68,7 +68,9 @@ class TestMemorySaver: }, "metadata": {"run_id": "my_run_id"}, } - self.memory_saver.put(config, self.chkpnt_2, self.metadata_2, {}) + self.memory_saver.put( + config, self.chkpnt_2, self.metadata_2, self.chkpnt_2["channel_versions"] + ) checkpoint = self.memory_saver.get_tuple(config) assert checkpoint is not None assert checkpoint.metadata == { @@ -80,9 +82,24 @@ class TestMemorySaver: async def test_search(self) -> None: # set up test # save checkpoints - self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {}) - self.memory_saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {}) - self.memory_saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {}) + self.memory_saver.put( + self.config_1, + self.chkpnt_1, + self.metadata_1, + self.chkpnt_1["channel_versions"], + ) + self.memory_saver.put( + self.config_2, + self.chkpnt_2, + self.metadata_2, + self.chkpnt_2["channel_versions"], + ) + self.memory_saver.put( + self.config_3, + self.chkpnt_3, + self.metadata_3, + self.chkpnt_3["channel_versions"], + ) # call method / assertions query_1 = {"source": "input"} # search by 1 key @@ -129,9 +146,24 @@ class TestMemorySaver: async def test_asearch(self) -> None: # set up test # save checkpoints - self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {}) - self.memory_saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {}) - self.memory_saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {}) + self.memory_saver.put( + self.config_1, + self.chkpnt_1, + self.metadata_1, + self.chkpnt_1["channel_versions"], + ) + self.memory_saver.put( + self.config_2, + self.chkpnt_2, + self.metadata_2, + self.chkpnt_2["channel_versions"], + ) + self.memory_saver.put( + self.config_3, + self.chkpnt_3, + self.metadata_3, + self.chkpnt_3["channel_versions"], + ) # call method / assertions query_1 = {"source": "input"} # search by 1 key diff --git a/libs/cli/README.md b/libs/cli/README.md index 595372ef7..69f9dec9b 100644 --- a/libs/cli/README.md +++ b/libs/cli/README.md @@ -79,7 +79,7 @@ The CLI uses a `langgraph.json` configuration file with these key settings: } ``` -See the [full documentation](https://langchain-ai.github.io/langgraph/docs/cloud/reference/cli.html) for detailed configuration options. +See the [full documentation](https://langchain-ai.github.io/langgraph/cloud/reference/cli/) for detailed configuration options. ## Development diff --git a/libs/cli/langgraph_cli/cli.py b/libs/cli/langgraph_cli/cli.py index f93a82077..e12e3e7f9 100644 --- a/libs/cli/langgraph_cli/cli.py +++ b/libs/cli/langgraph_cli/cli.py @@ -574,6 +574,12 @@ def dockerfile(save_path: str, config: pathlib.Path, add_docker_compose: bool) - help="Wait for a debugger client to connect to the debug port before starting the server", default=False, ) +@click.option( + "--studio_url", + type=str, + default=None, + help="URL of the LangGraph Studio instance to connect to. Defaults to https://smith.langchain.com", +) @cli.command( "dev", help="🏃‍♀️‍➡️ Run LangGraph API server in development mode with hot reloading and debugging support", @@ -588,6 +594,7 @@ def dev( no_browser: bool, debug_port: Optional[int], wait_for_client: bool, + studio_url: Optional[str], ): """CLI entrypoint for running the LangGraph API server.""" try: @@ -651,6 +658,7 @@ def dev( wait_for_client=wait_for_client, auth=config_json.get("auth"), http=config_json.get("http"), + studio_url=studio_url, ) diff --git a/libs/cli/poetry.lock b/libs/cli/poetry.lock index e710fb217..f18a5e33b 100644 --- a/libs/cli/poetry.lock +++ b/libs/cli/poetry.lock @@ -535,42 +535,42 @@ langgraph-sdk = ">=0.1.42,<0.2.0" [[package]] name = "langgraph-api" -version = "0.0.27" +version = "0.0.32" description = "" optional = true python-versions = "<4.0,>=3.11.0" files = [ - {file = "langgraph_api-0.0.27-py3-none-any.whl", hash = "sha256:9b21742238b15b8db9c2d3fd760a670332c8897d0bcbbd9d82e43b6ac15a7937"}, - {file = "langgraph_api-0.0.27.tar.gz", hash = "sha256:c21eb2b7fe3b93998379f7b13ad7d23b3ef06ab821b008c6b12b954acfb587ec"}, + {file = "langgraph_api-0.0.32-py3-none-any.whl", hash = "sha256:7990cedc65f784813aba867c5bde3fdfae3fa4588baef1aa346cbeac7c3aebf1"}, + {file = "langgraph_api-0.0.32.tar.gz", hash = "sha256:6f5b698ad8d136b73c2c53bcfa30670e9244a318b08b5e9cf00a707ea57c058c"}, ] [package.dependencies] cryptography = ">=43.0.3,<44.0.0" httpx = ">=0.25.0" -jsonschema-rs = ">=0.20.0,<0.21.0" +jsonschema-rs = ">=0.20.0,<0.30" langchain-core = ">=0.2.38,<0.4.0" langgraph = ">=0.2.56,<0.4.0" -langgraph-checkpoint = ">=2.0.15,<3.0" -langgraph-sdk = ">=0.1.53,<0.2.0" +langgraph-checkpoint = ">=2.0.21,<3.0" +langgraph-sdk = ">=0.1.58,<0.2.0" langsmith = ">=0.1.63,<0.4.0" orjson = ">=3.9.7" pyjwt = ">=2.9.0,<3.0.0" sse-starlette = ">=2.1.0,<2.2.0" starlette = ">=0.38.6" -structlog = ">=23.1.0,<24.0.0" +structlog = ">=24.1.0,<26" tenacity = ">=8.0.0" uvicorn = ">=0.26.0" watchfiles = ">=0.13" [[package]] name = "langgraph-checkpoint" -version = "2.0.16" +version = "2.0.21" description = "Library with base interfaces for LangGraph checkpoint savers." optional = true python-versions = "<4.0.0,>=3.9.0" files = [ - {file = "langgraph_checkpoint-2.0.16-py3-none-any.whl", hash = "sha256:dfab51076a6eddb5f9e146cfe1b977e3dd6419168b2afa23ff3f4e47973bf06f"}, - {file = "langgraph_checkpoint-2.0.16.tar.gz", hash = "sha256:49ba8cfa12b2aae845ccc3b1fbd1d7a8d3a6c4a2e387ab3a92fca40dd3d4baa5"}, + {file = "langgraph_checkpoint-2.0.21-py3-none-any.whl", hash = "sha256:ca89c2090cd9729f83f9782226935dc5ff9fe7756c24936f484ccb0ce367f87b"}, + {file = "langgraph_checkpoint-2.0.21.tar.gz", hash = "sha256:52beeb6dc1bd8c487b8315466cab271093b65eb97f54a0942dfe105cd20b237f"}, ] [package.dependencies] @@ -594,13 +594,13 @@ langgraph-checkpoint = ">=2.0.10,<3.0.0" [[package]] name = "langgraph-sdk" -version = "0.1.53" +version = "0.1.58" description = "SDK for interacting with LangGraph API" optional = true python-versions = "<4.0.0,>=3.9.0" files = [ - {file = "langgraph_sdk-0.1.53-py3-none-any.whl", hash = "sha256:4fab62caad73661ffe4c3ababedcd0d7bfaaba986bee4416b9c28948458a3af5"}, - {file = "langgraph_sdk-0.1.53.tar.gz", hash = "sha256:12906ed965905fa27e0c28d9fa07dc6fd89e6895ff321ff049fdf3965d057cc4"}, + {file = "langgraph_sdk-0.1.58-py3-none-any.whl", hash = "sha256:65f88cf5582da0c316714dc475126fa03c5f74d72bc0b9221dd42649de8e23d4"}, + {file = "langgraph_sdk-0.1.58.tar.gz", hash = "sha256:ef8b0e4c08af8c7efd3919497879c87a3627806b51e4ba5e8b06e0717e3d44cd"}, ] [package.dependencies] @@ -1357,18 +1357,18 @@ full = ["httpx (>=0.27.0,<0.29.0)", "itsdangerous", "jinja2", "python-multipart [[package]] name = "structlog" -version = "23.3.0" +version = "25.2.0" description = "Structured Logging for Python" optional = true python-versions = ">=3.8" files = [ - {file = "structlog-23.3.0-py3-none-any.whl", hash = "sha256:d6922a88ceabef5b13b9eda9c4043624924f60edbb00397f4d193bd754cde60a"}, - {file = "structlog-23.3.0.tar.gz", hash = "sha256:24b42b914ac6bc4a4e6f716e82ac70d7fb1e8c3b1035a765591953bfc37101a5"}, + {file = "structlog-25.2.0-py3-none-any.whl", hash = "sha256:0fecea2e345d5d491b72f3db2e5fcd6393abfc8cd06a4851f21fcd4d1a99f437"}, + {file = "structlog-25.2.0.tar.gz", hash = "sha256:d9f9776944207d1035b8b26072b9b140c63702fd7aa57c2f85d28ab701bd8e92"}, ] [package.extras] -dev = ["structlog[tests,typing]"] -docs = ["furo", "myst-parser", "sphinx", "sphinx-notfound-page", "sphinxcontrib-mermaid", "sphinxext-opengraph", "twisted"] +dev = ["freezegun (>=0.2.8)", "mypy (>=1.4)", "pretend", "pytest (>=6.0)", "pytest-asyncio (>=0.17)", "rich", "simplejson", "twisted"] +docs = ["cogapp", "furo", "myst-parser", "sphinx", "sphinx-notfound-page", "sphinxcontrib-mermaid", "sphinxext-opengraph", "twisted"] tests = ["freezegun (>=0.2.8)", "pretend", "pytest (>=6.0)", "pytest-asyncio (>=0.17)", "simplejson"] typing = ["mypy (>=1.4)", "rich", "twisted"] @@ -1717,4 +1717,4 @@ inmem = ["langgraph-api", "python-dotenv"] [metadata] lock-version = "2.0" python-versions = "^3.9.0,<4.0" -content-hash = "d0e2bdcb600ad031867413025fcc58bb162609209359d63ca99a77060cf8cbb4" +content-hash = "f5aa4d66f9c0b98b8321a70a82387dc6e5f3a3a7ecedd87ac00d6415199038f9" diff --git a/libs/cli/pyproject.toml b/libs/cli/pyproject.toml index aaeddd4a6..cbf0a7d14 100644 --- a/libs/cli/pyproject.toml +++ b/libs/cli/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-cli" -version = "0.1.77" +version = "0.1.78" description = "CLI for interacting with LangGraph API" authors = [] license = "MIT" @@ -14,7 +14,7 @@ langgraph = "langgraph_cli.cli:cli" [tool.poetry.dependencies] python = "^3.9.0,<4.0" click = "^8.1.7" -langgraph-api = { version = ">=0.0.27,<0.1.0", optional = true, python = ">=3.11,<4.0" } +langgraph-api = { version = ">=0.0.32,<0.1.0", optional = true, python = ">=3.11,<4.0" } python-dotenv = { version = ">=0.8.0", optional = true } [tool.poetry.group.dev.dependencies] diff --git a/libs/langgraph/README.md b/libs/langgraph/README.md index 5f5cd6acf..fcbf4dc1f 100644 --- a/libs/langgraph/README.md +++ b/libs/langgraph/README.md @@ -1,7 +1,7 @@ - - - LangGraph Logo + + + LangGraph Logo
diff --git a/libs/langgraph/bench/__main__.py b/libs/langgraph/bench/__main__.py index 61c8581f7..ad6c7e65d 100644 --- a/libs/langgraph/bench/__main__.py +++ b/libs/langgraph/bench/__main__.py @@ -420,4 +420,4 @@ compilation_benchmarks = ( ) for name, graph in compilation_benchmarks: - r.bench_func(name, "_compilation", compile_graph, graph) + r.bench_func(name + "_compilation", compile_graph, graph) diff --git a/libs/langgraph/langgraph/graph/branch.py b/libs/langgraph/langgraph/graph/branch.py index 8e90d9097..a1ae358d6 100644 --- a/libs/langgraph/langgraph/graph/branch.py +++ b/libs/langgraph/langgraph/graph/branch.py @@ -138,6 +138,7 @@ class Branch(NamedTuple): reader=reader, name=None, trace=False, + func_accepts_config=True, ) ) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 684e61c8e..0da072c94 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -859,8 +859,10 @@ class CompiledStateGraph(CompiledGraph): # subscribe to channel self.nodes[end].triggers.append(channel_name) # publish to channel - self.nodes[START] |= ChannelWrite( - [ChannelWriteEntry(channel_name, START)], tags=[TAG_HIDDEN] + self.nodes[START].writers.append( + ChannelWrite( + [ChannelWriteEntry(channel_name, START)], tags=[TAG_HIDDEN] + ) ) elif end != END: # subscribe to start channel @@ -873,8 +875,10 @@ class CompiledStateGraph(CompiledGraph): self.nodes[end].triggers.append(channel_name) # publish to channel for start in starts: - self.nodes[start] |= ChannelWrite( - [ChannelWriteEntry(channel_name, start)], tags=[TAG_HIDDEN] + self.nodes[start].writers.append( + ChannelWrite( + [ChannelWriteEntry(channel_name, start)], tags=[TAG_HIDDEN] + ) ) def attach_branch( @@ -910,28 +914,31 @@ class CompiledStateGraph(CompiledGraph): if start in self.builder.nodes else self.builder.schema ) - # attach branch publisher - self.nodes[start] |= branch.run( - branch_writer, - _get_state_reader(self.builder, schema) if with_reader else None, - ) - # attach branch subscribers - ends = ( - branch.ends.values() - if branch.ends - else [node for node in self.builder.nodes if node != branch.then] + # attach branch publisher + self.nodes[start].writers.append( + branch.run( + branch_writer, + _get_state_reader(self.builder, schema) if with_reader else None, + ) ) # attach then subscriber if branch.then and branch.then != END: + ends = ( + branch.ends.values() + if branch.ends + else [node for node in self.builder.nodes if node != branch.then] + ) channel_name = f"branch:{start}:{name}::then" self.channels[channel_name] = DynamicBarrierValue(str) self.nodes[branch.then].triggers.append(channel_name) for end in ends: if end != END: - self.nodes[end] |= ChannelWrite( - [ChannelWriteEntry(channel_name, end)], tags=[TAG_HIDDEN] + self.nodes[end].writers.append( + ChannelWrite( + [ChannelWriteEntry(channel_name, end)], tags=[TAG_HIDDEN] + ) ) @@ -1013,7 +1020,12 @@ async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]: CONTROL_BRANCH_PATH = RunnableCallable( - _control_branch, _acontrol_branch, tags=[TAG_HIDDEN], trace=False, recurse=False + _control_branch, + _acontrol_branch, + tags=[TAG_HIDDEN], + trace=False, + recurse=False, + func_accepts_config=False, ) CONTROL_BRANCH = Branch(CONTROL_BRANCH_PATH, None) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index f6479422c..45efe5c17 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -5,7 +5,7 @@ import concurrent import concurrent.futures import queue import weakref -from collections import deque +from collections import defaultdict, deque from functools import partial from typing import ( Any, @@ -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, ) @@ -503,6 +504,8 @@ class Pregel(PregelProtocol): name: str = "LangGraph" + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None + def __init__( self, *, @@ -524,6 +527,7 @@ class Pregel(PregelProtocol): config_type: Optional[Type[Any]] = None, input_model: Optional[Type[BaseModel]] = None, config: Optional[RunnableConfig] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, name: str = "LangGraph", ) -> None: self.nodes = nodes @@ -543,25 +547,28 @@ class Pregel(PregelProtocol): self.config_type = config_type self.input_model = input_model self.config = config + self.trigger_to_nodes = trigger_to_nodes self.name = name if auto_validate: self.validate() def get_graph( - self, config: RunnableConfig | None = None, *, xray: int | bool = False + self, config: Optional[RunnableConfig] = None, *, xray: Union[int, bool] = False ) -> Graph: raise NotImplementedError async def aget_graph( - self, config: RunnableConfig | None = None, *, xray: int | bool = False + self, config: Optional[RunnableConfig] = None, *, xray: Union[int, bool] = False ) -> Graph: raise NotImplementedError - def copy(self, update: dict[str, Any] | None = None) -> Self: + def copy(self, update: Optional[dict[str, Any]] = None) -> Self: attrs = {**self.__dict__, **(update or {})} return self.__class__(**attrs) - def with_config(self, config: RunnableConfig | None = None, **kwargs: Any) -> Self: + def with_config( + self, config: Optional[RunnableConfig] = None, **kwargs: Any + ) -> Self: return self.copy( {"config": merge_configs(self.config, config, cast(RunnableConfig, kwargs))} ) @@ -576,6 +583,7 @@ class Pregel(PregelProtocol): self.interrupt_after_nodes, self.interrupt_before_nodes, ) + self.trigger_to_nodes = _trigger_to_nodes(self.nodes) return self @property @@ -1162,22 +1170,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, "") @@ -1186,43 +1210,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, @@ -1249,182 +1454,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 + mv_writes, _ = apply_writes( + 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) @@ -1435,33 +1564,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, "") @@ -1470,46 +1623,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, @@ -1536,181 +1868,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 + mv_writes, _ = apply_writes( + 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) @@ -1722,18 +1976,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, @@ -1999,6 +2283,7 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after_, manager=run_manager, debug=debug, + trigger_to_nodes=self.trigger_to_nodes, ) as loop: # create runner runner = PregelRunner( @@ -2292,6 +2577,12 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after_, manager=run_manager, debug=debug, + # `self.nodes` can be modified after creation of `Pregel`. For example, + # that's how StateGraph compilation currently works. + # For now, we recompute the trigger_to_nodes mapping every time the + # loop is created. We could potentially memoize this if it becomes a + # performance issue. + trigger_to_nodes=_trigger_to_nodes(self.nodes), ) as loop: # create runner runner = PregelRunner( @@ -2460,3 +2751,12 @@ class Pregel(PregelProtocol): return latest else: return chunks + + +def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> Mapping[str, Sequence[str]]: + """Index from a trigger to nodes that depend on it.""" + trigger_to_nodes: defaultdict[str, list[str]] = defaultdict(list) + for name, node in nodes.items(): + for trigger in node.triggers: + trigger_to_nodes[trigger].append(name) + return dict(trigger_to_nodes) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 027f51ee7..19430f026 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -233,10 +233,21 @@ def apply_writes( channels: Mapping[str, BaseChannel], tasks: Iterable[WritesProtocol], get_next_version: Optional[GetNextVersion], -) -> dict[str, list[Any]]: +) -> tuple[dict[str, list[Any]], set[str]]: """Apply writes from a set of tasks (usually the tasks from a Pregel step) to the checkpoint and channels, and return managed values writes to be applied - externally.""" + externally. + + Args: + checkpoint: The checkpoint to update. + channels: The channels to update. + tasks: The tasks to apply writes from. + get_next_version: Optional function to determine the next version of a channel. + + Returns: + A tuple containing the managed values writes to be applied externally, and + the set of channels that were updated in this step. + """ # sort tasks on path, to ensure deterministic order for update application # any path parts after the 3rd are ignored for sorting # (we use them for eg. task ids which aren't good for sorting) @@ -312,15 +323,14 @@ def apply_writes( # Channels that weren't updated in this step are notified of a new step if bump_step: for chan in channels: - if chan not in updated_channels: - if channels[chan].update([]) and get_next_version is not None: + if channels[chan].is_available() and chan not in updated_channels: + if channels[chan].update(EMPTY_SEQ) and get_next_version is not None: checkpoint["channel_versions"][chan] = get_next_version( max_version, channels[chan], ) - # Return managed values writes to be applied externally - return pending_writes_by_managed + return pending_writes_by_managed, updated_channels @overload @@ -337,6 +347,8 @@ def prepare_next_tasks( store: Literal[None] = None, checkpointer: Literal[None] = None, manager: Literal[None] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, + updated_channels: Optional[set[str]] = None, ) -> dict[str, PregelTask]: ... @@ -354,6 +366,8 @@ def prepare_next_tasks( store: Optional[BaseStore], checkpointer: Optional[BaseCheckpointSaver], manager: Union[None, ParentRunManager, AsyncParentRunManager], + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, + updated_channels: Optional[set[str]] = None, ) -> dict[str, PregelExecutableTask]: ... @@ -370,10 +384,35 @@ def prepare_next_tasks( store: Optional[BaseStore] = None, checkpointer: Optional[BaseCheckpointSaver] = None, manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, + updated_channels: Optional[set[str]] = None, ) -> Union[dict[str, PregelTask], dict[str, PregelExecutableTask]]: """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).""" + + Args: + checkpoint: The current checkpoint. + pending_writes: The list of pending writes. + processes: The mapping of process names to PregelNode instances. + channels: The mapping of channel names to BaseChannel instances. + managed: The mapping of managed value names to functions. + config: The runnable configuration. + step: The current step. + for_execution: Whether the tasks are being prepared for execution. + store: An instance of BaseStore to make it available for usage within tasks. + checkpointer: Checkpointer instance used for saving checkpoints. + manager: The parent run manager to use for the tasks. + trigger_to_nodes: Optional: Mapping of channel names to the set of nodes + that are can be triggered by that channel. + updated_channels: Optional. Set of channel names that have been updated during + the previous step. Using in conjunction with trigger_to_nodes to speed + up the process of determining which nodes should be triggered in the next + step. + + Returns: + A dictionary of tasks to be executed. The keys are the task ids and the values + are the tasks themselves. 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]] = [] @@ -397,9 +436,30 @@ def prepare_next_tasks( manager=manager, ): tasks.append(task) + + # This section is an optimization that allows which nodes will be active + # during the next step. + # When there's information about: + # 1. Which channels were updated in the previous step + # 2. Which nodes are triggered by which channels + # Then we can determine which nodes should be triggered in the next step + # without having to cycle through all nodes. + if updated_channels and trigger_to_nodes: + triggered_nodes: set[str] = set() + # Get all nodes that have triggers associated with an updated channel + for channel in updated_channels: + if node_ids := trigger_to_nodes.get(channel): + triggered_nodes.update(node_ids) + # Sort the nodes to ensure deterministic order + candidate_nodes: Iterable[str] = sorted(triggered_nodes) + elif not checkpoint["channel_versions"]: + candidate_nodes = () + else: + candidate_nodes = processes.keys() + # Check if any processes should be run in next step # If so, prepare the values to be passed to them - for name in processes: + for name in candidate_nodes: if task := prepare_single_task( (PULL, name), None, @@ -517,7 +577,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, ), @@ -627,7 +687,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, ), @@ -655,13 +715,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) @@ -721,7 +782,7 @@ def prepare_single_task( CONFIG_KEY_SEND: partial( local_write, writes.extend, - processes.keys(), + tuple(processes.keys()), ), CONFIG_KEY_READ: partial( local_read, @@ -730,7 +791,10 @@ def prepare_single_task( channels, managed, PregelTaskWrites( - task_path[:3], name, writes, triggers + task_path[:3], + name, + writes, + triggers, ), config, ), @@ -748,7 +812,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, ), @@ -799,7 +863,7 @@ def _triggers( def _scratchpad( - config: RunnableConfig, + parent_scratchpad: Optional[PregelScratchpad], pending_writes: list[PendingWrite], task_id: str, ) -> PregelScratchpad: @@ -808,9 +872,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/loop.py b/libs/langgraph/langgraph/pregel/loop.py index edd69db01..3449ccde3 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -1,6 +1,7 @@ import asyncio import binascii import concurrent.futures +import dataclasses from collections import defaultdict, deque from contextlib import AsyncExitStack, ExitStack from inspect import signature @@ -209,6 +210,7 @@ class PregelLoop(LoopProtocol): manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: super().__init__( step=0, @@ -232,6 +234,7 @@ class PregelLoop(LoopProtocol): CONFIG_KEY_CHECKPOINT_ID not in config[CONF] or CONFIG_KEY_DEDUPE_TASKS in config[CONF] ) + self.trigger_to_nodes = trigger_to_nodes self.debug = debug if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) @@ -266,13 +269,13 @@ class PregelLoop(LoopProtocol): self.checkpoint_config = patch_configurable( self.config, { - CONFIG_KEY_CHECKPOINT_ID: config[CONF][CONFIG_KEY_CHECKPOINT_MAP][ - self.config[CONF][CONFIG_KEY_CHECKPOINT_NS] - ] + CONFIG_KEY_CHECKPOINT_ID: self.config[CONF][ + CONFIG_KEY_CHECKPOINT_MAP + ][self.config[CONF][CONFIG_KEY_CHECKPOINT_NS]] }, ) else: - self.checkpoint_config = config + self.checkpoint_config = self.config self.checkpoint_ns = ( tuple(cast(str, self.config[CONF][CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP)) if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS) @@ -406,8 +409,10 @@ class PregelLoop(LoopProtocol): if self.status != "pending": raise RuntimeError("Cannot tick when status is no longer 'pending'") + updated_channels: set[str] | None = None + if self.input not in (INPUT_DONE, INPUT_RESUMING, INPUT_SHOULD_VALIDATE): - self._first(input_keys=input_keys) + updated_channels = self._first(input_keys=input_keys) elif self.to_interrupt: # if we need to interrupt, do so self.status = "interrupt_before" @@ -427,7 +432,7 @@ class PregelLoop(LoopProtocol): ), ) # all tasks have finished - mv_writes = apply_writes( + mv_writes, updated_channels = apply_writes( self.checkpoint, self.channels, self.tasks.values(), @@ -493,6 +498,8 @@ class PregelLoop(LoopProtocol): manager=self.manager, store=self.store, checkpointer=self.checkpointer, + trigger_to_nodes=self.trigger_to_nodes, + updated_channels=updated_channels, ) self.to_interrupt = [] @@ -571,11 +578,11 @@ class PregelLoop(LoopProtocol): self.checkpoint["versions_seen"].get(INTERRUPT, {}).values(), default=None, ): - self.tasks[tid] = task._replace(scheduled=True) + self.tasks[tid] = dataclasses.replace(task, scheduled=True) else: task.writes.append((k, v)) - def _first(self, *, input_keys: Union[str, Sequence[str]]) -> None: + def _first(self, *, input_keys: Union[str, Sequence[str]]) -> Optional[set[str]]: # resuming from previous checkpoint requires # - finding a previous checkpoint # - receiving None input (outer graph) or RESUMING flag (subgraph) @@ -592,6 +599,8 @@ class PregelLoop(LoopProtocol): ), ) ) + # this can be set only when there are input_writes + updated_channels: Optional[set[str]] = None # map command to writes if isinstance(self.input, Command): @@ -612,7 +621,7 @@ class PregelLoop(LoopProtocol): if null_writes := [ w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID ]: - mv_writes = apply_writes( + mv_writes, _ = apply_writes( self.checkpoint, self.channels, [PregelTaskWrites((), INPUT, null_writes, [])], @@ -661,7 +670,7 @@ class PregelLoop(LoopProtocol): manager=None, ) # apply input writes - mv_writes = apply_writes( + mv_writes, updated_channels = apply_writes( self.checkpoint, self.channels, [ @@ -691,6 +700,7 @@ class PregelLoop(LoopProtocol): self.config = patch_configurable( self.config, {CONFIG_KEY_RESUMING: is_resuming} ) + return updated_channels def _put_checkpoint(self, metadata: CheckpointMetadata) -> None: for k, v in self.config["metadata"].items(): @@ -776,7 +786,7 @@ class PregelLoop(LoopProtocol): and self.checkpoint_pending_writes and any(task.writes for task in self.tasks.values()) ): - mv_writes = apply_writes( + mv_writes, _ = apply_writes( self.checkpoint, self.channels, self.tasks.values(), @@ -883,6 +893,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: super().__init__( input, @@ -899,6 +910,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): interrupt_before=interrupt_before, manager=manager, debug=debug, + trigger_to_nodes=trigger_to_nodes, ) self.stack = ExitStack() if checkpointer: @@ -1024,6 +1036,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: super().__init__( input, @@ -1040,6 +1053,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): interrupt_before=interrupt_before, manager=manager, debug=debug, + trigger_to_nodes=trigger_to_nodes, ) self.stack = AsyncExitStack() if checkpointer: 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/read.py b/libs/langgraph/langgraph/pregel/read.py index 728159156..05d0c6b60 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -62,7 +62,13 @@ class ChannelRead(RunnableCallable): mapper: Optional[Callable[[Any], Any]] = None, tags: Optional[list[str]] = None, ) -> None: - super().__init__(func=self._read, afunc=self._aread, tags=tags, name=None) + super().__init__( + func=self._read, + afunc=self._aread, + tags=tags, + name=None, + func_accepts_config=True, + ) self.fresh = fresh self.mapper = mapper self.channel = channel @@ -161,6 +167,7 @@ class PregelNode(Runnable): metadata: Optional[Mapping[str, Any]] = None, bound: Optional[Runnable[Any, Any]] = None, retry_policy: Optional[RetryPolicy] = None, + subgraphs: Optional[Sequence[PregelProtocol]] = None, ) -> None: self.channels = channels self.triggers = list(triggers) @@ -170,7 +177,9 @@ class PregelNode(Runnable): self.retry_policy = retry_policy self.tags = tags self.metadata = metadata - if self.bound is not DEFAULT_BOUND: + if subgraphs is not None: + self.subgraphs = subgraphs + elif self.bound is not DEFAULT_BOUND: try: subgraph = find_subgraph_pregel(self.bound) except Exception: @@ -184,7 +193,6 @@ class PregelNode(Runnable): def copy(self, update: dict[str, Any]) -> PregelNode: attrs = {**self.__dict__, **update} - attrs.pop("subgraphs") return PregelNode(**attrs) @cached_property 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/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 67f9c7b3e..0c312e145 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -149,7 +149,7 @@ class PregelRunner: configurable={ CONFIG_KEY_CALL: partial( _call, - t, + weakref.ref(t), retry=retry_policy, futures=weakref.ref(futures), schedule_task=self.schedule_task, @@ -185,7 +185,7 @@ class PregelRunner: configurable={ CONFIG_KEY_CALL: partial( _call, - t, + weakref.ref(t), retry=retry_policy, futures=weakref.ref(futures), schedule_task=self.schedule_task, @@ -263,7 +263,7 @@ class PregelRunner: configurable={ CONFIG_KEY_CALL: partial( _acall, - t, + weakref.ref(t), stream=self.use_astream, retry=retry_policy, futures=weakref.ref(futures), @@ -304,7 +304,7 @@ class PregelRunner: configurable={ CONFIG_KEY_CALL: partial( _acall, - t, + weakref.ref(t), retry=retry_policy, stream=self.use_astream, futures=weakref.ref(futures), @@ -469,7 +469,7 @@ def _panic_or_proceed( def _call( - task: PregelExecutableTask, + task: weakref.ref[PregelExecutableTask], func: Callable[[Any], Union[Awaitable[Any], Any]], input: Any, *, @@ -489,10 +489,10 @@ def _call( fut: Optional[concurrent.futures.Future] = None # schedule PUSH tasks, collect futures - scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] + scratchpad: PregelScratchpad = task().config[CONF][CONFIG_KEY_SCRATCHPAD] # type: ignore[union-attr] # schedule the next task, if the callback returns one if next_task := schedule_task()( # type: ignore[misc] - task, + task(), # type: ignore[arg-type] scratchpad.call_counter(), Call(func, input, retry=retry, callbacks=callbacks), ): @@ -528,7 +528,7 @@ def _call( configurable={ CONFIG_KEY_CALL: partial( _call, - next_task, + weakref.ref(next_task), futures=futures, retry=retry, callbacks=callbacks, @@ -550,7 +550,7 @@ def _call( def _acall( - task: PregelExecutableTask, + task: weakref.ref[PregelExecutableTask], func: Callable[[Any], Union[Awaitable[Any], Any]], input: Any, *, @@ -570,10 +570,10 @@ def _acall( ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]: fut: Optional[asyncio.Future] = None # schedule PUSH tasks, collect futures - scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] + scratchpad: PregelScratchpad = task().config[CONF][CONFIG_KEY_SCRATCHPAD] # type: ignore[union-attr] # schedule the next task, if the callback returns one if next_task := schedule_task()( # type: ignore[misc] - task, + task(), # type: ignore[arg-type] scratchpad.call_counter(), Call(func, input, retry=retry, callbacks=callbacks), ): @@ -614,7 +614,7 @@ def _acall( configurable={ CONFIG_KEY_CALL: partial( _acall, - next_task, + weakref.ref(next_task), stream=stream, futures=futures, schedule_task=schedule_task, @@ -623,7 +623,7 @@ def _acall( reraise=reraise, ), }, - __name__=task.name, + __name__=task().name, # type: ignore[union-attr] __cancel_on_exit__=True, __reraise_on_exit__=reraise, # starting a new task in the next tick ensures 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/pregel/write.py b/libs/langgraph/langgraph/pregel/write.py index 0744b4434..05fe5a388 100644 --- a/libs/langgraph/langgraph/pregel/write.py +++ b/libs/langgraph/langgraph/pregel/write.py @@ -57,7 +57,13 @@ class ChannelWrite(RunnableCallable): tags: Optional[Sequence[str]] = None, require_at_least_one_of: Optional[Sequence[str]] = None, # ignored ): - super().__init__(func=self._write, afunc=self._awrite, name=None, tags=tags) + super().__init__( + func=self._write, + afunc=self._awrite, + name=None, + tags=tags, + func_accepts_config=True, + ) self.writes = cast( list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]], writes ) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 4c8cf6da4..0a3e9f205 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 @@ -143,7 +148,14 @@ class PregelTask(NamedTuple): result: Optional[Any] = None -class PregelExecutableTask(NamedTuple): +if sys.version_info > (3, 11): + _T_DC_KWARGS = {"weakref_slot": True, "slots": True, "frozen": True} +else: + _T_DC_KWARGS = {"frozen": True} + + +@dataclasses.dataclass(**_T_DC_KWARGS) +class PregelExecutableTask: name: str input: Any proc: Runnable diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index cb8e2f447..f21420ad0 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -2,8 +2,8 @@ import asyncio import enum import inspect import sys -from contextlib import AsyncExitStack -from contextvars import copy_context +from contextlib import AsyncExitStack, contextmanager +from contextvars import Context, Token, copy_context from functools import partial, wraps from typing import ( Any, @@ -11,6 +11,7 @@ from typing import ( Awaitable, Callable, Coroutine, + Generator, Iterator, Optional, Protocol, @@ -53,13 +54,69 @@ from langgraph.utils.config import ( patch_config, ) -try: - from langchain_core.runnables.config import _set_config_context -except ImportError: - # For forwards compatibility - def _set_config_context(context: RunnableConfig) -> None: # type: ignore - """Set the context for the current thread.""" - var_child_runnable_config.set(context) + +def _set_config_context( + config: RunnableConfig, +) -> tuple[Token[Optional[RunnableConfig]], Optional[dict[str, Any]]]: + """Set the child Runnable config + tracing context. + + Args: + config (RunnableConfig): The config to set. + """ + from langchain_core.tracers.langchain import LangChainTracer + + config_token = var_child_runnable_config.set(config) + current_context = None + if ( + (callbacks := config.get("callbacks")) + and ( + parent_run_id := getattr(callbacks, "parent_run_id", None) + ) # Is callback manager + and ( + tracer := next( + ( + handler + for handler in getattr(callbacks, "handlers", []) + if isinstance(handler, LangChainTracer) + ), + None, + ) + ) + and (run := tracer.run_map.get(str(parent_run_id))) + ): + from langsmith.run_helpers import _set_tracing_context, get_tracing_context + + current_context = get_tracing_context() + _set_tracing_context({"parent": run}) + return config_token, current_context + + +@contextmanager +def set_config_context(config: RunnableConfig) -> Generator[Context, None, None]: + """Set the child Runnable config + tracing context. + + Args: + config (RunnableConfig): The config to set. + """ + from langsmith.run_helpers import _set_tracing_context + + ctx = copy_context() + config_token, _ = ctx.run(_set_config_context, config) + try: + yield ctx + finally: + ctx.run(var_child_runnable_config.reset, config_token) + ctx.run( + _set_tracing_context, + { + "parent": None, + "project_name": None, + "tags": None, + "metadata": None, + "enabled": None, + "client": None, + }, + ) # Before Python 3.11 native StrEnum is not available @@ -194,6 +251,7 @@ class RunnableCallable(Runnable): trace: bool = True, recurse: bool = True, explode_args: bool = False, + func_accepts_config: Optional[bool] = None, **kwargs: Any, ) -> None: self.name = name @@ -219,27 +277,32 @@ class RunnableCallable(Runnable): # check signature if func is None and afunc is None: raise ValueError("At least one of func or afunc must be provided.") - params = inspect.signature(cast(Callable, func or afunc)).parameters - self.func_accepts_config = "config" in params - # Mapping from kwarg name to (config key, default value) to be used. - # The default value is used if the config key is not found in the config. - self.func_accepts: dict[str, Tuple[str, Any]] = {} + if func_accepts_config is not None: + self.func_accepts_config = func_accepts_config + self.func_accepts: dict[str, Tuple[str, Any]] = {} + else: + params = inspect.signature(cast(Callable, func or afunc)).parameters - for kw, typ, config_key, default in KWARGS_CONFIG_KEYS: - p = params.get(kw) + self.func_accepts_config = "config" in params + # Mapping from kwarg name to (config key, default value) to be used. + # The default value is used if the config key is not found in the config. + self.func_accepts = {} - if p is None or p.kind not in VALID_KINDS: - # If parameter is not found or is not a valid kind, skip - continue + for kw, typ, config_key, default in KWARGS_CONFIG_KEYS: + p = params.get(kw) - if typ != (ANY_TYPE,) and p.annotation not in typ: - # A specific type is required, but the function annotation does - # not match the expected type. - continue + if p is None or p.kind not in VALID_KINDS: + # If parameter is not found or is not a valid kind, skip + continue - # If the kwarg is accepted by the function, store the default value - self.func_accepts[kw] = (config_key, default) + if typ != (ANY_TYPE,) and p.annotation not in typ: + # A specific type is required, but the function annotation does + # not match the expected type. + continue + + # If the kwarg is accepted by the function, store the default value + self.func_accepts[kw] = (config_key, default) def __repr__(self) -> str: repr_args = { @@ -286,7 +349,6 @@ class RunnableCallable(Runnable): kwargs[kw] = _conf.get(config_key, default_value) - context = copy_context() if self.trace: callback_manager = get_callback_manager_for_config(config, self.tags) run_manager = callback_manager.on_chain_start( @@ -297,17 +359,16 @@ class RunnableCallable(Runnable): ) try: child_config = patch_config(config, callbacks=run_manager.get_child()) - context = copy_context() - context.run(_set_config_context, child_config) - ret = context.run(self.func, *args, **kwargs) + with set_config_context(child_config) as context: + ret = context.run(self.func, *args, **kwargs) except BaseException as e: run_manager.on_chain_error(e) raise else: run_manager.on_chain_end(ret) else: - context.run(_set_config_context, config) - ret = context.run(self.func, *args, **kwargs) + with set_config_context(config) as context: + ret = context.run(self.func, *args, **kwargs) if isinstance(ret, Runnable) and self.recurse: return ret.invoke(input, config) return ret @@ -342,7 +403,6 @@ class RunnableCallable(Runnable): f"Missing required config key '{config_key}' for '{self.name}'." ) kwargs[kw] = _conf.get(config_key, default_value) - context = copy_context() if self.trace: callback_manager = get_async_callback_manager_for_config(config, self.tags) run_manager = await callback_manager.on_chain_start( @@ -353,24 +413,24 @@ class RunnableCallable(Runnable): ) try: child_config = patch_config(config, callbacks=run_manager.get_child()) - context.run(_set_config_context, child_config) - coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs)) - if ASYNCIO_ACCEPTS_CONTEXT: - ret = await asyncio.create_task(coro, context=context) - else: - ret = await coro + with set_config_context(child_config) as context: + coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs)) + if ASYNCIO_ACCEPTS_CONTEXT: + ret = await asyncio.create_task(coro, context=context) + else: + ret = await coro except BaseException as e: await run_manager.on_chain_error(e) raise else: await run_manager.on_chain_end(ret) else: - context.run(_set_config_context, config) - if ASYNCIO_ACCEPTS_CONTEXT: - coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs)) - ret = await asyncio.create_task(coro, context=context) - else: - ret = await self.afunc(*args, **kwargs) + with set_config_context(config) as context: + if ASYNCIO_ACCEPTS_CONTEXT: + coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs)) + ret = await asyncio.create_task(coro, context=context) + else: + ret = await self.afunc(*args, **kwargs) if isinstance(ret, Runnable) and self.recurse: return await ret.ainvoke(input, config) return ret diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index c79ef170f..c494889a7 100644 --- a/libs/langgraph/poetry.lock +++ b/libs/langgraph/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.0.0 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.0.1 and should not be changed by hand. [[package]] name = "aiosqlite" @@ -1324,14 +1324,14 @@ files = [ [[package]] name = "langchain-core" -version = "0.3.44" +version = "0.3.46" description = "Building applications with LLMs through composability" optional = false python-versions = "<4.0,>=3.9" groups = ["main", "dev"] files = [ - {file = "langchain_core-0.3.44-py3-none-any.whl", hash = "sha256:d989ce8bd62f1d07765acd575e6ec1254aec0cf7775aaea39fe4af8102377459"}, - {file = "langchain_core-0.3.44.tar.gz", hash = "sha256:7c0a01e78360f007cbca448178fe7e032404068e6431dbe8ce905f84febbdfa5"}, + {file = "langchain_core-0.3.46-py3-none-any.whl", hash = "sha256:28b5689fc347975ea520b5364ab4aee5567e661553bbee5e97cabf4596c28ce0"}, + {file = "langchain_core-0.3.46.tar.gz", hash = "sha256:5fca010eeb0a427be5aa8a8525e2112995dde790c584cef165be7c5e0ee1c2b5"}, ] [package.dependencies] @@ -1348,7 +1348,7 @@ typing-extensions = ">=4.7" [[package]] name = "langgraph-checkpoint" -version = "2.0.18" +version = "2.0.21" description = "Library with base interfaces for LangGraph checkpoint savers." optional = false python-versions = "^3.9.0,<4.0" @@ -1366,7 +1366,7 @@ url = "../checkpoint" [[package]] name = "langgraph-checkpoint-postgres" -version = "2.0.16" +version = "2.0.19" description = "Library with a Postgres implementation of LangGraph checkpoint saver." optional = false python-versions = "^3.9.0,<4.0" @@ -1375,7 +1375,7 @@ files = [] develop = true [package.dependencies] -langgraph-checkpoint = "^2.0.15" +langgraph-checkpoint = "^2.0.21" orjson = ">=3.10.1" psycopg = "^3.2.0" psycopg-pool = "^3.2.0" @@ -1404,7 +1404,7 @@ url = "../checkpoint-sqlite" [[package]] name = "langgraph-prebuilt" -version = "0.1.2" +version = "0.1.4" description = "Library with high-level APIs for creating and executing LangGraph agents and tools." optional = false python-versions = "^3.9.0,<4.0" @@ -1422,7 +1422,7 @@ url = "../prebuilt" [[package]] name = "langgraph-sdk" -version = "0.1.55" +version = "0.1.58" description = "SDK for interacting with LangGraph API" optional = false python-versions = "^3.9.0,<4.0" diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index c625f47bc..6b6f821db 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.18" description = "Building stateful, multi-actor applications with LLMs" authors = [] license = "MIT" diff --git a/libs/langgraph/tests/memory_assert.py b/libs/langgraph/tests/memory_assert.py index f88f0358f..3a1ef4536 100644 --- a/libs/langgraph/tests/memory_assert.py +++ b/libs/langgraph/tests/memory_assert.py @@ -1,4 +1,3 @@ -import asyncio import os import tempfile from collections import defaultdict @@ -13,7 +12,6 @@ from langgraph.checkpoint.base import ( CheckpointMetadata, CheckpointTuple, SerializerProtocol, - copy_checkpoint, ) from langgraph.checkpoint.memory import InMemorySaver, PersistentDict @@ -63,69 +61,14 @@ class MemorySaverAssertImmutable(InMemorySaver): self.storage_for_copies[thread_id][checkpoint_ns][saved["id"]] ) == saved - ) + ), config["configurable"]["checkpoint_ns"] self.storage_for_copies[thread_id][checkpoint_ns][checkpoint["id"]] = ( - self.serde.dumps_typed(copy_checkpoint(checkpoint)) + self.serde.dumps_typed(checkpoint) ) # call super to write checkpoint return super().put(config, checkpoint, metadata, new_versions) -class MemorySaverAssertCheckpointMetadata(InMemorySaver): - """This custom checkpointer is for verifying that a run's configurable - fields are merged with the previous checkpoint config for each step in - the run. This is the desired behavior. Because the checkpointer's (a)put() - method is called for each step, the implementation of this checkpointer - should produce a side effect that can be asserted. - """ - - def put( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> None: - """The implementation of put() merges config["configurable"] (a run's - configurable fields) with the metadata field. The state of the - checkpoint metadata can be asserted to confirm that the run's - configurable fields were merged with the previous checkpoint config. - """ - configurable = config["configurable"].copy() - - # remove checkpoint_id to make testing simpler - checkpoint_id = configurable.pop("checkpoint_id", None) - thread_id = config["configurable"]["thread_id"] - checkpoint_ns = config["configurable"]["checkpoint_ns"] - self.storage[thread_id][checkpoint_ns].update( - { - checkpoint["id"]: ( - self.serde.dumps_typed(checkpoint), - # merge configurable fields and metadata - self.serde.dumps_typed({**configurable, **metadata}), - checkpoint_id, - ) - } - ) - return { - "configurable": { - "thread_id": config["configurable"]["thread_id"], - "checkpoint_id": checkpoint["id"], - } - } - - async def aput( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - return await asyncio.get_running_loop().run_in_executor( - None, self.put, config, checkpoint, metadata, new_versions - ) - - class MemorySaverNoPending(InMemorySaver): def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: result = super().get_tuple(config) 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..9f547e940 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1,5 +1,6 @@ import enum import functools +import gc import json import logging import operator @@ -62,13 +63,16 @@ from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot +from langgraph.pregel.loop import SyncPregelLoop from langgraph.pregel.retry import RetryPolicy +from langgraph.pregel.runner import PregelRunner from langgraph.store.base import BaseStore from langgraph.types import ( Command, Interrupt, PregelTask, Send, + StateUpdate, StreamWriter, interrupt, ) @@ -80,7 +84,6 @@ from tests.conftest import ( REGULAR_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS, ) -from tests.memory_assert import MemorySaverAssertCheckpointMetadata from tests.messages import ( _AnyIdAIMessage, _AnyIdAIMessageChunk, @@ -4209,11 +4212,11 @@ def test_checkpoint_metadata() -> None: workflow.add_edge("tools", "agent") # graph w/o interrupt - checkpointer_1 = MemorySaverAssertCheckpointMetadata() + checkpointer_1 = InMemorySaver() app = workflow.compile(checkpointer=checkpointer_1) # graph w/ interrupt - checkpointer_2 = MemorySaverAssertCheckpointMetadata() + checkpointer_2 = InMemorySaver() app_w_interrupt = workflow.compile( checkpointer=checkpointer_2, interrupt_before=["tools"] ) @@ -4631,59 +4634,6 @@ def test_multiple_sinks_subgraphs(snapshot: SnapshotAssertion) -> None: assert app.get_graph(xray=True).draw_mermaid() == snapshot -def test_subgraph_retries(): - class State(TypedDict): - count: int - - class ChildState(State): - some_list: Annotated[list, operator.add] - - called_times = 0 - - class RandomError(ValueError): - """This will be retried on.""" - - def parent_node(state: State): - return {"count": state["count"] + 1} - - def child_node_a(state: ChildState): - nonlocal called_times - # We want it to retry only on node_b - # NOT re-compute the whole graph. - assert not called_times - called_times += 1 - return {"some_list": ["val"]} - - def child_node_b(state: ChildState): - raise RandomError("First attempt fails") - - child = StateGraph(ChildState) - child.add_node(child_node_a) - child.add_node(child_node_b) - child.add_edge("__start__", "child_node_a") - child.add_edge("child_node_a", "child_node_b") - - parent = StateGraph(State) - parent.add_node("parent_node", parent_node) - parent.add_node( - "child_graph", - child.compile(), - retry=RetryPolicy( - max_attempts=3, - retry_on=(RandomError,), - backoff_factor=0.0001, - initial_interval=0.0001, - ), - ) - parent.add_edge("parent_node", "child_graph") - parent.set_entry_point("parent_node") - - checkpointer = InMemorySaver() - app = parent.compile(checkpointer=checkpointer) - with pytest.raises(RandomError): - app.invoke({"count": 0}, {"configurable": {"thread_id": "foo"}}) - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) @pytest.mark.parametrize("store_name", ALL_STORES_SYNC) def test_store_injected( @@ -6290,6 +6240,7 @@ def test_double_interrupt_subgraph( def invoke_sub_agent(state: AgentState): return subgraph.invoke(state) + thread = {"configurable": {"thread_id": str(uuid.uuid4())}} parent_agent = ( StateGraph(AgentState) .add_node("invoke_sub_agent", invoke_sub_agent) @@ -6925,7 +6876,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 +7567,320 @@ def test_parallel_interrupts_double( assert invokes == 5 assert len(events) == 5 + + +def test_pregel_loop_refcount(): + gc.collect() + try: + gc.disable() + + class State(TypedDict): + messages: Annotated[list, add_messages] + + graph_builder = StateGraph(State) + + def chatbot(state: State): + return {"messages": [("ai", "HIYA")]} + + graph_builder.add_node("chatbot", chatbot) + graph_builder.set_entry_point("chatbot") + graph_builder.set_finish_point("chatbot") + graph = graph_builder.compile() + + for _ in range(5): + graph.invoke({"messages": [{"role": "user", "content": "hi"}]}) + assert ( + len( + [obj for obj in gc.get_objects() if isinstance(obj, SyncPregelLoop)] + ) + == 0 + ) + assert ( + len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)]) + == 0 + ) + finally: + gc.enable() + + +@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..77b5b19d3 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1,5 +1,7 @@ import asyncio +import enum import functools +import gc import logging import operator import random @@ -52,13 +54,16 @@ from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessagesState, add_messages from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot +from langgraph.pregel.loop import AsyncPregelLoop from langgraph.pregel.retry import RetryPolicy +from langgraph.pregel.runner import PregelRunner from langgraph.store.base import BaseStore from langgraph.types import ( Command, Interrupt, PregelTask, Send, + StateUpdate, StreamWriter, interrupt, ) @@ -73,10 +78,7 @@ from tests.conftest import ( awith_store, ) from tests.fake_tracer import FakeTracer -from tests.memory_assert import ( - MemorySaverAssertCheckpointMetadata, - MemorySaverNoPending, -) +from tests.memory_assert import MemorySaverNoPending from tests.messages import ( _AnyIdAIMessage, _AnyIdAIMessageChunk, @@ -4542,8 +4544,13 @@ async def test_nested_pydantic_models(version: str) -> None: name: str friends: list[str] = Field(default_factory=list) # IDs of friends + class MyEnum(enum.Enum): + A = 1 + B = 2 + class MyTypedDict(TypedDict): x: int + my_enum: MyEnum class State(BaseModel): # Basic nested model tests @@ -4552,6 +4559,7 @@ async def test_nested_pydantic_models(version: str) -> None: optional_nested: Optional[NestedModel] = None dict_nested: dict[str, NestedModel] my_set: set[int] + my_enum: MyEnum list_nested: Annotated[ Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y] ] @@ -4579,7 +4587,8 @@ async def test_nested_pydantic_models(version: str) -> None: "nested": {"value": 42, "name": "test"}, "optional_nested": {"value": 10, "name": "optional"}, "my_set": [1, 2, 7], - "my_typed_dict": {"x": 1}, + "my_enum": MyEnum.B, + "my_typed_dict": {"x": 1, "my_enum": MyEnum.A}, "dict_nested": {"a": {"value": 5, "name": "a"}}, "list_nested": [{"a": {"value": 6, "name": "b"}}], "list_nested_reversed": ["foo", "bar"], @@ -5758,11 +5767,11 @@ async def test_checkpoint_metadata() -> None: workflow.add_edge("tools", "agent") # graph w/o interrupt - checkpointer_1 = MemorySaverAssertCheckpointMetadata() + checkpointer_1 = InMemorySaver() app = workflow.compile(checkpointer=checkpointer_1) # graph w/ interrupt - checkpointer_2 = MemorySaverAssertCheckpointMetadata() + checkpointer_2 = InMemorySaver() app_w_interrupt = workflow.compile( checkpointer=checkpointer_2, interrupt_before=["tools"] ) @@ -7008,6 +7017,8 @@ async def test_double_interrupt_subgraph(checkpointer_name: str) -> None: def invoke_sub_agent(state: AgentState): return subgraph.invoke(state) + thread = {"configurable": {"thread_id": str(uuid.uuid4())}} + parent_agent = ( StateGraph(AgentState) .add_node("invoke_sub_agent", invoke_sub_agent) @@ -7582,7 +7593,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 +7852,329 @@ async def test_handles_multiple_interrupts_from_tasks() -> None: assert len(result) == 2 assert result[0] == "Added James!" assert result[1] == "Added Will!" + + +async def test_pregel_loop_refcount(): + gc.collect() + try: + gc.disable() + + class State(TypedDict): + messages: Annotated[list, add_messages] + + graph_builder = StateGraph(State) + + async def chatbot(state: State): + return {"messages": [("ai", "HIYA")]} + + graph_builder.add_node("chatbot", chatbot) + graph_builder.set_entry_point("chatbot") + graph_builder.set_finish_point("chatbot") + graph = graph_builder.compile() + + for _ in range(5): + await graph.ainvoke({"messages": [{"role": "user", "content": "hi"}]}) + assert ( + len( + [ + obj + for obj in gc.get_objects() + if isinstance(obj, AsyncPregelLoop) + ] + ) + == 0 + ) + assert ( + len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)]) + == 0 + ) + finally: + gc.enable() + + +@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/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index 2bff40fc4..0ca0d42b1 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -719,9 +719,7 @@ def create_react_agent( def generate_structured_response( state: StateSchema, config: RunnableConfig ) -> StateSchema: - # NOTE: we exclude the last message because there is enough information - # for the LLM to generate the structured response - messages = _get_state_value(state, "messages")[:-1] + messages = _get_state_value(state, "messages") structured_response_schema = response_format if isinstance(response_format, tuple): system_prompt, structured_response_schema = response_format @@ -736,9 +734,7 @@ def create_react_agent( async def agenerate_structured_response( state: StateSchema, config: RunnableConfig ) -> StateSchema: - # NOTE: we exclude the last message because there is enough information - # for the LLM to generate the structured response - messages = _get_state_value(state, "messages")[:-1] + messages = _get_state_value(state, "messages") structured_response_schema = response_format if isinstance(response_format, tuple): system_prompt, structured_response_schema = response_format diff --git a/libs/prebuilt/poetry.lock b/libs/prebuilt/poetry.lock index 3000e6ae7..99eb1c57a 100644 --- a/libs/prebuilt/poetry.lock +++ b/libs/prebuilt/poetry.lock @@ -435,7 +435,7 @@ typing-extensions = ">=4.7" [[package]] name = "langgraph" -version = "0.3.0" +version = "0.3.18" description = "Building stateful, multi-actor applications with LLMs" optional = false python-versions = ">=3.9.0,<4.0" @@ -446,6 +446,7 @@ develop = true [package.dependencies] langchain-core = ">=0.1,<0.4" langgraph-checkpoint = "^2.0.10" +langgraph-prebuilt = ">=0.1.1,<0.2" langgraph-sdk = "^0.1.42" [package.source] @@ -454,7 +455,7 @@ url = "../langgraph" [[package]] name = "langgraph-checkpoint" -version = "2.0.16" +version = "2.0.21" description = "Library with base interfaces for LangGraph checkpoint savers." optional = false python-versions = "^3.9.0,<4.0" @@ -472,7 +473,7 @@ url = "../checkpoint" [[package]] name = "langgraph-checkpoint-postgres" -version = "2.0.15" +version = "2.0.19" description = "Library with a Postgres implementation of LangGraph checkpoint saver." optional = false python-versions = "^3.9.0,<4.0" @@ -481,7 +482,7 @@ files = [] develop = true [package.dependencies] -langgraph-checkpoint = "^2.0.15" +langgraph-checkpoint = "^2.0.21" orjson = ">=3.10.1" psycopg = "^3.2.0" psycopg-pool = "^3.2.0" @@ -492,7 +493,7 @@ url = "../checkpoint-postgres" [[package]] name = "langgraph-checkpoint-sqlite" -version = "2.0.5" +version = "2.0.6" description = "Library with a SQLite implementation of LangGraph checkpoint saver." optional = false python-versions = "^3.9.0" diff --git a/libs/prebuilt/pyproject.toml b/libs/prebuilt/pyproject.toml index dbdf36bd6..3dd497c3b 100644 --- a/libs/prebuilt/pyproject.toml +++ b/libs/prebuilt/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-prebuilt" -version = "0.1.3" +version = "0.1.4" description = "Library with high-level APIs for creating and executing LangGraph agents and tools." authors = [] license = "MIT" diff --git a/libs/sdk-js/package.json b/libs/sdk-js/package.json index 2200fab34..6f1cd8d0b 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.60", "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..54acd5232 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"; @@ -481,15 +482,47 @@ export class ThreadsClient< * Metadata for the thread. */ metadata?: Metadata; + /** + * ID of the thread to create. + * + * If not provided, a random UUID will be generated. + */ threadId?: string; + /** + * How to handle duplicate creation. + * + * @default "raise" + */ ifExists?: OnConflictBehavior; + /** + * Graph ID to associate with the thread. + */ + graphId?: string; + /** + * Apply a list of supersteps when creating a thread, each containing a sequence of updates. + * + * Used for copying a thread between deployments. + */ + supersteps?: Array<{ + updates: Array<{ values: unknown; command?: Command; asNode: string }>; + }>; }): Promise> { return this.fetch>(`/threads`, { method: "POST", json: { - metadata: payload?.metadata, + metadata: { + ...payload?.metadata, + graph_id: payload?.graphId, + }, thread_id: payload?.threadId, if_exists: payload?.ifExists, + supersteps: payload?.supersteps?.map((s) => ({ + updates: s.updates.map((u) => ({ + values: u.values, + command: u.command, + as_node: u.asNode, + })), + })), }, }); } diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py index 6f4583b7f..4418e66de 100644 --- a/libs/sdk-py/langgraph_sdk/client.py +++ b/libs/sdk-py/langgraph_sdk/client.py @@ -839,6 +839,8 @@ class ThreadsClient: metadata: Json = None, thread_id: Optional[str] = None, if_exists: Optional[OnConflictBehavior] = None, + supersteps: Optional[Sequence[dict[str, Sequence[dict[str, Any]]]]] = None, + graph_id: Optional[str] = None, ) -> Thread: """Create a new thread. @@ -848,6 +850,9 @@ class ThreadsClient: If None, ID will be a randomly generated UUID. if_exists: How to handle duplicate creation. Defaults to 'raise' under the hood. Must be either 'raise' (raise error if duplicate), or 'do_nothing' (return existing thread). + supersteps: Apply a list of supersteps when creating a thread, each containing a sequence of updates. + Each update has `values` or `command` and `as_node`. Used for copying a thread between deployments. + graph_id: Optional graph ID to associate with the thread. Returns: Thread: The created thread. @@ -863,10 +868,28 @@ class ThreadsClient: payload: Dict[str, Any] = {} if thread_id: payload["thread_id"] = thread_id - if metadata: - payload["metadata"] = metadata + if metadata or graph_id: + payload["metadata"] = { + **(metadata or {}), + **({"graph_id": graph_id} if graph_id else {}), + } if if_exists: payload["if_exists"] = if_exists + if supersteps: + payload["supersteps"] = [ + { + "updates": [ + { + "values": u["values"], + "command": u.get("command"), + "as_node": u["as_node"], + } + for u in s["updates"] + ] + } + for s in supersteps + ] + return await self.http.post("/threads", json=payload) async def update(self, thread_id: str, *, metadata: dict[str, Any]) -> Thread: @@ -3036,6 +3059,8 @@ class SyncThreadsClient: metadata: Json = None, thread_id: Optional[str] = None, if_exists: Optional[OnConflictBehavior] = None, + supersteps: Optional[Sequence[dict[str, Sequence[dict[str, Any]]]]] = None, + graph_id: Optional[str] = None, ) -> Thread: """Create a new thread. @@ -3045,6 +3070,9 @@ class SyncThreadsClient: If None, ID will be a randomly generated UUID. if_exists: How to handle duplicate creation. Defaults to 'raise' under the hood. Must be either 'raise' (raise error if duplicate), or 'do_nothing' (return existing thread). + supersteps: Apply a list of supersteps when creating a thread, each containing a sequence of updates. + Each update has `values` or `command` and `as_node`. Used for copying a thread between deployments. + graph_id: Optional graph ID to associate with the thread. Returns: Thread: The created thread. @@ -3060,10 +3088,28 @@ class SyncThreadsClient: payload: Dict[str, Any] = {} if thread_id: payload["thread_id"] = thread_id - if metadata: - payload["metadata"] = metadata + if metadata or graph_id: + payload["metadata"] = { + **(metadata or {}), + **({"graph_id": graph_id} if graph_id else {}), + } if if_exists: payload["if_exists"] = if_exists + if supersteps: + payload["supersteps"] = [ + { + "updates": [ + { + "values": u["values"], + "command": u.get("command"), + "as_node": u["as_node"], + } + for u in s["updates"] + ] + } + for s in supersteps + ] + return self.http.post("/threads", json=payload) def update(self, thread_id: str, *, metadata: dict[str, Any]) -> Thread: @@ -3307,7 +3353,7 @@ class SyncThreadsClient: Example Usage: - response = client.threads.update_state( + response = await client.threads.update_state( thread_id="my_thread_id", values={"messages":[{"role": "user", "content": "hello!"}]}, as_node="my_node", diff --git a/libs/sdk-py/pyproject.toml b/libs/sdk-py/pyproject.toml index 75e402fee..92e942441 100644 --- a/libs/sdk-py/pyproject.toml +++ b/libs/sdk-py/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-sdk" -version = "0.1.57" +version = "0.1.58" description = "SDK for interacting with LangGraph API" authors = [] license = "MIT"