Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
39d9cdbef0 | ||
|
|
6dacd1aabe | ||
|
|
91ad8b803b | ||
|
|
a0e99f704f | ||
|
|
def3e06b4a | ||
|
|
a4fab2a867 | ||
|
|
cf1c7f3673 | ||
|
|
6c05b66c7c | ||
|
|
5c44dcef81 | ||
|
|
7fd6b1b4be | ||
|
|
6202e0f1d9 | ||
|
|
83238f51d8 | ||
|
|
05f008cbfb | ||
|
|
a8ae2a52a3 | ||
|
|
aa245a8e71 | ||
|
|
bdc75a22d5 | ||
|
|
69227daff3 | ||
|
|
dc8260bb72 | ||
|
|
62a5ec509d | ||
|
|
d32386f849 | ||
|
|
08a1ed38f1 | ||
|
|
037a95ff60 | ||
|
|
e294720ec5 | ||
|
|
f8a0b7a464 | ||
|
|
1121806ba4 | ||
|
|
dca200d6c4 | ||
|
|
6f236b5f2c | ||
|
|
916affa1b5 | ||
|
|
58cf0c6a6e | ||
|
|
0042889c31 | ||
|
|
2be012d8ed | ||
|
|
780285ef91 | ||
|
|
47c7b76aa1 | ||
|
|
9931f61525 | ||
|
|
0d81ad92f0 | ||
|
|
7d3f2ca3ed | ||
|
|
42648c88dd | ||
|
|
1aab758634 | ||
|
|
b5fbc7a7b8 | ||
|
|
b647dcb0f2 | ||
|
|
4df5680732 | ||
|
|
74a17a6d4c | ||
|
|
e2a3698250 | ||
|
|
4dfdb9a83e | ||
|
|
583d8c9499 | ||
|
|
15bbede7bc | ||
|
|
3ffdf4bb3f | ||
|
|
6578698414 | ||
|
|
048ae6c17b | ||
|
|
0e2c2eb13a | ||
|
|
515c4ffebe | ||
|
|
eefe057a47 | ||
|
|
18f34c30d8 | ||
|
|
2670bcf330 | ||
|
|
19ccb0c6af | ||
|
|
a277b86fcb | ||
|
|
2f819a6a9b |
@@ -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 }}
|
||||
|
||||
@@ -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 }}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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": {
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
After Width: | Height: | Size: 178 KiB |
|
Before Width: | Height: | Size: 193 KiB |
|
Before Width: | Height: | Size: 55 KiB |
|
After Width: | Height: | Size: 83 KiB |
|
After Width: | Height: | Size: 103 KiB |
|
Before Width: | Height: | Size: 97 KiB |
|
Before Width: | Height: | Size: 35 KiB |
@@ -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.
|
||||
|
||||
@@ -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"})
|
||||
```
|
||||
|
||||
@@ -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
|
||||

|
||||
|
||||

|
||||
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 it’s 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__"]
|
||||
|
||||

|
||||
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()
|
||||
|
||||

|
||||
# 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):
|
||||
...
|
||||
|
||||

|
||||
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.
|
||||
|
||||

|
||||
|
||||
#### 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, it’s 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)?
|
||||
|
||||

|
||||
|
||||
#### 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.
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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("."))
|
||||
@@ -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"
|
||||
@@ -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]}")
|
||||
@@ -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]}")
|
||||
@@ -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]:
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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,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
|
||||
|
||||
@@ -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__")
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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],
|
||||
|
||||
@@ -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",
|
||||
)
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -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'}}}%%
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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="",
|
||||
|
||||
@@ -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)
|
||||
@@ -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,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",
|
||||
|
||||
@@ -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";
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,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"
|
||||
|
||||