diff --git a/docs/docs/cloud/deployment/setup_javascript.md b/docs/docs/cloud/deployment/setup_javascript.md new file mode 100644 index 000000000..fa4c7b823 --- /dev/null +++ b/docs/docs/cloud/deployment/setup_javascript.md @@ -0,0 +1,200 @@ +# How to Set Up a LangGraph.js Application for Deployment + +A [LangGraph.js](https://langchain-ai.github.io/langgraphjs/) application must be configured with a [LangGraph API configuration file](../reference/cli.md#configuration-file) in order to be deployed to LangGraph Cloud (or to be self-hosted). This how-to guide discusses the basic steps to setup a LangGraph.js application for deployment using `package.json` to specify project dependencies. + +This walkthrough is based on [this repository](https://github.com/langchain-ai/langgraphjs-studio-starter), which you can play around with to learn more about how to setup your LangGraph application for deployment. + +The final repo structure will look something like this: + +```bash +my-app/ +├── src # all project code lies within here +│ ├── utils # optional utilities for your graph +│ │ ├── tools.ts # tools for your graph +│ │ ├── nodes.ts # node functions for you graph +│ │ └── state.ts # state definition of your graph +│   └── agent.ts # code for constructing your graph +├── package.json # package dependencies +├── .env # environment variables +└── langgraph.json # configuration file for LangGraph +``` + +After each step, an example file directory is provided to demonstrate how code can be organized. + +## Specify Dependencies + +Dependencies can be specified in a `package.json`. If none of these files is created, then dependencies can be specified later in the [LangGraph API configuration file](#create-langgraph-api-config). + +Example `package.json` file: + +```json +{ + "name": "langgraphjs-studio-starter", + "packageManager": "yarn@1.22.22", + "dependencies": { + "@langchain/community": "^0.2.31", + "@langchain/core": "^0.2.31", + "@langchain/langgraph": "^0.2.0", + "@langchain/openai": "^0.2.8" + } +} +``` + +Example file directory: + +```bash +my-app/ +└── package.json # package dependencies +``` + +## Specify Environment Variables + +Environment variables can optionally be specified in a file (e.g. `.env`). See the [Environment Variables reference](../reference/env_var.md) to configure additional variables for a deployment. + +Example `.env` file: + +``` +MY_ENV_VAR_1=foo +MY_ENV_VAR_2=bar +OPENAI_API_KEY=key +TAVILY_API_KEY=key_2 +``` + +Example file directory: + +```bash +my-app/ +├── package.json +└── .env # environment variables +``` + +## Define Graphs + +Implement your graphs! Graphs can be defined in a single file or multiple files. Make note of the variable names of each compiled graph to be included in the LangGraph application. The variable names will be used later when creating the [LangGraph API configuration file](../reference/cli.md#configuration-file). + +Here is an example `agent.ts`: + +```ts +import type { AIMessage } from "@langchain/core/messages"; +import { TavilySearchResults } from "@langchain/community/tools/tavily_search"; +import { ChatOpenAI } from "@langchain/openai"; + +import { MessagesAnnotation, StateGraph } from "@langchain/langgraph"; +import { ToolNode } from "@langchain/langgraph/prebuilt"; + +const tools = [ + new TavilySearchResults({ maxResults: 3, }), +]; + +// Define the function that calls the model +async function callModel( + state: typeof MessagesAnnotation.State, +) { + /** + * Call the LLM powering our agent. + * Feel free to customize the prompt, model, and other logic! + */ + const model = new ChatOpenAI({ + model: "gpt-4o", + }).bindTools(tools); + + const response = await model.invoke([ + { + role: "system", + content: `You are a helpful assistant. The current date is ${new Date().getTime()}.` + }, + ...state.messages + ]); + + // MessagesAnnotation supports returning a single message or array of messages + return { messages: response }; +} + +// Define the function that determines whether to continue or not +function routeModelOutput(state: typeof MessagesAnnotation.State) { + const messages = state.messages; + const lastMessage: AIMessage = messages[messages.length - 1]; + // If the LLM is invoking tools, route there. + if ((lastMessage?.tool_calls?.length ?? 0) > 0) { + return "tools"; + } + // Otherwise end the graph. + return "__end__"; +} + +// Define a new graph. +// See https://langchain-ai.github.io/langgraphjs/how-tos/define-state/#getting-started for +// more on defining custom graph states. +const workflow = new StateGraph(MessagesAnnotation) + // Define the two nodes we will cycle between + .addNode("callModel", callModel) + .addNode("tools", new ToolNode(tools)) + // Set the entrypoint as `callModel` + // This means that this node is the first one called + .addEdge("__start__", "callModel") + .addConditionalEdges( + // First, we define the edges' source node. We use `callModel`. + // This means these are the edges taken after the `callModel` node is called. + "callModel", + // Next, we pass in the function that will determine the sink node(s), which + // will be called after the source node is called. + routeModelOutput, + // List of the possible destinations the conditional edge can route to. + // Required for conditional edges to properly render the graph in Studio + [ + "tools", + "__end__" + ], + ) + // This means that after `tools` is called, `callModel` node is called next. + .addEdge("tools", "callModel"); + +// Finally, we compile it! +// This compiles it into a graph you can invoke and deploy. +export const graph = workflow.compile(); +``` + +!!! info "Assign `CompiledGraph` to Variable" + The build process for LangGraph Cloud requires that the `CompiledGraph` object be assigned to a variable at the top-level of a JavaScript module (alternatively, you can provide [a function that creates a graph](./graph_rebuild.md)). + +Example file directory: + +```bash +my-app/ +├── src # all project code lies within here +│ ├── utils # optional utilities for your graph +│ │ ├── tools.ts # tools for your graph +│ │ ├── nodes.ts # node functions for you graph +│ │ └── state.ts # state definition of your graph +│   └── agent.ts # code for constructing your graph +├── package.json # package dependencies +├── .env # environment variables +└── langgraph.json # configuration file for LangGraph +``` + +## Create LangGraph API Config + +Create a [LangGraph API configuration file](../reference/cli.md#configuration-file) called `langgraph.json`. See the [LangGraph CLI reference](../reference/cli.md#configuration-file) for detailed explanations of each key in the JSON object of the configuration file. + +Example `langgraph.json` file: + +```json +{ + "node_version": "20", + "dockerfile_lines": [], + "dependencies": ["."], + "graphs": { + "agent": "./src/agent.ts:graph" + }, + "env": ".env" +} +``` + +Note that the variable name of the `CompiledGraph` appears at the end of the value of each subkey in the top-level `graphs` key (i.e. `:`). + +!!! info "Configuration Location" + The LangGraph API configuration file must be placed in a directory that is at the same level or higher than the TypeScript files that contain compiled graphs and associated dependencies. + +## Next + +After you setup your project and place it in a github repo, it's time to [deploy your app](./cloud.md). diff --git a/docs/docs/cloud/faq/studio.md b/docs/docs/cloud/faq/studio.md index cef962f34..070014ce7 100644 --- a/docs/docs/cloud/faq/studio.md +++ b/docs/docs/cloud/faq/studio.md @@ -43,15 +43,23 @@ If you don't define your conditional edges carefully, you might notice extra edg ### Solution 1: Include a path map -The first way to solve this is to add path maps to your conditional edges. A path map is just a dictionary that maps the possible outputs of your router function with the names of the nodes that each output corresponds to. The path map is passed as the third argument to the `add_conditional_edges` function like so: +The first way to solve this is to add path maps to your conditional edges. A path map is just a dictionary or array that maps the possible outputs of your router function with the names of the nodes that each output corresponds to. The path map is passed as the third argument to the `add_conditional_edges` function like so: -```python -graph.add_conditional_edges("node_a", routing_function, {True: "node_b", False: "node_c"}) -``` +=== "Python" + + ```python + graph.add_conditional_edges("node_a", routing_function, {True: "node_b", False: "node_c"}) + ``` + +=== "Javascript" + + ```ts + graph.addConditionalEdges("node_a", routingFunction, ["node_b", "node_c"]); + ``` In this case, the routing function returns either True or False, which map to `node_b` and `node_c` respectively. -### Solution 2: Update the typing of the router +### Solution 2: Update the typing of the router (Python only) Instead of passing a path map, you can also be explicit about the typing of your routing function by specifying the nodes it can map to using the `Literal` python definition. Here is an example of how to define a routing function in that way: diff --git a/docs/docs/cloud/how-tos/index.md b/docs/docs/cloud/how-tos/index.md index 255490d32..90b13c35e 100644 --- a/docs/docs/cloud/how-tos/index.md +++ b/docs/docs/cloud/how-tos/index.md @@ -13,6 +13,7 @@ LangGraph Cloud gives you best in class observability, testing, and hosting serv - [How to set up app for deployment (requirements.txt)](../deployment/setup.md) - [How to set up app for deployment (pyproject.toml)](../deployment/setup_pyproject.md) +- [How to set up app for deployment (JavaScript)](../deployment/setup_javascript.md) - [How to test locally](../deployment/test_locally.md) - [How to deploy to LangGraph cloud](../deployment/cloud.md) - [How to self-host](../deployment/self_hosted.md) diff --git a/docs/docs/cloud/how-tos/interrupt_concurrent.md b/docs/docs/cloud/how-tos/interrupt_concurrent.md index e67fb4ce5..b09f908ef 100644 --- a/docs/docs/cloud/how-tos/interrupt_concurrent.md +++ b/docs/docs/cloud/how-tos/interrupt_concurrent.md @@ -44,7 +44,7 @@ Now, let's import our required packages and instantiate our client, assistant, a const thread = await client.threads.create(); ``` -Now we can start our two runs and join the second on euntil it has completed: +Now we can start our two runs and join the second one until it has completed: === "Python" diff --git a/docs/docs/cloud/quick_start.md b/docs/docs/cloud/quick_start.md index c59a6fb85..755ce0df4 100644 --- a/docs/docs/cloud/quick_start.md +++ b/docs/docs/cloud/quick_start.md @@ -14,13 +14,25 @@ This tutorial will use: 1. Create a new application with the following directory and files: +=== "Python" + / |-- agent.py # code for your LangGraph agent |-- requirements.txt # Python packages required for your graph |-- langgraph.json # configuration file for LangGraph |-- .env # environment files with API keys -2. The `agent.py` file should contain Python code for defining your graph. The following code is a simple example, the important thing is that at some point in your file you compile your graph and assign the compiled graph to a variable (in this case the `graph` variable). This example code uses `create_react_agent`, a prebuilt agent, read more about it [here](..//concepts/agentic_concepts.md#react-agent). +=== "Javascript" + + / + |-- agent.ts # code for your LangGraph agent + |-- package.json # Javascript packages required for your graph + |-- langgraph.json # configuration file for LangGraph + |-- .env # environment files with API keys + +2. The `agent.py`/`agent.ts` file should contain code for defining your graph. The following code is a simple example, the important thing is that at some point in your file you compile your graph and assign the compiled graph to a variable (in this case the `graph` variable). This example code uses `create_react_agent`, a prebuilt agent. You can read more about it [here](../concepts/agentic_concepts.md#react-agent). + +=== "Python" ```python from langchain_anthropic import ChatAnthropic @@ -34,14 +46,53 @@ This tutorial will use: graph = create_react_agent(model, tools) ``` -3. The `requirements.txt` file should contain any dependencies for your graph(s). In this case we only require four packages for our graph to run: +=== "Javascript" - langgraph - langchain_anthropic - tavily-python - langchain_community + ```ts + import { ChatAnthropic } from "@langchain/anthropic"; + import { TavilySearchResults } from "@langchain/community/tools/tavily_search"; + import { createReactAgent } from "@langchain/langgraph/prebuilt"; -4. The [`langgraph.json`][langgraph.json] file is a configuration file that describes what graph(s) you are going to host. In this case we only have one graph to host: the compiled `graph` object from `agent.py`. + const model = new ChatAnthropic({ + model: "claude-3-5-sonnet-20240620", + }); + + const tools = [ + new TavilySearchResults({ maxResults: 3, }), + ]; + + export const graph = createReactAgent({ llm: model, tools }); + ``` + +3. The `requirements.txt`/`package.json` file should contain any dependencies for your graph(s). In this case we only require four packages for our graph to run: + +=== "Python" + + ```python + langgraph + langchain_anthropic + tavily-python + langchain_community + ``` + +=== "Javascript" + + ```js + { + "name": "my-app", + "packageManager": "yarn@1.22.22", + "dependencies": { + "@langchain/community": "^0.2.31", + "@langchain/core": "^0.2.31", + "@langchain/langgraph": "0.2.0", + "@langchain/openai": "^0.2.8" + } + } + ``` + +4. The [`langgraph.json`][langgraph.json] file is a configuration file that describes what graph(s) you are going to host. In this case we only have one graph to host: the compiled `graph` object from `agent.py`/`agent.ts`. + +=== "Python" ```json { @@ -53,7 +104,21 @@ This tutorial will use: } ``` - Learn more about the LangGraph CLI configuration file [here](./reference/cli.md#configuration-file). +=== "Javascript" + + ```json + { + "node_version": "20", + "dockerfile_lines": [], + "dependencies": ["."], + "graphs": { + "agent": "./src/agent.ts:graph" + }, + "env": ".env" + } + ``` + +Learn more about the LangGraph CLI configuration file [here](./reference/cli.md#configuration-file). 5. The `.env` file should have any environment variables needed to run your graph. This will only be used for local testing, so if you are not testing locally you can skip this step. NOTE: if you do add this, you should NOT check this into git. For this graph, we need two environment variables: @@ -206,36 +271,108 @@ export LANGSMITH_API_KEY=... The first thing to do when using the SDK is to setup our client, access our assistant, and create a thread to execute a run on: -```python -from langgraph_sdk import get_client +=== "Python" -# Replace this with the URL of your own deployed graph -URL = "https://chatbot-23a570f3210f52a7b167f09f6158e3b3-ffoprvkqsa-uc.a.run.app" -client = get_client(url=URL) + ```python + from langgraph_sdk import get_client -# Search all hosted graphs -assistants = await client.assistants.search() -# In this example we select the first assistant since we are only hosting a single graph -assistant = assistants[0] + client = get_client(url=) + # get default assistant + assistants = await client.assistants.search() + assistant = [a for a in assistants if not a["config"]][0] + # create thread + thread = await client.threads.create() + print(thread) + ``` -# We create a thread for tracking the state of our run -thread = await client.threads.create() -``` +=== "Javascript" + + ```js + import { Client } from "@langchain/langgraph-sdk"; + + const client = new Client({ apiUrl: }); + // get default assistant + const assistants = await client.assistants.search(); + const assistant = assistants.find(a => !a.config); + // create thread + const thread = await client.threads.create(); + console.log(thread) + ``` + +=== "CURL" + + ```bash + curl --request POST \ + --url /assistants/search \ + --header 'Content-Type: application/json' \ + --data '{ + "limit": 10, + "offset": 0 + }' | jq -c 'map(select(.config == null or .config == {})) | .[0]' && \ + curl --request POST \ + --url /threads \ + --header 'Content-Type: application/json' \ + --data '{}' + ``` We can then execute a run on the thread: -```python -input = {"messages":[{"role": "user", "content": "Hello! My name is Bagatur and I am 26 years old."}]} +=== "Python" -async for chunk in client.runs.stream( - thread['thread_id'], - assistant["assistant_id"], - input=input, - stream_mode="updates", - ): - if chunk.data and chunk.event != "metadata": - print(chunk.data) -``` + ```python + input = {"messages":[{"role": "user", "content": "Hello! My name is Bagatur and I am 26 years old."}]} + + async for chunk in client.runs.stream( + thread['thread_id'], + assistant["assistant_id"], + input=input, + stream_mode="updates", + ): + if chunk.data and chunk.event != "metadata": + print(chunk.data) + ``` + +=== "Javascript" + + ```js + const input = { "messages":[{ "role": "user", "content": "Hello! My name is Bagatur and I am 26 years old." }] }; + + const streamResponse = client.runs.stream( + thread["thread_id"], + assistant["assistant_id"], + { + input, + } + ); + for await (const chunk of streamResponse) { + if (chunk.data && chunk.event !== "metadata" ) { + console.log(chunk.data); + } + } + ``` + +=== "CURL" + + ```bash + curl --request POST \ + --url /threads//runs/stream \ + --header 'Content-Type: application/json' \ + --data "{ + \"assistant_id\": , + \"input\": {\"messages\": [{\"role\": \"human\", \"content\": \"Hello! My name is Bagatur and I am 26 years old.\"}]}, + }" | sed 's/\r$//' | awk ' + /^event:/ { event = $2 } + /^data:/ { + json_data = substr($0, index($0, $2)) + + if (event != "metadata") { + print json_data + } + }' + ``` + + +Output: {'agent': {'messages': [{'content': "Hi Bagatur! It's nice to meet you. How can I assist you today?", 'additional_kwargs': {}, 'response_metadata': {'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_9cb5d38cf7'}, 'type': 'ai', 'name': None, 'id': 'run-c89118b7-1b1e-42b9-a85d-c43fe99881cd', 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None}]}} diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index eb6d5d36f..56a9435b2 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -197,6 +197,7 @@ nav: - Setup: - Setup App: "cloud/deployment/setup.md" - Setup App (pyproject.toml): "cloud/deployment/setup_pyproject.md" + - Setup App (JavaScript): "cloud/deployment/setup_javascript.md" - Rebuild Graph at Runtime: "cloud/deployment/graph_rebuild.md" - Test App Locally: "cloud/deployment/test_locally.md" - Deployment: diff --git a/examples/introduction.ipynb b/examples/introduction.ipynb index 3d0287d05..52cd1018b 100644 --- a/examples/introduction.ipynb +++ b/examples/introduction.ipynb @@ -1504,7 +1504,7 @@ "from langgraph.checkpoint.memory import MemorySaver\n", "from langgraph.graph import StateGraph\n", "from langgraph.graph.message import add_messages\n", - "from langgraph.prebuilt import ToolNode\n", + "from langgraph.prebuilt import ToolNode, tools_condition\n", "\n", "\n", "class State(TypedDict):\n", diff --git a/examples/llm-compiler/LLMCompiler.ipynb b/examples/llm-compiler/LLMCompiler.ipynb index 98f5416d8..3ae67e33b 100644 --- a/examples/llm-compiler/LLMCompiler.ipynb +++ b/examples/llm-compiler/LLMCompiler.ipynb @@ -646,11 +646,12 @@ "def _parse_joiner_output(decision: JoinOutputs) -> List[BaseMessage]:\n", " response = [AIMessage(content=f\"Thought: {decision.thought}\")]\n", " if isinstance(decision.action, Replan):\n", - " return response + [\n", + " return {\"messages\": response + [\n", " SystemMessage(\n", " content=f\"Context from last attempt: {decision.action.feedback}\"\n", " )\n", " ]\n", + " }\n", " else:\n", " return {\"messages\": response + [AIMessage(content=decision.action.response)]}\n", "\n", @@ -933,6 +934,46 @@ "print(step['join']['messages'][-1].content)" ] }, + { + "cell_type": "markdown", + "id": "f9487866", + "metadata": {}, + "source": [ + "#### Complex Replanning Example\n", + "\n", + "This question is likely to prompt the Replan functionality, but it may need to be run multiple times to see this in action." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "391d6931", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{'plan_and_schedule': {'messages': [FunctionMessage(content=\"[{'url': 'https://www.timeanddate.com/weather/japan/tokyo', 'content': '88 / 84 °F. 13. 87 / 82 °F. 14. 84 / 80 °F. Detailed forecast for 14 days. Need some help? Current weather in Tokyo and forecast for today, tomorrow, and next 14 days.'}]\", additional_kwargs={'idx': 1, 'args': {'query': 'current temperature in Tokyo'}}, name='tavily_search_results_json', tool_call_id=1), FunctionMessage(content='join', additional_kwargs={'idx': 2, 'args': ()}, name='join', tool_call_id=2)]}}\n", + "{'join': {'messages': [AIMessage(content=\"Thought: The search result provides the current temperature in Tokyo but does not explicitly state which temperature (88 / 84 °F) corresponds to the current condition. It seems to be a range, possibly the day's high and low. Without a clear indication of the exact current temperature, it's challenging to provide a precise flashcard summary.\", id='8ef2a131-69db-4180-a76e-fd9d6f4037c1'), SystemMessage(content='Context from last attempt: The information provided does not explicitly state the current temperature in Tokyo; it provides a temperature range without specifying which is the current temperature. Need to find a source that gives the exact current temperature in Tokyo for a precise flashcard summary.', id='f5bd752c-b068-459a-8d9e-bd1f1b5fa4fe')]}}\n", + "{'plan_and_schedule': {'messages': [FunctionMessage(content='join', additional_kwargs={'idx': 3, 'args': ()}, name='join', tool_call_id=3)]}}\n", + "{'join': {'messages': [AIMessage(content=\"Thought: The search result provides a temperature range for Tokyo but does not specify the current temperature. This makes it challenging to create a precise flashcard without an exact current temperature. The user's request cannot be fully satisfied without this detail.\", id='3cc41891-4f47-4453-8edf-b989926ab25e'), SystemMessage(content='Context from last attempt: The search did not provide an exact current temperature for Tokyo, making it impossible to create a precise flashcard. A source that explicitly states the current temperature is needed for an accurate response.', id='96290b41-a4c4-4ab5-829a-89cc31dfe6c8')]}}\n", + "{'plan_and_schedule': {'messages': [FunctionMessage(content='join', additional_kwargs={'idx': 4, 'args': ()}, name='join', tool_call_id=4)]}}\n", + "{'join': {'messages': [AIMessage(content=\"Thought: The search result provides a temperature range for Tokyo but does not specify the current temperature. This makes it challenging to create a precise flashcard without an exact current temperature. The user's request cannot be fully satisfied without this detail.\", id='4724b242-ddb8-47e6-b235-de25de54fe45'), AIMessage(content='I was unable to find the exact current temperature in Tokyo. However, the temperature range for today in Tokyo is between 88°F and 84°F. For the most accurate and up-to-date temperature, I recommend checking a reliable weather forecasting website or app.', id='40e29a47-a001-4f65-a18f-65c2931d1ae5')]}}\n" + ] + } + ], + "source": [ + "for step in chain.stream({\"messages\":\n", + " [\n", + " HumanMessage(\n", + " content=\"Find the current temperature in Tokyo, then, respond with a flashcard summarizing this information\"\n", + " )\n", + " ]}\n", + "):\n", + " print(step)" + ] + }, { "cell_type": "markdown", "id": "c647d5f3-5e00-4449-9cec-5a9f438c9cff", diff --git a/examples/multi_agent/agent_supervisor.ipynb b/examples/multi_agent/agent_supervisor.ipynb index bd2a7c6ce..b1ecfe31b 100644 --- a/examples/multi_agent/agent_supervisor.ipynb +++ b/examples/multi_agent/agent_supervisor.ipynb @@ -219,7 +219,7 @@ "workflow = StateGraph(AgentState)\n", "workflow.add_node(\"Researcher\", research_node)\n", "workflow.add_node(\"Coder\", code_node)\n", - "workflow.add_node(\"supervisor\", supervisor_chain)" + "workflow.add_node(\"supervisor\", supervisor_agent)" ] }, { diff --git a/examples/persistence_postgres.ipynb b/examples/persistence_postgres.ipynb index 5ea838cee..b4fa1e2fe 100644 --- a/examples/persistence_postgres.ipynb +++ b/examples/persistence_postgres.ipynb @@ -132,8 +132,6 @@ "metadata": {}, "outputs": [], "source": [ - "from psycopg.rows import dict_row\n", - "\n", "connection_kwargs = {\n", " \"autocommit\": True,\n", " \"prepare_threshold\": 0,\n", @@ -161,15 +159,13 @@ "source": [ "from psycopg_pool import ConnectionPool\n", "\n", - "pool = ConnectionPool(\n", + "with ConnectionPool(\n", " # Example configuration\n", " conninfo=DB_URI,\n", " max_size=20,\n", " kwargs=connection_kwargs,\n", - ")\n", - "\n", - "with pool.connection() as conn:\n", - " checkpointer = PostgresSaver(conn)\n", + ") as pool:\n", + " checkpointer = PostgresSaver(pool)\n", "\n", " # NOTE: you need to call .setup() the first time you're using your checkpointer\n", " checkpointer.setup()\n", @@ -394,8 +390,8 @@ " conninfo=DB_URI,\n", " max_size=20,\n", " kwargs=connection_kwargs,\n", - ") as pool, pool.connection() as conn:\n", - " checkpointer = AsyncPostgresSaver(conn)\n", + ") as pool:\n", + " checkpointer = AsyncPostgresSaver(pool)\n", "\n", " # NOTE: you need to call .setup() the first time you're using your checkpointer\n", " # await checkpointer.setup()\n", diff --git a/examples/rag/langgraph_crag_local.ipynb b/examples/rag/langgraph_crag_local.ipynb index dbac6f697..d6e791ae3 100644 --- a/examples/rag/langgraph_crag_local.ipynb +++ b/examples/rag/langgraph_crag_local.ipynb @@ -429,6 +429,14 @@ "]\n", "\n", "\n", + "def find_tool_calls_react(messages):\n", + " \"\"\"\n", + " Find all tool calls in the messages returned\n", + " \"\"\"\n", + " tool_calls = [tc['name'] for m in messages['messages'] for tc in getattr(m, 'tool_calls', [])]\n", + " return tool_calls\n", + "\n", + "\n", "def check_trajectory_react(root_run: Run, example: Example) -> dict:\n", " \"\"\"\n", " Check if all expected tools are called in exact order and without any additional tool calls.\n", diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 63a2038a3..6140af756 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -1,6 +1,6 @@ import asyncio from contextlib import asynccontextmanager -from typing import Any, AsyncIterator, Optional, Union +from typing import Any, AsyncIterator, Iterator, List, Optional, Union from langchain_core.runnables import RunnableConfig from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline @@ -51,6 +51,7 @@ class AsyncPostgresSaver(BasePostgresSaver): self.conn = conn self.pipe = pipe self.lock = asyncio.Lock() + self.loop = asyncio.get_running_loop() @classmethod @asynccontextmanager @@ -329,3 +330,96 @@ class AsyncPostgresSaver(BasePostgresSaver): binary=True, row_factory=dict_row ) as cur: yield cur + + def list( + self, + config: Optional[RunnableConfig], + *, + filter: Optional[dict[str, Any]] = None, + before: Optional[RunnableConfig] = None, + limit: Optional[int] = None, + ) -> Iterator[CheckpointTuple]: + """List checkpoints from the database. + + This method retrieves a list of checkpoint tuples from the Postgres database based + on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first). + + Args: + config (Optional[RunnableConfig]): Base configuration for filtering checkpoints. + filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. + before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None. + limit (Optional[int]): Maximum number of checkpoints to return. + + Yields: + Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples. + """ + aiter_ = self.alist(config, filter=filter, before=before, limit=limit) + while True: + try: + yield asyncio.run_coroutine_threadsafe( + anext(aiter_), self.loop + ).result() + except StopAsyncIteration: + break + + def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: + """Get a checkpoint tuple from the database. + + This method retrieves a checkpoint tuple from the Postgres database based on the + provided config. If the config contains a "checkpoint_id" key, the checkpoint with + the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint + for the given thread ID is retrieved. + + Args: + config (RunnableConfig): The config to use for retrieving the checkpoint. + + Returns: + Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. + """ + return asyncio.run_coroutine_threadsafe( + self.aget_tuple(config), self.loop + ).result() + + def put( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + new_versions: ChannelVersions, + ) -> RunnableConfig: + """Save a checkpoint to the database. + + This method saves a checkpoint to the Postgres database. The checkpoint is associated + with the provided config and its parent config (if any). + + Args: + config (RunnableConfig): The config to associate with the checkpoint. + checkpoint (Checkpoint): The checkpoint to save. + metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. + new_versions (ChannelVersions): New channel versions as of this write. + + Returns: + RunnableConfig: Updated configuration after storing the checkpoint. + """ + return asyncio.run_coroutine_threadsafe( + self.aput(config, checkpoint, metadata, new_versions), self.loop + ).result() + + def put_writes( + self, + config: RunnableConfig, + writes: List[tuple[str, Any]], + task_id: str, + ) -> None: + """Store intermediate writes linked to a checkpoint. + + This method saves intermediate writes associated with a checkpoint to the database. + + Args: + config (RunnableConfig): Configuration of the related checkpoint. + writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair. + task_id (str): Identifier for the task creating the writes. + """ + return asyncio.run_coroutine_threadsafe( + self.aput_writes(config, writes, task_id), self.loop + ).result() diff --git a/libs/checkpoint-postgres/poetry.lock b/libs/checkpoint-postgres/poetry.lock index f9ca0f97e..520300e91 100644 --- a/libs/checkpoint-postgres/poetry.lock +++ b/libs/checkpoint-postgres/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. [[package]] name = "annotated-types" @@ -266,7 +266,7 @@ tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<9.0.0" [[package]] name = "langgraph-checkpoint" -version = "1.0.6" +version = "1.0.8" description = "Library with base interfaces for LangGraph checkpoint savers." optional = false python-versions = "^3.9.0,<4.0" @@ -766,7 +766,6 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, @@ -972,4 +971,4 @@ watchmedo = ["PyYAML (>=3.10)"] [metadata] lock-version = "2.0" python-versions = "^3.9.0,<4.0" -content-hash = "b139531e8c6f4e24cea4bdfc29d111d1ea000a6a8a604ba81be4bff0977a2466" +content-hash = "e294b6996aa6c8f671e6aaf65be8b4aba94c18e2f237dd3ee0e1b777849ce8a8" diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml index 077f8dd09..6e3640e7c 100644 --- a/libs/checkpoint-postgres/pyproject.toml +++ b/libs/checkpoint-postgres/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint-postgres" -version = "1.0.5" +version = "1.0.6" 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 = "^1.0.1" +langgraph-checkpoint = "^1.0.8" orjson = ">=3.10.1" psycopg = "^3.0.0" psycopg-pool = "^3.0.0" diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py index 04493a381..31bffd19c 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py @@ -104,10 +104,12 @@ class SqliteSaver(BaseCheckpointSaver): with SqliteSaver.from_conn_string("checkpoints.sqlite") as memory: ... """ - with sqlite3.connect( - conn_string, - # https://ricardoanderegg.com/posts/python-sqlite-thread-safety/ - check_same_thread=False, + with closing( + sqlite3.connect( + conn_string, + # https://ricardoanderegg.com/posts/python-sqlite-thread-safety/ + check_same_thread=False, + ) ) as conn: yield SqliteSaver(conn) @@ -163,14 +165,15 @@ class SqliteSaver(BaseCheckpointSaver): Yields: sqlite3.Cursor: A cursor for the SQLite database. """ - self.setup() - cur = self.conn.cursor() - try: - yield cur - finally: - if transaction: - self.conn.commit() - cur.close() + with self.lock: + self.setup() + cur = self.conn.cursor() + try: + yield cur + finally: + if transaction: + self.conn.commit() + cur.close() def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the database. @@ -396,7 +399,7 @@ class SqliteSaver(BaseCheckpointSaver): checkpoint_ns = config["configurable"]["checkpoint_ns"] type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint) serialized_metadata = self.jsonplus_serde.dumps(metadata) - with self.lock, self.cursor() as cur: + with self.cursor() as cur: cur.execute( "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)", ( @@ -432,7 +435,7 @@ class SqliteSaver(BaseCheckpointSaver): writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair. task_id (str): Identifier for the task creating the writes. """ - with self.lock, self.cursor() as cur: + with self.cursor() as cur: cur.executemany( "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", [ diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index dc68a0403..28ad76109 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -1,11 +1,11 @@ import asyncio -import functools from contextlib import asynccontextmanager from typing import ( Any, AsyncIterator, Dict, Iterator, + List, Optional, Sequence, Tuple, @@ -31,20 +31,6 @@ from langgraph.checkpoint.sqlite.utils import search_where T = TypeVar("T", bound=callable) -def not_implemented_sync_method(func: T) -> T: - @functools.wraps(func) - def wrapper(*args, **kwargs): - raise NotImplementedError( - "The AsyncSqliteSaver does not support synchronous methods. " - "Consider using the SqliteSaver instead.\n" - "from langgraph.checkpoint.sqlite import SqliteSaver\n" - "See https://langchain-ai.github.io/langgraph/reference/checkpoints/langgraph.checkpoint.sqlite.SqliteSaver " - "for more information." - ) - - return wrapper - - class AsyncSqliteSaver(BaseCheckpointSaver): """An asynchronous checkpoint saver that stores checkpoints in a SQLite database. @@ -132,6 +118,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver): self.jsonplus_serde = JsonPlusSerializer() self.conn = conn self.lock = asyncio.Lock() + self.loop = asyncio.get_running_loop() self.is_setup = False @classmethod @@ -150,16 +137,24 @@ class AsyncSqliteSaver(BaseCheckpointSaver): async with aiosqlite.connect(conn_string) as conn: yield AsyncSqliteSaver(conn) - @not_implemented_sync_method def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the database. - Note: - This method is not implemented for the AsyncSqliteSaver. Use `aget` instead. - Or consider using the [SqliteSaver][sqlitesaver] checkpointer. - """ + This method retrieves a checkpoint tuple from the SQLite database based on the + provided config. If the config contains a "checkpoint_id" key, the checkpoint with + the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint + for the given thread ID is retrieved. + + Args: + config (RunnableConfig): The config to use for retrieving the checkpoint. + + Returns: + Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. + """ + return asyncio.run_coroutine_threadsafe( + self.aget_tuple(config), self.loop + ).result() - @not_implemented_sync_method def list( self, config: Optional[RunnableConfig], @@ -168,21 +163,60 @@ class AsyncSqliteSaver(BaseCheckpointSaver): before: Optional[RunnableConfig] = None, limit: Optional[int] = None, ) -> Iterator[CheckpointTuple]: - """List checkpoints from the database. + """List checkpoints from the database asynchronously. - Note: - This method is not implemented for the AsyncSqliteSaver. Use `alist` instead. - Or consider using the [SqliteSaver][sqlitesaver] checkpointer. + This method retrieves a list of checkpoint tuples from the SQLite database based + on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first). + + Args: + config (Optional[RunnableConfig]): Base configuration for filtering checkpoints. + filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. + before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None. + limit (Optional[int]): Maximum number of checkpoints to return. + + Yields: + Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples. """ + aiter_ = self.alist(config, filter=filter, before=before, limit=limit) + while True: + try: + yield asyncio.run_coroutine_threadsafe( + anext(aiter_), self.loop + ).result() + except StopAsyncIteration: + break - @not_implemented_sync_method def put( self, config: RunnableConfig, checkpoint: Checkpoint, metadata: CheckpointMetadata, + new_versions: ChannelVersions, ) -> RunnableConfig: - """Save a checkpoint to the database. FOO""" + """Save a checkpoint to the database. + + This method saves a checkpoint to the SQLite database. The checkpoint is associated + with the provided config and its parent config (if any). + + Args: + config (RunnableConfig): The config to associate with the checkpoint. + checkpoint (Checkpoint): The checkpoint to save. + metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. + new_versions (ChannelVersions): New channel versions as of this write. + + Returns: + RunnableConfig: Updated configuration after storing the checkpoint. + """ + return asyncio.run_coroutine_threadsafe( + self.aput(config, checkpoint, metadata, new_versions), self.loop + ).result() + + def put_writes( + self, config: RunnableConfig, writes: List[Tuple[str, Any]], task_id: str + ) -> None: + return asyncio.run_coroutine_threadsafe( + self.aput_writes(config, writes, task_id), self.loop + ).result() async def setup(self) -> None: """Set up the checkpoint database asynchronously. @@ -242,7 +276,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver): """ await self.setup() checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - async with self.conn.cursor() as cur: + async with self.lock, self.conn.cursor() as cur: # find the latest checkpoint for the thread_id if checkpoint_id := get_checkpoint_id(config): await cur.execute( @@ -337,7 +371,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver): ORDER BY checkpoint_id DESC""" if limit: query += f" LIMIT {limit}" - async with self.conn.execute(query, params) as cur, self.conn.cursor() as wcur: + async with self.lock, self.conn.execute( + query, params + ) as cur, self.conn.cursor() as wcur: async for ( thread_id, checkpoint_ns, @@ -404,7 +440,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver): checkpoint_ns = config["configurable"]["checkpoint_ns"] type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint) serialized_metadata = self.jsonplus_serde.dumps(metadata) - async with self.conn.execute( + async with self.lock, self.conn.execute( "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)", ( str(config["configurable"]["thread_id"]), @@ -441,7 +477,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver): task_id (str): Identifier for the task creating the writes. """ await self.setup() - async with self.conn.cursor() as cur: + async with self.lock, self.conn.cursor() as cur: await cur.executemany( "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", [ diff --git a/libs/checkpoint-sqlite/poetry.lock b/libs/checkpoint-sqlite/poetry.lock index 0ea747b2c..8d5dbe3af 100644 --- a/libs/checkpoint-sqlite/poetry.lock +++ b/libs/checkpoint-sqlite/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. [[package]] name = "aiosqlite" @@ -252,7 +252,7 @@ tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<9.0.0" [[package]] name = "langgraph-checkpoint" -version = "1.0.1" +version = "1.0.8" description = "Library with base interfaces for LangGraph checkpoint savers." optional = false python-versions = "^3.9.0,<4.0" @@ -651,7 +651,6 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, @@ -835,4 +834,4 @@ watchmedo = ["PyYAML (>=3.10)"] [metadata] lock-version = "2.0" python-versions = "^3.9.0" -content-hash = "d50ec7c6b55075d19193e080cc95d3b0998b9efa1a0ed407bd6055b5b5e867e0" +content-hash = "752a22dc2b57a0818a3a4d9bf5f62226ba7f8e0c458892551bf8b4c93723dbc1" diff --git a/libs/checkpoint-sqlite/pyproject.toml b/libs/checkpoint-sqlite/pyproject.toml index 536895f77..1b9462085 100644 --- a/libs/checkpoint-sqlite/pyproject.toml +++ b/libs/checkpoint-sqlite/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint-sqlite" -version = "1.0.1" +version = "1.0.3" description = "Library with a SQLite implementation of LangGraph checkpoint saver." authors = [] license = "MIT" @@ -10,7 +10,7 @@ packages = [{ include = "langgraph" }] [tool.poetry.dependencies] python = "^3.9.0" -langgraph-checkpoint = "^1.0.1" +langgraph-checkpoint = "^1.0.8" aiosqlite = "^0.20.0" [tool.poetry.group.dev.dependencies] diff --git a/libs/langgraph/Makefile b/libs/langgraph/Makefile index ccb8ef59a..d5c6d7ee3 100644 --- a/libs/langgraph/Makefile +++ b/libs/langgraph/Makefile @@ -15,7 +15,7 @@ coverage: --cov-report term-missing:skip-covered start-postgres: - docker compose -f tests/compose-postgres.yml up -V --force-recreate --wait + docker compose -f tests/compose-postgres.yml up -V --force-recreate --wait --remove-orphans stop-postgres: docker compose -f tests/compose-postgres.yml down -v diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index d466a3c96..570c06bad 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -424,6 +424,10 @@ class Graph: class CompiledGraph(Pregel): builder: Graph + def __init__(self, *, builder: Graph, **kwargs): + super().__init__(**kwargs) + self.builder = builder + def attach_node(self, key: str, node: NodeSpec) -> None: self.channels[key] = EphemeralValue(Any) self.nodes[key] = ( diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 1416d14d0..307b672aa 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -17,12 +17,10 @@ from typing import ( overload, ) -from langchain_core.pydantic_v1 import BaseModel from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.base import RunnableLike -from langchain_core.runnables.utils import ( - create_model, -) +from langchain_core.runnables.utils import create_model +from pydantic import BaseModel from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate @@ -33,14 +31,7 @@ from langgraph.channels.named_barrier_value import NamedBarrierValue from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.constants import NS_END, NS_SEP, TAG_HIDDEN from langgraph.errors import InvalidUpdateError -from langgraph.graph.graph import ( - END, - START, - Branch, - CompiledGraph, - Graph, - Send, -) +from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send from langgraph.managed.base import ( ChannelKeyPlaceholder, ChannelTypePlaceholder, @@ -53,7 +44,7 @@ from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.types import All, RetryPolicy from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.utils import RunnableCallable, coerce_to_runnable +from langgraph.utils import RunnableCallable, coerce_to_runnable, get_field_default logger = logging.getLogger(__name__) @@ -498,7 +489,16 @@ class CompiledStateGraph(CompiledGraph): return create_model( # type: ignore[call-overload] self.get_name("Input"), **{ - k: (self.channels[k].UpdateType, None) + k: ( + self.channels[k].UpdateType, + ( + get_field_default( + k, + self.channels[k].UpdateType, + self.builder.input, + ) + ), + ) for k in self.builder.schemas[self.builder.input] if isinstance(self.channels[k], BaseChannel) }, diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 01f0ad91d..d7c265b2d 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -25,12 +25,10 @@ from uuid import UUID, uuid5 from langchain_core.globals import get_debug from langchain_core.load.dump import dumpd -from langchain_core.pydantic_v1 import BaseModel, Field, root_validator from langchain_core.runnables import ( Runnable, RunnableLambda, RunnableSequence, - RunnableSerializable, ) from langchain_core.runnables.base import Input, Output, coerce_to_runnable from langchain_core.runnables.config import ( @@ -48,6 +46,7 @@ from langchain_core.runnables.utils import ( get_unique_config_specs, ) from langchain_core.tracers._streaming import _StreamingCallbackHandler +from pydantic import BaseModel from typing_extensions import Self from langgraph.channels.base import ( @@ -186,16 +185,10 @@ class Channel: ) -class Pregel( - RunnableSerializable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]] -): +class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): nodes: Mapping[str, PregelNode] - channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] = Field( - default_factory=dict - ) - - auto_validate: bool = True + channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] stream_mode: StreamMode = "values" """Mode to stream output, defaults to 'values'.""" @@ -205,16 +198,16 @@ class Pregel( stream_channels: Optional[Union[str, Sequence[str]]] = None """Channels to stream, defaults to all channels not in reserved channels""" - interrupt_after_nodes: Union[All, Sequence[str]] = Field(default_factory=list) + interrupt_after_nodes: Union[All, Sequence[str]] - interrupt_before_nodes: Union[All, Sequence[str]] = Field(default_factory=list) + interrupt_before_nodes: Union[All, Sequence[str]] input_channels: Union[str, Sequence[str]] step_timeout: Optional[float] = None """Maximum time to wait for a step to complete, in seconds. Defaults to None.""" - debug: bool = Field(default_factory=get_debug) + debug: bool """Whether to print debug information during execution. Defaults to False.""" checkpointer: Optional[BaseCheckpointSaver] = None @@ -232,36 +225,50 @@ class Pregel( name: str = "LangGraph" - class Config: - arbitrary_types_allowed = True + def __init__( + self, + *, + nodes: Mapping[str, PregelNode], + channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] = None, + auto_validate: bool = True, + stream_mode: StreamMode = "values", + output_channels: Union[str, Sequence[str]], + stream_channels: Optional[Union[str, Sequence[str]]] = None, + interrupt_after_nodes: Union[All, Sequence[str]] = (), + interrupt_before_nodes: Union[All, Sequence[str]] = (), + input_channels: Union[str, Sequence[str]], + step_timeout: Optional[float] = None, + debug: Optional[bool] = None, + checkpointer: Optional[BaseCheckpointSaver] = None, + store: Optional[BaseStore] = None, + retry_policy: Optional[RetryPolicy] = None, + config_type: Optional[Type[Any]] = None, + config: Optional[RunnableConfig] = None, + name: str = "LangGraph", + ) -> None: + self.nodes = nodes + self.channels = channels or {} + self.stream_mode = stream_mode + self.output_channels = output_channels + self.stream_channels = stream_channels + self.interrupt_after_nodes = interrupt_after_nodes + self.interrupt_before_nodes = interrupt_before_nodes + self.input_channels = input_channels + self.step_timeout = step_timeout + self.debug = debug if debug is not None else get_debug() + self.checkpointer = checkpointer + self.store = store + self.retry_policy = retry_policy + self.config_type = config_type + self.config = config + self.name = name + if auto_validate: + self.validate() def with_config(self, config: RunnableConfig | None = None, **kwargs: Any) -> Self: - return self.copy( - update={"config": cast(RunnableConfig, {**(config or {}), **kwargs})} - ) - - @classmethod - def is_lc_serializable(cls) -> bool: - """Return whether the graph can be serialized by Langchain.""" - return True - - @root_validator(skip_on_failure=True) - def validate_on_init(cls, values: dict[str, Any]) -> dict[str, Any]: - if not values["auto_validate"]: - return values - validate_graph( - values["nodes"], - values["channels"], - values["input_channels"], - values["output_channels"], - values["stream_channels"], - values["interrupt_after_nodes"], - values["interrupt_before_nodes"], - ) - if values["interrupt_after_nodes"] or values["interrupt_before_nodes"]: - if not values["checkpointer"]: - raise ValueError("Interrupts require a checkpointer") - return values + attrs = {**self.__dict__} + attrs["config"] = merge_configs(self.config, config, kwargs) + return self.__class__(**attrs) def validate(self) -> Self: validate_graph( diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index 163665e2e..ca0828fda 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -1,8 +1,16 @@ from __future__ import annotations -from typing import Any, Callable, Mapping, Optional, Sequence, Union +from typing import ( + Any, + AsyncIterator, + Callable, + Iterator, + Mapping, + Optional, + Sequence, + Union, +) -from langchain_core.pydantic_v1 import Field from langchain_core.runnables import ( Runnable, RunnableConfig, @@ -10,7 +18,7 @@ from langchain_core.runnables import ( RunnableSequence, RunnableSerializable, ) -from langchain_core.runnables.base import Other, RunnableBindingBase, coerce_to_runnable +from langchain_core.runnables.base import Input, Other, Output, coerce_to_runnable from langchain_core.runnables.config import merge_configs from langchain_core.runnables.utils import ConfigurableFieldSpec @@ -99,20 +107,47 @@ class ChannelRead(RunnableCallable): DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough() -class PregelNode(RunnableBindingBase): +class PregelNode(Runnable): channels: Union[list[str], Mapping[str, str]] - triggers: list[str] = Field(default_factory=list) + triggers: list[str] - mapper: Optional[Callable[[Any], Any]] = None + mapper: Optional[Callable[[Any], Any]] - writers: list[Runnable] = Field(default_factory=list) + writers: list[Runnable] - bound: Runnable[Any, Any] = Field(default=DEFAULT_BOUND) + bound: Runnable[Any, Any] - kwargs: Mapping[str, Any] = Field(default_factory=dict) + retry_policy: Optional[RetryPolicy] - retry_policy: Optional[RetryPolicy] = None + config: RunnableConfig + + def __init__( + self, + *, + channels: Union[list[str], Mapping[str, str]], + triggers: Sequence[str], + mapper: Optional[Callable[[Any], Any]] = None, + writers: Optional[list[Runnable]] = None, + tags: Optional[list[str]] = None, + metadata: Optional[Mapping[str, Any]] = None, + bound: Optional[Runnable[Any, Any]] = None, + retry_policy: Optional[RetryPolicy] = None, + config: Optional[RunnableConfig] = None, + ) -> None: + self.channels = channels + self.triggers = list(triggers) + self.mapper = mapper + self.writers = writers or [] + self.bound = bound if bound is not None else DEFAULT_BOUND + self.retry_policy = retry_policy + self.config = merge_configs( + config, {"tags": tags or [], "metadata": metadata or {}} + ) + + def copy(self, update: dict[str, Any]) -> PregelNode: + attrs = {**self.__dict__, **update} + return PregelNode(**attrs) def get_writers(self) -> list[Runnable]: """Get writers with optimizations applied.""" @@ -145,38 +180,6 @@ class PregelNode(RunnableBindingBase): else: return self.bound - def __init__( - self, - *, - channels: Union[list[str], Mapping[str, str]], - triggers: Sequence[str], - mapper: Optional[Callable[[Any], Any]] = None, - writers: Optional[list[Runnable]] = None, - tags: Optional[list[str]] = None, - metadata: Optional[Mapping[str, Any]] = None, - bound: Optional[Runnable[Any, Any]] = None, - kwargs: Optional[Mapping[str, Any]] = None, - config: Optional[RunnableConfig] = None, - retry_policy: Optional[RetryPolicy] = None, - **other_kwargs: Any, - ) -> None: - super().__init__( - channels=channels, - triggers=triggers, - mapper=mapper, - writers=writers or [], - bound=bound or DEFAULT_BOUND, - kwargs=kwargs or {}, - retry_policy=retry_policy, - config=merge_configs( - config, {"tags": tags or [], "metadata": metadata or {}} - ), - **other_kwargs, - ) - - def __repr_args__(self) -> Any: - return [(k, v) for k, v in super().__repr_args__() if k != "bound"] - def join(self, channels: Sequence[str]) -> PregelNode: assert isinstance(channels, list) or isinstance( channels, tuple @@ -226,3 +229,42 @@ class PregelNode(RunnableBindingBase): ], ) -> RunnableSerializable: raise NotImplementedError() + + def invoke( + self, + input: Input, + config: Optional[RunnableConfig] = None, + **kwargs: Optional[Any], + ) -> Output: + return self.bound.invoke(input, merge_configs(self.config, config), **kwargs) + + async def ainvoke( + self, + input: Input, + config: Optional[RunnableConfig] = None, + **kwargs: Optional[Any], + ) -> Output: + return await self.bound.ainvoke( + input, merge_configs(self.config, config), **kwargs + ) + + def stream( + self, + input: Input, + config: Optional[RunnableConfig] = None, + **kwargs: Optional[Any], + ) -> Iterator[Output]: + yield from self.bound.stream( + input, merge_configs(self.config, config), **kwargs + ) + + async def astream( + self, + input: Input, + config: Optional[RunnableConfig] = None, + **kwargs: Optional[Any], + ) -> AsyncIterator[Output]: + async for item in self.bound.astream( + input, merge_configs(self.config, config), **kwargs + ): + yield item diff --git a/libs/langgraph/langgraph/utils.py b/libs/langgraph/langgraph/utils.py index 64cb21128..cc3fd4760 100644 --- a/libs/langgraph/langgraph/utils.py +++ b/libs/langgraph/langgraph/utils.py @@ -4,7 +4,7 @@ import inspect import sys from contextvars import copy_context from functools import partial, wraps -from typing import Any, AsyncIterator, Awaitable, Callable, Optional +from typing import Any, AsyncIterator, Awaitable, Callable, Optional, Type, Union from langchain_core.runnables.base import ( Runnable, @@ -19,7 +19,14 @@ from langchain_core.runnables.config import ( var_child_runnable_config, ) from langchain_core.runnables.utils import accepts_config -from typing_extensions import TypeGuard +from typing_extensions import ( + Annotated, + NotRequired, + ReadOnly, + Required, + TypeGuard, + get_origin, +) try: from langchain_core.runnables.config import _set_config_context @@ -34,8 +41,6 @@ except ImportError: class StrEnum(str, enum.Enum): """A string enum.""" - pass - class RunnableCallable(Runnable): """A much simpler version of RunnableLambda that requires sync and async functions.""" @@ -183,3 +188,95 @@ def coerce_to_runnable(thing: RunnableLike, *, name: str, trace: bool) -> Runnab f"Expected a Runnable, callable or dict." f"Instead got an unsupported type: {type(thing)}" ) + + +def _is_optional_type(type_: Any) -> bool: + """Check if a type is Optional.""" + + if hasattr(type_, "__origin__") and hasattr(type_, "__args__"): + origin = get_origin(type_) + if origin is Optional: + return True + if origin is Union: + return any( + arg is type(None) or _is_optional_type(arg) for arg in type_.__args__ + ) + if origin is Annotated: + return _is_optional_type(type_.__args__[0]) + return origin is None + if hasattr(type_, "__bound__") and type_.__bound__ is not None: + return _is_optional_type(type_.__bound__) + return type_ is None + + +def _is_required_type(type_: Any) -> Optional[bool]: + """Check if an annotation is marked as Required/NotRequired. + + Returns: + - True if required + - False if not required + - None if not annotated with either + """ + origin = get_origin(type_) + if origin is Required: + return True + if origin is NotRequired: + return False + if origin is Annotated or getattr(origin, "__args__", None): + # See https://typing.readthedocs.io/en/latest/spec/typeddict.html#interaction-with-annotated + return _is_required_type(type_.__args__[0]) + return None + + +def _is_readonly_type(type_: Any) -> bool: + """Check if an annotation is marked as ReadOnly. + + Returns: + - True if is read only + - False if not read only + """ + + # See: https://typing.readthedocs.io/en/latest/spec/typeddict.html#typing-readonly-type-qualifier + origin = get_origin(type_) + if origin is Annotated: + return _is_readonly_type(type_.__args__[0]) + if origin is ReadOnly: + return True + return False + + +_DEFAULT_KEYS = frozenset() + + +def get_field_default(name: str, type_: Any, schema: Type[Any]) -> Any: + """Determine the default value for a field in a state schema. + + This is based on: + If TypedDict: + - Required/NotRequired + - total=False -> everything optional + - Type annotation (Optional/Union[None]) + """ + optional_keys = getattr(schema, "__optional_keys__", _DEFAULT_KEYS) + irq = _is_required_type(type_) + if name in optional_keys: + # Either total=False or explicit NotRequired. + # No type annotation trumps this. + if irq: + # Unless it's earlier versions of python & explicit Required + return ... + return None + if irq is not None: + if irq: + # Handle Required[] + # (we already handled NotRequired and total=False) + return ... + # Handle NotRequired[] for earlier versions of python + return None + # Note, we ignore ReadOnly attributes, + # as they don't make much sense. (we don't care if you mutate the state in your node) + # and mutating state in your node has no effect on our graph state. + # Base case is the annotation + if _is_optional_type(type_): + return None + return ... diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index 241d47b6e..50d6c6f89 100644 --- a/libs/langgraph/poetry.lock +++ b/libs/langgraph/poetry.lock @@ -2829,13 +2829,13 @@ diagrams = ["jinja2", "railroad-diagrams"] [[package]] name = "pytest" -version = "7.4.4" +version = "8.3.2" description = "pytest: simple powerful testing with Python" optional = false -python-versions = ">=3.7" +python-versions = ">=3.8" files = [ - {file = "pytest-7.4.4-py3-none-any.whl", hash = "sha256:b090cdf5ed60bf4c45261be03239c2c1c22df034fbffe691abe93cd80cea01d8"}, - {file = "pytest-7.4.4.tar.gz", hash = "sha256:2cf0005922c6ace4a3e2ec8b4080eb0d9753fdc93107415332f50ce9e7994280"}, + {file = "pytest-8.3.2-py3-none-any.whl", hash = "sha256:4ba08f9ae7dcf84ded419494d229b48d0903ea6407b030eaec46df5e6a73bba5"}, + {file = "pytest-8.3.2.tar.gz", hash = "sha256:c132345d12ce551242c87269de812483f5bcc87cdbb4722e48487ba194f9fdce"}, ] [package.dependencies] @@ -2843,29 +2843,11 @@ colorama = {version = "*", markers = "sys_platform == \"win32\""} exceptiongroup = {version = ">=1.0.0rc8", markers = "python_version < \"3.11\""} iniconfig = "*" packaging = "*" -pluggy = ">=0.12,<2.0" -tomli = {version = ">=1.0.0", markers = "python_version < \"3.11\""} +pluggy = ">=1.5,<2" +tomli = {version = ">=1", markers = "python_version < \"3.11\""} [package.extras] -testing = ["argcomplete", "attrs (>=19.2.0)", "hypothesis (>=3.56)", "mock", "nose", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"] - -[[package]] -name = "pytest-asyncio" -version = "0.20.3" -description = "Pytest support for asyncio" -optional = false -python-versions = ">=3.7" -files = [ - {file = "pytest-asyncio-0.20.3.tar.gz", hash = "sha256:83cbf01169ce3e8eb71c6c278ccb0574d1a7a3bb8eaaf5e50e0ad342afb33b36"}, - {file = "pytest_asyncio-0.20.3-py3-none-any.whl", hash = "sha256:f129998b209d04fcc65c96fc85c11e5316738358909a8399e93be553d7656442"}, -] - -[package.dependencies] -pytest = ">=6.1.0" - -[package.extras] -docs = ["sphinx (>=5.3)", "sphinx-rtd-theme (>=1.0)"] -testing = ["coverage (>=6.2)", "flaky (>=3.5.0)", "hypothesis (>=5.7.1)", "mypy (>=0.931)", "pytest-trio (>=0.7.0)"] +dev = ["argcomplete", "attrs (>=19.2)", "hypothesis (>=3.56)", "mock", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"] [[package]] name = "pytest-cov" @@ -3069,7 +3051,6 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, @@ -4310,4 +4291,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools", [metadata] lock-version = "2.0" python-versions = ">=3.9.0,<4.0" -content-hash = "7e0d6fd967987fcf46b038c47b5ca81d86cb8fff89704eac2074f2ce55336a2b" +content-hash = "e0787721e6cb80996a39284c74f1ddee4789c6060d6c2272cb1957f4686bd478" diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 36cdb7fdc..d6e04486f 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph" -version = "0.2.15" +version = "0.2.16" description = "Building stateful, multi-actor applications with LLMs" authors = [] license = "MIT" @@ -14,10 +14,9 @@ langgraph-checkpoint = "^1.0.2" [tool.poetry.group.dev.dependencies] -pytest = "^7.3.0" +pytest = "^8.3.2" pytest-cov = "^4.0.0" pytest-dotenv = "^0.5.2" -pytest-asyncio = "^0.20.3" pytest-mock = "^3.10.0" syrupy = "^4.0.2" httpx = "^0.26.0" @@ -73,7 +72,6 @@ requires = ["poetry-core>=1.0.0"] build-backend = "poetry.core.masonry.api" [tool.pytest.ini_options] -asyncio_mode = "auto" # --strict-markers will raise errors on unknown marks. # https://docs.pytest.org/en/7.1.x/how-to/mark.html#raising-errors-on-unknown-marks # diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 14cf25277..3937546f5 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -366,7 +366,7 @@ ''' # --- # name: test_conditional_entrypoint_to_multiple_state_graph - '{"title": "LangGraphInput", "type": "object", "properties": {"locations": {"title": "Locations", "type": "array", "items": {"type": "string"}}, "results": {"title": "Results", "type": "array", "items": {"type": "string"}}}}' + '{"title": "LangGraphInput", "type": "object", "properties": {"locations": {"title": "Locations", "type": "array", "items": {"type": "string"}}, "results": {"title": "Results", "type": "array", "items": {"type": "string"}}}, "required": ["locations", "results"]}' # --- # name: test_conditional_entrypoint_to_multiple_state_graph.1 '{"title": "LangGraphOutput", "type": "object", "properties": {"locations": {"title": "Locations", "type": "array", "items": {"type": "string"}}, "results": {"title": "Results", "type": "array", "items": {"type": "string"}}}}' @@ -3304,7 +3304,7 @@ 'query', 'inner', ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3323,11 +3323,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- @@ -3374,7 +3370,7 @@ 'query', 'inner', ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3393,11 +3389,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- @@ -3444,7 +3436,7 @@ 'query', 'inner', ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3463,11 +3455,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- @@ -3514,7 +3502,7 @@ 'query', 'inner', ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3533,11 +3521,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- @@ -3584,7 +3568,7 @@ 'query', 'inner', ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3603,11 +3587,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- @@ -4855,7 +4835,7 @@ ''' # --- # name: test_prebuilt_tool_chat - '{"title": "LangGraphInput", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}}}' + '{"title": "LangGraphInput", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"], "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}}}' # --- # name: test_prebuilt_tool_chat.1 '{"title": "LangGraphOutput", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}}}' @@ -5094,7 +5074,7 @@ '{"title": "LangGraphConfig", "type": "object", "properties": {"configurable": {"$ref": "#/definitions/Configurable"}}, "definitions": {"Configurable": {"title": "Configurable", "type": "object", "properties": {"tools": {"title": "Tools", "type": "array", "items": {"type": "string"}}}}}}' # --- # name: test_state_graph_w_config_inherited_state_keys.1 - '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' + '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "required": ["input"], "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' # --- # name: test_state_graph_w_config_inherited_state_keys.2 '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' diff --git a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr index c0d2e782c..3aaccffd1 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr @@ -36,6 +36,191 @@ +---------+ ''' # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[memory] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[postgres_aio] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[postgres_aio_pipe] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[postgres_aio_pool] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[sqlite_aio] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- # name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2 ''' graph TD; @@ -157,6 +342,611 @@ 'type': 'object', }) # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory].2 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio].2 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pipe] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pipe].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pipe].2 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pool] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pool].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pool].2 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio].2 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- # name: test_in_one_fan_out_state_graph_waiting_edge_via_branch ''' +-----------+ @@ -194,6 +984,191 @@ +---------+ ''' # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[memory] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_aio] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_aio_pipe] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_aio_pool] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[sqlite_aio] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- # name: test_nested_graph ''' +-----------+ @@ -219,3 +1194,128 @@ +---------+ ''' # --- +# name: test_weather_subgraph[memory] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node
__interrupt = before) + __end__([__end__]):::last + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[postgres_aio] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node
__interrupt = before) + __end__([__end__]):::last + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[postgres_aio_pipe] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node
__interrupt = before) + __end__([__end__]):::last + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[postgres_aio_pool] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node
__interrupt = before) + __end__([__end__]):::last + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[sqlite_aio] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node
__interrupt = before) + __end__([__end__]):::last + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- diff --git a/libs/langgraph/tests/checkpoint/__init__.py b/libs/langgraph/tests/checkpoint/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index bb6ef9603..ca6351c6b 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -1,8 +1,6 @@ -import asyncio import sys -from concurrent.futures import ThreadPoolExecutor -from contextlib import asynccontextmanager, contextmanager -from typing import AsyncIterator, Iterator, TypeVar +from contextlib import asynccontextmanager +from typing import AsyncIterator, Optional from uuid import UUID, uuid4 import pytest @@ -10,6 +8,7 @@ from psycopg import AsyncConnection, Connection from psycopg_pool import AsyncConnectionPool, ConnectionPool from pytest_mock import MockerFixture +from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.checkpoint.postgres import PostgresSaver from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver from langgraph.checkpoint.sqlite import SqliteSaver @@ -19,6 +18,11 @@ from tests.memory_assert import MemorySaverAssertImmutable DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/" +@pytest.fixture +def anyio_backend(): + return "asyncio" + + @pytest.fixture() def deterministic_uuids(mocker: MockerFixture) -> MockerFixture: side_effect = ( @@ -27,35 +31,6 @@ def deterministic_uuids(mocker: MockerFixture) -> MockerFixture: return mocker.patch("uuid.uuid4", side_effect=side_effect) -""" -pytest-asyncio doesn't support calling async fixtures with getfixturevalue -so we need to use ThreadPoolExecutor to run the async fixture in a thread -https://github.com/pytest-dev/pytest-asyncio/issues/112#issuecomment-462062890 -""" -T = TypeVar("T") - - -def close_loop(loop: asyncio.AbstractEventLoop) -> None: - loop.run_until_complete(loop.shutdown_asyncgens()) - loop.run_until_complete(loop.shutdown_default_executor()) - asyncio.set_event_loop(None) - loop.close() - - -@contextmanager -def agen_to_gen(agen: AsyncIterator[T]) -> Iterator[T]: - with ThreadPoolExecutor(1) as bg: - loop = asyncio.new_event_loop() - bg.submit(asyncio.set_event_loop, loop).result() - try: - yield bg.submit(loop.run_until_complete, agen.__aenter__()).result() - finally: - bg.submit( - loop.run_until_complete, agen.__aexit__(None, None, None) - ).result() - bg.submit(close_loop, loop).result() - - # checkpointer fixtures @@ -70,12 +45,6 @@ def checkpointer_sqlite(): yield checkpointer -@pytest.fixture(scope="function") -def checkpointer_sqlite_aio(): - with agen_to_gen(_checkpointer_sqlite_aio()) as checkpointer: - yield checkpointer - - @asynccontextmanager async def _checkpointer_sqlite_aio(): async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer: @@ -143,16 +112,10 @@ def checkpointer_postgres_pool(): conn.execute(f"DROP DATABASE {database}") -@pytest.fixture(scope="function") -def checkpointer_postgres_aio(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - with agen_to_gen(_checkpointer_postgres_aio()) as checkpointer: - yield checkpointer - - @asynccontextmanager async def _checkpointer_postgres_aio(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") database = f"test_{uuid4().hex[:16]}" # create unique db async with await AsyncConnection.connect( @@ -174,16 +137,10 @@ async def _checkpointer_postgres_aio(): await conn.execute(f"DROP DATABASE {database}") -@pytest.fixture(scope="function") -def checkpointer_postgres_aio_pipe(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - with agen_to_gen(_checkpointer_postgres_aio_pipe()) as checkpointer: - yield checkpointer - - @asynccontextmanager async def _checkpointer_postgres_aio_pipe(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") database = f"test_{uuid4().hex[:16]}" # create unique db async with await AsyncConnection.connect( @@ -208,16 +165,10 @@ async def _checkpointer_postgres_aio_pipe(): await conn.execute(f"DROP DATABASE {database}") -@pytest.fixture(scope="function") -def checkpointer_postgres_aio_pool(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - with agen_to_gen(_checkpointer_postgres_aio_pool()) as checkpointer: - yield checkpointer - - @asynccontextmanager async def _checkpointer_postgres_aio_pool(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") database = f"test_{uuid4().hex[:16]}" # create unique db async with await AsyncConnection.connect( @@ -240,6 +191,30 @@ async def _checkpointer_postgres_aio_pool(): await conn.execute(f"DROP DATABASE {database}") +@asynccontextmanager +async def awith_checkpointer( + checkpointer_name: Optional[str], +) -> AsyncIterator[BaseCheckpointSaver]: + if checkpointer_name is None: + yield None + elif checkpointer_name == "memory": + yield MemorySaverAssertImmutable() + elif checkpointer_name == "sqlite_aio": + async with _checkpointer_sqlite_aio() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_aio": + async with _checkpointer_postgres_aio() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_aio_pipe": + async with _checkpointer_postgres_aio_pipe() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_aio_pool": + async with _checkpointer_postgres_aio_pool() as checkpointer: + yield checkpointer + else: + raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}") + + ALL_CHECKPOINTERS_SYNC = [ "memory", "sqlite", @@ -254,3 +229,7 @@ ALL_CHECKPOINTERS_ASYNC = [ "postgres_aio_pipe", "postgres_aio_pool", ] +ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [ + *ALL_CHECKPOINTERS_ASYNC, + None, +] diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index ff393577e..69a624969 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -8,6 +8,8 @@ from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.errors import EmptyChannelError, InvalidUpdateError +pytestmark = pytest.mark.anyio + def test_last_value() -> None: with LastValue(int).from_checkpoint(None, {}) as channel: diff --git a/libs/langgraph/tests/test_interruption.py b/libs/langgraph/tests/test_interruption.py index e6682e964..d4e618ddf 100644 --- a/libs/langgraph/tests/test_interruption.py +++ b/libs/langgraph/tests/test_interruption.py @@ -4,12 +4,16 @@ import pytest from pytest_mock import MockerFixture from langgraph.graph import END, START, StateGraph - - -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], +from tests.conftest import ( + ALL_CHECKPOINTERS_ASYNC, + ALL_CHECKPOINTERS_SYNC, + awith_checkpointer, ) + +pytestmark = pytest.mark.anyio + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_interruption_without_state_updates( request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture ) -> None: @@ -47,12 +51,9 @@ def test_interruption_without_state_updates( assert graph.get_state(thread).next == () -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_interruption_without_state_updates_async( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture + checkpointer_name: str, mocker: MockerFixture ): """Test interruption without state updates. This test confirms that interrupting doesn't require a state key having been updated in the prev step""" @@ -72,17 +73,17 @@ async def test_interruption_without_state_updates_async( builder.add_edge("step_2", "step_3") builder.add_edge("step_3", END) - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - graph = builder.compile(checkpointer=checkpointer, interrupt_after="*") + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer, interrupt_after="*") - initial_input = {"input": "hello world"} - thread = {"configurable": {"thread_id": "1"}} + initial_input = {"input": "hello world"} + thread = {"configurable": {"thread_id": "1"}} - await graph.ainvoke(initial_input, thread, debug=True) - assert (await graph.aget_state(thread)).next == ("step_2",) + await graph.ainvoke(initial_input, thread, debug=True) + assert (await graph.aget_state(thread)).next == ("step_2",) - await graph.ainvoke(None, thread, debug=True) - assert (await graph.aget_state(thread)).next == ("step_3",) + await graph.ainvoke(None, thread, debug=True) + assert (await graph.aget_state(thread)).next == ("step_3",) - await graph.ainvoke(None, thread, debug=True) - assert (await graph.aget_state(thread)).next == () + await graph.ainvoke(None, thread, debug=True) + assert (await graph.aget_state(thread)).next == () diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index b29074dff..82fd5e613 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -23,8 +23,15 @@ from pydantic import BaseModel as BaseModelV2 from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.prebuilt import ToolNode, ValidationNode, create_react_agent from langgraph.prebuilt.tool_node import InjectedState +from tests.conftest import ( + ALL_CHECKPOINTERS_ASYNC, + ALL_CHECKPOINTERS_SYNC, + awith_checkpointer, +) from tests.messages import _AnyIdHumanMessage +pytestmark = pytest.mark.anyio + class FakeToolCallingModel(BaseChatModel): def _generate( @@ -53,10 +60,7 @@ class FakeToolCallingModel(BaseChatModel): return self -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_no_modifier(request: pytest.FixtureRequest, checkpointer_name: str) -> None: checkpointer: BaseCheckpointSaver = request.getfixturevalue( "checkpointer_" + checkpointer_name @@ -89,43 +93,35 @@ def test_no_modifier(request: pytest.FixtureRequest, checkpointer_name: str) -> assert saved.pending_writes == [] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) -async def test_no_modifier_async( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_no_modifier_async(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + model = FakeToolCallingModel() - model = FakeToolCallingModel() + agent = create_react_agent(model, [], checkpointer=checkpointer) + inputs = [HumanMessage("hi?")] + thread = {"configurable": {"thread_id": "123"}} + response = await agent.ainvoke({"messages": inputs}, thread, debug=True) + expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]} + assert response == expected_response - agent = create_react_agent(model, [], checkpointer=checkpointer) - inputs = [HumanMessage("hi?")] - thread = {"configurable": {"thread_id": "123"}} - response = await agent.ainvoke({"messages": inputs}, thread, debug=True) - expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]} - assert response == expected_response - - if checkpointer: - saved = await checkpointer.aget_tuple(thread) - assert saved is not None - assert saved.checkpoint["channel_values"] == { - "messages": [ - _AnyIdHumanMessage(content="hi?"), - AIMessage(content="hi?", id="0"), - ], - "agent": "agent", - } - assert saved.metadata == { - "parents": {}, - "source": "loop", - "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}}, - "step": 1, - } - assert saved.pending_writes == [] + if checkpointer: + saved = await checkpointer.aget_tuple(thread) + assert saved is not None + assert saved.checkpoint["channel_values"] == { + "messages": [ + _AnyIdHumanMessage(content="hi?"), + AIMessage(content="hi?", id="0"), + ], + "agent": "agent", + } + assert saved.metadata == { + "parents": {}, + "source": "loop", + "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}}, + "step": 1, + } + assert saved.pending_writes == [] def test_passing_two_modifiers(): diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 74a16b010..0a41b3dbc 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -74,10 +74,7 @@ from langgraph.store.memory import MemoryStore from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence from tests.conftest import ALL_CHECKPOINTERS_SYNC from tests.fake_tracer import FakeTracer -from tests.memory_assert import ( - MemorySaverAssertCheckpointMetadata, - MemorySaverNoPending, -) +from tests.memory_assert import MemorySaverAssertCheckpointMetadata from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage @@ -1626,29 +1623,6 @@ def test_cond_edge_after_send() -> None: assert graph.invoke(["0"]) == ["0", "1", "2", "2", "3"] -async def test_checkpointer_null_pending_writes() -> None: - class Node: - def __init__(self, name: str): - self.name = name - setattr(self, "__name__", name) - - def __call__(self, state): - return [self.name] - - builder = StateGraph(Annotated[list, operator.add]) - builder.add_node(Node("1")) - builder.add_edge(START, "1") - graph = builder.compile(checkpointer=MemorySaverNoPending()) - assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"] - assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"] * 2 - assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [ - "1" - ] * 3 - assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [ - "1" - ] * 4 - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_invoke_checkpoint_three( mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str @@ -1929,7 +1903,7 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None: one = Channel.subscribe_to("between") | add_one | Channel.write_to("output") two = Channel.subscribe_to("between") | add_one - with pytest.raises(ValueError): + with pytest.raises(TypeError): Pregel(nodes={"one": one, "two": two}) @@ -9969,10 +9943,12 @@ def test_send_to_nested_graphs( graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"}) # continue past interrupt - assert graph.invoke(None, config=config) == { - "subjects": ["cats", "dogs"], - "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - } + assert sorted( + graph.stream(None, config=config), key=lambda d: d["generate_joke"]["jokes"][0] + ) == [ + {"generate_joke": {"jokes": ["Joke about cats - hohoho"]}}, + {"generate_joke": {"jokes": ["Joke about turtles - hohoho"]}}, + ] actual_snapshot = graph.get_state(config) expected_snapshot = StateSnapshot( diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 70d90bb62..3da3b0886 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2,6 +2,7 @@ import asyncio import json import operator import re +import sys from collections import Counter from contextlib import asynccontextmanager, contextmanager from typing import ( @@ -17,6 +18,7 @@ from typing import ( Tuple, TypedDict, Union, + cast, ) from uuid import UUID @@ -43,7 +45,6 @@ from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.channels.untracked_value import UntrackedValue from langgraph.checkpoint.base import ( - BaseCheckpointSaver, ChannelVersions, Checkpoint, CheckpointMetadata, @@ -71,14 +72,20 @@ from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.types import PregelTask from langgraph.store.memory import MemoryStore from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence -from tests.conftest import ALL_CHECKPOINTERS_ASYNC +from tests.conftest import ( + ALL_CHECKPOINTERS_ASYNC, + ALL_CHECKPOINTERS_ASYNC_PLUS_NONE, + awith_checkpointer, +) from tests.fake_tracer import FakeTracer from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, - MemorySaverAssertImmutable, + MemorySaverNoPending, ) from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage +pytestmark = pytest.mark.anyio + async def test_checkpoint_errors() -> None: class FaultyGetCheckpointer(MemorySaver): @@ -217,11 +224,7 @@ async def test_node_cancellation_on_other_node_exception() -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_dynamic_interrupt( - checkpointer_name: str, snapshot: SnapshotAssertion, request: pytest.FixtureRequest -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +async def test_dynamic_interrupt(checkpointer_name: str) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -259,59 +262,60 @@ async def test_dynamic_interrupt( "market": "US", } - tool_two = tool_two_graph.compile(checkpointer=checkpointer) + async with awith_checkpointer(checkpointer_name) as checkpointer: + tool_two = tool_two_graph.compile(checkpointer=checkpointer) - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - thread1 = {"configurable": {"thread_id": "1"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == { - "my_key": "value ⛰️", - "market": "DE", - } - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - }, - ] - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=("tool_two",), - tasks=( - PregelTask( - AnyStr(), - "tool_two", - interrupts=(Interrupt("Just because..."),), + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert await tool_two.ainvoke( + {"my_key": "value ⛰️", "market": "DE"}, thread1 + ) == { + "my_key": "value ⛰️", + "market": "DE", + } + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + }, + ] + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + interrupts=(Interrupt("Just because..."),), + ), ), - ), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, - parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config, - ) - # TODO use aget_state_history + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + # TODO use aget_state_history @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_node_not_cancelled_on_other_node_interrupted( - checkpointer_name: str, request: pytest.FixtureRequest + checkpointer_name: str, ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict): hello: str @@ -339,23 +343,24 @@ async def test_node_not_cancelled_on_other_node_interrupted( builder.add_node("bad", iambad) builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END) - graph = builder.compile(checkpointer=checkpointer) - thread = {"configurable": {"thread_id": "1"}} + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + thread = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke({"hello": "world"}, thread) == {"hello": "world"} + assert await graph.ainvoke({"hello": "world"}, thread) == {"hello": "world"} - assert not inner_task_cancelled - assert awhiles == 1 + assert not inner_task_cancelled + assert awhiles == 1 - assert await graph.ainvoke(None, thread, debug=True) == {"hello": "world"} + assert await graph.ainvoke(None, thread, debug=True) == {"hello": "world"} - assert not inner_task_cancelled - assert awhiles == 1 + assert not inner_task_cancelled + assert awhiles == 1 - assert await graph.ainvoke({"hello": "bye"}, thread) == {"hello": "again"} + assert await graph.ainvoke({"hello": "bye"}, thread) == {"hello": "again"} - assert not inner_task_cancelled - assert awhiles == 2 + assert not inner_task_cancelled + assert awhiles == 2 async def test_step_timeout_on_stream_hang() -> None: @@ -388,12 +393,8 @@ async def test_step_timeout_on_stream_hang() -> None: assert inner_task_cancelled -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_cancel_graph_astream( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC_PLUS_NONE) +async def test_cancel_graph_astream(checkpointer_name: str) -> None: class State(TypedDict): value: Annotated[int, operator.add] @@ -426,53 +427,47 @@ async def test_cancel_graph_astream( builder.add_edge(START, "alittlewhile") builder.add_edge(START, "aparallelwhile") builder.add_edge("alittlewhile", "awhile") - graph = builder.compile(checkpointer=checkpointer) - # test interrupting astream - got_event = False - thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} - async with aclosing(graph.astream({"value": 1}, thread1)) as stream: - async for chunk in stream: - assert chunk == {"alittlewhile": {"value": 2}} - got_event = True - break + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) - assert got_event + # test interrupting astream + got_event = False + thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} + async with aclosing(graph.astream({"value": 1}, thread1)) as stream: + async for chunk in stream: + assert chunk == {"alittlewhile": {"value": 2}} + got_event = True + break - # node aparallelwhile should start, but be cancelled - assert aparallelwhile.started is True - assert aparallelwhile.cancelled is True + assert got_event - # node "awhile" should never start - assert awhile.started is False + # node aparallelwhile should start, but be cancelled + assert aparallelwhile.started is True + assert aparallelwhile.cancelled is True - # checkpoint with output of "alittlewhile" should not be saved - if checkpointer is not None: - state = await graph.aget_state(thread1) - assert state is not None - assert state.values == {"value": 1} - assert state.next == ( - "aparallelwhile", - "alittlewhile", - ) - assert state.metadata == { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - } + # node "awhile" should never start + assert awhile.started is False + + # checkpoint with output of "alittlewhile" should not be saved + if checkpointer is not None: + state = await graph.aget_state(thread1) + assert state is not None + assert state.values == {"value": 1} + assert state.next == ( + "aparallelwhile", + "alittlewhile", + ) + assert state.metadata == { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + } -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_cancel_graph_astream_events_v2( - request: pytest.FixtureRequest, checkpointer_name: Optional[str] -) -> None: - checkpointer = ( - request.getfixturevalue(f"checkpointer_{checkpointer_name}") - if checkpointer_name - else None - ) - +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC_PLUS_NONE) +async def test_cancel_graph_astream_events_v2(checkpointer_name: Optional[str]) -> None: class State(TypedDict): value: int @@ -505,43 +500,45 @@ async def test_cancel_graph_astream_events_v2( builder.add_edge(START, "alittlewhile") builder.add_edge("alittlewhile", "awhile") builder.add_edge("awhile", "anotherwhile") - graph = builder.compile(checkpointer=checkpointer) - # test interrupting astream_events v2 - got_event = False - thread2: RunnableConfig = {"configurable": {"thread_id": "2"}} - async with aclosing( - graph.astream_events({"value": 1}, thread2, version="v2") - ) as stream: - async for chunk in stream: - if chunk["event"] == "on_chain_stream" and not chunk["parent_ids"]: - got_event = True - assert chunk["data"]["chunk"] == {"alittlewhile": {"value": 2}} - break + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) - # did break - assert got_event + # test interrupting astream_events v2 + got_event = False + thread2: RunnableConfig = {"configurable": {"thread_id": "2"}} + async with aclosing( + graph.astream_events({"value": 1}, thread2, version="v2") + ) as stream: + async for chunk in stream: + if chunk["event"] == "on_chain_stream" and not chunk["parent_ids"]: + got_event = True + assert chunk["data"]["chunk"] == {"alittlewhile": {"value": 2}} + break - # node "awhile" maybe starts (impl detail of astream_events) - # if it does start, it must be cancelled - if awhile.started: - assert awhile.cancelled is True + # did break + assert got_event - # node "anotherwhile" should never start - assert anotherwhile.started is False + # node "awhile" maybe starts (impl detail of astream_events) + # if it does start, it must be cancelled + if awhile.started: + assert awhile.cancelled is True - # checkpoint with output of "alittlewhile" should not be saved - if checkpointer is not None: - state = await graph.aget_state(thread2) - assert state is not None - assert state.values == {"value": 2} - assert state.next == ("awhile",) - assert state.metadata == { - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"alittlewhile": {"value": 2}}, - } + # node "anotherwhile" should never start + assert anotherwhile.started is False + + # checkpoint with output of "alittlewhile" should not be saved + if checkpointer is not None: + state = await graph.aget_state(thread2) + assert state is not None + assert state.values == {"value": 2} + assert state.next == ("awhile",) + assert state.metadata == { + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"alittlewhile": {"value": 2}}, + } async def test_node_schemas_custom_output() -> None: @@ -819,424 +816,438 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_invoke_two_processes_in_out_interrupt( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture + checkpointer_name: str, mocker: MockerFixture ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=checkpointer, - interrupt_after_nodes=["one"], - ) - thread1 = {"configurable": {"thread_id": "1"}} - thread2 = {"configurable": {"thread_id": "2"}} + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=checkpointer, + interrupt_after_nodes=["one"], + ) + thread1 = {"configurable": {"thread_id": "1"}} + thread2 = {"configurable": {"thread_id": "2"}} - # start execution, stop at inbox - assert await app.ainvoke(2, thread1) is None + # start execution, stop at inbox + assert await app.ainvoke(2, thread1) is None - # inbox == 3 - checkpoint = await checkpointer.aget(thread1) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 3 + # inbox == 3 + checkpoint = await checkpointer.aget(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 3 - # resume execution, finish - assert await app.ainvoke(None, thread1) == 4 + # resume execution, finish + assert await app.ainvoke(None, thread1) == 4 - # start execution again, stop at inbox - assert await app.ainvoke(20, thread1) is None + # start execution again, stop at inbox + assert await app.ainvoke(20, thread1) is None - # inbox == 21 - checkpoint = await checkpointer.aget(thread1) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 21 + # inbox == 21 + checkpoint = await checkpointer.aget(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 21 - # send a new value in, interrupting the previous execution - assert await app.ainvoke(3, thread1) is None - assert await app.ainvoke(None, thread1) == 5 + # send a new value in, interrupting the previous execution + assert await app.ainvoke(3, thread1) is None + assert await app.ainvoke(None, thread1) == 5 - # start execution again, stopping at inbox - assert await app.ainvoke(20, thread2) is None + # start execution again, stopping at inbox + assert await app.ainvoke(20, thread2) is None - # inbox == 21 - snapshot = await app.aget_state(thread2) - assert snapshot.values["inbox"] == 21 - assert snapshot.next == ("two",) + # inbox == 21 + snapshot = await app.aget_state(thread2) + assert snapshot.values["inbox"] == 21 + assert snapshot.next == ("two",) - # update the state, resume - await app.aupdate_state(thread2, 25, as_node="one") - assert await app.ainvoke(None, thread2) == 26 + # update the state, resume + await app.aupdate_state(thread2, 25, as_node="one") + assert await app.ainvoke(None, thread2) == 26 - # no pending tasks - snapshot = await app.aget_state(thread2) - assert snapshot.next == () + # no pending tasks + snapshot = await app.aget_state(thread2) + assert snapshot.next == () - # list history - history = [c async for c in app.aget_state_history(thread1)] - assert history == [ - StateSnapshot( - values={"inbox": 4, "output": 5, "input": 3}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "step": 6, "writes": {"two": 5}}, - created_at=AnyStr(), - parent_config=history[1].config, - ), - StateSnapshot( - values={"inbox": 4, "output": 4, "input": 3}, - tasks=(PregelTask(AnyStr(), "two"),), - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 5, - "writes": {"one": None}, - }, - created_at=AnyStr(), - parent_config=history[2].config, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 3}, - tasks=(PregelTask(AnyStr(), "one"),), - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": 4, - "writes": {"input": 3}, - }, - created_at=AnyStr(), - parent_config=history[3].config, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 20}, - tasks=(PregelTask(AnyStr(), "two"),), - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"one": None}, - }, - created_at=AnyStr(), - parent_config=history[4].config, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 20}, - tasks=(PregelTask(AnyStr(), "one"),), - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": 2, - "writes": {"input": 20}, - }, - created_at=AnyStr(), - parent_config=history[5].config, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 2}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "step": 1, "writes": {"two": 4}}, - created_at=AnyStr(), - parent_config=history[6].config, - ), - StateSnapshot( - values={"inbox": 3, "input": 2}, - tasks=(PregelTask(AnyStr(), "two"),), - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": {"one": None}, - }, - created_at=AnyStr(), - parent_config=history[7].config, - ), - StateSnapshot( - values={"input": 2}, - tasks=(PregelTask(AnyStr(), "one"),), - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": -1, - "writes": {"input": 2}, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] + # list history + history = [c async for c in app.aget_state_history(thread1)] + assert history == [ + StateSnapshot( + values={"inbox": 4, "output": 5, "input": 3}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 6, + "writes": {"two": 5}, + }, + created_at=AnyStr(), + parent_config=history[1].config, + ), + StateSnapshot( + values={"inbox": 4, "output": 4, "input": 3}, + tasks=(PregelTask(AnyStr(), "two"),), + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 5, + "writes": {"one": None}, + }, + created_at=AnyStr(), + parent_config=history[2].config, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 3}, + tasks=(PregelTask(AnyStr(), "one"),), + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": 4, + "writes": {"input": 3}, + }, + created_at=AnyStr(), + parent_config=history[3].config, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 20}, + tasks=(PregelTask(AnyStr(), "two"),), + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"one": None}, + }, + created_at=AnyStr(), + parent_config=history[4].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 20}, + tasks=(PregelTask(AnyStr(), "one"),), + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": 2, + "writes": {"input": 20}, + }, + created_at=AnyStr(), + parent_config=history[5].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 2}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"two": 4}, + }, + created_at=AnyStr(), + parent_config=history[6].config, + ), + StateSnapshot( + values={"inbox": 3, "input": 2}, + tasks=(PregelTask(AnyStr(), "two"),), + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": {"one": None}, + }, + created_at=AnyStr(), + parent_config=history[7].config, + ), + StateSnapshot( + values={"input": 2}, + tasks=(PregelTask(AnyStr(), "one"),), + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": -1, + "writes": {"input": 2}, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] - # forking from any previous checkpoint should re-run nodes - assert [ - c async for c in app.astream(None, history[0].config, stream_mode="updates") - ] == [] - assert [ - c async for c in app.astream(None, history[1].config, stream_mode="updates") - ] == [ - {"two": {"output": 5}}, - ] - assert [ - c async for c in app.astream(None, history[2].config, stream_mode="updates") - ] == [ - {"one": {"inbox": 4}}, - ] + # forking from any previous checkpoint should re-run nodes + assert [ + c async for c in app.astream(None, history[0].config, stream_mode="updates") + ] == [] + assert [ + c async for c in app.astream(None, history[1].config, stream_mode="updates") + ] == [ + {"two": {"output": 5}}, + ] + assert [ + c async for c in app.astream(None, history[2].config, stream_mode="updates") + ] == [ + {"one": {"inbox": 4}}, + ] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_fork_always_re_runs_nodes( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture + checkpointer_name: str, mocker: MockerFixture ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") add_one = mocker.Mock(side_effect=lambda _: 1) builder = StateGraph(Annotated[int, operator.add]) builder.add_node("add_one", add_one) builder.add_edge(START, "add_one") builder.add_conditional_edges("add_one", lambda cnt: "add_one" if cnt < 6 else END) - graph = builder.compile(checkpointer=checkpointer) + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) - thread1 = {"configurable": {"thread_id": "1"}} + thread1 = {"configurable": {"thread_id": "1"}} - # start execution, stop at inbox - assert [ - c async for c in graph.astream(1, thread1, stream_mode=["values", "updates"]) - ] == [ - ("values", 1), - ("updates", {"add_one": 1}), - ("values", 2), - ("updates", {"add_one": 1}), - ("values", 3), - ("updates", {"add_one": 1}), - ("values", 4), - ("updates", {"add_one": 1}), - ("values", 5), - ("updates", {"add_one": 1}), - ("values", 6), - ] + # start execution, stop at inbox + assert [ + c + async for c in graph.astream(1, thread1, stream_mode=["values", "updates"]) + ] == [ + ("values", 1), + ("updates", {"add_one": 1}), + ("values", 2), + ("updates", {"add_one": 1}), + ("values", 3), + ("updates", {"add_one": 1}), + ("values", 4), + ("updates", {"add_one": 1}), + ("values", 5), + ("updates", {"add_one": 1}), + ("values", 6), + ] - # list history - history = [c async for c in graph.aget_state_history(thread1)] - assert history == [ - StateSnapshot( - values=6, - next=(), - tasks=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 5, - "writes": {"add_one": 1}, - }, - created_at=AnyStr(), - parent_config=history[1].config, - ), - StateSnapshot( - values=5, - tasks=(PregelTask(AnyStr(), "add_one"),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": {"add_one": 1}, - }, - created_at=AnyStr(), - parent_config=history[2].config, - ), - StateSnapshot( - values=4, - tasks=(PregelTask(AnyStr(), "add_one"),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"add_one": 1}, - }, - created_at=AnyStr(), - parent_config=history[3].config, - ), - StateSnapshot( - values=3, - tasks=(PregelTask(AnyStr(), "add_one"),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"add_one": 1}, - }, - created_at=AnyStr(), - parent_config=history[4].config, - ), - StateSnapshot( - values=2, - tasks=(PregelTask(AnyStr(), "add_one"),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"add_one": 1}, - }, - created_at=AnyStr(), - parent_config=history[5].config, - ), - StateSnapshot( - values=1, - tasks=(PregelTask(AnyStr(), "add_one"),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, - created_at=AnyStr(), - parent_config=history[6].config, - ), - StateSnapshot( - values=0, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": 1}, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] + # list history + history = [c async for c in graph.aget_state_history(thread1)] + assert history == [ + StateSnapshot( + values=6, + next=(), + tasks=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 5, + "writes": {"add_one": 1}, + }, + created_at=AnyStr(), + parent_config=history[1].config, + ), + StateSnapshot( + values=5, + tasks=(PregelTask(AnyStr(), "add_one"),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": {"add_one": 1}, + }, + created_at=AnyStr(), + parent_config=history[2].config, + ), + StateSnapshot( + values=4, + tasks=(PregelTask(AnyStr(), "add_one"),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"add_one": 1}, + }, + created_at=AnyStr(), + parent_config=history[3].config, + ), + StateSnapshot( + values=3, + tasks=(PregelTask(AnyStr(), "add_one"),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": {"add_one": 1}, + }, + created_at=AnyStr(), + parent_config=history[4].config, + ), + StateSnapshot( + values=2, + tasks=(PregelTask(AnyStr(), "add_one"),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"add_one": 1}, + }, + created_at=AnyStr(), + parent_config=history[5].config, + ), + StateSnapshot( + values=1, + tasks=(PregelTask(AnyStr(), "add_one"),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + created_at=AnyStr(), + parent_config=history[6].config, + ), + StateSnapshot( + values=0, + tasks=(PregelTask(AnyStr(), "__start__"),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": 1}, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] - # forking from any previous checkpoint should re-run nodes - assert [ - c async for c in graph.astream(None, history[0].config, stream_mode="updates") - ] == [] - assert [ - c async for c in graph.astream(None, history[1].config, stream_mode="updates") - ] == [ - {"add_one": 1}, - ] - assert [ - c async for c in graph.astream(None, history[2].config, stream_mode="updates") - ] == [ - {"add_one": 1}, - {"add_one": 1}, - ] + # forking from any previous checkpoint should re-run nodes + assert [ + c + async for c in graph.astream(None, history[0].config, stream_mode="updates") + ] == [] + assert [ + c + async for c in graph.astream(None, history[1].config, stream_mode="updates") + ] == [ + {"add_one": 1}, + ] + assert [ + c + async for c in graph.astream(None, history[2].config, stream_mode="updates") + ] == [ + {"add_one": 1}, + {"add_one": 1}, + ] async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: @@ -1508,7 +1519,8 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) assert await app.ainvoke(2) == [3, 3] -async def test_invoke_checkpoint(mocker: MockerFixture) -> None: +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_invoke_checkpoint(mocker: MockerFixture, checkpointer_name: str) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) errored_once = False @@ -1531,57 +1543,52 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None: | raise_if_above_10 ) - memory = MemorySaverAssertImmutable() + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = Pregel( + nodes={"one": one}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=checkpointer, + retry_policy=RetryPolicy(), + ) - app = Pregel( - nodes={"one": one}, - channels={ - "total": BinaryOperatorAggregate(int, operator.add), - "input": LastValue(int), - "output": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=memory, - retry_policy=RetryPolicy(), - ) - - # total starts out as 0, so output is 0+2=2 - assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2 - checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 2 - # total is now 2, so output is 2+3=5 - assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5 - assert errored_once, "errored and retried" - checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - # total is now 2+5=7, so output would be 7+4=11, but raises ValueError - with pytest.raises(ValueError): - await app.ainvoke(4, {"configurable": {"thread_id": "1"}}) - # checkpoint is not updated - checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - # on a new thread, total starts out as 0, so output is 0+5=5 - assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5 - checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - checkpoint = await memory.aget({"configurable": {"thread_id": "2"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 5 + # total starts out as 0, so output is 0+2=2 + assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2 + checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 2 + # total is now 2, so output is 2+3=5 + assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5 + assert errored_once, "errored and retried" + checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 7 + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + await app.ainvoke(4, {"configurable": {"thread_id": "1"}}) + # checkpoint is not updated + checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 7 + # on a new thread, total starts out as 0, so output is 0+5=5 + assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5 + checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 7 + checkpoint = await checkpointer.aget({"configurable": {"thread_id": "2"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 5 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_pending_writes_resume( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - class State(TypedDict): value: Annotated[int, operator.add] @@ -1609,214 +1616,219 @@ async def test_pending_writes_resume( builder.add_node("two", two) builder.add_edge(START, "one") builder.add_edge(START, "two") - graph = builder.compile(checkpointer=checkpointer) + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) - thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} - with pytest.raises(ValueError, match="I'm not good"): - await graph.ainvoke({"value": 1}, thread1) + thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} + with pytest.raises(ValueError, match="I'm not good"): + await graph.ainvoke({"value": 1}, thread1) - # both nodes should have been called once - assert one.calls == 1 - assert two.calls == 1 + # both nodes should have been called once + assert one.calls == 1 + assert two.calls == 1 - # latest checkpoint should be before nodes "one", "two" - state = await graph.aget_state(thread1) - assert state is not None - assert state.values == {"value": 1} - assert state.next == ("one", "two") - assert state.tasks == ( - PregelTask(AnyStr(), "one"), - PregelTask(AnyStr(), "two", 'ValueError("I\'m not good")'), - ) - assert state.metadata == { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - } - # should contain pending write of "one" - checkpoint = await checkpointer.aget_tuple(thread1) - assert checkpoint is not None - # should contain error from "two" - expected_writes = [ - (AnyStr(), "one", "one"), - (AnyStr(), "value", 2), - (AnyStr(), ERROR, 'ValueError("I\'m not good")'), - ] - assert len(checkpoint.pending_writes) == 3 - assert all(w in expected_writes for w in checkpoint.pending_writes) - # both non-error pending writes come from same task - non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR] - assert non_error_writes[0][0] == non_error_writes[1][0] - # error write is from the other task - error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR) - assert error_write[0] != non_error_writes[0][0] + # latest checkpoint should be before nodes "one", "two" + state = await graph.aget_state(thread1) + assert state is not None + assert state.values == {"value": 1} + assert state.next == ("one", "two") + assert state.tasks == ( + PregelTask(AnyStr(), "one"), + PregelTask(AnyStr(), "two", 'ValueError("I\'m not good")'), + ) + assert state.metadata == { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + } + # should contain pending write of "one" + checkpoint = await checkpointer.aget_tuple(thread1) + assert checkpoint is not None + # should contain error from "two" + expected_writes = [ + (AnyStr(), "one", "one"), + (AnyStr(), "value", 2), + (AnyStr(), ERROR, 'ValueError("I\'m not good")'), + ] + assert len(checkpoint.pending_writes) == 3 + assert all(w in expected_writes for w in checkpoint.pending_writes) + # both non-error pending writes come from same task + non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR] + assert non_error_writes[0][0] == non_error_writes[1][0] + # error write is from the other task + error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR) + assert error_write[0] != non_error_writes[0][0] - # TODO arguably this shouldn't even run the failed task again, - # and should require empty update_state (ie new checkpoint_id) - # in order to try again + # TODO arguably this shouldn't even run the failed task again, + # and should require empty update_state (ie new checkpoint_id) + # in order to try again - # resume execution - with pytest.raises(ValueError, match="I'm not good"): - await graph.ainvoke(None, thread1) + # resume execution + with pytest.raises(ValueError, match="I'm not good"): + await graph.ainvoke(None, thread1) - # node "one" succeeded previously, so shouldn't be called again - assert one.calls == 1 - # node "two" should have been called once again - assert two.calls == 2 + # node "one" succeeded previously, so shouldn't be called again + assert one.calls == 1 + # node "two" should have been called once again + assert two.calls == 2 - # confirm no new checkpoints saved - state_two = await graph.aget_state(thread1) - assert state_two.metadata == state.metadata + # confirm no new checkpoints saved + state_two = await graph.aget_state(thread1) + assert state_two.metadata == state.metadata - # resume execution, without exception - two.rtn = {"value": 3} - # both the pending write and the new write were applied, 1 + 2 + 3 = 6 - assert await graph.ainvoke(None, thread1) == {"value": 6} + # resume execution, without exception + two.rtn = {"value": 3} + # both the pending write and the new write were applied, 1 + 2 + 3 = 6 + assert await graph.ainvoke(None, thread1) == {"value": 6} - # check all final checkpoints - checkpoints = [c async for c in checkpointer.alist(thread1)] - # we should have 3 - assert len(checkpoints) == 3 - # the last one not too interesting for this test - assert checkpoints[0] == CheckpointTuple( - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - checkpoint={ - "v": 1, - "id": AnyStr(), - "ts": AnyStr(), - "pending_sends": [], - "versions_seen": { - "one": { - "start:one": AnyVersion(), + # check all final checkpoints + checkpoints = [c async for c in checkpointer.alist(thread1)] + # we should have 3 + assert len(checkpoints) == 3 + # the last one not too interesting for this test + assert checkpoints[0] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 1, + "id": AnyStr(), + "ts": AnyStr(), + "pending_sends": [], + "versions_seen": { + "one": { + "start:one": AnyVersion(), + }, + "two": { + "start:two": AnyVersion(), + }, + "__input__": {}, + "__start__": { + "__start__": AnyVersion(), + }, + "__interrupt__": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "start:one": AnyVersion(), + "start:two": AnyVersion(), + }, }, - "two": { - "start:two": AnyVersion(), - }, - "__input__": {}, - "__start__": { - "__start__": AnyVersion(), - }, - "__interrupt__": { + "channel_versions": { + "one": AnyVersion(), + "two": AnyVersion(), "value": AnyVersion(), "__start__": AnyVersion(), "start:one": AnyVersion(), "start:two": AnyVersion(), }, + "channel_values": {"one": "one", "two": "two", "value": 6}, }, - "channel_versions": { - "one": AnyVersion(), - "two": AnyVersion(), - "value": AnyVersion(), - "__start__": AnyVersion(), - "start:one": AnyVersion(), - "start:two": AnyVersion(), + metadata={ + "parents": {}, + "step": 1, + "source": "loop", + "writes": {"one": {"value": 2}, "two": {"value": 3}}, }, - "channel_values": {"one": "one", "two": "two", "value": 6}, - }, - metadata={ - "parents": {}, - "step": 1, - "source": "loop", - "writes": {"one": {"value": 2}, "two": {"value": 3}}, - }, - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": checkpoints[1].config["configurable"]["checkpoint_id"], - } - }, - pending_writes=[], - ) - # the previous one we assert that pending writes contains both - # - original error - # - successful writes from resuming after preventing error - assert checkpoints[1] == CheckpointTuple( - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - checkpoint={ - "v": 1, - "id": AnyStr(), - "ts": AnyStr(), - "pending_sends": [], - "versions_seen": { - "__input__": {}, - "__start__": { + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": checkpoints[1].config["configurable"][ + "checkpoint_id" + ], + } + }, + pending_writes=[], + ) + # the previous one we assert that pending writes contains both + # - original error + # - successful writes from resuming after preventing error + assert checkpoints[1] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 1, + "id": AnyStr(), + "ts": AnyStr(), + "pending_sends": [], + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": AnyVersion(), + }, + }, + "channel_versions": { + "value": AnyVersion(), "__start__": AnyVersion(), + "start:one": AnyVersion(), + "start:two": AnyVersion(), + }, + "channel_values": { + "value": 1, + "start:one": "__start__", + "start:two": "__start__", }, }, - "channel_versions": { - "value": AnyVersion(), - "__start__": AnyVersion(), - "start:one": AnyVersion(), - "start:two": AnyVersion(), + metadata={"parents": {}, "step": 0, "source": "loop", "writes": None}, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": checkpoints[2].config["configurable"][ + "checkpoint_id" + ], + } }, - "channel_values": { - "value": 1, - "start:one": "__start__", - "start:two": "__start__", + pending_writes=UnsortedSequence( + (AnyStr(), "one", "one"), + (AnyStr(), "value", 2), + (AnyStr(), "__error__", 'ValueError("I\'m not good")'), + (AnyStr(), "two", "two"), + (AnyStr(), "value", 3), + ), + ) + assert checkpoints[2] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - }, - metadata={"parents": {}, "step": 0, "source": "loop", "writes": None}, - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"], - } - }, - pending_writes=UnsortedSequence( - (AnyStr(), "one", "one"), - (AnyStr(), "value", 2), - (AnyStr(), "__error__", 'ValueError("I\'m not good")'), - (AnyStr(), "two", "two"), - (AnyStr(), "value", 3), - ), - ) - assert checkpoints[2] == CheckpointTuple( - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - checkpoint={ - "v": 1, - "id": AnyStr(), - "ts": AnyStr(), - "pending_sends": [], - "versions_seen": {"__input__": {}}, - "channel_versions": { - "__start__": AnyVersion(), + checkpoint={ + "v": 1, + "id": AnyStr(), + "ts": AnyStr(), + "pending_sends": [], + "versions_seen": {"__input__": {}}, + "channel_versions": { + "__start__": AnyVersion(), + }, + "channel_values": {"__start__": {"value": 1}}, }, - "channel_values": {"__start__": {"value": 1}}, - }, - metadata={ - "parents": {}, - "step": -1, - "source": "input", - "writes": {"__start__": {"value": 1}}, - }, - parent_config=None, - pending_writes=UnsortedSequence( - (AnyStr(), "value", 1), - (AnyStr(), "start:one", "__start__"), - (AnyStr(), "start:two", "__start__"), - ), - ) + metadata={ + "parents": {}, + "step": -1, + "source": "input", + "writes": {"__start__": {"value": 1}}, + }, + parent_config=None, + pending_writes=UnsortedSequence( + (AnyStr(), "value", 1), + (AnyStr(), "start:one", "__start__"), + (AnyStr(), "start:two", "__start__"), + ), + ) async def test_cond_edge_after_send() -> None: @@ -1848,9 +1860,8 @@ async def test_cond_edge_after_send() -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_invoke_checkpoint_three( - mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str + mocker: MockerFixture, checkpointer_name: str ) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -1865,118 +1876,121 @@ async def test_invoke_checkpoint_three( | raise_if_above_10 ) - app = Pregel( - nodes={"one": one}, - channels={ - "total": BinaryOperatorAggregate(int, operator.add), - "input": LastValue(int), - "output": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=checkpointer, - debug=True, - ) - - thread_1 = {"configurable": {"thread_id": "1"}} - # total starts out as 0, so output is 0+2=2 - assert await app.ainvoke(2, thread_1) == 2 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 2 - assert ( - state.config["configurable"]["checkpoint_id"] - == (await checkpointer.aget(thread_1))["id"] - ) - # total is now 2, so output is 2+3=5 - assert await app.ainvoke(3, thread_1) == 5 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 7 - assert ( - state.config["configurable"]["checkpoint_id"] - == (await checkpointer.aget(thread_1))["id"] - ) - # total is now 2+5=7, so output would be 7+4=11, but raises ValueError - with pytest.raises(ValueError): - await app.ainvoke(4, thread_1) - # checkpoint is not updated - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 7 - assert state.next == ("one",) - """we checkpoint inputs and it failed on "one", so the next node is one""" - # we can recover from error by sending new inputs - assert await app.ainvoke(2, thread_1) == 9 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 16, "total is now 7+9=16" - assert state.next == () - - thread_2 = {"configurable": {"thread_id": "2"}} - # on a new thread, total starts out as 0, so output is 0+5=5 - assert await app.ainvoke(5, thread_2) == 5 - state = await app.aget_state({"configurable": {"thread_id": "1"}}) - assert state is not None - assert state.values.get("total") == 16 - assert state.next == () - state = await app.aget_state(thread_2) - assert state is not None - assert state.values.get("total") == 5 - assert state.next == () - - assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1 - # list all checkpoints for thread 1 - thread_1_history = [c async for c in app.aget_state_history(thread_1)] - # there are 7 checkpoints - assert len(thread_1_history) == 7 - assert Counter(c.metadata["source"] for c in thread_1_history) == { - "input": 4, - "loop": 3, - } - # sorted descending - assert ( - thread_1_history[0].config["configurable"]["checkpoint_id"] - > thread_1_history[1].config["configurable"]["checkpoint_id"] - ) - # cursor pagination - cursored = [ - c - async for c in app.aget_state_history( - thread_1, limit=1, before=thread_1_history[0].config + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = Pregel( + nodes={"one": one}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=checkpointer, + debug=True, ) - ] - assert len(cursored) == 1 - assert cursored[0].config == thread_1_history[1].config - # the last checkpoint - assert thread_1_history[0].values["total"] == 16 - # the first "loop" checkpoint - assert thread_1_history[-2].values["total"] == 2 - # can get each checkpoint using aget with config - assert (await checkpointer.aget(thread_1_history[0].config))[ - "id" - ] == thread_1_history[0].config["configurable"]["checkpoint_id"] - assert (await checkpointer.aget(thread_1_history[1].config))[ - "id" - ] == thread_1_history[1].config["configurable"]["checkpoint_id"] - thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10) - # update creates a new checkpoint - assert ( - thread_1_next_config["configurable"]["checkpoint_id"] - > thread_1_history[0].config["configurable"]["checkpoint_id"] - ) - # 1 more checkpoint in history - assert len([c async for c in app.aget_state_history(thread_1)]) == 8 - assert Counter( - [c.metadata["source"] async for c in app.aget_state_history(thread_1)] - ) == { - "update": 1, - "input": 4, - "loop": 3, - } - # the latest checkpoint is the updated one - assert await app.aget_state(thread_1) == await app.aget_state(thread_1_next_config) + thread_1 = {"configurable": {"thread_id": "1"}} + # total starts out as 0, so output is 0+2=2 + assert await app.ainvoke(2, thread_1) == 2 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 2 + assert ( + state.config["configurable"]["checkpoint_id"] + == (await checkpointer.aget(thread_1))["id"] + ) + # total is now 2, so output is 2+3=5 + assert await app.ainvoke(3, thread_1) == 5 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert ( + state.config["configurable"]["checkpoint_id"] + == (await checkpointer.aget(thread_1))["id"] + ) + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + await app.ainvoke(4, thread_1) + # checkpoint is not updated + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert state.next == ("one",) + """we checkpoint inputs and it failed on "one", so the next node is one""" + # we can recover from error by sending new inputs + assert await app.ainvoke(2, thread_1) == 9 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 16, "total is now 7+9=16" + assert state.next == () + + thread_2 = {"configurable": {"thread_id": "2"}} + # on a new thread, total starts out as 0, so output is 0+5=5 + assert await app.ainvoke(5, thread_2) == 5 + state = await app.aget_state({"configurable": {"thread_id": "1"}}) + assert state is not None + assert state.values.get("total") == 16 + assert state.next == () + state = await app.aget_state(thread_2) + assert state is not None + assert state.values.get("total") == 5 + assert state.next == () + + assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1 + # list all checkpoints for thread 1 + thread_1_history = [c async for c in app.aget_state_history(thread_1)] + # there are 7 checkpoints + assert len(thread_1_history) == 7 + assert Counter(c.metadata["source"] for c in thread_1_history) == { + "input": 4, + "loop": 3, + } + # sorted descending + assert ( + thread_1_history[0].config["configurable"]["checkpoint_id"] + > thread_1_history[1].config["configurable"]["checkpoint_id"] + ) + # cursor pagination + cursored = [ + c + async for c in app.aget_state_history( + thread_1, limit=1, before=thread_1_history[0].config + ) + ] + assert len(cursored) == 1 + assert cursored[0].config == thread_1_history[1].config + # the last checkpoint + assert thread_1_history[0].values["total"] == 16 + # the first "loop" checkpoint + assert thread_1_history[-2].values["total"] == 2 + # can get each checkpoint using aget with config + assert (await checkpointer.aget(thread_1_history[0].config))[ + "id" + ] == thread_1_history[0].config["configurable"]["checkpoint_id"] + assert (await checkpointer.aget(thread_1_history[1].config))[ + "id" + ] == thread_1_history[1].config["configurable"]["checkpoint_id"] + + thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10) + # update creates a new checkpoint + assert ( + thread_1_next_config["configurable"]["checkpoint_id"] + > thread_1_history[0].config["configurable"]["checkpoint_id"] + ) + # 1 more checkpoint in history + assert len([c async for c in app.aget_state_history(thread_1)]) == 8 + assert Counter( + [c.metadata["source"] async for c in app.aget_state_history(thread_1)] + ) == { + "update": 1, + "input": 4, + "loop": 3, + } + # the latest checkpoint is the updated one + assert await app.aget_state(thread_1) == await app.aget_state( + thread_1_next_config + ) async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: @@ -2189,9 +2203,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_conditional_graph( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: +async def test_conditional_graph(checkpointer_name: str) -> None: from copy import deepcopy from langchain_core.agents import AgentAction, AgentFinish @@ -2200,8 +2212,6 @@ async def test_conditional_graph( from langchain_core.runnables import RunnablePassthrough from langchain_core.tools import tool - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - # Assemble the tools @tool() def search_api(query: str) -> str: @@ -2458,82 +2468,75 @@ async def test_conditional_graph( }, ] - # test state get/update methods with interrupt_after + async with awith_checkpointer(checkpointer_name) as checkpointer: + # test state get/update methods with interrupt_after - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["agent"], ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - } - ] + config = {"configurable": {"thread_id": "1"}} - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), } + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + }, }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { + "agent": { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { + await app_w_interrupt.aupdate_state( + config, + { "agent_outcome": AgentAction( tool="search_api", tool_input="query", @@ -2541,18 +2544,10 @@ async def test_conditional_graph( ), "input": "what is weather in sf", }, - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ "agent": { "agent_outcome": AgentAction( tool="search_api", @@ -2560,86 +2555,85 @@ async def test_conditional_graph( log="tool:search_api:a different query", ), "input": "what is weather in sf", - } + }, }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] + "input": "what is weather in sf", + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - await app_w_interrupt.aupdate_state( - config, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - ) + "input": "what is weather in sf", + }, + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { + await app_w_interrupt.aupdate_state( + config, + { "input": "what is weather in sf", "intermediate_steps": [ [ @@ -2656,18 +2650,10 @@ async def test_conditional_graph( log="finish:a really nice answer", ), }, - }, - tasks=(), - next=(), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 4, - "writes": { + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ "agent": { "input": "what is weather in sf", "intermediate_steps": [ @@ -2684,91 +2670,112 @@ async def test_conditional_graph( return_values={"answer": "a really nice answer"}, log="finish:a really nice answer", ), - } + }, }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - # test state get/update methods with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "2"}} - llm.i = 0 - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - } - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { + tasks=(), + next=(), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 4, + "writes": { "agent": { "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", ), } - } + }, }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - ) + # test state get/update methods with interrupt_before - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + llm.i = 0 + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + }, + }, + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { + "agent": { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + await app_w_interrupt.aupdate_state( + config, + { "agent_outcome": AgentAction( tool="search_api", tool_input="query", @@ -2776,18 +2783,10 @@ async def test_conditional_graph( ), "input": "what is weather in sf", }, - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ "agent": { "agent_outcome": AgentAction( tool="search_api", @@ -2795,86 +2794,85 @@ async def test_conditional_graph( log="tool:search_api:a different query", ), "input": "what is weather in sf", - } + }, }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] + "input": "what is weather in sf", + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - await app_w_interrupt.aupdate_state( - config, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - ) + "input": "what is weather in sf", + }, + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { + await app_w_interrupt.aupdate_state( + config, + { "input": "what is weather in sf", "intermediate_steps": [ [ @@ -2891,18 +2889,10 @@ async def test_conditional_graph( log="finish:a really nice answer", ), }, - }, - tasks=(), - next=(), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 4, - "writes": { + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ "agent": { "input": "what is weather in sf", "intermediate_steps": [ @@ -2919,196 +2909,234 @@ async def test_conditional_graph( return_values={"answer": "a really nice answer"}, log="finish:a really nice answer", ), - } + }, }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - # test re-invoke to continue with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "3"}} - llm.i = 0 # reset the llm - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - } - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { + tasks=(), + next=(), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 4, + "writes": { "agent": { "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", ), } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + # test re-invoke to continue with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "3"}} + llm.i = 0 # reset the llm + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + }, + }, + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { + "agent": { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + }, + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], } }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] + ] - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ], + [ + AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + "result for another", + ], ], - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ], + [ + AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + "result for another", + ], ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, - ] + "agent_outcome": AgentFinish( + return_values={"answer": "answer"}, log="finish:answer" + ), + } + }, + ] -async def test_conditional_graph_state(mocker: MockerFixture) -> None: +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_conditional_graph_state( + mocker: MockerFixture, checkpointer_name: str +) -> None: from langchain_core.agents import AgentAction, AgentFinish from langchain_core.language_models.fake import FakeStreamingListLLM from langchain_core.prompts import PromptTemplate @@ -3345,15 +3373,201 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: }, ] - # test state get/update methods with interrupt_after + async with awith_checkpointer(checkpointer_name) as checkpointer: + # test state get/update methods with interrupt_after - app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + async with assert_ctx_once(): + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + async with assert_ctx_once(): + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + async with assert_ctx_once(): + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + async with assert_ctx_once(): + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + }, + tasks=(), + next=(), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": { + "agent": { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + # test state get/update methods with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + llm.i = 0 # reset the llm - async with assert_ctx_once(): assert [ c async for c in app_w_interrupt.astream( @@ -3371,41 +3585,38 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: }, ] - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + "intermediate_steps": [], }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - async with assert_ctx_once(): await app_w_interrupt.aupdate_state( config, { @@ -3417,41 +3628,40 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: }, ) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - } + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "intermediate_steps": [], }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - async with assert_ctx_once(): assert [c async for c in app_w_interrupt.astream(None, config)] == [ { "tools": { @@ -3478,7 +3688,6 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: }, ] - async with assert_ctx_once(): await app_w_interrupt.aupdate_state( config, { @@ -3489,151 +3698,12 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: }, ) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - }, - tasks=(), - next=(), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - } - }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - # test state get/update methods with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "2"}} - llm.i = 0 # reset the llm - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", ), - } - }, - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - } - }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": { "intermediate_steps": [ [ AgentAction( @@ -3644,69 +3714,30 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: "result for query", ] ], - } - }, - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] - - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - }, - tasks=(), - next=(), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - } }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + tasks=(), + next=(), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": { + "agent": { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) async def test_conditional_entrypoint_graph() -> None: @@ -3983,9 +4014,7 @@ class ToolInput(BaseModel, arbitrary_types_allowed=True): @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_state_graph_packets( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: +async def test_state_graph_packets(checkpointer_name: str) -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, ) @@ -3997,8 +4026,6 @@ async def test_state_graph_packets( ) from langchain_core.tools import tool - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - class AgentState(TypedDict): messages: Annotated[list[BaseMessage], add_messages] session: Annotated[httpx.AsyncClient, Context(httpx.AsyncClient)] @@ -4220,63 +4247,20 @@ async def test_state_graph_packets( {"agent": {"messages": AIMessage(content="answer", id="ai3")}}, ] - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"messages": HumanMessage(content="what is weather in sf")}, config + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["agent"], ) - ] == [ - { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - } - }, - ] + config = {"configurable": {"thread_id": "1"}} - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - ] - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { + assert [ + c + async for c in app_w_interrupt.astream( + {"messages": HumanMessage(content="what is weather in sf")}, config + ) + ] == [ + { "agent": { "messages": AIMessage( id="ai1", @@ -4291,47 +4275,68 @@ async def test_state_graph_packets( ) } }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + ] - # modify ai message - last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1] - last_message.tool_calls[0]["args"]["query"] = "a different query" - await app_w_interrupt.aupdate_state(config, {"messages": last_message}) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - # message was replaced instead of appended - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - ] - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "messages": AIMessage( + # modify ai message + last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1] + last_message.tool_calls[0]["args"]["query"] = "a different query" + await app_w_interrupt.aupdate_state(config, {"messages": last_message}) + + # message was replaced instead of appended + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( id="ai1", content="", tool_calls=[ @@ -4341,97 +4346,50 @@ async def test_state_graph_packets( "args": {"query": "a different query"}, }, ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ) + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "messages": ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", ) } }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": { - "messages": ToolMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - tool_call_id="tool_call123", - ) - } - }, - { - "agent": { - "messages": AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ) - }, - }, - ] - - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - ToolMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - tool_call_id="tool_call123", - ), - AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ), - ] - }, - tasks=(PregelTask(AnyStr(), "tools"), PregelTask(AnyStr(), "tools")), - next=("tools", "tools"), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { + { "agent": { "messages": AIMessage( id="ai2", @@ -4451,63 +4409,133 @@ async def test_state_graph_packets( ) }, }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + ] - await app_w_interrupt.aupdate_state( - config, - {"messages": AIMessage(content="answer", id="ai2")}, - ) + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools"), PregelTask(AnyStr(), "tools")), + next=("tools", "tools"), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ) + }, + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - # replaces message even if object identity is different, as long as id is the same - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - ToolMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - tool_call_id="tool_call123", - ), - AIMessage(content="answer", id="ai2"), - ] - }, - tasks=(), - next=(), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": {"agent": {"messages": AIMessage(content="answer", id="ai2")}}, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + await app_w_interrupt.aupdate_state( + config, + {"messages": AIMessage(content="answer", id="ai2")}, + ) + + # replaces message even if object identity is different, as long as id is the same + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ] + }, + tasks=(), + next=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": { + "agent": {"messages": AIMessage(content="answer", id="ai2")} + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_message_graph( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: +async def test_message_graph(checkpointer_name: str) -> None: from langchain_core.agents import AgentAction from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, @@ -4515,8 +4543,6 @@ async def test_message_graph( from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage from langchain_core.tools import tool - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - class FakeFuntionChatModel(FakeMessagesListChatModel): def bind_functions(self, functions: list): return self @@ -4679,50 +4705,20 @@ async def test_message_graph( {"agent": AIMessage(content="answer", id="ai3")}, ] - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - HumanMessage(content="what is weather in sf"), config + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["agent"], ) - ] == [ - { - "agent": AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"query"'} - }, - id="ai1", - ) - }, - ] + config = {"configurable": {"thread_id": "1"}} - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"query"'} - }, - id="ai1", - ), - ], - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { + assert [ + c + async for c in app_w_interrupt.astream( + HumanMessage(content="what is weather in sf"), config + ) + ] == [ + { "agent": AIMessage( content="", additional_kwargs={ @@ -4731,43 +4727,59 @@ async def test_message_graph( id="ai1", ) }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + ] - # modify ai message - last_message = (await app_w_interrupt.aget_state(config)).values[-1] - last_message.additional_kwargs["function_call"]["arguments"] = '"a different query"' - await app_w_interrupt.aupdate_state(config, last_message) - - # message was replaced instead of appended - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"a different query"', - } + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"query"', + } + }, + id="ai1", + ) }, - id="ai1", - ), - ], - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": AIMessage( + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + # modify ai message + last_message = (await app_w_interrupt.aget_state(config)).values[-1] + last_message.additional_kwargs["function_call"]["arguments"] = ( + '"a different query"' + ) + await app_w_interrupt.aupdate_state(config, last_message) + + # message was replaced instead of appended + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( content="", additional_kwargs={ "function_call": { @@ -4776,69 +4788,43 @@ async def test_message_graph( } }, id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"a different query"', + } + }, + id="ai1", + ) + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": FunctionMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), ) }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": FunctionMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - ) - }, - { - "agent": AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"another"'} - }, - id="ai2", - ) - }, - ] - - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"a different query"', - } - }, - id="ai1", - ), - FunctionMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - ), - AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"another"'} - }, - id="ai2", - ), - ], - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { + { "agent": AIMessage( content="", additional_kwargs={ @@ -4850,53 +4836,105 @@ async def test_message_graph( id="ai2", ) }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + ] - await app_w_interrupt.aupdate_state( - config, - AIMessage(content="answer", id="ai2"), - ) - - # replaces message even if object identity is different, as long as id is the same - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"a different query"', - } + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"a different query"', + } + }, + id="ai1", + ), + FunctionMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + ), + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"another"', + } + }, + id="ai2", + ), + ], + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"another"', + } + }, + id="ai2", + ) }, - id="ai1", - ), - FunctionMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - ), + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + await app_w_interrupt.aupdate_state( + config, AIMessage(content="answer", id="ai2"), - ], - tasks=(), - next=(), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": {"agent": AIMessage(content="answer", id="ai2")}, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + ) + + # replaces message even if object identity is different, as long as id is the same + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"a different query"', + } + }, + id="ai1", + ), + FunctionMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + ), + AIMessage(content="answer", id="ai2"), + ], + tasks=(), + next=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": {"agent": AIMessage(content="answer", id="ai2")}, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) async def test_in_one_fan_out_out_one_graph_state() -> None: @@ -5126,11 +5164,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_start_branch_then( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - +async def test_start_branch_then(checkpointer_name: str) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -5175,172 +5209,183 @@ async def test_start_branch_then( "market": "US", } - tool_two = tool_two_graph.compile( - store=MemoryStore(), - checkpointer=checkpointer, - interrupt_before=["tool_two_fast", "tool_two_slow"], - ) + async with awith_checkpointer(checkpointer_name) as checkpointer: + tool_two = tool_two_graph.compile( + store=MemoryStore(), + checkpointer=checkpointer, + interrupt_before=["tool_two_fast", "tool_two_slow"], + ) - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value", - "market": "DE", - } - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - }, - ] - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), - next=("tool_two_slow",), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint["ts"], - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, - parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value slow", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value slow", "market": "DE"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config, - ) + thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value", + "market": "DE", + } + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value", "market": "DE"}}, + }, + ] + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value slow", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value slow", "market": "DE"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) - thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, - parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value fast", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value fast", "market": "US"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config, - ) + thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ + "ts" + ], + metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread2, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value fast", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value fast", "market": "US"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread2, limit=2) + ][-1].config, + ) - thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread3) == { - "my_key": "value", - "market": "US", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint["ts"], - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, - parent_config=[c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ - -1 - ].config, - ) - # update state - await tool_two.aupdate_state(thread3, {"my_key": "key"}) # appends to my_key - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "valuekey", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": {START: {"my_key": "key"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ - -1 - ].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread3, debug=1) == { - "my_key": "valuekey fast", - "market": "US", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "valuekey fast", "market": "US"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ - -1 - ].config, - ) + thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread3) == { + "my_key": "value", + "market": "US", + } + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ + "ts" + ], + metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread3, limit=2) + ][-1].config, + ) + # update state + await tool_two.aupdate_state(thread3, {"my_key": "key"}) # appends to my_key + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "valuekey", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": {START: {"my_key": "key"}}, + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread3, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread3, debug=1) == { + "my_key": "valuekey fast", + "market": "US", + } + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "valuekey fast", "market": "US"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread3, limit=2) + ][-1].config, + ) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_branch_then( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - +async def test_branch_then(checkpointer_name: str) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -5368,81 +5413,83 @@ async def test_branch_then( "market": "US", } - # test stream_mode=debug - tool_two = tool_two_graph.compile(checkpointer=checkpointer) - thread10 = {"configurable": {"thread_id": "10"}} - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + async with awith_checkpointer(checkpointer_name) as checkpointer: + # test stream_mode=debug + tool_two = tool_two_graph.compile(checkpointer=checkpointer) + thread10 = {"configurable": {"thread_id": "10"}} + assert [ + c + async for c in tool_two.astream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": -1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, }, - }, - "values": {"my_key": ""}, - "metadata": { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - }, - "next": ["__start__"], - "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + "values": {"my_key": ""}, + "metadata": { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value", "market": "DE"}}, }, + "next": ["__start__"], + "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, - "values": { - "my_key": "value", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - }, - "next": ["prepare"], - "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + }, + "next": ["prepare"], + "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, }, }, { @@ -5456,46 +5503,48 @@ async def test_branch_then( "error": None, "interrupts": [], }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, }, + "values": { + "my_key": "value prepared", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + "next": ["tool_two_slow"], + "tasks": [ + {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()} + ], }, - "values": { - "my_key": "value prepared", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - "next": ["tool_two_slow"], - "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}], }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", - "name": "tool_two_slow", - "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", + "name": "tool_two_slow", + "input": {"my_key": "value prepared", "market": "DE"}, + "triggers": ["branch:prepare:condition:tool_two_slow"], + }, }, }, { @@ -5509,46 +5558,46 @@ async def test_branch_then( "error": None, "interrupts": [], }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, }, + "values": { + "my_key": "value prepared slow", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 2, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, + "next": ["finish"], + "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], }, - "values": { - "my_key": "value prepared slow", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - }, - "next": ["finish"], - "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition::then"], + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition::then"], + }, }, }, { @@ -5562,121 +5611,122 @@ async def test_branch_then( "error": None, "interrupts": [], }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, }, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + "next": [], + "tasks": [], }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - "next": [], - "tasks": [], }, - }, - ] + ] - tool_two = tool_two_graph.compile( - checkpointer=checkpointer, interrupt_before=["tool_two_fast", "tool_two_slow"] - ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "11"}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread1, stream_mode="debug" + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, + interrupt_before=["tool_two_fast", "tool_two_slow"], ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "11"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream( + {"my_key": "value", "market": "DE"}, thread1, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": -1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, }, - }, - "values": {"my_key": ""}, - "metadata": { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - }, - "next": ["__start__"], - "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + "values": {"my_key": ""}, + "metadata": { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value", "market": "DE"}}, }, + "next": ["__start__"], + "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, - "values": { - "my_key": "value", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - }, - "next": ["prepare"], - "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "1a591be4-f85c-558f-8d00-1ccac0d1877f", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + }, + "next": ["prepare"], + "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "1a591be4-f85c-558f-8d00-1ccac0d1877f", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, }, }, { @@ -5690,278 +5740,302 @@ async def test_branch_then( "error": None, "interrupts": [], }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, }, + "values": { + "my_key": "value prepared", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + "next": ["tool_two_slow"], + "tasks": [ + {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()} + ], }, - "values": { - "my_key": "value prepared", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - "next": ["tool_two_slow"], - "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}], }, - }, - ] - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), - next=("tool_two_slow",), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config, - ) + ] + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) - thread2 = {"configurable": {"thread_id": "12"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config, - ) + thread2 = {"configurable": {"thread_id": "12"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread2, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread2, limit=2) + ][-1].config, + ) - tool_two = tool_two_graph.compile( - checkpointer=checkpointer, interrupt_after=["prepare"] - ) + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, interrupt_after=["prepare"] + ) - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - thread1 = {"configurable": {"thread_id": "21"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value prepared", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), - next=("tool_two_slow",), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config, - ) + thread1 = {"configurable": {"thread_id": "21"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value prepared", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) - thread2 = {"configurable": {"thread_id": "22"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config, - ) + thread2 = {"configurable": {"thread_id": "22"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread2, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread2, limit=2) + ][-1].config, + ) - thread3 = {"configurable": {"thread_id": "23"}} - # update an empty thread before first run - uconfig = await tool_two.aupdate_state(thread3, {"my_key": "key", "market": "DE"}) - # check current state - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "prepare"),), - next=("prepare",), - config=uconfig, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 0, - "writes": {START: {"my_key": "key", "market": "DE"}}, - }, - parent_config=None, - ) - # run from this point - assert await tool_two.ainvoke(None, thread3) == { - "my_key": "key prepared", - "market": "DE", - } - # get state after first node - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), - next=("tool_two_slow",), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=uconfig, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread3, debug=1) == { - "my_key": "key prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ - -1 - ].config, - ) + thread3 = {"configurable": {"thread_id": "23"}} + # update an empty thread before first run + uconfig = await tool_two.aupdate_state( + thread3, {"my_key": "key", "market": "DE"} + ) + # check current state + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "key", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "prepare"),), + next=("prepare",), + config=uconfig, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 0, + "writes": {START: {"my_key": "key", "market": "DE"}}, + }, + parent_config=None, + ) + # run from this point + assert await tool_two.ainvoke(None, thread3) == { + "my_key": "key prepared", + "market": "DE", + } + # get state after first node + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "key prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + parent_config=uconfig, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread3, debug=1) == { + "my_key": "key prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "key prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread3, limit=2) + ][-1].config, + ) -async def test_in_one_fan_out_state_graph_waiting_edge() -> None: +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_in_one_fan_out_state_graph_waiting_edge(checkpointer_name: str) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -6023,31 +6097,33 @@ async def test_in_one_fan_out_state_graph_waiting_edge() -> None: {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["retriever_one"], ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - ] + config = {"configurable": {"thread_id": "1"}} - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( - snapshot: SnapshotAssertion, + snapshot: SnapshotAssertion, checkpointer_name: str ) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -6114,31 +6190,33 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["retriever_one"], ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - ] + config = {"configurable": {"thread_id": "1"}} - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( - snapshot: SnapshotAssertion, mocker: MockerFixture + snapshot: SnapshotAssertion, mocker: MockerFixture, checkpointer_name: str ) -> None: from langchain_core.pydantic_v1 import BaseModel, ValidationError @@ -6256,75 +6334,77 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} - async with assert_ctx_once(): - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config - ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - ] + async with assert_ctx_once(): + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] - async with assert_ctx_once(): - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + async with assert_ctx_once(): + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "query": "analyzed: query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + "step": 4, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + + async with assert_ctx_once(): + assert await app_w_interrupt.aupdate_state( + config, {"docs": ["doc5"]}, as_node="rewrite_query" + ) == { + "configurable": { + "thread_id": "1", + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", + } } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - "step": 4, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - - async with assert_ctx_once(): - assert await app_w_interrupt.aupdate_state( - config, {"docs": ["doc5"]}, as_node="rewrite_query" - ) == { - "configurable": { - "thread_id": "1", - "checkpoint_id": AnyStr(), - "checkpoint_ns": "", - } - } +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( - snapshot: SnapshotAssertion, + snapshot: SnapshotAssertion, checkpointer_name: str ) -> None: from pydantic import BaseModel, ValidationError @@ -6419,40 +6499,44 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydant {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf", "inner": {"yo": 1}}, config + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["retriever_one"], ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - ] + config = {"configurable": {"thread_id": "1"}} - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf", "inner": {"yo": 1}}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] - assert await app_w_interrupt.aupdate_state( - config, {"docs": ["doc5"]}, as_node="rewrite_query" - ) == { - "configurable": { - "thread_id": "1", - "checkpoint_id": AnyStr(), - "checkpoint_ns": "", + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + assert await app_w_interrupt.aupdate_state( + config, {"docs": ["doc5"]}, as_node="rewrite_query" + ) == { + "configurable": { + "thread_id": "1", + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", + } } - } -async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( + checkpointer_name: str, +) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -6520,28 +6604,29 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["retriever_one"], ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"qa": {"answer": ""}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - ] + config = {"configurable": {"thread_id": "1"}} - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"qa": {"answer": ""}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: @@ -6839,11 +6924,7 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_nested_graph_interrupts_parallel( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: class InnerState(TypedDict): my_key: Annotated[str, operator.add] my_other_key: str @@ -6887,89 +6968,91 @@ async def test_nested_graph_interrupts_parallel( graph.add_edge(["inner", "outer_1"], "outer_2") graph.set_finish_point("outer_2") - app = graph.compile(checkpointer=checkpointer) + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = graph.compile(checkpointer=checkpointer) - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": ""}, config, debug=True) == { - "my_key": "", - } + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert await app.ainvoke({"my_key": ""}, config, debug=True) == { + "my_key": "", + } - assert await app.ainvoke(None, config, debug=True) == { - "my_key": "got here and there and parallel and back again", - } + assert await app.ainvoke(None, config, debug=True) == { + "my_key": "got here and there and parallel and back again", + } - # below combo of assertions is asserting two things - # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) - # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) - # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} - assert [c async for c in app.astream({"my_key": ""}, config, subgraphs=True)] == [ - # we got to parallel node first - ((), {"outer_1": {"my_key": " and parallel"}}), - ((AnyStr("inner:"),), {"inner_1": {"my_key": "got here", "my_other_key": ""}}), - ] - assert [c async for c in app.astream(None, config)] == [ - {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, - {"inner": {"my_key": "got here and there"}}, - {"outer_2": {"my_key": " and back again"}}, - ] + # below combo of assertions is asserting two things + # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) + # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [ + c async for c in app.astream({"my_key": ""}, config, subgraphs=True) + ] == [ + # we got to parallel node first + ((), {"outer_1": {"my_key": " and parallel"}}), + ( + (AnyStr("inner:"),), + {"inner_1": {"my_key": "got here", "my_other_key": ""}}, + ), + ] + assert [c async for c in app.astream(None, config)] == [ + {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, + {"inner": {"my_key": "got here and there"}}, + {"outer_2": {"my_key": " and back again"}}, + ] - # test stream values w/ nested interrupt - config = {"configurable": {"thread_id": "3"}} - assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") - ] == [ - {"my_key": ""}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": ""}, - {"my_key": "got here and there and parallel"}, - {"my_key": "got here and there and parallel and back again"}, - ] + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + ] == [ + {"my_key": ""}, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] - # # test interrupts BEFORE the parallel node - app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) - config = {"configurable": {"thread_id": "4"}} - assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") - ] == [ - {"my_key": ""}, - ] - # while we're waiting for the node w/ interrupt inside to finish - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": ""}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": ""}, - {"my_key": "got here and there and parallel"}, - {"my_key": "got here and there and parallel and back again"}, - ] + # # test interrupts BEFORE the parallel node + app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) + config = {"configurable": {"thread_id": "4"}} + assert [ + c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + ] == [ + {"my_key": ""}, + ] + # while we're waiting for the node w/ interrupt inside to finish + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + {"my_key": ""}, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] - # test interrupts AFTER the parallel node - app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) - config = {"configurable": {"thread_id": "5"}} - assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") - ] == [ - {"my_key": ""}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": ""}, - {"my_key": "got here and there and parallel"}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": "got here and there and parallel"}, - {"my_key": "got here and there and parallel and back again"}, - ] + # test interrupts AFTER the parallel node + app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) + config = {"configurable": {"thread_id": "5"}} + assert [ + c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + ] == [ + {"my_key": ""}, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_doubly_nested_graph_interrupts( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: class State(TypedDict): my_key: str @@ -7017,50 +7100,49 @@ async def test_doubly_nested_graph_interrupts( graph.add_edge("child", "parent_2") graph.set_finish_point("parent_2") - app = graph.compile(checkpointer=checkpointer) + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = graph.compile(checkpointer=checkpointer) - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { - "my_key": "hi my value", - } + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { + "my_key": "hi my value", + } - assert await app.ainvoke(None, config, debug=True) == { - "my_key": "hi my value here and there and back again", - } + assert await app.ainvoke(None, config, debug=True) == { + "my_key": "hi my value here and there and back again", + } - # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} - assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ - {"parent_1": {"my_key": "hi my value"}}, - ] - assert [c async for c in app.astream(None, config)] == [ - {"child": {"my_key": "hi my value here and there"}}, - {"parent_2": {"my_key": "hi my value here and there and back again"}}, - ] + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ + {"parent_1": {"my_key": "hi my value"}}, + ] + assert [c async for c in app.astream(None, config)] == [ + {"child": {"my_key": "hi my value here and there"}}, + {"parent_2": {"my_key": "hi my value here and there and back again"}}, + ] - # test stream values w/ nested interrupt - config = {"configurable": {"thread_id": "3"}} - assert [ - c - async for c in app.astream({"my_key": "my value"}, config, stream_mode="values") - ] == [ - {"my_key": "my value"}, - {"my_key": "hi my value"}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": "hi my value"}, - {"my_key": "hi my value here and there"}, - {"my_key": "hi my value here and there and back again"}, - ] + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + c + async for c in app.astream( + {"my_key": "my value"}, config, stream_mode="values" + ) + ] == [ + {"my_key": "my value"}, + {"my_key": "hi my value"}, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + {"my_key": "hi my value"}, + {"my_key": "hi my value here and there"}, + {"my_key": "hi my value here and there and back again"}, + ] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_nested_graph_state( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - +async def test_nested_graph_state(checkpointer_name: str) -> None: class InnerState(TypedDict): my_key: str my_other_key: str @@ -7106,388 +7188,14 @@ async def test_nested_graph_state( graph.add_edge("inner", "outer_2") graph.set_finish_point("outer_2") - app = graph.compile(checkpointer=checkpointer) + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = graph.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} - await app.ainvoke({"my_key": "my value"}, config, debug=True) - # test state w/ nested subgraph state (right after interrupt) - # first get_state without subgraph state - assert await app.aget_state(config) == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state={"configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()}}, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # now, get_state with subgraphs state - assert await app.aget_state(config, subgraphs=True) == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state=StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - name="inner_2", - error=None, - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "parents": { - "": AnyStr(), - }, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # get_state_history returns outer graph checkpoints - history = [c async for c in app.aget_state_history(config)] - assert history == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - } - }, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"my_key": "my value"}}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # get_state_history for a subgraph returns its checkpoints - child_history = [c async for c in app.aget_state_history(history[0].tasks[0].state)] - assert child_history == [ - StateSnapshot( - values={"my_key": "hi my value here", "my_other_key": "hi my value"}, - next=("inner_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("inner:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - "parents": {"": AnyStr()}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - } - }, - tasks=(PregelTask(id=AnyStr(), name="inner_2"),), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("inner:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": None, - "step": 0, - "parents": {"": AnyStr()}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - } - }, - tasks=(PregelTask(id=AnyStr(), name="inner_1"),), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("inner:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": {"": AnyStr()}, - }, - created_at=AnyStr(), - parent_config=None, - tasks=(PregelTask(id=AnyStr(), name="__start__"),), - ), - ] - - # resume - await app.ainvoke(None, config, debug=True) - # test state w/ nested subgraph state (after resuming from interrupt) - assert await app.aget_state(config) == StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # test full history at the end - actual_history = [c async for c in app.aget_state_history(config)] - expected_history = [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( + config = {"configurable": {"thread_id": "1"}} + await app.ainvoke({"my_key": "my value"}, config, debug=True) + # test state w/ nested subgraph state (right after interrupt) + # first get_state without subgraph state + assert await app.aget_state(config) == StateSnapshot( values={"my_key": "hi my value"}, tasks=( PregelTask( @@ -7520,11 +7228,62 @@ async def test_nested_graph_state( "checkpoint_id": AnyStr(), } }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), + ) + # now, get_state with subgraphs state + assert await app.aget_state(config, subgraphs=True) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + state=StateSnapshot( + values={ + "my_key": "hi my value here", + "my_other_key": "hi my value", + }, + tasks=( + PregelTask( + AnyStr(), + name="inner_2", + error=None, + ), + ), + next=("inner_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "parents": { + "": AnyStr(), + }, + "source": "loop", + "writes": { + "inner_1": { + "my_key": "hi my value here", + "my_other_key": "hi my value", + } + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + } + }, + ), + ), + ), + next=("inner",), config={ "configurable": { "thread_id": "1", @@ -7535,8 +7294,8 @@ async def test_nested_graph_state( metadata={ "parents": {}, "source": "loop", - "writes": None, - "step": 0, + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, }, created_at=AnyStr(), parent_config={ @@ -7546,11 +7305,194 @@ async def test_nested_graph_state( "checkpoint_id": AnyStr(), } }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), + ) + # get_state_history returns outer graph checkpoints + history = [c async for c in app.aget_state_history(config)] + assert history == [ + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + } + }, + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"my_key": "my value"}}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + # get_state_history for a subgraph returns its checkpoints + child_history = [ + c async for c in app.aget_state_history(history[0].tasks[0].state) + ] + assert child_history == [ + StateSnapshot( + values={"my_key": "hi my value here", "my_other_key": "hi my value"}, + next=("inner_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("inner:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "writes": { + "inner_1": { + "my_key": "hi my value here", + "my_other_key": "hi my value", + } + }, + "step": 1, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + } + }, + tasks=(PregelTask(id=AnyStr(), name="inner_2"),), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("inner:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + } + }, + tasks=(PregelTask(id=AnyStr(), name="inner_1"),), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("inner:"): AnyStr()} + ), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": {"my_key": "hi my value"}}, + "step": -1, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=(PregelTask(id=AnyStr(), name="__start__"),), + ), + ] + + # resume + await app.ainvoke(None, config, debug=True) + # test state w/ nested subgraph state (after resuming from interrupt) + assert await app.aget_state(config) == StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), config={ "configurable": { "thread_id": "1", @@ -7560,26 +7502,172 @@ async def test_nested_graph_state( }, metadata={ "parents": {}, - "source": "input", - "writes": {"__start__": {"my_key": "my value"}}, - "step": -1, + "source": "loop", + "writes": { + "outer_2": {"my_key": "hi my value here and there and back again"} + }, + "step": 3, }, created_at=AnyStr(), - parent_config=None, - ), - ] - assert actual_history == expected_history - # test looking up parent state by checkpoint ID - for actual_snapshot, expected_snapshot in zip(actual_history, expected_history): - assert await app.aget_state(actual_snapshot.config) == expected_snapshot + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + # test full history at the end + actual_history = [c async for c in app.aget_state_history(config)] + expected_history = [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": { + "outer_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + tasks=(PregelTask(AnyStr(), "outer_2"),), + next=("outer_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + } + }, + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"my_key": "my value"}}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert actual_history == expected_history + # test looking up parent state by checkpoint ID + for actual_snapshot, expected_snapshot in zip(actual_history, expected_history): + assert await app.aget_state(actual_snapshot.config) == expected_snapshot @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_doubly_nested_graph_state( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - +async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: class State(TypedDict): my_key: str @@ -7627,277 +7715,37 @@ async def test_doubly_nested_graph_state( graph.add_edge("child", "parent_2") graph.set_finish_point("parent_2") - app = graph.compile(checkpointer=checkpointer) + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = graph.compile(checkpointer=checkpointer) - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert [ - c async for c in app.astream({"my_key": "my value"}, config, subgraphs=True) - ] == [ - ((), {"parent_1": {"my_key": "hi my value"}}), - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_1": {"my_key": "hi my value here"}}, - ), - ] - # get state without subgraphs - outer_state = await app.aget_state(config) - assert outer_state == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child"), - } - }, + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert [ + c async for c in app.astream({"my_key": "my value"}, config, subgraphs=True) + ] == [ + ((), {"parent_1": {"my_key": "hi my value"}}), + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_1": {"my_key": "hi my value here"}}, ), - ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - child_state = await app.aget_state(outer_state.tasks[0].state) - assert ( - child_state.tasks[0] - == StateSnapshot( + ] + # get state without subgraphs + outer_state = await app.aget_state(config) + assert outer_state == StateSnapshot( values={"my_key": "hi my value"}, tasks=( PregelTask( AnyStr(), - "child_1", + "child", state={ "configurable": { "thread_id": "1", - "checkpoint_ns": AnyStr(), + "checkpoint_ns": AnyStr("child"), } }, ), ), - next=("child_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "parents": {"": AnyStr()}, - "source": "loop", - "writes": None, - "step": 0, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - } - }, - ).tasks[0] - ) - grandchild_state = await app.aget_state(child_state.tasks[0].state) - assert grandchild_state == StateSnapshot( - values={"my_key": "hi my value here"}, - tasks=( - PregelTask( - AnyStr(), - "grandchild_2", - ), - ), - next=("grandchild_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "source": "loop", - "writes": {"grandchild_1": {"my_key": "hi my value here"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ) - # get state with subgraphs - assert await app.aget_state(config, subgraphs=True) == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child", - state=StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child_1", - state=StateSnapshot( - values={"my_key": "hi my value here"}, - tasks=( - PregelTask( - AnyStr(), - "grandchild_2", - ), - ), - next=("grandchild_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr( - re.compile(r"child:.+|child1:") - ): AnyStr(), - } - ), - } - }, - metadata={ - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "source": "loop", - "writes": { - "grandchild_1": {"my_key": "hi my value here"} - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("child_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "parents": {"": AnyStr()}, - "source": "loop", - "writes": None, - "step": 0, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # resume - assert [c async for c in app.astream(None, config, subgraphs=True)] == [ - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_2": {"my_key": "hi my value here and there"}}, - ), - ((AnyStr("child:"),), {"child_1": {"my_key": "hi my value here and there"}}), - ((), {"child": {"my_key": "hi my value here and there"}}), - ((), {"parent_2": {"my_key": "hi my value here and there and back again"}}), - ] - # get state with and without subgraphs - assert ( - await app.aget_state(config) - == await app.aget_state(config, subgraphs=True) - == StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), + next=("child",), config={ "configurable": { "thread_id": "1", @@ -7908,10 +7756,8 @@ async def test_doubly_nested_graph_state( metadata={ "parents": {}, "source": "loop", - "writes": { - "parent_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, }, created_at=AnyStr(), parent_config={ @@ -7922,13 +7768,227 @@ async def test_doubly_nested_graph_state( } }, ) - ) - # get outer graph history - outer_history = [c async for c in app.aget_state_history(config)] - assert ( - outer_history[0] - == [ - StateSnapshot( + child_state = await app.aget_state(outer_state.tasks[0].state) + assert ( + child_state.tasks[0] + == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child_1", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + } + }, + ), + ), + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "parents": {"": AnyStr()}, + "source": "loop", + "writes": None, + "step": 0, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + } + }, + ).tasks[0] + ) + grandchild_state = await app.aget_state(child_state.tasks[0].state) + assert grandchild_state == StateSnapshot( + values={"my_key": "hi my value here"}, + tasks=( + PregelTask( + AnyStr(), + "grandchild_2", + ), + ), + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "source": "loop", + "writes": {"grandchild_1": {"my_key": "hi my value here"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + ) + # get state with subgraphs + assert await app.aget_state(config, subgraphs=True) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + state=StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child_1", + state=StateSnapshot( + values={"my_key": "hi my value here"}, + tasks=( + PregelTask( + AnyStr(), + "grandchild_2", + ), + ), + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr( + re.compile(r"child:.+|child1:") + ): AnyStr(), + } + ), + } + }, + metadata={ + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "source": "loop", + "writes": { + "grandchild_1": { + "my_key": "hi my value here" + } + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + ), + ), + ), + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "parents": {"": AnyStr()}, + "source": "loop", + "writes": None, + "step": 0, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + } + }, + ), + ), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + # resume + assert [c async for c in app.astream(None, config, subgraphs=True)] == [ + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_2": {"my_key": "hi my value here and there"}}, + ), + ( + (AnyStr("child:"),), + {"child_1": {"my_key": "hi my value here and there"}}, + ), + ((), {"child": {"my_key": "hi my value here and there"}}), + ((), {"parent_2": {"my_key": "hi my value here and there and back again"}}), + ] + # get state with and without subgraphs + assert ( + await app.aget_state(config) + == await app.aget_state(config, subgraphs=True) + == StateSnapshot( values={"my_key": "hi my value here and there and back again"}, tasks=(), next=(), @@ -7957,95 +8017,227 @@ async def test_doubly_nested_graph_state( "checkpoint_id": AnyStr(), } }, - ), + ) + ) + # get outer graph history + outer_history = [c async for c in app.aget_state_history(config)] + assert ( + outer_history[0] + == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": { + "parent_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("parent_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"child": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(PregelTask(id=AnyStr(), name="parent_2"),), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child"), + } + }, + ), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("parent_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(PregelTask(id=AnyStr(), name="parent_1"),), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + tasks=(PregelTask(id=AnyStr(), name="__start__"),), + ), + ][0] + ) + # get child graph history + child_history = [ + c async for c in app.aget_state_history(outer_history[2].tasks[0].state) + ] + assert child_history == [ StateSnapshot( values={"my_key": "hi my value here and there"}, - next=("parent_2",), + next=(), config={ "configurable": { "thread_id": "1", - "checkpoint_ns": "", + "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ - "parents": {}, "source": "loop", - "writes": {"child": {"my_key": "hi my value here and there"}}, - "step": 2, + "writes": {"child_1": {"my_key": "hi my value here and there"}}, + "step": 1, + "parents": {"": AnyStr()}, }, created_at=AnyStr(), parent_config={ "configurable": { "thread_id": "1", - "checkpoint_ns": "", + "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), } }, - tasks=(PregelTask(id=AnyStr(), name="parent_2"),), + tasks=(), ), StateSnapshot( values={"my_key": "hi my value"}, + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + } + }, tasks=( PregelTask( - AnyStr(), - "child", + id=AnyStr(), + name="child_1", state={ "configurable": { "thread_id": "1", - "checkpoint_ns": AnyStr("child"), + "checkpoint_ns": AnyStr("child:"), } }, ), ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("parent_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=(PregelTask(id=AnyStr(), name="parent_1"),), ), StateSnapshot( values={}, @@ -8053,290 +8245,195 @@ async def test_doubly_nested_graph_state( config={ "configurable": { "thread_id": "1", - "checkpoint_ns": "", + "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ - "parents": {}, "source": "input", - "writes": {"my_key": "my value"}, + "writes": {"__start__": {"my_key": "hi my value"}}, "step": -1, + "parents": {"": AnyStr()}, }, created_at=AnyStr(), parent_config=None, tasks=(PregelTask(id=AnyStr(), name="__start__"),), ), - ][0] - ) - # get child graph history - child_history = [ - c async for c in app.aget_state_history(outer_history[2].tasks[0].state) - ] - assert child_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": {"child_1": {"my_key": "hi my value here and there"}}, - "step": 1, - "parents": {"": AnyStr()}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - } - }, - tasks=(), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("child_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": None, - "step": 0, - "parents": {"": AnyStr()}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="child_1", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - } + ] + # get grandchild graph history + grandchild_history = [ + c async for c in app.aget_state_history(child_history[1].tasks[0].state) + ] + assert grandchild_history == [ + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "writes": { + "grandchild_2": {"my_key": "hi my value here and there"} }, - ), + "step": 2, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + tasks=(), ), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": {"": AnyStr()}, - }, - created_at=AnyStr(), - parent_config=None, - tasks=(PregelTask(id=AnyStr(), name="__start__"),), - ), - ] - # get grandchild graph history - grandchild_history = [ - c async for c in app.aget_state_history(child_history[1].tasks[0].state) - ] - assert grandchild_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( + StateSnapshot( + values={"my_key": "hi my value here"}, + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "writes": {"grandchild_1": {"my_key": "hi my value here"}}, + "step": 1, + "parents": AnyDict( { "": AnyStr(), AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), } ), - } - }, - metadata={ - "source": "loop", - "writes": {"grandchild_2": {"my_key": "hi my value here and there"}}, - "step": 2, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), } - ), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - tasks=(), - ), - StateSnapshot( - values={"my_key": "hi my value here"}, - next=("grandchild_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( + }, + tasks=(PregelTask(id=AnyStr(), name="grandchild_2"),), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("grandchild_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": AnyDict( { "": AnyStr(), AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), } ), - } - }, - metadata={ - "source": "loop", - "writes": {"grandchild_1": {"my_key": "hi my value here"}}, - "step": 1, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), } - ), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - tasks=(PregelTask(id=AnyStr(), name="grandchild_2"),), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("grandchild_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( + }, + tasks=(PregelTask(id=AnyStr(), name="grandchild_1"),), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": {"my_key": "hi my value"}}, + "step": -1, + "parents": AnyDict( { "": AnyStr(), AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), } ), - } - }, - metadata={ - "source": "loop", - "writes": None, - "step": 0, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - tasks=(PregelTask(id=AnyStr(), name="grandchild_1"),), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - }, - created_at=AnyStr(), - parent_config=None, - tasks=(PregelTask(id=AnyStr(), name="__start__"),), - ), - ] + }, + created_at=AnyStr(), + parent_config=None, + tasks=(PregelTask(id=AnyStr(), name="__start__"),), + ), + ] - # replay grandchild checkpoint - assert [ - c async for c in app.astream(None, grandchild_history[2].config, subgraphs=True) - ] == [ - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_1": {"my_key": "hi my value here"}}, - ) - ] + # replay grandchild checkpoint + assert [ + c + async for c in app.astream( + None, grandchild_history[2].config, subgraphs=True + ) + ] == [ + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_1": {"my_key": "hi my value here"}}, + ) + ] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_send_to_nested_graphs( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - +async def test_send_to_nested_graphs(checkpointer_name: str) -> None: class OverallState(TypedDict): subjects: list[str] jokes: Annotated[list[str], operator.add] @@ -8370,112 +8467,74 @@ async def test_send_to_nested_graphs( builder.add_conditional_edges(START, continue_to_jokes) builder.add_edge("generate_joke", END) - graph = builder.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + config = {"configurable": {"thread_id": "1"}} - # invoke and pause at nested interrupt - assert await graph.ainvoke({"subjects": ["cats", "dogs"]}, config=config) == { - "subjects": ["cats", "dogs"], - "jokes": [], - } - # check state - outer_state = await graph.aget_state(config) + # invoke and pause at nested interrupt + assert await graph.ainvoke({"subjects": ["cats", "dogs"]}, config=config) == { + "subjects": ["cats", "dogs"], + "jokes": [], + } + # check state + outer_state = await graph.aget_state(config) - assert outer_state == StateSnapshot( - values={"subjects": ["cats", "dogs"], "jokes": []}, - tasks=( - PregelTask( - AnyStr(), - "generate_joke", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, + assert outer_state == StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + tasks=( + PregelTask( + AnyStr(), + "generate_joke", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + ), + PregelTask( + AnyStr(), + "generate_joke", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + ), ), - PregelTask( - AnyStr(), - "generate_joke", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - ), - ), - next=("generate_joke", "generate_joke"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) + next=("generate_joke", "generate_joke"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) - # update state of dogs joke graph - await graph.aupdate_state( - outer_state.tasks[1].state, {"subject": "turtles - hohoho"} - ) + # update state of dogs joke graph + await graph.aupdate_state( + outer_state.tasks[1].state, {"subject": "turtles - hohoho"} + ) - # continue past interrupt - assert await graph.ainvoke(None, config=config) == { - "subjects": ["cats", "dogs"], - "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - } - - actual_snapshot = await graph.aget_state(config) - expected_snapshot = StateSnapshot( - values={ + # continue past interrupt + assert await graph.ainvoke(None, config=config) == { "subjects": ["cats", "dogs"], "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "generate_joke": [ - {"jokes": ["Joke about cats - hohoho"]}, - {"jokes": ["Joke about turtles - hohoho"]}, - ] - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - assert actual_snapshot == expected_snapshot + } - # test full history - actual_history = [c async for c in graph.aget_state_history(config)] - expected_history = [ - StateSnapshot( + actual_snapshot = await graph.aget_state(config) + expected_snapshot = StateSnapshot( values={ "subjects": ["cats", "dogs"], "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], @@ -8508,53 +8567,256 @@ async def test_send_to_nested_graphs( "checkpoint_id": AnyStr(), } }, - ), - StateSnapshot( - values={"subjects": ["cats", "dogs"], "jokes": []}, - next=("generate_joke", "generate_joke"), - tasks=( - PregelTask( - AnyStr(), - "generate_joke", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } + ) + assert actual_snapshot == expected_snapshot + + # test full history + actual_history = [c async for c in graph.aget_state_history(config)] + expected_history = [ + StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": [ + "Joke about cats - hohoho", + "Joke about turtles - hohoho", + ], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": { + "generate_joke": [ + {"jokes": ["Joke about cats - hohoho"]}, + {"jokes": ["Joke about turtles - hohoho"]}, + ] }, - ), - PregelTask( - AnyStr(), - "generate_joke", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - ), + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, ), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } + StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + next=("generate_joke", "generate_joke"), + tasks=( + PregelTask( + AnyStr(), + "generate_joke", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + ), + PregelTask( + AnyStr(), + "generate_joke", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + ), + ), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"jokes": []}, + tasks=(PregelTask(AnyStr(), "__start__"),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert actual_history == expected_history + + +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_weather_subgraph( + checkpointer_name: str, snapshot: SnapshotAssertion +) -> None: + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import AIMessage, HumanMessage, ToolCall + from langchain_core.tools import tool + + from langgraph.graph import MessagesState + + # setup subgraph + + @tool + def get_weather(city: str): + """Get the weather for a specific city""" + return f"I'ts sunny in {city}!" + + weather_model = FakeMessagesListChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="get_weather", + args={"city": "San Francisco"}, + ) + ], + ) + ] + ) + + class SubGraphState(MessagesState): + city: str + + def model_node(state: SubGraphState): + result = weather_model.invoke(state["messages"]) + return {"city": cast(AIMessage, result).tool_calls[0]["args"]["city"]} + + def weather_node(state: SubGraphState): + result = get_weather.invoke({"city": state["city"]}) + return {"messages": [{"role": "assistant", "content": result}]} + + subgraph = StateGraph(SubGraphState) + subgraph.add_node(model_node) + subgraph.add_node(weather_node) + subgraph.add_edge(START, "model_node") + subgraph.add_edge("model_node", "weather_node") + subgraph.add_edge("weather_node", END) + subgraph = subgraph.compile(interrupt_before=["weather_node"]) + + # setup main graph + + class RouterState(MessagesState): + route: Literal["weather", "other"] + + class Router(TypedDict): + route: Literal["weather", "other"] + + router_model = FakeMessagesListChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="router", + args={"dest": "weather"}, + ) + ], + ) + ] + ) + + def router_node(state: RouterState): + system_message = "Classify the incoming query as either about weather or not." + messages = [{"role": "system", "content": system_message}] + state["messages"] + route = router_model.invoke(messages) + return {"route": cast(AIMessage, route).tool_calls[0]["args"]["dest"]} + + def normal_llm_node(state: RouterState): + return {"messages": [AIMessage("Hello!")]} + + def route_after_prediction(state: RouterState): + if state["route"] == "weather": + return "weather_graph" + else: + return "normal_llm_node" + + def weather_graph(state: RouterState): + # this tests that all async checkpointers tested also implement sync methods + # as the subgraph called with sync invoke will use sync checkpointer methods + return subgraph.invoke(state) + + graph = StateGraph(RouterState) + graph.add_node(router_node) + graph.add_node(normal_llm_node) + graph.add_node("weather_graph", weather_graph) + graph.add_edge(START, "router_node") + graph.add_conditional_edges("router_node", route_after_prediction) + graph.add_edge("normal_llm_node", END) + graph.add_edge("weather_graph", END) + + def get_first_in_list(): + return [*graph.get_state_history(config, limit=1)][0] + + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = graph.compile(checkpointer=checkpointer) + + assert graph.get_graph(xray=1).draw_mermaid() == snapshot + + config = {"configurable": {"thread_id": "1"}} + inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} + + # run until interrupt + assert [ + c + async for c in graph.astream( + inputs, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ((), {"router_node": {"route": "weather"}}), + ((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}), + ] + + # check current state + state = await graph.aget_state(config) + assert state == StateSnapshot( + values={ + "messages": [ + HumanMessage(content="what's the weather in sf", id=AnyStr()) + ], + "route": "weather", }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"jokes": []}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), + next=("weather_graph",), config={ "configurable": { "thread_id": "1", @@ -8563,16 +8825,266 @@ async def test_send_to_nested_graphs( } }, metadata={ + "source": "loop", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, "parents": {}, - "source": "input", - "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, - "step": -1, }, created_at=AnyStr(), - parent_config=None, - ), - ] - assert actual_history == expected_history + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("weather_graph:"), + } + }, + ), + ), + ) + # confirm that list() delegates to alist() correctly + assert await asyncio.to_thread(get_first_in_list) == state + + # update + await graph.aupdate_state(state.tasks[0].state, {"city": "la"}) + + # run after update + assert [ + c + async for c in graph.astream( + None, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ( + (AnyStr("weather_graph:"),), + { + "weather_node": { + "messages": [ + {"role": "assistant", "content": "I'ts sunny in la!"} + ] + } + }, + ), + ( + (), + { + "weather_graph": { + "messages": [ + HumanMessage( + content="what's the weather in sf", id=AnyStr() + ), + AIMessage(content="I'ts sunny in la!", id=AnyStr()), + ] + } + }, + ), + ] + + # try updating acting as weather node + config = {"configurable": {"thread_id": "14"}} + inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} + assert [ + c + async for c in graph.astream( + inputs, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ((), {"router_node": {"route": "weather"}}), + ((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}), + ] + state = await graph.aget_state(config, subgraphs=True) + assert state == StateSnapshot( + values={ + "messages": [ + HumanMessage(content="what's the weather in sf", id=AnyStr()) + ], + "route": "weather", + }, + next=("weather_graph",), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + state=StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what's the weather in sf", id=AnyStr() + ) + ], + "city": "San Francisco", + }, + next=("weather_node",), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "writes": {"model_node": {"city": "San Francisco"}}, + "step": 1, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + } + }, + tasks=(PregelTask(id=AnyStr(), name="weather_node"),), + ), + ), + ), + ) + await graph.aupdate_state( + state.tasks[0].state.config, + {"messages": [{"role": "assistant", "content": "rainy"}]}, + as_node="weather_node", + ) + state = await graph.aget_state(config, subgraphs=True) + assert state == StateSnapshot( + values={ + "messages": [ + HumanMessage(content="what's the weather in sf", id=AnyStr()) + ], + "route": "weather", + }, + next=("weather_graph",), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + state=StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what's the weather in sf", id=AnyStr() + ), + AIMessage(content="rainy", id=AnyStr()), + ], + "city": "San Francisco", + }, + next=(), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), + } + }, + metadata={ + "source": "update", + "step": 2, + "writes": { + "weather_node": { + "messages": [ + {"role": "assistant", "content": "rainy"} + ] + } + }, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + } + }, + tasks=(), + ), + ), + ), + ) + assert [ + c + async for c in graph.astream( + None, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ( + (), + { + "weather_graph": { + "messages": [ + HumanMessage( + content="what's the weather in sf", id=AnyStr() + ), + AIMessage(content="rainy", id=AnyStr()), + ] + } + }, + ), + ] async def test_checkpoint_metadata() -> None: @@ -8731,3 +9243,26 @@ async def test_checkpoint_metadata() -> None: assert chkpnt_tuple.metadata["thread_id"] == "2" assert chkpnt_tuple.metadata["test_config_3"] == "foo" assert chkpnt_tuple.metadata["test_config_4"] == "bar" + + +async def test_checkpointer_null_pending_writes() -> None: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + + def __call__(self, state): + return [self.name] + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_edge(START, "1") + graph = builder.compile(checkpointer=MemorySaverNoPending()) + assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"] + assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"] * 2 + assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [ + "1" + ] * 3 + assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [ + "1" + ] * 4 diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py index 4088e38c9..12cd7f653 100644 --- a/libs/langgraph/tests/test_state.py +++ b/libs/langgraph/tests/test_state.py @@ -1,10 +1,11 @@ +import warnings from typing import Annotated as Annotated2 -from typing import Any +from typing import Any, Optional import pytest from langchain_core.runnables import RunnableConfig from pydantic.v1 import BaseModel -from typing_extensions import Annotated, TypedDict +from typing_extensions import Annotated, NotRequired, Required, TypedDict from langgraph.graph.state import StateGraph, _warn_invalid_state_schema @@ -45,7 +46,8 @@ def test_warns_invalid_schema(schema: Any): ) def test_doesnt_warn_valid_schema(schema: Any): # Assert the function does not raise a warning - with pytest.warns(None): + with warnings.catch_warnings(): + warnings.simplefilter("error") _warn_invalid_state_schema(schema) @@ -86,3 +88,45 @@ def test_state_schema_with_type_hint(): for i, c in enumerate(graph.stream(input_state, stream_mode="updates")): node_name = actions[i].__name__ assert c[node_name] == output_state + + +@pytest.mark.parametrize("total_", [True, False]) +def test_state_schema_optional_values(total_: bool): + class SomeParentState(TypedDict): + val0a: str + val0b: Optional[str] + + class InputState(SomeParentState, total=total_): # type: ignore + val1: str + val2: Optional[str] + val3: Required[str] + val4: NotRequired[dict] + val5: Annotated[Required[str], "foo"] + val6: Annotated[NotRequired[str], "bar"] + + class State(InputState): # this would be ignored + val4: dict + + builder = StateGraph(State, input=InputState) + builder.add_node("n", lambda x: x) + builder.add_edge("__start__", "n") + graph = builder.compile() + model = graph.input_schema + json_schema = model.schema() + + if total_ is False: + expected_required = set() + expected_optional = {"val2", "val1"} + else: + expected_required = {"val1"} + + expected_optional = {"val2"} + + # The others should always have precedence based on the required annotation + expected_required |= {"val0a", "val3", "val5"} + expected_optional |= {"val0b", "val4", "val6"} + + assert set(json_schema.get("required", set())) == expected_required + assert ( + set(json_schema["properties"].keys()) == expected_required | expected_optional + ) diff --git a/libs/langgraph/tests/test_store.py b/libs/langgraph/tests/test_store.py index cd53ee407..71494adaf 100644 --- a/libs/langgraph/tests/test_store.py +++ b/libs/langgraph/tests/test_store.py @@ -1,11 +1,14 @@ import asyncio from typing import Any, Optional +import pytest from pytest_mock import MockerFixture from langgraph.store.base import BaseStore from langgraph.store.batch import AsyncBatchedStore +pytestmark = pytest.mark.anyio + async def test_async_batch_store(mocker: MockerFixture) -> None: aget = mocker.stub() diff --git a/libs/langgraph/tests/test_utils.py b/libs/langgraph/tests/test_utils.py index 1e3bf8ea1..bee8fcd9c 100644 --- a/libs/langgraph/tests/test_utils.py +++ b/libs/langgraph/tests/test_utils.py @@ -1,15 +1,34 @@ import functools import sys import uuid -from typing import TypedDict +from typing import ( + Any, + Callable, + Dict, + ForwardRef, + List, + Literal, + Optional, + TypedDict, + TypeVar, + Union, +) from unittest.mock import patch import langsmith import pytest +from typing_extensions import Annotated, NotRequired, Required from langgraph.graph import END, StateGraph from langgraph.graph.graph import CompiledGraph -from langgraph.utils import is_async_callable, is_async_generator +from langgraph.utils import ( + _is_optional_type, + get_field_default, + is_async_callable, + is_async_generator, +) + +pytestmark = pytest.mark.anyio def test_is_async() -> None: @@ -119,3 +138,96 @@ async def test_runnable_callable_tracing_nested_async(rt_graph: CompiledGraph) - with langsmith.tracing_context(enabled=True): res = await rt_graph.ainvoke({"foo": 1}) assert isinstance(res["node_run_id"], uuid.UUID) + + +def test_is_optional_type(): + assert _is_optional_type(None) + assert not _is_optional_type(type(None)) + assert _is_optional_type(Optional[list]) + assert not _is_optional_type(int) + assert _is_optional_type(Optional[Literal[1, 2, 3]]) + assert not _is_optional_type(Literal[1, 2, 3]) + assert _is_optional_type(Optional[List[int]]) + assert _is_optional_type(Optional[Dict[str, int]]) + assert not _is_optional_type(List[Optional[int]]) + assert _is_optional_type(Union[Optional[str], Optional[int]]) + assert _is_optional_type( + Union[ + Union[Optional[str], Optional[int]], Union[Optional[float], Optional[dict]] + ] + ) + assert not _is_optional_type(Union[Union[str, int], Union[float, dict]]) + + assert _is_optional_type(Union[int, None]) + assert _is_optional_type(Union[str, None, int]) + assert _is_optional_type(Union[None, str, int]) + assert not _is_optional_type(Union[int, str]) + + assert not _is_optional_type(Any) # Do we actually want this? + assert _is_optional_type(Optional[Any]) + + class MyClass: + pass + + assert _is_optional_type(Optional[MyClass]) + assert not _is_optional_type(MyClass) + assert _is_optional_type(Optional[ForwardRef("MyClass")]) + assert not _is_optional_type(ForwardRef("MyClass")) + + assert _is_optional_type(Optional[Union[List[int], Dict[str, Optional[int]]]]) + assert not _is_optional_type(Union[List[int], Dict[str, Optional[int]]]) + + assert _is_optional_type(Optional[Callable[[int], str]]) + assert not _is_optional_type(Callable[[int], Optional[str]]) + + T = TypeVar("T") + assert _is_optional_type(Optional[T]) + assert not _is_optional_type(T) + + U = TypeVar("U", bound=Optional[T]) # type: ignore + assert _is_optional_type(U) + + +def test_is_required(): + class MyBaseTypedDict(TypedDict): + val_1: Required[Optional[str]] + val_2: Required[str] + val_3: NotRequired[str] + val_4: NotRequired[Optional[str]] + val_5: Annotated[NotRequired[int], "foo"] + val_6: NotRequired[Annotated[int, "foo"]] + val_7: Annotated[Required[int], "foo"] + val_8: Required[Annotated[int, "foo"]] + val_9: Optional[str] + val_10: str + + annos = MyBaseTypedDict.__annotations__ + assert get_field_default("val_1", annos["val_1"], MyBaseTypedDict) == ... + assert get_field_default("val_2", annos["val_2"], MyBaseTypedDict) == ... + assert get_field_default("val_3", annos["val_3"], MyBaseTypedDict) is None + assert get_field_default("val_4", annos["val_4"], MyBaseTypedDict) is None + # See https://peps.python.org/pep-0655/#interaction-with-annotated + assert get_field_default("val_5", annos["val_5"], MyBaseTypedDict) is None + assert get_field_default("val_6", annos["val_6"], MyBaseTypedDict) is None + assert get_field_default("val_7", annos["val_7"], MyBaseTypedDict) == ... + assert get_field_default("val_8", annos["val_8"], MyBaseTypedDict) == ... + assert get_field_default("val_9", annos["val_9"], MyBaseTypedDict) is None + assert get_field_default("val_10", annos["val_10"], MyBaseTypedDict) == ... + + class MyChildDict(MyBaseTypedDict): + val_11: int + val_11b: Optional[int] + val_11c: Union[int, None, str] + + class MyGrandChildDict(MyChildDict, total=False): + val_12: int + val_13: Required[str] + + cannos = MyChildDict.__annotations__ + gcannos = MyGrandChildDict.__annotations__ + assert get_field_default("val_11", cannos["val_11"], MyChildDict) == ... + assert get_field_default("val_11b", cannos["val_11b"], MyChildDict) is None + assert get_field_default("val_11c", cannos["val_11c"], MyChildDict) is None + assert get_field_default("val_12", gcannos["val_12"], MyGrandChildDict) is None + assert get_field_default("val_9", gcannos["val_9"], MyGrandChildDict) is None + assert get_field_default("val_13", gcannos["val_13"], MyGrandChildDict) == ... diff --git a/libs/sdk-js/src/client.mts b/libs/sdk-js/src/client.mts index 03024299d..459e9e6b1 100644 --- a/libs/sdk-js/src/client.mts +++ b/libs/sdk-js/src/client.mts @@ -133,6 +133,7 @@ export class CronsClient extends BaseClient { interrupt_before: payload?.interruptBefore, interrupt_after: payload?.interruptAfter, webhook: payload?.webhook, + multitask_strategy: payload?.multitaskStrategy, }; return this.fetch(`/threads/${threadId}/runs/crons`, { method: "POST", @@ -159,6 +160,7 @@ export class CronsClient extends BaseClient { interrupt_before: payload?.interruptBefore, interrupt_after: payload?.interruptAfter, webhook: payload?.webhook, + multitask_strategy: payload?.multitaskStrategy, }; return this.fetch(`/runs/crons`, { method: "POST", @@ -237,6 +239,8 @@ export class AssistantsClient extends BaseClient { graphId: string; config?: Config; metadata?: Metadata; + assistantId?: string; + ifExists?: OnConflictBehavior; }): Promise { return this.fetch("/assistants", { method: "POST", @@ -244,6 +248,8 @@ export class AssistantsClient extends BaseClient { graph_id: payload.graphId, config: payload.config, metadata: payload.metadata, + assistant_id: payload.assistantId, + if_exists: payload.ifExists, }, }); } @@ -257,7 +263,7 @@ export class AssistantsClient extends BaseClient { async update( assistantId: string, payload: { - graphId: string; + graphId?: string; config?: Config; metadata?: Metadata; }, @@ -518,7 +524,7 @@ export class RunsClient extends BaseClient { stream( threadId: null, assistantId: string, - payload?: Omit, + payload?: Omit, ): AsyncGenerator<{ event: StreamEvent; data: any; @@ -546,8 +552,6 @@ export class RunsClient extends BaseClient { payload?: RunsStreamPayload, ): AsyncGenerator<{ event: StreamEvent; - // TODO: figure out a better way to - // type this without any data: any; }> { const json: Record = { @@ -560,10 +564,11 @@ export class RunsClient extends BaseClient { interrupt_before: payload?.interruptBefore, interrupt_after: payload?.interruptAfter, checkpoint_id: payload?.checkpointId, + webhook: payload?.webhook, + multitask_strategy: payload?.multitaskStrategy, + on_completion: payload?.onCompletion, + on_disconnect: payload?.onDisconnect, }; - if (payload?.multitaskStrategy != null) { - json["multitask_strategy"] = payload?.multitaskStrategy; - } const endpoint = threadId == null ? `/runs/stream` : `/threads/${threadId}/runs/stream`; @@ -640,10 +645,8 @@ export class RunsClient extends BaseClient { interrupt_after: payload?.interruptAfter, webhook: payload?.webhook, checkpoint_id: payload?.checkpointId, + multitask_strategy: payload?.multitaskStrategy, }; - if (payload?.multitaskStrategy != null) { - json["multitask_strategy"] = payload?.multitaskStrategy; - } return this.fetch(`/threads/${threadId}/runs`, { method: "POST", json, @@ -654,7 +657,7 @@ export class RunsClient extends BaseClient { async wait( threadId: null, assistantId: string, - payload?: Omit, + payload?: Omit, ): Promise; async wait( @@ -684,10 +687,11 @@ export class RunsClient extends BaseClient { interrupt_before: payload?.interruptBefore, interrupt_after: payload?.interruptAfter, checkpoint_id: payload?.checkpointId, + webhook: payload?.webhook, + multitask_strategy: payload?.multitaskStrategy, + on_completion: payload?.onCompletion, + on_disconnect: payload?.onDisconnect, }; - if (payload?.multitaskStrategy != null) { - json["multitask_strategy"] = payload?.multitaskStrategy; - } const endpoint = threadId == null ? `/runs/wait` : `/threads/${threadId}/runs/wait`; return this.fetch(endpoint, { diff --git a/libs/sdk-js/src/types.mts b/libs/sdk-js/src/types.mts index 872cc5723..2a5d6c4f1 100644 --- a/libs/sdk-js/src/types.mts +++ b/libs/sdk-js/src/types.mts @@ -3,6 +3,8 @@ import { Config, Metadata } from "./schema.js"; export type StreamMode = "values" | "messages" | "updates" | "events" | "debug"; export type MultitaskStrategy = "reject" | "interrupt" | "rollback" | "enqueue"; export type OnConflictBehavior = "raise" | "do_nothing"; +export type OnCompletionBehavior = "complete" | "continue"; +export type DisconnectMode = "cancel" | "continue"; export type StreamEvent = | "events" | "metadata" @@ -61,6 +63,27 @@ interface RunsInvokePayload { * Abort controller signal to cancel the run. */ signal?: AbortController["signal"]; + + /** + * Behavior to handle run completion. Only relevant if + * there is a pending/inflight run on the same thread. One of: + * - "complete": Complete the run. + * - "continue": Continue the run. + */ + onCompletion?: OnCompletionBehavior; + + /** + * Webhook to call when the run is complete. + */ + webhook?: string; + + /** + * Behavior to handle disconnection. Only relevant if + * there is a pending/inflight run on the same thread. One of: + * - "cancel": Cancel the run. + * - "continue": Continue the run. + */ + onDisconnect?: DisconnectMode; } export interface RunsStreamPayload extends RunsInvokePayload { @@ -82,12 +105,7 @@ export interface RunsStreamPayload extends RunsInvokePayload { feedbackKeys?: string[]; } -export interface RunsCreatePayload extends RunsInvokePayload { - /** - * Webhook to call when the run is complete. - */ - webhook?: string; -} +export interface RunsCreatePayload extends RunsInvokePayload {} export interface CronsCreatePayload extends RunsCreatePayload { /**