Compare commits

..
Author SHA1 Message Date
William Fu-Hinthorn fba8665718 merge 2024-11-27 14:54:03 -08:00
William Fu-Hinthorn 70b42812b5 Merge branch 'wfh/store/base/add_vector_earch' into wfh/store/add_vector_search 2024-11-27 14:46:44 -08:00
William FHandGitHub 471b2fad51 Merge branch 'main' into wfh/store/base/add_vector_earch 2024-11-27 14:41:37 -08:00
William FHandGitHub 98db2c08d7 Merge branch 'main' into wfh/store/base/add_vector_earch 2024-11-27 14:36:01 -08:00
William Fu-Hinthorn 1f5cab1869 Format 2024-11-27 14:34:31 -08:00
William Fu-Hinthorn 33dc6c84bb Merge branch 'wfh/store/base/add_vector_earch' into wfh/store/add_vector_search 2024-11-27 14:16:08 -08:00
William Fu-Hinthorn 4455e4185c Improve docs 2024-11-27 14:16:00 -08:00
William Fu-Hinthorn d3b2e2dc95 merge 2024-11-27 14:13:43 -08:00
William Fu-Hinthorn 2bff330bd7 Merge branch 'wfh/store/base/add_vector_earch' into wfh/store/add_vector_search 2024-11-27 14:08:05 -08:00
William Fu-Hinthorn c2a1f3a30a Improve docs 2024-11-27 14:06:54 -08:00
William Fu-Hinthorn 112a4f6c12 Test put-time fields 2024-11-27 13:19:04 -08:00
William Fu-Hinthorn 34a4ca3eaf Rename file 2024-11-27 11:52:11 -08:00
William Fu-Hinthorn d75492c330 Rename config 2024-11-27 11:20:03 -08:00
William Fu-Hinthorn 0ed26a3af9 handle no emb situation 2024-11-27 06:40:21 -08:00
William Fu-Hinthorn 74229c03dc merge 2024-11-26 21:30:41 -08:00
William Fu-Hinthorn 9a83dd7900 Add test 2024-11-26 19:23:48 -08:00
William Fu-Hinthorn b03235f92d Add in-mem vector search 2024-11-26 18:59:44 -08:00
William Fu-Hinthorn 6ec2afe958 Handle conflicts 2024-11-26 18:56:42 -08:00
William Fu-Hinthorn 11dc4d2691 merge 2024-11-26 18:46:21 -08:00
William Fu-Hinthorn 01b4d15aca Add in-mem vector search 2024-11-26 18:43:16 -08:00
William Fu-Hinthorn d81e350105 feat: Add vector search 2024-11-25 17:09:32 -08:00
116 changed files with 2292 additions and 8780 deletions
+40 -9
View File
@@ -7,29 +7,35 @@ body:
value: >
Thank you for taking the time to file a bug report.
Use this to report BUGS in LangGraph. For usage questions, feature requests and general design questions, please use [GitHub Discussions](https://github.com/langchain-ai/langgraph/discussions).
Use this to report bugs in LangChain.
If you're not certain that your issue is due to a bug in LangChain, please use [GitHub Discussions](https://github.com/langchain-ai/langchain/discussions)
to ask for help with your issue.
Relevant links to check before filing a bug report to see if your issue has already been reported, fixed or
if there's another way to solve your problem:
[LangGraph Github Discussions](https://github.com/langchain-ai/langgraph/discussions),
[LangGraph Github Issues](https://github.com/langchain-ai/langgraph/issues),
[LangGraph how-to guides](https://langchain-ai.github.io/langgraph/how-tos/).
[LangGraph documentation](https://langchain-ai.github.io/langgraph/).
[LangChain documentation with the integrated search](https://python.langchain.com/docs/get_started/introduction),
[GitHub search](https://github.com/langchain-ai/langgraph),
[LangChain Github Discussions](https://github.com/langchain-ai/langgraph/discussions),
[LangChain Github Issues](https://github.com/langchain-ai/langgraph/issues),
[LangChain ChatBot](https://chat.langchain.com/)
- type: checkboxes
id: checks
attributes:
label: Checked other resources
description: Before submitting this issue, please confirm that you have completed all the steps below by checking each option. These steps help ensure your issue is well-defined, relevant, and actionable.
description: Please confirm and check all the following options.
options:
- label: This is a bug, not a usage question. For questions, please use GitHub Discussions.
- label: I added a very descriptive title to this issue.
required: true
- label: I added a clear and detailed title that summarizes the issue.
- label: I searched the [LangGraph](https://langchain-ai.github.io/langgraph/)/LangChain documentation with the integrated search.
required: true
- label: I read what a minimal reproducible example is (https://stackoverflow.com/help/minimal-reproducible-example).
- label: I used the GitHub search to find a similar question and didn't find it.
required: true
- label: I included a self-contained, minimal example that demonstrates the issue INCLUDING all the relevant imports. The code run AS IS to reproduce the issue.
- label: I am sure that this is a bug in LangGraph/LangChain rather than my code.
required: true
- label: I am sure this is better as an issue [rather than a GitHub discussion](https://github.com/langchain-ai/langgraph/discussions/new/choose), since this is a LangGraph bug and not a design question.
required: true
- type: textarea
id: reproduction
@@ -39,6 +45,14 @@ body:
label: Example Code
description: |
Please add a self-contained, [minimal, reproducible, example](https://stackoverflow.com/help/minimal-reproducible-example) with your use case.
If a maintainer can copy it, run it, and see it right away, there's a much higher chance that you'll be able to get help.
**Important!**
* Reduce your code to the minimum required to reproduce the issue if possible. This makes it much easier for others to help you.
* Avoid screenshots when possible, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.
placeholder: |
from langgraph.graph import StateGraph
@@ -78,8 +92,25 @@ body:
attributes:
label: System Info
description: |
Please share your system info with us.
"pip freeze | grep langchain"
platform (windows / linux / mac)
python version
OR if you're on a recent version of langchain-core you can paste the output of:
python -m langchain_core.sys_info
placeholder: |
"pip freeze | grep langgraph"
platform
python version
Alternatively, if you're on a recent version of langchain-core you can paste the output of:
python -m langchain_core.sys_info
These will only surface LangChain packages, don't forget to include any other relevant
packages you're using (if you're not sure what's relevant, you can paste the entire output of `pip freeze`).
validations:
required: true
+1
View File
@@ -42,6 +42,7 @@ jobs:
with:
python-version: ${{ matrix.python-version }}
poetry-version: ${{ env.POETRY_VERSION }}
working-directory: ${{ inputs.working-directory }}
cache-key: lint-${{ inputs.working-directory }}
- name: Check Poetry File
+1
View File
@@ -31,6 +31,7 @@ jobs:
with:
python-version: ${{ matrix.python-version }}
poetry-version: ${{ env.POETRY_VERSION }}
working-directory: ${{ inputs.working-directory }}
cache-key: test-${{ inputs.working-directory }}
- name: Login to Docker Hub
uses: docker/login-action@v3
+1 -1
View File
@@ -60,7 +60,7 @@ jobs:
env:
LANGGRAPH_FF_SEND_V2: ${{ matrix.ff-send-v2 }}
run: |
make test_parallel
make test
- name: Ensure the tests did not create any additional files
shell: bash
+1
View File
@@ -29,6 +29,7 @@ jobs:
with:
python-version: ${{ env.PYTHON_VERSION }}
poetry-version: ${{ env.POETRY_VERSION }}
working-directory: ${{ inputs.working-directory }}
cache-key: release
# We want to keep this build stage *separate* from the release stage,
+4
View File
@@ -31,6 +31,7 @@ jobs:
with:
python-version: ${{ env.PYTHON_VERSION }}
poetry-version: ${{ env.POETRY_VERSION }}
working-directory: ${{ inputs.working-directory }}
cache-key: release
# We want to keep this build stage *separate* from the release stage,
@@ -168,6 +169,7 @@ jobs:
with:
python-version: ${{ env.PYTHON_VERSION }}
poetry-version: ${{ env.POETRY_VERSION }}
working-directory: ${{ inputs.working-directory }}
- name: Import published package
shell: bash
@@ -254,6 +256,7 @@ jobs:
with:
python-version: ${{ env.PYTHON_VERSION }}
poetry-version: ${{ env.POETRY_VERSION }}
working-directory: ${{ inputs.working-directory }}
cache-key: release
- uses: actions/download-artifact@v4
@@ -295,6 +298,7 @@ jobs:
with:
python-version: ${{ env.PYTHON_VERSION }}
poetry-version: ${{ env.POETRY_VERSION }}
working-directory: ${{ inputs.working-directory }}
cache-key: release
- uses: actions/download-artifact@v4
+1 -1
View File
@@ -49,7 +49,7 @@ gain understanding of concepts and how they interact by showing one way to achie
They should **avoid** giving
multiple permutations of ways to achieve that goal in-depth. Choice is burdensome. Instead, they should guide a new user through a recommended path to accomplishing a concrete goal. While the end result of a tutorial does not necessarily need to
be completely production-ready, it should be useful and practically satisfy the goal that you clearly stated in the tutorial's introduction.
be completely production-ready, it should be useful and practically satisfy the the goal that you clearly stated in the tutorial's introduction.
To quote the Diataxis website:
+1 -1
View File
@@ -13,7 +13,7 @@ serve-clean-docs: clean-docs
poetry run python -m mkdocs serve -c -f docs/mkdocs.yml --strict -w ./libs/langgraph
serve-docs: build-typedoc
poetry run python -m mkdocs serve -f docs/mkdocs.yml -w ./libs/langgraph -w ./libs/checkpoint --dirty
poetry run python -m mkdocs serve -f docs/mkdocs.yml -w ./libs/langgraph --dirty
clean-docs:
find ./docs/docs -name "*.ipynb" -type f -delete
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -1,123 +0,0 @@
# How to add semantic search to your LangGraph deployment
This guide explains how to add semantic search to your LangGraph deployment's cross-thread [store](../../concepts/persistence.md#memory-store), so that your agent can search for memories and other documents by semantic similarity.
## Prerequisites
- A LangGraph deployment (see [how to deploy](setup_pyproject.md))
- API keys for your embedding provider (in this case, OpenAI)
- `langchain >= 0.3.8` (if you specify using the string format below)
## Steps
1. Update your `langgraph.json` configuration file to include the store configuration:
```json
{
...
"store": {
"index": {
"embed": "openai:text-embeddings-3-small",
"dims": 1536,
"fields": ["$"]
}
}
}
```
This configuration:
- Uses OpenAI's text-embeddings-3-small model for generating embeddings
- Sets the embedding dimension to 1536 (matching the model's output)
- Indexes all fields in your stored data (`["$"]` means index everything, or specify specific fields like `["text", "metadata.title"]`)
2. To use the string embedding format above, make sure your dependencies include `langchain >= 0.3.8`:
```toml
# In pyproject.toml
[project]
dependencies = [
"langchain>=0.3.8"
]
```
Or if using requirements.txt:
```
langchain>=0.3.8
```
## Usage
Once configured, you can use semantic search in your LangGraph nodes. The store requires a namespace tuple to organize memories:
```python
def search_memory(state: State, *, store: BaseStore):
# Search the store using semantic similarity
# The namespace tuple helps organize different types of memories
# e.g., ("user_facts", "preferences") or ("conversation", "summaries")
results = store.search(
namespace=("memory", "facts"), # Organize memories by type
query="your search query",
limit=3 # number of results to return
)
return results
```
## Custom Embeddings
If you want to use custom embeddings, you can pass a path to a custom embedding function:
```json
{
...
"store": {
"index": {
"embed": "path/to/embedding_function.py:embed",
"dims": 1536,
"fields": ["$"]
}
}
}
```
The deployment will look for the function in the specified path. The function must be async and accept a list of strings:
```python
# path/to/embedding_function.py
from openai import AsyncOpenAI
client = AsyncOpenAI()
async def aembed_texts(texts: list[str]) -> list[list[float]]:
"""Custom embedding function that must:
1. Be async
2. Accept a list of strings
3. Return a list of float arrays (embeddings)
"""
response = await client.embeddings.create(
model="text-embedding-3-small",
input=texts
)
return [e.embedding for e in response.data]
```
## Querying via the API
You can also query the store using the LangGraph SDK. Since the SDK uses async operations:
```python
from langgraph_sdk import get_client
async def search_store():
client = get_client()
results = await client.store.search_items(
("memory", "facts"),
query="your search query",
limit=3 # number of results to return
)
return results
# Use in an async context
results = await search_store()
```
+2 -2
View File
@@ -36,8 +36,8 @@ Dependencies can optionally be specified in one of the following files: `pyproje
The dependencies below will be included in the image, you can also use them in your code, as long as with a compatible version range:
```
langgraph>=0.2.56,<0.3.0
langgraph-checkpoint>=2.0.5,<3.0
langgraph>=0.2.30,<0.3.0
langgraph-checkpoint>=1.0.14
langchain-core>=0.2.38,<0.4.0
langsmith>=0.1.63
orjson>=3.9.7
@@ -36,8 +36,8 @@ Dependencies can optionally be specified in one of the following files: `pyproje
The dependencies below will be included in the image, you can also use them in your code, as long as with a compatible version range:
```
langgraph>=0.2.56,<0.3.0
langgraph-checkpoint>=2.0.5,<3.0
langgraph>=0.2.30,<0.3.0
langgraph-checkpoint>=1.0.14
langchain-core>=0.2.38,<0.4.0
langsmith>=0.1.63
orjson>=3.9.7
@@ -83,7 +83,7 @@ We can now call `.get_schemas` to get schemas associated with this graph:
assistant_id=assistant["assistant_id"]
)
# There are multiple types of schemas
# We can get the `config_schema` to look at the configurable parameters
# We can get the `config_schema` to look at the the configurable parameters
print(schemas["config_schema"])
```
@@ -94,7 +94,7 @@ We can now call `.get_schemas` to get schemas associated with this graph:
assistant["assistant_id"]
);
// There are multiple types of schemas
// We can get the `config_schema` to look at the configurable parameters
// We can get the `config_schema` to look at the the configurable parameters
console.log(schemas.config_schema);
```
+36 -93
View File
@@ -26,11 +26,10 @@ The LangGraph command line interface includes commands to build and run a LangGr
The LangGraph CLI requires a JSON configuration file with the following keys:
| Key | Description |
| ------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|--------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| `dependencies` | **Required**. Array of dependencies for LangGraph Cloud API server. Dependencies can be one of the following: (1) `"."`, which will look for local Python packages, (2) `pyproject.toml`, `setup.py` or `requirements.txt` in the app directory `"./local_package"`, or (3) a package name. |
| `graphs` | **Required**. Mapping from graph ID to path where the compiled graph or a function that makes a graph is defined. Example: <ul><li>`./your_package/your_file.py:variable`, where `variable` is an instance of `langgraph.graph.state.CompiledStateGraph`</li><li>`./your_package/your_file.py:make_graph`, where `make_graph` is a function that takes a config dictionary (`langchain_core.runnables.RunnableConfig`) and creates an instance of `langgraph.graph.state.StateGraph` / `langgraph.graph.state.CompiledStateGraph`.</li></ul> |
| `env` | Path to `.env` file or a mapping from environment variable to its value. |
| `store` | Configuration for adding semantic search to the BaseStore. Contains the following fields: <ul><li>`index`: Configuration for semantic search indexing with fields:<ul><li>`embed`: Embedding provider (e.g., "openai:text-embedding-3-small") or path to custom embedding function</li><li>`dims`: Dimension size of the embedding model. Used to initialize the vector table.</li><li>`fields` (optional): List of fields to index. Defaults to `["$"]`, meaningto index entire documents. Can be specific fields like `["text", "summary", "some.value"]`</li></ul></li></ul> |
| `python_version` | `3.11` or `3.12`. Defaults to `3.11`. |
| `pip_config_file` | Path to `pip` config file. |
| `dockerfile_lines` | Array of additional lines to add to Dockerfile following the import from parent image. |
@@ -42,84 +41,33 @@ The LangGraph CLI requires a JSON configuration file with the following keys:
</p>
</div>
### Examples
#### Basic Configuration
Example:
```json
{
"dependencies": ["."],
"dependencies": ["langchain_openai", "./your_package"],
"graphs": {
"chat": "./chat/graph.py:graph"
}
}
```
#### Adding semantic search to the store
All deployments come with a DB-backed BaseStore. Adding an "index" configuration to your `langgraph.json` will enable [semantic search](../deployment/semantic_search.md) within the BaseStore of your deployment.
The `fields` configuration determines which parts of your documents to embed:
- If omitted or set to `["$"]`, the entire document will be embedded
- To embed specific fields, use JSON path notation: `["metadata.title", "content.text"]`
- Documents missing specified fields will still be stored but won't have embeddings for those fields
- You can still override which fields to embed on a specific item at `put` time using the `index` parameter
```json
{
"dependencies": ["."],
"graphs": {
"memory_agent": "./agent/graph.py:graph"
"my_graph_id": "./your_package/your_file.py:variable"
},
"store": {
"index": {
"embed": "openai:text-embedding-3-small",
"dims": 1536,
"fields": ["$"]
}
}
"env": "./.env"
}
```
!!! note "Common model dimensions"
- openai:text-embedding-3-large: 3072
- openai:text-embedding-3-small: 1536
- openai:text-embedding-ada-002: 1536
- cohere:embed-english-v3.0: 1024
- cohere:embed-english-light-v3.0: 384
- cohere:embed-multilingual-v3.0: 1024
- cohere:embed-multilingual-light-v3.0: 384
#### Semantic search with a custom embedding function
If you want to use semantic search with a custom embedding function, you can pass a path to a custom embedding function:
Example with environment variables:
```json
{
"dependencies": ["."],
"python_version": "3.11",
"dependencies": ["langchain_openai", "."],
"graphs": {
"memory_agent": "./agent/graph.py:graph"
"my_graph_id": "./your_package/your_file.py:make_graph"
},
"store": {
"index": {
"embed": "./embeddings.py:embed_texts",
"dims": 768,
"fields": ["text", "summary"]
}
"env": {
"OPENAI_API_KEY": "secret-key"
}
}
```
The `embed` field in store configuration can reference a custom function that takes a list of strings and returns a list of embeddings. Example implementation:
```python
# embeddings.py
def embed_texts(texts: list[str]) -> list[list[float]]:
"""Custom embedding function for semantic search."""
# Implementation using your preferred embedding model
return [[0.1, 0.2, ...] for _ in texts] # dims-dimensional vectors
```
## Commands
The base command for the LangGraph CLI is `langgraph`.
@@ -134,11 +82,6 @@ langgraph [OPTIONS] COMMAND [ARGS]
Run LangGraph API server in development mode with hot reloading and debugging capabilities. This lightweight server requires no Docker installation and is suitable for development and testing. State is persisted to a local directory.
!!! note "Python only"
Currently, the CLI only supports Python >= 3.11.
JS support is coming soon.
**Installation**
This command requires the "inmem" extra to be installed:
@@ -155,16 +98,16 @@ langgraph dev [OPTIONS]
**Options**
| Option | Default | Description |
| ----------------------------- | ---------------- | ----------------------------------------------------------------------------------- |
| `-c, --config FILE` | `langgraph.json` | Path to configuration file declaring dependencies, graphs and environment variables |
| `--host TEXT` | `127.0.0.1` | Host to bind the server to |
| `--port INTEGER` | `2024` | Port to bind the server to |
| `--no-reload` | | Disable auto-reload |
| `--n-jobs-per-worker INTEGER` | | Number of jobs per worker. Default is 10 |
| `--no-browser` | | Disable automatic browser opening |
| `--debug-port INTEGER` | | Port for debugger to listen on |
| `--help` | | Display command documentation |
| Option | Default | Description |
|----------------------------|------------------|--------------------------------------------------------------------------------------------|
| `-c, --config FILE` | `langgraph.json` | Path to configuration file declaring dependencies, graphs and environment variables |
| `--host TEXT` | `127.0.0.1` | Host to bind the server to |
| `--port INTEGER` | `2024` | Port to bind the server to |
| `--no-reload` | | Disable auto-reload |
| `--n-jobs-per-worker INTEGER` | | Number of jobs per worker. Default is 10 |
| `--no-browser` | | Disable automatic browser opening |
| `--debug-port INTEGER` | | Port for debugger to listen on |
| `--help` | | Display command documentation |
### `build`
@@ -179,7 +122,7 @@ langgraph build [OPTIONS]
**Options**
| Option | Default | Description |
| -------------------- | ---------------- | ---------------------------------------------------------------------------------------------------------------------------- |
|----------------------|------------------|------------------------------------------------------------------------------------------------------------------------------|
| `--platform TEXT` | | Target platform(s) to build the Docker image for. Example: `langgraph build --platform linux/amd64,linux/arm64` |
| `-t, --tag TEXT` | | **Required**. Tag for the Docker image. Example: `langgraph build -t my-image` |
| `--pull / --no-pull` | `--pull` | Build with latest remote Docker image. Use `--no-pull` for running the LangGraph Cloud API server with locally built images. |
@@ -198,20 +141,20 @@ langgraph up [OPTIONS]
**Options**
| Option | Default | Description |
| ---------------------------- | ------------------------- | ----------------------------------------------------------------------------------------------------------------------- |
| `--wait` | | Wait for services to start before returning. Implies --detach |
| `--postgres-uri TEXT` | Local database | Postgres URI to use for the database. |
| `--watch` | | Restart on file changes |
| `--debugger-base-url TEXT` | `http://127.0.0.1:[PORT]` | URL used by the debugger to access LangGraph API. |
| `--debugger-port INTEGER` | | Pull the debugger image locally and serve the UI on specified port |
| `--verbose` | | Show more output from the server logs. |
| `-c, --config FILE` | `langgraph.json` | Path to configuration file declaring dependencies, graphs and environment variables. |
| `-d, --docker-compose FILE` | | Path to docker-compose.yml file with additional services to launch. |
| `-p, --port INTEGER` | `8123` | Port to expose. Example: `langgraph up --port 8000` |
| Option | Default | Description |
|------------------------------|---------------------------|-----------------------------------------------------------------------------------------------------------------------|
| `--wait` | | Wait for services to start before returning. Implies --detach |
| `--postgres-uri TEXT` | Local database | Postgres URI to use for the database. |
| `--watch` | | Restart on file changes |
| `--debugger-base-url TEXT` | `http://127.0.0.1:[PORT]` | URL used by the debugger to access LangGraph API. |
| `--debugger-port INTEGER` | | Pull the debugger image locally and serve the UI on specified port |
| `--verbose` | | Show more output from the server logs. |
| `-c, --config FILE` | `langgraph.json` | Path to configuration file declaring dependencies, graphs and environment variables. |
| `-d, --docker-compose FILE` | | Path to docker-compose.yml file with additional services to launch. |
| `-p, --port INTEGER` | `8123` | Port to expose. Example: `langgraph up --port 8000` |
| `--pull / --no-pull` | `pull` | Pull latest images. Use `--no-pull` for running the server with locally-built images. Example: `langgraph up --no-pull` |
| `--recreate / --no-recreate` | `no-recreate` | Recreate containers even if their configuration and image haven't changed |
| `--help` | | Display command documentation. |
| `--recreate / --no-recreate` | `no-recreate` | Recreate containers even if their configuration and image haven't changed |
| `--help` | | Display command documentation. |
### `dockerfile`
@@ -226,7 +169,7 @@ langgraph dockerfile [OPTIONS] SAVE_PATH
**Options**
| Option | Default | Description |
| ------------------- | ---------------- | --------------------------------------------------------------------------------------------------------------- |
|---------------------|------------------|-----------------------------------------------------------------------------------------------------------------|
| `-c, --config FILE` | `langgraph.json` | Path to the [configuration file](#configuration-file) declaring dependencies, graphs and environment variables. |
| `--help` | | Show this message and exit. |
+13 -13
View File
@@ -27,8 +27,8 @@ Adding a [breakpoint](./low_level.md#breakpoints) a specific location in the gra
Here, we compile our graph with a checkpointer and a breakpoint at the node we want to interrupt before, `step_for_human_in_the_loop`. We then perform one of the above interaction patterns, which will create a new checkpoint if a human edits the graph state. The new checkpoint is saved to the `thread` and we can resume the graph execution from there by passing in `None` as the input.
```python
# Compile our graph with a checkpointer and a breakpoint before "step_for_human_in_the_loop"
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["step_for_human_in_the_loop"])
# Compile our graph with a checkpoitner and a breakpoint before "step_for_human_in_the_loop"
graph = builder.compile(checkpointer=checkpoitner, interrupt_before=["step_for_human_in_the_loop"])
# Run the graph up to the breakpoint
thread_config = {"configurable": {"thread_id": "1"}}
@@ -98,8 +98,8 @@ With persistence, we can surface the current agent state as well as the next ste
If approved, the graph resumes execution from the last saved checkpoint, which is saved to the `thread`:
```python
# Compile our graph with a checkpointer and a breakpoint before the step to approve
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node_2"])
# Compile our graph with a checkpoitner and a breakpoint before the step to approve
graph = builder.compile(checkpointer=checkpoitner, interrupt_before=["node_2"])
# Run the graph up to the breakpoint
for event in graph.stream(inputs, thread, stream_mode="values"):
@@ -120,7 +120,7 @@ See [our guide](../how-tos/human_in_the_loop/breakpoints.ipynb) for a detailed h
Sometimes we want to review and edit the agent's state.
As with approval, we can interrupt our agent at a [breakpoint](./low_level.md#breakpoints) prior to the step we want to check.
As with approval, we can interrupt our agent at a [breakpoint](./low_level.md#breakpoints) prior the the step we want to check.
We can surface the current state to a user and allow the user to edit the agent state.
@@ -131,8 +131,8 @@ We can edit the graph state by forking the current checkpoint, which is saved to
We can then proceed with the graph from our forked checkpoint as done before.
```python
# Compile our graph with a checkpointer and a breakpoint before the step to review
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node_2"])
# Compile our graph with a checkpoitner and a breakpoint before the step to review
graph = builder.compile(checkpointer=checkpoitner, interrupt_before=["node_2"])
# Run the graph up to the breakpoint
for event in graph.stream(inputs, thread, stream_mode="values"):
@@ -170,11 +170,11 @@ With editing, the user makes a decision about whether or not to edit the graph s
With input, we explicitly define a node in our graph for collecting human input!
The state update with the human input then runs *as this node*.
The the state update with the human input then runs *as this node*.
```python
# Compile our graph with a checkpointer and a breakpoint before the step to to collect human input
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["human_input"])
# Compile our graph with a checkpoitner and a breakpoint before the step to to collect human input
graph = builder.compile(checkpointer=checkpoitner, interrupt_before=["human_input"])
# Run the graph up to the breakpoint
for event in graph.stream(inputs, thread, stream_mode="values"):
@@ -211,8 +211,8 @@ Even if the tool call is correct, we may also want to apply discretion:
With these points in mind, we can combine the above ideas to create a human-in-the-loop review of a tool call.
```python
# Compile our graph with a checkpointer and a breakpoint before the step to to review the tool call from the LLM
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["human_review"])
# Compile our graph with a checkpoitner and a breakpoint before the step to to review the tool call from the LLM
graph = builder.compile(checkpointer=checkpoitner, interrupt_before=["human_review"])
# Run the graph up to the breakpoint
for event in graph.stream(inputs, thread, stream_mode="values"):
@@ -319,4 +319,4 @@ for event in graph.stream(None, config, stream_mode="values"):
See [this additional conceptual guide](https://langchain-ai.github.io/langgraph/concepts/persistence/#update-state) for related context on forking.
See see [this guide](../how-tos/human_in_the_loop/time-travel.ipynb) for a detailed how-to on doing time-travel!
See see [this guide](../how-tos/human_in_the_loop/time-travel.ipynb) for a detailed how-to on doing time-travel!
-5
View File
@@ -33,11 +33,6 @@ The `langgraph build` command builds a Docker image for the [LangGraph API serve
!!! note "New in version 0.1.55"
The `langgraph dev` command was introduced in langgraph-cli version 0.1.55.
!!! note "Python only"
Currently, the CLI only supports Python >= 3.11.
JS support is coming soon.
The `langgraph dev` command starts a lightweight development server that requires no Docker installation. This server is ideal for rapid development and testing, with features like:
- Hot reloading: Changes to your code are automatically detected and reloaded
-61
View File
@@ -283,9 +283,6 @@ You can optionally provide a dictionary that maps the `routing_function`'s outpu
graph.add_conditional_edges("node_a", routing_function, {True: "node_b", False: "node_c"})
```
!!! tip
Use [`Command`](#command) instead of conditional edges if you want to combine state updates and routing in a single function.
### Entry Point
The entry point is the first node(s) that are run when the graph starts. You can use the [`add_edge`][langgraph.graph.StateGraph.add_edge] method from the virtual [`START`][langgraph.constants.START] node to the first node to execute to specify where to enter the graph.
@@ -325,64 +322,6 @@ def continue_to_jokes(state: OverallState):
graph.add_conditional_edges("node_a", continue_to_jokes)
```
## `Command`
It can be useful to combine control flow (edges) and state updates (nodes). For example, you might want to BOTH perform state updates AND decide which node to go to next in the SAME node. LangGraph provides a way to do so by returning a [`Command`][langgraph.types.Command] object from node functions:
```python
def my_node(state: State) -> Command[Literal["my_other_node"]]:
return Command(
# state update
update={"foo": "bar"},
# control flow
goto="my_other_node"
)
```
With `Command` you can also achieve dynamic control flow behavior (identical to [conditional edges](#conditional-edges)):
```python
def my_node(state: State) -> Command[Literal["my_other_node"]]:
if state["foo"] == "bar":
return Command(update={"foo": "baz"}, goto="my_other_node")
```
!!! important
When returning `Command` in your node functions, you must add return type annotations with the list of node names the node is routing to, e.g. `Command[Literal["my_other_node"]]`. This is necessary for the graph rendering and tells LangGraph that `my_node` can navigate to `my_other_node`.
Check out this [how-to guide](../how-tos/command.ipynb) for an end-to-end example of how to use `Command`.
### When should I use Command instead of conditional edges?
Use `Command` when you need to **both** update the graph state **and** route to a different node. For example, when implementing [multi-agent handoffs](./multi_agent.md#handoffs) where it's important to route to a different agent and pass some information to that agent.
Use [conditional edges](#conditional-edges) to route between nodes conditionally without updating the state.
### Using inside tools
A common use case is updating graph state from inside a tool. For example, in a customer support application you might want to look up customer information based on their account number or ID in the beginning of the conversation. To update the graph state from the tool, you can return `Command(update={"my_custom_key": "foo", "messages": [...]})` from the tool:
```python
@tool
def lookup_user_info(tool_call_id: Annotated[str, InjectedToolCallId], config: RunnableConfig):
"""Use this to look up user information to better assist them with their questions."""
user_info = get_user_info(config.get("configurable", {}).get("user_id"))
return Command(
update={
# update the state keys
"user_info": user_info,
# update the message history
"messages": [ToolMessage("Successfully looked up user information", tool_call_id=tool_call_id)]
}
)
```
!!! important
You MUST include `messages` (or any state key used for the message history) in `Command.update` when returning `Command` from a tool and the list of messages in `messages` MUST contain a `ToolMessage`. This is necessary for the resulting message history to be valid (LLM providers require AI messages with tool calls to be followed by the tool result messages).
If you are using tools that update state via `Command`, we recommend using prebuilt [`ToolNode`][langgraph.prebuilt.tool_node.ToolNode] which automatically handles tools returning `Command` objects and propagates them to the graph state. If you're writing a custom node that calls tools, you would need to manually propagate `Command` objects returned by the tools as the update from node.
## 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
+6 -27
View File
@@ -171,7 +171,7 @@ trim_messages(
## Long-term memory
Long-term memory in LangGraph allows systems to retain information across different conversations or sessions. Unlike short-term memory, which is **thread-scoped**, long-term memory is saved within custom "namespaces."
Long-term memory in LangGraph allows systems to retain information across different conversations or sessions. Unlike short-term memory, which is thread-scoped, long-term memory is saved within custom "namespaces."
### Storing memories
@@ -180,34 +180,16 @@ LangGraph stores long-term memories as JSON documents in a [store](persistence.m
```python
from langgraph.store.memory import InMemoryStore
def embed(texts: list[str]) -> list[list[float]]:
# Replace with an actual embedding function or LangChain embeddings object
return [[1.0, 2.0] * len(texts)]
# InMemoryStore saves data to an in-memory dictionary. Use a DB-backed store in production use.
store = InMemoryStore(index={"embed": embed, "dims": 2})
store = InMemoryStore()
user_id = "my-user"
application_context = "chitchat"
namespace = (user_id, application_context)
store.put(
namespace,
"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, "a-memory")
# search for "memories" within this namespace, filtering on content equivalence, sorted by vector similarity
items = store.search(
namespace, filter={"my-key": "my-value"}, query="language preferences"
)
# list "memories" within this namespace, filtering on content equivalence
items = store.search(namespace, filter={"my-key": "my-value"})
```
### Framework for thinking about long-term memory
@@ -236,9 +218,6 @@ Different applications require various types of memory. Although the analogy isn
[Semantic memory](https://en.wikipedia.org/wiki/Semantic_memory), both in humans and AI agents, involves the retention of specific facts and concepts. In humans, it can include information learned in school and the understanding of concepts and their relationships. For AI agents, semantic memory is often used to personalize applications by remembering facts or concepts from past interactions.
> Note: Not to be confused with "semantic search" which is a technique for finding similar content using "meaning" (usually as embeddings). Semantic memory is a term from psychology, referring to storing facts and knowledge, while semantic search is a method for retrieving information based on meaning rather than exact matches.
#### Profile
Semantic memories can be managed in different ways. For example, memories can be a single, continuously updated "profile" of well-scoped and specific information about a user, organization, or other entity (including the agent itself). A profile is generally just a JSON document with various key-value pairs you've selected to represent your domain.
@@ -253,7 +232,7 @@ Alternatively, memories can be a collection of documents that are continuously u
However, this shifts some complexity memory updating. The model must now _delete_ or _update_ existing items in the list, which can be tricky. In addition, some models may default to over-inserting and others may default to over-updating. See the [Trustcall](https://github.com/hinthornw/trustcall) package for one way to manage this and consider evaluation (e.g., with a tool like [LangSmith](https://docs.smith.langchain.com/tutorials/Developers/evaluation)) to help you tune the behavior.
Working with document collections also shifts complexity to memory **search** over the list. The `Store` currently supports both [semantic search](https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.SearchOp.query) and [filtering by content](https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.SearchOp.filter).
Working with document collections also shifts complexity to memory **search** over the list. The `Store` currently supports [filtering by metadata](https://langchain-ai.github.io/langgraph/reference/store/#storage) and will soon add [semantic search shortly](https://python.langchain.com/docs/concepts/vectorstores/), but selecting the most relevant documents can be tricky as the list grows.
Finally, using a collection of memories can make it challenging to provide comprehensive context to the model. While individual memories may follow a specific schema, this structure might not capture the full context or relationships between memories. As a result, when using these memories to generate responses, the model may lack important contextual information that would be more readily available in a unified profile approach.
+57 -154
View File
@@ -26,88 +26,18 @@ There are several ways to connect agents in a multi-agent system:
- **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.
### Handoffs
In multi-agent architectures, agents can be represented as graph nodes. Each agent node executes its step(s) and decides whether to finish execution or route to another agent, including potentially routing to itself (e.g., running in a loop). A common pattern in multi-agent interactions is handoffs, where one agent hands off control to another. Handoffs allow you to specify:
- __destination__: target agent to navigate to (e.g., name of the node to go to)
- __payload__: [information to pass to that agent](#communication-between-agents) (e.g., state update)
To implement handoffs in LangGraph, agent nodes can return [`Command`](./low_level.md#command) object that allows you to combine both control flow and state updates:
```python
def agent(state) -> Command[Literal["agent", "another_agent"]]:
# the condition for routing/halting can be anything, e.g. LLM tool call / structured output, etc.
goto = get_next_agent(...) # 'agent' / 'another_agent'
return Command(
# Specify which agent to call next
goto=goto,
# Update the graph state
update={"my_state_key": "my_state_value"}
)
```
In a more complex scenario where each agent node is itself a graph (i.e., a [subgraph](./low_level.md#subgraphs)), a node in one of the agent subgraphs might want to navigate to a different agent. For example, if you have two agents, `alice` and `bob` (subgraph nodes in a parent graph), and `alice` needs to navigate to `bob`, you can set `graph=Command.PARENT` in the `Command` object:
```python
def some_node_inside_alice(state)
return Command(
goto="bob",
update={"my_state_key": "my_state_value"},
# specify which graph to navigate to (defaults to the current graph)
graph=Command.PARENT,
)
```
!!! note
If you need to support visualization for subgraphs communicating using `Command(graph=Command.PARENT)` you would need to wrap them in a node function with `Command` annotation, e.g. instead of this:
```python
builder.add_node(alice)
```
you would need to do this:
```python
def call_alice(state) -> Command[Literal["bob"]]:
return alice.invoke(state)
builder.add_node("alice", call_alice)
```
#### Handoffs as tools
One of the most common agent types is a ReAct-style tool-calling agents. For those types of agents, a common pattern is wrapping a handoff in a tool call, e.g.:
```python
def transfer_to_bob(state):
"""Transfer to bob."""
return Command(
goto="bob",
update={"my_state_key": "my_state_value"},
graph=Command.PARENT,
)
```
This is a special case of updating the graph state from tools where in addition the state update, the control flow is included as well.
!!! important
If you want to use tools that return `Command`, you can either use prebuilt [`create_react_agent`][langgraph.prebuilt.chat_agent_executor.create_react_agent] / [`ToolNode`][langgraph.prebuilt.tool_node.ToolNode] components, or implement your own tool-executing node that collects `Command` objects returned by the tools and returns a list of them, e.g.:
```python
def call_tools(state):
...
commands = [tools_by_name[call["name"].invoke(call, config={"coerce_tool_content": False}) for tool_call in tool_calls]
return commands
```
Let's now take a closer look at the different multi-agent architectures.
### Network
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. This architecture is good for problems that do not have a clear hierarchy of agents or a specific sequence in which agents should be called.
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:
- hard to enforce which agent should be called next
- hard to determine how much [information](#shared-message-list) should be passed between the agents
We recommend avoiding this architecture in production and using one of the below architectures instead.
### Supervisor
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
@@ -116,83 +46,39 @@ from langgraph.graph import StateGraph, MessagesState, START
model = ChatOpenAI()
def agent_1(state: MessagesState) -> Command[Literal["agent_2", "agent_3", END]]:
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(...)
# route to one of the agents or exit based on the LLM's decision
# if the LLM returns "__end__", the graph will finish execution
return Command(
goto=response["next_agent"],
update={"messages": [response["content"]]},
)
# the "next" key will be used by the conditional edges to route execution
# to the appropriate agent
return {"next": response["next_agent"]}
def agent_2(state: MessagesState) -> Command[Literal["agent_1", "agent_3", END]]:
response = model.invoke(...)
return Command(
goto=response["next_agent"],
update={"messages": [response["content"]]},
)
def agent_3(state: MessagesState) -> Command[Literal["agent_1", "agent_2", END]]:
...
return Command(
goto=response["next_agent"],
update={"messages": [response["content"]]},
)
builder = StateGraph(MessagesState)
builder.add_node(agent_1)
builder.add_node(agent_2)
builder.add_node(agent_3)
builder.add_edge(START, "agent_1")
network = builder.compile()
```
### Supervisor
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 [`Command`](./low_level.md#command) 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 langgraph.graph import StateGraph, MessagesState, START, END
model = ChatOpenAI()
def supervisor(state: MessagesState) -> Command[Literal["agent_1", "agent_2", END]]:
# 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(...)
# route to one of the agents or exit based on the supervisor's decision
# if the supervisor returns "__end__", the graph will finish execution
return Command(goto=response["next_agent"])
def agent_1(state: MessagesState) -> Command[Literal["supervisor"]]:
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 Command(
goto="supervisor",
update={"messages": [response]},
)
return {"messages": [response]}
def agent_2(state: MessagesState) -> Command[Literal["supervisor"]]:
def agent_2(state: AgentState):
response = model.invoke(...)
return Command(
goto="supervisor",
update={"messages": [response]},
)
return {"messages": [response]}
builder = StateGraph(MessagesState)
builder = StateGraph(AgentState)
builder.add_node(supervisor)
builder.add_node(agent_1)
builder.add_node(agent_2)
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()
```
@@ -240,29 +126,37 @@ To address this, you can design your system _hierarchically_. For example, you c
```python
from typing import Literal
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, MessagesState, START, END
from langgraph.graph import StateGraph, MessagesState, START
model = ChatOpenAI()
# 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: MessagesState) -> Command[Literal["team_1_agent_1", "team_1_agent_2", END]]:
def team_1_supervisor(state: Team1State):
response = model.invoke(...)
return Command(goto=response["next_agent"])
return {"next": response["next_agent"]}
def team_1_agent_1(state: MessagesState) -> Command[Literal["team_1_supervisor"]]:
def team_1_agent_1(state: Team1State):
response = model.invoke(...)
return Command(goto="team_1_supervisor", update={"messages": [response]})
return {"messages": [response]}
def team_1_agent_2(state: MessagesState) -> Command[Literal["team_1_supervisor"]]:
def team_1_agent_2(state: Team1State):
response = model.invoke(...)
return Command(goto="team_1_supervisor", update={"messages": [response]})
return {"messages": [response]}
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")
team_1_graph = team_1_builder.compile()
# define team 2 (same as the single supervisor example above)
@@ -285,22 +179,31 @@ team_2_graph = team_2_builder.compile()
# define top-level supervisor
builder = StateGraph(MessagesState)
def top_level_supervisor(state: MessagesState):
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(...)
# route to one of the teams or exit based on the supervisor's decision
# if the supervisor returns "__end__", the graph will finish execution
return Command(goto=response["next_team"])
# the "next" key will be used by the conditional edges to route execution
# to the appropriate team
return {"next": response["next_team"]}
builder = StateGraph(MessagesState)
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()
```
@@ -310,7 +213,7 @@ In this architecture we add individual agents as graph nodes and define the orde
- **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 (Command)**: in LangGraph you can allow LLMs to decide parts of your application control flow. This can be achieved by using [`Command`](./low_level.md#command). 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.
- **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
+14 -88
View File
@@ -218,16 +218,13 @@ The final thing you can optionally specify when calling `update_state` is `as_no
## Memory Store
![Model of shared state](img/persistence/shared_state.png)
![Update](img/persistence/shared_state.png)
A [state schema](low_level.md#schema) specifies a set of keys that are populated as a graph is executed. As discussed above, state can be written by a checkpointer to a thread at each graph step, enabling state persistence.
But, what if we want to retrain some information *across threads*? Consider the case of a chatbot where we want to retain specific information about the user across *all* chat conversations (e.g., threads) with that user!
With checkpointers alone, we cannot share information across threads. This motivates the need for the [`Store`](../reference/store.md#langgraph.store.base.BaseStore) interface. As an illustration, we can define an `InMemoryStore` to store information about a user across threads. We simply compile our graph with a checkpointer, as before, and with our new `in_memory_store` variable.
### Basic Usage
With checkpointers alone, we cannot share information across threads. This motivates the need for the `Store` interface. As an illustration, we can define an `InMemoryStore` to store information about a user across threads. We simply compile our graph with a checkpointer, as before, and will our new `in_memory_store`.
First, let's showcase this in isolation without using LangGraph.
```python
@@ -242,7 +239,7 @@ user_id = "1"
namespace_for_memory = (user_id, "memories")
```
We use the `store.put` method to save memories to our namespace in the store. When we do this, we specify the namespace, as defined above, and a key-value pair for the memory: the key is simply a unique identifier for the memory (`memory_id`) and the value (a dictionary) is the memory itself.
We use the `store.put` to save memories to our namespace in the store. When we do this, we specify the namespace, as defined above, and a key-value pair for the memory: the key is simply a unique identifier for the memory (`memory_id`) and the value (a dictionary) is the memory itself.
```python
memory_id = str(uuid.uuid4())
@@ -250,7 +247,7 @@ memory = {"food_preference" : "I like pizza"}
in_memory_store.put(namespace_for_memory, memory_id, memory)
```
We can read out memories in our namespace using the `store.search` method, which will return all memories for a given user as a list. The most recent memory is the last in the list.
We can read out memories in our namespace using `store.search`, which will return all memories for a given user as a list. The most recent memory is the last in the list.
```python
memories = in_memory_store.search(namespace_for_memory)
@@ -262,69 +259,16 @@ memories[-1].dict()
'updated_at': '2024-10-02T17:22:31.590605+00:00'}
```
Each memory type is a Python class ([`Item`](https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.Item)) with certain attributes. We can access it as a dictionary by converting via `.dict` as above.
Each memory type is a Python class with certain attributes. We can access it as a dictionary by converting via `.dict` as above.
The attributes it has are:
- `value`: The value (itself a dictionary) of this memory
- `key`: A unique key for this memory in this namespace
- `key`: The UUID for this memory in this namespace
- `namespace`: A list of strings, the namespace of this memory type
- `created_at`: Timestamp for when this memory was created
- `updated_at`: Timestamp for when this memory was updated
### Semantic Search
Beyond simple retrieval, the store also supports semantic search, allowing you to find memories based on meaning rather than exact matches. To enable this, configure the store with an embedding model:
```python
from langchain.embeddings import init_embeddings
store = InMemoryStore(
index={
"embed": init_embeddings("openai:text-embedding-3-small"), # Embedding provider
"dims": 1536, # Embedding dimensions
"fields": ["food_preference", "$"] # Fields to embed
}
)
```
Now when searching, you can use natural language queries to find relevant memories:
```python
# Find memories about food preferences
# (This can be done after putting memories into the store)
memories = store.search(
namespace_for_memory,
query="What does the user like to eat?",
limit=3 # Return top 3 matches
)
```
You can control which parts of your memories get embedded by configuring the `fields` parameter or by specifying the `index` parameter when storing memories:
```python
# Store with specific fields to embed
store.put(
namespace_for_memory,
str(uuid.uuid4()),
{
"food_preference": "I love Italian cuisine",
"context": "Discussing dinner plans"
},
index=["food_preference"] # Only embed "food_preferences" field
)
# Store without embedding (still retrievable, but not searchable)
store.put(
namespace_for_memory,
str(uuid.uuid4()),
{"system_info": "Last updated: 2024-01-01"},
index=False
)
```
### Using in LangGraph
With this all in place, we use the `in_memory_store` in LangGraph. The `in_memory_store` works hand-in-hand with the checkpointer: the checkpointer saves state to threads, as discussed above, and the `in_memory_store` allows us to store arbitrary information for access *across* threads. We compile the graph with both the checkpointer and the `in_memory_store` as follows.
With this all in place, we use the `in_memory_store` in LangGraph. The `in_memory_store` works hand-in-hand with the checkpointer: the checkpointer saves state to threads, as discussed above, and the the `in_memory_store` allows us to store arbitrary information for access *across* threads. We compile the graph with both the checkpointer and the `in_memory_store` as follows.
```python
from langgraph.checkpoint.memory import MemorySaver
@@ -352,7 +296,7 @@ for update in graph.stream(
print(update)
```
We can access the `in_memory_store` and the `user_id` in *any node* by passing `store: BaseStore` and `config: RunnableConfig` as node arguments. Here's how we might use semantic search in a node to find relevant memories:
We can access the `in_memory_store` and the `user_id` in *any node* by passing `store: BaseStore` and `config: RunnableConfig` as node arguments. Just as we saw above, simply use the `put` method to save memories to the store.
```python
def update_memory(state: MessagesState, config: RunnableConfig, *, store: BaseStore):
@@ -373,7 +317,7 @@ def update_memory(state: MessagesState, config: RunnableConfig, *, store: BaseSt
```
As we showed above, we can also access the store in any node and use the `store.search` method to get memories. Recall the the memories are returned as a list of objects that can be converted to a dictionary.
As we showed above, we can also access the store in any node and use `search` to get memories. Recall the the memories are returned as a list of objects that can be converted to a dictionary.
```python
memories[-1].dict()
@@ -388,15 +332,12 @@ We can access the memories and use them in our model call.
```python
def call_model(state: MessagesState, config: RunnableConfig, *, store: BaseStore):
# Get the user id from the config
user_id = config["configurable"]["user_id"]
# Search based on the most recent message
memories = store.search(
namespace,
query=state["messages"][-1].content,
limit=3
)
# Get the memories for the user from the store
memories = store.search(("memories", user_id))
info = "\n".join([d.value["memory"] for d in memories])
# ... Use memories in the model call
@@ -415,22 +356,7 @@ for update in graph.stream(
print(update)
```
When we use the LangGraph Platform, either locally (e.g., in LangGraph Studio) or with LangGraph Cloud, the base store is available to use by default and does not need to be specified during graph compilation. To enable semantic search, however, you **do** need to configure the indexing settings in your `langgraph.json` file. For example:
```json
{
...
"store": {
"index": {
"embed": "openai:text-embeddings-3-small",
"dims": 1536,
"fields": ["$"]
}
}
}
```
See the [deployment guide](../cloud/deployment/semantic_search.md) for more details and configuration options.
When we use the LangGraph API, either locally (e.g., in LangGraph Studio) or with LangGraph Cloud, the memory store is available to use by default and does not need to be specified during graph compilation.
## Checkpointer libraries
@@ -479,4 +405,4 @@ Lastly, checkpointing also provides fault-tolerance and error recovery: if one o
#### Pending writes
Additionally, when a graph node fails mid-execution at a given superstep, LangGraph stores pending checkpoint writes from any other nodes that completed successfully at that superstep, so that whenever we resume graph execution from that superstep we don't re-run the successful nodes.
Additionally, when a graph node fails mid-execution at a given superstep, LangGraph stores pending checkpoint writes from any other nodes that completed successfully at that superstep, so that whenever we resume graph execution from that superstep we don't re-run the successful nodes.
+6 -49
View File
@@ -1,21 +1,14 @@
# Template Applications
!!! note Prerequisites
- [LangGraph Studio](./langgraph_studio.md)
Templates are open source reference applications designed to help you get started quickly when building with LangGraph. They provide working examples of common agentic workflows that can be customized to your needs.
You can create an application from a template using the LangGraph CLI.
Templates can be accessed via [LangGraph Studio (macOS only)](langgraph_studio.md), or cloned directly from Github. You can download LangGraph Studio and see available templates [here](https://studio.langchain.com/).
!!! info "Requirements"
- Python >= 3.11
- [LangGraph CLI](https://langchain-ai.github.io/langgraph/cloud/reference/cli/): Requires langchain-cli[inmem] >= 0.1.58
## Install the LangGraph CLI
```bash
pip install "langgraph-cli[inmem]==0.1.58" python-dotenv
```
## Available Templates
## Available templates
| Template | Description | Python | JS/TS |
|---------------------------|------------------------------------------------------------------------------------------|------------------------------------------------------------------|---------------------------------------------------------------------|
@@ -24,39 +17,3 @@ pip install "langgraph-cli[inmem]==0.1.58" python-dotenv
| **Memory Agent** | A ReAct-style agent with an additional tool to store memories for use across threads. | [Repo](https://github.com/langchain-ai/memory-agent) | [Repo](https://github.com/langchain-ai/memory-agent-js) |
| **Retrieval Agent** | An agent that includes a retrieval-based question-answering system. | [Repo](https://github.com/langchain-ai/retrieval-agent-template) | [Repo](https://github.com/langchain-ai/retrieval-agent-template-js) |
| **Data-Enrichment Agent** | An agent that performs web searches and organizes its findings into a structured format. | [Repo](https://github.com/langchain-ai/data-enrichment) | [Repo](https://github.com/langchain-ai/data-enrichment-js) |
## 🌱 Create a LangGraph App
To create a new app from a template, use the `langgraph new` command.
```bash
langgraph new
```
## Next Steps
Review the `README.md` file in the root of your new LangGraph app for more information about the template and how to customize it.
After configuring the app properly and adding your API keys, you can start the app using the LangGraph CLI:
```bash
langgraph dev
```
See the following guides for more information on how to deploy your app:
- **[Launch Local LangGraph Server](../tutorials/langgraph-platform/local-server.md)**: This quick start guide shows how to start a LangGraph Server locally for the **ReAct Agent** template. The steps are similar for other templates.
- **[Deploy to LangGraph Cloud](../cloud/quick_start.md)**: Deploy your LangGraph app using LangGraph Cloud.
### LangGraph Framework
- **[LangGraph Concepts](../concepts/index.md)**: Learn the foundational concepts of LangGraph.
- **[LangGraph How-to Guides](../how-tos/index.md)**: Guides for common tasks with LangGraph.
### 📚 Learn More about LangGraph Platform
Expand your knowledge with these resources:
- **[LangGraph Platform Concepts](../concepts/index.md#langgraph-platform)**: Understand the foundational concepts of the LangGraph Platform.
- **[LangGraph Platform How-to Guides](../how-tos/index.md#langgraph-platform)**: Discover step-by-step guides to build and deploy applications.
File diff suppressed because one or more lines are too long
@@ -41,9 +41,6 @@
" <p>\n",
" Support for the <code><a href=\"https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.BaseStore\">Store</a></code> API that is used in this guide was added in LangGraph <code>v0.2.32</code>.\n",
" </p>\n",
" <p>\n",
" Support for <b>index</b> and <b>query</b> arguments of the <code><a href=\"https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.BaseStore\">Store</a></code> API that is used in this guide was added in LangGraph <code>v0.2.54</code>.\n",
" </p>\n",
"</div>\n",
"\n",
"## Setup\n",
@@ -117,7 +114,7 @@
"\n",
"Importantly, to determine the user, we will be passing `user_id` via the config keyword argument of the node function.\n",
"\n",
"Let's first define an `InMemoryStore` already populated with some memories about the users."
"Let's first define an `InMemoryStore` which is already populated with some memories about the users."
]
},
{
@@ -128,14 +125,8 @@
"outputs": [],
"source": [
"from langgraph.store.memory import InMemoryStore\n",
"from langchain_openai import OpenAIEmbeddings\n",
"\n",
"in_memory_store = InMemoryStore(\n",
" index={\n",
" \"embed\": OpenAIEmbeddings(model=\"text-embedding-3-small\"),\n",
" \"dims\": 1536,\n",
" }\n",
")"
"in_memory_store = InMemoryStore()"
]
},
{
@@ -172,7 +163,7 @@
"def call_model(state: MessagesState, config: RunnableConfig, *, store: BaseStore):\n",
" user_id = config[\"configurable\"][\"user_id\"]\n",
" namespace = (\"memories\", user_id)\n",
" memories = store.search(namespace, query=str(state[\"messages\"][-1].content))\n",
" memories = store.search(namespace)\n",
" info = \"\\n\".join([d.value[\"data\"] for d in memories])\n",
" system_msg = f\"You are a helpful assistant talking to the user. User info: {info}\"\n",
"\n",
+9 -20
View File
@@ -20,7 +20,6 @@ These how-to guides show how to achieve that controllability.
- [How to create branches for parallel execution](branching.ipynb)
- [How to create map-reduce branches for parallel execution](map-reduce.ipynb)
- [How to control graph recursion limit](recursion-limit.ipynb)
- [How to combine control flow and state updates with Command](command.ipynb)
### Persistence
@@ -40,8 +39,6 @@ LangGraph makes it easy to manage conversation [memory](../concepts/memory.md) i
- [How to manage conversation history](memory/manage-conversation-history.ipynb)
- [How to delete messages](memory/delete-messages.ipynb)
- [How to add summary conversation memory](memory/add-summary-conversation-history.ipynb)
- [How to add long-term memory (cross-thread)](cross-thread-persistence.ipynb)
- [How to use semantic search for long-term memory](memory/semantic-search.ipynb)
### Human-in-the-loop
@@ -73,7 +70,7 @@ you to involve humans in the decision-making process of your graph. These how-to
### Tool calling
[Tool calling](https://python.langchain.com/docs/concepts/tool_calling/) is a type of chat model API that accepts tool schemas, along with messages, as input and returns invocations of those tools as part of the output message.
[Tool calling](https://python.langchain.com/docs/concepts/tool_calling/) is a type of chat model API that accepts tool schemas, along with messages, as input and returns invocations of those tools as part of the output message.
These how-to guides show common patterns for tool calling with LangGraph:
@@ -81,7 +78,6 @@ These how-to guides show common patterns for tool calling with LangGraph:
- [How to handle tool calling errors](tool-calling-errors.ipynb)
- [How to pass runtime values to tools](pass-run-time-values-to-tools.ipynb)
- [How to pass config to tools](pass-config-to-tools.ipynb)
- [How to update graph state from tools](update-state-from-tools.ipynb)
- [How to handle large numbers of tools](many-tools.ipynb)
### Subgraphs
@@ -92,12 +88,6 @@ These how-to guides show common patterns for tool calling with LangGraph:
- [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)
### Multi-agent
- [How to build a multi-agent network](multi-agent-network.ipynb)
See the [multi-agent tutorials](../tutorials/index.md#multi-agent-systems) for implementations of other multi-agent architectures.
### State Management
- [How to use Pydantic model as state](state-model.ipynb)
@@ -128,13 +118,12 @@ These guides show how to use the prebuilt ReAct agent:
- [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 add semantic search for long-term memory to a ReAct agent](memory/semantic-search.ipynb#using-in-create-react-agent)
## LangGraph Platform
This section includes how-to guides for LangGraph Platform.
LangGraph Platform is a commercial solution for deploying agentic applications in production, built on the open-source LangGraph framework.
LangGraph Platform is a commercial solution for deploying agentic applications in production, built on the open-source LangGraph framework.
The LangGraph Platform offers a few different deployment options described in the [deployment options guide](../concepts/deployment_options.md).
@@ -150,7 +139,6 @@ Learn how to set up your app for deployment to LangGraph Platform:
- [How to set up app for deployment (requirements.txt)](../cloud/deployment/setup.md)
- [How to set up app for deployment (pyproject.toml)](../cloud/deployment/setup_pyproject.md)
- [How to set up app for deployment (JavaScript)](../cloud/deployment/setup_javascript.md)
- [How to add semantic search](../cloud/deployment/semantic_search.md)
- [How to customize Dockerfile](../cloud/deployment/custom_docker.md)
- [How to test locally](../cloud/deployment/test_locally.md)
- [How to rebuild graph at runtime](../cloud/deployment/graph_rebuild.md)
@@ -162,8 +150,8 @@ LangGraph applications can be deployed using LangGraph Cloud, which provides a r
- [How to deploy to LangGraph cloud](../cloud/deployment/cloud.md)
- [How to deploy to a self-hosted environment](./deploy-self-hosted.md)
- [How to interact with the deployment using RemoteGraph](./use-remote-graph.md)
- [How to interact with the deployment using RemoteGraph](./use-remote-graph.md)
### Assistants
[Assistants](../concepts/assistants.md) is a configured instance of a template.
@@ -208,7 +196,7 @@ When designing complex graphs, relying entirely on the LLM for decision-making c
### Double-texting
Graph execution can take a while, and sometimes users may change their mind about the input they wanted to send before their original input has finished running. For example, a user might notice a typo in their original request and will edit the prompt and resend it. Deciding what to do in these cases is important for ensuring a smooth user experience and preventing your graphs from behaving in unexpected ways.
Graph execution can take a while, and sometimes users may change their mind about the input they wanted to send before their original input has finished running. For example, a user might notice a typo in their original request and will edit the prompt and resend it. Deciding what to do in these cases is important for ensuring a smooth user experience and preventing your graphs from behaving in unexpected ways.
- [How to use the interrupt option](../cloud/how-tos/interrupt_concurrent.md)
- [How to use the rollback option](../cloud/how-tos/rollback_concurrent.md)
@@ -228,9 +216,8 @@ Graph execution can take a while, and sometimes users may change their mind abou
LangGraph Studio is a built-in UI for visualizing, testing, and debugging your agents.
- [How to connect to a LangGraph Cloud deployment](../cloud/how-tos/test_deployment.md)
- [How to connect to a local dev server](../how-tos/local-studio.md)
- [How to connect to a local deployment (Docker)](../cloud/how-tos/test_local_deployment.md)
- [How to test your graph in LangGraph Studio (MacOS only)](../cloud/how-tos/invoke_studio.md)
- [How to connect to a local deployment](../cloud/how-tos/test_local_deployment.md)
- [How to test your graph in LangGraph Studio](../cloud/how-tos/invoke_studio.md)
- [How to interact with threads in LangGraph Studio](../cloud/how-tos/threads_studio.md)
## Troubleshooting
@@ -242,3 +229,5 @@ These are the guides for resolving common errors you may find while building wit
- [INVALID_GRAPH_NODE_RETURN_VALUE](../troubleshooting/errors/INVALID_GRAPH_NODE_RETURN_VALUE.md)
- [MULTIPLE_SUBGRAPHS](../troubleshooting/errors/MULTIPLE_SUBGRAPHS.md)
- [INVALID_CHAT_HISTORY](../troubleshooting/errors/INVALID_CHAT_HISTORY.md)
@@ -1,532 +0,0 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# How to add semantic search to your agent's memory\n",
"\n",
"This guide shows how to enable semantic search in your agent's memory store. This lets search for items in the store by semantic similarity.\n",
"\n",
"!!! tip Prerequisites\n",
" This guide assumes familiarity with the [memory in LangGraph](https://langchain-ai.github.io/langgraph/concepts/memory/).\n",
"\n",
"First, install this guide's prerequisites."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%%capture --no-stderr\n",
"%pip install -U langgraph langchain-openai langchain"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import getpass\n",
"import os\n",
"\n",
"\n",
"def _set_env(var: str):\n",
" if not os.environ.get(var):\n",
" os.environ[var] = getpass.getpass(f\"{var}: \")\n",
"\n",
"\n",
"_set_env(\"OPENAI_API_KEY\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, create the store with an [index configuration](https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.IndexConfig). By default, stores are configured without semantic/vector search. You can opt in to indexing items when creating the store by providing an [IndexConfig](https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.IndexConfig) to the store's constructor. If your store class does not implement this interface, or if you do not pass in an index configuration, semantic search is disabled, and all `index` arguments passed to `put` or `aput` will have no effect. Below is an example."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/var/folders/gf/6rnp_mbx5914kx7qmmh7xzmw0000gn/T/ipykernel_83572/2318027494.py:5: LangChainBetaWarning: The function `init_embeddings` is in beta. It is actively being worked on, so the API may change.\n",
" embeddings = init_embeddings(\"openai:text-embedding-3-small\")\n"
]
}
],
"source": [
"from langchain.embeddings import init_embeddings\n",
"from langgraph.store.memory import InMemoryStore\n",
"\n",
"# Create store with semantic search enabled\n",
"embeddings = init_embeddings(\"openai:text-embedding-3-small\")\n",
"store = InMemoryStore(\n",
" index={\n",
" \"embed\": embeddings,\n",
" \"dims\": 1536,\n",
" }\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now let's store some memories:"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"# Store some memories\n",
"store.put((\"user_123\", \"memories\"), \"1\", {\"text\": \"I love pizza\"})\n",
"store.put((\"user_123\", \"memories\"), \"2\", {\"text\": \"I prefer Italian food\"})\n",
"store.put((\"user_123\", \"memories\"), \"3\", {\"text\": \"I don't like spicy food\"})\n",
"store.put((\"user_123\", \"memories\"), \"3\", {\"text\": \"I am studying econometrics\"})\n",
"store.put((\"user_123\", \"memories\"), \"3\", {\"text\": \"I am a plumber\"})"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Search memories using natural language:"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Memory: I prefer Italian food (similarity: 0.46482669521168163)\n",
"Memory: I love pizza (similarity: 0.35514845174380766)\n",
"Memory: I am a plumber (similarity: 0.155698702336571)\n"
]
}
],
"source": [
"# Find memories about food preferences\n",
"memories = store.search((\"user_123\", \"memories\"), query=\"I like food?\", limit=5)\n",
"\n",
"for memory in memories:\n",
" print(f'Memory: {memory.value[\"text\"]} (similarity: {memory.score})')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Using in your agent\n",
"\n",
"Add semantic search to any node by injecting the store."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"What are you in the mood for? Since you love Italian food and pizza, would you like to order a pizza or try making one at home?"
]
}
],
"source": [
"from typing import Optional\n",
"\n",
"from langchain.chat_models import init_chat_model\n",
"from langgraph.store.base import BaseStore\n",
"\n",
"from langgraph.graph import START, MessagesState, StateGraph\n",
"\n",
"llm = init_chat_model(\"openai:gpt-4o-mini\")\n",
"\n",
"\n",
"def chat(state, *, store: BaseStore):\n",
" # Search based on user's last message\n",
" items = store.search(\n",
" (\"user_123\", \"memories\"), query=state[\"messages\"][-1].content, limit=2\n",
" )\n",
" memories = \"\\n\".join(item.value[\"text\"] for item in items)\n",
" memories = f\"## Memories of user\\n{memories}\" if memories else \"\"\n",
" response = llm.invoke(\n",
" [\n",
" {\"role\": \"system\", \"content\": f\"You are a helpful assistant.\\n{memories}\"},\n",
" *state[\"messages\"],\n",
" ]\n",
" )\n",
" return {\"messages\": [response]}\n",
"\n",
"\n",
"builder = StateGraph(MessagesState)\n",
"builder.add_node(chat)\n",
"builder.add_edge(START, \"chat\")\n",
"graph = builder.compile(store=store)\n",
"\n",
"for message, metadata in graph.stream(\n",
" input={\"messages\": [{\"role\": \"user\", \"content\": \"I'm hungry\"}]},\n",
" stream_mode=\"messages\",\n",
"):\n",
" print(message.content, end=\"\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Using in `create_react_agent`\n",
"\n",
"Add semantic search to your tool calling agent by injecting the store in the `state_modifier`. You can also use the store in a tool to let your agent manually store or search for memories."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"import uuid\n",
"from typing import Optional\n",
"\n",
"from langchain.chat_models import init_chat_model\n",
"from langchain_core.tools import InjectedToolArg\n",
"from langgraph.store.base import BaseStore\n",
"from typing_extensions import Annotated\n",
"\n",
"from langgraph.prebuilt import create_react_agent\n",
"\n",
"\n",
"def prepare_messages(state, *, store: BaseStore):\n",
" # Search based on user's last message\n",
" items = store.search(\n",
" (\"user_123\", \"memories\"), query=state[\"messages\"][-1].content, limit=2\n",
" )\n",
" memories = \"\\n\".join(item.value[\"text\"] for item in items)\n",
" memories = f\"## Memories of user\\n{memories}\" if memories else \"\"\n",
" return [\n",
" {\"role\": \"system\", \"content\": f\"You are a helpful assistant.\\n{memories}\"}\n",
" ] + state[\"messages\"]\n",
"\n",
"\n",
"# You can also use the store directly within a tool!\n",
"def upsert_memory(\n",
" content: str,\n",
" *,\n",
" memory_id: Optional[uuid.UUID] = None,\n",
" store: Annotated[BaseStore, InjectedToolArg],\n",
"):\n",
" \"\"\"Upsert a memory in the database.\"\"\"\n",
" # The LLM can use this tool to store a new memory\n",
" mem_id = memory_id or uuid.uuid4()\n",
" store.put(\n",
" (\"user_123\", \"memories\"),\n",
" key=str(mem_id),\n",
" value={\"text\": content},\n",
" )\n",
" return f\"Stored memory {mem_id}\"\n",
"\n",
"\n",
"agent = create_react_agent(\n",
" init_chat_model(\"openai:gpt-4o-mini\"),\n",
" tools=[upsert_memory],\n",
" # The state_modifier is run to prepare the messages for the LLM. It is called\n",
" # right before each LLM call\n",
" state_modifier=prepare_messages,\n",
" store=store,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"What are you in the mood for? Since you love Italian food and pizza, maybe something in that realm would be great! Would you like suggestions for a specific dish or restaurant?"
]
}
],
"source": [
"for message, metadata in agent.stream(\n",
" input={\"messages\": [{\"role\": \"user\", \"content\": \"I'm hungry\"}]},\n",
" stream_mode=\"messages\",\n",
"):\n",
" print(message.content, end=\"\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Advanced Usage\n",
"\n",
"#### Multi-vector indexing\n",
"\n",
"Store and search different aspects of memories separately to improve recall or omit certain fields from being indexed."
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Expect mem 2\n",
"Item: mem2; Score (0.5895009051396596)\n",
"Memory: Ate alone at home\n",
"Emotion: felt a bit lonely\n",
"\n",
"Expect mem1\n",
"Item: mem1; Score (0.6207546534134083)\n",
"Memory: Had pizza with friends at Mario's\n",
"Emotion: felt happy and connected\n",
"\n",
"Expect random lower score (ravioli not indexed)\n",
"Item: mem1; Score (0.2686278787315685)\n",
"Memory: Had pizza with friends at Mario's\n",
"Emotion: felt happy and connected\n",
"\n"
]
}
],
"source": [
"# Configure store to embed both memory content and emotional context\n",
"store = InMemoryStore(\n",
" index={\"embed\": embeddings, \"dims\": 1536, \"fields\": [\"memory\", \"emotional_context\"]}\n",
")\n",
"# Store memories with different content/emotion pairs\n",
"store.put(\n",
" (\"user_123\", \"memories\"),\n",
" \"mem1\",\n",
" {\n",
" \"memory\": \"Had pizza with friends at Mario's\",\n",
" \"emotional_context\": \"felt happy and connected\",\n",
" \"this_isnt_indexed\": \"I prefer ravioli though\",\n",
" },\n",
")\n",
"store.put(\n",
" (\"user_123\", \"memories\"),\n",
" \"mem2\",\n",
" {\n",
" \"memory\": \"Ate alone at home\",\n",
" \"emotional_context\": \"felt a bit lonely\",\n",
" \"this_isnt_indexed\": \"I like pie\",\n",
" },\n",
")\n",
"\n",
"# Search focusing on emotional state - matches mem2\n",
"results = store.search(\n",
" (\"user_123\", \"memories\"), query=\"times they felt isolated\", limit=1\n",
")\n",
"print(\"Expect mem 2\")\n",
"for r in results:\n",
" print(f\"Item: {r.key}; Score ({r.score})\")\n",
" print(f\"Memory: {r.value['memory']}\")\n",
" print(f\"Emotion: {r.value['emotional_context']}\\n\")\n",
"\n",
"# Search focusing on social eating - matches mem1\n",
"print(\"Expect mem1\")\n",
"results = store.search((\"user_123\", \"memories\"), query=\"fun pizza\", limit=1)\n",
"for r in results:\n",
" print(f\"Item: {r.key}; Score ({r.score})\")\n",
" print(f\"Memory: {r.value['memory']}\")\n",
" print(f\"Emotion: {r.value['emotional_context']}\\n\")\n",
"\n",
"print(\"Expect random lower score (ravioli not indexed)\")\n",
"results = store.search((\"user_123\", \"memories\"), query=\"ravioli\", limit=1)\n",
"for r in results:\n",
" print(f\"Item: {r.key}; Score ({r.score})\")\n",
" print(f\"Memory: {r.value['memory']}\")\n",
" print(f\"Emotion: {r.value['emotional_context']}\\n\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Override fields at storage time\n",
"You can override which fields to embed when storing a specific memory using `put(..., index=[...fields])`, regardless of the store's default configuration."
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Expect mem1\n",
"Item: mem1; Score (0.3374968677940555)\n",
"Memory: I love spicy food\n",
"Context: At a Thai restaurant\n",
"\n",
"Expect mem2\n",
"Item: mem2; Score (0.36784461593247436)\n",
"Memory: The restaurant was too loud\n",
"Context: Dinner at an Italian place\n",
"\n"
]
}
],
"source": [
"store = InMemoryStore(\n",
" index={\n",
" \"embed\": embeddings,\n",
" \"dims\": 1536,\n",
" \"fields\": [\"memory\"],\n",
" } # Default to embed memory field\n",
")\n",
"\n",
"# Store one memory with default indexing\n",
"store.put(\n",
" (\"user_123\", \"memories\"),\n",
" \"mem1\",\n",
" {\"memory\": \"I love spicy food\", \"context\": \"At a Thai restaurant\"},\n",
")\n",
"\n",
"# Store another overriding which fields to embed\n",
"store.put(\n",
" (\"user_123\", \"memories\"),\n",
" \"mem2\",\n",
" {\"memory\": \"The restaurant was too loud\", \"context\": \"Dinner at an Italian place\"},\n",
" index=[\"context\"], # Override: only embed the context\n",
")\n",
"\n",
"# Search about food - matches mem1 (using default field)\n",
"print(\"Expect mem1\")\n",
"results = store.search(\n",
" (\"user_123\", \"memories\"), query=\"what food do they like\", limit=1\n",
")\n",
"for r in results:\n",
" print(f\"Item: {r.key}; Score ({r.score})\")\n",
" print(f\"Memory: {r.value['memory']}\")\n",
" print(f\"Context: {r.value['context']}\\n\")\n",
"\n",
"# Search about restaurant atmosphere - matches mem2 (using overridden field)\n",
"print(\"Expect mem2\")\n",
"results = store.search(\n",
" (\"user_123\", \"memories\"), query=\"restaurant environment\", limit=1\n",
")\n",
"for r in results:\n",
" print(f\"Item: {r.key}; Score ({r.score})\")\n",
" print(f\"Memory: {r.value['memory']}\")\n",
" print(f\"Context: {r.value['context']}\\n\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Disable Indexing for Specific Memories\n",
"\n",
"Some memories shouldn't be searchable by content. You can disable indexing for these while still storing them using \n",
"`put(..., index=False)`. Example:"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Expect mem1\n",
"Item: mem1; Score (0.32269984224327286)\n",
"Memory: I love chocolate ice cream\n",
"Type: preference\n",
"\n",
"Expect low score (mem2 not indexed)\n",
"Item: mem1; Score (0.010241633698527089)\n",
"Memory: I love chocolate ice cream\n",
"Type: preference\n",
"\n"
]
}
],
"source": [
"store = InMemoryStore(index={\"embed\": embeddings, \"dims\": 1536, \"fields\": [\"memory\"]})\n",
"\n",
"# Store a normal indexed memory\n",
"store.put(\n",
" (\"user_123\", \"memories\"),\n",
" \"mem1\",\n",
" {\"memory\": \"I love chocolate ice cream\", \"type\": \"preference\"},\n",
")\n",
"\n",
"# Store a system memory without indexing\n",
"store.put(\n",
" (\"user_123\", \"memories\"),\n",
" \"mem2\",\n",
" {\"memory\": \"User completed onboarding\", \"type\": \"system\"},\n",
" index=False, # Disable indexing entirely\n",
")\n",
"\n",
"# Search about food preferences - finds mem1\n",
"print(\"Expect mem1\")\n",
"results = store.search((\"user_123\", \"memories\"), query=\"what food preferences\", limit=1)\n",
"for r in results:\n",
" print(f\"Item: {r.key}; Score ({r.score})\")\n",
" print(f\"Memory: {r.value['memory']}\")\n",
" print(f\"Type: {r.value['type']}\\n\")\n",
"\n",
"# Search about onboarding - won't find mem2 (not indexed)\n",
"print(\"Expect low score (mem2 not indexed)\")\n",
"results = store.search((\"user_123\", \"memories\"), query=\"onboarding status\", limit=1)\n",
"for r in results:\n",
" print(f\"Item: {r.key}; Score ({r.score})\")\n",
" print(f\"Memory: {r.value['memory']}\")\n",
" print(f\"Type: {r.value['type']}\\n\")"
]
}
],
"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.11.2"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
File diff suppressed because one or more lines are too long
@@ -1,383 +0,0 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "7c58c957-83d8-44ff-8580-a9b3dd39a0a9",
"metadata": {},
"source": [
"# How to update graph state from tools"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "95f30587-8dd2-40be-920d-59539089c09f",
"metadata": {},
"source": [
"!!! info \"Prerequisites\"\n",
" This guide assumes familiarity with the following:\n",
" \n",
" - [Command](../../concepts/low_level/#command)\n",
"\n",
"A common use case is updating graph state from inside a tool. For example, in a customer support application you might want to look up customer account number or ID in the beginning of the conversation. To update the graph state from the tool, you can return `Command(update={\"my_custom_key\": \"foo\", \"messages\": [...]})` from the tool:\n",
"\n",
"```python\n",
"@tool\n",
"def lookup_user_info(tool_call_id: Annotated[str, InjectedToolCallId], config: RunnableConfig):\n",
" \"\"\"Use this to look up user information to better assist them with their questions.\"\"\"\n",
" user_info = get_user_info(config)\n",
" return Command(\n",
" update={\n",
" # update the state keys\n",
" \"user_info\": user_info,\n",
" # update the message history\n",
" \"messages\": [ToolMessage(\"Successfully looked up user information\", tool_call_id=tool_call_id)]\n",
" }\n",
" )\n",
"```\n",
"\n",
"!!! important\n",
"\n",
" If you want to use tools that return `Command` and update graph state, you can either use prebuilt [`create_react_agent`][langgraph.prebuilt.chat_agent_executor.create_react_agent] / [`ToolNode`][langgraph.prebuilt.tool_node.ToolNode] components, or implement your own tool-executing node that collects `Command` objects returned by the tools and returns a list of them, e.g.:\n",
" \n",
" ```python\n",
" def call_tools(state):\n",
" ...\n",
" commands = [tools_by_name[call[\"name\"].invoke(call, config={\"coerce_tool_content\": False}) for tool_call in tool_calls]\n",
" return commands\n",
" ```\n",
"\n",
"This guide shows how you can do this using LangGraph's prebuilt components ([`create_react_agent`][langgraph.prebuilt.chat_agent_executor.create_react_agent] / [`ToolNode`][langgraph.prebuilt.tool_node.ToolNode]).\n",
"\n",
"!!! note\n",
"\n",
" Support for tools that return [`Command`][langgraph.types.Command] was added in LangGraph `v0.2.57`.\n",
"\n",
"## Setup\n",
"\n",
"First, let's install the required packages and set our API keys:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "64500eca-1cdc-43d9-9401-f4cd9999881f",
"metadata": {},
"outputs": [],
"source": [
"%%capture --no-stderr\n",
"%pip install -U langgraph langchain-openai"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "a3f92fb2-9175-47fa-9c7d-ad5f44bfd20e",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Please provide your OPENAI_API_KEY ········\n"
]
}
],
"source": [
"import os\n",
"import getpass\n",
"\n",
"\n",
"def _set_if_undefined(var: str):\n",
" if not os.environ.get(var):\n",
" os.environ[var] = getpass.getpass(f\"Please provide your {var}\")\n",
"\n",
"\n",
"_set_if_undefined(\"OPENAI_API_KEY\")"
]
},
{
"cell_type": "markdown",
"id": "caf6ff9f-c1e6-499e-a230-9fa231ea7d2f",
"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": "10e9a9c6-fa3f-416c-bac0-3e58d7259908",
"metadata": {},
"source": [
"Let's create a simple ReAct style agent that can look up user information and personalize the response based on the user info."
]
},
{
"cell_type": "markdown",
"id": "4255b9b9-cf67-4cc3-8018-1708f5dfcfd2",
"metadata": {},
"source": [
"## Define tool"
]
},
{
"cell_type": "markdown",
"id": "7de6b010-aab1-4fe8-8251-907fcae78583",
"metadata": {},
"source": [
"First, let's define the tool that we'll be using to look up user information. We'll use a naive implementation that simply looks user information up using a dictionary:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8d070c9f-6e61-4724-85dc-ac4531b9c79a",
"metadata": {},
"outputs": [],
"source": [
"USER_INFO = [\n",
" {\"user_id\": \"1\", \"name\": \"Bob Dylan\", \"location\": \"New York, NY\"},\n",
" {\"user_id\": \"2\", \"name\": \"Taylor Swift\", \"location\": \"Beverly Hills, CA\"},\n",
"]\n",
"\n",
"USER_ID_TO_USER_INFO = {info[\"user_id\"]: info for info in USER_INFO}"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "08d1ecca-ee57-4e97-b8d0-e09de85337d4",
"metadata": {},
"outputs": [],
"source": [
"from langgraph.prebuilt.chat_agent_executor import AgentState\n",
"from langgraph.types import Command\n",
"from langchain_core.tools import tool\n",
"from langchain_core.tools.base import InjectedToolCallId\n",
"from langchain_core.messages import ToolMessage\n",
"from langchain_core.runnables import RunnableConfig\n",
"\n",
"from typing_extensions import Any, Annotated\n",
"\n",
"\n",
"class State(AgentState):\n",
" # user provided\n",
" last_name: str\n",
" # updated by the tool\n",
" user_info: dict[str, Any]\n",
"\n",
"\n",
"@tool\n",
"def lookup_user_info(\n",
" tool_call_id: Annotated[str, InjectedToolCallId], config: RunnableConfig\n",
"):\n",
" \"\"\"Use this to look up user information to better assist them with their questions.\"\"\"\n",
" user_id = config.get(\"configurable\", {}).get(\"user_id\")\n",
" if user_id is None:\n",
" raise ValueError(\"Please provide user ID\")\n",
"\n",
" if user_id not in USER_ID_TO_USER_INFO:\n",
" raise ValueError(f\"User '{user_id}' not found\")\n",
"\n",
" user_info = USER_ID_TO_USER_INFO[user_id]\n",
" return Command(\n",
" update={\n",
" # update the state keys\n",
" \"user_info\": user_info,\n",
" # update the message history\n",
" \"messages\": [\n",
" ToolMessage(\n",
" \"Successfully looked up user information\", tool_call_id=tool_call_id\n",
" )\n",
" ],\n",
" }\n",
" )"
]
},
{
"cell_type": "markdown",
"id": "b99e5f24-5e5e-4a34-baae-467182675bb5",
"metadata": {},
"source": [
"## Define prompt"
]
},
{
"cell_type": "markdown",
"id": "cbb06aea-6654-4245-91f8-af6e8f2b5377",
"metadata": {},
"source": [
"Let's now add personalization: we'll respond differently to the user based on the state values AFTER the state has been updated from the tool. To achieve this, let's define a function that will dynamically construct the system prompt based on the graph state. It will be called ever time the LLM is called and the function output will be passed to the LLM:"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "c553d062-d145-4145-84bd-9b798f7c95c2",
"metadata": {},
"outputs": [],
"source": [
"def state_modifier(state: State):\n",
" user_info = state.get(\"user_info\")\n",
" if user_info is None:\n",
" return state[\"messages\"]\n",
"\n",
" system_msg = (\n",
" f\"User name is {user_info['name']}. User lives in {user_info['location']}\"\n",
" )\n",
" return [{\"role\": \"system\", \"content\": system_msg}] + state[\"messages\"]"
]
},
{
"cell_type": "markdown",
"id": "c5acdd5d-68be-466b-9c21-46cbed91d2bc",
"metadata": {},
"source": [
"## Define graph"
]
},
{
"cell_type": "markdown",
"id": "afb65028-0359-46c8-b09c-ffc90180f759",
"metadata": {},
"source": [
"Finally, let's combine this into a single graph using the prebuilt `create_react_agent`:"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "2d59db29-fd51-4d29-9854-21763a4855e3",
"metadata": {},
"outputs": [],
"source": [
"from langgraph.prebuilt import create_react_agent\n",
"from langchain_openai import ChatOpenAI\n",
"\n",
"model = ChatOpenAI(model=\"gpt-4o\")\n",
"\n",
"agent = create_react_agent(\n",
" model,\n",
" # pass the tool that can update state\n",
" [lookup_user_info],\n",
" state_schema=State,\n",
" # pass dynamic prompt function\n",
" state_modifier=state_modifier,\n",
")"
]
},
{
"cell_type": "markdown",
"id": "0782b8ab-a603-47b8-9a76-77f593402678",
"metadata": {},
"source": [
"## Use it!"
]
},
{
"cell_type": "markdown",
"id": "6165e153-ab28-4404-adea-796c7bd0701b",
"metadata": {},
"source": [
"Let's now try running our agent. We'll need to provide user ID in the config so that our tool knows what information to look up:"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "de34a58b-1765-4b63-a232-d46790aff884",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_7LSUh6ZDvGJAUvlWvXiCK4Gf', 'function': {'arguments': '{}', 'name': 'lookup_user_info'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 11, 'prompt_tokens': 56, 'total_tokens': 67, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_9d50cd990b', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-57eeb216-e35d-4501-aaac-b5c6b26fb17c-0', tool_calls=[{'name': 'lookup_user_info', 'args': {}, 'id': 'call_7LSUh6ZDvGJAUvlWvXiCK4Gf', 'type': 'tool_call'}], usage_metadata={'input_tokens': 56, 'output_tokens': 11, 'total_tokens': 67, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}}\n",
"\n",
"\n",
"{'tools': {'user_info': {'user_id': '1', 'name': 'Bob Dylan', 'location': 'New York, NY'}, 'messages': [ToolMessage(content='Successfully looked up user information', name='lookup_user_info', id='168d8ff8-b021-4c8b-a11a-3b50c30a072c', tool_call_id='call_7LSUh6ZDvGJAUvlWvXiCK4Gf')]}}\n",
"\n",
"\n",
"{'agent': {'messages': [AIMessage(content=\"Hi Bob! Since you're in New York, NY, there are plenty of exciting things to do over the weekend. Here are some suggestions:\\n\\n1. **Explore Central Park**: Take a leisurely walk, rent a bike, or have a picnic in this iconic park.\\n\\n2. **Visit a Museum**: Check out The Metropolitan Museum of Art or the Museum of Modern Art (MoMA) for an enriching cultural experience.\\n\\n3. **Broadway Show**: Catch a Broadway show or an off-Broadway performance for some world-class entertainment.\\n\\n4. **Food Tour**: Explore different neighborhoods like Greenwich Village or Williamsburg for diverse culinary experiences.\\n\\n5. **Brooklyn Bridge Walk**: Take a walk across the Brooklyn Bridge for stunning views of the city skyline.\\n\\n6. **Visit a Rooftop Bar**: Enjoy a drink with a view at one of New York’s many rooftop bars.\\n\\n7. **Explore a New Neighborhood**: Discover the unique charm of areas like SoHo, Chelsea, or Astoria.\\n\\n8. **Live Music**: Check out live music venues for a night of great performances.\\n\\n9. **Art Galleries**: Visit some of the smaller art galleries around Chelsea or the Lower East Side.\\n\\n10. **Attend a Local Event**: Look up any local events or festivals happening this weekend.\\n\\nFeel free to let me know if you want more details on any of these activities!\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 285, 'prompt_tokens': 95, 'total_tokens': 380, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_9d50cd990b', 'finish_reason': 'stop', 'logprobs': None}, id='run-f13ce15b-02b6-40e6-8264-c4d9edd0d03a-0', usage_metadata={'input_tokens': 95, 'output_tokens': 285, 'total_tokens': 380, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}}\n",
"\n",
"\n"
]
}
],
"source": [
"for chunk in agent.stream(\n",
" {\"messages\": [(\"user\", \"hi, what should i do this weekend?\")]},\n",
" # provide user ID in the config\n",
" {\"configurable\": {\"user_id\": \"1\"}},\n",
"):\n",
" print(chunk)\n",
" print(\"\\n\")"
]
},
{
"cell_type": "markdown",
"id": "d9b2281f-269c-41dd-b6b2-4c743f11ffc9",
"metadata": {},
"source": [
"We can see that the model correctly recommended some New York activities for Bob Dylan! Let's try getting recommendations for Taylor Swift:"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "9d71af94-572a-4961-88a7-665e792cf96a",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_5HLtJtzcgmKbtmK6By21wW5Y', 'function': {'arguments': '{}', 'name': 'lookup_user_info'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 11, 'prompt_tokens': 56, 'total_tokens': 67, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_c7ca0ebaca', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-bacacd7d-76cc-4f6b-9e9b-d9e6f00b9391-0', tool_calls=[{'name': 'lookup_user_info', 'args': {}, 'id': 'call_5HLtJtzcgmKbtmK6By21wW5Y', 'type': 'tool_call'}], usage_metadata={'input_tokens': 56, 'output_tokens': 11, 'total_tokens': 67, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}}\n",
"\n",
"\n",
"{'tools': {'user_info': {'user_id': '2', 'name': 'Taylor Swift', 'location': 'Beverly Hills, CA'}, 'messages': [ToolMessage(content='Successfully looked up user information', name='lookup_user_info', id='d81ef31e-6d77-4f13-ae86-e2e6ba567e3d', tool_call_id='call_5HLtJtzcgmKbtmK6By21wW5Y')]}}\n",
"\n",
"\n",
"{'agent': {'messages': [AIMessage(content=\"Hi Taylor! Since you're in Beverly Hills, here are a few suggestions for a fun weekend:\\n\\n1. **Hiking at Runyon Canyon**: Enjoy a scenic hike with beautiful views of Los Angeles. It's a great way to get some exercise and enjoy the outdoors.\\n\\n2. **Visit Rodeo Drive**: Spend some time shopping or window shopping at the famous Rodeo Drive. You might even spot some celebrities!\\n\\n3. **Explore the Getty Center**: Check out the art collections and beautiful gardens at the Getty Center. The architecture and views are stunning.\\n\\n4. **Relax at a Spa**: Treat yourself to a relaxing day at one of Beverly Hills' luxurious spas.\\n\\n5. **Dining Out**: Try a new restaurant or visit your favorite spot for a delicious meal. Beverly Hills has a fantastic dining scene.\\n\\n6. **Attend a Local Event**: Check out any local events or concerts happening this weekend. Beverly Hills often hosts exciting events.\\n\\nEnjoy your weekend!\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 198, 'prompt_tokens': 95, 'total_tokens': 293, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_c7ca0ebaca', 'finish_reason': 'stop', 'logprobs': None}, id='run-2057df76-f192-4c69-a66a-1f0a86bf5d66-0', usage_metadata={'input_tokens': 95, 'output_tokens': 198, 'total_tokens': 293, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}}\n",
"\n",
"\n"
]
}
],
"source": [
"for chunk in agent.stream(\n",
" {\"messages\": [(\"user\", \"hi, what should i do this weekend?\")]},\n",
" {\"configurable\": {\"user_id\": \"2\"}},\n",
"):\n",
" print(chunk)\n",
" print(\"\\n\")"
]
}
],
"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
}
-2
View File
@@ -13,5 +13,3 @@
- PregelExecutableTask
- StateSnapshot
- Send
- Command
- interrupt
-1
View File
@@ -13,7 +13,6 @@ New to LangGraph or LLM app development? Read this material to get up and runnin
- [LangGraph Quickstart](introduction.ipynb): Build a chatbot that can use tools and keep track of conversation history. Add human-in-the-loop capabilities and explore how time-travel works.
- [LangGraph Server Quickstart](langgraph-platform/local-server.md): Launch a LangGraph server locally and interact with it using the REST API and LangGraph Studio Web UI.
- [LangGraph Cloud QuickStart](../cloud/quick_start.md): Deploy a LangGraph app using LangGraph Cloud.
- [LangGraph Template Quickstart](../concepts/template_applications.md): Quickly start building with LangGraph Platform using a template application.
## Use cases 🛠️
+9 -1
View File
@@ -19,7 +19,7 @@
"\n",
"## Setup\n",
"\n",
"First, install the required packages and configure your environment:"
"First, install the required packages:"
]
},
{
@@ -33,6 +33,14 @@
"%pip install -U langgraph langsmith langchain_anthropic"
]
},
{
"cell_type": "markdown",
"id": "a6d1e870-1bc0-4d44-86c0-96681ccf6113",
"metadata": {},
"source": [
"In this tutorial, we'll be "
]
},
{
"cell_type": "code",
"execution_count": 2,
@@ -10,7 +10,7 @@ This is a quick start guide to help you get a LangGraph app up and running local
## Install the LangGraph CLI
```bash
pip install -U "langgraph-cli[inmem]" python-dotenv
pip install "langgraph-cli[inmem]==0.1.58" python-dotenv
```
## 🌱 Create a LangGraph App
@@ -35,10 +35,10 @@ Create a new app from the `react-agent` template. This template is a simple agen
## Install Dependencies
In the root of your new LangGraph app, install the dependencies in `edit` mode so your local changes are used by the server:
In the root of your new LangGraph app, install the dependencies:
```shell
pip install -e .
pip install .
```
## Create a `.env` file
@@ -250,4 +250,4 @@ Access detailed documentation for development and API usage:
- **[LangGraph Server API Reference](../../cloud/reference/api/api_ref.html)**: Explore the LangGraph Server API documentation.
- **[Python SDK Reference](../../cloud/reference/sdk/python_sdk_ref.md)**: Explore the Python SDK API Reference.
- **[JS/TS SDK Reference](../../cloud/reference/sdk/js_ts_sdk_ref.md)**: Explore the Python SDK API Reference.
- **[JS/TS SDK Reference](../../cloud/reference/sdk/js_ts_sdk_ref.md)**: Explore the Python SDK API Reference.
File diff suppressed because one or more lines are too long
@@ -289,10 +289,15 @@
"from langchain_core.language_models.chat_models import BaseChatModel\n",
"\n",
"from langgraph.graph import StateGraph, MessagesState, START, END\n",
"from langgraph.types import Command\n",
"from langchain_core.messages import HumanMessage, trim_messages\n",
"\n",
"\n",
"# The agent state is the input to each node in the graph\n",
"class AgentState(MessagesState):\n",
" # The 'next' field indicates where to route to next\n",
" next: str\n",
"\n",
"\n",
"def make_supervisor_node(llm: BaseChatModel, members: list[str]) -> str:\n",
" options = [\"FINISH\"] + members\n",
" system_prompt = (\n",
@@ -308,17 +313,17 @@
"\n",
" next: Literal[*options]\n",
"\n",
" def supervisor_node(state: MessagesState) -> Command[Literal[*members, \"__end__\"]]:\n",
" def supervisor_node(state: MessagesState) -> MessagesState:\n",
" \"\"\"An LLM-based router.\"\"\"\n",
" messages = [\n",
" {\"role\": \"system\", \"content\": system_prompt},\n",
" ] + state[\"messages\"]\n",
" response = llm.with_structured_output(Router).invoke(messages)\n",
" goto = response[\"next\"]\n",
" if goto == \"FINISH\":\n",
" goto = END\n",
" next_ = response[\"next\"]\n",
" if next_ == \"FINISH\":\n",
" next_ = END\n",
"\n",
" return Command(goto=goto)\n",
" return {\"next\": next_}\n",
"\n",
" return supervisor_node"
]
@@ -358,33 +363,25 @@
"search_agent = create_react_agent(llm, tools=[tavily_tool])\n",
"\n",
"\n",
"def search_node(state: MessagesState) -> Command[Literal[\"supervisor\"]]:\n",
"def search_node(state: AgentState) -> AgentState:\n",
" result = search_agent.invoke(state)\n",
" return Command(\n",
" update={\n",
" \"messages\": [\n",
" HumanMessage(content=result[\"messages\"][-1].content, name=\"search\")\n",
" ]\n",
" },\n",
" # We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
" goto=\"supervisor\",\n",
" )\n",
" return {\n",
" \"messages\": [\n",
" HumanMessage(content=result[\"messages\"][-1].content, name=\"search\")\n",
" ]\n",
" }\n",
"\n",
"\n",
"web_scraper_agent = create_react_agent(llm, tools=[scrape_webpages])\n",
"\n",
"\n",
"def web_scraper_node(state: MessagesState) -> Command[Literal[\"supervisor\"]]:\n",
"def web_scraper_node(state: AgentState) -> AgentState:\n",
" result = web_scraper_agent.invoke(state)\n",
" return Command(\n",
" update={\n",
" \"messages\": [\n",
" HumanMessage(content=result[\"messages\"][-1].content, name=\"web_scraper\")\n",
" ]\n",
" },\n",
" # We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
" goto=\"supervisor\",\n",
" )\n",
" return {\n",
" \"messages\": [\n",
" HumanMessage(content=result[\"messages\"][-1].content, name=\"web_scraper\")\n",
" ]\n",
" }\n",
"\n",
"\n",
"research_supervisor_node = make_supervisor_node(llm, [\"search\", \"web_scraper\"])"
@@ -415,7 +412,14 @@
"research_builder.add_node(\"search\", search_node)\n",
"research_builder.add_node(\"web_scraper\", web_scraper_node)\n",
"\n",
"# Define the control flow\n",
"research_builder.add_edge(START, \"supervisor\")\n",
"# We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
"research_builder.add_edge(\"search\", \"supervisor\")\n",
"research_builder.add_edge(\"web_scraper\", \"supervisor\")\n",
"# Add the edges where routing applies\n",
"research_builder.add_conditional_edges(\"supervisor\", lambda state: state[\"next\"])\n",
"\n",
"research_graph = research_builder.compile()"
]
},
@@ -528,17 +532,13 @@
")\n",
"\n",
"\n",
"def doc_writing_node(state: MessagesState) -> Command[Literal[\"supervisor\"]]:\n",
"def doc_writing_node(state: AgentState) -> AgentState:\n",
" result = doc_writer_agent.invoke(state)\n",
" return Command(\n",
" update={\n",
" \"messages\": [\n",
" HumanMessage(content=result[\"messages\"][-1].content, name=\"doc_writer\")\n",
" ]\n",
" },\n",
" # We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
" goto=\"supervisor\",\n",
" )\n",
" return {\n",
" \"messages\": [\n",
" HumanMessage(content=result[\"messages\"][-1].content, name=\"doc_writer\")\n",
" ]\n",
" }\n",
"\n",
"\n",
"note_taking_agent = create_react_agent(\n",
@@ -551,17 +551,13 @@
")\n",
"\n",
"\n",
"def note_taking_node(state: MessagesState) -> Command[Literal[\"supervisor\"]]:\n",
"def note_taking_node(state: AgentState) -> AgentState:\n",
" result = note_taking_agent.invoke(state)\n",
" return Command(\n",
" update={\n",
" \"messages\": [\n",
" HumanMessage(content=result[\"messages\"][-1].content, name=\"note_taker\")\n",
" ]\n",
" },\n",
" # We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
" goto=\"supervisor\",\n",
" )\n",
" return {\n",
" \"messages\": [\n",
" HumanMessage(content=result[\"messages\"][-1].content, name=\"note_taker\")\n",
" ]\n",
" }\n",
"\n",
"\n",
"chart_generating_agent = create_react_agent(\n",
@@ -569,19 +565,13 @@
")\n",
"\n",
"\n",
"def chart_generating_node(state: MessagesState) -> Command[Literal[\"supervisor\"]]:\n",
"def chart_generating_node(state: AgentState) -> AgentState:\n",
" result = chart_generating_agent.invoke(state)\n",
" return Command(\n",
" update={\n",
" \"messages\": [\n",
" HumanMessage(\n",
" content=result[\"messages\"][-1].content, name=\"chart_generator\"\n",
" )\n",
" ]\n",
" },\n",
" # We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
" goto=\"supervisor\",\n",
" )\n",
" return {\n",
" \"messages\": [\n",
" HumanMessage(content=result[\"messages\"][-1].content, name=\"chart_generator\")\n",
" ]\n",
" }\n",
"\n",
"\n",
"doc_writing_supervisor_node = make_supervisor_node(\n",
@@ -610,13 +600,21 @@
"outputs": [],
"source": [
"# Create the graph here\n",
"paper_writing_builder = StateGraph(MessagesState)\n",
"paper_writing_builder = StateGraph(AgentState)\n",
"paper_writing_builder.add_node(\"supervisor\", doc_writing_supervisor_node)\n",
"paper_writing_builder.add_node(\"doc_writer\", doc_writing_node)\n",
"paper_writing_builder.add_node(\"note_taker\", note_taking_node)\n",
"paper_writing_builder.add_node(\"chart_generator\", chart_generating_node)\n",
"\n",
"# Define the control flow\n",
"paper_writing_builder.add_edge(START, \"supervisor\")\n",
"# We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
"paper_writing_builder.add_edge(\"doc_writer\", \"supervisor\")\n",
"paper_writing_builder.add_edge(\"note_taker\", \"supervisor\")\n",
"paper_writing_builder.add_edge(\"chart_generator\", \"supervisor\")\n",
"# Add the edges where routing applies\n",
"paper_writing_builder.add_conditional_edges(\"supervisor\", lambda state: state[\"next\"])\n",
"\n",
"paper_writing_graph = paper_writing_builder.compile()"
]
},
@@ -730,41 +728,37 @@
},
"outputs": [],
"source": [
"def call_research_team(state: MessagesState) -> Command[Literal[\"supervisor\"]]:\n",
"def call_research_team(state: AgentState) -> AgentState:\n",
" response = research_graph.invoke({\"messages\": state[\"messages\"][-1]})\n",
" return Command(\n",
" update={\n",
" \"messages\": [\n",
" HumanMessage(\n",
" content=response[\"messages\"][-1].content, name=\"research_team\"\n",
" )\n",
" ]\n",
" },\n",
" goto=\"supervisor\",\n",
" )\n",
" return {\n",
" \"messages\": [\n",
" HumanMessage(content=response[\"messages\"][-1].content, name=\"research_team\")\n",
" ]\n",
" }\n",
"\n",
"\n",
"def call_paper_writing_team(state: MessagesState) -> Command[Literal[\"supervisor\"]]:\n",
"def call_paper_writing_team(state: AgentState) -> AgentState:\n",
" response = paper_writing_graph.invoke({\"messages\": state[\"messages\"][-1]})\n",
" return Command(\n",
" update={\n",
" \"messages\": [\n",
" HumanMessage(\n",
" content=response[\"messages\"][-1].content, name=\"writing_team\"\n",
" )\n",
" ]\n",
" },\n",
" goto=\"supervisor\",\n",
" )\n",
" return {\n",
" \"messages\": [\n",
" HumanMessage(content=response[\"messages\"][-1].content, name=\"writing_team\")\n",
" ]\n",
" }\n",
"\n",
"\n",
"# Define the graph.\n",
"super_builder = StateGraph(MessagesState)\n",
"super_builder = StateGraph(AgentState)\n",
"super_builder.add_node(\"supervisor\", teams_supervisor_node)\n",
"super_builder.add_node(\"research_team\", call_research_team)\n",
"super_builder.add_node(\"writing_team\", call_paper_writing_team)\n",
"\n",
"# Define the control flow\n",
"super_builder.add_edge(START, \"supervisor\")\n",
"# We want our teams to ALWAYS \"report back\" to the top-level supervisor when done\n",
"super_builder.add_edge(\"research_team\", \"supervisor\")\n",
"super_builder.add_edge(\"writing_team\", \"supervisor\")\n",
"# Add the edges where routing applies\n",
"super_builder.add_conditional_edges(\"supervisor\", lambda state: state[\"next\"])\n",
"super_graph = super_builder.compile()"
]
},
File diff suppressed because one or more lines are too long
+2 -2
View File
@@ -43,7 +43,7 @@
"outputs": [],
"source": [
"%%capture --no-stderr\n",
"%pip install -U langgraph langchain_anthropic langsmith langchain-community\n",
"%pip install -U langgraph langchain_anthropic langsmith\n",
"%pip install -U sklearn langchain_openai"
]
},
@@ -632,7 +632,7 @@
"metadata": {},
"outputs": [],
"source": [
"from langchain_community.cache import InMemoryCache\n",
"from langchain.cache import InMemoryCache\n",
"from langchain.globals import set_llm_cache\n",
"\n",
"# Optional. If you are running into errors or rate limits and want to avoid repeated computation,\n",
-6
View File
@@ -151,7 +151,6 @@ nav:
- how-tos/branching.ipynb
- how-tos/map-reduce.ipynb
- how-tos/recursion-limit.ipynb
- how-tos/command.ipynb
- Persistence:
- Persistence: how-tos#persistence
- how-tos/persistence.ipynb
@@ -165,7 +164,6 @@ nav:
- how-tos/memory/manage-conversation-history.ipynb
- how-tos/memory/delete-messages.ipynb
- how-tos/memory/add-summary-conversation-history.ipynb
- how-tos/memory/semantic-search.ipynb
- Human-in-the-loop:
- Human-in-the-loop: how-tos#human-in-the-loop
- how-tos/human_in_the_loop/breakpoints.ipynb
@@ -192,7 +190,6 @@ nav:
- how-tos/tool-calling.ipynb
- how-tos/tool-calling-errors.ipynb
- how-tos/pass-run-time-values-to-tools.ipynb
- how-tos/update-state-from-tools.ipynb
- how-tos/pass-config-to-tools.ipynb
- how-tos/many-tools.ipynb
- Subgraphs:
@@ -200,8 +197,6 @@ nav:
- how-tos/subgraph.ipynb
- how-tos/subgraphs-manage-state.ipynb
- how-tos/subgraph-transform-state.ipynb
- Multi-agent:
- how-tos/multi-agent-network.ipynb
- State Management:
- State Management: how-tos#state-management
- how-tos/state-model.ipynb
@@ -230,7 +225,6 @@ nav:
- cloud/deployment/setup.md
- cloud/deployment/setup_pyproject.md
- cloud/deployment/setup_javascript.md
- cloud/deployment/semantic_search.md
- cloud/deployment/custom_docker.md
- cloud/deployment/test_locally.md
- cloud/deployment/graph_rebuild.md
@@ -42,7 +42,7 @@ class DuckDBSaver(BaseDuckDBSaver):
DuckDBSaver: A new DuckDBSaver instance.
"""
with duckdb.connect(conn_string) as conn:
yield cls(conn)
yield DuckDBSaver(conn)
def setup(self) -> None:
"""Set up the checkpoint database asynchronously.
@@ -45,7 +45,7 @@ class AsyncDuckDBSaver(BaseDuckDBSaver):
AsyncDuckDBSaver: A new AsyncDuckDBSaver instance.
"""
with duckdb.connect(conn_string) as conn:
yield cls(conn)
yield AsyncDuckDBSaver(conn)
async def setup(self) -> None:
"""Set up the checkpoint database asynchronously.
@@ -156,7 +156,7 @@ class AsyncDuckDBStore(AsyncBatchedBaseStore, BaseDuckDBStore):
AsyncDuckDBStore: A new AsyncDuckDBStore instance.
"""
with duckdb.connect(conn_string) as conn:
yield cls(conn)
yield AsyncDuckDBStore(conn)
async def setup(self) -> None:
"""Set up the store database asynchronously.
@@ -5,6 +5,7 @@ from typing import Any, Optional
from langchain_core.runnables import RunnableConfig
from psycopg import Capabilities, Connection, Cursor, Pipeline
from psycopg.errors import UndefinedTable
from psycopg.rows import DictRow, dict_row
from psycopg.types.json import Jsonb
from psycopg_pool import ConnectionPool
@@ -75,15 +76,16 @@ class PostgresSaver(BasePostgresSaver):
the first time checkpointer is used.
"""
with self._cursor() as cur:
cur.execute(self.MIGRATIONS[0])
results = cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
)
row = results.fetchone()
if row is None:
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["v"]
except UndefinedTable:
version = -1
else:
version = row["v"]
for v, migration in zip(
range(version + 1, len(self.MIGRATIONS)),
self.MIGRATIONS[version + 1 :],
@@ -5,6 +5,7 @@ from typing import Any, Optional
from langchain_core.runnables import RunnableConfig
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
from psycopg.errors import UndefinedTable
from psycopg.rows import DictRow, dict_row
from psycopg.types.json import Jsonb
from psycopg_pool import AsyncConnectionPool
@@ -54,7 +55,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
pipeline: bool = False,
serde: Optional[SerializerProtocol] = None,
) -> AsyncIterator["AsyncPostgresSaver"]:
"""Create a new AsyncPostgresSaver instance from a connection string.
"""Create a new PostgresSaver instance from a connection string.
Args:
conn_string (str): The Postgres connection info string.
@@ -80,15 +81,17 @@ class AsyncPostgresSaver(BasePostgresSaver):
the first time checkpointer is used.
"""
async with self._cursor() as cur:
await cur.execute(self.MIGRATIONS[0])
results = await cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
)
row = await results.fetchone()
if row is None:
try:
results = await cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
)
row = await results.fetchone()
if row is None:
version = -1
else:
version = row["v"]
except UndefinedTable:
version = -1
else:
version = row["v"]
for v, migration in zip(
range(version + 1, len(self.MIGRATIONS)),
self.MIGRATIONS[version + 1 :],
@@ -57,17 +57,6 @@ MIGRATIONS = [
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
);""",
"ALTER TABLE checkpoint_blobs ALTER COLUMN blob DROP not null;",
"""
""",
"""
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoints_thread_id_idx ON checkpoints(thread_id);
""",
"""
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_blobs_thread_id_idx ON checkpoint_blobs(thread_id);
""",
"""
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id);
""",
]
SELECT_SQL = f"""
@@ -6,6 +6,7 @@ from typing import Any, Callable, Optional, Union, cast
import orjson
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
from psycopg.errors import UndefinedTable
from psycopg.rows import DictRow, dict_row
from psycopg_pool import AsyncConnectionPool
@@ -20,7 +21,6 @@ from langgraph.store.base import (
)
from langgraph.store.base.batch import AsyncBatchedBaseStore
from langgraph.store.postgres.base import (
_PLACEHOLDER,
BasePostgresStore,
PoolConfig,
PostgresIndexConfig,
@@ -36,75 +36,6 @@ logger = logging.getLogger(__name__)
class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
"""Asynchronous Postgres-backed store with optional vector search using pgvector.
!!! example "Examples"
Basic setup and key-value storage:
```python
from langgraph.store.postgres import AsyncPostgresStore
async with AsyncPostgresStore.from_conn_string(
"postgresql://user:pass@localhost:5432/dbname"
) as store:
await store.setup()
# Store and retrieve data
await store.aput(("users", "123"), "prefs", {"theme": "dark"})
item = await store.aget(("users", "123"), "prefs")
```
Vector search using LangChain embeddings:
```python
from langchain.embeddings import init_embeddings
from langgraph.store.postgres import AsyncPostgresStore
async with AsyncPostgresStore.from_conn_string(
"postgresql://user:pass@localhost:5432/dbname",
index={
"dims": 1536,
"embed": init_embeddings("openai:text-embedding-3-small"),
"fields": ["text"] # specify which fields to embed. Default is the whole serialized value
}
) as store:
await store.setup() # Do this once to run migrations
# Store documents
await store.aput(("docs",), "doc1", {"text": "Python tutorial"})
await store.aput(("docs",), "doc2", {"text": "TypeScript guide"})
# Don't index the following
await store.aput(("docs",), "doc3", {"text": "Other guide"}, index=False)
# Search by similarity
results = await store.asearch(("docs",), query="python programming")
```
Using connection pooling for better performance:
```python
from langgraph.store.postgres import AsyncPostgresStore, PoolConfig
async with AsyncPostgresStore.from_conn_string(
"postgresql://user:pass@localhost:5432/dbname",
pool_config=PoolConfig(
min_size=5,
max_size=20
)
) as store:
await store.setup()
# Use store with connection pooling...
```
Warning:
Make sure to:
1. Call `setup()` before first use to create necessary tables and indexes
2. Have the pgvector extension available to use vector search
3. Use Python 3.10+ for async functionality
Note:
Semantic search is disabled by default. You can enable it by providing an `index` configuration
when creating the store. Without this configuration, all `index` arguments passed to
`put` or `aput`will have no effect.
"""
__slots__ = (
"_deserializer",
"pipe",
@@ -155,6 +86,9 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
return results
def batch(self, ops: Iterable[Op]) -> list[Result]:
return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
@classmethod
@asynccontextmanager
async def from_conn_string(
@@ -213,34 +147,35 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
already exist and runs database migrations. It MUST be called directly by the user
the first time the store is used.
"""
async def _get_version(cur: AsyncCursor[DictRow], table: str) -> int:
await cur.execute(
f"""
CREATE TABLE IF NOT EXISTS {table} (
v INTEGER PRIMARY KEY
)
"""
)
await cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
row = cast(dict, await cur.fetchone())
if row is None:
version = -1
else:
version = row["v"]
return version
async with self._cursor() as cur:
version = await _get_version(cur, table="store_migrations")
for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1):
await cur.execute(sql)
await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
try:
await cur.execute(
"SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
)
row = await cur.fetchone()
if row is None:
version = -1
else:
version = row["v"]
except UndefinedTable:
version = -1
await 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
):
if isinstance(migration, str):
sql = migration
else:
if migration.condition and not migration.condition(self):
continue
if self.index_config:
version = await _get_version(cur, table="vector_migrations")
for v, migration in enumerate(
self.VECTOR_MIGRATIONS[version + 1 :], start=version + 1
):
sql = migration.sql
if migration.params:
params = {
@@ -248,10 +183,9 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
for k, v in migration.params.items()
}
sql = sql % params
await cur.execute(sql)
await cur.execute(
"INSERT INTO vector_migrations (v) VALUES (%s)", (v,)
)
await cur.execute(sql)
await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
async def _execute_batch(
self,
@@ -355,10 +289,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
[query for _, query in embedding_requests]
)
for (idx, _), vector in zip(embedding_requests, vectors):
_paramslist = queries[idx][1]
for i in range(len(_paramslist)):
if _paramslist[i] is _PLACEHOLDER:
_paramslist[i] = vector
queries[idx][1][0] = vector
for (idx, _), (query, params) in zip(search_ops, queries):
await cur.execute(query, params)
@@ -21,6 +21,7 @@ from typing import (
import orjson
from psycopg import Capabilities, Connection, Cursor, Pipeline
from psycopg.errors import UndefinedTable
from psycopg.rows import DictRow, dict_row
from psycopg.types.json import Jsonb
from psycopg_pool import ConnectionPool
@@ -54,11 +55,16 @@ class Migration(NamedTuple):
"""A database migration with optional conditions and parameters."""
sql: str
condition: Optional[Callable[[Any], bool]] = None
params: Optional[dict[str, Any]] = None
condition: Optional[Callable[["BasePostgresStore"], bool]] = None
MIGRATIONS: Sequence[str] = [
def _embedding_requested(store: Any) -> bool:
"""Check if vector operations are available in the database."""
return bool(store.index_config)
MIGRATIONS: Sequence[Union[str, Migration]] = [
"""
CREATE TABLE IF NOT EXISTS store (
-- 'prefix' represents the doc's 'namespace'
@@ -72,15 +78,13 @@ CREATE TABLE IF NOT EXISTS store (
""",
"""
-- For faster lookups by prefix
CREATE INDEX CONCURRENTLY IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
""",
]
VECTOR_MIGRATIONS: Sequence[Migration] = [
Migration(
"""
CREATE EXTENSION IF NOT EXISTS vector;
""",
condition=_embedding_requested,
),
Migration(
"""
@@ -95,23 +99,22 @@ CREATE TABLE IF NOT EXISTS store_vectors (
FOREIGN KEY (prefix, key) REFERENCES store(prefix, key) ON DELETE CASCADE
);
""",
condition=_embedding_requested,
params={
"dims": lambda store: store.index_config["dims"],
"vector_type": lambda store: (
cast(PostgresIndexConfig, store.index_config)
.get("ann_index_config", {})
.get("db_index_config", {})
.get("vector_type", "vector")
),
},
),
Migration(
"""
CREATE INDEX CONCURRENTLY IF NOT EXISTS store_vectors_embedding_idx ON store_vectors
CREATE INDEX IF NOT EXISTS store_vectors_embedding_idx ON store_vectors
USING %(index_type)s (embedding %(ops)s)%(index_params)s;
""",
condition=lambda store: bool(
store.index_config and _get_index_params(store)[0] != "flat"
),
condition=_embedding_requested,
params={
"index_type": lambda store: _get_index_params(store)[0],
"ops": lambda store: _get_vector_type_ops(store),
@@ -155,10 +158,10 @@ class PoolConfig(TypedDict, total=False):
"""
class ANNIndexConfig(TypedDict, total=False):
class DBIndexConfig(TypedDict, total=False):
"""Configuration for vector index in PostgreSQL store."""
kind: Literal["hnsw", "ivfflat", "flat"]
kind: Literal["hnsw", "ivfflat"]
"""Type of index to use: 'hnsw' for Hierarchical Navigable Small World, or 'ivfflat' for Inverted File Flat."""
vector_type: Literal["vector", "halfvec"]
"""Type of vector storage to use.
@@ -168,7 +171,7 @@ class ANNIndexConfig(TypedDict, total=False):
"""
class HNSWConfig(ANNIndexConfig, total=False):
class HNSWConfig(DBIndexConfig, total=False):
"""Configuration for HNSW (Hierarchical Navigable Small World) index."""
kind: Literal["hnsw"] # type: ignore[misc]
@@ -178,7 +181,7 @@ class HNSWConfig(ANNIndexConfig, total=False):
"""Size of dynamic candidate list for index construction. Default is 64."""
class IVFFlatConfig(ANNIndexConfig, total=False):
class IVFFlatConfig(DBIndexConfig, total=False):
"""IVFFlat index divides vectors into lists, and then searches a subset of those lists that are closest to the query vector. It has faster build times and uses less memory than HNSW, but has lower query performance (in terms of speed-recall tradeoff).
Three keys to achieving good recall are:
@@ -203,7 +206,7 @@ class PostgresIndexConfig(IndexConfig, total=False):
Extends EmbeddingConfig with additional configuration for pgvector index and vector types.
"""
ann_index_config: ANNIndexConfig
db_index_config: Union[HNSWConfig, IVFFlatConfig]
"""Specific configuration for the chosen index type (HNSW or IVF Flat)."""
distance_type: Literal["l2", "inner_product", "cosine"]
"""Distance metric to use for vector similarity search:
@@ -215,11 +218,17 @@ class PostgresIndexConfig(IndexConfig, total=False):
class BasePostgresStore(Generic[C]):
MIGRATIONS = MIGRATIONS
VECTOR_MIGRATIONS = VECTOR_MIGRATIONS
conn: C
_deserializer: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]]
index_config: Optional[PostgresIndexConfig]
@staticmethod
def _get_default_index_config() -> IndexConfig:
return HNSWConfig(
kind="hnsw",
vector_type="vector",
)
def _get_batch_GET_ops_queries(
self,
get_ops: Sequence[tuple[int, GetOp]],
@@ -289,12 +298,13 @@ class BasePostgresStore(Generic[C]):
[
_namespace_to_text(op.namespace),
op.key,
Jsonb(cast(dict, op.value)),
Jsonb(cast(dict, op.value).copy()),
]
)
# Then handle embeddings if configured
if self.index_config:
paths = self.index_config["__tokenized_fields"]
for op in inserts:
if op.index is False:
continue
@@ -302,11 +312,6 @@ class BasePostgresStore(Generic[C]):
ns = _namespace_to_text(op.namespace)
k = op.key
if op.index is None:
paths = self.index_config["__tokenized_fields"]
else:
paths = [(ix, tokenize_path(ix)) for ix in op.index]
for path, tokenized_path in paths:
texts = get_text_at_path(value, tokenized_path)
for i, text in enumerate(texts):
@@ -350,30 +355,22 @@ class BasePostgresStore(Generic[C]):
embedding_requests = []
for idx, (_, op) in enumerate(search_ops):
# Build filter conditions first
filter_params = []
filter_conditions = []
if op.filter:
for key, value in op.filter.items():
if isinstance(value, dict):
for op_name, val in value.items():
condition, filter_params_ = self._get_filter_condition(
key, op_name, val
)
filter_conditions.append(condition)
filter_params.extend(filter_params_)
else:
filter_conditions.append("value->%s = %s::jsonb")
filter_params.extend([key, json.dumps(value)])
base_query = """
SELECT prefix, key, value, created_at, updated_at
FROM store
WHERE prefix LIKE %s
"""
params: list = [f"{_namespace_to_text(op.namespace_prefix)}%"]
needs_vector_search = False
# Vector search branch
if op.query and self.index_config:
needs_vector_search = True
embedding_requests.append((idx, op.query))
score_operator, post_operator = _get_distance_operator(self)
score_expr = _get_distance_operator(self)
vector_type = (
cast(PostgresIndexConfig, self.index_config)
.get("ann_index_config", {})
.get("db_index_config", self._get_default_index_config())
.get("vector_type", "vector")
)
@@ -381,81 +378,61 @@ class BasePostgresStore(Generic[C]):
vector_type == "bit"
and self.index_config.get("distance_type") == "hamming"
):
score_operator = score_operator % (
"%s",
self.index_config["dims"],
)
score_expr = score_expr % ("%s", self.index_config["dims"])
else:
score_operator = score_operator % (
"%s",
vector_type,
)
score_expr = score_expr % ("%s", vector_type)
vectors_per_doc_estimate = self.index_config["__estimated_num_vectors"]
expanded_limit = (op.limit * vectors_per_doc_estimate * 2) + 1
# Vector search with CTE for proper score handling
filter_str = (
""
if not filter_conditions
else " AND " + " AND ".join(filter_conditions)
)
if op.namespace_prefix:
prefix_filter_str = f"WHERE s.prefix LIKE %s {filter_str} "
ns_args: Sequence = (f"{_namespace_to_text(op.namespace_prefix)}%",)
else:
ns_args = ()
if filter_str:
prefix_filter_str = f"WHERE {filter_str} "
else:
prefix_filter_str = ""
# Direct query with DISTINCT ON to get best score per document
base_query = f"""
WITH scored AS (
SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, {score_operator} AS neg_score
with scored as (
SELECT DISTINCT ON (s.prefix, s.key)
s.prefix, s.key, s.value, s.created_at, s.updated_at,
{score_expr} as score
FROM store s
JOIN store_vectors sv ON s.prefix = sv.prefix AND s.key = sv.key
{prefix_filter_str}
ORDER BY {score_operator} ASC
WHERE s.prefix LIKE %s
ORDER BY s.prefix, s.key, score DESC
LIMIT %s
)
SELECT * FROM (
SELECT DISTINCT ON (prefix, key)
prefix, key, value, created_at, updated_at, {post_operator} as score
FROM scored
ORDER BY prefix, key, score DESC
) AS unique_docs
ORDER BY score DESC
LIMIT %s
OFFSET %s
SELECT * FROM scored
"""
params = [
_PLACEHOLDER, # Vector placeholder
*ns_args,
*filter_params,
_PLACEHOLDER,
None, # Vector placeholder
f"{_namespace_to_text(op.namespace_prefix)}%",
expanded_limit,
op.limit,
op.offset,
]
# Regular search branch
else:
base_query = """
SELECT prefix, key, value, created_at, updated_at
FROM store
WHERE prefix LIKE %s
"""
params = [f"{_namespace_to_text(op.namespace_prefix)}%"]
if op.filter:
filter_conditions = []
for key, value in op.filter.items():
if isinstance(value, dict):
for op_name, val in value.items():
condition, filter_params = self._get_filter_condition(
key, op_name, val
)
filter_conditions.append(condition)
params.extend(filter_params)
else:
filter_conditions.append("value->%s = %s::jsonb")
params.extend([key, json.dumps(value)])
if filter_conditions:
params.extend(filter_params)
base_query += " AND " + " AND ".join(filter_conditions)
if needs_vector_search:
base_query += " WHERE " + " AND ".join(filter_conditions)
else:
base_query += " AND " + " AND ".join(filter_conditions)
if needs_vector_search:
base_query += " ORDER BY score DESC"
else:
base_query += " ORDER BY updated_at DESC"
base_query += " LIMIT %s OFFSET %s"
params.extend([op.limit, op.offset])
base_query += " LIMIT %s OFFSET %s"
params.extend([op.limit, op.offset])
queries.append((base_query, params))
return queries, embedding_requests
@@ -533,58 +510,6 @@ class BasePostgresStore(Generic[C]):
class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
"""Postgres-backed store with optional vector search using pgvector.
!!! example "Examples"
Basic setup and key-value storage:
```python
from langgraph.store.postgres import PostgresStore
store = PostgresStore(
connection_string="postgresql://user:pass@localhost:5432/dbname"
)
store.setup()
# Store and retrieve data
store.put(("users", "123"), "prefs", {"theme": "dark"})
item = store.get(("users", "123"), "prefs")
```
Vector search using LangChain embeddings:
```python
from langchain.embeddings import init_embeddings
from langgraph.store.postgres import PostgresStore
store = PostgresStore(
connection_string="postgresql://user:pass@localhost:5432/dbname",
index={
"dims": 1536,
"embed": init_embeddings("openai:text-embedding-3-small"),
"fields": ["text"] # specify which fields to embed. Default is the whole serialized value
}
)
store.setup() # Do this once to run migrations
# Store documents
store.put(("docs",), "doc1", {"text": "Python tutorial"})
store.put(("docs",), "doc2", {"text": "TypeScript guide"})
store.put(("docs",), "doc2", {"text": "Other guide"}, index=False) # don't index
# Search by similarity
results = store.search(("docs",), query="python programming")
```
Note:
Semantic search is disabled by default. You can enable it by providing an `index` configuration
when creating the store. Without this configuration, all `index` arguments passed to
`put` or `aput`will have no effect.
Warning:
Make sure to call `setup()` before first use to create necessary tables and indexes.
The pgvector extension must be available to use vector search.
"""
__slots__ = (
"_deserializer",
"pipe",
@@ -634,7 +559,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
pool_config (Optional[PoolArgs]): Configuration for the connection pool.
If provided, will create a connection pool and use it instead of a single connection.
This overrides the `pipeline` argument.
index (Optional[PostgresIndexConfig]): The index configuration for the store.
embedding (Optional[PostgresIndexConfig]): The embedding config.
Returns:
PostgresStore: A new PostgresStore instance.
@@ -780,6 +705,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
vectors = self.embeddings.embed_documents(
[param[-1] for param in txt_params]
)
queries.append(
(
query,
@@ -807,10 +733,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
[query for _, query in embedding_requests]
)
for (idx, _), embedding in zip(embedding_requests, embeddings):
_paramslist = queries[idx][1]
for i in range(len(_paramslist)):
if _paramslist[i] is _PLACEHOLDER:
_paramslist[i] = embedding
queries[idx][1][0] = embedding
for (idx, _), (query, params) in zip(search_ops, queries):
cur.execute(query, params)
@@ -844,36 +767,32 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
already exist and runs database migrations. It MUST be called directly by the user
the first time the store is used.
"""
def _get_version(cur: Cursor[dict[str, Any]], table: str) -> int:
cur.execute(
f"""
CREATE TABLE IF NOT EXISTS {table} (
v INTEGER PRIMARY KEY
)
"""
)
cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
row = cast(dict, cur.fetchone())
if row is None:
version = -1
else:
version = row["v"]
return version
with self._cursor() as cur:
version = _get_version(cur, table="store_migrations")
for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1):
cur.execute(sql)
cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
if self.index_config:
version = _get_version(cur, table="vector_migrations")
for v, migration in enumerate(
self.VECTOR_MIGRATIONS[version + 1 :], start=version + 1
):
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 UndefinedTable:
version = -1
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
):
if isinstance(migration, str):
sql = migration
else:
if migration.condition and not migration.condition(self):
continue
sql = migration.sql
if migration.params:
params = {
@@ -881,8 +800,8 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
for k, v in migration.params.items()
}
sql = sql % params
cur.execute(sql)
cur.execute("INSERT INTO vector_migrations (v) VALUES (%s)", (v,))
cur.execute(sql)
cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
class Row(TypedDict):
@@ -895,10 +814,6 @@ class Row(TypedDict):
# Private utilities
_DEFAULT_ANN_CONFIG = ANNIndexConfig(
vector_type="vector",
)
def _get_vector_type_ops(store: BasePostgresStore) -> str:
"""Get the vector type operator class based on config."""
@@ -906,7 +821,9 @@ def _get_vector_type_ops(store: BasePostgresStore) -> str:
return "vector_cosine_ops"
config = cast(PostgresIndexConfig, store.index_config)
index_config = config.get("ann_index_config", _DEFAULT_ANN_CONFIG).copy()
index_config = config.get(
"db_index_config", BasePostgresStore._get_default_index_config()
)
vector_type = cast(str, index_config.get("vector_type", "vector"))
if vector_type not in ("vector", "halfvec"):
raise ValueError(
@@ -938,7 +855,8 @@ def _get_index_params(store: Any) -> tuple[str, dict[str, Any]]:
return "hnsw", {}
config = cast(PostgresIndexConfig, store.index_config)
index_config = config.get("ann_index_config", _DEFAULT_ANN_CONFIG).copy()
default_config = BasePostgresStore._get_default_index_config()
index_config = config.get("db_index_config", default_config).copy()
kind = index_config.pop("kind", "hnsw")
index_config.pop("vector_type", None)
return kind, index_config
@@ -1007,6 +925,32 @@ def _row_to_search_item(
)
def _row_to_search_item(
namespace: tuple[str, ...],
row: Row,
*,
loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None,
) -> SearchItem:
"""Convert a row from the database into an Item."""
loader = loader or _json_loads
val = row["value"]
score = row.get("score")
if score is not None:
try:
score = float(score) # type: ignore[arg-type]
except ValueError:
logger.warning("Invalid score: %s", score)
score = None
return SearchItem(
value=val if isinstance(val, dict) else loader(val),
key=row["key"],
namespace=namespace,
created_at=row["created_at"],
updated_at=row["updated_at"],
score=score,
)
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
@@ -1036,17 +980,8 @@ def _decode_ns_bytes(namespace: Union[str, bytes, list]) -> tuple[str, ...]:
return tuple(namespace.split("."))
def _get_distance_operator(store: Any) -> tuple[str, str]:
def _get_distance_operator(store: Any) -> str:
"""Get the distance operator and score expression based on config."""
# Note: Today, we are not using ANN indices due to restrictions
# on PGVector's support for mixing vector and non-vector filters
# To use the index, PGVector expects:
# - ORDER BY the operator NOT an expression (even negation blocks it)
# - ASCENDING order
# - Any WHERE clause should be over a partial index.
# If we violate any of these, it will use a sequential scan
# See https://github.com/pgvector/pgvector/issues/216 and the
# pgvector documentation for more details.
if not store.index_config:
raise ValueError(
"Embedding configuration is required for vector operations "
@@ -1057,22 +992,12 @@ def _get_distance_operator(store: Any) -> tuple[str, str]:
config = cast(PostgresIndexConfig, store.index_config)
distance_type = config.get("distance_type", "cosine")
# Return the operator and the score expression
# The operator is used in the CTE and will be compatible with an ASCENDING ORDER
# sort clause.
# The score expression is used in the final query and will be compatible with
# a DESCENDING ORDER sort clause and the user's expectations of what the similarity score
# should be.
if distance_type == "l2":
# Final: "-(sv.embedding <-> %s::%s)"
# We return the "l2 similarity" so that the sorting order is the same
return "sv.embedding <-> %s::%s", "-scored.neg_score"
return "1 - (sv.embedding <-> %s::%s)"
elif distance_type == "inner_product":
# Final: "-(sv.embedding <#> %s::%s)"
return "sv.embedding <#> %s::%s", "-(scored.neg_score)"
else: # cosine similarity
# Final: "1 - (sv.embedding <=> %s::%s)"
return "sv.embedding <=> %s::%s", "1 - scored.neg_score"
return "-(sv.embedding <#> %s::%s)"
else: # cosine
return "1 - (sv.embedding <=> %s::%s)"
def _ensure_index_config(
@@ -1100,6 +1025,3 @@ def _ensure_index_config(
index_config.get("embed"),
)
return embeddings, index_config
_PLACEHOLDER = object()
+431 -533
View File
File diff suppressed because it is too large Load Diff
+4 -4
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-checkpoint-postgres"
version = "2.0.8"
version = "2.0.4"
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
authors = []
license = "MIT"
@@ -10,10 +10,10 @@ packages = [{ include = "langgraph" }]
[tool.poetry.dependencies]
python = "^3.9.0,<4.0"
langgraph-checkpoint = "^2.0.7"
langgraph-checkpoint = "^2.0.2"
orjson = ">=3.10.1"
psycopg = "^3.2.0"
psycopg-pool = "^3.2.0"
psycopg = "^3.0.0"
psycopg-pool = "^3.0.0"
[tool.poetry.group.dev.dependencies]
ruff = "^0.6.2"
+1 -1
View File
@@ -7,7 +7,6 @@ from psycopg.rows import DictRow, dict_row
from tests.embed_test_utils import CharacterEmbeddings
DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5441/"
DEFAULT_URI = "postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable"
@@ -41,4 +40,5 @@ def fake_embeddings() -> CharacterEmbeddings:
return CharacterEmbeddings(dims=500)
INDEX_TYPES = ["hnsw", "ivfflat"]
VECTOR_TYPES = ["vector", "halfvec"]
+85 -199
View File
@@ -1,14 +1,7 @@
# type: ignore
from contextlib import asynccontextmanager
from typing import Any
from uuid import uuid4
import pytest
from langchain_core.runnables import RunnableConfig
from psycopg import AsyncConnection
from psycopg.rows import dict_row
from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.base import (
Checkpoint,
@@ -17,211 +10,104 @@ from langgraph.checkpoint.base import (
empty_checkpoint,
)
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from tests.conftest import DEFAULT_POSTGRES_URI
from tests.conftest import DEFAULT_URI
@asynccontextmanager
async def _pool_saver():
"""Fixture for pool mode testing."""
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
async with AsyncConnectionPool(
DEFAULT_POSTGRES_URI + database,
max_size=10,
kwargs={"autocommit": True, "row_factory": dict_row},
) as pool:
checkpointer = AsyncPostgresSaver(pool)
await checkpointer.setup()
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _pipe_saver():
"""Fixture for pipeline mode testing."""
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI + database,
autocommit=True,
prepare_threshold=0,
row_factory=dict_row,
) as conn:
checkpointer = AsyncPostgresSaver(conn)
await checkpointer.setup()
async with conn.pipeline() as pipe:
checkpointer = AsyncPostgresSaver(conn, pipe=pipe)
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _base_saver():
"""Fixture for regular connection mode testing."""
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI + database,
autocommit=True,
prepare_threshold=0,
row_factory=dict_row,
) as conn:
checkpointer = AsyncPostgresSaver(conn)
await checkpointer.setup()
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _saver(name: str):
if name == "base":
async with _base_saver() as saver:
yield saver
elif name == "pool":
async with _pool_saver() as saver:
yield saver
elif name == "pipe":
async with _pipe_saver() as saver:
yield saver
@pytest.fixture
def test_data():
"""Fixture providing test data for checkpoint tests."""
config_1: RunnableConfig = {
"configurable": {
"thread_id": "thread-1",
# for backwards compatibility testing
"thread_ts": "1",
"checkpoint_ns": "",
class TestAsyncPostgresSaver:
@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": "",
}
}
}
config_2: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2",
"checkpoint_ns": "",
self.config_2: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2",
"checkpoint_ns": "",
}
}
}
config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
self.config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
}
}
}
chkpnt_1: Checkpoint = empty_checkpoint()
chkpnt_2: Checkpoint = create_checkpoint(chkpnt_1, {}, 1)
chkpnt_3: Checkpoint = empty_checkpoint()
self.chkpnt_1: Checkpoint = empty_checkpoint()
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
self.chkpnt_3: Checkpoint = empty_checkpoint()
metadata_1: CheckpointMetadata = {
"source": "input",
"step": 2,
"writes": {},
"score": 1,
}
metadata_2: CheckpointMetadata = {
"source": "loop",
"step": 1,
"writes": {"foo": "bar"},
"score": None,
}
metadata_3: CheckpointMetadata = {}
return {
"configs": [config_1, config_2, config_3],
"checkpoints": [chkpnt_1, chkpnt_2, chkpnt_3],
"metadata": [metadata_1, metadata_2, metadata_3],
}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
async def test_asearch(request, saver_name: str, test_data) -> None:
async with _saver(saver_name) as saver:
configs = test_data["configs"]
checkpoints = test_data["checkpoints"]
metadata = test_data["metadata"]
await saver.aput(configs[0], checkpoints[0], metadata[0], {})
await saver.aput(configs[1], checkpoints[1], metadata[1], {})
await saver.aput(configs[2], checkpoints[2], metadata[2], {})
# call method / assertions
query_1 = {"source": "input"} # search by 1 key
query_2 = {
self.metadata_1: CheckpointMetadata = {
"source": "input",
"step": 2,
"writes": {},
"score": 1,
}
self.metadata_2: CheckpointMetadata = {
"source": "loop",
"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
"score": None,
}
self.metadata_3: CheckpointMetadata = {}
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
await saver.setup()
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 == metadata[0]
async def test_asearch(self) -> None:
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
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, {})
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 == metadata[1]
# 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_3 = [c async for c in saver.alist(None, filter=query_3)]
assert len(search_results_3) == 3
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_4 = [c async for c in saver.alist(None, filter=query_4)]
assert len(search_results_4) == 0
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 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"}
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
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
async def test_null_chars(request, saver_name: str, test_data) -> None:
async with _saver(saver_name) as saver:
config = await saver.aput(
test_data["configs"][0],
test_data["checkpoints"][0],
{"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"
# 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 AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
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"
@@ -1,10 +1,7 @@
# type: ignore
import asyncio
import itertools
import sys
import uuid
from collections.abc import AsyncIterator
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from typing import Any, Optional
@@ -12,16 +9,11 @@ import pytest
from langchain_core.embeddings import Embeddings
from psycopg import AsyncConnection
from langgraph.store.base import (
GetOp,
Item,
ListNamespacesOp,
PutOp,
SearchOp,
)
from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
from langgraph.store.postgres import AsyncPostgresStore
from tests.conftest import (
DEFAULT_URI,
INDEX_TYPES,
VECTOR_TYPES,
CharacterEmbeddings,
)
@@ -71,128 +63,6 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]:
await conn.execute(f"DROP DATABASE {database}")
async def test_no_running_loop(store: AsyncPostgresStore) -> None:
with pytest.raises(asyncio.InvalidStateError):
store.put(("foo", "bar"), "baz", {"val": "baz"})
with pytest.raises(asyncio.InvalidStateError):
store.get(("foo", "bar"), "baz")
with pytest.raises(asyncio.InvalidStateError):
store.delete(("foo", "bar"), "baz")
with pytest.raises(asyncio.InvalidStateError):
store.search(("foo", "bar"))
with pytest.raises(asyncio.InvalidStateError):
store.list_namespaces(prefix=("foo",))
with pytest.raises(asyncio.InvalidStateError):
store.batch([PutOp(namespace=("foo", "bar"), key="baz", value={"val": "baz"})])
with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(store.put, ("foo", "bar"), "baz", {"val": "baz"})
result = await asyncio.wrap_future(future)
assert result is None
future = executor.submit(store.get, ("foo", "bar"), "baz")
result = await asyncio.wrap_future(future)
assert result.value == {"val": "baz"}
result = await asyncio.wrap_future(
executor.submit(store.list_namespaces, prefix=("foo",))
)
async def test_large_batches(request: Any, store: AsyncPostgresStore) -> None:
N = 100 # less important that we are performant here
M = 10
with ThreadPoolExecutor(max_workers=10) as executor:
futures = []
for m in range(M):
for i in range(N):
futures += [
executor.submit(
store.put,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
value={"foo": "bar" + str(i)},
),
executor.submit(
store.get,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
),
executor.submit(
store.list_namespaces,
prefix=None,
max_depth=m + 1,
),
executor.submit(
store.search,
("test",),
),
executor.submit(
store.put,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
value={"foo": "bar" + str(i)},
),
executor.submit(
store.put,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
None,
),
]
results = await asyncio.gather(
*(asyncio.wrap_future(future) for future in futures)
)
assert len(results) == M * N * 6
async def test_large_batches_async(store: AsyncPostgresStore) -> None:
N = 1000
M = 10
coros = []
for m in range(M):
for i in range(N):
coros.append(
store.aput(
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
value={"foo": "bar" + str(i)},
)
)
coros.append(
store.aget(
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
)
)
coros.append(
store.alist_namespaces(
prefix=None,
max_depth=m + 1,
)
)
coros.append(
store.asearch(
("test",),
)
)
coros.append(
store.aput(
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
value={"foo": "bar" + str(i)},
)
)
coros.append(
store.adelete(
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
)
)
results = await asyncio.gather(*coros)
assert len(results) == M * N * 6
async def test_abatch_order(store: AsyncPostgresStore) -> None:
# Setup test data
await store.aput(("test", "foo"), "key1", {"data": "value1"})
@@ -321,6 +191,7 @@ async def test_batch_list_namespaces_ops(store: AsyncPostgresStore) -> None:
@asynccontextmanager
async def _create_vector_store(
index_type: str,
vector_type: str,
distance_type: str,
fake_embeddings: CharacterEmbeddings,
@@ -344,7 +215,8 @@ async def _create_vector_store(
index_config = {
"dims": fake_embeddings.dims,
"embed": fake_embeddings,
"ann_index_config": {
"db_index_config": {
"kind": index_type,
"vector_type": vector_type,
},
"distance_type": distance_type,
@@ -372,22 +244,25 @@ async def _create_vector_store(
@pytest.fixture(
scope="function",
params=[
(vector_type, distance_type)
(index_type, vector_type, distance_type)
for index_type in INDEX_TYPES
for vector_type in VECTOR_TYPES
for distance_type in (
["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
(["hamming"] if index_type == "ivfflat" else ["hamming", "jaccard"])
if vector_type == "bit"
else ["l2", "inner_product", "cosine"]
)
],
ids=lambda p: f"{p[0]}_{p[1]}",
ids=lambda p: f"{p[0]}_{p[1]}_{p[2]}",
)
async def vector_store(
request,
fake_embeddings: CharacterEmbeddings,
) -> AsyncIterator[AsyncPostgresStore]:
"""Create a store with vector search enabled."""
vector_type, distance_type = request.param
index_type, vector_type, distance_type = request.param
async with _create_vector_store(
vector_type, distance_type, fake_embeddings
index_type, vector_type, distance_type, fake_embeddings
) as store:
yield store
@@ -542,19 +417,24 @@ async def test_vector_search_edge_cases(vector_store: AsyncPostgresStore) -> Non
@pytest.mark.parametrize(
"vector_type,distance_type",
"index_type,vector_type,distance_type",
[
*itertools.product(["vector", "halfvec"], ["cosine", "inner_product", "l2"]),
("ivfflat", "vector", "cosine"),
("hnsw", "vector", "cosine"),
("hnsw", "halfvec", "cosine"),
("hnsw", "halfvec", "inner_product"),
],
)
async def test_embed_with_path(
request: Any,
fake_embeddings: CharacterEmbeddings,
index_type: str,
vector_type: str,
distance_type: str,
) -> None:
"""Test vector search with specific text fields in Postgres store."""
async with _create_vector_store(
index_type,
vector_type,
distance_type,
fake_embeddings,
@@ -598,40 +478,3 @@ async def test_embed_with_path(
assert len(results) == 2
assert results[0].score < ascore
assert results[1].score < ascore
@pytest.mark.parametrize(
"vector_type,distance_type",
[
*itertools.product(["vector", "halfvec"], ["cosine", "inner_product", "l2"]),
],
)
async def test_search_sorting(
request: Any,
fake_embeddings: CharacterEmbeddings,
vector_type: str,
distance_type: str,
) -> None:
"""Test operation-level field configuration for vector search."""
async with _create_vector_store(
vector_type,
distance_type,
fake_embeddings,
text_fields=["key1"], # Default fields that won't match our test data
) as store:
amatch = {
"key1": "mmm",
}
await store.aput(("test", "M"), "M", amatch)
N = 100
for i in range(N):
await store.aput(("test", "A"), f"A{i}", {"key1": "no"})
for i in range(N):
await store.aput(("test", "Z"), f"Z{i}", {"key1": "no"})
results = await store.asearch(("test",), query="mmm", limit=10)
assert len(results) == 10
assert len(set(r.key for r in results)) == 10
assert results[0].key == "M"
assert results[0].score > results[1].score
+21 -177
View File
@@ -19,6 +19,7 @@ from langgraph.store.base import (
from langgraph.store.postgres import PostgresStore
from tests.conftest import (
DEFAULT_URI,
INDEX_TYPES,
VECTOR_TYPES,
CharacterEmbeddings,
)
@@ -351,6 +352,7 @@ class TestPostgresStore:
@contextmanager
def _create_vector_store(
index_type: str,
vector_type: str,
distance_type: str,
fake_embeddings: Embeddings,
@@ -371,7 +373,8 @@ def _create_vector_store(
index_config = {
"dims": fake_embeddings.dims,
"embed": fake_embeddings,
"ann_index_config": {
"db_index_config": {
"kind": index_type,
"vector_type": vector_type,
},
"distance_type": distance_type,
@@ -395,21 +398,26 @@ def _create_vector_store(
@pytest.fixture(
scope="function",
params=[
(vector_type, distance_type)
(index_type, vector_type, distance_type)
for index_type in INDEX_TYPES
for vector_type in VECTOR_TYPES
for distance_type in (
["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
(["hamming"] if index_type == "ivfflat" else ["hamming", "jaccard"])
if vector_type == "bit"
else ["l2", "inner_product", "cosine"]
)
],
ids=lambda p: f"{p[0]}_{p[1]}",
ids=lambda p: f"{p[0]}_{p[1]}_{p[2]}",
)
def vector_store(
request,
fake_embeddings: Embeddings,
) -> PostgresStore:
"""Create a store with vector search enabled."""
vector_type, distance_type = request.param
with _create_vector_store(vector_type, distance_type, fake_embeddings) as store:
index_type, vector_type, distance_type = request.param
with _create_vector_store(
index_type, vector_type, distance_type, fake_embeddings
) as store:
yield store
@@ -547,22 +555,24 @@ def test_vector_search_edge_cases(vector_store: PostgresStore) -> None:
@pytest.mark.parametrize(
"vector_type,distance_type",
"index_type,vector_type,distance_type",
[
("vector", "cosine"),
("vector", "inner_product"),
("halfvec", "cosine"),
("halfvec", "inner_product"),
("ivfflat", "vector", "cosine"),
("hnsw", "vector", "cosine"),
("hnsw", "halfvec", "cosine"),
("hnsw", "halfvec", "inner_product"),
],
)
def test_embed_with_path_sync(
request: Any,
fake_embeddings: CharacterEmbeddings,
index_type: str,
vector_type: str,
distance_type: str,
) -> None:
"""Test vector search with specific text fields in Postgres store."""
with _create_vector_store(
index_type,
vector_type,
distance_type,
fake_embeddings,
@@ -616,169 +626,3 @@ def test_embed_with_path_sync(
assert results[0].key != results[1].key
assert results[0].score < ascore
assert results[1].score < ascore
@pytest.mark.parametrize(
"vector_type,distance_type",
[
("vector", "cosine"),
("vector", "inner_product"),
("halfvec", "cosine"),
("halfvec", "inner_product"),
],
)
def test_embed_with_path_operation_config(
request: Any,
fake_embeddings: CharacterEmbeddings,
vector_type: str,
distance_type: str,
) -> None:
"""Test operation-level field configuration for vector search."""
with _create_vector_store(
vector_type,
distance_type,
fake_embeddings,
text_fields=["key17"], # Default fields that won't match our test data
) as store:
doc3 = {
"key0": "aaa",
"key1": "bbb",
"key2": "ccc",
"key3": "ddd",
}
doc4 = {
"key0": "eee",
"key1": "bbb", # Same as doc3.key1
"key2": "fff",
"key3": "ggg",
}
store.put(("test",), "doc3", doc3, index=["key0", "key1"])
store.put(("test",), "doc4", doc4, index=["key1", "key3"])
results = store.search(("test",), query="aaa")
assert len(results) == 2
assert results[0].key == "doc3"
assert len(set(r.key for r in results)) == 2
assert results[0].score > results[1].score
results = store.search(("test",), query="ggg")
assert len(results) == 2
assert results[0].key == "doc4"
assert results[0].score > results[1].score
results = store.search(("test",), query="bbb")
assert len(results) == 2
assert results[0].key != results[1].key
assert results[0].score == pytest.approx(results[1].score, abs=1e-3)
results = store.search(("test",), query="ccc")
assert len(results) == 2
assert all(
r.score < 0.9 for r in results
) # Unindexed field should have low scores
# Test index=False behavior
doc5 = {
"key0": "hhh",
"key1": "iii",
}
store.put(("test",), "doc5", doc5, index=False)
results = store.search(("test",))
assert len(results) == 3
assert all(r.score is None for r in results)
assert any(r.key == "doc5" for r in results)
results = store.search(("test",), query="hhh")
# TODO: We don't currently fill in additional results if there are not enough
# returned during vector search.
# assert len(results) == 3
# doc5_result = next(r for r in results if r.key == "doc5")
# assert doc5_result.score is None
def _cosine_similarity(X: list[float], Y: list[list[float]]) -> list[float]:
"""
Compute cosine similarity between a vector X and a matrix Y.
Lazy import numpy for efficiency.
"""
similarities = []
for y in Y:
dot_product = sum(a * b for a, b in zip(X, y))
norm1 = sum(a * a for a in X) ** 0.5
norm2 = sum(a * a for a in y) ** 0.5
similarity = dot_product / (norm1 * norm2) if norm1 > 0 and norm2 > 0 else 0.0
similarities.append(similarity)
return similarities
def _inner_product(X: list[float], Y: list[list[float]]) -> list[float]:
"""
Compute inner product between a vector X and a matrix Y.
Lazy import numpy for efficiency.
"""
similarities = []
for y in Y:
similarity = sum(a * b for a, b in zip(X, y))
similarities.append(similarity)
return similarities
def _neg_l2_distance(X: list[float], Y: list[list[float]]) -> list[float]:
"""
Compute l2 distance between a vector X and a matrix Y.
Lazy import numpy for efficiency.
"""
similarities = []
for y in Y:
similarity = sum((a - b) ** 2 for a, b in zip(X, y)) ** 0.5
similarities.append(-similarity)
return similarities
@pytest.mark.parametrize(
"vector_type,distance_type",
[
("vector", "cosine"),
("vector", "inner_product"),
("halfvec", "l2"),
],
)
@pytest.mark.parametrize("query", ["aaa", "bbb", "ccc", "abcd", "poisson"])
def test_scores(
fake_embeddings: CharacterEmbeddings,
vector_type: str,
distance_type: str,
query: str,
) -> None:
"""Test operation-level field configuration for vector search."""
with _create_vector_store(
vector_type,
distance_type,
fake_embeddings,
text_fields=["key0"],
) as store:
doc = {
"key0": "aaa",
}
store.put(("test",), "doc", doc, index=["key0", "key1"])
results = store.search((), query=query)
vec0 = fake_embeddings.embed_query(doc["key0"])
vec1 = fake_embeddings.embed_query(query)
if distance_type == "cosine":
similarities = _cosine_similarity(vec1, [vec0])
elif distance_type == "inner_product":
similarities = _inner_product(vec1, [vec0])
elif distance_type == "l2":
similarities = _neg_l2_distance(vec1, [vec0])
assert len(results) == 1
assert results[0].score == pytest.approx(similarities[0], abs=1e-3)
+84 -186
View File
@@ -1,14 +1,7 @@
# type: ignore
from contextlib import contextmanager
from typing import Any
from uuid import uuid4
import pytest
from langchain_core.runnables import RunnableConfig
from psycopg import Connection
from psycopg.rows import dict_row
from psycopg_pool import ConnectionPool
from langgraph.checkpoint.base import (
Checkpoint,
@@ -17,198 +10,103 @@ from langgraph.checkpoint.base import (
empty_checkpoint,
)
from langgraph.checkpoint.postgres import PostgresSaver
from tests.conftest import DEFAULT_POSTGRES_URI
from tests.conftest import DEFAULT_URI
@contextmanager
def _pool_saver():
"""Fixture for pool mode testing."""
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
with ConnectionPool(
DEFAULT_POSTGRES_URI + database,
max_size=10,
kwargs={"autocommit": True, "row_factory": dict_row},
) as pool:
checkpointer = PostgresSaver(pool)
checkpointer.setup()
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
@contextmanager
def _pipe_saver():
"""Fixture for pipeline mode testing."""
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
with Connection.connect(
DEFAULT_POSTGRES_URI + database,
autocommit=True,
prepare_threshold=0,
row_factory=dict_row,
) as conn:
checkpointer = PostgresSaver(conn)
checkpointer.setup()
with conn.pipeline() as pipe:
checkpointer = PostgresSaver(conn, pipe=pipe)
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
@contextmanager
def _base_saver():
"""Fixture for regular connection mode testing."""
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
with Connection.connect(
DEFAULT_POSTGRES_URI + database,
autocommit=True,
prepare_threshold=0,
row_factory=dict_row,
) as conn:
checkpointer = PostgresSaver(conn)
checkpointer.setup()
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
@contextmanager
def _saver(name: str):
if name == "base":
with _base_saver() as saver:
yield saver
elif name == "pool":
with _pool_saver() as saver:
yield saver
elif name == "pipe":
with _pipe_saver() as saver:
yield saver
@pytest.fixture
def test_data():
"""Fixture providing test data for checkpoint tests."""
config_1: RunnableConfig = {
"configurable": {
"thread_id": "thread-1",
# for backwards compatibility testing
"thread_ts": "1",
"checkpoint_ns": "",
class TestPostgresSaver:
@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": "",
}
}
}
config_2: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2",
"checkpoint_ns": "",
self.config_2: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2",
"checkpoint_ns": "",
}
}
}
config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
self.config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
}
}
}
chkpnt_1: Checkpoint = empty_checkpoint()
chkpnt_2: Checkpoint = create_checkpoint(chkpnt_1, {}, 1)
chkpnt_3: Checkpoint = empty_checkpoint()
self.chkpnt_1: Checkpoint = empty_checkpoint()
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
self.chkpnt_3: Checkpoint = empty_checkpoint()
metadata_1: CheckpointMetadata = {
"source": "input",
"step": 2,
"writes": {},
"score": 1,
}
metadata_2: CheckpointMetadata = {
"source": "loop",
"step": 1,
"writes": {"foo": "bar"},
"score": None,
}
metadata_3: CheckpointMetadata = {}
return {
"configs": [config_1, config_2, config_3],
"checkpoints": [chkpnt_1, chkpnt_2, chkpnt_3],
"metadata": [metadata_1, metadata_2, metadata_3],
}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
def test_search(saver_name: str, test_data) -> None:
with _saver(saver_name) as saver:
configs = test_data["configs"]
checkpoints = test_data["checkpoints"]
metadata = test_data["metadata"]
saver.put(configs[0], checkpoints[0], metadata[0], {})
saver.put(configs[1], checkpoints[1], metadata[1], {})
saver.put(configs[2], checkpoints[2], metadata[2], {})
# call method / assertions
query_1 = {"source": "input"} # search by 1 key
query_2 = {
self.metadata_1: CheckpointMetadata = {
"source": "input",
"step": 2,
"writes": {},
"score": 1,
}
self.metadata_2: CheckpointMetadata = {
"source": "loop",
"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
"score": None,
}
self.metadata_3: CheckpointMetadata = {}
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
saver.setup()
search_results_1 = list(saver.list(None, filter=query_1))
assert len(search_results_1) == 1
assert search_results_1[0].metadata == metadata[0]
def test_search(self) -> None:
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
# 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, {})
search_results_2 = list(saver.list(None, filter=query_2))
assert len(search_results_2) == 1
assert search_results_2[0].metadata == metadata[1]
# 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_3 = list(saver.list(None, filter=query_3))
assert len(search_results_3) == 3
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_4 = list(saver.list(None, filter=query_4))
assert len(search_results_4) == 0
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 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"}
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
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
def test_null_chars(saver_name: str, test_data) -> None:
with _saver(saver_name) as saver:
config = saver.put(
test_data["configs"][0],
test_data["checkpoints"][0],
{"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"]
== "abc"
)
# 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 PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
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"
)
@@ -110,7 +110,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
check_same_thread=False,
)
) as conn:
yield cls(conn)
yield SqliteSaver(conn)
def setup(self) -> None:
"""Set up the checkpoint database.
@@ -137,7 +137,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
AsyncSqliteSaver: A new AsyncSqliteSaver instance.
"""
async with aiosqlite.connect(conn_string) as conn:
yield cls(conn)
yield AsyncSqliteSaver(conn)
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Get a checkpoint tuple from the database.
@@ -1,3 +1,4 @@
import asyncio
import logging
import os
import pickle
@@ -5,6 +6,7 @@ import random
import shutil
from collections import defaultdict
from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack
from functools import partial
from types import TracebackType
from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple, Type
@@ -393,7 +395,9 @@ class MemorySaver(
Returns:
Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
"""
return self.get_tuple(config)
return await asyncio.get_running_loop().run_in_executor(
None, self.get_tuple, config
)
async def alist(
self,
@@ -414,8 +418,24 @@ class MemorySaver(
Yields:
AsyncIterator[CheckpointTuple]: An asynchronous iterator of checkpoint tuples.
"""
for item in self.list(config, filter=filter, before=before, limit=limit):
yield item
loop = asyncio.get_running_loop()
iter = await loop.run_in_executor(
None,
partial(
self.list,
before=before,
limit=limit,
filter=filter,
),
config,
)
while True:
# handling StopIteration exception inside coroutine won't work
# as expected, so using next() with default value to break the loop
if item := await loop.run_in_executor(None, next, iter, None):
yield item
else:
break
async def aput(
self,
@@ -435,7 +455,9 @@ class MemorySaver(
Returns:
RunnableConfig: The updated config containing the saved checkpoint's timestamp.
"""
return self.put(config, checkpoint, metadata, new_versions)
return await asyncio.get_running_loop().run_in_executor(
None, self.put, config, checkpoint, metadata, new_versions
)
async def aput_writes(
self,
@@ -452,9 +474,10 @@ class MemorySaver(
config (RunnableConfig): The config to associate with the writes.
writes (List[Tuple[str, Any]]): The writes to save, each as a (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
return self.put_writes(config, writes, task_id)
"""
return self.put_writes(config, writes, task_id)
return await asyncio.get_running_loop().run_in_executor(
None, self.put_writes, config, writes, task_id
)
def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
if current is None:
+61 -249
View File
@@ -4,9 +4,9 @@ Stores provide long-term memory that persists across threads and conversations.
Supports hierarchical namespaces, key-value storage, and optional vector search.
Core types:
- BaseStore: Store interface with sync/async operations
- Item: Stored key-value pairs with metadata
- Op: Get/Put/Search/List operations
- BaseStore: Store interface with sync/async operations
- Item: Stored key-value pairs with metadata
- Op: Get/Put/Search/List operations
"""
from abc import ABC, abstractmethod
@@ -80,19 +80,16 @@ class Item:
def dict(self) -> dict:
return {
"namespace": list(self.namespace),
"key": self.key,
"value": self.value,
"key": self.key,
"namespace": list(self.namespace),
"created_at": self.created_at.isoformat(),
"updated_at": self.updated_at.isoformat(),
}
def __repr__(self) -> str:
return f"Item({', '.join(f'{k}={v!r}' for k, v in self.dict().items())})"
class SearchItem(Item):
"""Represents an item returned from a search operation with additional metadata."""
"""Represents a result item with additional response metadata."""
__slots__ = ("score",)
@@ -136,7 +133,7 @@ class GetOp(NamedTuple):
This operation allows precise retrieval of stored items using their full path
(namespace) and unique identifier (key) combination.
???+ example "Examples"
??? example "Examples"
Basic item retrieval:
```python
@@ -148,7 +145,7 @@ class GetOp(NamedTuple):
namespace: tuple[str, ...]
"""Hierarchical path that uniquely identifies the item's location.
???+ example "Examples"
??? example "Examples"
```python
("users",) # Root level users namespace
@@ -159,7 +156,7 @@ class GetOp(NamedTuple):
key: str
"""Unique identifier for the item within its specific namespace.
???+ example "Examples"
??? example "Examples"
```python
"user123" # For a user profile
@@ -178,7 +175,7 @@ class SearchOp(NamedTuple):
Note:
Natural language search support depends on your store implementation.
???+ example "Examples"
??? example "Examples"
Search with filters and pagination:
```python
SearchOp(
@@ -202,7 +199,7 @@ class SearchOp(NamedTuple):
namespace_prefix: tuple[str, ...]
"""Hierarchical path prefix defining the search scope.
???+ example "Examples"
??? example "Examples"
```python
() # Search entire store
@@ -224,7 +221,8 @@ class SearchOp(NamedTuple):
- $lt: Less than
- $lte: Less than or equal to
???+ example "Examples"
??? example "Examples"
Simple exact match:
```python
@@ -245,6 +243,9 @@ class SearchOp(NamedTuple):
"color": "red"
}
```
Note:
Comparison operator support depends on your store implementation.
"""
limit: int = 10
@@ -256,7 +257,7 @@ class SearchOp(NamedTuple):
query: Optional[str] = None
"""Natural language search query for semantic search capabilities.
???+ example "Examples"
??? example "Examples"
- "technical documentation about REST APIs"
- "machine learning papers from 2023"
"""
@@ -266,12 +267,10 @@ class SearchOp(NamedTuple):
NamespacePath = tuple[Union[str, Literal["*"]], ...]
"""A tuple representing a namespace path that can include wildcards.
???+ example "Examples"
```python
Examples:
("users",) # Exact users namespace
("documents", "*") # Any sub-namespace under documents
("cache", "*", "v1") # Any cache category with v1 version
```
"""
# Type for specifying how to match namespaces
@@ -291,7 +290,7 @@ class MatchCondition(NamedTuple):
pattern that can include wildcards to flexibly match different namespace
hierarchies.
???+ example "Examples"
??? example "Examples"
Prefix matching:
```python
MatchCondition(match_type="prefix", path=("users", "profiles"))
@@ -321,7 +320,7 @@ class ListNamespacesOp(NamedTuple):
This operation allows exploring the organization of data, finding specific
collections, and navigating the namespace hierarchy.
???+ example "Examples"
??? example "Examples"
List all namespaces under the "documents" path:
```python
@@ -344,7 +343,7 @@ class ListNamespacesOp(NamedTuple):
match_conditions: Optional[tuple[MatchCondition, ...]] = None
"""Optional conditions for filtering namespaces.
???+ example "Examples"
??? example "Examples"
All user namespaces:
```python
(MatchCondition(match_type="prefix", path=("users",)),)
@@ -386,7 +385,7 @@ class PutOp(NamedTuple):
The namespace acts as a folder-like structure to organize items.
Each element in the tuple represents one level in the hierarchy.
???+ example "Examples"
??? example "Examples"
Root level documents
```python
("documents",)
@@ -432,9 +431,9 @@ class PutOp(NamedTuple):
"""Controls how the item's fields are indexed for search operations.
Indexing configuration determines how the item can be found through search:
- None (default): Uses the store's default indexing configuration (if provided)
- False: Disables indexing for this item
- list[str]: Specifies which json path fields to index for search
- None (default): Uses the store's default indexing configuration (if provided)
- False: Disables indexing for this item
- list[str]: Specifies which json path fields to index for search
The item remains accessible through direct get() operations regardless of indexing.
When indexed, fields can be searched using natural language queries through
@@ -448,14 +447,15 @@ class PutOp(NamedTuple):
- Last element: "array[-1]"
- All elements (each individually): "array[*]"
???+ example "Examples"
- None - Use store defaults (whole item)
??? example "Examples"
- None - Use store defaults
- False - Don't index this item
- list[str] - List of fields to index
```python
[
"metadata.title", # Nested field access
"context[*].content", # Index content from all context as separate vectors
"chapters[*].content", # Index content from all chapters as separate vectors
"authors[0].name", # First author's name
"revisions[-1].changes", # Most recent revision's changes
"sections[*].paragraphs[*].text", # All text from all paragraphs in all sections
@@ -474,126 +474,28 @@ class InvalidNamespaceError(ValueError):
class IndexConfig(TypedDict, total=False):
"""Configuration for indexing documents for semantic search in the store.
If not provided to the store, the store will not support vector search.
In that case, all `index` arguments to put() and `aput()` operations will be ignored.
"""
"""Configuration for indexing documents for semantic search in the store."""
dims: int
"""Number of dimensions in the embedding vectors.
Common embedding models have the following dimensions:
- openai:text-embedding-3-large: 3072
- openai:text-embedding-3-small: 1536
- openai:text-embedding-ada-002: 1536
- cohere:embed-english-v3.0: 1024
- cohere:embed-english-light-v3.0: 384
- cohere:embed-multilingual-v3.0: 1024
- cohere:embed-multilingual-light-v3.0: 384
- OpenAI text-embedding-3-large: 256, 1024, or 3072
- OpenAI text-embedding-3-small: 512 or 1536
- OpenAI text-embedding-ada-002: 1536
- Cohere embed-english-v3.0: 1024
- Cohere embed-english-light-v3.0: 384
- Cohere embed-multilingual-v3.0: 1024
- Cohere embed-multilingual-light-v3.0: 384
"""
embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc]
"""Optional function to generate embeddings from text.
Can be specified in three ways:
1. A LangChain Embeddings instance
2. A synchronous embedding function (EmbeddingsFunc)
3. An asynchronous embedding function (AEmbeddingsFunc)
???+ example "Examples"
Using LangChain's initialization with InMemoryStore:
```python
from langchain.embeddings import init_embeddings
from langgraph.store.memory import InMemoryStore
store = InMemoryStore(
index={
"dims": 1536,
"embed": init_embeddings("openai:text-embedding-3-small")
}
)
```
Using a custom embedding function with InMemoryStore:
```python
from openai import OpenAI
from langgraph.store.memory import InMemoryStore
client = OpenAI()
def embed_texts(texts: list[str]) -> list[list[float]]:
response = client.embeddings.create(
model="text-embedding-3-small",
input=texts
)
return [e.embedding for e in response.data]
store = InMemoryStore(
index={
"dims": 1536,
"embed": embed_texts
}
)
```
Using an asynchronous embedding function with InMemoryStore:
```python
from openai import AsyncOpenAI
from langgraph.store.memory import InMemoryStore
client = AsyncOpenAI()
async def aembed_texts(texts: list[str]) -> list[list[float]]:
response = await client.embeddings.create(
model="text-embedding-3-small",
input=texts
)
return [e.embedding for e in response.data]
store = InMemoryStore(
index={
"dims": 1536,
"embed": aembed_texts
}
)
```
"""
"""Optional function to generate embeddings from text."""
fields: Optional[list[str]]
"""Fields to extract text from for embedding generation.
Controls which parts of stored items are embedded for semantic search. Follows JSON path syntax:
- ["$"]: Embeds the entire JSON object as one vector (default)
- ["field1", "field2"]: Embeds specific top-level fields
- ["parent.child"]: Embeds nested fields using dot notation
- ["array[*].field"]: Embeds field from each array element separately
Note:
You can always override this behavior when storing an item using the
`index` parameter in the `put` or `aput` operations.
???+ example "Examples"
```python
# Embed entire document (default)
fields=["$"]
# Embed specific fields
fields=["text", "summary"]
# Embed nested fields
fields=["metadata.title", "content.body"]
# Embed from arrays
fields=["messages[*].content"] # Each message content separately
fields=["context[0].text"] # First context item's text
```
Note:
- Fields missing from a document are skipped
- Array notation creates separate embeddings for each element
- Complex nested paths are supported (e.g., "a.b[*].c.d")
Defaults to the root ["$"], which embeds the json object as a whole.
"""
@@ -602,15 +504,6 @@ class BaseStore(ABC):
Stores enable persistence and memory that can be shared across threads,
scoped to user IDs, assistant IDs, or other arbitrary namespaces.
Some implementations may support semantic search capabilities through
an optional `index` configuration.
Note:
Semantic search capabilities vary by implementation and are typically
disabled by default. Stores that support this feature can be configured
by providing an `index` configuration at creation time. Without this
configuration, semantic search is disabled and any `index` arguments
to storage operations will have no effect.
"""
__slots__ = ("__weakref__",)
@@ -672,39 +565,6 @@ class BaseStore(ABC):
Returns:
List of items matching the search criteria.
???+ example "Examples"
Basic filtering:
```python
# Search for documents with specific metadata
results = store.search(
("docs",),
filter={"type": "article", "status": "published"}
)
```
Natural language search (requires vector store implementation):
```python
# Initialize store with embedding configuration
store = YourStore( # e.g., InMemoryStore, AsyncPostgresStore
index={
"dims": 1536, # embedding dimensions
"embed": your_embedding_function, # function to create embeddings
"fields": ["text"] # fields to embed. Defaults to ["$"]
}
)
# Search for semantically similar documents
results = store.search(
("docs",),
query="machine learning applications in healthcare",
filter={"type": "research_paper"},
limit=5
)
```
Note: Natural language search support depends on your store implementation
and requires proper embedding configuration.
"""
return self.batch([SearchOp(namespace_prefix, filter, limit, offset, query)])[0]
@@ -725,10 +585,7 @@ class BaseStore(ABC):
value: Dictionary containing the item's data. Must contain string keys
and JSON-serializable values.
index: Controls how the item's fields are indexed for search:
- None (default): Use `fields` you configured when creating the store (if any)
If you do not initialize the store with indexing capabilities,
the `index` parameter will be ignored
- None (default): Use store's default indexing configuration
- False: Disable indexing for this item
- list[str]: List of field paths to index, supporting:
- Nested fields: "metadata.title"
@@ -736,25 +593,23 @@ class BaseStore(ABC):
- Specific indices: "authors[0].name"
Note:
Indexing support depends on your store implementation.
If you do not initialize the store with indexing capabilities,
the `index` parameter will be ignored.
Indexing capabilities depend on your store implementation.
Some implementations may support only a subset of indexing features.
???+ example "Examples"
Store item. Indexing depends on how you configure the store.
??? example "Examples"
Simple storage without special indexing (respects store defaults)
```python
store.put(("docs",), "report", {"memory": "Will likes ai"})
store.put(("docs",), "report", {"title": "Annual Report"})
```
Do not index item for semantic search. Still accessible through get()
and search() operations but won't have a vector representation.
Index specific fields for search
```python
store.put(("docs",), "report", {"memory": "Will likes ai"}, index=False)
store.put(("docs",), "report", {"title": "Annual Report"}, index=["title"])
```
Index specific fields for search.
Do not index for semantic search
```python
store.put(("docs",), "report", {"memory": "Will likes ai"}, index=["memory"])
store.put(("docs",), "report", {"title": "Annual Report"}, index=False)
```
"""
_validate_namespace(namespace)
@@ -795,7 +650,7 @@ class BaseStore(ABC):
List[Tuple[str, ...]]: A list of namespace tuples that match the criteria.
Each tuple represents a full namespace path up to `max_depth`.
???+ example "Examples":
??? example "Examples":
Setting max_depth=3. Given the namespaces:
```python
# Example if you have the following namespaces:
@@ -855,39 +710,6 @@ class BaseStore(ABC):
Returns:
List of items matching the search criteria.
???+ example "Examples"
Basic filtering:
```python
# Search for documents with specific metadata
results = await store.asearch(
("docs",),
filter={"type": "article", "status": "published"}
)
```
Natural language search (requires vector store implementation):
```python
# Initialize store with embedding configuration
store = YourStore( # e.g., InMemoryStore, AsyncPostgresStore
index={
"dims": 1536, # embedding dimensions
"embed": your_embedding_function, # function to create embeddings
"fields": ["text"] # fields to embed
}
)
# Search for semantically similar documents
results = await store.asearch(
("docs",),
query="machine learning applications in healthcare",
filter={"type": "research_paper"},
limit=5
)
```
Note: Natural language search support depends on your store implementation
and requires proper embedding configuration.
"""
return (
await self.abatch(
@@ -912,10 +734,7 @@ class BaseStore(ABC):
value: Dictionary containing the item's data. Must contain string keys
and JSON-serializable values.
index: Controls how the item's fields are indexed for search:
- None (default): Use `fields` you configured when creating the store (if any)
If you do not initialize the store with indexing capabilities,
the `index` parameter will be ignored
- None (default): Use store's default indexing configuration
- False: Disable indexing for this item
- list[str]: List of field paths to index, supporting:
- Nested fields: "metadata.title"
@@ -923,32 +742,25 @@ class BaseStore(ABC):
- Specific indices: "authors[0].name"
Note:
Indexing support depends on your store implementation.
If you do not initialize the store with indexing capabilities,
the `index` parameter will be ignored.
Indexing capabilities depend on your store implementation.
Some implementations may support only a subset of indexing features.
???+ example "Examples"
Store item. Indexing depends on how you configure the store.
??? example "Examples"
Simple storage without special indexing:
```python
await store.aput(("docs",), "report", {"memory": "Will likes ai"})
await store.aput(("docs",), "report", {"title": "Annual Report"})
```
Do not index item for semantic search. Still accessible through get()
and search() operations but won't have a vector representation.
```python
await store.aput(("docs",), "report", {"memory": "Will likes ai"}, index=False)
```
Index specific fields for search (if store configured to index items):
Index specific fields for search:
```python
await store.aput(
("docs",),
"report",
{
"memory": "Will likes ai",
"context": [{"content": "..."}, {"content": "..."}]
"title": "Q4 Report",
"chapters": [{"content": "..."}, {"content": "..."}]
},
index=["memory", "context[*].content"]
index=["title", "chapters[*].content"]
)
```
"""
@@ -990,7 +802,7 @@ class BaseStore(ABC):
List[Tuple[str, ...]]: A list of namespace tuples that match the criteria.
Each tuple represents a full namespace path up to `max_depth`.
???+ example "Examples"
??? example "Examples"
Setting max_depth=3 with existing namespaces:
```python
# Given the following namespaces:
+2 -109
View File
@@ -1,7 +1,6 @@
import asyncio
import functools
import weakref
from typing import Any, Callable, Iterable, Literal, Optional, TypeVar, Union
from typing import Any, Literal, Optional, Union
from langgraph.store.base import (
BaseStore,
@@ -12,39 +11,11 @@ from langgraph.store.base import (
NamespacePath,
Op,
PutOp,
Result,
SearchItem,
SearchOp,
_validate_namespace,
)
F = TypeVar("F", bound=Callable)
def _check_loop(func: F) -> F:
@functools.wraps(func)
def wrapper(store: "AsyncBatchedBaseStore", *args: Any, **kwargs: Any) -> Any:
method_name: str = func.__name__
try:
current_loop = asyncio.get_running_loop()
if current_loop is store._loop:
replacement_str = (
f"Specifically, replace `store.{method_name}(...)` with `await store.a{method_name}(...)"
if method_name
else "For example, replace `store.get(...)` with `await store.aget(...)`"
)
raise asyncio.InvalidStateError(
f"Synchronous calls to {store.__class__.__name__} detected in the main event loop. "
"This can lead to deadlocks or performance issues. "
"Please use the asynchronous interface for main thread operations. "
f"{replacement_str} "
)
except RuntimeError:
pass
return func(store, *args, **kwargs)
return wrapper
class AsyncBatchedBaseStore(BaseStore):
"""Efficiently batch operations in a background task."""
@@ -52,7 +23,6 @@ class AsyncBatchedBaseStore(BaseStore):
__slots__ = ("_loop", "_aqueue", "_task")
def __init__(self) -> None:
super().__init__()
self._loop = asyncio.get_running_loop()
self._aqueue: dict[asyncio.Future, Op] = {}
self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self)))
@@ -129,82 +99,6 @@ class AsyncBatchedBaseStore(BaseStore):
self._aqueue[fut] = op
return await fut
@_check_loop
def batch(self, ops: Iterable[Op]) -> list[Result]:
return asyncio.run_coroutine_threadsafe(self.abatch(ops), self._loop).result()
@_check_loop
def get(
self,
namespace: tuple[str, ...],
key: str,
) -> Optional[Item]:
return asyncio.run_coroutine_threadsafe(
self.aget(namespace, key=key), self._loop
).result()
@_check_loop
def search(
self,
namespace_prefix: tuple[str, ...],
/,
*,
query: Optional[str] = None,
filter: Optional[dict[str, Any]] = None,
limit: int = 10,
offset: int = 0,
) -> list[SearchItem]:
return asyncio.run_coroutine_threadsafe(
self.asearch(
namespace_prefix, query=query, filter=filter, limit=limit, offset=offset
),
self._loop,
).result()
@_check_loop
def put(
self,
namespace: tuple[str, ...],
key: str,
value: dict[str, Any],
index: Optional[Union[Literal[False], list[str]]] = None,
) -> None:
_validate_namespace(namespace)
asyncio.run_coroutine_threadsafe(
self.aput(namespace, key=key, value=value, index=index), self._loop
).result()
@_check_loop
def delete(
self,
namespace: tuple[str, ...],
key: str,
) -> None:
asyncio.run_coroutine_threadsafe(
self.adelete(namespace, key=key), self._loop
).result()
@_check_loop
def list_namespaces(
self,
*,
prefix: Optional[NamespacePath] = None,
suffix: Optional[NamespacePath] = None,
max_depth: Optional[int] = None,
limit: int = 100,
offset: int = 0,
) -> list[tuple[str, ...]]:
return asyncio.run_coroutine_threadsafe(
self.alist_namespaces(
prefix=prefix,
suffix=suffix,
max_depth=max_depth,
limit=limit,
offset=offset,
),
self._loop,
).result()
def _dedupe_ops(values: list[Op]) -> tuple[Optional[list[int]], list[Op]]:
"""Dedupe operations while preserving order for results.
@@ -250,8 +144,7 @@ def _dedupe_ops(values: list[Op]) -> tuple[Optional[list[int]], list[Op]]:
async def _run(
aqueue: dict[asyncio.Future, Op],
store: weakref.ReferenceType[BaseStore],
aqueue: dict[asyncio.Future, Op], store: weakref.ReferenceType[BaseStore]
) -> None:
while True:
await asyncio.sleep(0)
@@ -1,102 +1,31 @@
"""In-memory dictionary-backed store with optional vector search.
"""In-memory key-value store.
!!! example "Examples"
A lightweight store implementation using Python dictionaries. Supports basic
key-value operations and vector search when configured with embeddings.
Examples:
Basic key-value storage:
```python
from langgraph.store.memory import InMemoryStore
store = InMemoryStore()
store.put(("users", "123"), "prefs", {"theme": "dark"})
item = store.get(("users", "123"), "prefs")
store = InMemoryStore()
store.put(("users", "123"), "prefs", {"theme": "dark"})
item = store.get(("users", "123"), "prefs")
```
Vector search using LangChain embeddings:
```python
from langchain.embeddings import init_embeddings
from langgraph.store.memory import InMemoryStore
store = InMemoryStore(
index={
Vector search with embeddings:
from langchain_openai import OpenAIEmbeddings
store = InMemoryStore(index={
"dims": 1536,
"embed": init_embeddings("openai:text-embedding-3-small")
}
)
"embed": OpenAIEmbeddings(model="text-embedding-3-small"),
})
# Store documents
store.put(("docs",), "doc1", {"text": "Python tutorial"})
store.put(("docs",), "doc2", {"text": "TypeScript guide"})
# Store documents
store.put(("docs",), "doc1", {"text": "Python tutorial"})
store.put(("docs",), "doc2", {"text": "TypeScript guide"})
# Search by similarity
results = store.search(("docs",), query="python programming")
```
# Search by similarity
results = store.search(("docs",), query="python programming")
Vector search using OpenAI SDK directly:
```python
from openai import OpenAI
from langgraph.store.memory import InMemoryStore
client = OpenAI()
def embed_texts(texts: list[str]) -> list[list[float]]:
response = client.embeddings.create(
model="text-embedding-3-small",
input=texts
)
return [e.embedding for e in response.data]
store = InMemoryStore(
index={
"dims": 1536,
"embed": embed_texts
}
)
# Store documents
store.put(("docs",), "doc1", {"text": "Python tutorial"})
store.put(("docs",), "doc2", {"text": "TypeScript guide"})
# Search by similarity
results = store.search(("docs",), query="python programming")
```
Async vector search using OpenAI SDK:
```python
from openai import AsyncOpenAI
from langgraph.store.memory import InMemoryStore
client = AsyncOpenAI()
async def aembed_texts(texts: list[str]) -> list[list[float]]:
response = await client.embeddings.create(
model="text-embedding-3-small",
input=texts
)
return [e.embedding for e in response.data]
store = InMemoryStore(
index={
"dims": 1536,
"embed": aembed_texts
}
)
# Store documents
await store.aput(("docs",), "doc1", {"text": "Python tutorial"})
await store.aput(("docs",), "doc2", {"text": "TypeScript guide"})
# Search by similarity
results = await store.asearch(("docs",), query="python programming")
```
Warning:
This store keeps all data in memory. Data is lost when the process exits.
For persistence, use a database-backed store like PostgresStore.
Tip:
For vector search, install numpy for better performance:
```bash
pip install numpy
```
Note:
For production use cases requiring persistence, use a database-backed store instead.
"""
import asyncio
@@ -133,18 +62,17 @@ logger = logging.getLogger(__name__)
class InMemoryStore(BaseStore):
"""In-memory dictionary-backed store with optional vector search.
!!! example "Examples"
Examples:
Basic key-value storage:
store = InMemoryStore()
store.put(("users", "123"), "prefs", {"theme": "dark"})
item = store.get(("users", "123"), "prefs")
Vector search with embeddings:
from langchain.embeddings import init_embeddings
from langchain_openai import OpenAIEmbeddings
store = InMemoryStore(index={
"dims": 1536,
"embed": init_embeddings("openai:text-embedding-3-small"),
"fields": ["text"],
"embed": OpenAIEmbeddings(model="text-embedding-3-small"),
})
# Store documents
@@ -154,11 +82,6 @@ class InMemoryStore(BaseStore):
# Search by similarity
results = store.search(("docs",), query="python programming")
Note:
Semantic search is disabled by default. You can enable it by providing an `index` configuration
when creating the store. Without this configuration, all `index` arguments passed to
`put` or `aput`will have no effect.
Warning:
This store keeps all data in memory. Data is lost when the process exits.
For persistence, use a database-backed store like PostgresStore.
@@ -490,8 +413,6 @@ def _cosine_similarity(X: list[float], Y: list[list[float]]) -> list[float]:
Compute cosine similarity between a vector X and a matrix Y.
Lazy import numpy for efficiency.
"""
if not Y:
return []
if _check_numpy():
import numpy as np # type: ignore
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-checkpoint"
version = "2.0.8"
version = "2.0.5"
description = "Library with base interfaces for LangGraph checkpoint savers."
authors = []
license = "MIT"
+14 -24
View File
@@ -511,6 +511,19 @@ def dockerfile(save_path: str, config: pathlib.Path, add_docker_compose: bool) -
)
@click.argument("path", required=False)
@click.option(
"--template",
type=str,
help=TEMPLATE_HELP_STRING,
)
@cli.command("new", help="🌱 Create a new LangGraph project from a template.")
@log_command
def new(path: Optional[str], template: Optional[str]) -> None:
"""Create a new LangGraph project from a template."""
return create_new(path, template)
@click.option(
"--host",
default="127.0.0.1",
@@ -550,12 +563,6 @@ def dockerfile(save_path: str, config: pathlib.Path, add_docker_compose: bool) -
type=int,
help="Enable remote debugging by listening on specified port. Requires debugpy to be installed",
)
@click.option(
"--wait-for-client",
is_flag=True,
help="Wait for a debugger client to connect to the debug port before starting the server",
default=False,
)
@cli.command(
"dev",
help="🏃‍♀️‍➡️ Run LangGraph API server in development mode with hot reloading and debugging support",
@@ -569,7 +576,6 @@ def dev(
n_jobs_per_worker: Optional[int],
no_browser: bool,
debug_port: Optional[int],
wait_for_client: bool,
):
"""CLI entrypoint for running the LangGraph API server."""
try:
@@ -602,7 +608,6 @@ def dev(
sys.path.append(str(dep_path))
graphs = config_json.get("graphs", {})
run_server(
host,
port,
@@ -611,25 +616,10 @@ def dev(
n_jobs_per_worker=n_jobs_per_worker,
open_browser=not no_browser,
debug_port=debug_port,
env=config_json.get("env"),
store=config_json.get("store"),
wait_for_client=wait_for_client,
env=config_json.get("env", None),
)
@click.argument("path", required=False)
@click.option(
"--template",
type=str,
help=TEMPLATE_HELP_STRING,
)
@cli.command("new", help="🌱 Create a new LangGraph project from a template.")
@log_command
def new(path: Optional[str], template: Optional[str]) -> None:
"""Create a new LangGraph project from a template."""
return create_new(path, template)
def prepare_args_and_stdin(
*,
capabilities: DockerCapabilities,
+5 -59
View File
@@ -10,44 +10,7 @@ MIN_NODE_VERSION = "20"
MIN_PYTHON_VERSION = "3.11"
class IndexConfig(TypedDict, total=False):
"""Configuration for indexing documents for semantic search in the store."""
dims: int
"""Number of dimensions in the embedding vectors.
Common embedding models have the following dimensions:
- openai:text-embedding-3-large: 3072
- openai:text-embedding-3-small: 1536
- openai:text-embedding-ada-002: 1536
- cohere:embed-english-v3.0: 1024
- cohere:embed-english-light-v3.0: 384
- cohere:embed-multilingual-v3.0: 1024
- cohere:embed-multilingual-light-v3.0: 384
"""
embed: str
"""Optional model (string) to generate embeddings from text or path to model or function.
Examples:
- "openai:text-embedding-3-large"
- "cohere:embed-multilingual-v3.0"
- "src/app.py:embeddings
"""
fields: Optional[list[str]]
"""Fields to extract text from for embedding generation.
Defaults to the root ["$"], which embeds the json object as a whole.
"""
class StoreConfig(TypedDict, total=False):
embed: Optional[IndexConfig]
"""Configuration for vector embeddings in store."""
class Config(TypedDict, total=False):
class Config(TypedDict):
python_version: str
node_version: Optional[str]
pip_config_file: Optional[str]
@@ -55,7 +18,6 @@ class Config(TypedDict, total=False):
dependencies: list[str]
graphs: dict[str, str]
env: Union[dict[str, str], str]
store: Optional[StoreConfig]
def _parse_version(version_str: str) -> tuple[int, int]:
@@ -87,7 +49,6 @@ def validate_config(config: Config) -> Config:
"dockerfile_lines": config.get("dockerfile_lines", []),
"graphs": config.get("graphs", {}),
"env": config.get("env", {}),
"store": config.get("store"),
}
if config.get("node_version")
else {
@@ -97,7 +58,6 @@ def validate_config(config: Config) -> Config:
"dependencies": config.get("dependencies", []),
"graphs": config.get("graphs", {}),
"env": config.get("env", {}),
"store": config.get("store"),
}
)
@@ -392,14 +352,7 @@ RUN set -ex && \\
],
)
)
store_config = config.get("store")
env_additional_config = (
""
if not store_config
else f"""
ENV LANGGRAPH_STORE='{json.dumps(store_config)}'
"""
)
return f"""FROM {base_image}:{config['python_version']}
{os.linesep.join(config["dockerfile_lines"])}
@@ -407,7 +360,7 @@ ENV LANGGRAPH_STORE='{json.dumps(store_config)}'
{installs}
RUN {pip_install} -e /deps/*
{env_additional_config}
ENV LANGSERVE_GRAPHS='{json.dumps(config["graphs"])}'
{f"WORKDIR {local_deps.working_dir}" if local_deps.working_dir else ""}"""
@@ -437,14 +390,7 @@ def node_config_to_docker(config_path: pathlib.Path, config: Config, base_image:
install_cmd = "npm ci"
else:
install_cmd = "npm i"
store_config = config.get("store")
env_additional_config = (
""
if not store_config
else f"""
ENV LANGGRAPH_STORE='{json.dumps(store_config)}'
"""
)
return f"""FROM {base_image}:{config['node_version']}
{os.linesep.join(config["dockerfile_lines"])}
@@ -452,7 +398,7 @@ ENV LANGGRAPH_STORE='{json.dumps(store_config)}'
ADD . {faux_path}
RUN cd {faux_path} && {install_cmd}
{env_additional_config}
ENV LANGSERVE_GRAPHS='{json.dumps(config["graphs"])}'
WORKDIR {faux_path}
+21 -54
View File
@@ -565,13 +565,13 @@ langgraph-sdk = ">=0.1.32,<0.2.0"
[[package]]
name = "langgraph-api"
version = "0.0.6"
version = "0.0.2"
description = ""
optional = true
python-versions = "<4.0,>=3.11.0"
files = [
{file = "langgraph_api-0.0.6-py3-none-any.whl", hash = "sha256:f64b13959d721143f6a023af5b9ffc9aa054064af98d21d5d8090cda7e7bffd2"},
{file = "langgraph_api-0.0.6.tar.gz", hash = "sha256:badac44fa1ec979509e56fc0da57eeb5f278ee5871f27803f73ea6d8822c21b9"},
{file = "langgraph_api-0.0.2-py3-none-any.whl", hash = "sha256:7a30fb21987572eacc93dd1c69c2155c17957afed71dde18d6f47992b3124d65"},
{file = "langgraph_api-0.0.2.tar.gz", hash = "sha256:b751afca96cb6db67fe2f48e798ada27a4df068f0df86b36d2b8eee52344bbf0"},
]
[package.dependencies]
@@ -579,8 +579,8 @@ cryptography = ">=43.0.3,<44.0.0"
httpx = ">=0.27.0"
jsonschema-rs = ">=0.25.0,<0.26.0"
langchain-core = ">=0.2.38,<0.4.0"
langgraph = ">=0.2.52,<0.3.0"
langgraph-checkpoint = ">=2.0.7,<3.0"
langgraph = ">=0.2.52"
langgraph-checkpoint = ">=2.0.5,<3.0"
langsmith = ">=0.1.63,<0.2.0"
orjson = ">=3.10.1"
pyjwt = ">=2.9.0,<3.0.0"
@@ -593,13 +593,13 @@ watchfiles = ">=0.13"
[[package]]
name = "langgraph-checkpoint"
version = "2.0.7"
version = "2.0.6"
description = "Library with base interfaces for LangGraph checkpoint savers."
optional = true
python-versions = "<4.0.0,>=3.9.0"
files = [
{file = "langgraph_checkpoint-2.0.7-py3-none-any.whl", hash = "sha256:9709f672e1c5a47e13352067c2ffa114dd91d443967b7ce8a1d36d6fc170370e"},
{file = "langgraph_checkpoint-2.0.7.tar.gz", hash = "sha256:88d648a331d20aa8ce65280de34a34a9190380b004f6afcc5f9894fe3abeed08"},
{file = "langgraph_checkpoint-2.0.6-py3-none-any.whl", hash = "sha256:2878283c3ee2519bf180df9b7b7155b73fa05eb63b1af9600a03e03a930d8c53"},
{file = "langgraph_checkpoint-2.0.6.tar.gz", hash = "sha256:69ab9c61c4e2992264671f55579c24070b7b6cedc105a33da3fba6526df248cf"},
]
[package.dependencies]
@@ -608,13 +608,13 @@ msgpack = ">=1.1.0,<2.0.0"
[[package]]
name = "langgraph-sdk"
version = "0.1.40"
version = "0.1.36"
description = "SDK for interacting with LangGraph API"
optional = true
python-versions = "<4.0.0,>=3.9.0"
files = [
{file = "langgraph_sdk-0.1.40-py3-none-any.whl", hash = "sha256:8810cca5e4144cf3a5441fc76b4ee6e658ec95f932d3a0bf9ad63de117e925b9"},
{file = "langgraph_sdk-0.1.40.tar.gz", hash = "sha256:ab2719ac7274612a791a7a0ad9395d250357106cba8ba81bca9968fc91009af2"},
{file = "langgraph_sdk-0.1.36-py3-none-any.whl", hash = "sha256:b11e1f0bc67631134d09d50c812dc73f9eb30394764ae1144d7d2a786a715355"},
{file = "langgraph_sdk-0.1.36.tar.gz", hash = "sha256:2a2c651b7851ba15aeaab7e4e3ea7fd8357ef1cb0b592f264916fa990cdda6e7"},
]
[package.dependencies]
@@ -624,13 +624,13 @@ orjson = ">=3.10.1"
[[package]]
name = "langsmith"
version = "0.1.147"
version = "0.1.146"
description = "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform."
optional = true
python-versions = "<4.0,>=3.8.1"
files = [
{file = "langsmith-0.1.147-py3-none-any.whl", hash = "sha256:7166fc23b965ccf839d64945a78e9f1157757add228b086141eb03a60d699a15"},
{file = "langsmith-0.1.147.tar.gz", hash = "sha256:2e933220318a4e73034657103b3b1a3a6109cc5db3566a7e8e03be8d6d7def7a"},
{file = "langsmith-0.1.146-py3-none-any.whl", hash = "sha256:9d062222f1a32c9b047dab0149b24958f988989cd8d4a5f9139ff959a51e59d8"},
{file = "langsmith-0.1.146.tar.gz", hash = "sha256:ead8b0b9d5b6cd3ac42937ec48bdf09d4afe7ca1bba22dc05eb65591a18106f8"},
]
[package.dependencies]
@@ -643,9 +643,6 @@ pydantic = [
requests = ">=2,<3"
requests-toolbelt = ">=1.0.0,<2.0.0"
[package.extras]
langsmith-pyo3 = ["langsmith-pyo3 (>=0.1.0rc2,<0.2.0)"]
[[package]]
name = "msgpack"
version = "1.1.0"
@@ -1038,13 +1035,13 @@ typing-extensions = ">=4.6.0,<4.7.0 || >4.7.0"
[[package]]
name = "pyjwt"
version = "2.10.1"
version = "2.10.0"
description = "JSON Web Token implementation in Python"
optional = true
python-versions = ">=3.9"
files = [
{file = "PyJWT-2.10.1-py3-none-any.whl", hash = "sha256:dcdd193e30abefd5debf142f9adfcdd2b58004e644f25406ffaebd50bd98dacb"},
{file = "pyjwt-2.10.1.tar.gz", hash = "sha256:3cc5772eb20009233caf06e9d8a0577824723b44e6648ee0a2aedb6cf9381953"},
{file = "PyJWT-2.10.0-py3-none-any.whl", hash = "sha256:543b77207db656de204372350926bed5a86201c4cbff159f623f79c7bb487a15"},
{file = "pyjwt-2.10.0.tar.gz", hash = "sha256:7628a7eb7938959ac1b26e819a1df0fd3259505627b575e4bad6d08f76db695c"},
]
[package.extras]
@@ -1345,43 +1342,13 @@ test = ["pytest", "tornado (>=4.5)", "typeguard"]
[[package]]
name = "tomli"
version = "2.2.1"
version = "2.1.0"
description = "A lil' TOML parser"
optional = false
python-versions = ">=3.8"
files = [
{file = "tomli-2.2.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:678e4fa69e4575eb77d103de3df8a895e1591b48e740211bd1067378c69e8249"},
{file = "tomli-2.2.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:023aa114dd824ade0100497eb2318602af309e5a55595f76b626d6d9f3b7b0a6"},
{file = "tomli-2.2.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ece47d672db52ac607a3d9599a9d48dcb2f2f735c6c2d1f34130085bb12b112a"},
{file = "tomli-2.2.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6972ca9c9cc9f0acaa56a8ca1ff51e7af152a9f87fb64623e31d5c83700080ee"},
{file = "tomli-2.2.1-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c954d2250168d28797dd4e3ac5cf812a406cd5a92674ee4c8f123c889786aa8e"},
{file = "tomli-2.2.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:8dd28b3e155b80f4d54beb40a441d366adcfe740969820caf156c019fb5c7ec4"},
{file = "tomli-2.2.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:e59e304978767a54663af13c07b3d1af22ddee3bb2fb0618ca1593e4f593a106"},
{file = "tomli-2.2.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:33580bccab0338d00994d7f16f4c4ec25b776af3ffaac1ed74e0b3fc95e885a8"},
{file = "tomli-2.2.1-cp311-cp311-win32.whl", hash = "sha256:465af0e0875402f1d226519c9904f37254b3045fc5084697cefb9bdde1ff99ff"},
{file = "tomli-2.2.1-cp311-cp311-win_amd64.whl", hash = "sha256:2d0f2fdd22b02c6d81637a3c95f8cd77f995846af7414c5c4b8d0545afa1bc4b"},
{file = "tomli-2.2.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:4a8f6e44de52d5e6c657c9fe83b562f5f4256d8ebbfe4ff922c495620a7f6cea"},
{file = "tomli-2.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8d57ca8095a641b8237d5b079147646153d22552f1c637fd3ba7f4b0b29167a8"},
{file = "tomli-2.2.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4e340144ad7ae1533cb897d406382b4b6fede8890a03738ff1683af800d54192"},
{file = "tomli-2.2.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:db2b95f9de79181805df90bedc5a5ab4c165e6ec3fe99f970d0e302f384ad222"},
{file = "tomli-2.2.1-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:40741994320b232529c802f8bc86da4e1aa9f413db394617b9a256ae0f9a7f77"},
{file = "tomli-2.2.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:400e720fe168c0f8521520190686ef8ef033fb19fc493da09779e592861b78c6"},
{file = "tomli-2.2.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:02abe224de6ae62c19f090f68da4e27b10af2b93213d36cf44e6e1c5abd19fdd"},
{file = "tomli-2.2.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:b82ebccc8c8a36f2094e969560a1b836758481f3dc360ce9a3277c65f374285e"},
{file = "tomli-2.2.1-cp312-cp312-win32.whl", hash = "sha256:889f80ef92701b9dbb224e49ec87c645ce5df3fa2cc548664eb8a25e03127a98"},
{file = "tomli-2.2.1-cp312-cp312-win_amd64.whl", hash = "sha256:7fc04e92e1d624a4a63c76474610238576942d6b8950a2d7f908a340494e67e4"},
{file = "tomli-2.2.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:f4039b9cbc3048b2416cc57ab3bda989a6fcf9b36cf8937f01a6e731b64f80d7"},
{file = "tomli-2.2.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:286f0ca2ffeeb5b9bd4fcc8d6c330534323ec51b2f52da063b11c502da16f30c"},
{file = "tomli-2.2.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a92ef1a44547e894e2a17d24e7557a5e85a9e1d0048b0b5e7541f76c5032cb13"},
{file = "tomli-2.2.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9316dc65bed1684c9a98ee68759ceaed29d229e985297003e494aa825ebb0281"},
{file = "tomli-2.2.1-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e85e99945e688e32d5a35c1ff38ed0b3f41f43fad8df0bdf79f72b2ba7bc5272"},
{file = "tomli-2.2.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ac065718db92ca818f8d6141b5f66369833d4a80a9d74435a268c52bdfa73140"},
{file = "tomli-2.2.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:d920f33822747519673ee656a4b6ac33e382eca9d331c87770faa3eef562aeb2"},
{file = "tomli-2.2.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:a198f10c4d1b1375d7687bc25294306e551bf1abfa4eace6650070a5c1ae2744"},
{file = "tomli-2.2.1-cp313-cp313-win32.whl", hash = "sha256:d3f5614314d758649ab2ab3a62d4f2004c825922f9e370b29416484086b264ec"},
{file = "tomli-2.2.1-cp313-cp313-win_amd64.whl", hash = "sha256:a38aa0308e754b0e3c67e344754dff64999ff9b513e691d0e786265c93583c69"},
{file = "tomli-2.2.1-py3-none-any.whl", hash = "sha256:cb55c73c5f4408779d0cf3eef9f762b9c9f147a77de7b258bef0a5628adc85cc"},
{file = "tomli-2.2.1.tar.gz", hash = "sha256:cd45e1dc79c835ce60f7404ec8119f2eb06d38b1deba146f07ced3bbc44505ff"},
{file = "tomli-2.1.0-py3-none-any.whl", hash = "sha256:a5c57c3d1c56f5ccdf89f6523458f60ef716e210fc47c4cfb188c5ba473e0391"},
{file = "tomli-2.1.0.tar.gz", hash = "sha256:3f646cae2aec94e17d04973e4249548320197cfabdf130015d023de4b74d8ab8"},
]
[[package]]
@@ -1561,4 +1528,4 @@ inmem = ["langgraph-api", "python-dotenv"]
[metadata]
lock-version = "2.0"
python-versions = "^3.9.0,<4.0"
content-hash = "8eaaa66d9e6e447699e3bcee336dfe779b58c956f8c2ad6678008a07be935838"
content-hash = "3d655bb578e20219e19152d4a3d86be370fe3be61b5559847f0204dfff499b4a"
+2 -2
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-cli"
version = "0.1.61"
version = "0.1.59"
description = "CLI for interacting with LangGraph API"
authors = []
license = "MIT"
@@ -14,7 +14,7 @@ langgraph = "langgraph_cli.cli:cli"
[tool.poetry.dependencies]
python = "^3.9.0,<4.0"
click = "^8.1.7"
langgraph-api = { version = ">=0.0.6,<0.1.0", optional = true, python = ">=3.11,<4.0" }
langgraph-api = { version = ">=0.0.2,<0.1.0", optional = true, python = ">=3.11,<4.0" }
python-dotenv = { version = ">=0.8.0", optional = true }
[tool.poetry.group.dev.dependencies]
-2
View File
@@ -30,7 +30,6 @@ def test_validate_config():
"pip_config_file": None,
"dockerfile_lines": [],
"env": {},
"store": None,
**expected_config,
}
actual_config = validate_config(expected_config)
@@ -47,7 +46,6 @@ def test_validate_config():
"agent": "./agent.py:graph",
},
"env": env,
"store": None,
}
actual_config = validate_config(expected_config)
assert actual_config == expected_config
-6
View File
@@ -48,12 +48,6 @@ test:
make stop-postgres; \
exit $$EXIT_CODE
test_parallel:
make start-postgres && poetry run pytest -n auto --dist worksteal $(TEST); \
EXIT_CODE=$$?; \
make stop-postgres; \
exit $$EXIT_CODE
WORKERS ?= auto
XDIST_ARGS := $(if $(WORKERS),-n $(WORKERS) --dist worksteal,)
MAXFAIL ?=
+2 -2
View File
@@ -1,5 +1,5 @@
from abc import ABC, abstractmethod
from typing import Any, Generic, Optional, Sequence, TypeVar
from typing import Any, Generic, Optional, Sequence, Type, TypeVar
from typing_extensions import Self
@@ -13,7 +13,7 @@ C = TypeVar("C")
class BaseChannel(Generic[Value, Update, C], ABC):
__slots__ = ("key", "typ")
def __init__(self, typ: Any, key: str = "") -> None:
def __init__(self, typ: Type[Any], key: str = "") -> None:
self.typ = typ
self.key = key
+2 -8
View File
@@ -40,16 +40,12 @@ SCHEDULED = sys.intern("__scheduled__")
# marker to signal node was scheduled (in distributed mode)
TASKS = sys.intern("__pregel_tasks")
# for Send objects returned by nodes/edges, corresponds to PUSH below
RETURN = sys.intern("__return__")
# for writes of a task where we simply record the return value
# --- Reserved config.configurable keys ---
CONFIG_KEY_SEND = sys.intern("__pregel_send")
# holds the `write` function that accepts writes to state/edges/reserved keys
CONFIG_KEY_READ = sys.intern("__pregel_read")
# holds the `read` function that returns a copy of the current state
CONFIG_KEY_CALL = sys.intern("__pregel_call")
# holds the `call` function that accepts a node/func, args and returns a future
CONFIG_KEY_CHECKPOINTER = sys.intern("__pregel_checkpointer")
# holds a `BaseCheckpointSaver` passed from parent graph to child graphs
CONFIG_KEY_STREAM = sys.intern("__pregel_stream")
@@ -76,11 +72,9 @@ CONFIG_KEY_CHECKPOINT_ID = sys.intern("checkpoint_id")
CONFIG_KEY_CHECKPOINT_NS = sys.intern("checkpoint_ns")
# holds the current checkpoint_ns, "" for root graph
CONFIG_KEY_NODE_FINISHED = sys.intern("__pregel_node_finished")
# callback to be called when a node is finished
CONFIG_KEY_RESUME_VALUE = sys.intern("__pregel_resume_value")
# holds the value that "answers" an interrupt() call
CONFIG_KEY_WRITES = sys.intern("__pregel_writes")
# read-only list of existing task writes
CONFIG_KEY_SCRATCHPAD = sys.intern("__pregel_scratchpad")
# holds a mutable dict for temporary storage scoped to the current task
# --- Other constants ---
PUSH = sys.intern("__pregel_push")
-120
View File
@@ -1,120 +0,0 @@
import asyncio
import concurrent
import concurrent.futures
import inspect
import types
from functools import partial, update_wrapper
from typing import (
Any,
Awaitable,
Callable,
Optional,
TypeVar,
Union,
overload,
)
from typing_extensions import ParamSpec
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import END, START, TAG_HIDDEN
from langgraph.pregel import Pregel
from langgraph.pregel.call import get_runnable_for_func
from langgraph.pregel.read import PregelNode
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
from langgraph.store.base import BaseStore
from langgraph.types import RetryPolicy, StreamMode, StreamWriter
P = ParamSpec("P")
P1 = TypeVar("P1")
T = TypeVar("T")
def call(
func: Callable[[P1], T],
input: P1,
*,
retry: Optional[RetryPolicy] = None,
) -> concurrent.futures.Future[T]:
from langgraph.constants import CONFIG_KEY_CALL
from langgraph.utils.config import get_configurable
conf = get_configurable()
impl = conf[CONFIG_KEY_CALL]
fut = impl(func, input, retry=retry)
return fut
@overload
def task(
*, retry: Optional[RetryPolicy] = None
) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, asyncio.Future[T]]]: ...
@overload
def task( # type: ignore[overload-cannot-match]
*, retry: Optional[RetryPolicy] = None
) -> Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]]: ...
def task(
*, retry: Optional[RetryPolicy] = None
) -> Union[
Callable[[Callable[P, Awaitable[T]]], Callable[P, asyncio.Future[T]]],
Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]],
]:
def _task(func: Callable[P, T]) -> Callable[P, concurrent.futures.Future[T]]:
return update_wrapper(partial(call, func, retry=retry), func)
return _task
def entrypoint(
*,
checkpointer: Optional[BaseCheckpointSaver] = None,
store: Optional[BaseStore] = None,
) -> Callable[[types.FunctionType], Pregel]:
def _imp(func: types.FunctionType) -> Pregel:
if inspect.isgeneratorfunction(func):
def gen_wrapper(*args: Any, writer: StreamWriter, **kwargs: Any) -> Any:
for chunk in func(*args, **kwargs):
writer(chunk)
bound = get_runnable_for_func(gen_wrapper)
stream_mode: StreamMode = "custom"
elif inspect.isasyncgenfunction(func):
async def agen_wrapper(
*args: Any, writer: StreamWriter, **kwargs: Any
) -> Any:
async for chunk in func(*args, **kwargs):
writer(chunk)
bound = get_runnable_for_func(agen_wrapper)
stream_mode = "custom"
else:
bound = get_runnable_for_func(func)
stream_mode = "updates"
return Pregel(
nodes={
func.__name__: PregelNode(
bound=bound,
triggers=[START],
channels=[START],
writers=[ChannelWrite([ChannelWriteEntry(END)], tags=[TAG_HIDDEN])],
)
},
channels={START: EphemeralValue(Any), END: LastValue(Any, END)},
input_channels=START,
output_channels=END,
stream_channels=END,
stream_mode=stream_mode,
checkpointer=checkpointer,
store=store,
)
return _imp
+2 -1
View File
@@ -1,12 +1,13 @@
from langgraph.graph.graph import END, START, Graph
from langgraph.graph.message import MessageGraph, MessagesState, add_messages
from langgraph.graph.state import StateGraph
from langgraph.graph.state import GraphCommand, StateGraph
__all__ = [
"END",
"START",
"Graph",
"StateGraph",
"GraphCommand",
"MessageGraph",
"add_messages",
"MessagesState",
+4 -12
View File
@@ -374,11 +374,6 @@ class Graph:
if source not in self.nodes and source != START:
raise ValueError(f"Found edge starting at unknown node '{source}'")
if START not in all_sources:
raise ValueError(
"Graph must have an entrypoint: add at least one edge from START to another node"
)
# assemble targets
all_targets = {end for _, end in self._all_edges}
for start, branches in self.branches.items():
@@ -400,6 +395,10 @@ class Graph:
for name, spec in self.nodes.items():
if spec.ends:
all_targets.update(spec.ends)
# validate targets
for node in self.nodes:
if node not in all_targets:
raise ValueError(f"Node `{node}` is not reachable")
for target in all_targets:
if target not in self.nodes and target != END:
raise ValueError(f"Found edge ending at unknown node `{target}`")
@@ -629,10 +628,3 @@ class CompiledGraph(Pregel):
add_edge(key, end, conditional=True)
return graph
def _repr_mimebundle_(self, **kwargs: Any) -> dict[str, Any]:
"""Mime bundle used by Jupyter to display the graph"""
return {
"text/plain": repr(self),
"image/png": self.get_graph().draw_mermaid_png(),
}
+73 -99
View File
@@ -1,3 +1,4 @@
import dataclasses
import inspect
import logging
import typing
@@ -8,6 +9,7 @@ from types import FunctionType
from typing import (
Any,
Callable,
Generic,
Literal,
NamedTuple,
Optional,
@@ -51,13 +53,9 @@ from langgraph.managed.base import (
is_writable_managed_value,
)
from langgraph.pregel.read import ChannelRead, PregelNode
from langgraph.pregel.write import (
ChannelWrite,
ChannelWriteEntry,
ChannelWriteTupleEntry,
)
from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
from langgraph.store.base import BaseStore
from langgraph.types import All, Checkpointer, Command, RetryPolicy
from langgraph.types import _DC_KWARGS, All, Checkpointer, Command, N, RetryPolicy
from langgraph.utils.fields import get_field_default
from langgraph.utils.pydantic import create_model
from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable
@@ -86,6 +84,22 @@ def _get_node_name(node: RunnableLike) -> str:
raise TypeError(f"Unsupported node type: {type(node)}")
@dataclasses.dataclass(**_DC_KWARGS)
class GraphCommand(Generic[N], Command[N]):
"""One or more commands to update a StateGraph's state and go to, or send messages to nodes."""
goto: Union[str, Sequence[str]] = ()
def __repr__(self) -> str:
# get all non-None values
contents = ", ".join(
f"{key}={value!r}"
for key, value in dataclasses.asdict(self).items()
if value
)
return f"Command({contents})"
class StateNodeSpec(NamedTuple):
runnable: Runnable
metadata: Optional[dict[str, Any]]
@@ -378,7 +392,7 @@ class StateGraph(Graph):
input = input_hint
if (
(rtn := hints.get("return"))
and get_origin(rtn) is Command
and get_origin(rtn) in (Command, GraphCommand)
and (rargs := get_args(rtn))
and get_origin(rargs[0]) is Literal
and (vals := get_args(rargs[0]))
@@ -559,7 +573,6 @@ class StateGraph(Graph):
for key, node in self.nodes.items():
compiled.attach_node(key, node)
compiled.attach_branch(START, SELF, CONTROL_BRANCH, with_reader=False)
for key, node in self.nodes.items():
compiled.attach_branch(key, SELF, CONTROL_BRANCH, with_reader=False)
@@ -613,59 +626,33 @@ class CompiledStateGraph(CompiledGraph):
if is_writable_managed_value(v)
]
def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]:
def _get_root(input: Any) -> Any:
if isinstance(input, Command):
if input.graph == Command.PARENT:
return ()
return input._update_as_tuples()
elif (
isinstance(input, (list, tuple))
and input
and any(isinstance(i, Command) for i in input)
):
updates: list[tuple[str, Any]] = []
for i in input:
if isinstance(i, Command):
if i.graph == Command.PARENT:
continue
updates.extend(i._update_as_tuples())
else:
updates.append(("__root__", i))
return updates
elif input is not None:
return [("__root__", input)]
return SKIP_WRITE
return input.update
else:
return input
def _get_updates(
input: Union[None, dict, Any],
) -> Optional[Sequence[tuple[str, Any]]]:
# to avoid name collision below
node_key = key
def _get_state_key(input: Union[None, dict, Any], *, key: str) -> Any:
if input is None:
return None
return SKIP_WRITE
elif isinstance(input, dict):
return [(k, v) for k, v in input.items() if k in output_keys]
if all(k not in output_keys for k in input):
raise InvalidUpdateError(
f"Expected node {node_key} to update at least one of {output_keys}, got {input}"
)
return input.get(key, SKIP_WRITE)
elif isinstance(input, Command):
if input.graph == Command.PARENT:
return None
return input._update_as_tuples()
elif (
isinstance(input, (list, tuple))
and input
and any(isinstance(i, Command) for i in input)
):
updates: list[tuple[str, Any]] = []
for i in input:
if isinstance(i, Command):
if i.graph == Command.PARENT:
continue
updates.extend(i._update_as_tuples())
else:
updates.extend(_get_updates(i) or ())
return updates
return SKIP_WRITE
return _get_state_key(input.update, key=key)
elif get_type_hints(type(input)):
return [
(k, getattr(input, k))
for k in output_keys
if getattr(input, k, None) is not None
]
value = getattr(input, key, SKIP_WRITE)
return value if value is not None else SKIP_WRITE
else:
msg = create_error_message(
message=f"Expected dict, got {input}",
@@ -674,11 +661,14 @@ class CompiledStateGraph(CompiledGraph):
raise InvalidUpdateError(msg)
# state updaters
write_entries: list[Union[ChannelWriteEntry, ChannelWriteTupleEntry]] = [
ChannelWriteTupleEntry(
mapper=_get_root if output_keys == ["__root__"] else _get_updates
)
]
write_entries = (
[ChannelWriteEntry("__root__", skip_none=True, mapper=_get_root)]
if output_keys == ["__root__"]
else [
ChannelWriteEntry(key, mapper=partial(_get_state_key, key=key))
for key in output_keys
]
)
# add node and output channel
if key == START:
@@ -713,7 +703,7 @@ class CompiledStateGraph(CompiledGraph):
writers=[
# publish to this channel and state keys
ChannelWrite(
write_entries + [ChannelWriteEntry(key, key)],
[ChannelWriteEntry(key, key)] + write_entries,
tags=[TAG_HIDDEN],
),
],
@@ -839,54 +829,38 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
def _control_branch(value: Any) -> Sequence[Union[str, Send]]:
if isinstance(value, Send):
return [value]
commands: list[Command] = []
if isinstance(value, Command):
commands.append(value)
elif (
isinstance(value, (list, tuple))
and value
and all(isinstance(i, Command) for i in value)
):
commands.extend(value)
else:
if not isinstance(value, GraphCommand):
return EMPTY_SEQ
if value.graph == Command.PARENT:
raise ParentCommand(value)
rtn: list[Union[str, Send]] = []
for command in commands:
if command.graph == Command.PARENT:
raise ParentCommand(command)
if isinstance(command.goto, Send):
rtn.append(command.goto)
elif isinstance(command.goto, str):
rtn.append(command.goto)
else:
rtn.extend(command.goto)
if isinstance(value.goto, str):
rtn.append(value.goto)
else:
rtn.extend(value.goto)
if isinstance(value.send, Send):
rtn.append(value.send)
else:
rtn.extend(value.send)
return rtn
async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]:
if isinstance(value, Send):
return [value]
commands: list[Command] = []
if isinstance(value, Command):
commands.append(value)
elif (
isinstance(value, (list, tuple))
and value
and all(isinstance(i, Command) for i in value)
):
commands.extend(value)
else:
if not isinstance(value, GraphCommand):
return EMPTY_SEQ
if value.graph == Command.PARENT:
raise ParentCommand(value)
rtn: list[Union[str, Send]] = []
for command in commands:
if command.graph == Command.PARENT:
raise ParentCommand(command)
if isinstance(command.goto, Send):
rtn.append(command.goto)
elif isinstance(command.goto, str):
rtn.append(command.goto)
else:
rtn.extend(command.goto)
if isinstance(value.goto, str):
rtn.append(value.goto)
else:
rtn.extend(value.goto)
if isinstance(value.send, Send):
rtn.append(value.send)
else:
rtn.extend(value.send)
return rtn
@@ -959,12 +933,12 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
if hasattr(typ, "__metadata__"):
meta = typ.__metadata__
if len(meta) >= 1 and callable(meta[-1]):
sig = signature(meta[-1])
sig = signature(meta[0])
params = list(sig.parameters.values())
if len(params) == 2 and all(
p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) for p in params
):
return BinaryOperatorAggregate(typ, meta[-1])
return BinaryOperatorAggregate(typ, meta[0])
else:
raise ValueError(
f"Invalid reducer signature. Expected (a, b) -> c. Got {sig}"
+80 -200
View File
@@ -1,10 +1,15 @@
from __future__ import annotations
import asyncio
import inspect
import json
from copy import copy, deepcopy
from copy import copy
from typing import (
TYPE_CHECKING,
Any,
Callable,
Dict,
List,
Literal,
Optional,
Sequence,
@@ -20,7 +25,6 @@ from langchain_core.messages import (
AnyMessage,
ToolCall,
ToolMessage,
convert_to_messages,
)
from langchain_core.runnables import RunnableConfig
from langchain_core.runnables.config import (
@@ -31,25 +35,26 @@ from langchain_core.runnables.utils import Input
from langchain_core.tools import BaseTool, InjectedToolArg
from langchain_core.tools import tool as create_tool
from langchain_core.tools.base import get_all_basemodel_annotations
from pydantic import BaseModel
from typing_extensions import Annotated, get_args, get_origin
from langgraph.errors import GraphBubbleUp
from langgraph.store.base import BaseStore
from langgraph.types import Command
from langgraph.utils.runnable import RunnableCallable
if TYPE_CHECKING:
from pydantic import BaseModel
INVALID_TOOL_NAME_ERROR_TEMPLATE = (
"Error: {requested_tool} is not a valid tool, try one of [{available_tools}]."
)
TOOL_CALL_ERROR_TEMPLATE = "Error: {error}\n Please fix your mistakes."
def msg_content_output(output: Any) -> Union[str, list[dict]]:
def msg_content_output(output: Any) -> str | List[dict]:
recognized_content_block_types = ("image", "image_url", "text", "json")
if isinstance(output, str):
return output
elif isinstance(output, list) and all(
elif all(
[
isinstance(x, dict) and x.get("type") in recognized_content_block_types
for x in output
@@ -90,7 +95,7 @@ def _handle_tool_error(
return content
def _infer_handled_types(handler: Callable[..., str]) -> tuple[type[Exception], ...]:
def _infer_handled_types(handler: Callable[..., str]) -> tuple[type[Exception]]:
sig = inspect.signature(handler)
params = list(sig.parameters.values())
if params:
@@ -189,9 +194,9 @@ class ToolNode(RunnableCallable):
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.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:
@@ -212,31 +217,12 @@ class ToolNode(RunnableCallable):
*,
store: BaseStore,
) -> Any:
tool_calls, input_type = self._parse_input(input, store)
tool_calls, output_type = self._parse_input(input, store)
config_list = get_config_list(config, len(tool_calls))
input_types = [input_type] * len(tool_calls)
with get_executor_for_config(config) as executor:
outputs = [
*executor.map(self._run_one, tool_calls, input_types, config_list)
]
# preserve existing behavior for non-command tool outputs for backwards compatibility
if not any(isinstance(output, Command) for output in outputs):
# TypedDict, pydantic, dataclass, etc. should all be able to load from dict
return outputs if input_type == "list" else {self.messages_key: outputs}
# LangGraph will automatically handle list of Command and non-command node updates
combined_outputs: list[
Command | list[ToolMessage] | dict[str, list[ToolMessage]]
] = []
for output in outputs:
if isinstance(output, Command):
combined_outputs.append(output)
else:
combined_outputs.append(
[output] if input_type == "list" else {self.messages_key: [output]}
)
return combined_outputs
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 {self.messages_key: outputs}
def invoke(
self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
@@ -263,97 +249,67 @@ class ToolNode(RunnableCallable):
*,
store: BaseStore,
) -> Any:
tool_calls, input_type = self._parse_input(input, store)
tool_calls, output_type = self._parse_input(input, store)
outputs = await asyncio.gather(
*(self._arun_one(call, input_type, config) for call in tool_calls)
*(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 {self.messages_key: outputs}
def _run_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 = self.tools_by_name[call["name"]].invoke(
input, config
)
tool_message.content = cast(
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 GraphBubbleUp as e:
raise e
except Exception as e:
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
# 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"
)
# preserve existing behavior for non-command tool outputs for backwards compatibility
if not any(isinstance(output, Command) for output in outputs):
# TypedDict, pydantic, dataclass, etc. should all be able to load from dict
return outputs if input_type == "list" else {self.messages_key: outputs}
# LangGraph will automatically handle list of Command and non-command node updates
combined_outputs: list[
Command | list[ToolMessage] | dict[str, list[ToolMessage]]
] = []
for output in outputs:
if isinstance(output, Command):
combined_outputs.append(output)
else:
combined_outputs.append(
[output] if input_type == "list" else {self.messages_key: [output]}
)
return combined_outputs
def _run_one(
self,
call: ToolCall,
input_type: Literal["list", "dict"],
config: RunnableConfig,
) -> ToolMessage:
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"}}
response = self.tools_by_name[call["name"]].invoke(input)
# 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 GraphBubbleUp as e:
raise e
except Exception as e:
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
# 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",
tool_message: ToolMessage = await self.tools_by_name[call["name"]].ainvoke(
input, config
)
if isinstance(response, Command):
return self._validate_tool_command(response, call, input_type)
elif isinstance(response, ToolMessage):
response.content = cast(
Union[str, list], msg_content_output(response.content)
tool_message.content = cast(
Union[str, list], msg_content_output(tool_message.content)
)
return response
else:
raise TypeError(
f"Tool {call['name']} returned unexpected type: {type(response)}"
)
async def _arun_one(
self,
call: ToolCall,
input_type: Literal["list", "dict"],
config: RunnableConfig,
) -> ToolMessage:
if invalid_tool_message := self._validate_tool_call(call):
return invalid_tool_message
try:
input = {**call, **{"type": "tool_call"}}
response = await self.tools_by_name[call["name"]].ainvoke(input)
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
@@ -378,24 +334,9 @@ class ToolNode(RunnableCallable):
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",
)
if isinstance(response, Command):
return self._validate_tool_command(response, call, input_type)
elif isinstance(response, ToolMessage):
response.content = cast(
Union[str, list], msg_content_output(response.content)
)
return response
else:
raise TypeError(
f"Tool {call['name']} returned unexpected type: {type(response)}"
)
return ToolMessage(
content=content, name=call["name"], tool_call_id=call["id"], status="error"
)
def _parse_input(
self,
@@ -405,16 +346,16 @@ class ToolNode(RunnableCallable):
BaseModel,
],
store: BaseStore,
) -> Tuple[list[ToolCall], Literal["list", "dict"]]:
) -> Tuple[List[ToolCall], Literal["list", "dict"]]:
if isinstance(input, list):
input_type = "list"
output_type = "list"
message: AnyMessage = input[-1]
elif isinstance(input, dict) and (messages := input.get(self.messages_key, [])):
input_type = "dict"
output_type = "dict"
message = messages[-1]
elif messages := getattr(input, self.messages_key, None):
# Assume dataclass-like state that can coerce from dict
input_type = "dict"
output_type = "dict"
message = messages[-1]
else:
raise ValueError("No message found in input")
@@ -425,7 +366,7 @@ class ToolNode(RunnableCallable):
tool_calls = [
self._inject_tool_args(call, input, store) for call in message.tool_calls
]
return tool_calls, input_type
return tool_calls, output_type
def _validate_tool_call(self, call: ToolCall) -> Optional[ToolMessage]:
if (requested_tool := call["name"]) not in self.tools_by_name:
@@ -519,67 +460,6 @@ class ToolNode(RunnableCallable):
tool_call_with_store = self._inject_store(tool_call_with_state, store)
return tool_call_with_store
def _validate_tool_command(
self, command: Command, call: ToolCall, input_type: Literal["list", "dict"]
) -> Command:
if isinstance(command.update, dict):
# input type is dict when ToolNode is invoked with a dict input (e.g. {"messages": [AIMessage(..., tool_calls=[...])]})
if input_type != "dict":
raise ValueError(
f"Tools can provide a dict in Command.update only when using dict with '{self.messages_key}' key as ToolNode input, "
f"got: {command.update} for tool '{call['name']}'"
)
updated_command = deepcopy(command)
state_update = cast(dict[str, Any], updated_command.update) or {}
messages_update = state_update.get(self.messages_key, [])
elif isinstance(command.update, list):
# input type is list when ToolNode is invoked with a list input (e.g. [AIMessage(..., tool_calls=[...])])
if input_type != "list":
raise ValueError(
f"Tools can provide a list of messages in Command.update only when using list of messages as ToolNode input, "
f"got: {command.update} for tool '{call['name']}'"
)
updated_command = deepcopy(command)
messages_update = updated_command.update
else:
return command
# convert to message objects if updates are in a dict format
messages_update = convert_to_messages(messages_update)
have_seen_tool_messages = False
for message in messages_update:
if not isinstance(message, ToolMessage):
continue
if have_seen_tool_messages:
raise ValueError(
f"Expected at most one ToolMessage in Command.update for tool '{call['name']}', got multiple: {messages_update}."
)
if message.tool_call_id != call["id"]:
raise ValueError(
f"ToolMessage.tool_call_id must match the tool call id. Expected: {call['id']}, got: {message.tool_call_id} for tool '{call['name']}'."
)
message.name = call["name"]
have_seen_tool_messages = True
# validate that we always have exactly one ToolMessage in Command.update if command is sent to the CURRENT graph
if updated_command.graph is None and not have_seen_tool_messages:
example_update = (
'`Command(update={"messages": [ToolMessage("Success", tool_call_id=tool_call_id), ...]}, ...)`'
if input_type == "dict"
else '`Command(update=[ToolMessage("Success", tool_call_id=tool_call_id), ...], ...)`'
)
raise ValueError(
f"Expected exactly one message (ToolMessage) in Command.update for tool '{call['name']}', got: {messages_update}. "
"Every tool call (LLM requesting to call a tool) in the message history MUST have a corresponding ToolMessage. "
f"You can fix it by modifying the tool to return {example_update}."
)
return updated_command
def tools_condition(
state: Union[list[AnyMessage], dict[str, Any], BaseModel],
@@ -776,9 +656,9 @@ def _is_injection(
return False
def _get_state_args(tool: BaseTool) -> dict[str, Optional[str]]:
def _get_state_args(tool: BaseTool) -> Dict[str, Optional[str]]:
full_schema = tool.get_input_schema()
tool_args_to_state_fields: dict = {}
tool_args_to_state_fields: Dict = {}
for name, type_ in get_all_basemodel_annotations(full_schema).items():
injections = [
+9 -18
View File
@@ -18,6 +18,7 @@ from typing import (
Type,
Union,
cast,
get_type_hints,
overload,
)
from uuid import UUID, uuid5
@@ -116,7 +117,6 @@ from langgraph.utils.config import (
patch_config,
patch_configurable,
)
from langgraph.utils.fields import get_enhanced_type_hints
from langgraph.utils.pydantic import create_model
from langgraph.utils.queue import AsyncQueue, SyncQueue # type: ignore[attr-defined]
@@ -319,15 +319,8 @@ class Pregel(PregelProtocol):
)
+ (
[
ConfigurableFieldSpec(
id=name,
annotation=typ,
default=default,
description=description,
)
for name, typ, default, description in get_enhanced_type_hints(
self.config_type
)
ConfigurableFieldSpec(id=name, annotation=typ)
for name, typ in get_type_hints(self.config_type).items()
]
if self.config_type is not None
else []
@@ -680,7 +673,7 @@ class Pregel(PregelProtocol):
self, config: RunnableConfig, *, subgraphs: bool = False
) -> StateSnapshot:
"""Get the current state of the graph."""
checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
@@ -717,7 +710,7 @@ class Pregel(PregelProtocol):
self, config: RunnableConfig, *, subgraphs: bool = False
) -> StateSnapshot:
"""Get the current state of the graph."""
checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
@@ -758,9 +751,8 @@ class Pregel(PregelProtocol):
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[StateSnapshot]:
config = ensure_config(config)
"""Get the history of the state of the graph."""
checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
@@ -808,9 +800,8 @@ class Pregel(PregelProtocol):
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[StateSnapshot]:
config = ensure_config(config)
"""Get the history of the state of the graph."""
checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
@@ -864,7 +855,7 @@ class Pregel(PregelProtocol):
node `as_node`. If `as_node` is not provided, it will be set to the last node
that updated the state, if not ambiguous.
"""
checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
@@ -1139,7 +1130,7 @@ class Pregel(PregelProtocol):
values: dict[str, Any] | Any,
as_node: Optional[str] = None,
) -> RunnableConfig:
checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
+39 -150
View File
@@ -37,14 +37,13 @@ from langgraph.constants import (
CONFIG_KEY_CHECKPOINT_NS,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_READ,
CONFIG_KEY_SCRATCHPAD,
CONFIG_KEY_RESUME_VALUE,
CONFIG_KEY_SEND,
CONFIG_KEY_STORE,
CONFIG_KEY_TASK_ID,
CONFIG_KEY_WRITES,
EMPTY_SEQ,
ERROR,
INTERRUPT,
MISSING,
NO_WRITES,
NS_END,
NS_SEP,
@@ -53,26 +52,18 @@ from langgraph.constants import (
PUSH,
RESERVED,
RESUME,
RETURN,
TAG_HIDDEN,
TASKS,
Send,
)
from langgraph.errors import EmptyChannelError, InvalidUpdateError
from langgraph.managed.base import ManagedValueMapping
from langgraph.pregel.call import get_runnable_for_func
from langgraph.pregel.io import read_channel, read_channels
from langgraph.pregel.log import logger
from langgraph.pregel.manager import ChannelsManager
from langgraph.pregel.read import PregelNode
from langgraph.store.base import BaseStore
from langgraph.types import (
All,
LoopProtocol,
PregelExecutableTask,
PregelTask,
RetryPolicy,
)
from langgraph.types import All, LoopProtocol, PregelExecutableTask, PregelTask
from langgraph.utils.config import merge_configs, patch_config
GetNextVersion = Callable[[Optional[V], BaseChannel], V]
@@ -106,21 +97,6 @@ class PregelTaskWrites(NamedTuple):
triggers: Sequence[str]
class Call:
__slots__ = ("func", "input", "retry")
func: Callable
input: Any
retry: Optional[RetryPolicy]
def __init__(
self, func: Callable, input: Any, *, retry: Optional[RetryPolicy]
) -> None:
self.func = func
self.input = input
self.retry = retry
def should_interrupt(
checkpoint: Checkpoint,
interrupt_nodes: Union[All, Sequence[str]],
@@ -203,7 +179,7 @@ def local_write(
"""Function injected under CONFIG_KEY_SEND in task config, to write to channels.
Validates writes and forwards them to `commit` function."""
for chan, value in writes:
if chan in (PUSH, TASKS) and value is not None:
if chan in (PUSH, TASKS):
if not isinstance(value, Send):
raise InvalidUpdateError(f"Expected Send, got {value}")
if value.node not in process_keys:
@@ -271,7 +247,7 @@ def apply_writes(
pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list)
for task in tasks:
for chan, val in task.writes:
if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR):
if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT):
pass
elif chan == TASKS: # TODO: remove branch in 1.0
checkpoint["pending_sends"].append(val)
@@ -462,7 +438,7 @@ def prepare_next_tasks(
def prepare_single_task(
task_path: tuple[Any, ...],
task_path: tuple[Union[str, int, tuple], ...],
task_id_checksum: Optional[str],
*,
checkpoint: Checkpoint,
@@ -483,94 +459,7 @@ def prepare_single_task(
configurable = config.get(CONF, {})
parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "")
if task_path[0] == PUSH and isinstance(task_path[-1], Call):
# (PUSH, parent task path, idx of PUSH write, id of parent task, Call)
task_path_t = cast(tuple[str, tuple, int, str, Call], task_path)
call = task_path_t[-1]
proc_ = get_runnable_for_func(call.func)
name = proc_.name
if name is None:
raise ValueError("`call` functions must have a `__name__` attribute")
# create task id
triggers = [PUSH]
checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name
task_id = _uuid5_str(
checkpoint_id,
checkpoint_ns,
str(step),
name,
PUSH,
_tuple_str(task_path[1]),
str(task_path[2]),
)
task_checkpoint_ns = f"{checkpoint_ns}:{task_id}"
metadata = {
"langgraph_step": step,
"langgraph_node": name,
"langgraph_triggers": triggers,
"langgraph_path": task_path[:3],
"langgraph_checkpoint_ns": task_checkpoint_ns,
}
if task_id_checksum is not None:
assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}"
if for_execution:
writes: deque[tuple[str, Any]] = deque()
return PregelExecutableTask(
name,
call.input,
proc_,
writes,
patch_config(
merge_configs(config, {"metadata": metadata}),
run_name=name,
callbacks=(
manager.get_child(f"graph:step:{step}") if manager else None
),
configurable={
CONFIG_KEY_TASK_ID: task_id,
# deque.extend is thread-safe
CONFIG_KEY_SEND: partial(
local_write,
writes.extend,
processes.keys(),
),
CONFIG_KEY_READ: partial(
local_read,
step,
checkpoint,
channels,
managed,
PregelTaskWrites(task_path[:3], name, writes, triggers),
config,
),
CONFIG_KEY_STORE: (store or configurable.get(CONFIG_KEY_STORE)),
CONFIG_KEY_CHECKPOINTER: (
checkpointer or configurable.get(CONFIG_KEY_CHECKPOINTER)
),
CONFIG_KEY_CHECKPOINT_MAP: {
**configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}),
parent_ns: checkpoint["id"],
},
CONFIG_KEY_CHECKPOINT_ID: None,
CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
CONFIG_KEY_WRITES: [
w
for w in pending_writes
+ configurable.get(CONFIG_KEY_WRITES, [])
if w[0] in (NULL_TASK_ID, task_id)
],
CONFIG_KEY_SCRATCHPAD: {},
},
),
triggers,
call.retry,
None,
task_id,
task_path[:3],
)
else:
return PregelTask(task_id, name, task_path[:3])
elif task_path[0] == PUSH:
if task_path[0] == PUSH:
if len(task_path) == 2: # TODO: remove branch in 1.0
# legacy SEND tasks, executed in superstep n+1
# (PUSH, idx of pending send)
@@ -601,19 +490,17 @@ def prepare_single_task(
PUSH,
str(idx),
)
elif len(task_path) >= 4:
elif len(task_path) == 4:
# new PUSH tasks, executed in superstep n
# (PUSH, parent task path, idx of PUSH write, id of parent task)
task_path_tt = cast(tuple[str, tuple, int, str], task_path)
writes_for_path = [w for w in pending_writes if w[0] == task_path_tt[3]]
if task_path_tt[2] >= len(writes_for_path):
task_path_t = cast(tuple[str, tuple, int, str], task_path)
writes_for_path = [w for w in pending_writes if w[0] == task_path_t[3]]
if task_path_t[2] >= len(writes_for_path):
logger.warning(
f"Ignoring invalid write index {task_path[2]} in pending writes"
)
return
packet = writes_for_path[task_path_tt[2]][2]
if packet is None:
return
packet = writes_for_path[task_path_t[2]][2]
if not isinstance(packet, Send):
logger.warning(
f"Ignoring invalid packet type {type(packet)} in pending writes"
@@ -646,7 +533,7 @@ def prepare_single_task(
"langgraph_step": step,
"langgraph_node": packet.node,
"langgraph_triggers": triggers,
"langgraph_path": task_path[:3],
"langgraph_path": task_path,
"langgraph_checkpoint_ns": task_checkpoint_ns,
}
if task_id_checksum is not None:
@@ -656,7 +543,7 @@ def prepare_single_task(
if node := proc.node:
if proc.metadata:
metadata.update(proc.metadata)
writes = deque()
writes: deque[tuple[str, Any]] = deque()
return PregelExecutableTask(
packet.node,
packet.arg,
@@ -685,7 +572,7 @@ def prepare_single_task(
channels,
managed,
PregelTaskWrites(
task_path[:3], packet.node, writes, triggers
task_path, packet.node, writes, triggers
),
config,
),
@@ -702,24 +589,26 @@ def prepare_single_task(
},
CONFIG_KEY_CHECKPOINT_ID: None,
CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
CONFIG_KEY_WRITES: [
w
for w in pending_writes
+ configurable.get(CONFIG_KEY_WRITES, [])
if w[0] in (NULL_TASK_ID, task_id)
],
CONFIG_KEY_SCRATCHPAD: {},
CONFIG_KEY_RESUME_VALUE: next(
(
v
for tid, c, v in pending_writes
if tid in (NULL_TASK_ID, task_id) and c == RESUME
),
MISSING,
),
},
),
triggers,
proc.retry_policy,
None,
task_id,
task_path[:3],
task_path,
writers=proc.flat_writers,
)
else:
return PregelTask(task_id, packet.node, task_path[:3])
return PregelTask(task_id, packet.node, task_path)
elif task_path[0] == PULL:
# (PULL, node name)
name = cast(str, task_path[1])
@@ -769,7 +658,7 @@ def prepare_single_task(
"langgraph_step": step,
"langgraph_node": name,
"langgraph_triggers": triggers,
"langgraph_path": task_path[:3],
"langgraph_path": task_path,
"langgraph_checkpoint_ns": task_checkpoint_ns,
}
if task_id_checksum is not None:
@@ -808,9 +697,7 @@ def prepare_single_task(
checkpoint,
channels,
managed,
PregelTaskWrites(
task_path[:3], name, writes, triggers
),
PregelTaskWrites(task_path, name, writes, triggers),
config,
),
CONFIG_KEY_STORE: (
@@ -826,24 +713,26 @@ def prepare_single_task(
},
CONFIG_KEY_CHECKPOINT_ID: None,
CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
CONFIG_KEY_WRITES: [
w
for w in pending_writes
+ configurable.get(CONFIG_KEY_WRITES, [])
if w[0] in (NULL_TASK_ID, task_id)
],
CONFIG_KEY_SCRATCHPAD: {},
CONFIG_KEY_RESUME_VALUE: next(
(
v
for tid, c, v in pending_writes
if tid in (NULL_TASK_ID, task_id)
and c == RESUME
),
MISSING,
),
},
),
triggers,
proc.retry_policy,
None,
task_id,
task_path[:3],
task_path,
writers=proc.flat_writers,
)
else:
return PregelTask(task_id, name, task_path[:3])
return PregelTask(task_id, name, task_path)
def _proc_input(
-123
View File
@@ -1,123 +0,0 @@
import sys
import types
from typing import Any, Callable, Optional
from langgraph.constants import RETURN
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
from langgraph.utils.runnable import RunnableSeq, coerce_to_runnable
"""
Utilities borrowed from cloudpickle.
https://github.com/cloudpipe/cloudpickle/blob/6220b0ce83ffee5e47e06770a1ee38ca9e47c850/cloudpickle/cloudpickle.py#L265
"""
def _getattribute(obj: Any, name: str) -> Any:
for subpath in name.split("."):
if subpath == "<locals>":
raise AttributeError(
"Can't get local attribute {!r} on {!r}".format(name, obj)
)
try:
parent = obj
obj = getattr(obj, subpath)
except AttributeError:
raise AttributeError(
"Can't get attribute {!r} on {!r}".format(name, obj)
) from None
return obj, parent
def _whichmodule(obj: Any, name: str) -> Optional[str]:
"""Find the module an object belongs to.
This function differs from ``pickle.whichmodule`` in two ways:
- it does not mangle the cases where obj's module is __main__ and obj was
not found in any module.
- Errors arising during module introspection are ignored, as those errors
are considered unwanted side effects.
"""
module_name = getattr(obj, "__module__", None)
if module_name is not None:
return module_name
# Protect the iteration by using a copy of sys.modules against dynamic
# modules that trigger imports of other modules upon calls to getattr or
# other threads importing at the same time.
for module_name, module in sys.modules.copy().items():
# Some modules such as coverage can inject non-module objects inside
# sys.modules
if (
module_name == "__main__"
or module_name == "__mp_main__"
or module is None
or not isinstance(module, types.ModuleType)
):
continue
try:
if _getattribute(module, name)[0] is obj:
return module_name
except Exception:
pass
return None
def _lookup_module_and_qualname(
obj: Any, name: Optional[str] = None
) -> Optional[tuple[types.ModuleType, str]]:
if name is None:
name = getattr(obj, "__qualname__", None)
if name is None: # pragma: no cover
# This used to be needed for Python 2.7 support but is probably not
# needed anymore. However we keep the __name__ introspection in case
# users of cloudpickle rely on this old behavior for unknown reasons.
name = getattr(obj, "__name__", None)
if name is None:
return None
module_name = _whichmodule(obj, name)
if module_name is None:
# In this case, obj.__module__ is None AND obj was not found in any
# imported module. obj is thus treated as dynamic.
return None
if module_name == "__main__":
return None
# Note: if module_name is in sys.modules, the corresponding module is
# assumed importable at unpickling time. See #357
module = sys.modules.get(module_name, None)
if module is None:
# The main reason why obj's module would not be imported is that this
# module has been dynamically created, using for example
# types.ModuleType. The other possibility is that module was removed
# from sys.modules after obj was created/imported. But this case is not
# supported, as the standard pickle does not support it either.
return None
try:
obj2, parent = _getattribute(module, name)
except AttributeError:
# obj was not found inside the module it points to
return None
if obj2 is not obj:
return None
return module, name
def get_runnable_for_func(func: Callable[..., Any]) -> RunnableSeq:
if func in CACHE:
return CACHE[func]
else:
seq = RunnableSeq(
coerce_to_runnable(func, name=None, trace=False),
ChannelWrite([ChannelWriteEntry(RETURN)]),
name=func.__name__,
)
if not _lookup_module_and_qualname(func):
return seq
return CACHE.setdefault(func, seq)
CACHE: dict[Callable[..., Any], RunnableSeq] = {}
+1 -22
View File
@@ -1,7 +1,6 @@
import asyncio
import concurrent.futures
import sys
import time
from contextlib import ExitStack
from contextvars import copy_context
from types import TracebackType
@@ -35,7 +34,6 @@ class Submit(Protocol[P, T]):
__name__: Optional[str] = None,
__cancel_on_exit__: bool = False,
__reraise_on_exit__: bool = True,
__next_tick__: bool = False,
**kwargs: P.kwargs,
) -> concurrent.futures.Future[T]: ...
@@ -60,13 +58,9 @@ class BackgroundExecutor(ContextManager):
__name__: Optional[str] = None, # currently not used in sync version
__cancel_on_exit__: bool = False, # for sync, can cancel only if not started
__reraise_on_exit__: bool = True,
__next_tick__: bool = False,
**kwargs: P.kwargs,
) -> concurrent.futures.Future[T]:
if __next_tick__:
task = self.executor.submit(next_tick, fn, *args, **kwargs)
else:
task = self.executor.submit(fn, *args, **kwargs)
task = self.executor.submit(fn, *args, **kwargs)
self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__)
task.add_done_callback(self.done)
return task
@@ -143,14 +137,11 @@ class AsyncBackgroundExecutor(AsyncContextManager):
__name__: Optional[str] = None,
__cancel_on_exit__: bool = False,
__reraise_on_exit__: bool = True,
__next_tick__: bool = False,
**kwargs: P.kwargs,
) -> asyncio.Task[T]:
coro = cast(Coroutine[None, None, T], fn(*args, **kwargs))
if self.semaphore:
coro = gated(self.semaphore, coro)
if __next_tick__:
coro = anext_tick(coro)
if self.context_not_supported:
task = self.loop.create_task(coro, name=__name__)
else:
@@ -206,15 +197,3 @@ async def gated(semaphore: asyncio.Semaphore, coro: Coroutine[None, None, T]) ->
"""A coroutine that waits for a semaphore before running another coroutine."""
async with semaphore:
return await coro
def next_tick(fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
"""A function that yields control to other threads before running another function."""
time.sleep(0)
return fn(*args, **kwargs)
async def anext_tick(coro: Coroutine[None, None, T]) -> T:
"""A coroutine that yields control to event loop before running another coroutine."""
await asyncio.sleep(0)
return await coro
+29 -35
View File
@@ -4,7 +4,6 @@ from uuid import UUID
from langchain_core.runnables.utils import AddableDict
from langgraph.channels.base import BaseChannel, EmptyChannelError
from langgraph.checkpoint.base import PendingWrite
from langgraph.constants import (
EMPTY_SEQ,
ERROR,
@@ -13,9 +12,6 @@ from langgraph.constants import (
NULL_TASK_ID,
PUSH,
RESUME,
RETURN,
SELF,
START,
TAG_HIDDEN,
TASKS,
)
@@ -70,37 +66,34 @@ def read_channels(
def map_command(
cmd: Command, pending_writes: list[PendingWrite]
cmd: Command,
) -> Iterator[tuple[str, str, Any]]:
"""Map input chunk to a sequence of pending writes in the form (channel, value)."""
if cmd.graph == Command.PARENT:
raise InvalidUpdateError("There is not parent graph")
if cmd.goto:
if isinstance(cmd.goto, (tuple, list)):
sends = cmd.goto
if cmd.send:
if isinstance(cmd.send, (tuple, list)):
sends = cmd.send
else:
sends = [cmd.goto]
sends = [cmd.send]
for send in sends:
if isinstance(send, Send):
yield (NULL_TASK_ID, PUSH if FF_SEND_V2 else TASKS, send)
elif isinstance(send, str):
yield (NULL_TASK_ID, f"branch:{START}:{SELF}:{send}", START)
else:
if not isinstance(send, Send):
raise TypeError(
f"In Command.goto, expected Send/str, got {type(send).__name__}"
f"In Command.send, expected Send, got {type(send).__name__}"
)
yield (NULL_TASK_ID, PUSH if FF_SEND_V2 else TASKS, send)
if cmd.resume:
if isinstance(cmd.resume, dict) and all(is_task_id(k) for k in cmd.resume):
for tid, resume in cmd.resume.items():
existing: list[Any] = next(
(w[2] for w in pending_writes if w[0] == tid and w[1] == RESUME), []
)
existing.append(resume)
yield (tid, RESUME, existing)
yield (tid, RESUME, resume)
else:
yield (NULL_TASK_ID, RESUME, cmd.resume)
if cmd.update:
for k, v in cmd._update_as_tuples():
if not isinstance(cmd.update, dict):
raise TypeError(
f"Expected cmd.update to be a dict mapping channel names to update values, got {type(cmd.update).__name__}"
)
for k, v in cmd.update.items():
yield (NULL_TASK_ID, k, v)
@@ -172,21 +165,22 @@ def map_output_updates(
]
if not output_tasks:
return
updated: list[tuple[str, Any]] = []
for task, writes in output_tasks:
if rtn := next((value for chan, value in writes if chan == RETURN), None):
updated.append((task.name, rtn))
elif isinstance(output_channels, str):
updated.extend(
(task.name, value) for chan, value in writes if chan == output_channels
)
elif any(chan in output_channels for chan, _ in writes):
updated.append(
(
task.name,
{chan: value for chan, value in writes if chan in output_channels},
)
if isinstance(output_channels, str):
updated = (
(task.name, value)
for task, writes in output_tasks
for chan, value in writes
if chan == output_channels
)
else:
updated = (
(
task.name,
{chan: value for chan, value in writes if chan in output_channels},
)
for task, writes in output_tasks
if any(chan in output_channels for chan, _ in writes)
)
grouped: dict[str, list[Any]] = {t.name: [] for t, _ in output_tasks}
for node, value in updated:
grouped[node].append(value)
+42 -66
View File
@@ -26,7 +26,6 @@ from typing_extensions import ParamSpec, Self
from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
Checkpoint,
@@ -73,7 +72,6 @@ from langgraph.managed.base import (
WritableManagedValue,
)
from langgraph.pregel.algo import (
Call,
GetNextVersion,
PregelTaskWrites,
apply_writes,
@@ -265,40 +263,21 @@ class PregelLoop(LoopProtocol):
"""Put writes for a task, to be read by the next tick."""
if not writes:
return
# deduplicate writes to special channels, last write wins
if all(w[0] in WRITES_IDX_MAP for w in writes):
writes = list({w[0]: w for w in writes}.values())
# save writes
for c, v in writes:
if (
c in WRITES_IDX_MAP
and (
idx := next(
(
i
for i, w in enumerate(self.checkpoint_pending_writes)
if w[0] == task_id and w[1] == c
),
None,
)
)
is not None
):
self.checkpoint_pending_writes[idx] = (task_id, c, v)
else:
self.checkpoint_pending_writes.append((task_id, c, v))
self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes)
if self.checkpointer_put_writes is not None:
self.submit(
self.checkpointer_put_writes,
patch_configurable(
self.checkpoint_config,
{
{
**self.checkpoint_config,
CONF: {
**self.checkpoint_config[CONF],
CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get(
CONFIG_KEY_CHECKPOINT_NS, ""
),
CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
},
),
},
writes,
task_id,
)
@@ -307,19 +286,20 @@ class PregelLoop(LoopProtocol):
self._output_writes(task_id, writes)
def accept_push(
self, task: PregelExecutableTask, write_idx: int, call: Optional[Call] = None
self, task: PregelExecutableTask, write_idx: int
) -> Optional[PregelExecutableTask]:
"""Accept a PUSH from a task, potentially returning a new task to start."""
# don't start if an earlier PUSH has already triggered an interrupt
if self.to_interrupt:
return
# don't start if we should interrupt *after* the original task
if self.interrupt_after and should_interrupt(
self.checkpoint, self.interrupt_after, [task]
):
if should_interrupt(self.checkpoint, self.interrupt_after, [task]):
self.to_interrupt.append(task)
return
if pushed := cast(
Optional[PregelExecutableTask],
prepare_single_task(
(PUSH, task.path, write_idx, task.id, call),
(PUSH, task.path, write_idx, task.id),
None,
checkpoint=self.checkpoint,
pending_writes=[(task.id, *w) for w in task.writes],
@@ -335,9 +315,7 @@ class PregelLoop(LoopProtocol):
),
):
# don't start if we should interrupt *before* the new task
if self.interrupt_before and should_interrupt(
self.checkpoint, self.interrupt_before, [pushed]
):
if should_interrupt(self.checkpoint, self.interrupt_before, [pushed]):
self.to_interrupt.append(pushed)
return
# produce debug output
@@ -350,8 +328,9 @@ class PregelLoop(LoopProtocol):
# match any pending writes to the new task
if self.skip_done_tasks:
self._match_writes({pushed.id: pushed})
# return the new task, to be started if not run before
return pushed
# return the new task, to be started, if not run before
if not pushed.writes:
return pushed
def tick(
self,
@@ -413,7 +392,7 @@ class PregelLoop(LoopProtocol):
}
)
# after execution, check if we should interrupt
if self.interrupt_after and should_interrupt(
if should_interrupt(
self.checkpoint, self.interrupt_after, self.tasks.values()
):
self.status = "interrupt_after"
@@ -426,6 +405,18 @@ class PregelLoop(LoopProtocol):
self.status = "out_of_steps"
return False
# apply NULL writes
if null_writes := [
w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID
]:
mv_writes = apply_writes(
self.checkpoint,
self.channels,
[PregelTaskWrites((), INPUT, null_writes, [])],
self.checkpointer_get_next_version,
)
for key, values in mv_writes.items():
self._update_mv(key, values)
# prepare next tasks
self.tasks = prepare_next_tasks(
self.checkpoint,
@@ -485,7 +476,7 @@ class PregelLoop(LoopProtocol):
return self.tick(input_keys=input_keys)
# before execution, check if we should interrupt
if self.interrupt_before and should_interrupt(
if should_interrupt(
self.checkpoint, self.interrupt_before, self.tasks.values()
):
self.status = "interrupt_before"
@@ -527,35 +518,9 @@ class PregelLoop(LoopProtocol):
# - receiving None input (outer graph) or RESUMING flag (subgraph)
configurable = self.config.get(CONF, {})
is_resuming = bool(self.checkpoint["channel_versions"]) and bool(
configurable.get(
CONFIG_KEY_RESUMING,
self.input is None or isinstance(self.input, Command),
)
configurable.get(CONFIG_KEY_RESUMING, self.input is None)
)
# map command to writes
if isinstance(self.input, Command):
writes: defaultdict[str, list[tuple[str, Any]]] = defaultdict(list)
# group writes by task ID
for tid, c, v in map_command(self.input, self.checkpoint_pending_writes):
writes[tid].append((c, v))
if not writes:
raise EmptyInputError("Received empty Command input")
# save writes
for tid, ws in writes.items():
self.put_writes(tid, ws)
# apply NULL writes
if null_writes := [
w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID
]:
mv_writes = apply_writes(
self.checkpoint,
self.channels,
[PregelTaskWrites((), INPUT, null_writes, [])],
self.checkpointer_get_next_version,
)
for key, values in mv_writes.items():
self._update_mv(key, values)
# proceed past previous checkpoint
if is_resuming:
self.checkpoint["versions_seen"].setdefault(INTERRUPT, {})
@@ -567,6 +532,17 @@ class PregelLoop(LoopProtocol):
self._emit(
"values", map_output_values, self.output_keys, True, self.channels
)
# map command to writes
elif isinstance(self.input, Command):
writes: defaultdict[str, list[tuple[str, Any]]] = defaultdict(list)
# group writes by task ID
for tid, c, v in map_command(self.input):
writes[tid].append((c, v))
if not writes:
raise EmptyInputError("Received empty Command input")
# save writes
for tid, ws in writes.items():
self.put_writes(tid, ws)
# map inputs to channel updates
elif input_writes := deque(map_input(input_keys, self.input)):
# TODO shouldn't these writes be passed to put_writes too?
+9 -27
View File
@@ -1,4 +1,3 @@
from dataclasses import asdict
from typing import (
Any,
AsyncIterator,
@@ -28,7 +27,6 @@ from langgraph_sdk.client import (
get_sync_client,
)
from langgraph_sdk.schema import Checkpoint, ThreadState
from langgraph_sdk.schema import Command as CommandSDK
from langgraph_sdk.schema import StreamMode as StreamModeSDK
from typing_extensions import Self
@@ -43,7 +41,7 @@ from langgraph.constants import (
from langgraph.errors import GraphInterrupt
from langgraph.pregel.protocol import PregelProtocol
from langgraph.pregel.types import All, PregelTask, StateSnapshot, StreamMode
from langgraph.types import Command, Interrupt, StreamProtocol
from langgraph.types import Interrupt, StreamProtocol
from langgraph.utils.config import merge_configs
@@ -575,7 +573,6 @@ class RemoteGraph(PregelProtocol):
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
subgraphs: bool = False,
**kwargs: Any,
) -> Iterator[Union[dict[str, Any], Any]]:
"""Create a run and stream the results.
@@ -590,7 +587,6 @@ class RemoteGraph(PregelProtocol):
interrupt_before: Interrupt the graph before these nodes.
interrupt_after: Interrupt the graph after these nodes.
subgraphs: Stream from subgraphs.
**kwargs: Additional params to pass to client.runs.stream.
Yields:
The output of the graph.
@@ -601,24 +597,17 @@ class RemoteGraph(PregelProtocol):
stream_modes, requested, req_single, stream = self._get_stream_modes(
stream_mode, config
)
if isinstance(input, Command):
command: Optional[CommandSDK] = cast(CommandSDK, asdict(input))
input = None
else:
command = None
for chunk in sync_client.runs.stream(
thread_id=sanitized_config["configurable"].get("thread_id"),
assistant_id=self.name,
input=input,
command=command,
config=sanitized_config,
stream_mode=stream_modes,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
stream_subgraphs=subgraphs or stream is not None,
if_not_exists="create",
**kwargs,
):
# split mode and ns
if NS_SEP in chunk.event:
@@ -667,7 +656,6 @@ class RemoteGraph(PregelProtocol):
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
subgraphs: bool = False,
**kwargs: Any,
) -> AsyncIterator[Union[dict[str, Any], Any]]:
"""Create a run and stream the results.
@@ -682,7 +670,6 @@ class RemoteGraph(PregelProtocol):
interrupt_before: Interrupt the graph before these nodes.
interrupt_after: Interrupt the graph after these nodes.
subgraphs: Stream from subgraphs.
**kwargs: Additional params to pass to client.runs.stream.
Yields:
The output of the graph.
@@ -693,24 +680,17 @@ class RemoteGraph(PregelProtocol):
stream_modes, requested, req_single, stream = self._get_stream_modes(
stream_mode, config
)
if isinstance(input, Command):
command: Optional[CommandSDK] = cast(CommandSDK, asdict(input))
input = None
else:
command = None
async for chunk in client.runs.stream(
thread_id=sanitized_config["configurable"].get("thread_id"),
assistant_id=self.name,
input=input,
command=command,
config=sanitized_config,
stream_mode=stream_modes,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
stream_subgraphs=subgraphs or stream is not None,
if_not_exists="create",
**kwargs,
):
# split mode and ns
if NS_SEP in chunk.event:
@@ -773,16 +753,18 @@ class RemoteGraph(PregelProtocol):
*,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
**kwargs: Any,
) -> Union[dict[str, Any], Any]:
"""Create a run, wait until it finishes and return the final state.
This method calls `POST /threads/{thread_id}/runs/wait` if a `thread_id`
is speciffed in the `configurable` field of the config or
`POST /runs/wait` otherwise.
Args:
input: Input to the graph.
config: A `RunnableConfig` for graph invocation.
interrupt_before: Interrupt the graph before these nodes.
interrupt_after: Interrupt the graph after these nodes.
**kwargs: Additional params to pass to RemoteGraph.stream.
Returns:
The output of the graph.
@@ -793,7 +775,6 @@ class RemoteGraph(PregelProtocol):
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
stream_mode="values",
**kwargs,
):
pass
try:
@@ -808,16 +789,18 @@ class RemoteGraph(PregelProtocol):
*,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
**kwargs: Any,
) -> Union[dict[str, Any], Any]:
"""Create a run, wait until it finishes and return the final state.
This method calls `POST /threads/{thread_id}/runs/wait` if a `thread_id`
is speciffed in the `configurable` field of the config or
`POST /runs/wait` otherwise.
Args:
input: Input to the graph.
config: A `RunnableConfig` for graph invocation.
interrupt_before: Interrupt the graph before these nodes.
interrupt_after: Interrupt the graph after these nodes.
**kwargs: Additional params to pass to RemoteGraph.astream.
Returns:
The output of the graph.
@@ -828,7 +811,6 @@ class RemoteGraph(PregelProtocol):
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
stream_mode="values",
**kwargs,
):
pass
try:
+19 -11
View File
@@ -4,12 +4,14 @@ import random
import sys
import time
from dataclasses import replace
from typing import Any, Optional, Sequence
from functools import partial
from typing import Any, Callable, Optional, Sequence
from langgraph.constants import (
CONF,
CONFIG_KEY_CHECKPOINT_NS,
CONFIG_KEY_RESUMING,
CONFIG_KEY_SEND,
NS_SEP,
)
from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphBubbleUp, ParentCommand
@@ -23,21 +25,25 @@ SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11)
def run_with_retry(
task: PregelExecutableTask,
retry_policy: Optional[RetryPolicy],
configurable: Optional[dict[str, Any]] = None,
writer: Optional[
Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None]
] = None,
) -> None:
"""Run a task with retries."""
retry_policy = task.retry_policy or retry_policy
interval = retry_policy.initial_interval if retry_policy else 0
attempts = 0
config = task.config
if configurable is not None:
config = patch_configurable(config, configurable)
if writer is not None:
config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)})
while True:
try:
# clear any writes from previous attempts
task.writes.clear()
# run the task
return task.proc.invoke(task.input, config)
task.proc.invoke(task.input, config)
# if successful, end
break
except ParentCommand as exc:
ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
cmd = exc.args[0]
@@ -109,15 +115,17 @@ async def arun_with_retry(
task: PregelExecutableTask,
retry_policy: Optional[RetryPolicy],
stream: bool = False,
configurable: Optional[dict[str, Any]] = None,
writer: Optional[
Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None]
] = None,
) -> None:
"""Run a task asynchronously with retries."""
retry_policy = task.retry_policy or retry_policy
interval = retry_policy.initial_interval if retry_policy else 0
attempts = 0
config = task.config
if configurable is not None:
config = patch_configurable(config, configurable)
if writer is not None:
config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)})
while True:
try:
# clear any writes from previous attempts
@@ -126,10 +134,10 @@ async def arun_with_retry(
if stream:
async for _ in task.proc.astream(task.input, config):
pass
# if successful, end
break
else:
return await task.proc.ainvoke(task.input, config)
await task.proc.ainvoke(task.input, config)
# if successful, end
break
except ParentCommand as exc:
ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
cmd = exc.args[0]
+91 -261
View File
@@ -1,12 +1,9 @@
import asyncio
import concurrent.futures
import threading
import time
from functools import partial
from typing import (
Any,
AsyncIterator,
Awaitable,
Callable,
Iterable,
Iterator,
@@ -19,22 +16,17 @@ from typing import (
from langgraph.constants import (
CONF,
CONFIG_KEY_CALL,
CONFIG_KEY_SEND,
ERROR,
INTERRUPT,
NO_WRITES,
PUSH,
RESUME,
RETURN,
TAG_HIDDEN,
)
from langgraph.errors import GraphBubbleUp, GraphInterrupt
from langgraph.pregel.algo import Call
from langgraph.pregel.executor import Submit
from langgraph.pregel.retry import arun_with_retry, run_with_retry
from langgraph.types import PregelExecutableTask, RetryPolicy
from langgraph.utils.future import chain_future
class PregelRunner:
@@ -48,7 +40,7 @@ class PregelRunner:
submit: Submit,
put_writes: Callable[[str, Sequence[tuple[str, Any]]], None],
schedule_task: Callable[
[PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask]
[PregelExecutableTask, int], Optional[PregelExecutableTask]
],
use_astream: bool = False,
node_finished: Optional[Callable[[str], None]] = None,
@@ -68,143 +60,73 @@ class PregelRunner:
retry_policy: Optional[RetryPolicy] = None,
get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None,
) -> Iterator[None]:
locks: dict[str, threading.Lock] = {}
def writer(
task: PregelExecutableTask,
writes: Sequence[tuple[str, Any]],
*,
calls: Optional[Sequence[Call]] = None,
) -> Sequence[Optional[concurrent.futures.Future]]:
if all(w[0] != PUSH for w in writes):
return task.config[CONF][CONFIG_KEY_SEND](writes)
if task.id not in locks:
locks[task.id] = threading.Lock()
with locks[task.id]:
prev_length = len(task.writes)
# delegate to the underlying writer
task.config[CONF][CONFIG_KEY_SEND](writes)
# confirm no other concurrent writes were added
assert len(task.writes) == prev_length + len(writes)
# schedule PUSH tasks, collect futures
rtn: dict[int, Optional[concurrent.futures.Future]] = {}
for idx, w in enumerate(writes, start=prev_length):
task: PregelExecutableTask, writes: Sequence[tuple[str, Any]]
) -> None:
prev_length = len(task.writes)
# delegate to the underlying writer
task.config[CONF][CONFIG_KEY_SEND](writes)
for idx, w in enumerate(task.writes):
# find the index for the newly inserted writes
if idx < prev_length:
continue
assert writes[idx - prev_length] is w
# bail if not a PUSH write
if w[0] != PUSH:
continue
# schedule the next task, if the callback returns one
if next_task := self.schedule_task(
task, idx, calls[idx - prev_length] if calls else None
):
if fut := next(
(
f
for f, t in futures.items()
if t is not None and t == next_task.id
),
None,
if next_task := self.schedule_task(task, idx):
# if the parent task was retried,
# the next task might already be running
if any(
t == next_task.id for t in futures.values() if t is not None
):
# if the parent task was retried,
# the next task might already be running
rtn[idx - prev_length] = fut
elif next_task.writes:
# if it already ran, return the result
fut = concurrent.futures.Future()
if val := next(v for c, v in next_task.writes if c == RETURN):
fut.set_result(val)
elif exc := next(v for c, v in next_task.writes if c == ERROR):
fut.set_exception(
exc
if isinstance(exc, BaseException)
else Exception(exc)
)
else:
fut.set_result(None)
rtn[idx - prev_length] = fut
else:
# schedule the next task
fut = self.submit(
continue
# schedule the next task
futures[
self.submit(
run_with_retry,
next_task,
retry_policy,
configurable={
CONFIG_KEY_SEND: partial(writer, next_task),
CONFIG_KEY_CALL: partial(call, next_task),
},
writer=writer,
__reraise_on_exit__=reraise,
# starting a new task in the next tick ensures
# updates from this tick are committed/streamed first
__next_tick__=True,
)
fut.add_done_callback(partial(self.commit, next_task))
futures[fut] = next_task
rtn[idx - prev_length] = fut
return [rtn.get(i) for i in range(len(writes))]
def call(
task: PregelExecutableTask,
func: Callable[[Any], Union[Awaitable[Any], Any]],
input: Any,
*,
retry: Optional[RetryPolicy] = None,
) -> concurrent.futures.Future[Any]:
(fut,) = writer(
task, [(PUSH, None)], calls=[Call(func, input, retry=retry)]
)
assert fut is not None, "writer did not return a future for call"
return fut
] = next_task
tasks = tuple(tasks)
futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = {}
done_futures: set[concurrent.futures.Future] = set()
# give control back to the caller
yield
# fast path if single task with no timeout and no waiter
if len(tasks) == 1 and timeout is None and get_waiter is None:
t = tasks[0]
try:
run_with_retry(
t,
retry_policy,
configurable={
CONFIG_KEY_SEND: partial(writer, t),
CONFIG_KEY_CALL: partial(call, t),
},
)
run_with_retry(t, retry_policy, writer=writer)
self.commit(t, None)
except Exception as exc:
self.commit(t, None, exc)
if reraise and futures:
# will be re-raised after futures are done
fut: concurrent.futures.Future = concurrent.futures.Future()
fut.set_exception(exc)
done_futures.add(fut)
elif reraise:
self.commit(t, exc)
if reraise:
raise
if not futures: # maybe `t` schuduled another task
return
# add waiter task if requested
if get_waiter is not None:
futures[get_waiter()] = None
# schedule tasks
for t in tasks:
if not t.writes:
fut = self.submit(
run_with_retry,
t,
retry_policy,
configurable={
CONFIG_KEY_SEND: partial(writer, t),
CONFIG_KEY_CALL: partial(call, t),
},
__reraise_on_exit__=reraise,
)
fut.add_done_callback(partial(self.commit, t))
futures[fut] = t
# execute tasks, and wait for one to fail or all to finish.
# each task is independent from all other concurrent tasks
# yield updates/debug output as each task finishes
for t in tasks:
if not t.writes:
futures[
self.submit(
run_with_retry,
t,
retry_policy,
writer=writer,
__reraise_on_exit__=reraise,
)
] = t
done_futures: set[concurrent.futures.Future] = set()
end_time = timeout + time.monotonic() if timeout else None
while len(futures) > (1 if get_waiter is not None else 0):
done, inflight = concurrent.futures.wait(
@@ -223,6 +145,8 @@ class PregelRunner:
else:
# store for panic check
done_futures.add(fut)
# task finished, commit writes
self.commit(task, _exception(fut))
else:
# remove references to loop vars
del fut, task
@@ -231,10 +155,6 @@ class PregelRunner:
break
# give control back to the caller
yield
# wait for pending done callbacks
# if a 2nd future finishes while `wait` is returning, it's possible
# that done callbacks for the 2nd future aren't called until next tick
time.sleep(0)
# panic on failure or timeout
_panic_or_proceed(
done_futures.union(f for f, t in futures.items() if t is not None),
@@ -250,109 +170,48 @@ class PregelRunner:
retry_policy: Optional[RetryPolicy] = None,
get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None,
) -> AsyncIterator[None]:
locks: dict[str, threading.Lock] = {}
def writer(
task: PregelExecutableTask,
writes: Sequence[tuple[str, Any]],
*,
calls: Optional[Sequence[Call]] = None,
) -> Sequence[Optional[asyncio.Future]]:
if all(w[0] != PUSH for w in writes):
return task.config[CONF][CONFIG_KEY_SEND](writes)
if task.id not in locks:
locks[task.id] = threading.Lock()
with locks[task.id]:
prev_length = len(task.writes)
# delegate to the underlying writer
task.config[CONF][CONFIG_KEY_SEND](writes)
# confirm no other concurrent writes were added
assert len(task.writes) == prev_length + len(writes)
# schedule PUSH tasks, collect futures
rtn: dict[int, Optional[asyncio.Future]] = {}
for idx, w in enumerate(writes, start=prev_length):
task: PregelExecutableTask, writes: Sequence[tuple[str, Any]]
) -> None:
prev_length = len(task.writes)
# delegate to the underlying writer
task.config[CONF][CONFIG_KEY_SEND](writes)
for idx, w in enumerate(task.writes):
# find the index for the newly inserted writes
if idx < prev_length:
continue
assert writes[idx - prev_length] is w
# bail if not a PUSH write
if w[0] != PUSH:
continue
# schedule the next task, if the callback returns one
wcall = calls[idx - prev_length] if calls is not None else None
if next_task := self.schedule_task(task, idx, wcall):
if next_task := self.schedule_task(task, idx):
# if the parent task was retried,
# the next task might already be running
if fut := next(
(
f
for f, t in futures.items()
if t is not None and t == next_task.id
),
None,
if any(
t == next_task.id for t in futures.values() if t is not None
):
# if the parent task was retried,
# the next task might already be running
rtn[idx - prev_length] = fut
elif next_task.writes:
# if it already ran, return the result
fut = asyncio.Future()
if val := next(v for c, v in next_task.writes if c == RETURN):
fut.set_result(val)
elif exc := next(v for c, v in next_task.writes if c == ERROR):
fut.set_exception(
exc
if isinstance(exc, BaseException)
else Exception(exc)
)
else:
fut.set_result(None)
rtn[idx - prev_length] = fut
else:
# schedule the next task
fut = cast(
continue
# schedule the next task
futures[
cast(
asyncio.Future,
self.submit(
arun_with_retry,
next_task,
retry_policy,
stream=self.use_astream,
configurable={
CONFIG_KEY_SEND: partial(writer, next_task),
CONFIG_KEY_CALL: partial(call, next_task),
},
writer=writer,
__name__=t.name,
__cancel_on_exit__=True,
__reraise_on_exit__=reraise,
# starting a new task in the next tick ensures
# updates from this tick are committed/streamed first
__next_tick__=True,
),
)
fut.add_done_callback(partial(self.commit, next_task))
futures[fut] = next_task
rtn[idx - prev_length] = fut
return [rtn.get(i) for i in range(len(writes))]
def call(
task: PregelExecutableTask,
func: Callable[[Any], Union[Awaitable[Any], Any]],
input: Any,
*,
retry: Optional[RetryPolicy] = None,
) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]:
(fut,) = writer(
task, [(PUSH, None)], calls=[Call(func, input, retry=retry)]
)
assert fut is not None, "writer did not return a future for call"
if asyncio.iscoroutinefunction(func):
return fut
# adapted from asyncio.run_coroutine_threadsafe
sfut: concurrent.futures.Future = concurrent.futures.Future()
loop.call_soon_threadsafe(chain_future, fut, sfut)
return sfut
] = next_task
loop = asyncio.get_event_loop()
tasks = tuple(tasks)
futures: dict[asyncio.Future, Optional[PregelExecutableTask]] = {}
done_futures: set[asyncio.Future] = set()
# give control back to the caller
yield
# fast path if single task with no waiter and no timeout
@@ -360,53 +219,39 @@ class PregelRunner:
t = tasks[0]
try:
await arun_with_retry(
t,
retry_policy,
stream=self.use_astream,
configurable={
CONFIG_KEY_SEND: partial(writer, t),
CONFIG_KEY_CALL: partial(call, t),
},
t, retry_policy, stream=self.use_astream, writer=writer
)
self.commit(t, None)
except Exception as exc:
self.commit(t, None, exc)
if reraise and futures:
# will be re-raised after futures are done
fut: asyncio.Future = loop.create_future()
fut.set_exception(exc)
done_futures.add(fut)
elif reraise:
self.commit(t, exc)
if reraise:
raise
if not futures: # maybe `t` schuduled another task
return
# add waiter task if requested
if get_waiter is not None:
futures[get_waiter()] = None
# schedule tasks
for t in tasks:
if not t.writes:
fut = cast(
asyncio.Future,
self.submit(
arun_with_retry,
t,
retry_policy,
stream=self.use_astream,
configurable={
CONFIG_KEY_SEND: partial(writer, t),
CONFIG_KEY_CALL: partial(call, t),
},
__name__=t.name,
__cancel_on_exit__=True,
__reraise_on_exit__=reraise,
),
)
fut.add_done_callback(partial(self.commit, t))
futures[fut] = t
# execute tasks, and wait for one to fail or all to finish.
# each task is independent from all other concurrent tasks
# yield updates/debug output as each task finishes
for t in tasks:
if not t.writes:
futures[
cast(
asyncio.Future,
self.submit(
arun_with_retry,
t,
retry_policy,
stream=self.use_astream,
writer=writer,
__name__=t.name,
__cancel_on_exit__=True,
__reraise_on_exit__=reraise,
),
)
] = t
done_futures: set[asyncio.Future] = set()
end_time = timeout + loop.time() if timeout else None
while len(futures) > (1 if get_waiter is not None else 0):
done, inflight = await asyncio.wait(
@@ -425,6 +270,8 @@ class PregelRunner:
else:
# store for panic check
done_futures.add(fut)
# task finished, commit writes
self.commit(task, _exception(fut))
else:
# remove references to loop vars
del fut, task
@@ -433,10 +280,6 @@ class PregelRunner:
break
# give control back to the caller
yield
# wait for pending done callbacks
# if a 2nd future finishes while `wait` is returning, it's possible
# that done callbacks for the 2nd future aren't called until next tick
await asyncio.sleep(0)
# cancel waiter task
for fut in futures:
fut.cancel()
@@ -448,24 +291,12 @@ class PregelRunner:
)
def commit(
self,
task: PregelExecutableTask,
fut: Union[None, concurrent.futures.Future[Any], asyncio.Future[Any]],
exception: Optional[BaseException] = None,
self, task: PregelExecutableTask, exception: Optional[BaseException]
) -> None:
if fut is not None:
exception = _exception(fut)
if isinstance(exception, asyncio.CancelledError):
# for cancelled tasks, also save error in task,
# so loop can finish super-step
task.writes.append((ERROR, exception))
self.put_writes(task.id, task.writes)
elif exception:
if exception:
if isinstance(exception, GraphInterrupt):
# save interrupt to checkpointer
if interrupts := [(INTERRUPT, i) for i in exception.args[0]]:
if resumes := [w for w in task.writes if w[0] == RESUME]:
interrupts.extend(resumes)
self.put_writes(task.id, interrupts)
elif isinstance(exception, GraphBubbleUp):
raise exception
@@ -491,12 +322,11 @@ def _should_stop_others(
GraphInterrupts are not considered failures."""
for fut in done:
if fut.cancelled():
continue
elif exc := fut.exception():
if not isinstance(exc, GraphBubbleUp):
return True
return False
return True
if exc := fut.exception():
return not isinstance(exc, GraphBubbleUp)
else:
return False
def _exception(
@@ -522,9 +352,7 @@ def _panic_or_proceed(
done: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set()
inflight: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set()
for fut in futs:
if fut.cancelled():
continue
elif fut.done():
if fut.done():
done.add(fut)
else:
inflight.add(fut)
@@ -537,6 +365,8 @@ def _panic_or_proceed(
# raise the exception
if panic:
raise exc
else:
return
if inflight:
# if we got here means we timed out
while inflight:
+27 -42
View File
@@ -36,40 +36,31 @@ class ChannelWriteEntry(NamedTuple):
"""Function to transform the value before writing."""
class ChannelWriteTupleEntry(NamedTuple):
mapper: Callable[[Any], Optional[Sequence[tuple[str, Any]]]]
"""Function to extract tuples from value."""
value: Any = PASSTHROUGH
"""Value to write, or PASSTHROUGH to use the input."""
class ChannelWrite(RunnableCallable):
"""Implements the logic for sending writes to CONFIG_KEY_SEND.
"""Implements th logic for sending writes to CONFIG_KEY_SEND.
Can be used as a runnable or as a static method to call imperatively."""
writes: list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]]
writes: list[Union[ChannelWriteEntry, Send]]
"""Sequence of write entries or Send objects to write."""
require_at_least_one_of: Optional[Sequence[str]]
"""If defined, at least one of these channels must be written to."""
def __init__(
self,
writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
writes: Sequence[Union[ChannelWriteEntry, Send]],
*,
tags: Optional[Sequence[str]] = None,
require_at_least_one_of: Optional[Sequence[str]] = None,
):
super().__init__(func=self._write, afunc=self._awrite, name=None, tags=tags)
self.writes = cast(
list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]], writes
)
self.writes = cast(list[Union[ChannelWriteEntry, Send]], writes)
self.require_at_least_one_of = require_at_least_one_of
def get_name(
self, suffix: Optional[str] = None, *, name: Optional[str] = None
) -> str:
if not name:
name = f"ChannelWrite<{','.join(w.channel if isinstance(w, ChannelWriteEntry) else '...' if isinstance(w, ChannelWriteTupleEntry) else w.node for w in self.writes)}>"
name = f"ChannelWrite<{','.join(w.channel if isinstance(w, ChannelWriteEntry) else w.node for w in self.writes)}>"
return super().get_name(suffix, name=name)
@property
@@ -88,8 +79,6 @@ class ChannelWrite(RunnableCallable):
writes = [
ChannelWriteEntry(write.channel, input, write.skip_none, write.mapper)
if isinstance(write, ChannelWriteEntry) and write.value is PASSTHROUGH
else ChannelWriteTupleEntry(write.mapper, input)
if isinstance(write, ChannelWriteTupleEntry) and write.value is PASSTHROUGH
else write
for write in self.writes
]
@@ -104,8 +93,6 @@ class ChannelWrite(RunnableCallable):
writes = [
ChannelWriteEntry(write.channel, input, write.skip_none, write.mapper)
if isinstance(write, ChannelWriteEntry) and write.value is PASSTHROUGH
else ChannelWriteTupleEntry(write.mapper, input)
if isinstance(write, ChannelWriteTupleEntry) and write.value is PASSTHROUGH
else write
for write in self.writes
]
@@ -119,7 +106,7 @@ class ChannelWrite(RunnableCallable):
@staticmethod
def do_write(
config: RunnableConfig,
writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
writes: Sequence[Union[ChannelWriteEntry, Send]],
require_at_least_one_of: Optional[Sequence[str]] = None,
) -> None:
# validate
@@ -131,34 +118,32 @@ class ChannelWrite(RunnableCallable):
)
if w.value is PASSTHROUGH:
raise InvalidUpdateError("PASSTHROUGH value must be replaced")
if isinstance(w, ChannelWriteTupleEntry):
if w.value is PASSTHROUGH:
raise InvalidUpdateError("PASSTHROUGH value must be replaced")
# assemble writes
tuples: list[tuple[str, Any]] = []
for w in writes:
if isinstance(w, Send):
tuples.append((PUSH if FF_SEND_V2 else TASKS, w))
elif isinstance(w, ChannelWriteTupleEntry):
if ww := w.mapper(w.value):
tuples.extend(ww)
elif isinstance(w, ChannelWriteEntry):
value = w.mapper(w.value) if w.mapper is not None else w.value
if value is SKIP_WRITE:
continue
if w.skip_none and value is None:
continue
tuples.append((w.channel, value))
else:
raise ValueError(f"Invalid write entry: {w}")
# assert required channels
# split packets and entries
sends = [
(PUSH if FF_SEND_V2 else TASKS, packet)
for packet in writes
if isinstance(packet, Send)
]
entries = [write for write in writes if isinstance(write, ChannelWriteEntry)]
# process entries into values
values = [
write.mapper(write.value) if write.mapper is not None else write.value
for write in entries
]
values = [
(write.channel, val)
for val, write in zip(values, entries)
if not write.skip_none or val is not None
]
# filter out SKIP_WRITE values
filtered = [(chan, val) for chan, val in values if val is not SKIP_WRITE]
if require_at_least_one_of is not None:
if not {chan for chan, _ in tuples} & set(require_at_least_one_of):
if not {chan for chan, _ in filtered} & set(require_at_least_one_of):
raise InvalidUpdateError(
f"Must write to at least one of {require_at_least_one_of}"
)
write: TYPE_SEND = config[CONF][CONFIG_KEY_SEND]
write(tuples)
write(sends + filtered)
@staticmethod
def is_writer(runnable: Runnable) -> bool:
+17 -90
View File
@@ -13,7 +13,6 @@ from typing import (
Optional,
Sequence,
Type,
TypedDict,
TypeVar,
Union,
cast,
@@ -22,24 +21,11 @@ from typing import (
from langchain_core.runnables import Runnable, RunnableConfig
from typing_extensions import Self
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
CheckpointMetadata,
PendingWrite,
)
from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
if TYPE_CHECKING:
from langgraph.store.base import BaseStore
try:
from langchain_core.messages.tool import ToolOutputMixin
except ImportError:
class ToolOutputMixin: # type: ignore[no-redef]
pass
All = Literal["*"]
"""Special value to indicate that graph should interrupt on all nodes."""
@@ -252,28 +238,13 @@ N = TypeVar("N", bound=Hashable)
@dataclasses.dataclass(**_DC_KWARGS)
class Command(Generic[N], ToolOutputMixin):
"""One or more commands to update the graph's state and send messages to nodes.
Args:
graph: graph to send the command to. Supported values are:
- None: the current graph (default)
- Command.PARENT: closest parent graph
update: update to apply to the graph's state.
resume: value to resume execution with. To be used together with [`interrupt()`][langgraph.types.interrupt].
goto: can be one of the following:
- name of the node to navigate to next (any node that belongs to the specified `graph`)
- sequence of node names to navigate to next
- `Send` object (to execute a node with the input provided)
- sequence of `Send` objects
"""
class Command(Generic[N]):
"""One or more commands to update the graph's state and send messages to nodes."""
graph: Optional[str] = None
update: Union[dict[str, Any], Sequence[tuple[str, Any]]] = ()
update: Optional[dict[str, Any]] = None
send: Union[Send, Sequence[Send]] = ()
resume: Optional[Union[Any, dict[str, Any]]] = None
goto: Union[Send, Sequence[Union[Send, str]], str] = ()
def __repr__(self) -> str:
# get all non-None values
@@ -284,17 +255,6 @@ class Command(Generic[N], ToolOutputMixin):
)
return f"Command({contents})"
def _update_as_tuples(self) -> Sequence[tuple[str, Any]]:
if isinstance(self.update, dict):
return list(self.update.items())
elif isinstance(self.update, (list, tuple)) and all(
isinstance(t, tuple) and len(t) == 2 and isinstance(t[0], str)
for t in self.update
):
return self.update
else:
return [("__root__", self.update)]
PARENT: ClassVar[Literal["__parent__"]] = "__parent__"
@@ -340,59 +300,26 @@ class LoopProtocol:
self.stop = stop
class PregelScratchpad(TypedDict, total=False):
interrupt_counter: int
used_null_resume: bool
resume: list[Any]
def interrupt(value: Any) -> Any:
from langgraph.constants import (
CONFIG_KEY_CHECKPOINT_NS,
CONFIG_KEY_SCRATCHPAD,
CONFIG_KEY_SEND,
CONFIG_KEY_TASK_ID,
CONFIG_KEY_WRITES,
CONFIG_KEY_RESUME_VALUE,
MISSING,
NS_SEP,
NULL_TASK_ID,
RESUME,
)
from langgraph.errors import GraphInterrupt
from langgraph.utils.config import get_configurable
conf = get_configurable()
# track interrupt index
scratchpad: PregelScratchpad = conf[CONFIG_KEY_SCRATCHPAD]
if "interrupt_counter" not in scratchpad:
scratchpad["interrupt_counter"] = 0
if (resume := conf.get(CONFIG_KEY_RESUME_VALUE, MISSING)) and resume is not MISSING:
return resume
else:
scratchpad["interrupt_counter"] += 1
idx = scratchpad["interrupt_counter"]
# find previous resume values
task_id = conf[CONFIG_KEY_TASK_ID]
writes: list[PendingWrite] = conf[CONFIG_KEY_WRITES]
scratchpad.setdefault(
"resume", next((w[2] for w in writes if w[0] == task_id and w[1] == RESUME), [])
)
if scratchpad["resume"]:
if idx < len(scratchpad["resume"]):
return scratchpad["resume"][idx]
# find current resume value
if not scratchpad.get("used_null_resume"):
scratchpad["used_null_resume"] = True
for tid, c, v in sorted(writes, key=lambda x: x[0], reverse=True):
if tid == NULL_TASK_ID and c == RESUME:
assert len(scratchpad["resume"]) == idx, (scratchpad["resume"], idx)
scratchpad["resume"].append(v)
conf[CONFIG_KEY_SEND]([(RESUME, scratchpad["resume"])])
return v
# no resume value found
raise GraphInterrupt(
(
Interrupt(
value=value,
resumable=True,
ns=cast(str, conf[CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP),
),
raise GraphInterrupt(
(
Interrupt(
value=value,
resumable=True,
ns=cast(str, conf[CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP),
),
)
)
)
+1 -42
View File
@@ -1,5 +1,5 @@
import dataclasses
from typing import Any, Generator, Optional, Type, Union, get_type_hints
from typing import Any, Optional, Type, Union
from typing_extensions import Annotated, NotRequired, ReadOnly, Required, get_origin
@@ -106,44 +106,3 @@ def get_field_default(name: str, type_: Any, schema: Type[Any]) -> Any:
if _is_optional_type(type_):
return None
return ...
def get_enhanced_type_hints(
type: Type[Any],
) -> Generator[tuple[str, Any, Any, Optional[str]], None, None]:
"""Attempt to extract default values and descriptions from provided type, used for config schema."""
for name, typ in get_type_hints(type).items():
default = None
description = None
# Pydantic models
try:
if hasattr(type, "__fields__") and name in type.__fields__:
field = type.__fields__[name]
if hasattr(field, "description") and field.description is not None:
description = field.description
if hasattr(field, "default") and field.default is not None:
default = field.default
if (
hasattr(default, "__class__")
and getattr(default.__class__, "__name__", "")
== "PydanticUndefinedType"
):
default = None
except (AttributeError, KeyError, TypeError):
pass
# TypedDict, dataclass
try:
if hasattr(type, "__dict__"):
type_dict = getattr(type, "__dict__")
if name in type_dict:
default = type_dict[name]
except (AttributeError, KeyError, TypeError):
pass
yield name, typ, default, description
-124
View File
@@ -1,124 +0,0 @@
import asyncio
import concurrent.futures
from typing import Union
AnyFuture = Union[asyncio.Future, concurrent.futures.Future]
def _get_loop(fut: asyncio.Future) -> asyncio.AbstractEventLoop:
# Tries to call Future.get_loop() if it's available.
# Otherwise fallbacks to using the old '_loop' property.
try:
get_loop = fut.get_loop
except AttributeError:
pass
else:
return get_loop()
return fut._loop
def _convert_future_exc(exc: BaseException) -> BaseException:
exc_class = type(exc)
if exc_class is concurrent.futures.CancelledError:
return asyncio.CancelledError(*exc.args)
elif exc_class is concurrent.futures.TimeoutError:
return asyncio.TimeoutError(*exc.args)
elif exc_class is concurrent.futures.InvalidStateError:
return asyncio.InvalidStateError(*exc.args)
else:
return exc
def _set_concurrent_future_state(
concurrent: concurrent.futures.Future,
source: AnyFuture,
) -> None:
"""Copy state from a future to a concurrent.futures.Future."""
assert source.done()
if source.cancelled():
concurrent.cancel()
if not concurrent.set_running_or_notify_cancel():
return
exception = source.exception()
if exception is not None:
concurrent.set_exception(_convert_future_exc(exception))
else:
result = source.result()
concurrent.set_result(result)
def _copy_future_state(source: AnyFuture, dest: asyncio.Future) -> None:
"""Internal helper to copy state from another Future.
The other Future may be a concurrent.futures.Future.
"""
assert source.done()
if dest.cancelled():
return
assert not dest.done()
if source.cancelled():
dest.cancel()
else:
exception = source.exception()
if exception is not None:
dest.set_exception(_convert_future_exc(exception))
else:
result = source.result()
dest.set_result(result)
def _chain_future(source: AnyFuture, destination: AnyFuture) -> None:
"""Chain two futures so that when one completes, so does the other.
The result (or exception) of source will be copied to destination.
If destination is cancelled, source gets cancelled too.
Compatible with both asyncio.Future and concurrent.futures.Future.
"""
if not asyncio.isfuture(source) and not isinstance(
source, concurrent.futures.Future
):
raise TypeError("A future is required for source argument")
if not asyncio.isfuture(destination) and not isinstance(
destination, concurrent.futures.Future
):
raise TypeError("A future is required for destination argument")
source_loop = _get_loop(source) if asyncio.isfuture(source) else None
dest_loop = _get_loop(destination) if asyncio.isfuture(destination) else None
def _set_state(future: AnyFuture, other: AnyFuture) -> None:
if asyncio.isfuture(future):
_copy_future_state(other, future)
else:
_set_concurrent_future_state(future, other)
def _call_check_cancel(destination: AnyFuture) -> None:
if destination.cancelled():
if source_loop is None or source_loop is dest_loop:
source.cancel()
else:
source_loop.call_soon_threadsafe(source.cancel)
def _call_set_state(source: AnyFuture) -> None:
if destination.cancelled() and dest_loop is not None and dest_loop.is_closed():
return
if dest_loop is None or dest_loop is source_loop:
_set_state(destination, source)
else:
if dest_loop.is_closed():
return
dest_loop.call_soon_threadsafe(_set_state, destination, source)
destination.add_done_callback(_call_check_cancel)
source.add_done_callback(_call_set_state)
def chain_future(source: AnyFuture, destination: concurrent.futures.Future) -> None:
# adapted from asyncio.run_coroutine_threadsafe
try:
_chain_future(source, destination)
except (SystemExit, KeyboardInterrupt):
raise
except BaseException as exc:
if destination.set_running_or_notify_cancel():
destination.set_exception(exc)
raise
+12 -4
View File
@@ -404,10 +404,12 @@ class RunnableSeq(Runnable):
config = patch_config(
config, callbacks=run_manager.get_child(f"seq:step:{i+1}")
)
context = copy_context()
context.run(_set_config_context, config)
if i == 0:
input = step.invoke(input, config, **kwargs)
input = context.run(step.invoke, input, config, **kwargs)
else:
input = step.invoke(input, config)
input = context.run(step.invoke, input, config)
# finish the root run
except BaseException as e:
run_manager.on_chain_error(e)
@@ -441,10 +443,16 @@ class RunnableSeq(Runnable):
config = patch_config(
config, callbacks=run_manager.get_child(f"seq:step:{i+1}")
)
context = copy_context()
context.run(_set_config_context, config)
if i == 0:
input = await step.ainvoke(input, config, **kwargs)
coro = step.ainvoke(input, config, **kwargs)
else:
input = await step.ainvoke(input, config)
coro = step.ainvoke(input, config)
if ASYNCIO_ACCEPTS_CONTEXT:
input = await asyncio.create_task(coro, context=context)
else:
input = await asyncio.create_task(coro)
# finish the root run
except BaseException as e:
await run_manager.on_chain_error(e)
+22 -10
View File
@@ -791,6 +791,17 @@ cli = ["click (==8.*)", "pygments (==2.*)", "rich (>=10,<14)"]
http2 = ["h2 (>=3,<5)"]
socks = ["socksio (==1.*)"]
[[package]]
name = "httpx-sse"
version = "0.4.0"
description = "Consume Server-Sent Event (SSE) messages with HTTPX."
optional = false
python-versions = ">=3.8"
files = [
{file = "httpx-sse-0.4.0.tar.gz", hash = "sha256:1e81a3a3070ce322add1d3529ed42eb5f70817f45ed6ec915ab753f961139721"},
{file = "httpx_sse-0.4.0-py3-none-any.whl", hash = "sha256:f329af6eae57eaa2bdfd962b42524764af68075ea87370a2de920af5341e318f"},
]
[[package]]
name = "idna"
version = "3.10"
@@ -1325,13 +1336,13 @@ files = [
[[package]]
name = "langchain-core"
version = "0.3.23"
version = "0.3.15"
description = "Building applications with LLMs through composability"
optional = false
python-versions = "<4.0,>=3.9"
files = [
{file = "langchain_core-0.3.23-py3-none-any.whl", hash = "sha256:550c0b996990830fa6515a71a1192a8a0343367999afc36d4ede14222941e420"},
{file = "langchain_core-0.3.23.tar.gz", hash = "sha256:f9e175e3b82063cc3b160c2ca2b155832e1c6f915312e1204828f97d4aabf6e1"},
{file = "langchain_core-0.3.15-py3-none-any.whl", hash = "sha256:3d4ca6dbb8ed396a6ee061063832a2451b0ce8c345570f7b086ffa7288e4fa29"},
{file = "langchain_core-0.3.15.tar.gz", hash = "sha256:b1a29787a4ffb7ec2103b4e97d435287201da7809b369740dd1e32f176325aba"},
]
[package.dependencies]
@@ -1348,7 +1359,7 @@ typing-extensions = ">=4.7"
[[package]]
name = "langgraph-checkpoint"
version = "2.0.8"
version = "2.0.4"
description = "Library with base interfaces for LangGraph checkpoint savers."
optional = false
python-versions = "^3.9.0,<4.0"
@@ -1382,7 +1393,7 @@ url = "../checkpoint-duckdb"
[[package]]
name = "langgraph-checkpoint-postgres"
version = "2.0.8"
version = "2.0.2"
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
optional = false
python-versions = "^3.9.0,<4.0"
@@ -1390,10 +1401,10 @@ files = []
develop = true
[package.dependencies]
langgraph-checkpoint = "^2.0.7"
langgraph-checkpoint = "^2.0.2"
orjson = ">=3.10.1"
psycopg = "^3.2.0"
psycopg-pool = "^3.2.0"
psycopg = "^3.0.0"
psycopg-pool = "^3.0.0"
[package.source]
type = "directory"
@@ -1418,7 +1429,7 @@ url = "../checkpoint-sqlite"
[[package]]
name = "langgraph-sdk"
version = "0.1.43"
version = "0.1.36"
description = "SDK for interacting with LangGraph API"
optional = false
python-versions = "^3.9.0,<4.0"
@@ -1427,6 +1438,7 @@ develop = true
[package.dependencies]
httpx = ">=0.25.2"
httpx-sse = ">=0.4.0"
orjson = ">=3.10.1"
[package.source]
@@ -3413,4 +3425,4 @@ type = ["pytest-mypy"]
[metadata]
lock-version = "2.0"
python-versions = ">=3.9.0,<4.0"
content-hash = "936530a5f00f329aeff2e6e921fe64480be317fad2c0a59cd54ea9018d089304"
content-hash = "9bf5668d3f70f3b77457906732404a6401583a5966f70a72ef10a68f2a5b27ad"
+3 -3
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph"
version = "0.2.58"
version = "0.2.53"
description = "Building stateful, multi-actor applications with LLMs"
authors = []
license = "MIT"
@@ -9,9 +9,9 @@ repository = "https://www.github.com/langchain-ai/langgraph"
[tool.poetry.dependencies]
python = ">=3.9.0,<4.0"
langchain-core = ">=0.2.43,<0.4.0,!=0.3.0,!=0.3.1,!=0.3.2,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.20,!=0.3.21,!=0.3.22"
langchain-core = ">=0.2.43,<0.4.0,!=0.3.0,!=0.3.1,!=0.3.2,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14"
langgraph-checkpoint = "^2.0.4"
langgraph-sdk = "^0.1.42"
langgraph-sdk = "^0.1.32"
[tool.poetry.group.dev.dependencies]
pytest = "^8.3.2"
+2 -641
View File
@@ -46,7 +46,7 @@ from langgraph.prebuilt import (
create_react_agent,
tools_condition,
)
from langgraph.prebuilt.chat_agent_executor import AgentState, _validate_chat_history
from langgraph.prebuilt.chat_agent_executor import _validate_chat_history
from langgraph.prebuilt.tool_node import (
TOOL_CALL_ERROR_TEMPLATE,
InjectedState,
@@ -56,7 +56,7 @@ from langgraph.prebuilt.tool_node import (
)
from langgraph.store.base import BaseStore
from langgraph.store.memory import InMemoryStore
from langgraph.types import Command, Interrupt, interrupt
from langgraph.types import Interrupt
from tests.conftest import (
ALL_CHECKPOINTERS_ASYNC,
ALL_CHECKPOINTERS_SYNC,
@@ -988,645 +988,6 @@ def test_tool_node_node_interrupt():
assert task.interrupts == (Interrupt(value="foo", when="during"),)
@pytest.mark.skipif(
not IS_LANGCHAIN_CORE_030_OR_GREATER,
reason="Langchain core 0.3.0 or greater is required",
)
async def test_tool_node_command():
from langchain_core.tools.base import InjectedToolCallId
@dec_tool
def transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
"""Transfer to Bob"""
return Command(
update={
"messages": [
ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
]
},
goto="bob",
graph=Command.PARENT,
)
@dec_tool
async def async_transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
"""Transfer to Bob"""
return Command(
update={
"messages": [
ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
]
},
goto="bob",
graph=Command.PARENT,
)
class CustomToolSchema(BaseModel):
tool_call_id: Annotated[str, InjectedToolCallId]
class MyCustomTool(BaseTool):
def _run(*args: Any, **kwargs: Any):
return Command(
update={
"messages": [
ToolMessage(
content="Transferred to Bob",
tool_call_id=kwargs["tool_call_id"],
)
]
},
goto="bob",
graph=Command.PARENT,
)
async def _arun(*args: Any, **kwargs: Any):
return Command(
update={
"messages": [
ToolMessage(
content="Transferred to Bob",
tool_call_id=kwargs["tool_call_id"],
)
]
},
goto="bob",
graph=Command.PARENT,
)
custom_tool = MyCustomTool(
name="custom_transfer_to_bob",
description="Transfer to bob",
args_schema=CustomToolSchema,
)
async_custom_tool = MyCustomTool(
name="async_custom_transfer_to_bob",
description="Transfer to bob",
args_schema=CustomToolSchema,
)
# test mixing regular tools and tools returning commands
def add(a: int, b: int) -> int:
"""Add two numbers"""
return a + b
result = ToolNode([add, transfer_to_bob]).invoke(
{
"messages": [
AIMessage(
"",
tool_calls=[
{"args": {"a": 1, "b": 2}, "id": "1", "name": "add"},
{"args": {}, "id": "2", "name": "transfer_to_bob"},
],
)
]
}
)
assert result == [
{
"messages": [
ToolMessage(
content="3",
tool_call_id="1",
name="add",
)
]
},
Command(
update={
"messages": [
ToolMessage(
content="Transferred to Bob",
tool_call_id="2",
name="transfer_to_bob",
)
]
},
goto="bob",
graph=Command.PARENT,
),
]
# test tools returning commands
# test sync tools
for tool in [transfer_to_bob, custom_tool]:
result = ToolNode([tool]).invoke(
{
"messages": [
AIMessage(
"", tool_calls=[{"args": {}, "id": "1", "name": tool.name}]
)
]
}
)
assert result == [
Command(
update={
"messages": [
ToolMessage(
content="Transferred to Bob",
tool_call_id="1",
name=tool.name,
)
]
},
goto="bob",
graph=Command.PARENT,
)
]
# test async tools
for tool in [async_transfer_to_bob, async_custom_tool]:
result = await ToolNode([tool]).ainvoke(
{
"messages": [
AIMessage(
"", tool_calls=[{"args": {}, "id": "1", "name": tool.name}]
)
]
}
)
assert result == [
Command(
update={
"messages": [
ToolMessage(
content="Transferred to Bob",
tool_call_id="1",
name=tool.name,
)
]
},
goto="bob",
graph=Command.PARENT,
)
]
# test multiple commands
result = ToolNode([transfer_to_bob, custom_tool]).invoke(
{
"messages": [
AIMessage(
"",
tool_calls=[
{"args": {}, "id": "1", "name": "transfer_to_bob"},
{"args": {}, "id": "2", "name": "custom_transfer_to_bob"},
],
)
]
}
)
assert result == [
Command(
update={
"messages": [
ToolMessage(
content="Transferred to Bob",
tool_call_id="1",
name="transfer_to_bob",
)
]
},
goto="bob",
graph=Command.PARENT,
),
Command(
update={
"messages": [
ToolMessage(
content="Transferred to Bob",
tool_call_id="2",
name="custom_transfer_to_bob",
)
]
},
goto="bob",
graph=Command.PARENT,
),
]
# test validation (mismatch between input type and command.update type)
with pytest.raises(ValueError):
@dec_tool
def list_update_tool(tool_call_id: Annotated[str, InjectedToolCallId]):
"""My tool"""
return Command(
update=[ToolMessage(content="foo", tool_call_id=tool_call_id)]
)
ToolNode([list_update_tool]).invoke(
{
"messages": [
AIMessage(
"",
tool_calls=[
{"args": {}, "id": "1", "name": "list_update_tool"}
],
)
]
}
)
# test validation (missing tool message in the update for current graph)
with pytest.raises(ValueError):
@dec_tool
def no_update_tool():
"""My tool"""
return Command(update={"messages": []})
ToolNode([no_update_tool]).invoke(
{
"messages": [
AIMessage(
"",
tool_calls=[{"args": {}, "id": "1", "name": "no_update_tool"}],
)
]
}
)
# test validation (missing tool message in the update for parent graph is OK)
@dec_tool
def node_update_parent_tool():
"""No update"""
return Command(update={"messages": []}, graph=Command.PARENT)
assert ToolNode([node_update_parent_tool]).invoke(
{
"messages": [
AIMessage(
"",
tool_calls=[
{"args": {}, "id": "1", "name": "node_update_parent_tool"}
],
)
]
}
) == [Command(update={"messages": []}, graph=Command.PARENT)]
# test validation (multiple tool messages)
with pytest.raises(ValueError):
for graph in (None, Command.PARENT):
@dec_tool
def multiple_tool_messages_tool():
"""My tool"""
return Command(
update={
"messages": [
ToolMessage(content="foo", tool_call_id=""),
ToolMessage(content="bar", tool_call_id=""),
]
},
graph=graph,
)
ToolNode([multiple_tool_messages_tool]).invoke(
{
"messages": [
AIMessage(
"",
tool_calls=[
{
"args": {},
"id": "1",
"name": "multiple_tool_messages_tool",
}
],
)
]
}
)
@pytest.mark.skipif(
not IS_LANGCHAIN_CORE_030_OR_GREATER,
reason="Langchain core 0.3.0 or greater is required",
)
async def test_tool_node_command_list_input():
from langchain_core.tools.base import InjectedToolCallId
@dec_tool
def transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
"""Transfer to Bob"""
return Command(
update=[
ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
],
goto="bob",
graph=Command.PARENT,
)
@dec_tool
async def async_transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
"""Transfer to Bob"""
return Command(
update=[
ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
],
goto="bob",
graph=Command.PARENT,
)
class CustomToolSchema(BaseModel):
tool_call_id: Annotated[str, InjectedToolCallId]
class MyCustomTool(BaseTool):
def _run(*args: Any, **kwargs: Any):
return Command(
update=[
ToolMessage(
content="Transferred to Bob",
tool_call_id=kwargs["tool_call_id"],
)
],
goto="bob",
graph=Command.PARENT,
)
async def _arun(*args: Any, **kwargs: Any):
return Command(
update=[
ToolMessage(
content="Transferred to Bob",
tool_call_id=kwargs["tool_call_id"],
)
],
goto="bob",
graph=Command.PARENT,
)
custom_tool = MyCustomTool(
name="custom_transfer_to_bob",
description="Transfer to bob",
args_schema=CustomToolSchema,
)
async_custom_tool = MyCustomTool(
name="async_custom_transfer_to_bob",
description="Transfer to bob",
args_schema=CustomToolSchema,
)
# test mixing regular tools and tools returning commands
def add(a: int, b: int) -> int:
"""Add two numbers"""
return a + b
result = ToolNode([add, transfer_to_bob]).invoke(
[
AIMessage(
"",
tool_calls=[
{"args": {"a": 1, "b": 2}, "id": "1", "name": "add"},
{"args": {}, "id": "2", "name": "transfer_to_bob"},
],
)
]
)
assert result == [
[
ToolMessage(
content="3",
tool_call_id="1",
name="add",
)
],
Command(
update=[
ToolMessage(
content="Transferred to Bob",
tool_call_id="2",
name="transfer_to_bob",
)
],
goto="bob",
graph=Command.PARENT,
),
]
# test tools returning commands
# test sync tools
for tool in [transfer_to_bob, custom_tool]:
result = ToolNode([tool]).invoke(
[AIMessage("", tool_calls=[{"args": {}, "id": "1", "name": tool.name}])]
)
assert result == [
Command(
update=[
ToolMessage(
content="Transferred to Bob",
tool_call_id="1",
name=tool.name,
)
],
goto="bob",
graph=Command.PARENT,
)
]
# test async tools
for tool in [async_transfer_to_bob, async_custom_tool]:
result = await ToolNode([tool]).ainvoke(
[AIMessage("", tool_calls=[{"args": {}, "id": "1", "name": tool.name}])]
)
assert result == [
Command(
update=[
ToolMessage(
content="Transferred to Bob",
tool_call_id="1",
name=tool.name,
)
],
goto="bob",
graph=Command.PARENT,
)
]
# test multiple commands
result = ToolNode([transfer_to_bob, custom_tool]).invoke(
[
AIMessage(
"",
tool_calls=[
{"args": {}, "id": "1", "name": "transfer_to_bob"},
{"args": {}, "id": "2", "name": "custom_transfer_to_bob"},
],
)
]
)
assert result == [
Command(
update=[
ToolMessage(
content="Transferred to Bob",
tool_call_id="1",
name="transfer_to_bob",
)
],
goto="bob",
graph=Command.PARENT,
),
Command(
update=[
ToolMessage(
content="Transferred to Bob",
tool_call_id="2",
name="custom_transfer_to_bob",
)
],
goto="bob",
graph=Command.PARENT,
),
]
# test validation (mismatch between input type and command.update type)
with pytest.raises(ValueError):
@dec_tool
def list_update_tool(tool_call_id: Annotated[str, InjectedToolCallId]):
"""My tool"""
return Command(
update={
"messages": [ToolMessage(content="foo", tool_call_id=tool_call_id)]
}
)
ToolNode([list_update_tool]).invoke(
[
AIMessage(
"",
tool_calls=[{"args": {}, "id": "1", "name": "list_update_tool"}],
)
]
)
# test validation (missing tool message in the update for current graph)
with pytest.raises(ValueError):
@dec_tool
def no_update_tool():
"""My tool"""
return Command(update=[])
ToolNode([no_update_tool]).invoke(
[
AIMessage(
"",
tool_calls=[{"args": {}, "id": "1", "name": "no_update_tool"}],
)
]
)
# test validation (missing tool message in the update for parent graph is OK)
@dec_tool
def node_update_parent_tool():
"""No update"""
return Command(update=[], graph=Command.PARENT)
assert ToolNode([node_update_parent_tool]).invoke(
[
AIMessage(
"",
tool_calls=[{"args": {}, "id": "1", "name": "node_update_parent_tool"}],
)
]
) == [Command(update=[], graph=Command.PARENT)]
# test validation (multiple tool messages)
with pytest.raises(ValueError):
for graph in (None, Command.PARENT):
@dec_tool
def multiple_tool_messages_tool():
"""My tool"""
return Command(
update=[
ToolMessage(content="foo", tool_call_id=""),
ToolMessage(content="bar", tool_call_id=""),
],
graph=graph,
)
ToolNode([multiple_tool_messages_tool]).invoke(
[
AIMessage(
"",
tool_calls=[
{
"args": {},
"id": "1",
"name": "multiple_tool_messages_tool",
}
],
)
]
)
@pytest.mark.skipif(
not IS_LANGCHAIN_CORE_030_OR_GREATER,
reason="Langchain core 0.3.0 or greater is required",
)
def test_react_agent_update_state():
from langchain_core.tools.base import InjectedToolCallId
class State(AgentState):
user_name: str
@dec_tool
def get_user_name(tool_call_id: Annotated[str, InjectedToolCallId]):
"""Retrieve user name"""
user_name = interrupt("Please provider user name:")
return Command(
update={
"user_name": user_name,
"messages": [
ToolMessage(
"Successfully retrieved user name", tool_call_id=tool_call_id
)
],
}
)
def state_modifier(state: State):
user_name = state.get("user_name")
if user_name is None:
return state["messages"]
system_msg = f"User name is {user_name}"
return [{"role": "system", "content": system_msg}] + state["messages"]
checkpointer = MemorySaver()
tool_calls = [[{"args": {}, "id": "1", "name": "get_user_name"}]]
model = FakeToolCallingModel(tool_calls=tool_calls)
agent = create_react_agent(
model,
[get_user_name],
state_schema=State,
state_modifier=state_modifier,
checkpointer=checkpointer,
)
config = {"configurable": {"thread_id": "1"}}
# run until interrpupted
agent.invoke({"messages": [("user", "what's my name")]}, config)
# supply the value for the interrupt
response = agent.invoke(Command(resume="Archibald"), config)
# confirm that the state was updated
assert response["user_name"] == "Archibald"
assert len(response["messages"]) == 4
tool_message: ToolMessage = response["messages"][-2]
assert tool_message.content == "Successfully retrieved user name"
assert tool_message.tool_call_id == "1"
assert tool_message.name == "get_user_name"
def my_function(some_val: int, some_other_val: str) -> str:
return f"{some_val} - {some_other_val}"
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+1 -59
View File
@@ -21,11 +21,7 @@ from typing_extensions import Annotated, NotRequired, Required
from langgraph.graph import END, StateGraph
from langgraph.graph.graph import CompiledGraph
from langgraph.utils.fields import (
_is_optional_type,
get_enhanced_type_hints,
get_field_default,
)
from langgraph.utils.fields import _is_optional_type, get_field_default
from langgraph.utils.runnable import is_async_callable, is_async_generator
pytestmark = pytest.mark.anyio
@@ -231,57 +227,3 @@ def test_is_required():
assert get_field_default("val_12", gcannos["val_12"], MyGrandChildDict) is None
assert get_field_default("val_9", gcannos["val_9"], MyGrandChildDict) is None
assert get_field_default("val_13", gcannos["val_13"], MyGrandChildDict) == ...
def test_enhanced_type_hints() -> None:
from dataclasses import dataclass
from typing import Annotated
from pydantic import BaseModel, Field
class MyTypedDict(TypedDict):
val_1: str
val_2: int = 42
val_3: str = "default"
hints = list(get_enhanced_type_hints(MyTypedDict))
assert len(hints) == 3
assert hints[0] == ("val_1", str, None, None)
assert hints[1] == ("val_2", int, 42, None)
assert hints[2] == ("val_3", str, "default", None)
@dataclass
class MyDataclass:
val_1: str
val_2: int = 42
val_3: str = "default"
hints = list(get_enhanced_type_hints(MyDataclass))
assert len(hints) == 3
assert hints[0] == ("val_1", str, None, None)
assert hints[1] == ("val_2", int, 42, None)
assert hints[2] == ("val_3", str, "default", None)
class MyPydanticModel(BaseModel):
val_1: str
val_2: int = 42
val_3: str = Field(default="default", description="A description")
hints = list(get_enhanced_type_hints(MyPydanticModel))
assert len(hints) == 3
assert hints[0] == ("val_1", str, None, None)
assert hints[1] == ("val_2", int, 42, None)
assert hints[2] == ("val_3", str, "default", "A description")
class MyPydanticModelWithAnnotated(BaseModel):
val_1: Annotated[str, Field(description="A description")]
val_2: Annotated[int, Field(default=42)]
val_3: Annotated[
str, Field(default="default", description="Another description")
]
hints = list(get_enhanced_type_hints(MyPydanticModelWithAnnotated))
assert len(hints) == 3
assert hints[0] == ("val_1", str, None, "A description")
assert hints[1] == ("val_2", int, 42, None)
assert hints[2] == ("val_3", str, "default", "Another description")

Some files were not shown because too many files have changed in this diff Show More