Compare commits

...
Author SHA1 Message Date
Vadym BardaandGitHub 39d9cdbef0 checkpoint-postgres: release 2.0.2 (#2183) 2024-10-24 17:06:36 -04:00
Vadym BardaandGitHub 6dacd1aabe langgraph: always raise NodeInterrupt in ToolNode if raised from a tool (#2175) 2024-10-24 18:21:18 +00:00
Vadym BardaandGitHub 91ad8b803b checkpoint-duckdb: release 2.0.1 (#2180) 2024-10-24 13:58:48 -04:00
Vadym BardaandGitHub a0e99f704f checkpoint: release 2.0.2 (#2179) 2024-10-24 13:52:02 -04:00
Vadym BardaandGitHub def3e06b4a move py.typed to submodules for namespace packages (#2177) 2024-10-24 13:43:43 -04:00
David DuongandGitHub a4fab2a867 Merge pull request #2178 from langchain-ai/dqbd/sdk-bump-24oct
feat(sdk): bump SDK to js@0.0.18 and py@0.1.34
2024-10-24 18:58:07 +02:00
Tat Dat Duong cf1c7f3673 feat(sdk): bump SDK to js@0.0.18 and py@0.1.34 2024-10-24 18:51:56 +02:00
Nuno CamposandGitHub 6c05b66c7c Merge pull request #2176 from langchain-ai/dqbd/error-thread-state
fix(sdk): add error thread state
2024-10-24 09:23:20 -07:00
Tat Dat Duong 5c44dcef81 fix(sdk): add error thread state 2024-10-24 18:16:21 +02:00
Nuno Campos 7fd6b1b4be Fix 2024-10-24 08:48:46 -07:00
Vadym BardaandGitHub 6202e0f1d9 docs: update branching how-to notebook (#2174) 2024-10-24 14:56:19 +00:00
Nuno CamposandGitHub 83238f51d8 Merge pull request #2166 from langchain-ai/nc/23oct/remote-graph-interop
Interop of RemoteGraph w core lib
2024-10-23 20:36:38 -07:00
Nuno Campos 05f008cbfb Lint 2024-10-23 20:30:32 -07:00
Nuno Campos a8ae2a52a3 Lint 2024-10-23 20:27:18 -07:00
Nuno Campos aa245a8e71 Fix up 2024-10-23 20:22:04 -07:00
bdc75a22d5 langgraph: expand handle_tool_errors in ToolNode (#1667)
This change expands error-handling functionality of the `ToolNode` by
introducing more options for `handle_tool_errors`. Default behavior of
the `ToolNode` is unchanged -- all errors are handled and wrapped in a
`ToolMessage` to be sent back to LLM.

With this change, users have flexibility to only handle the exceptions
that they need to pass back to the LLM:

* they can specify exceptions to handle by passing a tuple of exceptions
in `handle_tool_errors`
* specify `handle_tool_errors=True/str/callable`
* when `handle_tool_errors` is a callable, the signature will be
inspected and exceptions from the signature will be handled

---------

Co-authored-by: vbarda <vadym@langchain.dev>
2024-10-24 00:58:05 +00:00
Nuno Campos 69227daff3 Lint 2024-10-23 17:01:58 -07:00
Nuno Campos dc8260bb72 Interop of RemoteGraph w core lib 2024-10-23 15:37:17 -07:00
Vadym BardaandGitHub 62a5ec509d checkpoint: add DuckDB store (#2154) 2024-10-23 22:20:10 +00:00
Vadym BardaandGitHub d32386f849 checkpoint: add DuckDB checkpointer (#2145) 2024-10-23 21:11:03 +00:00
Nuno CamposandGitHub 08a1ed38f1 Merge pull request #2092 from langchain-ai/an/11oct/remote-graph-interrupt
Update `stream()` and `astream()` methods in `RemoteGraph` to process `updates` event types
2024-10-23 13:48:43 -07:00
Nuno Campos 037a95ff60 Update tests 2024-10-23 13:43:48 -07:00
Nuno Campos e294720ec5 Lint 2024-10-23 13:29:59 -07:00
Nuno Campos f8a0b7a464 Use if_not_exists 2024-10-23 13:24:19 -07:00
Nuno Campos 1121806ba4 Add if_not_exists 2024-10-23 13:23:49 -07:00
Nuno Campos dca200d6c4 Finish 2024-10-23 13:23:42 -07:00
Vadym BardaandGitHub 6f236b5f2c docs: update multi-agent concept examples (#2151) 2024-10-23 13:40:24 -04:00
Vadym BardaandGitHub 916affa1b5 langgraph: add 'messages_key' param to ToolNode / tools_condition (#2049) 2024-10-22 17:32:05 -04:00
Nuno Campos 58cf0c6a6e chore: Switch s3 client utils from httpx client to curl client 2024-10-22 10:46:37 -07:00
gbaian10andGitHub 0042889c31 Support read type hints from the method in add_node (#2014)
Add function to read type hints from the `__call__` method to resolve issue #1950.
2024-10-22 17:11:59 +00:00
nikhildigdeandGitHub 2be012d8ed docs: fix typo in concept docs 2024-10-21 21:40:47 +00:00
Yuki OshimaandGitHub 780285ef91 Fix(docs): InMemoryStore example error (#2148) 2024-10-21 18:39:43 +00:00
Nuno CamposandGitHub 47c7b76aa1 Merge pull request #2152 from langchain-ai/nc/21oct/skip-docker-login-for-forks
ci: Skip docker login for PRs from forks
2024-10-21 10:09:54 -07:00
Nuno Campos 9931f61525 ci: Skip docker login for PRs from forks 2024-10-21 10:04:40 -07:00
Nuno CamposandGitHub 0d81ad92f0 Merge pull request #2143 from langchain-ai/nc/19oct/async-max-concurrency
lib: Add max_concurrency for async executions
2024-10-19 15:16:15 -07:00
Nuno Campos 7d3f2ca3ed Lint 2024-10-19 15:11:18 -07:00
Nuno Campos 42648c88dd lib: Add max_concurrency for async executions 2024-10-19 12:45:01 -07:00
Vadym BardaandGitHub 1aab758634 langgraph: release 0.2.39 (#2139) 2024-10-18 14:47:50 -04:00
Vadym BardaandGitHub b5fbc7a7b8 docs: temporarily disable some link checks (#2138) 2024-10-18 14:37:19 -04:00
b647dcb0f2 feat: Add LangGraph error pages (#2136)
Co-authored-by: Erick Friis <erick@langchain.dev>
Co-authored-by: vbarda <vadym@langchain.dev>
2024-10-18 18:20:22 +00:00
Nuno CamposandGitHub 4df5680732 Merge pull request #2132 from langchain-ai/nc/17oct/stream-messages-nostream-tag
For stream_mode=messages skip any nodes/llms with tag nostream
2024-10-17 15:57:09 -07:00
Nuno Campos 74a17a6d4c Update tests 2024-10-17 15:52:17 -07:00
Nuno Campos e2a3698250 For stream_mode=messages skip any nodes/llms with tag nostream 2024-10-17 15:28:14 -07:00
Vadym BardaandGitHub 4dfdb9a83e docs: update tags for store endpoints in API docs (#2127) 2024-10-16 16:34:54 +00:00
Vadym BardaandGitHub 583d8c9499 docs: update tutorial names/links (#2126) 2024-10-16 15:17:28 +00:00
vbarda 15bbede7bc update image in multi-agent concepts 2024-10-16 10:59:07 -04:00
Vadym BardaandGitHub 3ffdf4bb3f docs: update image in concepts (#2125) 2024-10-16 13:59:17 +00:00
Nuno CamposandGitHub 6578698414 Merge pull request #2124 from langchain-ai/dqbd/js-bump-0.0.17
feat(sdk-js): bump to 0.0.17
2024-10-16 06:21:22 -07:00
Tat Dat Duong 048ae6c17b feat(sdk-js): bump to 0.0.17 2024-10-16 15:19:39 +02:00
Nuno CamposandGitHub 0e2c2eb13a Merge pull request #2120 from langchain-ai/nc/15oct/executor-dict
fix: Avoid errors from executor modifying tasks dict during exit routine
2024-10-15 16:15:55 -07:00
Nuno Campos 515c4ffebe Fix 2024-10-15 16:11:08 -07:00
Nuno Campos eefe057a47 fix: Avoid errors from executor modifying tasks dict during exit routine
- This could happen if a task happened to finish while the exit routine is running
2024-10-15 15:13:46 -07:00
18f34c30d8 docs: update subgraph how-to (#2079)
Co-authored-by: vbarda <vadym@langchain.dev>
2024-10-15 17:30:34 -04:00
2670bcf330 docs: add concepts for subgraphs and multi-agent (#2069)
Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
Co-authored-by: Nuno Campos <nuno@langchain.dev>
2024-10-15 17:17:22 -04:00
Andrew Nguonly 19ccb0c6af Update astream_events() to process interrupt. 2024-10-11 19:06:52 -07:00
Andrew Nguonly a277b86fcb Fix unit test. 2024-10-11 18:53:45 -07:00
Andrew Nguonly 2f819a6a9b Update stream() and astream() to process 'updates' event types. 2024-10-11 18:40:40 -07:00
93 changed files with 7692 additions and 1203 deletions
+1
View File
@@ -35,6 +35,7 @@ jobs:
cache-key: test-${{ inputs.working-directory }}
- name: Login to Docker Hub
uses: docker/login-action@v3
if: ${{ !github.event.pull_request.head.repo.fork }}
with:
username: ${{ secrets.DOCKERHUB_USERNAME }}
password: ${{ secrets.DOCKERHUB_RO_TOKEN }}
+1
View File
@@ -37,6 +37,7 @@ jobs:
cache-key: test-langgraph
- name: Login to Docker Hub
uses: docker/login-action@v3
if: ${{ !github.event.pull_request.head.repo.fork }}
with:
username: ${{ secrets.DOCKERHUB_USERNAME }}
password: ${{ secrets.DOCKERHUB_RO_TOKEN }}
@@ -29,6 +29,7 @@ jobs:
cache-key: test-scheduler-kafka
- name: Login to Docker Hub
uses: docker/login-action@v3
if: ${{ !github.event.pull_request.head.repo.fork }}
with:
username: ${{ secrets.DOCKERHUB_USERNAME }}
password: ${{ secrets.DOCKERHUB_RO_TOKEN }}
+2
View File
@@ -31,6 +31,7 @@ jobs:
"libs/cli",
"libs/checkpoint",
"libs/checkpoint-sqlite",
"libs/checkpoint-duckdb",
"libs/checkpoint-postgres",
"libs/scheduler-kafka",
]
@@ -47,6 +48,7 @@ jobs:
"libs/cli",
"libs/checkpoint",
"libs/checkpoint-sqlite",
"libs/checkpoint-duckdb",
"libs/checkpoint-postgres"
]
uses: ./.github/workflows/_test.yml
+1
View File
@@ -81,6 +81,7 @@ jobs:
--check-links-ignore "https://x.com/.*" \
--check-links-ignore "https://github\.com/.*" \
--check-links-ignore "/.*\.(ipynb|html)$" \
--check-links-ignore "https://python\.langchain\.com/.*" \
--check-links $(find docs/site -name "index.html" | grep -v 'storm/index.html')
else
echo "Fetching changes from origin/main..."
@@ -1867,6 +1867,7 @@
"/store/items": {
"put": {
"tags": ["store/manage"],
"summary": "Store or update an item.",
"operationId": "put_item",
"requestBody": {
@@ -1892,6 +1893,7 @@
}
},
"delete": {
"tags": ["store/manage"],
"summary": "Delete an item.",
"operationId": "delete_item",
"requestBody": {
@@ -1917,6 +1919,7 @@
}
},
"get": {
"tags": ["store/manage"],
"summary": "Retrieve a single item.",
"operationId": "get_item",
"parameters": [
@@ -1962,6 +1965,7 @@
},
"/store/items/search": {
"post": {
"tags": ["store/manage"],
"summary": "Search for items within a namespace prefix.",
"operationId": "search_items",
"requestBody": {
@@ -1994,6 +1998,7 @@
},
"/store/namespaces": {
"post": {
"tags": ["store/manage"],
"summary": "List namespaces with optional match conditions.",
"operationId": "list_namespaces",
"requestBody": {
+3 -3
View File
@@ -103,15 +103,15 @@ Parallel processing is vital for efficient multi-agent systems and complex tasks
For practical implementation, see our [map-reduce tutorial](../how-tos/map-reduce.ipynb).
### Sub-graphs
### Subgraphs
Sub-graphs are essential for managing complex agent architectures, particularly in multi-agent systems. They allow:
[Subgraphs](./low_level.md#subgraphs) are essential for managing complex agent architectures, particularly in [multi-agent systems](./multi_agent.md). They allow:
- Isolated state management for individual agents
- Hierarchical organization of agent teams
- Controlled communication between agents and the main system
Sub-graphs communicate with the parent graph through overlapping keys in the state schema. This enables flexible, modular agent design. For implementation details, refer to our [sub-graph tutorial](../how-tos/subgraph.ipynb).
Subgraphs communicate with the parent graph through overlapping keys in the state schema. This enables flexible, modular agent design. For implementation details, refer to our [subgraph how-to guide](../how-tos/subgraph.ipynb).
### Reflection
Binary file not shown.

After

Width:  |  Height:  |  Size: 178 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 193 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 55 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 83 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 103 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 97 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 35 KiB

+112 -9
View File
@@ -20,7 +20,7 @@ A super-step can be considered a single iteration over the graph nodes. Nodes th
### StateGraph
The `StateGraph` class is the main graph class to uses. This is parameterized by a user defined `State` object.
The `StateGraph` class is the main graph class to use. This is parameterized by a user defined `State` object.
### MessageGraph
@@ -52,12 +52,12 @@ By default, the graph will have the same input and output schemas. If you want t
Typically, all graph nodes communicate with a single schema. This means that they will read and write to the same state channels. But, there are cases where we want more control over this:
* Internal nodes can pass information that is not required in the graph's input / output.
* We may also want to use different input / output schemas for the graph. The output might, for example, only contain a single relevant output key.
- Internal nodes can pass information that is not required in the graph's input / output.
- We may also want to use different input / output schemas for the graph. The output might, for example, only contain a single relevant output key.
It is possible to have nodes write to private state channels inside the graph for internal node communication. We can simply define a private schema, `PrivateState`. See [this notebook](../how-tos/pass_private_state.ipynb) for more detail.
It is possible to have nodes write to private state channels inside the graph for internal node communication. We can simply define a private schema, `PrivateState`. See [this notebook](../how-tos/pass_private_state.ipynb) for more detail.
It is also possible to define explicit input and output schemas for a graph. In these cases, we define an "internal" schema that contains *all* keys relevant to graph operations. But, we also define `input` and `output` schemas that are sub-sets of the "internal" schema to constrain the input and output of the graph. See [this notebook](../how-tos/input_output_schema.ipynb) for more detail.
It is also possible to define explicit input and output schemas for a graph. In these cases, we define an "internal" schema that contains _all_ keys relevant to graph operations. But, we also define `input` and `output` schemas that are sub-sets of the "internal" schema to constrain the input and output of the graph. See [this notebook](../how-tos/input_output_schema.ipynb) for more detail.
Let's look at an example:
@@ -101,11 +101,12 @@ graph = builder.compile()
graph.invoke({"user_input":"My"})
{'graph_output': 'My name is Lance'}
```
There are two subtle and important points to note here:
1. We pass `state: InputState` as the input schema to `node_1`. But, we write out to `foo`, a channel in `OverallState`. How can we write out to a state channel that is not included in the input schema? This is because a node *can write to any state channel in the graph state.* The graph state is the union of of the state channels defined at initialization, which includes `OverallState` and the filters `InputState` and `OutputState`.
1. We pass `state: InputState` as the input schema to `node_1`. But, we write out to `foo`, a channel in `OverallState`. How can we write out to a state channel that is not included in the input schema? This is because a node _can write to any state channel in the graph state._ The graph state is the union of of the state channels defined at initialization, which includes `OverallState` and the filters `InputState` and `OutputState`.
2. We initialize the graph with `StateGraph(OverallState,input=InputState,output=OutputState)`. So, how can we write to `PrivateState` in `node_2`? How does the graph gain access to this schema if it was not passed in the `StateGraph` initialization? We can do this because *nodes can also declare additional state channels* as long as the state schema definition exists. In this case, the `PrivateState` schema is defined, so we can add `bar` as a new state channel in the graph and write to it.
2. We initialize the graph with `StateGraph(OverallState,input=InputState,output=OutputState)`. So, how can we write to `PrivateState` in `node_2`? How does the graph gain access to this schema if it was not passed in the `StateGraph` initialization? We can do this because _nodes can also declare additional state channels_ as long as the state schema definition exists. In this case, the `PrivateState` schema is defined, so we can add `bar` as a new state channel in the graph and write to it.
### Reducers
@@ -323,7 +324,7 @@ graph.add_conditional_edges("node_a", continue_to_jokes)
## Persistence
LangGraph provides built-in persistence for your agent's state using [checkpointers][langgraph.checkpoint.base.BaseCheckpointSaver]. Checkpointers save snapshots of the graph state at every superstep, allowing resumption at any time. This enables features like human-in-the-loop interactions, memory management, and fault-tolerance. You can even directly manipulate a graph's state after its execution using the
LangGraph provides built-in persistence for your agent's state using [checkpointers][langgraph.checkpoint.base.BaseCheckpointSaver]. Checkpointers save snapshots of the graph state at every superstep, allowing resumption at any time. This enables features like human-in-the-loop interactions, memory management, and fault-tolerance. You can even directly manipulate a graph's state after its execution using the
appropriate `get` and `update` methods. For more details, see the [persistence conceptual guide](./persistence.md).
## Threads
@@ -416,10 +417,112 @@ def my_node(state: State) -> State:
return state
```
## Subgraphs
A subgraph is a [graph](#graphs) that is used as a [node](#nodes) in another graph. This is nothing more than the age-old concept of encapsulation, applied to LangGraph. Some reasons for using subgraphs are:
- building [multi-agent systems](./multi_agent.md)
- when you want to reuse a set of nodes in multiple graphs, which maybe share some state, you can define them once in a subgraph and then use them in multiple parent graphs
- when you want different teams to work on different parts of the graph independently, you can define each part as a subgraph, and as long as the subgraph interface (the input and output schemas) is respected, the parent graph can be built without knowing any details of the subgraph
There are two ways to add subgraphs to a parent graph:
- add a node with the compiled subgraph: this is useful when the parent graph and the subgraph share state keys and you don't need to transform state on the way in or out
```python
builder.add_node("subgraph", subgraph_builder.compile())
```
- add a node with a function that invokes the subgraph: this is useful when the parent graph and the subgraph have different state schemas and you need to transform state before or after calling the subgraph
```python
subgraph = subgraph_builder.compile()
def call_subgraph(state: State):
return subgraph.invoke({"subgraph_key": state["parent_key"]})
builder.add_node("subgraph", call_subgraph)
```
Let's take a look at examples for each.
### As a compiled graph
The simplest way to create subgraph nodes is by using a [compiled subgraph](#compiling-your-graph) directly. When doing so, it is **important** that the parent graph and the subgraph [state schemas](#state) share at least one key which they can use to communicate. If your graph and subgraph do not share any keys, you should use write a function [invoking the subgraph](#as-a-function) instead.
!!! Note
If you pass extra keys to the subgraph node (i.e., in addition to the shared keys), they will be ignored by the subgraph node. Similarly, if you return extra keys from the subgraph, they will be ignored by the parent graph.
```python
from langgraph.graph import START, StateGraph
from typing import TypedDict
class State(TypedDict):
foo: str
class SubgraphState(TypedDict):
foo: str # note that this key is shared with the parent graph state
bar: str
# Define subgraph
def subgraph_node(state: SubgraphState):
# note that this subgraph node can communicate with the parent graph via the shared "foo" key
return {"foo": state["foo"] + "bar"}
subgraph_builder = StateGraph(SubgraphState)
subgraph_builder.add_node(subgraph_node)
...
subgraph = subgraph_builder.compile()
# Define parent graph
builder = StateGraph(State)
builder.add_node("subgraph", subgraph)
...
graph = builder.compile()
```
### As a function
You might want to define a subgraph with a completely different schema. In this case, you can create a node function that invokes the subgraph. This function will need to [transform](../how-tos/subgraph-transform-state.ipynb) the input (parent) state to the subgraph state before invoking the subgraph, and transform the results back to the parent state before returning the state update from the node.
```python
class State(TypedDict):
foo: str
class SubgraphState(TypedDict):
# note that none of these keys are shared with the parent graph state
bar: str
baz: str
# Define subgraph
def subgraph_node(state: SubgraphState):
return {"bar": state["bar"] + "baz"}
subgraph_builder = StateGraph(SubgraphState)
subgraph_builder.add_node(subgraph_node)
...
subgraph = subgraph_builder.compile()
# Define parent graph
def node(state: State):
# transform the state to the subgraph state
response = subgraph.invoke({"bar": state["foo"]})
# transform response back to the parent state
return {"foo": response["bar"]}
builder = StateGraph(State)
# note that we are using `node` function instead of a compiled subgraph
builder.add_node(node)
...
graph = builder.compile()
```
## Visualization
It's often nice to be able to visualize graphs, especially as they get more complex. LangGraph comes with several built-in ways to visualize graphs. See [this how-to guide](../how-tos/visualization.ipynb) for more info.
## Streaming
LangGraph is built with first class support for streaming, including streaming updates from graph nodes during the execution, streaming tokens from LLM calls and more. See this [conceptual guide](./streaming.md) for more information.
LangGraph is built with first class support for streaming, including streaming updates from graph nodes during the execution, streaming tokens from LLM calls and more. See this [conceptual guide](./streaming.md) for more information.
+2 -2
View File
@@ -183,9 +183,9 @@ store = InMemoryStore()
user_id = "my-user"
application_context = "chitchat"
namespace = (user_id, application_context)
store.put(namespace, key="a-memory", {"rules": ["User likes short, direct language", "User only speaks English & python"], "my-key": "my-value"})
store.put(namespace, "a-memory", {"rules": ["User likes short, direct language", "User only speaks English & python"], "my-key": "my-value"})
# get the "memory" by ID
item = store.get(namespace)
item = store.get(namespace, "a-memory")
# list "memories" within this namespace, filtering on content equivalence
items = store.search(namespace, filter={"my-key": "my-value"})
```
+225 -82
View File
@@ -1,138 +1,281 @@
# Multi-agent Systems
A multi-agent system is a system with multiple independent actors powered by LLMs that are connected in a specific way. These actors can be as simple as a prompt and an LLM call, or as complex as a [ReAct](./agentic_concepts.md#react-implementation) agent.
An [agent](./agentic_concepts.md#agent-architectures) is _a system that uses an LLM to decide the control flow of an application_. As you develop these systems, they might grow more complex over time, making them harder to manage and scale. For example, you might run into the following problems:
The primary benefits of this architecture are:
- agent has too many tools at its disposal and makes poor decisions about which tool to call next
- context grows too complex for a single agent to keep track of
- there is a need for multiple specialization areas in the system (e.g. planner, researcher, math expert, etc.)
* **Modularity**: Separate agents facilitate easier development, testing, and maintenance of agentic systems.
* **Specialization**: You can create expert agents focused on specific domains, and compose them into more complex applications
* **Control**: You can explicitly control how agents communicate (as opposed to relying on function calling)
To tackle these, you might consider breaking your application into multiple smaller, independent agents and composing them into a **multi-agent system**. These independent agents can be as simple as a prompt and an LLM call, or as complex as a [ReAct](./agentic_concepts.md#react-implementation) agent (and more!).
## Multi-agent systems in LangGraph
The primary benefits of using multi-agent systems are:
### Agents as nodes
- **Modularity**: Separate agents make it easier to develop, test, and maintain agentic systems.
- **Specialization**: You can create expert agents focused on specific domains, which helps with the overall system performance.
- **Control**: You can explicitly control how agents communicate (as opposed to relying on function calling).
Agents can be defined as nodes in LangGraph. As any other node in the LangGraph, these agent nodes receive the graph state as an input and return an update to the state as their output.
## Multi-agent architectures
* Simple **LLM nodes**: single LLMs with custom prompts
* **Subgraph nodes**: complex graphs called inside the orchestrator graph node
![](./img/multi_agent/architectures.png)
![](./img/multi_agent/subgraph.png)
There are several ways to connect agents in a multi-agent system:
### Agents as tools
- **Network**: each agent can communicate with [every other agent](https://langchain-ai.github.io/langgraph/tutorials/multi_agent/multi-agent-collaboration/). Any agent can decide which other agent to call next.
- **Supervisor**: each agent communicates with a single [supervisor](https://langchain-ai.github.io/langgraph/tutorials/multi_agent/agent_supervisor/) agent. Supervisor agent makes decisions on which agent should be called next.
- **Supervisor (tool-calling)**: this is a special case of supervisor architecture. Individual agents can be represented as tools. In this case, a supervisor agent uses a tool-calling LLM to decide which of the agent tools to call, as well as the arguments to pass to those agents.
- **Hierarchical**: you can define a multi-agent system with [a supervisor of supervisors](https://langchain-ai.github.io/langgraph/tutorials/multi_agent/hierarchical_agent_teams/). This is a generalization of the supervisor architecture and allows for more complex control flows.
- **Custom multi-agent workflow**: each agent communicates with only a subset of agents. Parts of the flow are deterministic, and only some agents can decide which other agents to call next.
Agents can also be defined as tools. In this case, the orchestrator agent (e.g. ReAct agent) would use a tool-calling LLM to decide which of the agent tools to call, as well as the arguments to pass to those agents.
### Network
You could also take a "mega-graph" approach incorporating subordinate agents' nodes directly into the parent, orchestrator graph. However, this is not recommended for complex subordinate agents, as it would make the overall system harder to scale, maintain and debug you should use subgraphs or tools in those cases.
In this architecture, agents are defined as graph nodes. Each agent can communicate with every other agent (many-to-many connections) and can decide which agent to call next. While very flexible, this architecture doesn't scale well as the number of agents grows:
## Communication in multi-agent systems
- hard to enforce which agent should be called next
- hard to determine how much [information](#shared-message-list) should be passed between the agents
A big question in multi-agent systems is how the agents communicate amongst themselves and with the orchestrator agent. This involves both the schema of how they communicate, as well as the sequence in which they communicate. LangGraph is perfect for orchestrating these types of systems and allows you to define both.
We recommend avoiding this architecture in production and using one of the below architectures instead.
### Schema
### Supervisor
LangGraph provides a lot of flexibility for how to communicate within multi-agent architectures.
* A node in LangGraph can have a [private input state schema](https://langchain-ai.github.io/langgraph/how-tos/pass_private_state/) that is distinct from the graph state schema. This allows passing additional information during the graph execution that is only needed for executing a particular node.
* Subgraph node agents can have independent [input / output state schemas](https://langchain-ai.github.io/langgraph/how-tos/input_output_schema/). In this case its important to [add input / output transformations](https://langchain-ai.github.io/langgraph/how-tos/subgraph-transform-state/) so that the parent graph knows how to communicate with the subgraphs.
* For tool-based subordinate agents, the orchestrator determines the inputs based on the tool schema. Additionally, LangGraph allows passing state to individual tools at runtime, so subordinate agents can access parent state, if needed.
### Sequence
LangGraph provides multiple methods to control agent communication sequence:
* **Explicit control flow (graph edges)**: LangGraph allows you to define the control flow of your application (i.e. the sequence of how agents communicate) explicitly, via [graph edges](./low_level.md#edges).
In this architecture, we define agents as nodes and add a supervisor node (LLM) that decides which agent nodes should be called next. We use [conditional edges](./low_level.md#conditional-edges) to route execution to the appropriate agent node based on supervisor's decision. This architecture also lends itself well to running multiple agents in parallel or using [map-reduce](../how-tos/map-reduce.ipynb) pattern.
```python
from typing import Literal
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage
from langgraph.graph import StateGraph, MessagesState, START, END
from langgraph.graph import StateGraph, MessagesState, START
model = ChatOpenAI(model="gpt-4o-mini")
model = ChatOpenAI()
def research_agent(state: MessagesState):
"""Call research agent"""
messages = [SystemMessage(content="You are a research assistant. Given a topic, provide key facts and information.")] + state["messages"]
response = model.invoke(messages)
class AgentState(MessagesState):
next: Literal["agent_1", "agent_2", "__end__"]
def supervisor(state: AgentState):
# you can pass relevant parts of the state to the LLM (e.g., state["messages"])
# to determine which agent to call next. a common pattern is to call the model
# with a structured output (e.g. force it to return an output with a "next_agent" field)
response = model.invoke(...)
# the "next" key will be used by the conditional edges to route execution
# to the appropriate agent
return {"next": response["next_agent"]}
def agent_1(state: AgentState):
# you can pass relevant parts of the state to the LLM (e.g., state["messages"])
# and add any additional logic (different models, custom prompts, structured output, etc.)
response = model.invoke(...)
return {"messages": [response]}
def summarize_agent(state: MessagesState):
"""Call summarization agent"""
messages = [SystemMessage(content="You are a summarization expert. Condense the given information into a brief summary.")] + state["messages"]
response = model.invoke(messages)
def agent_2(state: AgentState):
response = model.invoke(...)
return {"messages": [response]}
graph = StateGraph(MessagesState)
graph.add_node("research", research_agent)
graph.add_node("summarize", summarize_agent)
builder = StateGraph(AgentState)
builder.add_node(supervisor)
builder.add_node(agent_1)
builder.add_node(agent_2)
# define the flow explicitly
graph.add_edge(START, "research")
graph.add_edge("research", "summarize")
graph.add_edge("summarize", END)
builder.add_edge(START, "supervisor")
# route to one of the agents or exit based on the supervisor's decisiion
# if the supervisor returns "__end__", the graph will finish execution
builder.add_conditional_edges("supervisor", lambda state: state["next"])
builder.add_edge("agent_1", "supervisor")
builder.add_edge("agent_2", "supervisor")
supervisor = builder.compile()
```
* **Dynamic control flow (conditional edges)**: LangGraph also allows you to define [conditional edges](./low_level.md#conditional-edges), where the control flow is dependent on satisfying a given condition. In such cases, you can use an LLM to decide which subordinate agent to call next.
Check out this [tutorial](https://langchain-ai.github.io/langgraph/tutorials/multi_agent/agent_supervisor/) for an example of supervisor multi-agent architecture.
### Supervisor (tool-calling)
* **Implicit control flow (tool calling)**: if the orchestrator agent treats subordinate agents as tools, the tool-calling LLM powering the orchestrator will make decisions about the order in which the tools (agents) are being called.
In this variant of the [supervisor](#supervisor) architecture, we define individual agents as **tools** and use a tool-calling LLM in the supervisor node. This can be implemented as a [ReAct](./agentic_concepts.md#react-implementation)-style agent with two nodes — an LLM node (supervisor) and a tool-calling node that executes tools (agents in this case).
```python
from typing import Annotated
from langchain_core.messages import SystemMessage, ToolMessage
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import ToolNode, InjectedState, create_react_agent
from langgraph.prebuilt import InjectedState, create_react_agent
model = ChatOpenAI(model="gpt-4o-mini")
model = ChatOpenAI()
def research_agent(state: Annotated[dict, InjectedState]):
"""Call research agent"""
messages = [SystemMessage(content="You are a research assistant. Given a topic, provide key facts and information.")] + state["messages"][:-1]
response = model.invoke(messages)
tool_call = state["messages"][-1].tool_calls[0]
return {"messages": [ToolMessage(response.content, tool_call_id=tool_call["id"])]}
# this is the agent function that will be called as tool
# notice that you can pass the state to the tool via InjectedState annotation
def agent_1(state: Annotated[dict, InjectedState]):
# you can pass relevant parts of the state to the LLM (e.g., state["messages"])
# and add any additional logic (different models, custom prompts, structured output, etc.)
response = model.invoke(...)
# return the LLM response as a string (expected tool response format)
# this will be automatically turned to ToolMessage
# by the prebuilt create_react_agent (supervisor)
return response.content
def summarize_agent(state: Annotated[dict, InjectedState]):
"""Call summarization agent"""
messages = [SystemMessage(content="You are a summarization expert. Condense the given information into a brief summary.")] + state["messages"][:-1]
response = model.invoke(messages)
tool_call = state["messages"][-1].tool_calls[0]
return {"messages": [ToolMessage(response.content, tool_call_id=tool_call["id"])]}
def agent_2(state: Annotated[dict, InjectedState]):
response = model.invoke(...)
return response.content
tool_node = ToolNode([research_agent, summarize_agent])
graph = create_react_agent(model, [research_agent, summarize_agent], state_modifier="First research and then summarize information on a given topic.")
tools = [agent_1, agent_2]
# the simplest way to build a supervisor w/ tool-calling is to use prebuilt ReAct agent graph
# that consists of a tool-calling LLM node (i.e. supervisor) and a tool-executing node
supervisor = create_react_agent(model, tools)
```
## Example architectures
### Hierarchical
Below are several examples of complex multi-agent architectures that can be implemented in LangGraph.
As you add more agents to your system, it might become too hard for the supervisor to manage all of them. The supervisor might start making poor decisions about which agent to call next, the context might become too complex for a single supervisor to keep track of. In other words, you end up with the same problems that motivated the multi-agent architecture in the first place.
### Multi-Agent Collaboration
To address this, you can design your system _hierarchically_. For example, you can create separate, specialized teams of agents managed by individual supervisors, and a top-level supervisor to manage the teams.
In this example, different agents collaborate on a **shared** scratchpad of messages (i.e. shared graph state). This means that all the work any of them do is visible to the other ones. The benefit is that the other agents can see all the individual steps done. The downside is that sometimes is it overly verbose and unnecessary to pass ALL this information along, and sometimes only the final answer from an agent is needed. We call this **collaboration** because of the shared nature the scratchpad.
```python
from typing import Literal
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, MessagesState, START
In this case, the independent agents are actually just a single LLM call with a custom system message.
model = ChatOpenAI()
Here is a visualization of how these agents are connected:
# define team 1 (same as the single supervisor example above)
class Team1State(MessagesState):
next: Literal["team_1_agent_1", "team_1_agent_2", "__end__"]
![](./img/multi_agent/collaboration.png)
def team_1_supervisor(state: Team1State):
response = model.invoke(...)
return {"next": response["next_agent"]}
See full code example in this [tutorial](https://langchain-ai.github.io/langgraph/tutorials/multi_agent/multi-agent-collaboration/).
def team_1_agent_1(state: Team1State):
response = model.invoke(...)
return {"messages": [response]}
### Agent Supervisor
def team_1_agent_2(state: Team1State):
response = model.invoke(...)
return {"messages": [response]}
In this example, multiple agents are connected, but compared to above they do NOT share a shared scratchpad. Rather, they have their own independent scratchpads (i.e. their own state), and then their final responses are appended to a global scratchpad.
team_1_builder = StateGraph(Team1State)
team_1_builder.add_node(team_1_supervisor)
team_1_builder.add_node(team_1_agent_1)
team_1_builder.add_node(team_1_agent_2)
team_1_builder.add_edge(START, "team_1_supervisor")
# route to one of the agents or exit based on the supervisor's decisiion
# if the supervisor returns "__end__", the graph will finish execution
team_1_builder.add_conditional_edges("team_1_supervisor", lambda state: state["next"])
team_1_builder.add_edge("team_1_agent_1", "team_1_supervisor")
team_1_builder.add_edge("team_1_agent_2", "team_1_supervisor")
In this case, the independent agents are a LangGraph ReAct agent (graph). This means they have their own individual prompt, LLM, and tools. When called, it's not just a single LLM call, but rather an invocation of the graph powering the ReAct agent.
team_1_graph = team_1_builder.compile()
![](./img/multi_agent/supervisor.png)
# define team 2 (same as the single supervisor example above)
class Team2State(MessagesState):
next: Literal["team_2_agent_1", "team_2_agent_2", "__end__"]
See full code example in this [tutorial](https://langchain-ai.github.io/langgraph/tutorials/multi_agent/agent_supervisor/).
def team_2_supervisor(state: Team2State):
...
### Hierarchical Agent Teams
def team_2_agent_1(state: Team2State):
...
What if the job for a single worker in agent supervisor example becomes too complex? What if the number of workers becomes too large? For some applications, the system may be more effective if work is distributed hierarchically. You can do this by creating additional level of subgraphs and creating a top-level supervisor, along with mid-level supervisors:
def team_2_agent_2(state: Team2State):
...
![](./img/multi_agent/hierarchical.png)
team_2_builder = StateGraph(Team2State)
...
team_2_graph = team_2_builder.compile()
See full code example in this [tutorial](https://langchain-ai.github.io/langgraph/tutorials/multi_agent/hierarchical_agent_teams/).
# define top-level supervisor
class TopLevelState(MessagesState):
next: Literal["team_1", "team_2", "__end__"]
builder = StateGraph(TopLevelState)
def top_level_supervisor(state: TopLevelState):
# you can pass relevant parts of the state to the LLM (e.g., state["messages"])
# to determine which team to call next. a common pattern is to call the model
# with a structured output (e.g. force it to return an output with a "next_team" field)
response = model.invoke(...)
# the "next" key will be used by the conditional edges to route execution
# to the appropriate team
return {"next": response["next_team"]}
builder = StateGraph(TopLevelState)
builder.add_node(top_level_supervisor)
builder.add_node(team_1_graph)
builder.add_node(team_2_graph)
builder.add_edge(START, "top_level_supervisor")
# route to one of the teams or exit based on the supervisor's decision
# if the top-level supervisor returns "__end__", the graph will finish execution
builder.add_conditional_edges("top_level_supervisor", lambda state: state["next"])
builder.add_edge("team_1_graph", "top_level_supervisor")
builder.add_edge("team_2_graph", "top_level_supervisor")
graph = builder.compile()
```
### Custom multi-agent workflow
In this architecture we add individual agents as graph nodes and define the order in which agents are called ahead of time, in a custom workflow. In LangGraph the workflow can be defined in two ways:
- **Explicit control flow (normal edges)**: LangGraph allows you to explicitly define the control flow of your application (i.e. the sequence of how agents communicate) explicitly, via [normal graph edges](./low_level.md#normal-edges). This is the most deterministic variant of this architecture above — we always know which agent will be called next ahead of time.
- **Dynamic control flow (conditional edges)**: in LangGraph you can allow LLMs to decide parts of your application control flow. This can be achieved by using [conditional edges](./low_level.md#conditional-edges). A special case of this is a [supervisor tool-calling](#supervisor-tool-calling) architecture. In that case, the tool-calling LLM powering the supervisor agent will make decisions about the order in which the tools (agents) are being called.
```python
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, MessagesState, START
model = ChatOpenAI()
def agent_1(state: MessagesState):
response = model.invoke(...)
return {"messages": [response]}
def agent_2(state: MessagesState):
response = model.invoke(...)
return {"messages": [response]}
builder = StateGraph(MessagesState)
builder.add_node(agent_1)
builder.add_node(agent_2)
# define the flow explicitly
builder.add_edge(START, "agent_1")
builder.add_edge("agent_1", "agent_2")
```
## Communication between agents
The most important thing when building multi-agent systems is figuring out how the agents communicate. There are few different considerations:
- Do agents communicate via [**via graph state or via tool calls**](#graph-state-vs-tool-calls)?
- What if two agents have [**different state schemas**](#different-state-schemas)?
- How to communicate over a [**shared message list**](#shared-message-list)?
### Graph state vs tool calls
What is the "payload" that is being passed around between agents? In most of the architectures discussed above the agents communicate via the [graph state](./low_level.md#state). In the case of the [supervisor with tool-calling](#supervisor-tool-calling), the payloads are tool call arguments.
![](./img/multi_agent/request.png)
#### Graph state
To communicate via graph state, individual agents need to be defined as [graph nodes](./low_level.md#nodes). These can be added as functions or as entire [subgraphs](./low_level.md#subgraphs). At each step of the graph execution, agent node receives the current state of the graph, executes the agent code and then passes the updated state to the next nodes.
Typically agent nodes share a single [state schema](./low_level.md#schema). However, you might want to design agent nodes with [different state schemas](#different-state-schemas).
### Different state schemas
An agent might need to have a different state schema from the rest of the agents. For example, a search agent might only need to keep track of queries and retrieved documents. There are two ways to achieve this in LangGraph:
- Define [subgraph](./low_level.md#subgraphs) agents with a separate state schema. If there are no shared state keys (channels) between the subgraph and the parent graph, its important to [add input / output transformations](https://langchain-ai.github.io/langgraph/how-tos/subgraph-transform-state/) so that the parent graph knows how to communicate with the subgraphs.
- Define agent node functions with a [private input state schema](https://langchain-ai.github.io/langgraph/how-tos/pass_private_state/) that is distinct from the overall graph state schema. This allows passing information that is only needed for executing that particular agent.
### Shared message list
The most common way for the agents to communicate is via a shared state channel, typically a list of messages. This assumes that there is always at least a single channel (key) in the state that is shared by the agents. When communicating via a shared message list there is an additional consideration: should the agents [share the full history](#share-full-history) of their thought process or only [the final result](#share-final-result)?
![](./img/multi_agent/response.png)
#### Share full history
Agents can **share the full history** of their thought process (i.e. "scratchpad") with all other agents. This "scratchpad" would typically look like a [list of messages](./low_level.md#why-use-messages). The benefit of sharing full thought process is that it might help other agents make better decisions and improve reasoning ability for the system as a whole. The downside is that as the number of agents and their complexity grows, the "scratchpad" will grow quickly and might require additional strategies for [memory management](./memory.md/#managing-long-conversation-history).
#### Share final result
Agents can have their own private "scratchpad" and only **share the final result** with the rest of the agents. This approach might work better for systems with many agents or agents that are more complex. In this case, you would need to define agents with [different state schemas](#different-state-schemas)
For agents called as tools, the supervisor determines the inputs based on the tool schema. Additionally, LangGraph allows [passing state](https://langchain-ai.github.io/langgraph/how-tos/pass-run-time-values-to-tools/#pass-graph-state-to-tools) to individual tools at runtime, so subordinate agents can access parent state, if needed.
+39 -21
View File
@@ -89,7 +89,7 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": 2,
"id": "09372b8b-edea-4b9d-9ec3-3d93ce1ba819",
"metadata": {},
"outputs": [],
@@ -132,7 +132,7 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": 3,
"id": "66f52a20",
"metadata": {},
"outputs": [
@@ -163,7 +163,7 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": 4,
"id": "38846b01",
"metadata": {},
"outputs": [
@@ -183,7 +183,7 @@
"{'aggregate': [\"I'm A\", \"I'm B\", \"I'm C\", \"I'm D\"]}"
]
},
"execution_count": 5,
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
@@ -220,7 +220,7 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": 5,
"id": "259a7704-5aa0-4e4c-aeef-cca04e8be0ff",
"metadata": {},
"outputs": [],
@@ -238,6 +238,15 @@
" aggregate: Annotated[list, operator.add]\n",
"\n",
"\n",
"class ReturnNodeValue:\n",
" def __init__(self, node_secret: str):\n",
" self._value = node_secret\n",
"\n",
" def __call__(self, state: State) -> Any:\n",
" print(f\"Adding {self._value} to {state['aggregate']}\")\n",
" return {\"aggregate\": [self._value]}\n",
"\n",
"\n",
"builder = StateGraph(State)\n",
"builder.add_node(\"a\", ReturnNodeValue(\"I'm A\"))\n",
"builder.add_edge(START, \"a\")\n",
@@ -255,7 +264,7 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": 6,
"id": "83320227-8ab3-44c0-b6cf-064a7a425b9f",
"metadata": {},
"outputs": [
@@ -278,7 +287,7 @@
},
{
"cell_type": "code",
"execution_count": 8,
"execution_count": 7,
"id": "3f971fa3-29e4-466f-a85e-2863bfecf7fe",
"metadata": {},
"outputs": [
@@ -299,7 +308,7 @@
"{'aggregate': [\"I'm A\", \"I'm B\", \"I'm C\", \"I'm B2\", \"I'm D\"]}"
]
},
"execution_count": 8,
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
@@ -322,7 +331,7 @@
},
{
"cell_type": "code",
"execution_count": 9,
"execution_count": 8,
"id": "95f5e026",
"metadata": {},
"outputs": [],
@@ -341,6 +350,15 @@
" which: str\n",
"\n",
"\n",
"class ReturnNodeValue:\n",
" def __init__(self, node_secret: str):\n",
" self._value = node_secret\n",
"\n",
" def __call__(self, state: State) -> Any:\n",
" print(f\"Adding {self._value} to {state['aggregate']}\")\n",
" return {\"aggregate\": [self._value]}\n",
"\n",
"\n",
"builder = StateGraph(State)\n",
"builder.add_node(\"a\", ReturnNodeValue(\"I'm A\"))\n",
"builder.add_edge(START, \"a\")\n",
@@ -372,7 +390,7 @@
},
{
"cell_type": "code",
"execution_count": 10,
"execution_count": 9,
"id": "1d0e6c56",
"metadata": {},
"outputs": [
@@ -395,7 +413,7 @@
},
{
"cell_type": "code",
"execution_count": 11,
"execution_count": 10,
"id": "7134f652",
"metadata": {},
"outputs": [
@@ -415,7 +433,7 @@
"{'aggregate': [\"I'm A\", \"I'm B\", \"I'm C\", \"I'm E\"], 'which': 'bc'}"
]
},
"execution_count": 11,
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
@@ -426,7 +444,7 @@
},
{
"cell_type": "code",
"execution_count": 12,
"execution_count": 11,
"id": "b130e694",
"metadata": {},
"outputs": [
@@ -446,7 +464,7 @@
"{'aggregate': [\"I'm A\", \"I'm C\", \"I'm D\", \"I'm E\"], 'which': 'cd'}"
]
},
"execution_count": 12,
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
@@ -471,7 +489,7 @@
},
{
"cell_type": "code",
"execution_count": 13,
"execution_count": 12,
"id": "836bc12d",
"metadata": {},
"outputs": [],
@@ -564,7 +582,7 @@
},
{
"cell_type": "code",
"execution_count": 14,
"execution_count": 13,
"id": "932c497e",
"metadata": {},
"outputs": [
@@ -587,7 +605,7 @@
},
{
"cell_type": "code",
"execution_count": 15,
"execution_count": 14,
"id": "933b3afd",
"metadata": {},
"outputs": [
@@ -608,7 +626,7 @@
" 'which': 'bc'}"
]
},
"execution_count": 15,
"execution_count": 14,
"metadata": {},
"output_type": "execute_result"
}
@@ -619,7 +637,7 @@
},
{
"cell_type": "code",
"execution_count": 16,
"execution_count": 15,
"id": "e30531bf",
"metadata": {},
"outputs": [
@@ -640,7 +658,7 @@
" 'which': 'cd'}"
]
},
"execution_count": 16,
"execution_count": 15,
"metadata": {},
"output_type": "execute_result"
}
@@ -666,7 +684,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.4"
"version": "3.12.3"
}
},
"nbformat": 4,
+10 -3
View File
@@ -21,6 +21,7 @@ These how-to guides show how to achieve that controllability.
LangGraph makes it easy to persist state across graph runs (thread-level persistence) and across threads (cross-thread persistence). These how-to guides show how to add persistence to your graph.
- [How to add thread-level persistence to your graph](persistence.ipynb)
- [How to add thread-level persistence to subgraphs](subgraph-persistence.ipynb)
- [How to add cross-thread persistence to your graph](cross-thread-persistence.ipynb)
- [How to use Postgres checkpointer for persistence](persistence_postgres.ipynb)
- [How to create a custom checkpointer using MongoDB](persistence_mongodb.ipynb)
@@ -73,8 +74,8 @@ These guides show how to use different streaming modes.
## Subgraphs
- [How to create subgraphs](subgraph.ipynb)
- [How to manage state in subgraphs](subgraphs-manage-state.ipynb)
- [How to add and use subgraphs](subgraph.ipynb)
- [How to view and update state in subgraphs](subgraphs-manage-state.ipynb)
- [How to transform inputs and outputs of a subgraph](subgraph-transform-state.ipynb)
## State Management
@@ -103,4 +104,10 @@ Please note that here will we use a **prebuilt agent**. One of the big benefits
- [How to add memory to a ReAct agent](create-react-agent-memory.ipynb)
- [How to add a custom system prompt to a ReAct agent](create-react-agent-system-prompt.ipynb)
- [How to add human-in-the-loop processes to a ReAct agent](create-react-agent-hitl.ipynb)
- [How to create prebuilt ReAct agent from scratch](react-agent-from-scratch.ipynb)
- [How to create prebuilt ReAct agent from scratch](react-agent-from-scratch.ipynb)
## Troubleshooting
### Errors
- [Error reference](../troubleshooting/errors/index.md)
@@ -0,0 +1,379 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "176e8dbb-1a0a-49ce-a10e-2417e8ea17a0",
"metadata": {},
"source": [
"# How to add thread-level persistence to subgraphs"
]
},
{
"cell_type": "markdown",
"id": "8c67581a-49fb-4597-a7fc-6774581c2160",
"metadata": {},
"source": [
"<div class=\"admonition tip\">\n",
" <p class=\"admonition-title\">Prerequisites</p>\n",
" <p>\n",
" This guide assumes familiarity with the following:\n",
" <ul>\n",
" <li> \n",
" <a href=\"https://langchain-ai.github.io/langgraph/concepts/low_level/#subgraphs\">\n",
" Subgraphs\n",
" </a>\n",
" </li>\n",
" <li>\n",
" <a href=\"https://langchain-ai.github.io/langgraph/concepts/persistence/\">\n",
" Persistence\n",
" </a>\n",
" </li>\n",
" </ul>\n",
" </p>\n",
"</div>\n",
"\n",
"This guide shows how you can add [thread-level](https://langchain-ai.github.io/langgraph/how-tos/persistence/) persistence to graphs that use [subgraphs](https://langchain-ai.github.io/langgraph/how-tos/subgraph/)."
]
},
{
"cell_type": "markdown",
"id": "8f83b855-ab23-4de7-9559-702cad9a29c6",
"metadata": {},
"source": [
"## Setup\n",
"\n",
"First, let's install the required packages"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "77d1eafa-3252-45f6-9af0-d94e1f9c5c9e",
"metadata": {},
"outputs": [],
"source": [
"%%capture --no-stderr\n",
"%pip install -U langgraph"
]
},
{
"cell_type": "markdown",
"id": "2e60c6cd-bf4e-46af-9761-b872d0fbe3b6",
"metadata": {},
"source": [
"<div class=\"admonition tip\">\n",
" <p class=\"admonition-title\">Set up <a href=\"https://smith.langchain.com\">LangSmith</a> for LangGraph development</p>\n",
" <p style=\"padding-top: 5px;\">\n",
" Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started <a href=\"https://docs.smith.langchain.com\">here</a>. \n",
" </p>\n",
"</div>"
]
},
{
"cell_type": "markdown",
"id": "871b9056-fec7-4683-8c22-f56c91f5b13b",
"metadata": {},
"source": [
"## Define the graph with persistence"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "9f1303ef-df37-48e0-8a59-8ff169c52c5b",
"metadata": {},
"source": [
"To add persistence to a graph with subgraphs, all you need to do is pass a [checkpointer](https://langchain-ai.github.io/langgraph/reference/checkpoints/#langgraph.checkpoint.base.BaseCheckpointSaver) when **compiling the parent graph**. LangGraph will automatically propagate the checkpointer to the child subgraphs."
]
},
{
"cell_type": "markdown",
"id": "c74cde2e-c127-4326-8d36-b6acef987f0a",
"metadata": {},
"source": [
"!!! note\n",
" You **shouldn't provide** a checkpointer when compiling a subgraph. Instead, you must define a **single** checkpointer that you pass to `parent_graph.compile()`, and LangGraph will automatically propagate the checkpointer to the child subgraphs. If you pass the checkpointer to the `subgraph.compile()`, it will simply be ignored. This also applies when you [add a node function that invokes the subgraph](../subgraph#add-a-node-function-that-invokes-the-subgraph)."
]
},
{
"cell_type": "markdown",
"id": "c3a1fe22-1ca9-45eb-a35b-71b9c905e8c5",
"metadata": {},
"source": [
"Let's define a simple graph with a single subgraph node to show how to do this."
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "0d76f0c0-bd77-4eca-9527-27bcdf85dd42",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"<langgraph.graph.state.StateGraph at 0x106d2fa10>"
]
},
"execution_count": 2,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from langgraph.graph import START, StateGraph\n",
"from langgraph.checkpoint.memory import MemorySaver\n",
"from typing import TypedDict\n",
"\n",
"\n",
"# subgraph\n",
"\n",
"\n",
"class SubgraphState(TypedDict):\n",
" foo: str # note that this key is shared with the parent graph state\n",
" bar: str\n",
"\n",
"\n",
"def subgraph_node_1(state: SubgraphState):\n",
" return {\"bar\": \"bar\"}\n",
"\n",
"\n",
"def subgraph_node_2(state: SubgraphState):\n",
" # note that this node is using a state key ('bar') that is only available in the subgraph\n",
" # and is sending update on the shared state key ('foo')\n",
" return {\"foo\": state[\"foo\"] + state[\"bar\"]}\n",
"\n",
"\n",
"subgraph_builder = StateGraph(SubgraphState)\n",
"subgraph_builder.add_node(subgraph_node_1)\n",
"subgraph_builder.add_node(subgraph_node_2)\n",
"subgraph_builder.add_edge(START, \"subgraph_node_1\")\n",
"subgraph_builder.add_edge(\"subgraph_node_1\", \"subgraph_node_2\")\n",
"subgraph = subgraph_builder.compile()\n",
"\n",
"\n",
"# parent graph\n",
"\n",
"\n",
"class State(TypedDict):\n",
" foo: str\n",
"\n",
"\n",
"def node_1(state: State):\n",
" return {\"foo\": \"hi! \" + state[\"foo\"]}\n",
"\n",
"\n",
"builder = StateGraph(State)\n",
"builder.add_node(\"node_1\", node_1)\n",
"# note that we're adding the compiled subgraph as a node to the parent graph\n",
"builder.add_node(\"node_2\", subgraph)\n",
"builder.add_edge(START, \"node_1\")\n",
"builder.add_edge(\"node_1\", \"node_2\")"
]
},
{
"cell_type": "markdown",
"id": "47084b1f-9fd5-40a9-9d75-89eb5f853d02",
"metadata": {},
"source": [
"We can now compile the graph with an in-memory checkpointer (`MemorySaver`)."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "7657d285-c896-40c9-a569-b4a3b9c230c7",
"metadata": {},
"outputs": [],
"source": [
"checkpointer = MemorySaver()\n",
"# You must only pass checkpointer when compiling the parent graph.\n",
"# LangGraph will automatically propagate the checkpointer to the child subgraphs.\n",
"graph = builder.compile(checkpointer=checkpointer)"
]
},
{
"cell_type": "markdown",
"id": "0d193e3c-4ec3-4034-beed-8e5550c6542c",
"metadata": {},
"source": [
"## Verify persistence works"
]
},
{
"cell_type": "markdown",
"id": "eb69a5f0-b92e-4d4e-9aa9-c4c4ec7de91a",
"metadata": {},
"source": [
"Let's now run the graph and inspect the persisted state for both the parent graph and the subgraph to verify that persistence works. We should expect to see the final execution results for both the parent and subgraph in `state.values`."
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "13da686e-6ed6-4b83-93e8-1631fcc8c2a9",
"metadata": {},
"outputs": [],
"source": [
"config = {\"configurable\": {\"thread_id\": \"1\"}}"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "8721f045-2e82-4bf0-9d85-5ba6ecf899d6",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{'node_1': {'foo': 'hi! foo'}}\n",
"{'subgraph_node_1': {'bar': 'bar'}}\n",
"{'subgraph_node_2': {'foo': 'hi! foobar'}}\n",
"{'node_2': {'foo': 'hi! foobar'}}\n"
]
}
],
"source": [
"for _, chunk in graph.stream({\"foo\": \"foo\"}, config, subgraphs=True):\n",
" print(chunk)"
]
},
{
"cell_type": "markdown",
"id": "ec6b5ce4-becc-4910-8a6d-d6b60d9d6f60",
"metadata": {},
"source": [
"We can now view the parent graph state by calling `graph.get_state()` with the same config that we used to invoke the graph."
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "3e817283-142d-4fda-8cb1-8de34717f833",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'foo': 'hi! foobar'}"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"graph.get_state(config).values"
]
},
{
"cell_type": "markdown",
"id": "fbc4f30b-941e-4140-8bfa-3b8cc670489c",
"metadata": {},
"source": [
"To view the subgraph state, we need to do two things:\n",
"\n",
"1. Find the most recent config value for the subgraph\n",
"2. Use `graph.get_state()` to retrieve that value for the most recent subgraph config.\n",
"\n",
"To find the correct config, we can examine the state history from the parent graph and find the state snapshot before we return results from `node_2` (the node with subgraph):"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "e896628f-36b2-45eb-b7c5-c64c1098f328",
"metadata": {},
"outputs": [],
"source": [
"state_with_subgraph = [\n",
" s for s in graph.get_state_history(config) if s.next == (\"node_2\",)\n",
"][0]"
]
},
{
"cell_type": "markdown",
"id": "7af49977-42b1-40a1-88f1-f07437f8b7f9",
"metadata": {},
"source": [
"The state snapshot will include the list of `tasks` to be executed next. When using subgraphs, the `tasks` will contain the config that we can use to retrieve the subgraph state:"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "21e96df3-946d-40f8-8d6d-055ae4177452",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'configurable': {'thread_id': '1',\n",
" 'checkpoint_ns': 'node_2:6ef111a6-f290-7376-0dfc-a4152307bc5b'}}"
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"subgraph_config = state_with_subgraph.tasks[0].state\n",
"subgraph_config"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "1d2401b3-d52b-4895-a5d1-dccf015ba216",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'foo': 'hi! foobar', 'bar': 'bar'}"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"graph.get_state(subgraph_config).values"
]
},
{
"cell_type": "markdown",
"id": "40aded92-99dd-427b-932d-aa78f474c271",
"metadata": {},
"source": [
"If you want to learn more about how to modify the subgraph state for human-in-the-loop workflows, check out this [how-to guide](https://langchain-ai.github.io/langgraph/how-tos/subgraphs-manage-state/)."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.3"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
File diff suppressed because one or more lines are too long
+47 -8
View File
@@ -5,15 +5,48 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"# How to manage state in subgraphs\n",
"# How to view and update state in subgraphs\n",
"\n",
"For more complex systems, sub-graphs are a useful design principle. Sub-graphs allow you to create and manage different states in different parts of your graph. This allows you build things like [multi-agent teams](https://langchain-ai.github.io/langgraph/tutorials/multi_agent/hierarchical_agent_teams/), where each team can track its own separate state.\n",
"<div class=\"admonition tip\">\n",
" <p class=\"admonition-title\">Prerequisites</p>\n",
" <p>\n",
" This guide assumes familiarity with the following:\n",
" <ul>\n",
" <li> \n",
" <a href=\"https://langchain-ai.github.io/langgraph/concepts/low_level/#subgraphs\">\n",
" Subgraphs\n",
" </a>\n",
" </li>\n",
" <li> \n",
" <a href=\"https://langchain-ai.github.io/langgraph/concepts/human_in_the_loop/\">\n",
" Human-in-the-loop\n",
" </a>\n",
" </li>\n",
" <li> \n",
" <a href=\"https://langchain-ai.github.io/langgraph/concepts/low_level/#state\">\n",
" State\n",
" </a>\n",
" </li>\n",
" </ul>\n",
" </p>\n",
"</div> \n",
"\n",
"In this how-to guide we will cover how to manage the persisted state in subgraphs. This will enable a lot of the human-in-the-loop interaction patterns.\n",
"Once you add [persistence](../subgraph-persistence), you can easily view and update the state of the subgraph at any point in time. This enables a lot of the human-in-the-loop interaction patterns:\n",
"\n",
"* You can surface a state during an interrupt to a user to let them accept an action.\n",
"* You can rewind the subgraph to reproduce or avoid issues.\n",
"* You can modify the state to let the user better control its actions.\n",
"\n",
"This guide shows how you can do this."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setup\n",
"\n",
"First we need to install the packages required"
"First, let's install the required packages"
]
},
{
@@ -68,7 +101,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"## Define SubGraph\n",
"## Define subgraph\n",
"\n",
"First, let's set up our subgraph. For this, we will create a simple graph that can get the weather for a specific city. We will compile this graph with a [breakpoint](https://langchain-ai.github.io/langgraph/how-tos/human_in_the_loop/breakpoints/) before the `weather_node`:"
]
@@ -121,7 +154,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"## Define Parent Graph\n",
"## Define parent graph\n",
"\n",
"We can now setup the overall graph. This graph will first route to the subgraph if it needs to get the weather, otherwise it will route to a normal LLM."
]
@@ -444,7 +477,7 @@
" if h.next == (\"model_node\",)\n",
")\n",
"\n",
"# This pattern can be extended no matter how many levels deep - image model node was another subgraph in this case\n",
"# This pattern can be extended no matter how many levels deep\n",
"# subsubgraph_stat_history = next(h for h in graph.get_state_history(subgraph_state_before_model_node.tasks[0].state) if h.next == ('my_subsubgraph_node',))"
]
},
@@ -660,7 +693,9 @@
" print(update)\n",
"# Graph execution should stop before the weather node\n",
"print(\"interrupted!\")\n",
"\n",
"state = graph.get_state(config, subgraphs=True)\n",
"\n",
"# We update the state by passing in the message we want returned from the weather node, and make sure to use as_node\n",
"graph.update_state(\n",
" state.tasks[0].state.config,\n",
@@ -669,6 +704,7 @@
")\n",
"for update in graph.stream(None, config=config, stream_mode=\"updates\", subgraphs=True):\n",
" print(update)\n",
"\n",
"print(graph.get_state(config).values[\"messages\"])"
]
},
@@ -708,6 +744,7 @@
" print(update)\n",
"# Graph execution should stop before the weather node\n",
"print(\"interrupted!\")\n",
"\n",
"# We update the state by passing in the message we want returned from the weather graph, making sure to use as_node\n",
"# Note that we don't need to pass in the subgraph config, since we aren't updating the state inside the subgraph\n",
"graph.update_state(\n",
@@ -717,6 +754,7 @@
")\n",
"for update in graph.stream(None, config=config, stream_mode=\"updates\"):\n",
" print(update)\n",
"\n",
"print(graph.get_state(config).values[\"messages\"])"
]
},
@@ -947,6 +985,7 @@
" None, config=config, stream_mode=\"updates\", subgraphs=True\n",
"):\n",
" print(update)\n",
"\n",
"print(grandparent_graph.get_state(config).values[\"messages\"])"
]
},
@@ -1002,7 +1041,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.9"
"version": "3.12.3"
}
},
"nbformat": 4,
@@ -0,0 +1,29 @@
# GRAPH_RECURSION_LIMIT
Your LangGraph [`StateGraph`](https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.state.StateGraph) reached the maximum number of steps before hitting a stop condition.
This is often due to an infinite loop caused by code like the example below:
```python
class State(TypedDict):
some_key: str
builder = StateGraph(State)
builder.add_node("a", ...)
builder.add_node("b", ...)
builder.add_edge("a", "b")
builder.add_edge("b", "a")
...
graph = builder.compile()
```
However, complex graphs may hit the default limit naturally.
## Troubleshooting
- If you are not expecting your graph to go through many iterations, you likely have a cycle. Check your logic for infinite loops.
- If you have a complex graph, you can pass in a higher `recursion_limit` value into your `config` object when invoking your graph like this:
```python
graph.invoke({...}, {"recursion_limit": 100})
```
@@ -0,0 +1,49 @@
# INVALID_CONCURRENT_GRAPH_UPDATE
A LangGraph [`StateGraph`](https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.state.StateGraph) received concurrent updates to its state from multiple nodes to a state property that doesn't
support it.
One way this can occur is if you are using a [fanout](https://langchain-ai.github.io/langgraph/how-tos/map-reduce/)
or other parallel execution in your graph and you have defined a graph like this:
```python
class State(TypedDict):
some_key: str
def node(state: State):
return {"some_key": "some_string_value"}
def other_node(state: State):
return {"some_key": "some_string_value"}
builder = StateGraph(State)
builder.add_node(node)
builder.add_node(other_node)
builder.add_edge(START, "node")
builder.add_edge(START, "other_node")
graph = builder.compile()
```
If a node in the above graph returns `{ "some_key": "some_string_value" }`, this will overwrite the state value for `"some_key"` with `"some_string_value"`.
However, if multiple nodes in e.g. a fanout within a single step return values for `"some_key"`, the graph will throw this error because
there is uncertainty around how to update the internal state.
To get around this, you can define a reducer that combines multiple values:
```python
import operator
from typing import Annotated
class State(TypedDict):
# The operator.add reducer fn makes this append-only
some_key: Annotated[list, operator.add]
```
This will allow you to define logic that handles the same key returned from multiple nodes executed in parallel.
## Troubleshooting
The following may help resolve this error:
- If your graph executes nodes in parallel, make sure you have defined relevant state keys with a reducer.
@@ -0,0 +1,38 @@
# INVALID_GRAPH_NODE_RETURN_VALUE
A LangGraph [`StateGraph`](https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.state.StateGraph)
received a non-dict return type from a node. Here's an example:
```python
class State(TypedDict):
some_key: str
def bad_node(state: State):
# Should return an dict with a value for "some_key", not a list
return ["whoops"]
builder = StateGraph(State)
builder.add_node(bad_node)
...
graph = builder.compile()
```
Invoking the above graph will result in an error like this:
```python
graph.invoke({ "some_key": "someval" });
```
```
InvalidUpdateError: Expected dict, got ['whoops']
For troubleshooting, visit: https://python.langchain.com/docs/troubleshooting/errors/INVALID_GRAPH_NODE_RETURN_VALUE
```
Nodes in your graph must return an dict containing one or more keys defined in your state.
## Troubleshooting
The following may help resolve this error:
- If you have complex logic in your node, make sure all code paths return an appropriate dict for your defined state.
@@ -0,0 +1,12 @@
# MULTIPLE_SUBGRAPHS
You are calling the same subgraph multiple times within a single LangGraph node with checkpointing enabled for each subgraph.
This is currently not allowed due to internal restrictions on how checkpoint namespacing for subgraphs works.
## Troubleshooting
The following may help resolve this error:
- If you don't need to interrupt/resume from a subgraph, pass `checkpointer=False` when compiling it like this: `.compile(checkpointer=False)`
- Don't imperatively call graphs multiple times in the same node, and instead use the [`Send`](https://langchain-ai.github.io/langgraph/concepts/low_level/#send) API.
@@ -0,0 +1,9 @@
# Error reference
This page contains guides around resolving common errors you may find while building with LangChain.
Errors referenced below will have an `lc_error_code` property corresponding to one of the below codes when they are thrown in code.
- [GRAPH_RECURSION_LIMIT](./GRAPH_RECURSION_LIMIT.md)
- [INVALID_CONCURRENT_GRAPH_UPDATE](./INVALID_CONCURRENT_GRAPH_UPDATE.md)
- [INVALID_GRAPH_NODE_RETURN_VALUE](./INVALID_GRAPH_NODE_RETURN_VALUE.md)
- [MULTIPLE_SUBGRAPHS](./MULTIPLE_SUBGRAPHS.md)
+2 -2
View File
@@ -25,8 +25,8 @@ Learn from example implementations of graphs designed for specific scenarios and
#### Multi-Agent Systems
- [Collaboration](multi_agent/multi-agent-collaboration.ipynb): Enable two agents to collaborate on a task
- [Supervision](multi_agent/agent_supervisor.ipynb): Use an LLM to orchestrate and delegate to individual agents
- [Network](multi_agent/multi-agent-collaboration.ipynb): Enable two or more agents to collaborate on a task
- [Supervisor](multi_agent/agent_supervisor.ipynb): Use an LLM to orchestrate and delegate to individual agents
- [Hierarchical Teams](multi_agent/hierarchical_agent_teams.ipynb): Orchestrate nested teams of agents to solve problems
#### RAG
@@ -10,11 +10,11 @@
"id": "a3e3ebc4-57af-4fe4-bdd3-36aff67bf276",
"metadata": {},
"source": [
"# Agent Supervisor\n",
"# Multi-agent supervisor\n",
"\n",
"The [previous example](../multi-agent-collaboration) routed messages automatically based on the output of the initial researcher agent.\n",
"\n",
"We can also choose to use an LLM to orchestrate the different agents.\n",
"We can also choose to use an [LLM to orchestrate](https://langchain-ai.github.io/langgraph/concepts/multi_agent/#supervisor) the different agents.\n",
"\n",
"Below, we will create an agent group, with an agent supervisor to help delegate tasks.\n",
"\n",
@@ -376,7 +376,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.9"
"version": "3.12.3"
}
},
"nbformat": 4,
@@ -12,7 +12,7 @@
"source": [
"# Hierarchical Agent Teams\n",
"\n",
"In our previous example ([Agent Supervisor](../agent_supervisor)), we introduced the concept of a single supervisor node to route work between different worker nodes.\n",
"In our previous example ([Agent Supervisor](../agent_supervisor)), we introduced the concept of a single [supervisor node](https://langchain-ai.github.io/langgraph/concepts/multi_agent/#supervisor) to route work between different worker nodes.\n",
"\n",
"But what if the job for a single worker becomes too complex? What if the number of workers becomes too large?\n",
"\n",
@@ -1117,7 +1117,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.9"
"version": "3.12.3"
}
},
"nbformat": 4,
@@ -10,11 +10,11 @@
"id": "39fd1948-b5c3-48c4-b10e-2ae7e8c83334",
"metadata": {},
"source": [
"# Basic Multi-agent Collaboration\n",
"# Multi-agent network\n",
"\n",
"A single agent can usually operate effectively using a handful of tools within a single domain, but even using powerful models like `gpt-4`, it can be less effective at using many tools. \n",
"\n",
"One way to approach complicated tasks is through a \"divide-and-conquer\" approach: create an specialized agent for each task or domain and route tasks to the correct \"expert\".\n",
"One way to approach complicated tasks is through a \"divide-and-conquer\" approach: create an specialized agent for each task or domain and route tasks to the correct \"expert\". This is an example of a [multi-agent network](https://langchain-ai.github.io/langgraph/concepts/multi_agent/#network) architecture.\n",
"\n",
"This notebook (inspired by the paper [AutoGen: Enabling Next-Gen LLM Applications via Multi-Agent Conversation](https://arxiv.org/abs/2308.08155), by Wu, et. al.) shows one way to do this using LangGraph.\n",
"\n",
@@ -535,7 +535,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.9"
"version": "3.12.3"
}
},
"nbformat": 4,
+16 -9
View File
@@ -97,8 +97,8 @@ nav:
- SQL Agent: tutorials/sql-agent.ipynb
- Agent Architectures:
- Multi-Agent Systems:
- Collaboration: tutorials/multi_agent/multi-agent-collaboration.ipynb
- Supervision: tutorials/multi_agent/agent_supervisor.ipynb
- Network: tutorials/multi_agent/multi-agent-collaboration.ipynb
- Supervisor: tutorials/multi_agent/agent_supervisor.ipynb
- Hierarchical Teams: tutorials/multi_agent/hierarchical_agent_teams.ipynb
- Planning Agents:
- Plan-and-Execute: tutorials/plan-and-execute/plan-and-execute.ipynb
@@ -128,6 +128,7 @@ nav:
- Control graph recursion limit: how-tos/recursion-limit.ipynb
- Persistence:
- Add thread-level persistence: how-tos/persistence.ipynb
- Add thread-level persistence to subgraphs: how-tos/subgraph-persistence.ipynb
- Add cross-thread persistence: how-tos/cross-thread-persistence.ipynb
- Use Postgres checkpointer for persistence: how-tos/persistence_postgres.ipynb
- Create custom checkpointer using MongoDB: how-tos/persistence_mongodb.ipynb
@@ -162,8 +163,8 @@ nav:
- Pass config to tools: how-tos/pass-config-to-tools.ipynb
- Handle many tools: how-tos/many-tools.ipynb
- Subgraphs:
- Create subgraphs: how-tos/subgraph.ipynb
- Manage state in subgraphs: how-tos/subgraphs-manage-state.ipynb
- Add and use subgraphs: how-tos/subgraph.ipynb
- View and update state in subgraphs: how-tos/subgraphs-manage-state.ipynb
- Transform inputs and outputs of a subgraph: how-tos/subgraph-transform-state.ipynb
- State Management:
- Use Pydantic model as state: how-tos/state-model.ipynb
@@ -177,6 +178,12 @@ nav:
- Return structured output from a ReAct agent: how-tos/react-agent-structured-output.ipynb
- Pass custom LangSmith run ID for graph runs: how-tos/run-id-langsmith.ipynb
- Return state before hitting recursion limit: how-tos/return-when-recursion-limit-hits.ipynb
- Error reference:
- "troubleshooting/errors/index.md"
- GRAPH_RECURSION_LIMIT: "troubleshooting/errors/GRAPH_RECURSION_LIMIT.md"
- INVALID_CONCURRENT_GRAPH_UPDATE: "troubleshooting/errors/INVALID_CONCURRENT_GRAPH_UPDATE.md"
- INVALID_GRAPH_NODE_RETURN_VALUE: "troubleshooting/errors/INVALID_GRAPH_NODE_RETURN_VALUE.md"
- MULTIPLE_SUBGRAPHS: "troubleshooting/errors/MULTIPLE_SUBGRAPHS.md"
- Prebuilt ReAct Agent:
- Create a ReAct agent: how-tos/create-react-agent.ipynb
- Add memory to a ReAct agent: how-tos/create-react-agent-memory.ipynb
@@ -206,7 +213,7 @@ nav:
- "cloud/index.md"
- Tutorials:
- Quick Start: "cloud/quick_start.md"
- How-to Guides:
- How-to Guides:
- "cloud/how-tos/index.md"
- Setup:
- Setup App: "cloud/deployment/setup.md"
@@ -229,7 +236,7 @@ nav:
- Rollback: "cloud/how-tos/rollback_concurrent.md"
- Reject: "cloud/how-tos/reject_concurrent.md"
- Enqueue: "cloud/how-tos/enqueue_concurrent.md"
- Human-in-the-Loop:
- Human-in-the-Loop:
- Add Breakpoint: "cloud/how-tos/human_in_the_loop_breakpoint.md"
- Wait for User Input: "cloud/how-tos/human_in_the_loop_user_input.md"
- Edit Graph State: "cloud/how-tos/human_in_the_loop_edit_state.md"
@@ -249,8 +256,8 @@ nav:
- Configure Agents: "cloud/how-tos/configuration_cloud.md"
- Versioning Assistants: "cloud/how-tos/assistant_versioning.md"
- Convert LangGraph calls to LangGraph Cloud calls: "cloud/how-tos/langgraph_to_langgraph_cloud.ipynb"
- Integrate Webhooks: 'cloud/how-tos/webhooks.md'
- Copy Threads: 'cloud/how-tos/copy_threads.md'
- Integrate Webhooks: "cloud/how-tos/webhooks.md"
- Copy Threads: "cloud/how-tos/copy_threads.md"
- Check Status of Threads: "cloud/how-tos/check_thread_status.md"
- Conceptual Guides:
- API Concepts: "cloud/concepts/api.md"
@@ -351,4 +358,4 @@ validation:
# it's only an issue for tutorials/storm/storm.ipynb
# because it creates anchors in the generated report
# and those anchors are not available in the actual doc
anchors: info
anchors: info
+35
View File
@@ -0,0 +1,35 @@
.PHONY: test test_watch lint format
######################
# TESTING AND COVERAGE
######################
test:
poetry run pytest tests
test_watch:
poetry run ptw .
######################
# LINTING AND FORMATTING
######################
# Define a variable for Python and notebook files.
PYTHON_FILES=.
MYPY_CACHE=.mypy_cache
lint format: PYTHON_FILES=.
lint_diff format_diff: PYTHON_FILES=$(shell git diff --name-only --relative --diff-filter=d main . | grep -E '\.py$$|\.ipynb$$')
lint_package: PYTHON_FILES=langgraph
lint_tests: PYTHON_FILES=tests
lint_tests: MYPY_CACHE=.mypy_cache_test
lint lint_diff lint_package lint_tests:
poetry run ruff check .
[ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff
[ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES)
[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE)
[ "$(PYTHON_FILES)" = "" ] || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
format format_diff:
poetry run ruff format $(PYTHON_FILES)
poetry run ruff check --select I --fix $(PYTHON_FILES)
+95
View File
@@ -0,0 +1,95 @@
# LangGraph Checkpoint DuckDB
Implementation of LangGraph CheckpointSaver that uses DuckDB.
## Usage
> [!IMPORTANT]
> When using DuckDB checkpointers for the first time, make sure to call `.setup()` method on them to create required tables. See example below.
```python
from langgraph.checkpoint.duckdb import DuckDBSaver
write_config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
read_config = {"configurable": {"thread_id": "1"}}
with DuckDBSaver.from_conn_string(":memory:") as checkpointer:
# call .setup() the first time you're using the checkpointer
checkpointer.setup()
checkpoint = {
"v": 1,
"ts": "2024-07-31T20:14:19.804150+00:00",
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
"channel_values": {
"my_key": "meow",
"node": "node"
},
"channel_versions": {
"__start__": 2,
"my_key": 3,
"start:node": 3,
"node": 3
},
"versions_seen": {
"__input__": {},
"__start__": {
"__start__": 1
},
"node": {
"start:node": 2
}
},
"pending_sends": [],
}
# store checkpoint
checkpointer.put(write_config, checkpoint, {}, {})
# load checkpoint
checkpointer.get(read_config)
# list checkpoints
list(checkpointer.list(read_config))
```
### Async
```python
from langgraph.checkpoint.duckdb.aio import AsyncDuckDBSaver
async with AsyncDuckDBSaver.from_conn_string(":memory:") as checkpointer:
checkpoint = {
"v": 1,
"ts": "2024-07-31T20:14:19.804150+00:00",
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
"channel_values": {
"my_key": "meow",
"node": "node"
},
"channel_versions": {
"__start__": 2,
"my_key": 3,
"start:node": 3,
"node": 3
},
"versions_seen": {
"__input__": {},
"__start__": {
"__start__": 1
},
"node": {
"start:node": 2
}
},
"pending_sends": [],
}
# store checkpoint
await checkpointer.aput(write_config, checkpoint, {}, {})
# load checkpoint
await checkpointer.aget(read_config)
# list checkpoints
[c async for c in checkpointer.alist(read_config)]
```
@@ -0,0 +1,356 @@
import threading
from contextlib import contextmanager
from typing import Any, Iterator, Optional, Sequence
from langchain_core.runnables import RunnableConfig
import duckdb
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
CheckpointTuple,
get_checkpoint_id,
)
from langgraph.checkpoint.duckdb.base import BaseDuckDBSaver
from langgraph.checkpoint.serde.base import SerializerProtocol
class DuckDBSaver(BaseDuckDBSaver):
lock: threading.Lock
def __init__(
self,
conn: duckdb.DuckDBPyConnection,
serde: Optional[SerializerProtocol] = None,
) -> None:
super().__init__(serde=serde)
self.conn = conn
self.lock = threading.Lock()
@classmethod
@contextmanager
def from_conn_string(cls, conn_string: str) -> Iterator["DuckDBSaver"]:
"""Create a new DuckDBSaver instance from a connection string.
Args:
conn_string (str): The DuckDB connection info string.
Returns:
DuckDBSaver: A new DuckDBSaver instance.
"""
with duckdb.connect(conn_string) as conn:
yield DuckDBSaver(conn)
def setup(self) -> None:
"""Set up the checkpoint database asynchronously.
This method creates the necessary tables in the DuckDB database if they don't
already exist and runs database migrations. It MUST be called directly by the user
the first time checkpointer is used.
"""
with self.lock, self.conn.cursor() as cur:
try:
row = cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
).fetchone()
if row is None:
version = -1
else:
version = row[0]
except duckdb.CatalogException:
version = -1
for v, migration in zip(
range(version + 1, len(self.MIGRATIONS)),
self.MIGRATIONS[version + 1 :],
):
cur.execute(migration)
cur.execute("INSERT INTO checkpoint_migrations (v) VALUES (?)", [v])
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 DuckDB database based
on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
Args:
config (RunnableConfig): The config to use for listing the checkpoints.
filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. Defaults to None.
before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
Yields:
Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
Examples:
>>> from langgraph.checkpoint.duckdb import DuckDBSaver
>>> with DuckDBSaver.from_conn_string(":memory:") as memory:
... # Run a graph, then list the checkpoints
>>> config = {"configurable": {"thread_id": "1"}}
>>> checkpoints = list(memory.list(config, limit=2))
>>> print(checkpoints)
[CheckpointTuple(...), CheckpointTuple(...)]
>>> config = {"configurable": {"thread_id": "1"}}
>>> before = {"configurable": {"checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875"}}
>>> with DuckDBSaver.from_conn_string(":memory:") as memory:
... # Run a graph, then list the checkpoints
>>> checkpoints = list(memory.list(config, before=before))
>>> print(checkpoints)
[CheckpointTuple(...), ...]
"""
where, args = self._search_where(config, filter, before)
query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
if limit:
query += f" LIMIT {limit}"
# if we change this to use .stream() we need to make sure to close the cursor
with self._cursor() as cur:
cur.execute(query, args)
for value in cur.fetchall():
(
thread_id,
checkpoint,
checkpoint_ns,
checkpoint_id,
parent_checkpoint_id,
metadata,
channel_values,
pending_writes,
pending_sends,
) = value
yield CheckpointTuple(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
self._load_checkpoint(
checkpoint,
channel_values,
pending_sends,
),
self._load_metadata(metadata),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_checkpoint_id
else None
),
self._load_writes(pending_writes),
)
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Get a checkpoint tuple from the database.
This method retrieves a checkpoint tuple from the DuckDB database based on the
provided config. If the config contains a "checkpoint_id" key, the checkpoint with
the matching thread ID and timestamp 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.
Examples:
Basic:
>>> config = {"configurable": {"thread_id": "1"}}
>>> checkpoint_tuple = memory.get_tuple(config)
>>> print(checkpoint_tuple)
CheckpointTuple(...)
With timestamp:
>>> config = {
... "configurable": {
... "thread_id": "1",
... "checkpoint_ns": "",
... "checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875",
... }
... }
>>> checkpoint_tuple = memory.get_tuple(config)
>>> print(checkpoint_tuple)
CheckpointTuple(...)
""" # noqa
thread_id = config["configurable"]["thread_id"]
checkpoint_id = get_checkpoint_id(config)
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
if checkpoint_id:
args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
where = "WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?"
else:
args = (thread_id, checkpoint_ns)
where = "WHERE thread_id = ? AND checkpoint_ns = ? ORDER BY checkpoint_id DESC LIMIT 1"
with self._cursor() as cur:
cur.execute(
self.SELECT_SQL + where,
args,
)
value = cur.fetchone()
if value:
(
thread_id,
checkpoint,
checkpoint_ns,
checkpoint_id,
parent_checkpoint_id,
metadata,
channel_values,
pending_writes,
pending_sends,
) = value
return CheckpointTuple(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
self._load_checkpoint(
checkpoint,
channel_values,
pending_sends,
),
self._load_metadata(metadata),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_checkpoint_id
else None
),
self._load_writes(pending_writes),
)
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 DuckDB 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.
Examples:
>>> from langgraph.checkpoint.duckdb import DuckDBSaver
>>> with DuckDBSaver.from_conn_string(":memory:") as memory:
>>> config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
>>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}}
>>> saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {})
>>> print(saved_config)
{'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}}
"""
configurable = config["configurable"].copy()
thread_id = configurable.pop("thread_id")
checkpoint_ns = configurable.pop("checkpoint_ns")
checkpoint_id = configurable.pop(
"checkpoint_id", configurable.pop("thread_ts", None)
)
copy = checkpoint.copy()
next_config = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint["id"],
}
}
checkpoint_blobs = self._dump_blobs(
thread_id,
checkpoint_ns,
copy.pop("channel_values"), # type: ignore[misc]
new_versions,
)
with self._cursor() as cur:
if checkpoint_blobs:
cur.executemany(self.UPSERT_CHECKPOINT_BLOBS_SQL, checkpoint_blobs)
cur.execute(
self.UPSERT_CHECKPOINTS_SQL,
(
thread_id,
checkpoint_ns,
checkpoint["id"],
checkpoint_id,
self._dump_checkpoint(copy),
self._dump_metadata(metadata),
),
)
return next_config
def put_writes(
self,
config: RunnableConfig,
writes: Sequence[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 DuckDB database.
Args:
config (RunnableConfig): Configuration of the related checkpoint.
writes (List[Tuple[str, Any]]): List of writes to store.
task_id (str): Identifier for the task creating the writes.
"""
query = (
self.UPSERT_CHECKPOINT_WRITES_SQL
if all(w[0] in WRITES_IDX_MAP for w in writes)
else self.INSERT_CHECKPOINT_WRITES_SQL
)
with self._cursor() as cur:
cur.executemany(
query,
self._dump_writes(
config["configurable"]["thread_id"],
config["configurable"]["checkpoint_ns"],
config["configurable"]["checkpoint_id"],
task_id,
writes,
),
)
@contextmanager
def _cursor(self) -> Iterator[duckdb.DuckDBPyConnection]:
with self.lock, self.conn.cursor() as cur:
yield cur
__all__ = ["DuckDBSaver", "Conn"]
@@ -0,0 +1,431 @@
import asyncio
from contextlib import asynccontextmanager
from typing import Any, AsyncIterator, Iterator, Optional, Sequence
from langchain_core.runnables import RunnableConfig
import duckdb
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
CheckpointTuple,
get_checkpoint_id,
)
from langgraph.checkpoint.duckdb.base import BaseDuckDBSaver
from langgraph.checkpoint.serde.base import SerializerProtocol
class AsyncDuckDBSaver(BaseDuckDBSaver):
lock: asyncio.Lock
def __init__(
self,
conn: duckdb.DuckDBPyConnection,
serde: Optional[SerializerProtocol] = None,
) -> None:
super().__init__(serde=serde)
self.conn = conn
self.lock = asyncio.Lock()
self.loop = asyncio.get_running_loop()
@classmethod
@asynccontextmanager
async def from_conn_string(
cls,
conn_string: str,
) -> AsyncIterator["AsyncDuckDBSaver"]:
"""Create a new AsyncDuckDBSaver instance from a connection string.
Args:
conn_string (str): The DuckDB connection info string.
Returns:
AsyncDuckDBSaver: A new AsyncDuckDBSaver instance.
"""
with duckdb.connect(conn_string) as conn:
yield AsyncDuckDBSaver(conn)
async def setup(self) -> None:
"""Set up the checkpoint database asynchronously.
This method creates the necessary tables in the DuckDB database if they don't
already exist and runs database migrations. It MUST be called directly by the user
the first time checkpointer is used.
"""
async with self.lock:
with self.conn.cursor() as cur:
try:
await asyncio.to_thread(
cur.execute,
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1",
)
row = await asyncio.to_thread(cur.fetchone)
if row is None:
version = -1
else:
version = row[0]
except duckdb.CatalogException:
version = -1
for v, migration in zip(
range(version + 1, len(self.MIGRATIONS)),
self.MIGRATIONS[version + 1 :],
):
await asyncio.to_thread(cur.execute, migration)
await asyncio.to_thread(
cur.execute,
"INSERT INTO checkpoint_migrations (v) VALUES (?)",
[v],
)
async def alist(
self,
config: Optional[RunnableConfig],
*,
filter: Optional[dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
"""List checkpoints from the database asynchronously.
This method retrieves a list of checkpoint tuples from the DuckDB 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:
AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.
"""
where, args = self._search_where(config, filter, before)
query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
if limit:
query += f" LIMIT {limit}"
# if we change this to use .stream() we need to make sure to close the cursor
async with self._cursor() as cur:
await asyncio.to_thread(cur.execute, query, args)
results = await asyncio.to_thread(cur.fetchall)
for value in results:
(
thread_id,
checkpoint,
checkpoint_ns,
checkpoint_id,
parent_checkpoint_id,
metadata,
channel_values,
pending_writes,
pending_sends,
) = value
yield CheckpointTuple(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
await asyncio.to_thread(
self._load_checkpoint,
checkpoint,
channel_values,
pending_sends,
),
self._load_metadata(metadata),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_checkpoint_id
else None
),
await asyncio.to_thread(self._load_writes, pending_writes),
)
async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Get a checkpoint tuple from the database asynchronously.
This method retrieves a checkpoint tuple from the DuckDBdatabase 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.
"""
thread_id = config["configurable"]["thread_id"]
checkpoint_id = get_checkpoint_id(config)
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
if checkpoint_id:
args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
where = "WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?"
else:
args = (thread_id, checkpoint_ns)
where = "WHERE thread_id = ? AND checkpoint_ns = ? ORDER BY checkpoint_id DESC LIMIT 1"
async with self._cursor() as cur:
await asyncio.to_thread(
cur.execute,
self.SELECT_SQL + where,
args,
)
value = await asyncio.to_thread(cur.fetchone)
if value:
(
thread_id,
checkpoint,
checkpoint_ns,
checkpoint_id,
parent_checkpoint_id,
metadata,
channel_values,
pending_writes,
pending_sends,
) = value
return CheckpointTuple(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
await asyncio.to_thread(
self._load_checkpoint,
checkpoint,
channel_values,
pending_sends,
),
self._load_metadata(metadata),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_checkpoint_id
else None
),
await asyncio.to_thread(self._load_writes, pending_writes),
)
async def aput(
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
new_versions: ChannelVersions,
) -> RunnableConfig:
"""Save a checkpoint to the database asynchronously.
This method saves a checkpoint to the DuckDB 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.
"""
configurable = config["configurable"].copy()
thread_id = configurable.pop("thread_id")
checkpoint_ns = configurable.pop("checkpoint_ns")
checkpoint_id = configurable.pop(
"checkpoint_id", configurable.pop("thread_ts", None)
)
copy = checkpoint.copy()
next_config = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint["id"],
}
}
checkpoint_blobs = await asyncio.to_thread(
self._dump_blobs,
thread_id,
checkpoint_ns,
copy.pop("channel_values"), # type: ignore[misc]
new_versions,
)
async with self._cursor() as cur:
if checkpoint_blobs:
await asyncio.to_thread(
cur.executemany, self.UPSERT_CHECKPOINT_BLOBS_SQL, checkpoint_blobs
)
await asyncio.to_thread(
cur.execute,
self.UPSERT_CHECKPOINTS_SQL,
(
thread_id,
checkpoint_ns,
checkpoint["id"],
checkpoint_id,
self._dump_checkpoint(copy),
self._dump_metadata(metadata),
),
)
return next_config
async def aput_writes(
self,
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
) -> None:
"""Store intermediate writes linked to a checkpoint asynchronously.
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.
"""
query = (
self.UPSERT_CHECKPOINT_WRITES_SQL
if all(w[0] in WRITES_IDX_MAP for w in writes)
else self.INSERT_CHECKPOINT_WRITES_SQL
)
params = await asyncio.to_thread(
self._dump_writes,
config["configurable"]["thread_id"],
config["configurable"]["checkpoint_ns"],
config["configurable"]["checkpoint_id"],
task_id,
writes,
)
async with self._cursor() as cur:
await asyncio.to_thread(cur.executemany, query, params)
@asynccontextmanager
async def _cursor(self) -> AsyncIterator[duckdb.DuckDBPyConnection]:
async with self.lock:
with self.conn.cursor() 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 DuckDB 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 DuckDB 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.
"""
try:
# check if we are in the main thread, only bg threads can block
# we don't check in other methods to avoid the overhead
if asyncio.get_running_loop() is self.loop:
raise asyncio.InvalidStateError(
"Synchronous calls to AsyncDuckDBSaver are only allowed from a "
"different thread. From the main thread, use the async interface."
"For example, use `await checkpointer.aget_tuple(...)` or `await "
"graph.ainvoke(...)`."
)
except RuntimeError:
pass
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 DuckDB 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: Sequence[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()
@@ -0,0 +1,290 @@
import json
import random
from typing import Any, List, Optional, Sequence, Tuple, cast
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
get_checkpoint_id,
)
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.serde.types import TASKS, ChannelProtocol
MetadataInput = Optional[dict[str, Any]]
"""
To add a new migration, add a new string to the MIGRATIONS list.
The position of the migration in the list is the version number.
"""
MIGRATIONS = [
"""CREATE TABLE IF NOT EXISTS checkpoint_migrations (
v INTEGER PRIMARY KEY
);""",
"""CREATE TABLE IF NOT EXISTS checkpoints (
thread_id TEXT NOT NULL,
checkpoint_ns TEXT NOT NULL DEFAULT '',
checkpoint_id TEXT NOT NULL,
parent_checkpoint_id TEXT,
type TEXT,
checkpoint JSON NOT NULL,
metadata JSON NOT NULL DEFAULT '{}',
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
);""",
"""CREATE TABLE IF NOT EXISTS checkpoint_blobs (
thread_id TEXT NOT NULL,
checkpoint_ns TEXT NOT NULL DEFAULT '',
channel TEXT NOT NULL,
version TEXT NOT NULL,
type TEXT NOT NULL,
blob BLOB,
PRIMARY KEY (thread_id, checkpoint_ns, channel, version)
);""",
"""CREATE TABLE IF NOT EXISTS checkpoint_writes (
thread_id TEXT NOT NULL,
checkpoint_ns TEXT NOT NULL DEFAULT '',
checkpoint_id TEXT NOT NULL,
task_id TEXT NOT NULL,
idx INTEGER NOT NULL,
channel TEXT NOT NULL,
type TEXT,
blob BLOB NOT NULL,
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
);""",
]
SELECT_SQL = f"""
select
thread_id,
checkpoint,
checkpoint_ns,
checkpoint_id,
parent_checkpoint_id,
metadata,
(
select array_agg(array[bl.channel::bytea, bl.type::bytea, bl.blob])
from (
SELECT unnest(json_keys(json_extract(checkpoint, '$.channel_versions'))) as key
) cv
inner join checkpoint_blobs bl
on bl.thread_id = checkpoints.thread_id
and bl.checkpoint_ns = checkpoints.checkpoint_ns
and bl.channel = cv.key
and bl.version = json_extract_string(checkpoint, '$.channel_versions.' || cv.key)
) as channel_values,
(
select
array_agg(array[cw.task_id::blob, cw.channel::blob, cw.type::blob, cw.blob])
from checkpoint_writes cw
where cw.thread_id = checkpoints.thread_id
and cw.checkpoint_ns = checkpoints.checkpoint_ns
and cw.checkpoint_id = checkpoints.checkpoint_id
) as pending_writes,
(
select array_agg(array[cw.type::blob, cw.blob])
from checkpoint_writes cw
where cw.thread_id = checkpoints.thread_id
and cw.checkpoint_ns = checkpoints.checkpoint_ns
and cw.checkpoint_id = checkpoints.parent_checkpoint_id
and cw.channel = '{TASKS}'
) as pending_sends
from checkpoints """
UPSERT_CHECKPOINT_BLOBS_SQL = """
INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob)
VALUES (?, ?, ?, ?, ?, ?)
ON CONFLICT (thread_id, checkpoint_ns, channel, version) DO NOTHING
"""
UPSERT_CHECKPOINTS_SQL = """
INSERT INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, checkpoint, metadata)
VALUES (?, ?, ?, ?, ?, ?)
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id)
DO UPDATE SET
checkpoint = EXCLUDED.checkpoint,
metadata = EXCLUDED.metadata;
"""
UPSERT_CHECKPOINT_WRITES_SQL = """
INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO UPDATE SET
channel = EXCLUDED.channel,
type = EXCLUDED.type,
blob = EXCLUDED.blob;
"""
INSERT_CHECKPOINT_WRITES_SQL = """
INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
"""
class BaseDuckDBSaver(BaseCheckpointSaver[str]):
SELECT_SQL = SELECT_SQL
MIGRATIONS = MIGRATIONS
UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL
jsonplus_serde = JsonPlusSerializer()
def _load_checkpoint(
self,
checkpoint_json_str: str,
channel_values: list[tuple[bytes, bytes, bytes]],
pending_sends: list[tuple[bytes, bytes]],
) -> Checkpoint:
checkpoint = json.loads(checkpoint_json_str)
return {
**checkpoint,
"pending_sends": [
self.serde.loads_typed((c.decode(), b)) for c, b in pending_sends or []
],
"channel_values": self._load_blobs(channel_values),
}
def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]:
return {**checkpoint, "pending_sends": []}
def _load_blobs(
self, blob_values: list[tuple[bytes, bytes, bytes]]
) -> dict[str, Any]:
if not blob_values:
return {}
return {
k.decode(): self.serde.loads_typed((t.decode(), v))
for k, t, v in blob_values
if t.decode() != "empty"
}
def _dump_blobs(
self,
thread_id: str,
checkpoint_ns: str,
values: dict[str, Any],
versions: ChannelVersions,
) -> list[tuple[str, str, str, str, str, Optional[bytes]]]:
if not versions:
return []
return [
(
thread_id,
checkpoint_ns,
k,
cast(str, ver),
*(
self.serde.dumps_typed(values[k])
if k in values
else ("empty", None)
),
)
for k, ver in versions.items()
]
def _load_writes(
self, writes: list[tuple[bytes, bytes, bytes, bytes]]
) -> list[tuple[str, str, Any]]:
return (
[
(
tid.decode(),
channel.decode(),
self.serde.loads_typed((t.decode(), v)),
)
for tid, channel, t, v in writes
]
if writes
else []
)
def _dump_writes(
self,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
task_id: str,
writes: Sequence[tuple[str, Any]],
) -> list[tuple[str, str, str, str, int, str, str, bytes]]:
return [
(
thread_id,
checkpoint_ns,
checkpoint_id,
task_id,
WRITES_IDX_MAP.get(channel, idx),
channel,
*self.serde.dumps_typed(value),
)
for idx, (channel, value) in enumerate(writes)
]
def _load_metadata(self, metadata_json_str: str) -> CheckpointMetadata:
return self.jsonplus_serde.loads(metadata_json_str.encode())
def _dump_metadata(self, metadata: CheckpointMetadata) -> str:
serialized_metadata = self.jsonplus_serde.dumps(metadata)
# NOTE: we're using JSON serializer (not msgpack), so we need to remove null characters before writing
return serialized_metadata.decode().replace("\\u0000", "")
def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
if current is None:
current_v = 0
elif isinstance(current, int):
current_v = current
else:
current_v = int(current.split(".")[0])
next_v = current_v + 1
next_h = random.random()
return f"{next_v:032}.{next_h:016}"
def _search_where(
self,
config: Optional[RunnableConfig],
filter: MetadataInput,
before: Optional[RunnableConfig] = None,
) -> Tuple[str, List[Any]]:
"""Return WHERE clause predicates for alist() given config, filter, before.
This method returns a tuple of a string and a tuple of values. The string
is the parametered WHERE clause predicate (including the WHERE keyword):
"WHERE column1 = $1 AND column2 IS $2". The list of values contains the
values for each of the corresponding parameters.
"""
wheres = []
param_values = []
# construct predicate for config filter
if config:
wheres.append("thread_id = ?")
param_values.append(config["configurable"]["thread_id"])
checkpoint_ns = config["configurable"].get("checkpoint_ns")
if checkpoint_ns is not None:
wheres.append("checkpoint_ns = ?")
param_values.append(checkpoint_ns)
if checkpoint_id := get_checkpoint_id(config):
wheres.append("checkpoint_id = ?")
param_values.append(checkpoint_id)
# construct predicate for metadata filter
if filter:
wheres.append("json_contains(metadata, ?)")
param_values.append(json.dumps(filter))
# construct predicate for `before`
if before is not None:
wheres.append("checkpoint_id < ?")
param_values.append(get_checkpoint_id(before))
return (
"WHERE " + " AND ".join(wheres) if wheres else "",
param_values,
)
@@ -0,0 +1,4 @@
from langgraph.store.duckdb.aio import AsyncDuckDBStore
from langgraph.store.duckdb.base import DuckDBStore
__all__ = ["AsyncDuckDBStore", "DuckDBStore"]
@@ -0,0 +1,195 @@
import asyncio
import logging
from contextlib import asynccontextmanager
from typing import (
AsyncIterator,
Iterable,
Sequence,
cast,
)
import duckdb
from langgraph.store.base import GetOp, ListNamespacesOp, Op, PutOp, Result, SearchOp
from langgraph.store.base.batch import AsyncBatchedBaseStore
from langgraph.store.duckdb.base import (
BaseDuckDBStore,
_convert_ns,
_group_ops,
_row_to_item,
)
logger = logging.getLogger(__name__)
class AsyncDuckDBStore(AsyncBatchedBaseStore, BaseDuckDBStore):
def __init__(
self,
conn: duckdb.DuckDBPyConnection,
) -> None:
super().__init__()
self.conn = conn
self.loop = asyncio.get_running_loop()
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
grouped_ops, num_ops = _group_ops(ops)
results: list[Result] = [None] * num_ops
tasks = []
if GetOp in grouped_ops:
tasks.append(
self._batch_get_ops(
cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results
)
)
if PutOp in grouped_ops:
tasks.append(
self._batch_put_ops(
cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp])
)
)
if SearchOp in grouped_ops:
tasks.append(
self._batch_search_ops(
cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
results,
)
)
if ListNamespacesOp in grouped_ops:
tasks.append(
self._batch_list_namespaces_ops(
cast(
Sequence[tuple[int, ListNamespacesOp]],
grouped_ops[ListNamespacesOp],
),
results,
)
)
await asyncio.gather(*tasks)
return results
def batch(self, ops: Iterable[Op]) -> list[Result]:
return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
async def _batch_get_ops(
self,
get_ops: Sequence[tuple[int, GetOp]],
results: list[Result],
) -> None:
cursors = []
for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
cur = self.conn.cursor()
await asyncio.to_thread(cur.execute, query, params)
cursors.append((cur, namespace, items))
for cur, namespace, items in cursors:
rows = await asyncio.to_thread(cur.fetchall)
key_to_row = {row[1]: row for row in rows}
for idx, key in items:
row = key_to_row.get(key)
if row:
results[idx] = _row_to_item(namespace, row)
else:
results[idx] = None
async def _batch_put_ops(
self,
put_ops: Sequence[tuple[int, PutOp]],
) -> None:
queries = self._get_batch_PUT_queries(put_ops)
for query, params in queries:
cur = self.conn.cursor()
await asyncio.to_thread(cur.execute, query, params)
async def _batch_search_ops(
self,
search_ops: Sequence[tuple[int, SearchOp]],
results: list[Result],
) -> None:
queries = self._get_batch_search_queries(search_ops)
cursors: list[tuple[duckdb.DuckDBPyConnection, int]] = []
for (query, params), (idx, _) in zip(queries, search_ops):
cur = self.conn.cursor()
await asyncio.to_thread(cur.execute, query, params)
cursors.append((cur, idx))
for cur, idx in cursors:
rows = await asyncio.to_thread(cur.fetchall)
items = [_row_to_item(_convert_ns(row[0]), row) for row in rows]
results[idx] = items
async def _batch_list_namespaces_ops(
self,
list_ops: Sequence[tuple[int, ListNamespacesOp]],
results: list[Result],
) -> None:
queries = self._get_batch_list_namespaces_queries(list_ops)
cursors: list[tuple[duckdb.DuckDBPyConnection, int]] = []
for (query, params), (idx, _) in zip(queries, list_ops):
cur = self.conn.cursor()
await asyncio.to_thread(cur.execute, query, params)
cursors.append((cur, idx))
for cur, idx in cursors:
rows = cast(list[tuple], await asyncio.to_thread(cur.fetchall))
namespaces = [_convert_ns(row[0]) for row in rows]
results[idx] = namespaces
@classmethod
@asynccontextmanager
async def from_conn_string(
cls,
conn_string: str,
) -> AsyncIterator["AsyncDuckDBStore"]:
"""Create a new AsyncDuckDBStore instance from a connection string.
Args:
conn_string (str): The DuckDB connection info string.
Returns:
AsyncDuckDBStore: A new AsyncDuckDBStore instance.
"""
with duckdb.connect(conn_string) as conn:
yield AsyncDuckDBStore(conn)
async def setup(self) -> None:
"""Set up the store database asynchronously.
This method creates the necessary tables in the DuckDB database if they don't
already exist and runs database migrations. It is called automatically when needed and should not be called
directly by the user.
"""
cur = self.conn.cursor()
try:
await asyncio.to_thread(
cur.execute, "SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
)
row = await asyncio.to_thread(cur.fetchone)
if row is None:
version = -1
else:
version = row[0]
except duckdb.CatalogException:
version = -1
# Create store_migrations table if it doesn't exist
await asyncio.to_thread(
cur.execute,
"""
CREATE TABLE IF NOT EXISTS store_migrations (
v INTEGER PRIMARY KEY
)
""",
)
for v, migration in enumerate(
self.MIGRATIONS[version + 1 :], start=version + 1
):
await asyncio.to_thread(cur.execute, migration)
await asyncio.to_thread(
cur.execute, "INSERT INTO store_migrations (v) VALUES (?)", (v,)
)
@@ -0,0 +1,391 @@
import asyncio
import json
import logging
from collections import defaultdict
from contextlib import contextmanager
from typing import (
Any,
Generic,
Iterable,
Iterator,
Sequence,
TypeVar,
Union,
cast,
)
import duckdb
from langgraph.store.base import (
BaseStore,
GetOp,
Item,
ListNamespacesOp,
Op,
PutOp,
Result,
SearchOp,
)
logger = logging.getLogger(__name__)
MIGRATIONS = [
"""
CREATE TABLE IF NOT EXISTS store (
prefix TEXT NOT NULL,
key TEXT NOT NULL,
value JSON NOT NULL,
created_at TIMESTAMP DEFAULT now(),
updated_at TIMESTAMP DEFAULT now(),
PRIMARY KEY (prefix, key)
);
""",
"""
CREATE INDEX IF NOT EXISTS store_prefix_idx ON store (prefix);
""",
]
C = TypeVar("C", bound=duckdb.DuckDBPyConnection)
class BaseDuckDBStore(Generic[C]):
MIGRATIONS = MIGRATIONS
conn: C
def _get_batch_GET_ops_queries(
self,
get_ops: Sequence[tuple[int, GetOp]],
) -> list[tuple[str, tuple, tuple[str, ...], list]]:
namespace_groups = defaultdict(list)
for idx, op in get_ops:
namespace_groups[op.namespace].append((idx, op.key))
results = []
for namespace, items in namespace_groups.items():
_, keys = zip(*items)
keys_to_query = ",".join(["?"] * len(keys))
query = f"""
SELECT prefix, key, value, created_at, updated_at
FROM store
WHERE prefix = ? AND key IN ({keys_to_query})
"""
params = (_namespace_to_text(namespace), *keys)
results.append((query, params, namespace, items))
return results
def _get_batch_PUT_queries(
self,
put_ops: Sequence[tuple[int, PutOp]],
) -> list[tuple[str, Sequence]]:
inserts: list[PutOp] = []
deletes: list[PutOp] = []
for _, op in put_ops:
if op.value is None:
deletes.append(op)
else:
inserts.append(op)
queries: list[tuple[str, Sequence]] = []
if deletes:
namespace_groups: dict[tuple[str, ...], list[str]] = defaultdict(list)
for op in deletes:
namespace_groups[op.namespace].append(op.key)
for namespace, keys in namespace_groups.items():
placeholders = ",".join(["?"] * len(keys))
query = (
f"DELETE FROM store WHERE prefix = ? AND key IN ({placeholders})"
)
params = (_namespace_to_text(namespace), *keys)
queries.append((query, params))
if inserts:
values = []
insertion_params = []
for op in inserts:
values.append("(?, ?, ?, now(), now())")
insertion_params.extend(
[
_namespace_to_text(op.namespace),
op.key,
json.dumps(op.value),
]
)
values_str = ",".join(values)
query = f"""
INSERT INTO store (prefix, key, value, created_at, updated_at)
VALUES {values_str}
ON CONFLICT (prefix, key) DO UPDATE
SET value = EXCLUDED.value, updated_at = now()
"""
queries.append((query, insertion_params))
return queries
def _get_batch_search_queries(
self,
search_ops: Sequence[tuple[int, SearchOp]],
) -> list[tuple[str, Sequence]]:
queries: list[tuple[str, Sequence]] = []
for _, op in search_ops:
query = """
SELECT prefix, key, value, created_at, updated_at
FROM store
WHERE prefix LIKE ?
"""
params: list = [f"{_namespace_to_text(op.namespace_prefix)}%"]
if op.filter:
filter_conditions = []
for key, value in op.filter.items():
filter_conditions.append(f"json_extract(value, '$.{key}') = ?")
params.append(json.dumps(value))
query += " AND " + " AND ".join(filter_conditions)
query += " ORDER BY updated_at DESC LIMIT ? OFFSET ?"
params.extend([op.limit, op.offset])
queries.append((query, params))
return queries
def _get_batch_list_namespaces_queries(
self,
list_ops: Sequence[tuple[int, ListNamespacesOp]],
) -> list[tuple[str, Sequence]]:
queries: list[tuple[str, Sequence]] = []
for _, op in list_ops:
query = """
WITH split_prefix AS (
SELECT
prefix,
string_split(prefix, '.') AS parts
FROM store
)
SELECT DISTINCT ON (truncated_prefix)
CASE
WHEN ? IS NOT NULL THEN
array_to_string(array_slice(parts, 1, ?), '.')
ELSE prefix
END AS truncated_prefix,
prefix
FROM split_prefix
"""
params: list[Any] = [op.max_depth, op.max_depth]
conditions = []
if op.match_conditions:
for condition in op.match_conditions:
if condition.match_type == "prefix":
conditions.append("prefix LIKE ?")
params.append(
f"{_namespace_to_text(condition.path, handle_wildcards=True)}%"
)
elif condition.match_type == "suffix":
conditions.append("prefix LIKE ?")
params.append(
f"%{_namespace_to_text(condition.path, handle_wildcards=True)}"
)
else:
logger.warning(
f"Unknown match_type in list_namespaces: {condition.match_type}"
)
if conditions:
query += " WHERE " + " AND ".join(conditions)
query += " ORDER BY prefix LIMIT ? OFFSET ?"
params.extend([op.limit, op.offset])
queries.append((query, params))
return queries
class DuckDBStore(BaseStore, BaseDuckDBStore[duckdb.DuckDBPyConnection]):
def __init__(
self,
conn: duckdb.DuckDBPyConnection,
) -> None:
super().__init__()
self.conn = conn
def batch(self, ops: Iterable[Op]) -> list[Result]:
grouped_ops, num_ops = _group_ops(ops)
results: list[Result] = [None] * num_ops
if GetOp in grouped_ops:
self._batch_get_ops(
cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results
)
if PutOp in grouped_ops:
self._batch_put_ops(cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]))
if SearchOp in grouped_ops:
self._batch_search_ops(
cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
results,
)
if ListNamespacesOp in grouped_ops:
self._batch_list_namespaces_ops(
cast(
Sequence[tuple[int, ListNamespacesOp]],
grouped_ops[ListNamespacesOp],
),
results,
)
return results
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
return await asyncio.get_running_loop().run_in_executor(None, self.batch, ops)
def _batch_get_ops(
self,
get_ops: Sequence[tuple[int, GetOp]],
results: list[Result],
) -> None:
cursors = []
for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
cur = self.conn.cursor()
cur.execute(query, params)
cursors.append((cur, namespace, items))
for cur, namespace, items in cursors:
rows = cur.fetchall()
key_to_row = {row[1]: row for row in rows}
for idx, key in items:
row = key_to_row.get(key)
if row:
results[idx] = _row_to_item(namespace, row)
else:
results[idx] = None
def _batch_put_ops(
self,
put_ops: Sequence[tuple[int, PutOp]],
) -> None:
queries = self._get_batch_PUT_queries(put_ops)
for query, params in queries:
cur = self.conn.cursor()
cur.execute(query, params)
def _batch_search_ops(
self,
search_ops: Sequence[tuple[int, SearchOp]],
results: list[Result],
) -> None:
queries = self._get_batch_search_queries(search_ops)
cursors: list[tuple[duckdb.DuckDBPyConnection, int]] = []
for (query, params), (idx, _) in zip(queries, search_ops):
cur = self.conn.cursor()
cur.execute(query, params)
cursors.append((cur, idx))
for cur, idx in cursors:
rows = cur.fetchall()
items = [_row_to_item(_convert_ns(row[0]), row) for row in rows]
results[idx] = items
def _batch_list_namespaces_ops(
self,
list_ops: Sequence[tuple[int, ListNamespacesOp]],
results: list[Result],
) -> None:
queries = self._get_batch_list_namespaces_queries(list_ops)
cursors: list[tuple[duckdb.DuckDBPyConnection, int]] = []
for (query, params), (idx, _) in zip(queries, list_ops):
cur = self.conn.cursor()
cur.execute(query, params)
cursors.append((cur, idx))
for cur, idx in cursors:
rows = cast(list[dict], cur.fetchall())
namespaces = [_convert_ns(row[0]) for row in rows]
results[idx] = namespaces
@classmethod
@contextmanager
def from_conn_string(
cls,
conn_string: str,
) -> Iterator["DuckDBStore"]:
"""Create a new BaseDuckDBStore instance from a connection string.
Args:
conn_string (str): The DuckDB connection info string.
Returns:
DuckDBStore: A new DuckDBStore instance.
"""
with duckdb.connect(conn_string) as conn:
yield cls(conn=conn)
def setup(self) -> None:
"""Set up the store database.
This method creates the necessary tables in the DuckDB database if they don't
already exist and runs database migrations. It is called automatically when needed and should not be called
directly by the user.
"""
with self.conn.cursor() as cur:
try:
cur.execute("SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1")
row = cast(dict, cur.fetchone())
if row is None:
version = -1
else:
version = row["v"]
except duckdb.CatalogException:
version = -1
# Create store_migrations table if it doesn't exist
cur.execute(
"""
CREATE TABLE IF NOT EXISTS store_migrations (
v INTEGER PRIMARY KEY
)
"""
)
for v, migration in enumerate(
self.MIGRATIONS[version + 1 :], start=version + 1
):
cur.execute(migration)
cur.execute("INSERT INTO store_migrations (v) VALUES (?)", (v,))
def _namespace_to_text(
namespace: tuple[str, ...], handle_wildcards: bool = False
) -> str:
"""Convert namespace tuple to text string."""
if handle_wildcards:
namespace = tuple("%" if val == "*" else val for val in namespace)
return ".".join(namespace)
def _row_to_item(
namespace: tuple[str, ...],
row: tuple,
) -> Item:
"""Convert a row from the database into an Item."""
_, key, val, created_at, updated_at = row
return Item(
value=val if isinstance(val, dict) else json.loads(val),
key=key,
namespace=namespace,
created_at=created_at,
updated_at=updated_at,
)
def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int]:
grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list)
tot = 0
for idx, op in enumerate(ops):
grouped_ops[type(op)].append((idx, op))
tot += 1
return grouped_ops, tot
def _convert_ns(namespace: Union[str, list]) -> tuple[str, ...]:
if isinstance(namespace, list):
return tuple(namespace)
return tuple(namespace.split("."))
+1058
View File
File diff suppressed because it is too large Load Diff
+60
View File
@@ -0,0 +1,60 @@
[tool.poetry]
name = "langgraph-checkpoint-duckdb"
version = "2.0.1"
description = "Library with a DuckDB implementation of LangGraph checkpoint saver."
authors = []
license = "MIT"
readme = "README.md"
repository = "https://www.github.com/langchain-ai/langgraph"
packages = [{ include = "langgraph" }]
[tool.poetry.dependencies]
python = "^3.9.0,<4.0"
langgraph-checkpoint = "^2.0.2"
duckdb = ">=1.1.2"
[tool.poetry.group.dev.dependencies]
ruff = "^0.6.2"
codespell = "^2.2.0"
pytest = "^7.2.1"
anyio = "^4.4.0"
pytest-asyncio = "^0.21.1"
pytest-mock = "^3.11.1"
pytest-watch = "^4.2.0"
mypy = "^1.10.0"
langgraph-checkpoint = {path = "../checkpoint", develop = true}
[tool.pytest.ini_options]
# --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
#
# https://docs.pytest.org/en/7.1.x/reference/reference.html
# --strict-config any warnings encountered while parsing the `pytest`
# section of the configuration file raise errors.
addopts = "--strict-markers --strict-config --durations=5 -vv"
asyncio_mode = "auto"
[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.ruff]
lint.select = [
"E", # pycodestyle
"F", # Pyflakes
"UP", # pyupgrade
"B", # flake8-bugbear
"I", # isort
]
lint.ignore = ["E501", "B008", "UP007", "UP006"]
[tool.mypy]
# https://mypy.readthedocs.io/en/stable/config_file.html
disallow_untyped_defs = "True"
explicit_package_bases = "True"
warn_no_return = "False"
warn_unused_ignores = "True"
warn_redundant_casts = "True"
allow_redefinition = "True"
disable_error_code = "typeddict-item, return-value"
+112
View File
@@ -0,0 +1,112 @@
from typing import Any
import pytest
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
Checkpoint,
CheckpointMetadata,
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.duckdb.aio import AsyncDuckDBSaver
class TestAsyncDuckDBSaver:
@pytest.fixture(autouse=True)
async def setup(self) -> None:
# objects for test setup
self.config_1: RunnableConfig = {
"configurable": {
"thread_id": "thread-1",
# for backwards compatibility testing
"thread_ts": "1",
"checkpoint_ns": "",
}
}
self.config_2: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2",
"checkpoint_ns": "",
}
}
self.config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
}
}
self.chkpnt_1: Checkpoint = empty_checkpoint()
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
self.chkpnt_3: Checkpoint = empty_checkpoint()
self.metadata_1: CheckpointMetadata = {
"source": "input",
"step": 2,
"writes": {},
"score": 1,
}
self.metadata_2: CheckpointMetadata = {
"source": "loop",
"step": 1,
"writes": {"foo": "bar"},
"score": None,
}
self.metadata_3: CheckpointMetadata = {}
async def test_asearch(self) -> None:
async with AsyncDuckDBSaver.from_conn_string(":memory:") as saver:
await saver.setup()
await saver.aput(self.config_1, self.chkpnt_1, self.metadata_1, {})
await saver.aput(self.config_2, self.chkpnt_2, self.metadata_2, {})
await saver.aput(self.config_3, self.chkpnt_3, self.metadata_3, {})
# call method / assertions
query_1 = {"source": "input"} # search by 1 key
query_2 = {
"step": 1,
"writes": {"foo": "bar"},
} # search by multiple keys
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
query_4 = {"source": "update", "step": 1} # no match
search_results_1 = [c async for c in saver.alist(None, filter=query_1)]
assert len(search_results_1) == 1
assert search_results_1[0].metadata == self.metadata_1
search_results_2 = [c async for c in saver.alist(None, filter=query_2)]
assert len(search_results_2) == 1
assert search_results_2[0].metadata == self.metadata_2
search_results_3 = [c async for c in saver.alist(None, filter=query_3)]
assert len(search_results_3) == 3
search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
assert len(search_results_4) == 0
# search by config (defaults to checkpoints across all namespaces)
search_results_5 = [
c
async for c in saver.alist({"configurable": {"thread_id": "thread-2"}})
]
assert len(search_results_5) == 2
assert {
search_results_5[0].config["configurable"]["checkpoint_ns"],
search_results_5[1].config["configurable"]["checkpoint_ns"],
} == {"", "inner"}
# TODO: test before and limit params
async def test_null_chars(self) -> None:
async with AsyncDuckDBSaver.from_conn_string(":memory:") as saver:
await saver.setup()
config = await saver.aput(
self.config_1, self.chkpnt_1, {"my_key": "\x00abc"}, {}
)
assert (await saver.aget_tuple(config)).metadata["my_key"] == "abc" # type: ignore
assert [c async for c in saver.alist(None, filter={"my_key": "abc"})][
0
].metadata["my_key"] == "abc"
@@ -0,0 +1,517 @@
# type: ignore
import uuid
from datetime import datetime
from typing import Any
from unittest.mock import MagicMock
import pytest
from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
from langgraph.store.duckdb import AsyncDuckDBStore
class MockCursor:
def __init__(self, fetch_result: Any) -> None:
self.fetch_result = fetch_result
self.execute = MagicMock()
self.fetchall = MagicMock(return_value=self.fetch_result)
class MockConnection:
def __init__(self) -> None:
self.cursor = MagicMock()
@pytest.fixture
def mock_connection() -> MockConnection:
return MockConnection()
@pytest.fixture
async def store(mock_connection: MockConnection) -> AsyncDuckDBStore:
duck_db_store = AsyncDuckDBStore(mock_connection)
await duck_db_store.setup()
return duck_db_store
async def test_abatch_order(store: AsyncDuckDBStore) -> None:
mock_connection = store.conn
mock_get_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
(
"test.bar",
"key2",
'{"data": "value2"}',
datetime.now(),
datetime.now(),
),
]
)
mock_search_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
]
)
mock_list_namespaces_cursor = MockCursor(
[
("test",),
]
)
failures = []
def cursor_side_effect() -> Any:
cursor = MagicMock()
def execute_side_effect(query: str, *params: Any) -> None:
# My super sophisticated database.
if "WHERE prefix = ? AND key" in query:
cursor.fetchall = mock_get_cursor.fetchall
elif "SELECT prefix, key, value" in query:
cursor.fetchall = mock_search_cursor.fetchall
elif "SELECT DISTINCT ON (truncated_prefix)" in query:
cursor.fetchall = mock_list_namespaces_cursor.fetchall
elif "INSERT INTO " in query:
pass
else:
e = ValueError(f"Unmatched query: {query}")
failures.append(e)
raise e
cursor.execute = MagicMock(side_effect=execute_side_effect)
return cursor
mock_connection.cursor.side_effect = cursor_side_effect # type: ignore
ops = [
GetOp(namespace=("test",), key="key1"),
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
SearchOp(
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
),
ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0),
GetOp(namespace=("test",), key="key3"),
]
results = await store.abatch(ops)
assert not failures
assert len(results) == 5
assert isinstance(results[0], Item)
assert isinstance(results[0].value, dict)
assert results[0].value == {"data": "value1"}
assert results[0].key == "key1"
assert results[1] is None
assert isinstance(results[2], list)
assert len(results[2]) == 1
assert isinstance(results[3], list)
assert results[3] == [("test",)]
assert results[4] is None
ops_reordered = [
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
GetOp(namespace=("test",), key="key2"),
ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0),
PutOp(namespace=("test",), key="key3", value={"data": "value3"}),
GetOp(namespace=("test",), key="key1"),
]
results_reordered = await store.abatch(ops_reordered)
assert not failures
assert len(results_reordered) == 5
assert isinstance(results_reordered[0], list)
assert len(results_reordered[0]) == 1
assert isinstance(results_reordered[1], Item)
assert results_reordered[1].value == {"data": "value2"}
assert results_reordered[1].key == "key2"
assert isinstance(results_reordered[2], list)
assert results_reordered[2] == [("test",)]
assert results_reordered[3] is None
assert isinstance(results_reordered[4], Item)
assert results_reordered[4].value == {"data": "value1"}
assert results_reordered[4].key == "key1"
async def test_batch_get_ops(store: AsyncDuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
(
"test.bar",
"key2",
'{"data": "value2"}',
datetime.now(),
datetime.now(),
),
]
)
mock_connection.cursor.return_value = mock_cursor
ops = [
GetOp(namespace=("test",), key="key1"),
GetOp(namespace=("test",), key="key2"),
GetOp(namespace=("test",), key="key3"),
]
results = await store.abatch(ops)
assert len(results) == 3
assert results[0] is not None
assert results[1] is not None
assert results[2] is None
assert results[0].key == "key1"
assert results[1].key == "key2"
async def test_batch_put_ops(store: AsyncDuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor([])
mock_connection.cursor.return_value = mock_cursor
ops = [
PutOp(namespace=("test",), key="key1", value={"data": "value1"}),
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
PutOp(namespace=("test",), key="key3", value=None),
]
results = await store.abatch(ops)
assert len(results) == 3
assert all(result is None for result in results)
assert mock_cursor.execute.call_count == 2
async def test_batch_search_ops(store: AsyncDuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
(
"test.bar",
"key2",
'{"data": "value2"}',
datetime.now(),
datetime.now(),
),
]
)
mock_connection.cursor.return_value = mock_cursor
ops = [
SearchOp(
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
),
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
]
results = await store.abatch(ops)
assert len(results) == 2
assert len(results[0]) == 2
assert len(results[1]) == 2
async def test_batch_list_namespaces_ops(store: AsyncDuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor([("test.namespace1",), ("test.namespace2",)])
mock_connection.cursor.return_value = mock_cursor
ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)]
results = await store.abatch(ops)
assert len(results) == 1
assert results[0] == [("test", "namespace1"), ("test", "namespace2")]
# The following use the actual DB connection
async def test_basic_store_ops() -> None:
async with AsyncDuckDBStore.from_conn_string(":memory:") as store:
await store.setup()
namespace = ("test", "documents")
item_id = "doc1"
item_value = {"title": "Test Document", "content": "Hello, World!"}
await store.aput(namespace, item_id, item_value)
item = await store.aget(namespace, item_id)
assert item
assert item.namespace == namespace
assert item.key == item_id
assert item.value == item_value
updated_value = {
"title": "Updated Test Document",
"content": "Hello, LangGraph!",
}
await store.aput(namespace, item_id, updated_value)
updated_item = await store.aget(namespace, item_id)
assert updated_item.value == updated_value
assert updated_item.updated_at > item.updated_at
different_namespace = ("test", "other_documents")
item_in_different_namespace = await store.aget(different_namespace, item_id)
assert item_in_different_namespace is None
new_item_id = "doc2"
new_item_value = {"title": "Another Document", "content": "Greetings!"}
await store.aput(namespace, new_item_id, new_item_value)
search_results = await store.asearch(["test"], limit=10)
items = search_results
assert len(items) == 2
assert any(item.key == item_id for item in items)
assert any(item.key == new_item_id for item in items)
namespaces = await store.alist_namespaces(prefix=["test"])
assert ("test", "documents") in namespaces
await store.adelete(namespace, item_id)
await store.adelete(namespace, new_item_id)
deleted_item = await store.aget(namespace, item_id)
assert deleted_item is None
deleted_item = await store.aget(namespace, new_item_id)
assert deleted_item is None
empty_search_results = await store.asearch(["test"], limit=10)
assert len(empty_search_results) == 0
async def test_list_namespaces() -> None:
async with AsyncDuckDBStore.from_conn_string(":memory:") as store:
await store.setup()
test_pref = str(uuid.uuid4())
test_namespaces = [
(test_pref, "test", "documents", "public", test_pref),
(test_pref, "test", "documents", "private", test_pref),
(test_pref, "test", "images", "public", test_pref),
(test_pref, "test", "images", "private", test_pref),
(test_pref, "prod", "documents", "public", test_pref),
(
test_pref,
"prod",
"documents",
"some",
"nesting",
"public",
test_pref,
),
(test_pref, "prod", "documents", "private", test_pref),
]
for namespace in test_namespaces:
await store.aput(namespace, "dummy", {"content": "dummy"})
prefix_result = await store.alist_namespaces(prefix=[test_pref, "test"])
assert len(prefix_result) == 4
assert all([ns[1] == "test" for ns in prefix_result])
specific_prefix_result = await store.alist_namespaces(
prefix=[test_pref, "test", "documents"]
)
assert len(specific_prefix_result) == 2
assert all([ns[1:3] == ("test", "documents") for ns in specific_prefix_result])
suffix_result = await store.alist_namespaces(suffix=["public", test_pref])
assert len(suffix_result) == 4
assert all(ns[-2] == "public" for ns in suffix_result)
prefix_suffix_result = await store.alist_namespaces(
prefix=[test_pref, "test"], suffix=["public", test_pref]
)
assert len(prefix_suffix_result) == 2
assert all(
ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result
)
wildcard_prefix_result = await store.alist_namespaces(
prefix=[test_pref, "*", "documents"]
)
assert len(wildcard_prefix_result) == 5
assert all(ns[2] == "documents" for ns in wildcard_prefix_result)
wildcard_suffix_result = await store.alist_namespaces(
suffix=["*", "public", test_pref]
)
assert len(wildcard_suffix_result) == 4
assert all(ns[-2] == "public" for ns in wildcard_suffix_result)
wildcard_single = await store.alist_namespaces(
suffix=["some", "*", "public", test_pref]
)
assert len(wildcard_single) == 1
assert wildcard_single[0] == (
test_pref,
"prod",
"documents",
"some",
"nesting",
"public",
test_pref,
)
max_depth_result = await store.alist_namespaces(max_depth=3)
assert all([len(ns) <= 3 for ns in max_depth_result])
max_depth_result = await store.alist_namespaces(
max_depth=4, prefix=[test_pref, "*", "documents"]
)
assert (
len(set(tuple(res) for res in max_depth_result))
== len(max_depth_result)
== 5
)
limit_result = await store.alist_namespaces(prefix=[test_pref], limit=3)
assert len(limit_result) == 3
offset_result = await store.alist_namespaces(prefix=[test_pref], offset=3)
assert len(offset_result) == len(test_namespaces) - 3
empty_prefix_result = await store.alist_namespaces(prefix=[test_pref])
assert len(empty_prefix_result) == len(test_namespaces)
assert set(tuple(ns) for ns in empty_prefix_result) == set(
tuple(ns) for ns in test_namespaces
)
for namespace in test_namespaces:
await store.adelete(namespace, "dummy")
async def test_search():
async with AsyncDuckDBStore.from_conn_string(":memory:") as store:
await store.setup()
test_namespaces = [
("test_search", "documents", "user1"),
("test_search", "documents", "user2"),
("test_search", "reports", "department1"),
("test_search", "reports", "department2"),
]
test_items = [
{"title": "Doc 1", "author": "John Doe", "tags": ["important"]},
{"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]},
{"title": "Report A", "author": "John Doe", "tags": ["final"]},
{"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]},
]
empty = await store.asearch(
(
"scoped",
"assistant_id",
"shared",
"6c5356f6-63ab-4158-868d-cd9fd14c736e",
),
limit=10,
offset=0,
)
assert len(empty) == 0
for namespace, item in zip(test_namespaces, test_items):
await store.aput(namespace, f"item_{namespace[-1]}", item)
docs_result = await store.asearch(["test_search", "documents"])
assert len(docs_result) == 2
assert all([item.namespace[1] == "documents" for item in docs_result]), [
item.namespace for item in docs_result
]
reports_result = await store.asearch(["test_search", "reports"])
assert len(reports_result) == 2
assert all(item.namespace[1] == "reports" for item in reports_result)
limited_result = await store.asearch(["test_search"], limit=2)
assert len(limited_result) == 2
offset_result = await store.asearch(["test_search"])
assert len(offset_result) == 4
offset_result = await store.asearch(["test_search"], offset=2)
assert len(offset_result) == 2
assert all(item not in limited_result for item in offset_result)
john_doe_result = await store.asearch(
["test_search"], filter={"author": "John Doe"}
)
assert len(john_doe_result) == 2
assert all(item.value["author"] == "John Doe" for item in john_doe_result)
draft_result = await store.asearch(["test_search"], filter={"tags": ["draft"]})
assert len(draft_result) == 2
assert all("draft" in item.value["tags"] for item in draft_result)
page1 = await store.asearch(["test_search"], limit=2, offset=0)
page2 = await store.asearch(["test_search"], limit=2, offset=2)
all_items = page1 + page2
assert len(all_items) == 4
assert len(set(item.key for item in all_items)) == 4
empty = await store.asearch(
(
"scoped",
"assistant_id",
"shared",
"again",
"maybe",
"some-long",
"6be5cb0e-2eb4-42e6-bb6b-fba3c269db25",
),
limit=10,
offset=0,
)
assert len(empty) == 0
# Test with a namespace beginning with a number (like a UUID)
uuid_namespace = (str(uuid.uuid4()), "documents")
uuid_item_id = "uuid_doc"
uuid_item_value = {
"title": "UUID Document",
"content": "This document has a UUID namespace.",
}
# Insert the item with the UUID namespace
await store.aput(uuid_namespace, uuid_item_id, uuid_item_value)
# Retrieve the item to verify it was stored correctly
retrieved_item = await store.aget(uuid_namespace, uuid_item_id)
assert retrieved_item is not None
assert retrieved_item.namespace == uuid_namespace
assert retrieved_item.key == uuid_item_id
assert retrieved_item.value == uuid_item_value
# Search for the item using the UUID namespace
search_result = await store.asearch([uuid_namespace[0]])
assert len(search_result) == 1
assert search_result[0].key == uuid_item_id
assert search_result[0].value == uuid_item_value
# Clean up: delete the item with the UUID namespace
await store.adelete(uuid_namespace, uuid_item_id)
# Verify the item was deleted
deleted_item = await store.aget(uuid_namespace, uuid_item_id)
assert deleted_item is None
for namespace in test_namespaces:
await store.adelete(namespace, f"item_{namespace[-1]}")
+457
View File
@@ -0,0 +1,457 @@
# type: ignore
import uuid
from datetime import datetime
from typing import Any
from unittest.mock import MagicMock
import pytest
from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
from langgraph.store.duckdb import DuckDBStore
class MockCursor:
def __init__(self, fetch_result: Any) -> None:
self.fetch_result = fetch_result
self.execute = MagicMock()
self.fetchall = MagicMock(return_value=self.fetch_result)
class MockConnection:
def __init__(self) -> None:
self.cursor = MagicMock()
@pytest.fixture
def mock_connection() -> MockConnection:
return MockConnection()
@pytest.fixture
def store(mock_connection: MockConnection) -> DuckDBStore:
duck_db_store = DuckDBStore(mock_connection)
duck_db_store.setup()
return duck_db_store
def test_batch_order(store: DuckDBStore) -> None:
mock_connection = store.conn
mock_get_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
(
"test.bar",
"key2",
'{"data": "value2"}',
datetime.now(),
datetime.now(),
),
]
)
mock_search_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
]
)
mock_list_namespaces_cursor = MockCursor(
[
("test",),
]
)
failures = []
def cursor_side_effect() -> Any:
cursor = MagicMock()
def execute_side_effect(query: str, *params: Any) -> None:
# My super sophisticated database.
if "WHERE prefix = ? AND key" in query:
cursor.fetchall = mock_get_cursor.fetchall
elif "SELECT prefix, key, value" in query:
cursor.fetchall = mock_search_cursor.fetchall
elif "SELECT DISTINCT ON (truncated_prefix)" in query:
cursor.fetchall = mock_list_namespaces_cursor.fetchall
elif "INSERT INTO " in query:
pass
else:
e = ValueError(f"Unmatched query: {query}")
failures.append(e)
raise e
cursor.execute = MagicMock(side_effect=execute_side_effect)
return cursor
mock_connection.cursor.side_effect = cursor_side_effect
ops = [
GetOp(namespace=("test",), key="key1"),
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
SearchOp(
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
),
ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0),
GetOp(namespace=("test",), key="key3"),
]
results = store.batch(ops)
assert not failures
assert len(results) == 5
assert isinstance(results[0], Item)
assert isinstance(results[0].value, dict)
assert results[0].value == {"data": "value1"}
assert results[0].key == "key1"
assert results[1] is None
assert isinstance(results[2], list)
assert len(results[2]) == 1
assert isinstance(results[3], list)
assert results[3] == [("test",)]
assert results[4] is None
ops_reordered = [
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
GetOp(namespace=("test",), key="key2"),
ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0),
PutOp(namespace=("test",), key="key3", value={"data": "value3"}),
GetOp(namespace=("test",), key="key1"),
]
results_reordered = store.batch(ops_reordered)
assert not failures
assert len(results_reordered) == 5
assert isinstance(results_reordered[0], list)
assert len(results_reordered[0]) == 1
assert isinstance(results_reordered[1], Item)
assert results_reordered[1].value == {"data": "value2"}
assert results_reordered[1].key == "key2"
assert isinstance(results_reordered[2], list)
assert results_reordered[2] == [("test",)]
assert results_reordered[3] is None
assert isinstance(results_reordered[4], Item)
assert results_reordered[4].value == {"data": "value1"}
assert results_reordered[4].key == "key1"
def test_batch_get_ops(store: DuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
(
"test.bar",
"key2",
'{"data": "value2"}',
datetime.now(),
datetime.now(),
),
]
)
mock_connection.cursor.return_value = mock_cursor
ops = [
GetOp(namespace=("test",), key="key1"),
GetOp(namespace=("test",), key="key2"),
GetOp(namespace=("test",), key="key3"),
]
results = store.batch(ops)
assert len(results) == 3
assert results[0] is not None
assert results[1] is not None
assert results[2] is None
assert results[0].key == "key1"
assert results[1].key == "key2"
def test_batch_put_ops(store: DuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor([])
mock_connection.cursor.return_value = mock_cursor
ops = [
PutOp(namespace=("test",), key="key1", value={"data": "value1"}),
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
PutOp(namespace=("test",), key="key3", value=None),
]
results = store.batch(ops)
assert len(results) == 3
assert all(result is None for result in results)
assert mock_cursor.execute.call_count == 2
def test_batch_search_ops(store: DuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
(
"test.bar",
"key2",
'{"data": "value2"}',
datetime.now(),
datetime.now(),
),
]
)
mock_connection.cursor.return_value = mock_cursor
ops = [
SearchOp(
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
),
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
]
results = store.batch(ops)
assert len(results) == 2
assert len(results[0]) == 2
assert len(results[1]) == 2
def test_batch_list_namespaces_ops(store: DuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor([("test.namespace1",), ("test.namespace2",)])
mock_connection.cursor.return_value = mock_cursor
ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)]
results = store.batch(ops)
assert len(results) == 1
assert results[0] == [("test", "namespace1"), ("test", "namespace2")]
def test_basic_store_ops() -> None:
with DuckDBStore.from_conn_string(":memory:") as store:
store.setup()
namespace = ("test", "documents")
item_id = "doc1"
item_value = {"title": "Test Document", "content": "Hello, World!"}
store.put(namespace, item_id, item_value)
item = store.get(namespace, item_id)
assert item
assert item.namespace == namespace
assert item.key == item_id
assert item.value == item_value
updated_value = {
"title": "Updated Test Document",
"content": "Hello, LangGraph!",
}
store.put(namespace, item_id, updated_value)
updated_item = store.get(namespace, item_id)
assert updated_item.value == updated_value
assert updated_item.updated_at > item.updated_at
different_namespace = ("test", "other_documents")
item_in_different_namespace = store.get(different_namespace, item_id)
assert item_in_different_namespace is None
new_item_id = "doc2"
new_item_value = {"title": "Another Document", "content": "Greetings!"}
store.put(namespace, new_item_id, new_item_value)
search_results = store.search(["test"], limit=10)
items = search_results
assert len(items) == 2
assert any(item.key == item_id for item in items)
assert any(item.key == new_item_id for item in items)
namespaces = store.list_namespaces(prefix=["test"])
assert ("test", "documents") in namespaces
store.delete(namespace, item_id)
store.delete(namespace, new_item_id)
deleted_item = store.get(namespace, item_id)
assert deleted_item is None
deleted_item = store.get(namespace, new_item_id)
assert deleted_item is None
empty_search_results = store.search(["test"], limit=10)
assert len(empty_search_results) == 0
def test_list_namespaces() -> None:
with DuckDBStore.from_conn_string(":memory:") as store:
store.setup()
test_pref = str(uuid.uuid4())
test_namespaces = [
(test_pref, "test", "documents", "public", test_pref),
(test_pref, "test", "documents", "private", test_pref),
(test_pref, "test", "images", "public", test_pref),
(test_pref, "test", "images", "private", test_pref),
(test_pref, "prod", "documents", "public", test_pref),
(
test_pref,
"prod",
"documents",
"some",
"nesting",
"public",
test_pref,
),
(test_pref, "prod", "documents", "private", test_pref),
]
for namespace in test_namespaces:
store.put(namespace, "dummy", {"content": "dummy"})
prefix_result = store.list_namespaces(prefix=[test_pref, "test"])
assert len(prefix_result) == 4
assert all([ns[1] == "test" for ns in prefix_result])
specific_prefix_result = store.list_namespaces(
prefix=[test_pref, "test", "documents"]
)
assert len(specific_prefix_result) == 2
assert all([ns[1:3] == ("test", "documents") for ns in specific_prefix_result])
suffix_result = store.list_namespaces(suffix=["public", test_pref])
assert len(suffix_result) == 4
assert all(ns[-2] == "public" for ns in suffix_result)
prefix_suffix_result = store.list_namespaces(
prefix=[test_pref, "test"], suffix=["public", test_pref]
)
assert len(prefix_suffix_result) == 2
assert all(
ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result
)
wildcard_prefix_result = store.list_namespaces(
prefix=[test_pref, "*", "documents"]
)
assert len(wildcard_prefix_result) == 5
assert all(ns[2] == "documents" for ns in wildcard_prefix_result)
wildcard_suffix_result = store.list_namespaces(
suffix=["*", "public", test_pref]
)
assert len(wildcard_suffix_result) == 4
assert all(ns[-2] == "public" for ns in wildcard_suffix_result)
wildcard_single = store.list_namespaces(
suffix=["some", "*", "public", test_pref]
)
assert len(wildcard_single) == 1
assert wildcard_single[0] == (
test_pref,
"prod",
"documents",
"some",
"nesting",
"public",
test_pref,
)
max_depth_result = store.list_namespaces(max_depth=3)
assert all([len(ns) <= 3 for ns in max_depth_result])
max_depth_result = store.list_namespaces(
max_depth=4, prefix=[test_pref, "*", "documents"]
)
assert (
len(set(tuple(res) for res in max_depth_result))
== len(max_depth_result)
== 5
)
limit_result = store.list_namespaces(prefix=[test_pref], limit=3)
assert len(limit_result) == 3
offset_result = store.list_namespaces(prefix=[test_pref], offset=3)
assert len(offset_result) == len(test_namespaces) - 3
empty_prefix_result = store.list_namespaces(prefix=[test_pref])
assert len(empty_prefix_result) == len(test_namespaces)
assert set(tuple(ns) for ns in empty_prefix_result) == set(
tuple(ns) for ns in test_namespaces
)
for namespace in test_namespaces:
store.delete(namespace, "dummy")
def test_search():
with DuckDBStore.from_conn_string(":memory:") as store:
store.setup()
test_namespaces = [
("test_search", "documents", "user1"),
("test_search", "documents", "user2"),
("test_search", "reports", "department1"),
("test_search", "reports", "department2"),
]
test_items = [
{"title": "Doc 1", "author": "John Doe", "tags": ["important"]},
{"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]},
{"title": "Report A", "author": "John Doe", "tags": ["final"]},
{"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]},
]
for namespace, item in zip(test_namespaces, test_items):
store.put(namespace, f"item_{namespace[-1]}", item)
docs_result = store.search(["test_search", "documents"])
assert len(docs_result) == 2
assert all(
[item.namespace[1] == "documents" for item in docs_result]
), docs_result
reports_result = store.search(["test_search", "reports"])
assert len(reports_result) == 2
assert all(item.namespace[1] == "reports" for item in reports_result)
limited_result = store.search(["test_search"], limit=2)
assert len(limited_result) == 2
offset_result = store.search(["test_search"])
assert len(offset_result) == 4
offset_result = store.search(["test_search"], offset=2)
assert len(offset_result) == 2
assert all(item not in limited_result for item in offset_result)
john_doe_result = store.search(["test_search"], filter={"author": "John Doe"})
assert len(john_doe_result) == 2
assert all(item.value["author"] == "John Doe" for item in john_doe_result)
draft_result = store.search(["test_search"], filter={"tags": ["draft"]})
assert len(draft_result) == 2
assert all("draft" in item.value["tags"] for item in draft_result)
page1 = store.search(["test_search"], limit=2, offset=0)
page2 = store.search(["test_search"], limit=2, offset=2)
all_items = page1 + page2
assert len(all_items) == 4
assert len(set(item.key for item in all_items)) == 4
for namespace in test_namespaces:
store.delete(namespace, f"item_{namespace[-1]}")
+111
View File
@@ -0,0 +1,111 @@
from typing import Any
import pytest
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
Checkpoint,
CheckpointMetadata,
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.duckdb import DuckDBSaver
class TestDuckDBSaver:
@pytest.fixture(autouse=True)
def setup(self) -> None:
# objects for test setup
self.config_1: RunnableConfig = {
"configurable": {
"thread_id": "thread-1",
# for backwards compatibility testing
"thread_ts": "1",
"checkpoint_ns": "",
}
}
self.config_2: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2",
"checkpoint_ns": "",
}
}
self.config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
}
}
self.chkpnt_1: Checkpoint = empty_checkpoint()
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
self.chkpnt_3: Checkpoint = empty_checkpoint()
self.metadata_1: CheckpointMetadata = {
"source": "input",
"step": 2,
"writes": {},
"score": 1,
}
self.metadata_2: CheckpointMetadata = {
"source": "loop",
"step": 1,
"writes": {"foo": "bar"},
"score": None,
}
self.metadata_3: CheckpointMetadata = {}
def test_search(self) -> None:
with DuckDBSaver.from_conn_string(":memory:") as saver:
saver.setup()
# save checkpoints
saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {})
saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
# call method / assertions
query_1 = {"source": "input"} # search by 1 key
query_2 = {
"step": 1,
"writes": {"foo": "bar"},
} # search by multiple keys
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
query_4 = {"source": "update", "step": 1} # no match
search_results_1 = list(saver.list(None, filter=query_1))
assert len(search_results_1) == 1
assert search_results_1[0].metadata == self.metadata_1
search_results_2 = list(saver.list(None, filter=query_2))
assert len(search_results_2) == 1
assert search_results_2[0].metadata == self.metadata_2
search_results_3 = list(saver.list(None, filter=query_3))
assert len(search_results_3) == 3
search_results_4 = list(saver.list(None, filter=query_4))
assert len(search_results_4) == 0
# search by config (defaults to checkpoints across all namespaces)
search_results_5 = list(
saver.list({"configurable": {"thread_id": "thread-2"}})
)
assert len(search_results_5) == 2
assert {
search_results_5[0].config["configurable"]["checkpoint_ns"],
search_results_5[1].config["configurable"]["checkpoint_ns"],
} == {"", "inner"}
# TODO: test before and limit params
def test_null_chars(self) -> None:
with DuckDBSaver.from_conn_string(":memory:") as saver:
saver.setup()
config = saver.put(self.config_1, self.chkpnt_1, {"my_key": "\x00abc"}, {})
assert saver.get_tuple(config).metadata["my_key"] == "abc" # type: ignore
assert (
list(saver.list(None, filter={"my_key": "abc"}))[0].metadata["my_key"] # type: ignore
== "abc"
)
@@ -287,7 +287,7 @@ class PostgresSaver(BasePostgresSaver):
>>> DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
>>> with PostgresSaver.from_conn_string(DB_URI) as memory:
>>> config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
>>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "data": {"key": "value"}}
>>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}}
>>> saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {})
>>> print(saved_config)
{'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}}
@@ -249,7 +249,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
filter: MetadataInput,
before: Optional[RunnableConfig] = None,
) -> Tuple[str, List[Any]]:
"""Return WHERE clause predicates for alist() given config, filter, cursor.
"""Return WHERE clause predicates for alist() given config, filter, before.
This method returns a tuple of a string and a tuple of values. The string
is the parametered WHERE clause predicate (including the WHERE keyword):
@@ -44,7 +44,6 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnectio
super().__init__()
self._deserializer = deserializer
self.conn = conn
self.conn = conn
self.loop = asyncio.get_running_loop()
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
+3 -3
View File
@@ -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"
@@ -324,7 +324,7 @@ typing-extensions = ">=4.7"
[[package]]
name = "langgraph-checkpoint"
version = "2.0.0"
version = "2.0.2"
description = "Library with base interfaces for LangGraph checkpoint savers."
optional = false
python-versions = "^3.9.0,<4.0"
@@ -1116,4 +1116,4 @@ watchmedo = ["PyYAML (>=3.10)"]
[metadata]
lock-version = "2.0"
python-versions = "^3.9.0,<4.0"
content-hash = "8f763cd1727287f8c8b5ad2b4d8df00fb446e68d0cd4e88c278e4007969b83fd"
content-hash = "6bd85ce8ee1192995c1ff03d5fa65af8ee7872214d71b84559a6192cadf82be6"
+2 -2
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-checkpoint-postgres"
version = "2.0.1"
version = "2.0.2"
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
authors = []
license = "MIT"
@@ -10,7 +10,7 @@ packages = [{ include = "langgraph" }]
[tool.poetry.dependencies]
python = "^3.9.0,<4.0"
langgraph-checkpoint = "^2.0.0"
langgraph-checkpoint = "^2.0.2"
orjson = ">=3.10.1"
psycopg = "^3.0.0"
psycopg-pool = "^3.0.0"
@@ -389,7 +389,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
>>> from langgraph.checkpoint.sqlite import SqliteSaver
>>> with SqliteSaver.from_conn_string(":memory:") as memory:
>>> config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
>>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "data": {"key": "value"}}
>>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}}
>>> saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {})
>>> print(saved_config)
{'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}}
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-checkpoint"
version = "2.0.1"
version = "2.0.2"
description = "Library with base interfaces for LangGraph checkpoint savers."
authors = []
license = "MIT"
@@ -3,7 +3,12 @@ from typing import Generic, Optional, Sequence, Type
from typing_extensions import Self
from langgraph.channels.base import BaseChannel, Value
from langgraph.errors import EmptyChannelError, InvalidUpdateError
from langgraph.errors import (
EmptyChannelError,
ErrorCode,
InvalidUpdateError,
create_error_message,
)
class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
@@ -35,9 +40,11 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
if len(values) == 0:
return False
if len(values) != 1:
raise InvalidUpdateError(
f"At key '{self.key}': Can receive only one value per step. Use an Annotated key to handle multiple values."
msg = create_error_message(
message=f"At key '{self.key}': Can receive only one value per step. Use an Annotated key to handle multiple values.",
error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE,
)
raise InvalidUpdateError(msg)
self.value = values[-1]
return True
+2
View File
@@ -12,6 +12,8 @@ EMPTY_MAP: Mapping[str, Any] = MappingProxyType({})
EMPTY_SEQ: tuple[str, ...] = tuple()
# --- Public constants ---
TAG_NOSTREAM = sys.intern("langsmith:nostream")
"""Tag to disable streaming for a chat model."""
TAG_HIDDEN = sys.intern("langsmith:hidden")
"""Tag to hide a node/edge from certain tracing/streaming environments."""
START = sys.intern("__start__")
+33 -2
View File
@@ -1,3 +1,4 @@
from enum import Enum
from typing import Any, Sequence
from langgraph.checkpoint.base import EmptyChannelError # noqa: F401
@@ -6,12 +7,31 @@ from langgraph.types import Interrupt
# EmptyChannelError re-exported for backwards compatibility
class ErrorCode(Enum):
GRAPH_RECURSION_LIMIT = "GRAPH_RECURSION_LIMIT"
INVALID_CONCURRENT_GRAPH_UPDATE = "INVALID_CONCURRENT_GRAPH_UPDATE"
INVALID_GRAPH_NODE_RETURN_VALUE = "INVALID_GRAPH_NODE_RETURN_VALUE"
MULTIPLE_SUBGRAPHS = "MULTIPLE_SUBGRAPHS"
def create_error_message(*, message: str, error_code: ErrorCode) -> str:
return (
f"{message}\n"
"For troubleshooting, visit: https://python.langchain.com/docs/"
f"troubleshooting/errors/{error_code.value}"
)
class GraphRecursionError(RecursionError):
"""Raised when the graph has exhausted the maximum number of steps.
This prevents infinite loops. To increase the maximum number of steps,
run your graph with a config specifying a higher `recursion_limit`.
Troubleshooting Guides:
- [GRAPH_RECURSION_LIMIT](https://python.langchain.com/docs/troubleshooting/errors/GRAPH_RECURSION_LIMIT)
Examples:
graph = builder.compile()
@@ -26,7 +46,13 @@ class GraphRecursionError(RecursionError):
class InvalidUpdateError(Exception):
"""Raised when attempting to update a channel with an invalid set of updates."""
"""Raised when attempting to update a channel with an invalid set of updates.
Troubleshooting Guides:
- [INVALID_CONCURRENT_GRAPH_UPDATE](https://python.langchain.com/docs/troubleshooting/errors/INVALID_CONCURRENT_GRAPH_UPDATE)
- [INVALID_GRAPH_NODE_RETURN_VALUE](https://python.langchain.com/docs/troubleshooting/errors/INVALID_GRAPH_NODE_RETURN_VALUE)
"""
pass
@@ -72,7 +98,12 @@ class CheckpointNotLatest(Exception):
class MultipleSubgraphsError(Exception):
"""Raised when multiple subgraphs are called inside the same node."""
"""Raised when multiple subgraphs are called inside the same node.
Troubleshooting guides:
- [MULTIPLE_SUBGRAPHS](https://python.langchain.com/docs/troubleshooting/errors/MULTIPLE_SUBGRAPHS)
"""
pass
+8
View File
@@ -514,6 +514,14 @@ class CompiledGraph(Pregel):
self.nodes[end].triggers.append(channel_name)
cast(list[str], self.nodes[end].channels).append(channel_name)
async def aget_graph(
self,
config: Optional[RunnableConfig] = None,
*,
xray: Union[int, bool] = False,
) -> DrawableGraph:
return self.get_graph(config, xray=xray)
def get_graph(
self,
config: Optional[RunnableConfig] = None,
+27 -16
View File
@@ -3,7 +3,7 @@ import logging
import typing
import warnings
from functools import partial
from inspect import isclass, isfunction, signature
from inspect import isclass, isfunction, ismethod, signature
from typing import (
Any,
Callable,
@@ -33,7 +33,7 @@ from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue
from langgraph.channels.named_barrier_value import NamedBarrierValue
from langgraph.constants import NS_END, NS_SEP, TAG_HIDDEN
from langgraph.errors import InvalidUpdateError
from langgraph.errors import ErrorCode, InvalidUpdateError, create_error_message
from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send
from langgraph.managed.base import (
ChannelKeyPlaceholder,
@@ -338,19 +338,8 @@ class StateGraph(Graph):
f"'{character}' is a reserved character and is not allowed in the node names."
)
try:
if isfunction(action) and (
hints := get_type_hints(action.__call__) or get_type_hints(action)
):
if input is None:
first_parameter_name = next(
iter(inspect.signature(action).parameters.keys())
)
if input_hint := hints.get(first_parameter_name):
if isinstance(input_hint, type) and get_type_hints(input_hint):
input = input_hint
except (TypeError, StopIteration):
pass
if input is None:
input = _get_input_schema_from_type_hint(action)
if input is not None:
self._add_schema(input)
self.nodes[cast(str, node)] = StateNodeSpec(
@@ -538,7 +527,11 @@ class CompiledStateGraph(CompiledGraph):
value = getattr(input, key, SKIP_WRITE)
return value if value is not None else SKIP_WRITE
else:
raise InvalidUpdateError(f"Expected dict, got {input}")
msg = create_error_message(
message=f"Expected dict, got {input}",
error_code=ErrorCode.INVALID_GRAPH_NODE_RETURN_VALUE,
)
raise InvalidUpdateError(msg)
# state updaters
write_entries = (
@@ -830,3 +823,21 @@ def _get_schema(
if k in channels and isinstance(channels[k], BaseChannel)
},
)
def _get_input_schema_from_type_hint(
action: Optional[RunnableLike],
) -> Optional[Type[Any]]:
if not isfunction(action) and not ismethod(getattr(action, "__call__", None)):
return None
action = cast(Callable, action)
try:
hints = get_type_hints(getattr(action, "__call__")) or get_type_hints(action)
first_parameter_name = next(iter(inspect.signature(action).parameters.keys()))
input_hint = hints.get(first_parameter_name)
if isinstance(input_hint, type) and get_type_hints(input_hint):
return input_hint
except (TypeError, StopIteration):
pass
return None
+156 -18
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
import inspect
import json
from copy import copy
from typing import (
@@ -16,6 +17,7 @@ from typing import (
Type,
Union,
cast,
get_type_hints,
)
from langchain_core.messages import (
@@ -34,6 +36,7 @@ from langchain_core.tools import BaseTool, InjectedToolArg
from langchain_core.tools import tool as create_tool
from typing_extensions import Annotated, get_args, get_origin
from langgraph.errors import GraphInterrupt
from langgraph.store.base import BaseStore
from langgraph.utils.runnable import RunnableCallable
@@ -67,13 +70,96 @@ def msg_content_output(output: Any) -> str | List[dict]:
return str(output)
def _handle_tool_error(
e: Exception,
*,
flag: Union[
bool,
str,
Callable[..., str],
tuple[type[Exception], ...],
],
) -> str:
if isinstance(flag, (bool, tuple)):
content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
elif isinstance(flag, str):
content = flag
elif callable(flag):
content = flag(e)
else:
raise ValueError(
f"Got unexpected type of `handle_tool_error`. Expected bool, str "
f"or callable. Received: {flag}"
)
return content
def _infer_handled_types(handler: Callable[..., str]) -> tuple[type[Exception]]:
sig = inspect.signature(handler)
params = list(sig.parameters.values())
if params:
# If it's a method, the first argument is typically 'self' or 'cls'
if params[0].name in ["self", "cls"] and len(params) == 2:
first_param = params[1]
else:
first_param = params[0]
type_hints = get_type_hints(handler)
if first_param.name in type_hints:
origin = get_origin(first_param.annotation)
if origin is Union:
args = get_args(first_param.annotation)
if all(issubclass(arg, Exception) for arg in args):
return tuple(args)
else:
raise ValueError(
"All types in the error handler error annotation must be Exception types. "
"For example, `def custom_handler(e: Union[ValueError, TypeError])`. "
f"Got '{first_param.annotation}' instead."
)
exception_type = type_hints[first_param.name]
if Exception in exception_type.__mro__:
return (exception_type,)
else:
raise ValueError(
f"Arbitrary types are not supported in the error handler signature. "
"Please annotate the error with either a specific Exception type or a union of Exception types. "
"For example, `def custom_handler(e: ValueError)` or `def custom_handler(e: Union[ValueError, TypeError])`. "
f"Got '{exception_type}' instead."
)
# If no type information is available, return (Exception,) for backwards compatibility.
return (Exception,)
class ToolNode(RunnableCallable):
"""A node that runs the tools called in the last AIMessage.
It can be used either in StateGraph with a "messages" key or in MessageGraph. If
multiple tool calls are requested, they will be run in parallel. The output will be
It can be used either in StateGraph with a "messages" state key (or a custom key passed via ToolNode's 'messages_key').
If multiple tool calls are requested, they will be run in parallel. The output will be
a list of ToolMessages, one for each tool call.
Args:
tools: A sequence of tools that can be invoked by the ToolNode.
name: The name of the ToolNode in the graph. Defaults to "tools".
tags: Optional tags to associate with the node. Defaults to None.
handle_tool_errors: How to handle tool errors raised by tools inside the node. Defaults to True.
Must be one of the following:
- True: all errors will be caught and
a ToolMessage with a default error message (TOOL_CALL_ERROR_TEMPLATE) will be returned.
- str: all errors will be caught and
a ToolMessage with the string value of 'handle_tool_errors' will be returned.
- tuple[type[Exception], ...]: exceptions in the tuple will be caught and
a ToolMessage with a default error message (TOOL_CALL_ERROR_TEMPLATE) will be returned.
- Callable[..., str]: exceptions from the signature of the callable will be caught and
a ToolMessage with the string value of the result of the 'handle_tool_errors' callable will be returned.
- False: none of the errors raised by the tools will be caught
messages_key: The state key in the input that contains the list of messages.
The same key will be used for the output from the ToolNode.
Defaults to "messages".
The `ToolNode` is roughly analogous to:
```python
@@ -101,13 +187,17 @@ class ToolNode(RunnableCallable):
*,
name: str = "tools",
tags: Optional[list[str]] = None,
handle_tool_errors: Optional[bool] = True,
handle_tool_errors: Union[
bool, str, Callable[..., str], tuple[type[Exception], ...]
] = True,
messages_key: str = "messages",
) -> None:
super().__init__(self._func, self._afunc, name=name, tags=tags, trace=False)
self.tools_by_name: Dict[str, BaseTool] = {}
self.tool_to_state_args: Dict[str, Dict[str, Optional[str]]] = {}
self.tool_to_store_arg: Dict[str, Optional[str]] = {}
self.handle_tool_errors = handle_tool_errors
self.messages_key = messages_key
for tool_ in tools:
if not isinstance(tool_, BaseTool):
tool_ = cast(BaseTool, create_tool(tool_))
@@ -131,7 +221,7 @@ class ToolNode(RunnableCallable):
with get_executor_for_config(config) as executor:
outputs = [*executor.map(self._run_one, tool_calls, config_list)]
# TypedDict, pydantic, dataclass, etc. should all be able to load from dict
return outputs if output_type == "list" else {"messages": outputs}
return outputs if output_type == "list" else {self.messages_key: outputs}
def invoke(
self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
@@ -163,7 +253,7 @@ class ToolNode(RunnableCallable):
*(self._arun_one(call, config) for call in tool_calls)
)
# TypedDict, pydantic, dataclass, etc. should all be able to load from dict
return outputs if output_type == "list" else {"messages": outputs}
return outputs if output_type == "list" else {self.messages_key: outputs}
def _run_one(self, call: ToolCall, config: RunnableConfig) -> ToolMessage:
if invalid_tool_message := self._validate_tool_call(call):
@@ -178,15 +268,38 @@ class ToolNode(RunnableCallable):
Union[str, list], msg_content_output(tool_message.content)
)
return tool_message
# GraphInterrupt is a special exception that will always be raised.
# It can be triggered in the following scenarios:
# (1) a NodeInterrupt is raised inside a tool
# (2) a NodeInterrupt is raised inside a graph node for a graph called as a tool
# (3) a GraphInterrupt is raised when a subgraph is interrupted inside a graph called as a tool
# (2 and 3 can happen in a "supervisor w/ tools" multi-agent architecture)
except GraphInterrupt as e:
raise e
except Exception as e:
if not self.handle_tool_errors:
if isinstance(self.handle_tool_errors, tuple):
handled_types: tuple = self.handle_tool_errors
elif callable(self.handle_tool_errors):
handled_types = _infer_handled_types(self.handle_tool_errors)
else:
# default behavior is catching all exceptions
handled_types = (Exception,)
# Unhandled
if not self.handle_tool_errors or not isinstance(e, handled_types):
raise e
content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
return ToolMessage(content, name=call["name"], tool_call_id=call["id"])
# Handled
else:
content = _handle_tool_error(e, flag=self.handle_tool_errors)
return ToolMessage(
content=content, name=call["name"], tool_call_id=call["id"], status="error"
)
async def _arun_one(self, call: ToolCall, config: RunnableConfig) -> ToolMessage:
if invalid_tool_message := self._validate_tool_call(call):
return invalid_tool_message
try:
input = {**call, **{"type": "tool_call"}}
tool_message: ToolMessage = await self.tools_by_name[call["name"]].ainvoke(
@@ -196,11 +309,33 @@ class ToolNode(RunnableCallable):
Union[str, list], msg_content_output(tool_message.content)
)
return tool_message
# GraphInterrupt is a special exception that will always be raised.
# It can be triggered in the following scenarios:
# (1) a NodeInterrupt is raised inside a tool
# (2) a NodeInterrupt is raised inside a graph node for a graph called as a tool
# (3) a GraphInterrupt is raised when a subgraph is interrupted inside a graph called as a tool
# (2 and 3 can happen in a "supervisor w/ tools" multi-agent architecture)
except GraphInterrupt as e:
raise e
except Exception as e:
if not self.handle_tool_errors:
if isinstance(self.handle_tool_errors, tuple):
handled_types: tuple = self.handle_tool_errors
elif callable(self.handle_tool_errors):
handled_types = _infer_handled_types(self.handle_tool_errors)
else:
# default behavior is catching all exceptions
handled_types = (Exception,)
# Unhandled
if not self.handle_tool_errors or not isinstance(e, handled_types):
raise e
content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
return ToolMessage(content, name=call["name"], tool_call_id=call["id"])
# Handled
else:
content = _handle_tool_error(e, flag=self.handle_tool_errors)
return ToolMessage(
content=content, name=call["name"], tool_call_id=call["id"], status="error"
)
def _parse_input(
self,
@@ -214,10 +349,10 @@ class ToolNode(RunnableCallable):
if isinstance(input, list):
output_type = "list"
message: AnyMessage = input[-1]
elif isinstance(input, dict) and (messages := input.get("messages", [])):
elif isinstance(input, dict) and (messages := input.get(self.messages_key, [])):
output_type = "dict"
message = messages[-1]
elif messages := getattr(input, "messages", None):
elif messages := getattr(input, self.messages_key, None):
# Assume dataclass-like state that can coerce from dict
output_type = "dict"
message = messages[-1]
@@ -238,7 +373,9 @@ class ToolNode(RunnableCallable):
requested_tool=requested_tool,
available_tools=", ".join(self.tools_by_name.keys()),
)
return ToolMessage(content, name=requested_tool, tool_call_id=call["id"])
return ToolMessage(
content, name=requested_tool, tool_call_id=call["id"], status="error"
)
else:
return None
@@ -256,10 +393,10 @@ class ToolNode(RunnableCallable):
required_fields = list(state_args.values())
if (
len(required_fields) == 1
and required_fields[0] == "messages"
and required_fields[0] == self.messages_key
or required_fields[0] is None
):
input = {"messages": input}
input = {self.messages_key: input}
else:
err_msg = (
f"Invalid input to ToolNode. Tool {tool_call['name']} requires "
@@ -325,6 +462,7 @@ class ToolNode(RunnableCallable):
def tools_condition(
state: Union[list[AnyMessage], dict[str, Any], BaseModel],
messages_key: str = "messages",
) -> Literal["tools", "__end__"]:
"""Use in the conditional_edge to route to the ToolNode if the last message
@@ -377,9 +515,9 @@ def tools_condition(
"""
if isinstance(state, list):
ai_message = state[-1]
elif isinstance(state, dict) and (messages := state.get("messages", [])):
elif isinstance(state, dict) and (messages := state.get(messages_key, [])):
ai_message = messages[-1]
elif messages := getattr(state, "messages", []):
elif messages := getattr(state, messages_key, []):
ai_message = messages[-1]
else:
raise ValueError(f"No messages found in input state to tool_edge: {state}")
+42 -14
View File
@@ -25,7 +25,6 @@ from uuid import UUID, uuid5
from langchain_core.globals import get_debug
from langchain_core.runnables import (
Runnable,
RunnableSequence,
)
from langchain_core.runnables.base import Input, Output
@@ -34,6 +33,7 @@ from langchain_core.runnables.config import (
get_async_callback_manager_for_config,
get_callback_manager_for_config,
)
from langchain_core.runnables.graph import Graph
from langchain_core.runnables.utils import (
ConfigurableFieldSpec,
get_unique_config_specs,
@@ -67,7 +67,12 @@ from langgraph.constants import (
NS_END,
NS_SEP,
)
from langgraph.errors import GraphRecursionError, InvalidUpdateError
from langgraph.errors import (
ErrorCode,
GraphRecursionError,
InvalidUpdateError,
create_error_message,
)
from langgraph.managed.base import ManagedValueSpec
from langgraph.pregel.algo import (
PregelTaskWrites,
@@ -81,6 +86,7 @@ from langgraph.pregel.io import read_channels
from langgraph.pregel.loop import AsyncPregelLoop, StreamProtocol, SyncPregelLoop
from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
from langgraph.pregel.messages import StreamMessagesHandler
from langgraph.pregel.protocol import PregelProtocol
from langgraph.pregel.read import PregelNode
from langgraph.pregel.retry import RetryPolicy
from langgraph.pregel.runner import PregelRunner
@@ -164,15 +170,17 @@ class Channel:
return ChannelWrite(
[ChannelWriteEntry(c) for c in channels]
+ [
ChannelWriteEntry(k, mapper=v)
if callable(v)
else ChannelWriteEntry(k, value=v)
(
ChannelWriteEntry(k, mapper=v)
if callable(v)
else ChannelWriteEntry(k, value=v)
)
for k, v in kwargs.items()
]
)
class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
class Pregel(PregelProtocol):
nodes: dict[str, PregelNode]
channels: dict[str, Union[BaseChannel, ManagedValueSpec]]
@@ -252,6 +260,16 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
if auto_validate:
self.validate()
def get_graph(
self, config: RunnableConfig | None = None, *, xray: int | bool = False
) -> Graph:
raise NotImplementedError
async def aget_graph(
self, config: RunnableConfig | None = None, *, xray: int | bool = False
) -> Graph:
raise NotImplementedError
def copy(self, update: dict[str, Any] | None = None) -> Self:
attrs = {**self.__dict__, **(update or {})}
return self.__class__(**attrs)
@@ -1291,6 +1309,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
return waiter
else:
return waiter
else:
get_waiter = None # type: ignore[assignment]
# Similarly to Bulk Synchronous Parallel / Pregel model
@@ -1316,11 +1335,15 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
yield from output()
# handle exit
if loop.status == "out_of_steps":
raise GraphRecursionError(
f"Recursion limit of {config['recursion_limit']} reached "
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
msg = create_error_message(
message=(
f"Recursion limit of {config['recursion_limit']} reached "
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
),
error_code=ErrorCode.GRAPH_RECURSION_LIMIT,
)
raise GraphRecursionError(msg)
# set final channel values as run output
run_manager.on_chain_end(loop.output)
except BaseException as e:
@@ -1495,6 +1518,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
def get_waiter() -> asyncio.Task[None]:
return aioloop.create_task(stream.wait())
else:
get_waiter = None # type: ignore[assignment]
# Similarly to Bulk Synchronous Parallel / Pregel model
@@ -1522,11 +1546,15 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
yield o
# handle exit
if loop.status == "out_of_steps":
raise GraphRecursionError(
f"Recursion limit of {config['recursion_limit']} reached "
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
msg = create_error_message(
message=(
f"Recursion limit of {config['recursion_limit']} reached "
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
),
error_code=ErrorCode.GRAPH_RECURSION_LIMIT,
)
raise GraphRecursionError(msg)
# set final channel values as run output
await run_manager.on_chain_end(loop.output)
except BaseException as e:
+27 -9
View File
@@ -86,19 +86,21 @@ class BackgroundExecutor(ContextManager):
exc_value: Optional[BaseException],
traceback: Optional[TracebackType],
) -> Optional[bool]:
# copy the tasks as done() callback may modify the dict
tasks = self.tasks.copy()
# cancel all tasks that should be cancelled
for task, (cancel, _) in self.tasks.items():
for task, (cancel, _) in tasks.items():
if cancel:
task.cancel()
# wait for all tasks to finish
if tasks := {t for t in self.tasks if not t.done()}:
concurrent.futures.wait(tasks)
if pending := {t for t in tasks if not t.done()}:
concurrent.futures.wait(pending)
# shutdown the executor
self.stack.__exit__(exc_type, exc_value, traceback)
# re-raise the first exception that occurred in a task
if exc_type is None:
# if there's already an exception being raised, don't raise another one
for task, (_, reraise) in self.tasks.items():
for task, (_, reraise) in tasks.items():
if not reraise:
continue
try:
@@ -116,11 +118,17 @@ class AsyncBackgroundExecutor(AsyncContextManager):
- re-raises the first exception from tasks with `__reraise_on_exit__=True`
ignoring CancelledError"""
def __init__(self) -> None:
def __init__(self, config: RunnableConfig) -> None:
self.context_not_supported = sys.version_info < (3, 11)
self.tasks: dict[asyncio.Task, tuple[bool, bool]] = {}
self.sentinel = object()
self.loop = asyncio.get_running_loop()
if max_concurrency := config.get("max_concurrency"):
self.semaphore: Optional[asyncio.Semaphore] = asyncio.Semaphore(
max_concurrency
)
else:
self.semaphore = None
def submit( # type: ignore[valid-type]
self,
@@ -132,6 +140,8 @@ class AsyncBackgroundExecutor(AsyncContextManager):
**kwargs: P.kwargs,
) -> asyncio.Task[T]:
coro = cast(Coroutine[None, None, T], fn(*args, **kwargs))
if self.semaphore:
coro = gated(self.semaphore, coro)
if self.context_not_supported:
task = self.loop.create_task(coro, name=__name__)
else:
@@ -161,17 +171,19 @@ class AsyncBackgroundExecutor(AsyncContextManager):
exc_value: Optional[BaseException],
traceback: Optional[TracebackType],
) -> None:
# copy the tasks as done() callback may modify the dict
tasks = self.tasks.copy()
# cancel all tasks that should be cancelled
for task, (cancel, _) in self.tasks.items():
for task, (cancel, _) in tasks.items():
if cancel:
task.cancel(self.sentinel)
# wait for all tasks to finish
if self.tasks:
await asyncio.wait(self.tasks)
if tasks:
await asyncio.wait(tasks)
# if there's already an exception being raised, don't raise another one
if exc_type is None:
# re-raise the first exception that occurred in a task
for task, (_, reraise) in self.tasks.items():
for task, (_, reraise) in tasks.items():
if not reraise:
continue
try:
@@ -179,3 +191,9 @@ class AsyncBackgroundExecutor(AsyncContextManager):
raise exc
except asyncio.CancelledError:
pass
async def gated(semaphore: asyncio.Semaphore, coro: Coroutine[None, None, T]) -> T:
"""A coroutine that waits for a semaphore before running another coroutine."""
async with semaphore:
return await coro
+18 -8
View File
@@ -209,7 +209,11 @@ class PregelLoop(LoopProtocol):
)
if check_subgraphs and self.is_nested and self.checkpointer is not None:
if self.config[CONF][CONFIG_KEY_CHECKPOINT_NS] in _SEEN_CHECKPOINT_NS:
raise MultipleSubgraphsError
raise MultipleSubgraphsError(
"Multiple subgraphs called inside the same node\n\n"
"Troubleshooting URL: https://python.langchain.com/docs"
"/troubleshooting/errors/MULTIPLE_SUBGRAPHS/"
)
else:
_SEEN_CHECKPOINT_NS.add(self.config[CONF][CONFIG_KEY_CHECKPOINT_NS])
if (
@@ -281,9 +285,11 @@ class PregelLoop(LoopProtocol):
print_step_writes(
self.step,
writes,
[self.stream_keys]
if isinstance(self.stream_keys, str)
else self.stream_keys,
(
[self.stream_keys]
if isinstance(self.stream_keys, str)
else self.stream_keys
),
)
# all tasks have finished
mv_writes = apply_writes(
@@ -493,9 +499,11 @@ class PregelLoop(LoopProtocol):
print_step_checkpoint(
metadata,
self.channels,
[self.stream_keys]
if isinstance(self.stream_keys, str)
else self.stream_keys,
(
[self.stream_keys]
if isinstance(self.stream_keys, str)
else self.stream_keys
),
)
# create new checkpoint
self.checkpoint = create_checkpoint(self.checkpoint, self.channels, self.step)
@@ -839,7 +847,9 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
else []
)
self.submit = await self.stack.enter_async_context(AsyncBackgroundExecutor())
self.submit = await self.stack.enter_async_context(
AsyncBackgroundExecutor(self.config)
)
self.channels, self.managed = await self.stack.enter_async_context(
AsyncChannelsManager(self.specs, self.checkpoint, self)
)
+7 -3
View File
@@ -17,7 +17,7 @@ from langchain_core.messages import BaseMessage
from langchain_core.outputs import ChatGenerationChunk, LLMResult
from langchain_core.tracers._streaming import T, _StreamingCallbackHandler
from langgraph.constants import NS_SEP
from langgraph.constants import NS_SEP, TAG_HIDDEN, TAG_NOSTREAM
from langgraph.pregel.loop import StreamChunk
Meta = tuple[tuple[str, ...], dict[str, Any]]
@@ -63,7 +63,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
metadata: Optional[dict[str, Any]] = None,
**kwargs: Any,
) -> Any:
if metadata:
if metadata and (not tags or TAG_NOSTREAM not in tags):
self.metadata[run_id] = (
tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)),
metadata,
@@ -114,7 +114,11 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
metadata: Optional[Dict[str, Any]] = None,
**kwargs: Any,
) -> Any:
if metadata and kwargs.get("name") == metadata.get("langgraph_node"):
if (
metadata
and kwargs.get("name") == metadata.get("langgraph_node")
and (not tags or TAG_HIDDEN not in tags)
):
self.metadata[run_id] = (
tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)),
metadata,
+18 -13
View File
@@ -1,27 +1,29 @@
from abc import ABC, abstractmethod
from typing import (
Any,
AsyncIterator,
Iterator,
Optional,
Protocol,
Sequence,
Union,
runtime_checkable,
)
from langchain_core.runnables import RunnableConfig
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.runnables.graph import Graph as DrawableGraph
from typing_extensions import Self
from langgraph.pregel.types import All, StateSnapshot, StreamMode
@runtime_checkable
class PregelProtocol(Protocol):
class PregelProtocol(
Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]], ABC
):
@abstractmethod
def with_config(
self, config: Optional[RunnableConfig] = None, **kwargs: Any
) -> Self: ...
@abstractmethod
def get_graph(
self,
config: Optional[RunnableConfig] = None,
@@ -29,6 +31,7 @@ class PregelProtocol(Protocol):
xray: Union[int, bool] = False,
) -> DrawableGraph: ...
@abstractmethod
async def aget_graph(
self,
config: Optional[RunnableConfig] = None,
@@ -36,22 +39,17 @@ class PregelProtocol(Protocol):
xray: Union[int, bool] = False,
) -> DrawableGraph: ...
def get_subgraphs(
self, namespace: Optional[str] = None, recurse: bool = False
) -> Iterator[tuple[str, "PregelProtocol"]]: ...
def aget_subgraphs(
self, namespace: Optional[str] = None, recurse: bool = False
) -> AsyncIterator[tuple[str, "PregelProtocol"]]: ...
@abstractmethod
def get_state(
self, config: RunnableConfig, *, subgraphs: bool = False
) -> StateSnapshot: ...
@abstractmethod
async def aget_state(
self, config: RunnableConfig, *, subgraphs: bool = False
) -> StateSnapshot: ...
@abstractmethod
def get_state_history(
self,
config: RunnableConfig,
@@ -61,6 +59,7 @@ class PregelProtocol(Protocol):
limit: Optional[int] = None,
) -> Iterator[StateSnapshot]: ...
@abstractmethod
def aget_state_history(
self,
config: RunnableConfig,
@@ -70,6 +69,7 @@ class PregelProtocol(Protocol):
limit: Optional[int] = None,
) -> AsyncIterator[StateSnapshot]: ...
@abstractmethod
def update_state(
self,
config: RunnableConfig,
@@ -77,6 +77,7 @@ class PregelProtocol(Protocol):
as_node: Optional[str] = None,
) -> RunnableConfig: ...
@abstractmethod
async def aupdate_state(
self,
config: RunnableConfig,
@@ -84,6 +85,7 @@ class PregelProtocol(Protocol):
as_node: Optional[str] = None,
) -> RunnableConfig: ...
@abstractmethod
def stream(
self,
input: Union[dict[str, Any], Any],
@@ -95,6 +97,7 @@ class PregelProtocol(Protocol):
subgraphs: bool = False,
) -> Iterator[Union[dict[str, Any], Any]]: ...
@abstractmethod
def astream(
self,
input: Union[dict[str, Any], Any],
@@ -106,6 +109,7 @@ class PregelProtocol(Protocol):
subgraphs: bool = False,
) -> AsyncIterator[Union[dict[str, Any], Any]]: ...
@abstractmethod
def invoke(
self,
input: Union[dict[str, Any], Any],
@@ -115,6 +119,7 @@ class PregelProtocol(Protocol):
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
) -> Union[dict[str, Any], Any]: ...
@abstractmethod
async def ainvoke(
self,
input: Union[dict[str, Any], Any],
+134 -77
View File
@@ -2,6 +2,7 @@ from typing import (
Any,
AsyncIterator,
Iterator,
Literal,
Optional,
Sequence,
Union,
@@ -9,7 +10,7 @@ from typing import (
)
import orjson
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.runnables import RunnableConfig
from langchain_core.runnables.graph import (
Edge as DrawableEdge,
)
@@ -19,7 +20,6 @@ from langchain_core.runnables.graph import (
from langchain_core.runnables.graph import (
Node as DrawableNode,
)
from langchain_core.runnables.schema import StandardStreamEvent, StreamEvent
from langgraph_sdk.client import (
LangGraphClient,
SyncLangGraphClient,
@@ -27,32 +27,45 @@ from langgraph_sdk.client import (
get_sync_client,
)
from langgraph_sdk.schema import Checkpoint, ThreadState
from langgraph_sdk.schema import StreamMode as StreamModeSDK
from typing_extensions import Self
from langgraph.checkpoint.base import CheckpointMetadata
from langgraph.constants import INTERRUPT
from langgraph.errors import GraphInterrupt
from langgraph.pregel.protocol import PregelProtocol
from langgraph.pregel.types import All, PregelTask, StateSnapshot, StreamMode
from langgraph.types import Interrupt
from langgraph.utils.config import merge_configs
class RemoteGraph(PregelProtocol, Runnable):
class RemoteException(Exception):
"""Exception raised when an error occurs in the remote graph."""
pass
class RemoteGraph(PregelProtocol):
name: str
def __init__(
self,
graph_id: str,
config: Optional[RunnableConfig] = None,
name: str, # graph_id
/,
*,
url: Optional[str] = None,
api_key: Optional[str] = None,
headers: Optional[dict[str, str]] = None,
client: Optional[LangGraphClient] = None,
sync_client: Optional[SyncLangGraphClient] = None,
config: Optional[RunnableConfig] = None,
):
"""Specify `url`, `api_key`, and/or `headers` to create default sync and async clients.
If `client` or `sync_client` are provided, they will be used instead of the default clients.
See `LangGraphClient` and `SyncLangGraphClient` for details on the default clients.
"""
self.graph_id = graph_id
self.name = name
self.config = config
self.client = client or get_client(url=url, api_key=api_key, headers=headers)
self.sync_client = sync_client or get_sync_client(
@@ -61,7 +74,7 @@ class RemoteGraph(PregelProtocol, Runnable):
def copy(self, update: dict[str, Any]) -> Self:
attrs = {**self.__dict__, **update}
return self.__class__(**attrs)
return self.__class__(attrs.pop("name"), **attrs)
def with_config(
self, config: Optional[RunnableConfig] = None, **kwargs: Any
@@ -91,7 +104,7 @@ class RemoteGraph(PregelProtocol, Runnable):
xray: Union[int, bool] = False,
) -> DrawableGraph:
graph = self.sync_client.assistants.get_graph(
assistant_id=self.graph_id,
assistant_id=self.name,
xray=xray,
)
return DrawableGraph(
@@ -106,7 +119,7 @@ class RemoteGraph(PregelProtocol, Runnable):
xray: Union[int, bool] = False,
) -> DrawableGraph:
graph = await self.client.assistants.get_graph(
assistant_id=self.graph_id,
assistant_id=self.name,
xray=xray,
)
return DrawableGraph(
@@ -114,30 +127,6 @@ class RemoteGraph(PregelProtocol, Runnable):
edges=[DrawableEdge(**edge) for edge in graph["edges"]],
)
def get_subgraphs(
self, namespace: Optional[str] = None, recurse: bool = False
) -> Iterator[tuple[str, "PregelProtocol"]]:
subgraphs = self.sync_client.assistants.get_subgraphs(
assistant_id=self.graph_id,
namespace=namespace,
recurse=recurse,
)
for namespace, graph_schema in subgraphs.items():
remote_subgraph = self.copy({"graph_id": graph_schema["graph_id"]})
yield (namespace, remote_subgraph)
async def aget_subgraphs(
self, namespace: Optional[str] = None, recurse: bool = False
) -> AsyncIterator[tuple[str, "PregelProtocol"]]:
subgraphs = await self.client.assistants.get_subgraphs(
assistant_id=self.graph_id,
namespace=namespace,
recurse=recurse,
)
for namespace, graph_schema in subgraphs.items():
remote_subgraph = self.copy({"graph_id": graph_schema["graph_id"]})
yield (namespace, remote_subgraph)
def _create_state_snapshot(self, state: ThreadState) -> StateSnapshot:
tasks = []
for task in state["tasks"]:
@@ -250,7 +239,11 @@ class RemoteGraph(PregelProtocol, Runnable):
if k not in reserved_configurable_keys and not k.startswith("__pregel_")
}
return {"configurable": new_configurable}
return {
"tags": config.get("tags") or [],
"metadata": config.get("metadata") or {},
"configurable": new_configurable,
}
def get_state(
self, config: RunnableConfig, *, subgraphs: bool = False
@@ -348,6 +341,37 @@ class RemoteGraph(PregelProtocol, Runnable):
)
return self._get_config(response["checkpoint"])
def _get_stream_modes(
self,
stream_mode: Optional[Union[StreamMode, list[StreamMode]]],
default: StreamMode = "updates",
) -> tuple[list[StreamModeSDK], bool, bool]:
"""Return a tuple of the final list of stream modes sent to the
remote graph and a boolean flag indicating if stream mode 'updates'
was present in the original list of stream modes.
'updates' mode is added to the list of stream modes so that interrupts
can be detected in the remote graph.
"""
updated_stream_modes: list[StreamMode] = []
req_updates = False
req_single = True
# coerce to list, or add default stream mode
if stream_mode:
if isinstance(stream_mode, str):
updated_stream_modes.append(stream_mode)
else:
req_single = False
updated_stream_modes.extend(stream_mode)
else:
updated_stream_modes.append(default)
# add 'updates' mode if not present
if "updates" in updated_stream_modes:
req_updates = True
else:
updated_stream_modes.append("updates")
return (updated_stream_modes, req_updates, req_single)
def stream(
self,
input: Union[dict[str, Any], Any],
@@ -360,18 +384,40 @@ class RemoteGraph(PregelProtocol, Runnable):
) -> Iterator[Union[dict[str, Any], Any]]:
merged_config = merge_configs(self.config, config)
sanitized_config = self._sanitize_config(merged_config)
stream_modes, req_updates, req_single = self._get_stream_modes(stream_mode)
for chunk in self.sync_client.runs.stream(
thread_id=sanitized_config["configurable"]["thread_id"],
assistant_id=self.graph_id,
thread_id=sanitized_config["configurable"].get("thread_id"),
assistant_id=self.name,
input=input,
config=sanitized_config,
stream_mode=stream_mode, # type: ignore
interrupt_before=interrupt_before, # type: ignore
interrupt_after=interrupt_after, # type: ignore
stream_mode=stream_modes,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
stream_subgraphs=subgraphs,
if_not_exists="create",
):
yield chunk
if chunk.event.startswith("updates"):
if isinstance(chunk.data, dict) and INTERRUPT in chunk.data:
raise GraphInterrupt(chunk.data[INTERRUPT])
if not req_updates:
continue
elif chunk.event.startswith("error"):
raise RemoteException(chunk.data)
if subgraphs:
if "|" in chunk.event:
mode, ns_ = chunk.event.split("|", 1)
ns = tuple(ns_.split("|"))
else:
mode, ns = chunk.event, ()
if req_single:
yield ns, chunk.data
else:
yield ns, mode, chunk.data
elif req_single:
yield chunk.data
else:
yield chunk
async def astream(
self,
@@ -385,47 +431,56 @@ class RemoteGraph(PregelProtocol, Runnable):
) -> AsyncIterator[Union[dict[str, Any], Any]]:
merged_config = merge_configs(self.config, config)
sanitized_config = self._sanitize_config(merged_config)
stream_modes, req_updates, req_single = self._get_stream_modes(stream_mode)
async for chunk in self.client.runs.stream(
thread_id=sanitized_config["configurable"]["thread_id"],
assistant_id=self.graph_id,
thread_id=sanitized_config["configurable"].get("thread_id"),
assistant_id=self.name,
input=input,
config=sanitized_config,
stream_mode=stream_mode if stream_mode else "values", # type: ignore
interrupt_before=interrupt_before, # type: ignore
interrupt_after=interrupt_after, # type: ignore
stream_mode=stream_modes,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
stream_subgraphs=subgraphs,
if_not_exists="create",
):
yield chunk
if chunk.event.startswith("updates"):
if isinstance(chunk.data, dict) and INTERRUPT in chunk.data:
raise GraphInterrupt(chunk.data[INTERRUPT])
if not req_updates:
continue
elif chunk.event.startswith("error"):
raise RemoteException(chunk.data)
if subgraphs:
if "|" in chunk.event:
mode, ns_ = chunk.event.split("|", 1)
ns = tuple(ns_.split("|"))
else:
mode, ns = chunk.event, ()
if req_single:
yield ns, chunk.data
else:
yield ns, mode, chunk.data
elif req_single:
yield chunk.data
else:
yield chunk
async def astream_events(
self,
input: Any,
config: Optional[RunnableConfig] = None,
*,
version: Literal["v1", "v2"],
include_names: Optional[Sequence[All]] = None,
include_types: Optional[Sequence[All]] = None,
include_tags: Optional[Sequence[All]] = None,
exclude_names: Optional[Sequence[All]] = None,
exclude_types: Optional[Sequence[All]] = None,
exclude_tags: Optional[Sequence[All]] = None,
**kwargs: Any,
) -> AsyncIterator[StreamEvent]:
merged_config = merge_configs(self.config, config)
sanitized_config = self._sanitize_config(merged_config)
# manually add 'events' to stream modes list
stream_mode: list[str] = kwargs.get("stream_mode", [])
if "events" not in stream_mode:
stream_mode.append("events")
async for chunk in self.client.runs.stream(
thread_id=sanitized_config["configurable"]["thread_id"],
assistant_id=self.graph_id,
input=input,
config=sanitized_config,
stream_mode=stream_mode, # type: ignore
interrupt_before=kwargs.get("interrupt_before"),
interrupt_after=kwargs.get("interrupt_after"),
stream_subgraphs=kwargs.get("subgraphs", False),
):
yield StandardStreamEvent(
event=chunk.event,
data=chunk.data,
)
) -> AsyncIterator[dict[str, Any]]:
raise NotImplementedError
def invoke(
self,
@@ -439,12 +494,13 @@ class RemoteGraph(PregelProtocol, Runnable):
sanitized_config = self._sanitize_config(merged_config)
return self.sync_client.runs.wait(
thread_id=sanitized_config["configurable"]["thread_id"],
assistant_id=self.graph_id,
thread_id=sanitized_config["configurable"].get("thread_id"),
assistant_id=self.name,
input=input,
config=sanitized_config,
interrupt_before=interrupt_before, # type: ignore
interrupt_after=interrupt_after, # type: ignore
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
if_not_exists="create",
)
async def ainvoke(
@@ -459,10 +515,11 @@ class RemoteGraph(PregelProtocol, Runnable):
sanitized_config = self._sanitize_config(merged_config)
return await self.client.runs.wait(
thread_id=sanitized_config["configurable"]["thread_id"],
assistant_id=self.graph_id,
thread_id=sanitized_config["configurable"].get("thread_id"),
assistant_id=self.name,
input=input,
config=sanitized_config,
interrupt_before=interrupt_before, # type: ignore
interrupt_after=interrupt_after, # type: ignore
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
if_not_exists="create",
)
+3 -2
View File
@@ -4,6 +4,7 @@ from langchain_core.runnables import RunnableLambda, RunnableSequence
from langchain_core.runnables.utils import get_function_nonlocals
from langgraph.checkpoint.base import ChannelVersions
from langgraph.pregel.protocol import PregelProtocol
from langgraph.utils.runnable import Runnable, RunnableCallable, RunnableSeq
@@ -32,9 +33,9 @@ def find_subgraph_pregel(candidate: Runnable) -> Optional[Runnable]:
for c in candidates:
if (
isinstance(c, Pregel)
isinstance(c, PregelProtocol)
# subgraphs that disabled checkpointing are not considered
and c.checkpointer is not False
and (not isinstance(c, Pregel) or c.checkpointer is not False)
):
return c
elif isinstance(c, RunnableSequence) or isinstance(c, RunnableSeq):
+81 -2
View File
@@ -564,6 +564,67 @@ files = [
{file = "defusedxml-0.7.1.tar.gz", hash = "sha256:1bb3032db185915b62d7c6209c5a8792be6a32ab2fedacc84e01b52c51aa3e69"},
]
[[package]]
name = "duckdb"
version = "1.1.2"
description = "DuckDB in-process database"
optional = false
python-versions = ">=3.7.0"
files = [
{file = "duckdb-1.1.2-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:91e7f99cf5cab1d26f92cb014429153497d805e79689baa44f4c4585a8cb243f"},
{file = "duckdb-1.1.2-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:0107de622fe208142a1108263a03c43956048dcc99be3702d8e5d2aeaf99554c"},
{file = "duckdb-1.1.2-cp310-cp310-macosx_12_0_x86_64.whl", hash = "sha256:8a09610f780857677725897856f8cdf3cafd8a991f871e6cb8ba88b2dbc8d737"},
{file = "duckdb-1.1.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c0f0ddac0482f0f3fece54d720d13819e82ae26c01a939ffa66a87be53f7f665"},
{file = "duckdb-1.1.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:84103373e818758dfa361d27781d0f096553843c5ffb9193260a0786c5248270"},
{file = "duckdb-1.1.2-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bfdfd23e2bf58014ad0673973bd0ed88cd048dfe8e82420814a71d7d52ef2288"},
{file = "duckdb-1.1.2-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:25889e6e29b87047b1dd56385ac08156e4713c59326cc6fff89657d01b2c417b"},
{file = "duckdb-1.1.2-cp310-cp310-win_amd64.whl", hash = "sha256:312570fa5277c3079de18388b86c2d87cbe1044838bb152b235c0227581d5d42"},
{file = "duckdb-1.1.2-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:568439ea4fce8cb72ec1f767cd510686a9e7e29a011fc7c56d990059a6e94e48"},
{file = "duckdb-1.1.2-cp311-cp311-macosx_12_0_universal2.whl", hash = "sha256:74974f2d7210623a5d61b1fb0cb589c6e5ffcbf7dbb757a04c5ba24adcfc8cac"},
{file = "duckdb-1.1.2-cp311-cp311-macosx_12_0_x86_64.whl", hash = "sha256:e26422a3358c816d764639070945b73eef55d1b4df990989e3492c85ef725c21"},
{file = "duckdb-1.1.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:87e972bd452eeeab197fe39dcaeecdb7c264b1f75a0ee67e532e235fe45b84df"},
{file = "duckdb-1.1.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9a6b73e70b73c8df85da383f6e557c03cad5c877868b9a7e41715761e8166c1e"},
{file = "duckdb-1.1.2-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:623cb1952466aae5907af84107bcdec25a5ca021a8b6441e961f41edc724f6f2"},
{file = "duckdb-1.1.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:d9fc0b550f96901fa7e76dc70a13f6477ad3e18ef1cb21d414c3a5569de3f27e"},
{file = "duckdb-1.1.2-cp311-cp311-win_amd64.whl", hash = "sha256:181edb1973bd8f493bcb6ecfa035f1a592dff4667758592f300619012ba251c0"},
{file = "duckdb-1.1.2-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:83372b1b411086cac01ab2071122772fa66170b1b41ddbc37527464066083668"},
{file = "duckdb-1.1.2-cp312-cp312-macosx_12_0_universal2.whl", hash = "sha256:db37441deddfee6ac35a0c742d2f9e90e4e50b9e76d586a060d122b8fc56dada"},
{file = "duckdb-1.1.2-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:19142a77e72874aeaa6fda30aeb13612c6de5e8c60fbcc3392cea6ef0694eeaf"},
{file = "duckdb-1.1.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:099d99dd48d6e4682a3dd6233ceab73d977ebe1a87afaac54cf77c844e24514a"},
{file = "duckdb-1.1.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:be86e586ca7af7e807f72479a2b8d0983565360b19dbda4ef8a9d7b3909b8e2c"},
{file = "duckdb-1.1.2-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:578e0953e4d8ba8da0cd69fb2930c45f51ce47d213b77d8a4cd461f9c0960b87"},
{file = "duckdb-1.1.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:72b5eb5762c1a5e68849c7143f3b3747a9f15c040e34e41559f233a1569ad16f"},
{file = "duckdb-1.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:9b4c6b6a08180261d98330d97355503961a25ca31cd9ef296e0681f7895b4a2c"},
{file = "duckdb-1.1.2-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:695dcbc561374b126e86659709feadf883c9969ed718e94713edd4ba15d16619"},
{file = "duckdb-1.1.2-cp313-cp313-macosx_12_0_universal2.whl", hash = "sha256:ada29be1e889f486c6cf1f6dffd15463e748faf361f33996f2e862779edc24a9"},
{file = "duckdb-1.1.2-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:6ca722738fa9eb6218619740631de29acfdd132de6f6a6350fee5e291c2f6117"},
{file = "duckdb-1.1.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c796d33f1e5a0c8c570d22da0c0b1db8578687e427029e1ce2c8ce3f9fffa6a3"},
{file = "duckdb-1.1.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f5c0996988a70dd3bc8111d9b9aeab7e38ed1999a52607c5f1b528e362b4dd1c"},
{file = "duckdb-1.1.2-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6c37b039f6d6fed14d89450f5ccf54922b3304192d7412e12d6cc8d9e757f7a2"},
{file = "duckdb-1.1.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e8c766b87f675c76d6d17103bf6fb9fb1a9e2fcb3d9b25c28bbc634bde31223e"},
{file = "duckdb-1.1.2-cp313-cp313-win_amd64.whl", hash = "sha256:e3e6300b7ccaf64b609f4f0780a6e1d25ab8cf34cceed46e62c35b6c4c5cb63b"},
{file = "duckdb-1.1.2-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8a01fae9604a54ecbc26e7503c522311f15afbd2870e6d8f6fbef4545dfae550"},
{file = "duckdb-1.1.2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:492b1d86a696428bd3f14dc1c7c3230e2dbca8978f288be64b04a26e0e00fad5"},
{file = "duckdb-1.1.2-cp37-cp37m-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1bba58459ad897a78c4e478a097626fc266459a40338cecc68a49a8d5dc72fb7"},
{file = "duckdb-1.1.2-cp37-cp37m-win_amd64.whl", hash = "sha256:d395a3bf510bf24686821eec15802624797dcb33e8f14f8a7cc8e17d909474af"},
{file = "duckdb-1.1.2-cp38-cp38-macosx_12_0_arm64.whl", hash = "sha256:fd800f75728727fe699ed1eb22b636867cf48c9dd105ee88b977e20c89df4509"},
{file = "duckdb-1.1.2-cp38-cp38-macosx_12_0_universal2.whl", hash = "sha256:d8caaf43909e49537e26df51d80d075ae2b25a610d28ed8bd31d6ccebeaf3c65"},
{file = "duckdb-1.1.2-cp38-cp38-macosx_12_0_x86_64.whl", hash = "sha256:564166811c68d9c7f9911eb707ad32ec9c2507b98336d894fbe658b85bf1c697"},
{file = "duckdb-1.1.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:19386aa09f0d6f97634ba2972096d1c80d880176dfb0e949eadc91c98262a663"},
{file = "duckdb-1.1.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e9e8387bcc9a591ad14011ddfec0d408d1d9b1889c6c9b495a04c7016a24b9b3"},
{file = "duckdb-1.1.2-cp38-cp38-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f8c5ff4970403ed3ff0ac71fe0ce1e6be3199df9d542afc84c424b444ba4ffe8"},
{file = "duckdb-1.1.2-cp38-cp38-win_amd64.whl", hash = "sha256:9283dcca87c3260eb631a99d738fa72b8545ed45b475bc72ad254f7310e14284"},
{file = "duckdb-1.1.2-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:f87edaf20001530e63a4f7bda13b55dc3152d7171226915f2bf34e0813c8759e"},
{file = "duckdb-1.1.2-cp39-cp39-macosx_12_0_universal2.whl", hash = "sha256:efec169b3fe0b821e3207ba3e445f227d42dd62b4440ff79c37fa168a4fc5a71"},
{file = "duckdb-1.1.2-cp39-cp39-macosx_12_0_x86_64.whl", hash = "sha256:89164a2d29d56605a95ee5032aa415dd487028c4fd3e06d971497840e74c56e7"},
{file = "duckdb-1.1.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6858e10c60ff7e70e61d3dd53d2545c8b2609942e45fd6de38cd0dee52932de3"},
{file = "duckdb-1.1.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7ca967c5a57b1d0cb0fd5e539ab24110e5a59dcbedd365bb2dc80533d6e44a8d"},
{file = "duckdb-1.1.2-cp39-cp39-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4ce949f1d7999aa6a046eb64067eee41d4c5c2872ba4fa408c9947742d0c7231"},
{file = "duckdb-1.1.2-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:9ba6d1f918e6ca47a368a0c32806016405cb9beb2c245806b0ca998f569d2bdf"},
{file = "duckdb-1.1.2-cp39-cp39-win_amd64.whl", hash = "sha256:7111fd3e7b334a7be383313ce29918b7c643e4f6ef44d6d63c3ab3fa6716c114"},
{file = "duckdb-1.1.2.tar.gz", hash = "sha256:c8232861dc8ec6daa29067056d5a0e5789919f2ab22ab792787616d7cd52f02a"},
]
[[package]]
name = "exceptiongroup"
version = "1.2.1"
@@ -1264,6 +1325,24 @@ msgpack = "^1.1.0"
type = "directory"
url = "../checkpoint"
[[package]]
name = "langgraph-checkpoint-duckdb"
version = "1.0.0"
description = "Library with a DuckDB implementation of LangGraph checkpoint saver."
optional = false
python-versions = "^3.9.0,<4.0"
files = []
develop = true
[package.dependencies]
duckdb = ">=1.1.2"
langgraph-checkpoint = "^2.0.0"
orjson = ">=3.10.1"
[package.source]
type = "directory"
url = "../checkpoint-duckdb"
[[package]]
name = "langgraph-checkpoint-postgres"
version = "2.0.1"
@@ -1302,7 +1381,7 @@ url = "../checkpoint-sqlite"
[[package]]
name = "langgraph-sdk"
version = "0.1.32"
version = "0.1.33"
description = "SDK for interacting with LangGraph API"
optional = false
python-versions = "^3.9.0,<4.0"
@@ -3281,4 +3360,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
[metadata]
lock-version = "2.0"
python-versions = ">=3.9.0,<4.0"
content-hash = "fefcf32c107aa6384115fc90b5dc628ca784667a970f2390a49750e66b334f8b"
content-hash = "16d0de53dfe8b9a3ae450608c7af600d2b4f00484506897d9516c3113807ee7c"
+2 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph"
version = "0.2.38"
version = "0.2.39"
description = "Building stateful, multi-actor applications with LLMs"
authors = []
license = "MIT"
@@ -27,6 +27,7 @@ jupyter = "^1.0.0"
pytest-xdist = {extras = ["psutil"], version = "^3.6.1"}
pytest-repeat = "^0.9.3"
langgraph-checkpoint = {path = "../checkpoint", develop = true}
langgraph-checkpoint-duckdb = {path = "../checkpoint-duckdb", develop = true}
langgraph-checkpoint-sqlite = {path = "../checkpoint-sqlite", develop = true}
langgraph-checkpoint-postgres = {path = "../checkpoint-postgres", develop = true}
langgraph-sdk = {path = "../sdk-py", develop = true}
File diff suppressed because one or more lines are too long
@@ -221,19 +221,6 @@
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2
'''
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.1
dict({
'$defs': dict({
@@ -342,6 +329,127 @@
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[duckdb_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[duckdb_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[duckdb_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[memory]
'''
graph TD;
@@ -1194,6 +1302,31 @@
+---------+
'''
# ---
# name: test_weather_subgraph[duckdb_aio]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
__end__([<p>__end__</p>]):::last
__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__;
subgraph weather_graph
weather_graph_model_node --> weather_graph_weather_node;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[memory]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
+41 -2
View File
@@ -11,11 +11,14 @@ from psycopg_pool import AsyncConnectionPool, ConnectionPool
from pytest_mock import MockerFixture
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.checkpoint.duckdb import DuckDBSaver
from langgraph.checkpoint.duckdb.aio import AsyncDuckDBSaver
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from langgraph.store.base import BaseStore
from langgraph.store.duckdb import AsyncDuckDBStore, DuckDBStore
from langgraph.store.memory import InMemoryStore
from langgraph.store.postgres import AsyncPostgresStore, PostgresStore
from tests.memory_assert import MemorySaverAssertImmutable
@@ -61,6 +64,20 @@ async def _checkpointer_sqlite_aio():
yield checkpointer
@pytest.fixture(scope="function")
def checkpointer_duckdb():
with DuckDBSaver.from_conn_string(":memory:") as checkpointer:
checkpointer.setup()
yield checkpointer
@asynccontextmanager
async def _checkpointer_duckdb_aio():
async with AsyncDuckDBSaver.from_conn_string(":memory:") as checkpointer:
await checkpointer.setup()
yield checkpointer
@pytest.fixture(scope="function")
def checkpointer_postgres():
database = f"test_{uuid4().hex[:16]}"
@@ -212,6 +229,9 @@ async def awith_checkpointer(
elif checkpointer_name == "sqlite_aio":
async with _checkpointer_sqlite_aio() as checkpointer:
yield checkpointer
elif checkpointer_name == "duckdb_aio":
async with _checkpointer_duckdb_aio() as checkpointer:
yield checkpointer
elif checkpointer_name == "postgres_aio":
async with _checkpointer_postgres_aio() as checkpointer:
yield checkpointer
@@ -247,6 +267,13 @@ async def _store_postgres_aio():
await conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _store_duckdb_aio():
async with AsyncDuckDBStore.from_conn_string(":memory:") as store:
await store.setup()
yield store
@pytest.fixture(scope="function")
def store_postgres():
database = f"test_{uuid4().hex[:16]}"
@@ -264,6 +291,13 @@ def store_postgres():
conn.execute(f"DROP DATABASE {database}")
@pytest.fixture(scope="function")
def store_duckdb():
with DuckDBStore.from_conn_string(":memory:") as store:
store.setup()
yield store
@pytest.fixture(scope="function")
def store_in_memory():
yield InMemoryStore()
@@ -278,6 +312,9 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
elif store_name == "postgres_aio":
async with _store_postgres_aio() as store:
yield store
elif store_name == "duckdb_aio":
async with _store_duckdb_aio() as store:
yield store
else:
raise NotImplementedError(f"Unknown store {store_name}")
@@ -285,6 +322,7 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
ALL_CHECKPOINTERS_SYNC = [
"memory",
"sqlite",
"duckdb",
"postgres",
"postgres_pipe",
"postgres_pool",
@@ -292,6 +330,7 @@ ALL_CHECKPOINTERS_SYNC = [
ALL_CHECKPOINTERS_ASYNC = [
"memory",
"sqlite_aio",
"duckdb_aio",
"postgres_aio",
"postgres_aio_pipe",
"postgres_aio_pool",
@@ -300,5 +339,5 @@ ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [
*ALL_CHECKPOINTERS_ASYNC,
None,
]
ALL_STORES_SYNC = ["in_memory", "postgres"]
ALL_STORES_ASYNC = ["in_memory", "postgres_aio"]
ALL_STORES_SYNC = ["in_memory", "postgres", "duckdb"]
ALL_STORES_ASYNC = ["in_memory", "postgres_aio", "duckdb_aio"]
+514 -103
View File
@@ -1,5 +1,6 @@
import dataclasses
import json
from functools import partial
from typing import (
Annotated,
Any,
@@ -28,25 +29,39 @@ from langchain_core.messages import (
)
from langchain_core.outputs import ChatGeneration, ChatResult
from langchain_core.runnables import Runnable, RunnableLambda
from langchain_core.tools import BaseTool
from langchain_core.tools import BaseTool, ToolException
from langchain_core.tools import tool as dec_tool
from pydantic import BaseModel
from pydantic import BaseModel, ValidationError
from pydantic.v1 import BaseModel as BaseModelV1
from pydantic.v1 import ValidationError as ValidationErrorV1
from typing_extensions import TypedDict
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.graph import START, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode, ValidationNode, create_react_agent
from langgraph.prebuilt.tool_node import InjectedState, InjectedStore
from langgraph.checkpoint.memory import MemorySaver
from langgraph.errors import NodeInterrupt
from langgraph.graph import START, MessagesState, StateGraph, add_messages
from langgraph.prebuilt import (
ToolNode,
ValidationNode,
create_react_agent,
tools_condition,
)
from langgraph.prebuilt.tool_node import (
TOOL_CALL_ERROR_TEMPLATE,
InjectedState,
InjectedStore,
_infer_handled_types,
)
from langgraph.store.base import BaseStore
from langgraph.store.memory import InMemoryStore
from langgraph.types import Interrupt
from tests.conftest import (
ALL_CHECKPOINTERS_ASYNC,
ALL_CHECKPOINTERS_SYNC,
IS_LANGCHAIN_CORE_030_OR_GREATER,
awith_checkpointer,
)
from tests.messages import _AnyIdHumanMessage
from tests.messages import _AnyIdHumanMessage, _AnyIdToolMessage
pytestmark = pytest.mark.anyio
@@ -362,32 +377,107 @@ def test_model_with_tools(tool_style: str):
create_react_agent(model.bind_tools([tool1]), [tool2])
def test__infer_handled_types() -> None:
def handle(e): # type: ignore
return ""
def handle2(e: Exception) -> str:
return ""
def handle3(e: Union[ValueError, ToolException]) -> str:
return ""
class Handler:
def handle(self, e: ValueError) -> str:
return ""
handle4 = Handler().handle
def handle5(e: Union[Union[TypeError, ValueError], ToolException]):
return ""
expected: tuple = (Exception,)
actual = _infer_handled_types(handle)
assert expected == actual
expected = (Exception,)
actual = _infer_handled_types(handle2)
assert expected == actual
expected = (ValueError, ToolException)
actual = _infer_handled_types(handle3)
assert expected == actual
expected = (ValueError,)
actual = _infer_handled_types(handle4)
assert expected == actual
expected = (TypeError, ValueError, ToolException)
actual = _infer_handled_types(handle5)
assert expected == actual
with pytest.raises(ValueError):
def handler(e: str):
return ""
_infer_handled_types(handler)
with pytest.raises(ValueError):
def handler(e: list[Exception]):
return ""
_infer_handled_types(handler)
with pytest.raises(ValueError):
def handler(e: Union[str, int]):
return ""
_infer_handled_types(handler)
# tools for testing Too
def tool1(some_val: int, some_other_val: str) -> str:
"""Tool 1 docstring."""
if some_val == 0:
raise ValueError("Test error")
return f"{some_val} - {some_other_val}"
async def tool2(some_val: int, some_other_val: str) -> str:
"""Tool 2 docstring."""
if some_val == 0:
raise ToolException("Test error")
return f"tool2: {some_val} - {some_other_val}"
async def tool3(some_val: int, some_other_val: str) -> str:
"""Tool 3 docstring."""
return [
{"key_1": some_val, "key_2": "foo"},
{"key_1": some_other_val, "key_2": "baz"},
]
async def tool4(some_val: int, some_other_val: str) -> str:
"""Tool 4 docstring."""
return [
{"type": "image_url", "image_url": {"url": "abdc"}},
]
@dec_tool
def tool5(some_val: int):
"""Tool 5 docstring."""
raise ToolException("Test error")
tool5.handle_tool_error = "foo"
async def test_tool_node():
def tool1(some_val: int, some_other_val: str) -> str:
"""Tool 1 docstring."""
if some_val == 0:
raise ValueError("Test error")
return f"{some_val} - {some_other_val}"
async def tool2(some_val: int, some_other_val: str) -> str:
"""Tool 2 docstring."""
if some_val == 0:
raise ValueError("Test error")
return f"tool2: {some_val} - {some_other_val}"
async def tool3(some_val: int, some_other_val: str) -> str:
"""Tool 3 docstring."""
return [
{"key_1": some_val, "key_2": "foo"},
{"key_1": some_other_val, "key_2": "baz"},
]
async def tool4(some_val: int, some_other_val: str) -> str:
"""Tool 4 docstring."""
return [
{"type": "image_url", "image_url": {"url": "abdc"}},
]
result = ToolNode([tool1]).invoke(
{
"messages": [
@@ -410,31 +500,6 @@ async def test_tool_node():
assert tool_message.content == "1 - foo"
assert tool_message.tool_call_id == "some 0"
result_error = ToolNode([tool1]).invoke(
{
"messages": [
AIMessage(
"hi?",
tool_calls=[
{
"name": "tool1",
"args": {"some_val": 0, "some_other_val": "foo"},
"id": "some 0",
}
],
)
]
}
)
tool_message: ToolMessage = result_error["messages"][-1]
assert tool_message.type == "tool"
assert (
tool_message.content
== f"Error: {repr(ValueError('Test error'))}\n Please fix your mistakes."
)
assert tool_message.tool_call_id == "some 0"
result2 = await ToolNode([tool2]).ainvoke(
{
"messages": [
@@ -451,11 +516,232 @@ async def test_tool_node():
]
}
)
tool_message: ToolMessage = result2["messages"][-1]
assert tool_message.type == "tool"
assert tool_message.content == "tool2: 2 - bar"
with pytest.raises(ValueError):
# list of dicts tool content
result3 = await ToolNode([tool3]).ainvoke(
{
"messages": [
AIMessage(
"hi?",
tool_calls=[
{
"name": "tool3",
"args": {"some_val": 2, "some_other_val": "bar"},
"id": "some 2",
}
],
)
]
}
)
tool_message: ToolMessage = result3["messages"][-1]
assert tool_message.type == "tool"
assert (
tool_message.content
== '[{"key_1": 2, "key_2": "foo"}, {"key_1": "bar", "key_2": "baz"}]'
)
assert tool_message.tool_call_id == "some 2"
# list of content blocks tool content
result4 = await ToolNode([tool4]).ainvoke(
{
"messages": [
AIMessage(
"hi?",
tool_calls=[
{
"name": "tool4",
"args": {"some_val": 2, "some_other_val": "bar"},
"id": "some 3",
}
],
)
]
}
)
tool_message: ToolMessage = result4["messages"][-1]
assert tool_message.type == "tool"
assert tool_message.content == [{"type": "image_url", "image_url": {"url": "abdc"}}]
assert tool_message.tool_call_id == "some 3"
async def test_tool_node_error_handling():
def handle_all(e: Union[ValueError, ToolException, ValidationError]):
return TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
# test catching all exceptions, via:
# - handle_tool_errors = True
# - passing a tuple of all exceptions
# - passing a callable with all exceptions in the signature
for handle_tool_errors in (
True,
(ValueError, ToolException, ValidationError),
handle_all,
):
result_error = await ToolNode(
[tool1, tool2, tool3], handle_tool_errors=handle_tool_errors
).ainvoke(
{
"messages": [
AIMessage(
"hi?",
tool_calls=[
{
"name": "tool1",
"args": {"some_val": 0, "some_other_val": "foo"},
"id": "some id",
},
{
"name": "tool2",
"args": {"some_val": 0, "some_other_val": "bar"},
"id": "some other id",
},
{
"name": "tool3",
"args": {"some_val": 0},
"id": "another id",
},
],
)
]
}
)
assert all(m.type == "tool" for m in result_error["messages"])
assert all(m.status == "error" for m in result_error["messages"])
assert (
result_error["messages"][0].content
== f"Error: {repr(ValueError('Test error'))}\n Please fix your mistakes."
)
assert (
result_error["messages"][1].content
== f"Error: {repr(ToolException('Test error'))}\n Please fix your mistakes."
)
assert (
"ValidationError" in result_error["messages"][2].content
or "validation error" in result_error["messages"][2].content
)
assert result_error["messages"][0].tool_call_id == "some id"
assert result_error["messages"][1].tool_call_id == "some other id"
assert result_error["messages"][2].tool_call_id == "another id"
async def test_tool_node_error_handling_callable():
def handle_value_error(e: ValueError):
return "Value error"
def handle_tool_exception(e: ToolException):
return "Tool exception"
for handle_tool_errors in ("Value error", handle_value_error):
result_error = await ToolNode(
[tool1], handle_tool_errors=handle_tool_errors
).ainvoke(
{
"messages": [
AIMessage(
"hi?",
tool_calls=[
{
"name": "tool1",
"args": {"some_val": 0, "some_other_val": "foo"},
"id": "some id",
},
],
)
]
}
)
tool_message: ToolMessage = result_error["messages"][-1]
assert tool_message.type == "tool"
assert tool_message.status == "error"
assert tool_message.content == "Value error"
# test raising for an unhandled exception, via:
# - passing a tuple of all exceptions
# - passing a callable with all exceptions in the signature
for handle_tool_errors in ((ValueError,), handle_value_error):
with pytest.raises(ToolException) as exc_info:
await ToolNode(
[tool1, tool2], handle_tool_errors=handle_tool_errors
).ainvoke(
{
"messages": [
AIMessage(
"hi?",
tool_calls=[
{
"name": "tool1",
"args": {"some_val": 0, "some_other_val": "foo"},
"id": "some id",
},
{
"name": "tool2",
"args": {"some_val": 0, "some_other_val": "bar"},
"id": "some other id",
},
],
)
]
}
)
assert str(exc_info.value) == "Test error"
for handle_tool_errors in ((ToolException,), handle_tool_exception):
with pytest.raises(ValueError) as exc_info:
await ToolNode(
[tool1, tool2], handle_tool_errors=handle_tool_errors
).ainvoke(
{
"messages": [
AIMessage(
"hi?",
tool_calls=[
{
"name": "tool1",
"args": {"some_val": 0, "some_other_val": "foo"},
"id": "some id",
},
{
"name": "tool2",
"args": {"some_val": 0, "some_other_val": "bar"},
"id": "some other id",
},
],
)
]
}
)
assert str(exc_info.value) == "Test error"
async def test_tool_node_handle_tool_errors_false():
with pytest.raises(ValueError) as exc_info:
ToolNode([tool1], handle_tool_errors=False).invoke(
{
"messages": [
AIMessage(
"hi?",
tool_calls=[
{
"name": "tool1",
"args": {"some_val": 0, "some_other_val": "foo"},
"id": "some id",
}
],
)
]
}
)
assert str(exc_info.value) == "Test error"
with pytest.raises(ToolException):
await ToolNode([tool2], handle_tool_errors=False).ainvoke(
{
"messages": [
@@ -465,7 +751,7 @@ async def test_tool_node():
{
"name": "tool2",
"args": {"some_val": 0, "some_other_val": "bar"},
"id": "some 1",
"id": "some id",
}
],
)
@@ -473,7 +759,57 @@ async def test_tool_node():
}
)
# incorrect tool name
assert str(exc_info.value) == "Test error"
# test validation errors get raised if handle_tool_errors is False
with pytest.raises((ValidationError, ValidationErrorV1)):
ToolNode([tool1], handle_tool_errors=False).invoke(
{
"messages": [
AIMessage(
"hi?",
tool_calls=[
{
"name": "tool1",
"args": {"some_val": 0},
"id": "some id",
}
],
)
]
}
)
def test_tool_node_individual_tool_error_handling():
# test error handling on individual tools (and that it overrides overall error handling!)
result_individual_tool_error_handler = ToolNode(
[tool5], handle_tool_errors="bar"
).invoke(
{
"messages": [
AIMessage(
"hi?",
tool_calls=[
{
"name": "tool5",
"args": {"some_val": 0},
"id": "some 0",
}
],
)
]
}
)
tool_message: ToolMessage = result_individual_tool_error_handler["messages"][-1]
assert tool_message.type == "tool"
assert tool_message.status == "error"
assert tool_message.content == "foo"
assert tool_message.tool_call_id == "some 0"
def test_tool_node_incorrect_tool_name():
result_incorrect_name = ToolNode([tool1, tool2]).invoke(
{
"messages": [
@@ -490,60 +826,94 @@ async def test_tool_node():
]
}
)
tool_message: ToolMessage = result_incorrect_name["messages"][-1]
assert tool_message.type == "tool"
assert tool_message.status == "error"
assert (
tool_message.content
== "Error: tool3 is not a valid tool, try one of [tool1, tool2]."
)
assert tool_message.tool_call_id == "some 0"
# list of dicts tool content
result3 = await ToolNode([tool3]).ainvoke(
{
"messages": [
AIMessage(
"hi?",
tool_calls=[
{
"name": "tool3",
"args": {"some_val": 2, "some_other_val": "bar"},
"id": "some 0",
}
],
)
]
}
)
tool_message: ToolMessage = result3["messages"][-1]
assert tool_message.type == "tool"
assert (
tool_message.content
== '[{"key_1": 2, "key_2": "foo"}, {"key_1": "bar", "key_2": "baz"}]'
)
assert tool_message.tool_call_id == "some 0"
# list of content blocks tool content
result4 = await ToolNode([tool4]).ainvoke(
{
"messages": [
AIMessage(
"hi?",
tool_calls=[
{
"name": "tool4",
"args": {"some_val": 2, "some_other_val": "bar"},
"id": "some 0",
}
],
)
]
}
def test_tool_node_node_interrupt():
def tool_normal(some_val: int) -> str:
"""Tool docstring."""
return "normal"
def tool_interrupt(some_val: int) -> str:
"""Tool docstring."""
raise NodeInterrupt("foo")
def handle(e: NodeInterrupt):
return "handled"
for handle_tool_errors in (True, (NodeInterrupt,), "handled", handle, False):
node = ToolNode([tool_interrupt], handle_tool_errors=handle_tool_errors)
with pytest.raises(NodeInterrupt) as exc_info:
node.invoke(
{
"messages": [
AIMessage(
"hi?",
tool_calls=[
{
"name": "tool_interrupt",
"args": {"some_val": 0},
"id": "some 0",
}
],
)
]
}
)
assert exc_info.value == "foo"
# test inside react agent
model = FakeToolCallingModel(
tool_calls=[
[
ToolCall(name="tool_interrupt", args={"some_val": 0}, id="1"),
ToolCall(name="tool_normal", args={"some_val": 1}, id="2"),
],
[],
]
)
tool_message: ToolMessage = result4["messages"][-1]
assert tool_message.type == "tool"
assert tool_message.content == [{"type": "image_url", "image_url": {"url": "abdc"}}]
assert tool_message.tool_call_id == "some 0"
checkpointer = MemorySaver()
config = {"configurable": {"thread_id": "1"}}
agent = create_react_agent(
model, [tool_interrupt, tool_normal], checkpointer=checkpointer
)
result = agent.invoke({"messages": [HumanMessage("hi?")]}, config)
assert result["messages"] == [
_AnyIdHumanMessage(
content="hi?",
),
AIMessage(
content="hi?",
id="0",
tool_calls=[
{
"name": "tool_interrupt",
"args": {"some_val": 0},
"id": "1",
"type": "tool_call",
},
{
"name": "tool_normal",
"args": {"some_val": 1},
"id": "2",
"type": "tool_call",
},
],
),
]
state = agent.get_state(config)
assert state.next == ("tools",)
task = state.tasks[0]
assert task.name == "tools"
assert task.interrupts == (Interrupt(value="foo", when="during"),)
def my_function(some_val: int, some_other_val: str) -> str:
@@ -826,6 +1196,47 @@ def test_tool_node_ensure_utf8() -> None:
assert outputs[0].content == json.dumps(data, ensure_ascii=False)
def test_tool_node_messages_key() -> None:
@dec_tool
def add(a: int, b: int):
"""Adds a and b."""
return a + b
model = FakeToolCallingModel(
tool_calls=[[ToolCall(name=add.name, args={"a": 1, "b": 2}, id="test_id")]]
)
class State(TypedDict):
subgraph_messages: Annotated[list[AnyMessage], add_messages]
def call_model(state: State):
response = model.invoke(state["subgraph_messages"])
model.tool_calls = []
return {"subgraph_messages": response}
builder = StateGraph(State)
builder.add_node("agent", call_model)
builder.add_node("tools", ToolNode([add], messages_key="subgraph_messages"))
builder.add_conditional_edges(
"agent", partial(tools_condition, messages_key="subgraph_messages")
)
builder.add_edge(START, "agent")
builder.add_edge("tools", "agent")
graph = builder.compile()
result = graph.invoke({"subgraph_messages": [HumanMessage(content="hi")]})
assert result["subgraph_messages"] == [
_AnyIdHumanMessage(content="hi"),
AIMessage(
content="hi",
id="0",
tool_calls=[ToolCall(name=add.name, args={"a": 1, "b": 2}, id="test_id")],
),
_AnyIdToolMessage(content="3", name=add.name, tool_call_id="test_id"),
AIMessage(content="hi-hi-3", id="1"),
]
async def test_return_direct() -> None:
@dec_tool(return_direct=True)
def tool_return_direct(input: str) -> str:
-12
View File
@@ -4078,18 +4078,6 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
stream_mode="messages",
)
] == [
(
_AnyIdHumanMessage(
content="what is weather in sf",
),
{
"langgraph_step": 0,
"langgraph_node": "__start__",
"langgraph_triggers": ["__start__"],
"langgraph_path": ("__pregel_pull", "__start__"),
"langgraph_checkpoint_ns": AnyStr("__start__:"),
},
),
(
_AnyIdAIMessageChunk(
content="",
+47 -12
View File
@@ -1922,6 +1922,53 @@ async def test_cond_edge_after_send() -> None:
assert await graph.ainvoke(["0"]) == ["0", "1", "2", "2", "3"]
async def test_max_concurrency() -> None:
class Node:
def __init__(self, name: str):
self.name = name
setattr(self, "__name__", name)
self.currently = 0
self.max_currently = 0
async def __call__(self, state):
self.currently += 1
if self.currently > self.max_currently:
self.max_currently = self.currently
await asyncio.sleep(0.1)
self.currently -= 1
return [self.name]
async def send_to_many(state):
return [Send("2", state)] * 100
async def route_to_three(state) -> Literal["3"]:
return "3"
node2 = Node("2")
builder = StateGraph(Annotated[list, operator.add])
builder.add_node(Node("1"))
builder.add_node(node2)
builder.add_node(Node("3"))
builder.add_edge(START, "1")
builder.add_conditional_edges("1", send_to_many)
builder.add_conditional_edges("2", route_to_three)
graph = builder.compile()
assert await graph.ainvoke(["0"]) == ["0", "1", *(["2"] * 100), "3"]
assert node2.max_currently == 100
assert node2.currently == 0
node2.max_currently = 0
assert await graph.ainvoke(["0"], {"max_concurrency": 10}) == [
"0",
"1",
*(["2"] * 100),
"3",
]
assert node2.max_currently == 10
assert node2.currently == 0
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_invoke_checkpoint_three(
mocker: MockerFixture, checkpointer_name: str
@@ -3999,18 +4046,6 @@ async def test_prebuilt_tool_chat() -> None:
stream_mode="messages",
)
] == [
(
_AnyIdHumanMessage(
content="what is weather in sf",
),
{
"langgraph_step": 0,
"langgraph_node": "__start__",
"langgraph_triggers": ["__start__"],
"langgraph_path": ("__pregel_pull", "__start__"),
"langgraph_checkpoint_ns": AnyStr("__start__:"),
},
),
(
_AnyIdAIMessageChunk(
content="",
@@ -7,7 +7,9 @@ from langchain_core.runnables.graph import (
from langchain_core.runnables.graph import (
Node as DrawableNode,
)
from langgraph_sdk.schema import StreamPart
from langgraph.errors import GraphInterrupt
from langgraph.pregel.remote import RemoteGraph
from langgraph.pregel.types import StateSnapshot
@@ -15,7 +17,7 @@ from langgraph.pregel.types import StateSnapshot
def test_with_config():
# set up test
remote_pregel = RemoteGraph(
graph_id="test_graph_id",
"test_graph_id",
config={
"configurable": {
"foo": "bar",
@@ -62,7 +64,7 @@ def test_get_graph():
],
}
remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id")
remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client)
# call method / assertions
drawable_graph = remote_pregel.get_graph()
@@ -109,7 +111,7 @@ async def test_aget_graph():
],
}
remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id")
remote_pregel = RemoteGraph("test_graph_id", client=mock_async_client)
# call method / assertions
drawable_graph = await remote_pregel.aget_graph()
@@ -133,92 +135,6 @@ async def test_aget_graph():
]
def test_get_subgraphs():
# set up test
mock_sync_client = MagicMock()
mock_sync_client.assistants.get_subgraphs.return_value = {
"namespace_1": {
"graph_id": "test_graph_id_2",
"input_schema": {},
"output_schema": {},
"state_schema": {},
"config_schema": {},
},
"namespace_2": {
"graph_id": "test_graph_id_3",
"input_schema": {},
"output_schema": {},
"state_schema": {},
"config_schema": {},
},
}
remote_pregel = RemoteGraph(
sync_client=mock_sync_client, graph_id="test_graph_id_1"
)
# call method / assertions
subgraphs = list(remote_pregel.get_subgraphs())
assert len(subgraphs) == 2
subgraph_1 = subgraphs[0]
ns_1 = subgraph_1[0]
remote_pregel_1: RemoteGraph = subgraph_1[1]
assert ns_1 == "namespace_1"
assert remote_pregel_1.graph_id == "test_graph_id_2"
subgraph_2 = subgraphs[1]
ns_2 = subgraph_2[0]
remote_pregel_2: RemoteGraph = subgraph_2[1]
assert ns_2 == "namespace_2"
assert remote_pregel_2.graph_id == "test_graph_id_3"
@pytest.mark.anyio
async def test_aget_subgraphs():
# set up test
mock_async_client = AsyncMock()
mock_async_client.assistants.get_subgraphs.return_value = {
"namespace_1": {
"graph_id": "test_graph_id_2",
"input_schema": {},
"output_schema": {},
"state_schema": {},
"config_schema": {},
},
"namespace_2": {
"graph_id": "test_graph_id_3",
"input_schema": {},
"output_schema": {},
"state_schema": {},
"config_schema": {},
},
}
remote_pregel = RemoteGraph(
client=mock_async_client,
graph_id="test_graph_id_1",
)
# call method / assertions
subgraphs = []
async for subgraph in remote_pregel.aget_subgraphs():
subgraphs.append(subgraph)
assert len(subgraphs) == 2
subgraph_1 = subgraphs[0]
ns_1 = subgraph_1[0]
remote_pregel_1: RemoteGraph = subgraph_1[1]
assert ns_1 == "namespace_1"
assert remote_pregel_1.graph_id == "test_graph_id_2"
subgraph_2 = subgraphs[1]
ns_2 = subgraph_2[0]
remote_pregel_2: RemoteGraph = subgraph_2[1]
assert ns_2 == "namespace_2"
assert remote_pregel_2.graph_id == "test_graph_id_3"
def test_get_state():
# set up test
mock_sync_client = MagicMock()
@@ -238,7 +154,10 @@ def test_get_state():
}
# call method / assertions
remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id")
remote_pregel = RemoteGraph(
"test_graph_id",
sync_client=mock_sync_client,
)
config = {"configurable": {"thread_id": "thread1"}}
state_snapshot = remote_pregel.get_state(config)
@@ -285,7 +204,10 @@ async def test_aget_state():
}
# call method / assertions
remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id")
remote_pregel = RemoteGraph(
"test_graph_id",
client=mock_async_client,
)
config = {"configurable": {"thread_id": "thread1"}}
state_snapshot = await remote_pregel.aget_state(config)
@@ -336,7 +258,10 @@ def test_get_state_history():
]
# call method / assertions
remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id")
remote_pregel = RemoteGraph(
"test_graph_id",
sync_client=mock_sync_client,
)
config = {"configurable": {"thread_id": "thread1"}}
state_history_snapshot = list(
@@ -384,7 +309,10 @@ async def test_aget_state_history():
]
# call method / assertions
remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id")
remote_pregel = RemoteGraph(
"test_graph_id",
client=mock_async_client,
)
config = {"configurable": {"thread_id": "thread1"}}
state_history_snapshot = []
@@ -425,7 +353,10 @@ def test_update_state():
}
# call method / assertions
remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id")
remote_pregel = RemoteGraph(
"test_graph_id",
sync_client=mock_sync_client,
)
config = {"configurable": {"thread_id": "thread1"}}
response = remote_pregel.update_state(config, {"key": "value"})
@@ -454,7 +385,10 @@ async def test_aupdate_state():
}
# call method / assertions
remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id")
remote_pregel = RemoteGraph(
"test_graph_id",
client=mock_async_client,
)
config = {"configurable": {"thread_id": "thread1"}}
response = await remote_pregel.aupdate_state(config, {"key": "value"})
@@ -473,17 +407,100 @@ def test_stream():
# set up test
mock_sync_client = MagicMock()
mock_sync_client.runs.stream.return_value = [
StreamPart(event="values", data={"chunk": "data1"}),
StreamPart(event="values", data={"chunk": "data2"}),
StreamPart(event="values", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
StreamPart(event="updates", data={"__interrupt__": ()}),
]
# call method / assertions
remote_pregel = RemoteGraph(
"test_graph_id",
sync_client=mock_sync_client,
)
# stream modes doesn't include 'updates'
stream_parts = []
with pytest.raises(GraphInterrupt):
for stream_part in remote_pregel.stream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode="values",
):
stream_parts.append(stream_part)
assert stream_parts == [
{"chunk": "data1"},
{"chunk": "data2"},
{"chunk": "data3"},
]
# call method / assertions
remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id")
mock_sync_client.runs.stream.return_value = [
StreamPart(event="updates", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
StreamPart(event="updates", data={"__interrupt__": ()}),
]
config = {"configurable": {"thread_id": "thread_1"}}
result = list(remote_pregel.stream({"input": "data"}, config))
assert result == [{"chunk": "data1"}, {"chunk": "data2"}, {"chunk": "data3"}]
# default stream_mode is updates
stream_parts = []
with pytest.raises(GraphInterrupt):
for stream_part in remote_pregel.stream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
):
stream_parts.append(stream_part)
assert stream_parts == [
{"chunk": "data3"},
{"chunk": "data4"},
]
# list stream_mode includes mode names
stream_parts = []
with pytest.raises(GraphInterrupt):
for stream_part in remote_pregel.stream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode=["updates"],
):
stream_parts.append(stream_part)
assert stream_parts == [
("updates", {"chunk": "data3"}),
("updates", {"chunk": "data4"}),
]
# subgraphs + list modes
stream_parts = []
with pytest.raises(GraphInterrupt):
for stream_part in remote_pregel.stream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode=["updates"],
subgraphs=True,
):
stream_parts.append(stream_part)
assert stream_parts == [
((), "updates", {"chunk": "data3"}),
((), "updates", {"chunk": "data4"}),
]
# subgraphs + single mode
stream_parts = []
with pytest.raises(GraphInterrupt):
for stream_part in remote_pregel.stream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
subgraphs=True,
):
stream_parts.append(stream_part)
assert stream_parts == [
((), {"chunk": "data3"}),
((), {"chunk": "data4"}),
]
@pytest.mark.anyio
@@ -492,20 +509,142 @@ async def test_astream():
mock_async_client = MagicMock()
async_iter = MagicMock()
async_iter.__aiter__.return_value = [
{"chunk": "data1"},
{"chunk": "data2"},
{"chunk": "data3"},
StreamPart(event="values", data={"chunk": "data1"}),
StreamPart(event="values", data={"chunk": "data2"}),
StreamPart(event="values", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
StreamPart(event="updates", data={"__interrupt__": ()}),
]
mock_async_client.runs.stream.return_value = async_iter
# call method / assertions
remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id")
remote_pregel = RemoteGraph(
"test_graph_id",
client=mock_async_client,
)
config = {"configurable": {"thread_id": "thread_1"}}
chunks = []
async for chunk in remote_pregel.astream({"input": "data"}, config):
chunks.append(chunk)
assert chunks == [{"chunk": "data1"}, {"chunk": "data2"}, {"chunk": "data3"}]
# stream modes doesn't include 'updates'
stream_parts = []
with pytest.raises(GraphInterrupt):
async for stream_part in remote_pregel.astream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode="values",
):
stream_parts.append(stream_part)
assert stream_parts == [
{"chunk": "data1"},
{"chunk": "data2"},
{"chunk": "data3"},
]
async_iter = MagicMock()
async_iter.__aiter__.return_value = [
StreamPart(event="updates", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
StreamPart(event="updates", data={"__interrupt__": ()}),
]
mock_async_client.runs.stream.return_value = async_iter
# default stream_mode is updates
stream_parts = []
with pytest.raises(GraphInterrupt):
async for stream_part in remote_pregel.astream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
):
stream_parts.append(stream_part)
assert stream_parts == [
{"chunk": "data3"},
{"chunk": "data4"},
]
# list stream_mode includes mode names
stream_parts = []
with pytest.raises(GraphInterrupt):
async for stream_part in remote_pregel.astream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode=["updates"],
):
stream_parts.append(stream_part)
assert stream_parts == [
("updates", {"chunk": "data3"}),
("updates", {"chunk": "data4"}),
]
# subgraphs + list modes
stream_parts = []
with pytest.raises(GraphInterrupt):
async for stream_part in remote_pregel.astream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode=["updates"],
subgraphs=True,
):
stream_parts.append(stream_part)
assert stream_parts == [
((), "updates", {"chunk": "data3"}),
((), "updates", {"chunk": "data4"}),
]
# subgraphs + single mode
stream_parts = []
with pytest.raises(GraphInterrupt):
async for stream_part in remote_pregel.astream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
subgraphs=True,
):
stream_parts.append(stream_part)
assert stream_parts == [
((), {"chunk": "data3"}),
((), {"chunk": "data4"}),
]
async_iter = MagicMock()
async_iter.__aiter__.return_value = [
StreamPart(event="updates|my|subgraph", data={"chunk": "data3"}),
StreamPart(event="updates|hello|subgraph", data={"chunk": "data4"}),
StreamPart(event="updates|bye|subgraph", data={"__interrupt__": ()}),
]
mock_async_client.runs.stream.return_value = async_iter
# subgraphs + list modes
stream_parts = []
with pytest.raises(GraphInterrupt):
async for stream_part in remote_pregel.astream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode=["updates"],
subgraphs=True,
):
stream_parts.append(stream_part)
assert stream_parts == [
(("my", "subgraph"), "updates", {"chunk": "data3"}),
(("hello", "subgraph"), "updates", {"chunk": "data4"}),
]
# subgraphs + single mode
stream_parts = []
with pytest.raises(GraphInterrupt):
async for stream_part in remote_pregel.astream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
subgraphs=True,
):
stream_parts.append(stream_part)
assert stream_parts == [
(("my", "subgraph"), {"chunk": "data3"}),
(("hello", "subgraph"), {"chunk": "data4"}),
]
def test_invoke():
@@ -516,7 +655,10 @@ def test_invoke():
}
# call method / assertions
remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id")
remote_pregel = RemoteGraph(
"test_graph_id",
sync_client=mock_sync_client,
)
config = {"configurable": {"thread_id": "thread_1"}}
result = remote_pregel.invoke(
@@ -535,7 +677,10 @@ async def test_ainvoke():
}
# call method / assertions
remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id")
remote_pregel = RemoteGraph(
"test_graph_id",
client=mock_async_client,
)
config = {"configurable": {"thread_id": "thread_1"}}
result = await remote_pregel.ainvoke(
@@ -557,7 +702,9 @@ async def test_langgraph_cloud_integration():
client = get_client()
sync_client = get_sync_client()
remote_pregel = RemoteGraph(
client=client, sync_client=sync_client, graph_id="agent"
"agent",
client=client,
sync_client=sync_client,
)
# define graph
@@ -572,7 +719,7 @@ async def test_langgraph_cloud_integration():
"messages": [
{
"role": "human",
"content": "Hello world!",
"content": "What's the weather in SF?",
}
]
}
@@ -580,7 +727,8 @@ async def test_langgraph_cloud_integration():
# test invoke
response = app.invoke(
input,
config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}},
config={"configurable": {"thread_id": "39a6104a-34e7-4f83-929c-d9eb163003c9"}},
interrupt_before=["agent"],
)
print("response:", response["messages"][-1].content)
@@ -634,9 +782,3 @@ async def test_langgraph_cloud_integration():
remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca" # must be UUID
graph = await remote_pregel.aget_graph(xray=True)
print("graph:", graph)
# test get subgraphs
remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca" # must be UUID
async for name, pregel in remote_pregel.aget_subgraphs():
print("name:", name)
print("pregel:", pregel)
+31 -6
View File
@@ -61,6 +61,9 @@ def test_state_schema_with_type_hint():
class OutputState(TypedDict):
input_state: InputState
class FooState(InputState):
foo: str
def complete_hint(state: InputState) -> OutputState:
return {"input_state": state}
@@ -73,24 +76,46 @@ def test_state_schema_with_type_hint():
def miss_all_hint(state, config):
return {"input_state": state}
def pre_foo(_) -> FooState:
return {"foo": "bar"}
class Foo:
def __call__(self, state: FooState) -> OutputState:
assert state.pop("foo") == "bar"
return {"input_state": state}
graph = StateGraph(InputState, output=OutputState)
actions = [complete_hint, miss_first_hint, only_return_hint, miss_all_hint]
actions = [
complete_hint,
miss_first_hint,
only_return_hint,
miss_all_hint,
pre_foo,
Foo(),
]
for action in actions:
graph.add_node(action)
graph.set_entry_point(actions[0].__name__)
def get_name(action) -> str:
return getattr(action, "__name__", action.__class__.__name__)
graph.set_entry_point(get_name(actions[0]))
for i in range(len(actions) - 1):
graph.add_edge(actions[i].__name__, actions[i + 1].__name__)
graph.set_finish_point(actions[-1].__name__)
graph.add_edge(get_name(actions[i]), get_name(actions[i + 1]))
graph.set_finish_point(get_name(actions[-1]))
graph = graph.compile()
input_state = InputState(question="Hello World!")
output_state = OutputState(input_state=input_state)
foo_state = FooState(foo="bar")
for i, c in enumerate(graph.stream(input_state, stream_mode="updates")):
node_name = actions[i].__name__
assert c[node_name] == output_state
node_name = get_name(actions[i])
if node_name == get_name(pre_foo):
assert c[node_name] == foo_state
else:
assert c[node_name] == output_state
@pytest.mark.parametrize("total_", [True, False])
@@ -191,7 +191,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
step=saved.metadata["step"] + 1,
stop=saved.metadata["step"] + 2,
),
) as (channels, managed), AsyncBackgroundExecutor() as submit:
) as (channels, managed), AsyncBackgroundExecutor(msg["config"]) as submit:
if task := await asyncio.to_thread(
prepare_single_task,
msg["task"]["path"],
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@langchain/langgraph-sdk",
"version": "0.0.16",
"version": "0.0.18",
"description": "Client library for interacting with the LangGraph API",
"type": "module",
"packageManager": "yarn@1.22.19",
+1 -1
View File
@@ -10,7 +10,7 @@ type RunStatus =
| "timeout"
| "interrupted";
type ThreadStatus = "idle" | "busy" | "interrupted";
type ThreadStatus = "idle" | "busy" | "interrupted" | "error";
type MultitaskStrategy = "reject" | "interrupt" | "rollback" | "enqueue";
+91 -54
View File
@@ -39,6 +39,7 @@ from langgraph_sdk.schema import (
Cron,
DisconnectMode,
GraphSchema,
IfNotExists,
Item,
Json,
ListNamespaceResponse,
@@ -1169,18 +1170,19 @@ class RunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
stream_mode: Union[StreamMode, list[StreamMode]] = "values",
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
checkpoint_id: Optional[str] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
feedback_keys: Optional[list[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
feedback_keys: Optional[Sequence[str]] = None,
on_disconnect: Optional[DisconnectMode] = None,
webhook: Optional[str] = None,
multitask_strategy: Optional[MultitaskStrategy] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> AsyncIterator[StreamPart]: ...
@@ -1191,15 +1193,16 @@ class RunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
stream_mode: Union[StreamMode, list[StreamMode]] = "values",
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
feedback_keys: Optional[list[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
feedback_keys: Optional[Sequence[str]] = None,
on_disconnect: Optional[DisconnectMode] = None,
on_completion: Optional[OnCompletionBehavior] = None,
if_not_exists: Optional[IfNotExists] = None,
webhook: Optional[str] = None,
after_seconds: Optional[int] = None,
) -> AsyncIterator[StreamPart]: ...
@@ -1210,19 +1213,20 @@ class RunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
stream_mode: Union[StreamMode, list[StreamMode]] = "values",
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
checkpoint_id: Optional[str] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
feedback_keys: Optional[list[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
feedback_keys: Optional[Sequence[str]] = None,
on_disconnect: Optional[DisconnectMode] = None,
on_completion: Optional[OnCompletionBehavior] = None,
webhook: Optional[str] = None,
multitask_strategy: Optional[MultitaskStrategy] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> AsyncIterator[StreamPart]:
"""Create a run and stream the results.
@@ -1248,6 +1252,8 @@ class RunsClient:
webhook: Webhook to call after LangGraph API call is done.
multitask_strategy: Multitask strategy to use.
Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
if_not_exists: How to handle missing thread. Defaults to 'reject'.
Must be either 'reject' (raise error if missing), or 'create' (create new thread).
after_seconds: The number of seconds to wait before starting the run.
Use to schedule future runs.
@@ -1293,6 +1299,7 @@ class RunsClient:
"checkpoint": checkpoint,
"checkpoint_id": checkpoint_id,
"multitask_strategy": multitask_strategy,
"if_not_exists": if_not_exists,
"on_disconnect": on_disconnect,
"on_completion": on_completion,
"after_seconds": after_seconds,
@@ -1313,14 +1320,15 @@ class RunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
stream_mode: Union[StreamMode, list[StreamMode]] = "values",
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
webhook: Optional[str] = None,
on_completion: Optional[OnCompletionBehavior] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> Run: ...
@@ -1331,16 +1339,17 @@ class RunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
stream_mode: Union[StreamMode, list[StreamMode]] = "values",
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
checkpoint_id: Optional[str] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
webhook: Optional[str] = None,
multitask_strategy: Optional[MultitaskStrategy] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> Run: ...
@@ -1350,16 +1359,17 @@ class RunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
stream_mode: Union[StreamMode, list[StreamMode]] = "values",
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
checkpoint_id: Optional[str] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
webhook: Optional[str] = None,
multitask_strategy: Optional[MultitaskStrategy] = None,
if_not_exists: Optional[IfNotExists] = None,
on_completion: Optional[OnCompletionBehavior] = None,
after_seconds: Optional[int] = None,
) -> Run:
@@ -1383,6 +1393,8 @@ class RunsClient:
Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
on_completion: Whether to delete or keep the thread created for a stateless run.
Must be one of 'delete' or 'keep'.
if_not_exists: How to handle missing thread. Defaults to 'reject'.
Must be either 'reject' (raise error if missing), or 'create' (create new thread).
after_seconds: The number of seconds to wait before starting the run.
Use to schedule future runs.
@@ -1466,6 +1478,7 @@ class RunsClient:
"checkpoint": checkpoint,
"checkpoint_id": checkpoint_id,
"multitask_strategy": multitask_strategy,
"if_not_exists": if_not_exists,
"on_completion": on_completion,
"after_seconds": after_seconds,
}
@@ -1495,11 +1508,12 @@ class RunsClient:
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
checkpoint_id: Optional[str] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
webhook: Optional[str] = None,
on_disconnect: Optional[DisconnectMode] = None,
multitask_strategy: Optional[MultitaskStrategy] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> Union[list[dict], dict[str, Any]]: ...
@@ -1512,11 +1526,12 @@ class RunsClient:
input: Optional[dict] = None,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
webhook: Optional[str] = None,
on_disconnect: Optional[DisconnectMode] = None,
on_completion: Optional[OnCompletionBehavior] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> Union[list[dict], dict[str, Any]]: ...
@@ -1530,12 +1545,13 @@ class RunsClient:
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
checkpoint_id: Optional[str] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
webhook: Optional[str] = None,
on_disconnect: Optional[DisconnectMode] = None,
on_completion: Optional[OnCompletionBehavior] = None,
multitask_strategy: Optional[MultitaskStrategy] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> Union[list[dict], dict[str, Any]]:
"""Create a run, wait until it finishes and return the final state.
@@ -1558,6 +1574,8 @@ class RunsClient:
Must be one of 'delete' or 'keep'.
multitask_strategy: Multitask strategy to use.
Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
if_not_exists: How to handle missing thread. Defaults to 'reject'.
Must be either 'reject' (raise error if missing), or 'create' (create new thread).
after_seconds: The number of seconds to wait before starting the run.
Use to schedule future runs.
@@ -1619,6 +1637,7 @@ class RunsClient:
"checkpoint": checkpoint,
"checkpoint_id": checkpoint_id,
"multitask_strategy": multitask_strategy,
"if_not_exists": if_not_exists,
"on_disconnect": on_disconnect,
"on_completion": on_completion,
"after_seconds": after_seconds,
@@ -3232,18 +3251,19 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
stream_mode: Union[StreamMode, list[StreamMode]] = "values",
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
checkpoint_id: Optional[str] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
feedback_keys: Optional[list[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
feedback_keys: Optional[Sequence[str]] = None,
on_disconnect: Optional[DisconnectMode] = None,
webhook: Optional[str] = None,
multitask_strategy: Optional[MultitaskStrategy] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> Iterator[StreamPart]: ...
@@ -3254,15 +3274,16 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
stream_mode: Union[StreamMode, list[StreamMode]] = "values",
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
feedback_keys: Optional[list[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
feedback_keys: Optional[Sequence[str]] = None,
on_disconnect: Optional[DisconnectMode] = None,
on_completion: Optional[OnCompletionBehavior] = None,
if_not_exists: Optional[IfNotExists] = None,
webhook: Optional[str] = None,
after_seconds: Optional[int] = None,
) -> Iterator[StreamPart]: ...
@@ -3273,19 +3294,20 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
stream_mode: Union[StreamMode, list[StreamMode]] = "values",
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
checkpoint_id: Optional[str] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
feedback_keys: Optional[list[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
feedback_keys: Optional[Sequence[str]] = None,
on_disconnect: Optional[DisconnectMode] = None,
on_completion: Optional[OnCompletionBehavior] = None,
webhook: Optional[str] = None,
multitask_strategy: Optional[MultitaskStrategy] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> Iterator[StreamPart]:
"""Create a run and stream the results.
@@ -3311,6 +3333,8 @@ class SyncRunsClient:
webhook: Webhook to call after LangGraph API call is done.
multitask_strategy: Multitask strategy to use.
Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
if_not_exists: How to handle missing thread. Defaults to 'reject'.
Must be either 'reject' (raise error if missing), or 'create' (create new thread).
after_seconds: The number of seconds to wait before starting the run.
Use to schedule future runs.
@@ -3356,6 +3380,7 @@ class SyncRunsClient:
"checkpoint": checkpoint,
"checkpoint_id": checkpoint_id,
"multitask_strategy": multitask_strategy,
"if_not_exists": if_not_exists,
"on_disconnect": on_disconnect,
"on_completion": on_completion,
"after_seconds": after_seconds,
@@ -3376,14 +3401,15 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
stream_mode: Union[StreamMode, list[StreamMode]] = "values",
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
webhook: Optional[str] = None,
on_completion: Optional[OnCompletionBehavior] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> Run: ...
@@ -3394,16 +3420,17 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
stream_mode: Union[StreamMode, list[StreamMode]] = "values",
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
checkpoint_id: Optional[str] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
webhook: Optional[str] = None,
multitask_strategy: Optional[MultitaskStrategy] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> Run: ...
@@ -3413,17 +3440,18 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
stream_mode: Union[StreamMode, list[StreamMode]] = "values",
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
checkpoint_id: Optional[str] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
webhook: Optional[str] = None,
multitask_strategy: Optional[MultitaskStrategy] = None,
on_completion: Optional[OnCompletionBehavior] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> Run:
"""Create a background run.
@@ -3446,6 +3474,8 @@ class SyncRunsClient:
Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
on_completion: Whether to delete or keep the thread created for a stateless run.
Must be one of 'delete' or 'keep'.
if_not_exists: How to handle missing thread. Defaults to 'reject'.
Must be either 'reject' (raise error if missing), or 'create' (create new thread).
after_seconds: The number of seconds to wait before starting the run.
Use to schedule future runs.
@@ -3529,6 +3559,7 @@ class SyncRunsClient:
"checkpoint": checkpoint,
"checkpoint_id": checkpoint_id,
"multitask_strategy": multitask_strategy,
"if_not_exists": if_not_exists,
"on_completion": on_completion,
"after_seconds": after_seconds,
}
@@ -3558,11 +3589,12 @@ class SyncRunsClient:
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
checkpoint_id: Optional[str] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
webhook: Optional[str] = None,
on_disconnect: Optional[DisconnectMode] = None,
multitask_strategy: Optional[MultitaskStrategy] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> Union[list[dict], dict[str, Any]]: ...
@@ -3575,11 +3607,12 @@ class SyncRunsClient:
input: Optional[dict] = None,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
webhook: Optional[str] = None,
on_disconnect: Optional[DisconnectMode] = None,
on_completion: Optional[OnCompletionBehavior] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> Union[list[dict], dict[str, Any]]: ...
@@ -3593,12 +3626,13 @@ class SyncRunsClient:
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
checkpoint_id: Optional[str] = None,
interrupt_before: Optional[Union[All, list[str]]] = None,
interrupt_after: Optional[Union[All, list[str]]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
webhook: Optional[str] = None,
on_disconnect: Optional[DisconnectMode] = None,
on_completion: Optional[OnCompletionBehavior] = None,
multitask_strategy: Optional[MultitaskStrategy] = None,
if_not_exists: Optional[IfNotExists] = None,
after_seconds: Optional[int] = None,
) -> Union[list[dict], dict[str, Any]]:
"""Create a run, wait until it finishes and return the final state.
@@ -3621,6 +3655,8 @@ class SyncRunsClient:
Must be one of 'delete' or 'keep'.
multitask_strategy: Multitask strategy to use.
Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
if_not_exists: How to handle missing thread. Defaults to 'reject'.
Must be either 'reject' (raise error if missing), or 'create' (create new thread).
after_seconds: The number of seconds to wait before starting the run.
Use to schedule future runs.
@@ -3682,6 +3718,7 @@ class SyncRunsClient:
"checkpoint": checkpoint,
"checkpoint_id": checkpoint_id,
"multitask_strategy": multitask_strategy,
"if_not_exists": if_not_exists,
"on_disconnect": on_disconnect,
"on_completion": on_completion,
"after_seconds": after_seconds,
+9 -1
View File
@@ -17,12 +17,13 @@ Represents the status of a run:
- "interrupted": The run was manually stopped or interrupted.
"""
ThreadStatus = Literal["idle", "busy", "interrupted"]
ThreadStatus = Literal["idle", "busy", "interrupted", "error"]
"""
Represents the status of a thread:
- "idle": The thread is not currently processing any task.
- "busy": The thread is actively processing a task.
- "interrupted": The thread's execution was interrupted.
- "error": An exception occurred during task processing.
"""
StreamMode = Literal["values", "messages", "updates", "events", "debug", "custom"]
@@ -69,6 +70,13 @@ Defines action after completion:
All = Literal["*"]
"""Represents a wildcard or 'all' selector."""
IfNotExists = Literal["create", "reject"]
"""
Specifies behavior if the thread doesn't exist:
- "create": Create a new thread if it doesn't exist.
- "reject": Reject the operation if the thread doesn't exist.
"""
class Config(TypedDict, total=False):
"""Configuration options for a call."""
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-sdk"
version = "0.1.33"
version = "0.1.34"
description = "SDK for interacting with LangGraph API"
authors = []
license = "MIT"