mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-18 05:35:43 +02:00
Compare commits
98
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1130c3accb | ||
|
|
ee8653d1c5 | ||
|
|
12486d977a | ||
|
|
c87f9ab6b1 | ||
|
|
855a3d21ff | ||
|
|
d767af421b | ||
|
|
07ac016e60 | ||
|
|
4576a259dd | ||
|
|
53ec7c41b2 | ||
|
|
769f6a1925 | ||
|
|
62a36befd5 | ||
|
|
dfaff2511b | ||
|
|
1d9a0d1e4e | ||
|
|
35c7eb18ee | ||
|
|
dc09b13400 | ||
|
|
b2d8acffc4 | ||
|
|
1031e54860 | ||
|
|
7ac365ea84 | ||
|
|
5144b8f374 | ||
|
|
4b1b3cecb4 | ||
|
|
c6a953c02a | ||
|
|
16b955dee2 | ||
|
|
877124f7df | ||
|
|
d3a4865c0e | ||
|
|
a3761ac522 | ||
|
|
376c58ff3b | ||
|
|
58b99c899e | ||
|
|
2ee279a977 | ||
|
|
f04ce5d1ee | ||
|
|
8f649abd0a | ||
|
|
1febec7c0d | ||
|
|
a4eb4c6942 | ||
|
|
8e1cd0e225 | ||
|
|
98935e1ffd | ||
|
|
328ef609af | ||
|
|
486d5412af | ||
|
|
abc0c8c223 | ||
|
|
5bbb9dae57 | ||
|
|
fed60e713c | ||
|
|
4f4e7a6981 | ||
|
|
3351d4f6c5 | ||
|
|
b4900341e4 | ||
|
|
65f515e020 | ||
|
|
0d0665a6e3 | ||
|
|
93b8525dc1 | ||
|
|
aeb6f784e1 | ||
|
|
3eedeac0d4 | ||
|
|
b09e7b20b0 | ||
|
|
26ce731eab | ||
|
|
55593446f8 | ||
|
|
7082e2613e | ||
|
|
ceeb9636ee | ||
|
|
f7788abbb6 | ||
|
|
9bd430142a | ||
|
|
72dac006f4 | ||
|
|
54d848913f | ||
|
|
7021e81150 | ||
|
|
b977045679 | ||
|
|
a933776436 | ||
|
|
588373c2d5 | ||
|
|
267962bece | ||
|
|
4ae29b6e2a | ||
|
|
3c0de26914 | ||
|
|
a570662773 | ||
|
|
9766068896 | ||
|
|
7e8eef88ca | ||
|
|
e3e63c70c9 | ||
|
|
312f0982bc | ||
|
|
153245145e | ||
|
|
c95abd88a1 | ||
|
|
a2b357bed5 | ||
|
|
7090d7e9a8 | ||
|
|
b3fa43e4a6 | ||
|
|
b1779cf348 | ||
|
|
26d18d3ca5 | ||
|
|
12052d7d26 | ||
|
|
e3a30a9b69 | ||
|
|
ff1370a9a5 | ||
|
|
679a7365da | ||
|
|
b2522ffe19 | ||
|
|
4212a795a0 | ||
|
|
517d67aa32 | ||
|
|
feaf14765a | ||
|
|
cc6063c729 | ||
|
|
013397042e | ||
|
|
9a775d9c9f | ||
|
|
2c945ceb68 | ||
|
|
39eabd0fb8 | ||
|
|
e5cc2e2044 | ||
|
|
f00c0515e7 | ||
|
|
d87c0d4d53 | ||
|
|
fb40a974c8 | ||
|
|
d63bfc6879 | ||
|
|
97dd30711a | ||
|
|
7866bd2718 | ||
|
|
433c382280 | ||
|
|
7352ab14a2 | ||
|
|
85a76912d3 |
@@ -22,8 +22,7 @@ def test(
|
||||
# check docker available
|
||||
capabilities = langgraph_cli.docker.check_capabilities(runner)
|
||||
# open config
|
||||
with open(config) as f:
|
||||
config_json = langgraph_cli.config.validate_config(json.load(f))
|
||||
config_json = langgraph_cli.config.validate_config_file(config)
|
||||
|
||||
set("Running...")
|
||||
args = [
|
||||
|
||||
@@ -88,6 +88,7 @@ jobs:
|
||||
--check-links-ignore "https://(api|web|docs)\.smith\.langchain\.com/.*" \
|
||||
--check-links-ignore "https://x.com/.*" \
|
||||
--check-links-ignore "https://github\.com/.*" \
|
||||
--check-links-ignore "http://localhost:8123/.*" \
|
||||
--check-links-ignore "/.*\.(ipynb|html)$" \
|
||||
--check-links-ignore "https://python\.langchain\.com/.*" \
|
||||
--check-links-ignore "https://openai\.com/.*" \
|
||||
@@ -104,6 +105,7 @@ jobs:
|
||||
echo "Running link check on HTML files matching changed notebook files..."
|
||||
poetry run pytest -v \
|
||||
--check-links-ignore "https://(api|web|docs)\.smith\.langchain\.com/.*" \
|
||||
--check-links-ignore "http://localhost:8123/.*" \
|
||||
--check-links-ignore "https://x.com/.*" \
|
||||
--check-links-ignore "https://github\.com/.*" \
|
||||
--check-links-ignore "/.*\.(ipynb|html)$" \
|
||||
|
||||
@@ -238,7 +238,7 @@ final_state["messages"][-1].content
|
||||
* [How-to Guides](https://langchain-ai.github.io/langgraph/how-tos/): Accomplish specific things within LangGraph, from streaming, to adding memory & persistence, to common design patterns (branching, subgraphs, etc.), these are the place to go if you want to copy and run a specific code snippet.
|
||||
* [Conceptual Guides](https://langchain-ai.github.io/langgraph/concepts/high_level/): In-depth explanations of the key concepts and principles behind LangGraph, such as nodes, edges, state and more.
|
||||
* [API Reference](https://langchain-ai.github.io/langgraph/reference/graphs/): Review important classes and methods, simple examples of how to use the graph and checkpointing APIs, higher-level prebuilt components and more.
|
||||
* [Cloud (beta)](https://langchain-ai.github.io/langgraph/cloud/): With one click, deploy LangGraph applications to LangGraph Cloud.
|
||||
* [LangGraph Platform](https://langchain-ai.github.io/langgraph/concepts/#langgraph-platform): LangGraph Platform is a commercial solution for deploying agentic applications in production, built on the open-source LangGraph framework.
|
||||
|
||||
## Contributing
|
||||
|
||||
|
||||
@@ -36,10 +36,11 @@ NOTEBOOKS_NO_EXECUTION = [
|
||||
"docs/docs/tutorials/rag/langgraph_self_rag_local.ipynb",
|
||||
# this loads a massive dataset from gcp
|
||||
"docs/docs/tutorials/usaco/usaco.ipynb",
|
||||
# TODO: figure out why autogen notebook is not runnable (they are just hanging. possible due to code execution?)
|
||||
"docs/docs/how-tos/autogen-integration.ipynb",
|
||||
# TODO: need to update these notebooks to make sure they are runnable in CI
|
||||
"docs/docs/tutorials/storm/storm.ipynb", # issues only when running with VCR
|
||||
"docs/docs/tutorials/lats/lats.ipynb", # issues only when running with VCR
|
||||
"docs/docs/tutorials/multi_agent/hierarchical_agent_teams.ipynb", # taking a very long time to run
|
||||
"docs/docs/tutorials/rag/langgraph_crag.ipynb", # flakiness from tavily
|
||||
"docs/docs/tutorials/rag/langgraph_adaptive_rag.ipynb", # Cannot create a consistent method resolution error from VCR
|
||||
"docs/docs/how-tos/map-reduce.ipynb" # flakiness from structured output, only when running with VCR
|
||||
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
+1
File diff suppressed because one or more lines are too long
+1
File diff suppressed because one or more lines are too long
+1
File diff suppressed because one or more lines are too long
-1
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
+1
File diff suppressed because one or more lines are too long
@@ -21,7 +21,7 @@ Install the proper packages:
|
||||
Ensure you have an API key, which you can create from the [LangSmith UI](https://smith.langchain.com) (Settings > API Keys). This is required to authenticate that you have LangGraph Cloud access. After you have saved the key to a safe place, place the following line in your `.env` file:
|
||||
|
||||
```python
|
||||
LANGCHAIN_API_KEY = *********
|
||||
LANGSMITH_API_KEY = *********
|
||||
```
|
||||
|
||||
## Start the API server
|
||||
@@ -54,7 +54,7 @@ You can either initialize by passing authentication or by setting an environment
|
||||
from langgraph_sdk import get_client
|
||||
|
||||
# only pass the url argument to get_client() if you changed the default port when calling langgraph up
|
||||
client = get_client(url=<DEPLOYMENT_URL>,api_key=<LANGCHAIN_API_KEY>)
|
||||
client = get_client(url=<DEPLOYMENT_URL>,api_key=<LANGSMITH_API_KEY>)
|
||||
# Using the graph deployed with the name "agent"
|
||||
assistant_id = "agent"
|
||||
thread = await client.threads.create()
|
||||
@@ -66,7 +66,7 @@ You can either initialize by passing authentication or by setting an environment
|
||||
import { Client } from "@langchain/langgraph-sdk";
|
||||
|
||||
// only set the apiUrl if you changed the default port when calling langgraph up
|
||||
const client = new Client({ apiUrl: <DEPLOYMENT_URL>, apiKey: <LANGCHAIN_API_KEY> });
|
||||
const client = new Client({ apiUrl: <DEPLOYMENT_URL>, apiKey: <LANGSMITH_API_KEY> });
|
||||
// Using the graph deployed with the name "agent"
|
||||
const assistantId = "agent";
|
||||
const thread = await client.threads.create();
|
||||
@@ -78,13 +78,13 @@ You can either initialize by passing authentication or by setting an environment
|
||||
curl --request POST \
|
||||
--url <DEPLOYMENT_URL>/threads \
|
||||
--header 'Content-Type: application/json'
|
||||
--header 'x-api-key: <LANGCHAIN_API_KEY>'
|
||||
--header 'x-api-key: <LANGSMITH_API_KEY>'
|
||||
```
|
||||
|
||||
|
||||
#### Initialize with environment variables
|
||||
|
||||
If you have a `LANGCHAIN_API_KEY` set in your environment, you do not need to explicitly pass authentication to the client
|
||||
If you have a `LANGSMITH_API_KEY` set in your environment, you do not need to explicitly pass authentication to the client
|
||||
|
||||
=== "Python"
|
||||
|
||||
@@ -154,7 +154,7 @@ Now we can invoke our graph to ensure it is working. Make sure to change the inp
|
||||
}
|
||||
```
|
||||
|
||||
=== "CURL"
|
||||
=== "CURL"
|
||||
|
||||
```bash
|
||||
curl --request POST \
|
||||
|
||||
@@ -8,9 +8,9 @@ If you want to learn how to build an agent like this from scratch, take a look a
|
||||
|
||||
This tutorial will use:
|
||||
|
||||
- Anthropic for the LLM - sign up and get an API key [here](https://console.anthropic.com/)
|
||||
- Tavily for the search engine - sign up and get an API key [here](https://app.tavily.com/)
|
||||
- LangSmith for hosting - sign up and get an API key [here](https://smith.langchain.com/)
|
||||
- Anthropic for the LLM - sign up and get an API key [here](https://console.anthropic.com/).
|
||||
- Tavily for the search engine - sign up and get an API key [here](https://app.tavily.com/).
|
||||
- LangSmith for hosting - sign up and get an API key [here](https://smith.langchain.com/).
|
||||
|
||||
## Create and configure your app
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,19 +0,0 @@
|
||||
<!doctype html>
|
||||
<html>
|
||||
<head>
|
||||
<title>Open Assistants API Specification</title>
|
||||
<meta charset="utf-8" />
|
||||
<meta
|
||||
name="viewport"
|
||||
content="width=device-width, initial-scale=1" />
|
||||
</head>
|
||||
<body>
|
||||
<script id="api-reference" data-url="./open_agent_api.json"></script>
|
||||
<script>
|
||||
var configuration = {}
|
||||
document.getElementById('api-reference').dataset.configuration =
|
||||
JSON.stringify(configuration)
|
||||
</script>
|
||||
<script src="https://cdn.jsdelivr.net/npm/@scalar/api-reference"></script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -1557,8 +1557,11 @@
|
||||
"200": {
|
||||
"description": "Success",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {}
|
||||
"text/event-stream": {
|
||||
"schema": {
|
||||
"type": "string",
|
||||
"description": "The server will send a stream of events in SSE format.\n\n**Example event**:\n\nid: 1\n\nevent: message\n\ndata: {}"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -1905,8 +1908,11 @@
|
||||
"200": {
|
||||
"description": "Success",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {}
|
||||
"text/event-stream": {
|
||||
"schema": {
|
||||
"type": "string",
|
||||
"description": "The server will send a stream of events in SSE format.\n\n**Example event**:\n\nid: 1\n\nevent: message\n\ndata: {}"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -2143,8 +2149,11 @@
|
||||
"200": {
|
||||
"description": "Success",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {}
|
||||
"text/event-stream": {
|
||||
"schema": {
|
||||
"type": "string",
|
||||
"description": "The server will send a stream of events in SSE format.\n\n**Example event**:\n\nid: 1\n\nevent: message\n\ndata: {}"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
|
||||
@@ -53,7 +53,7 @@ Example:
|
||||
}
|
||||
```
|
||||
|
||||
Example:
|
||||
Example with environment variables:
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -78,6 +78,37 @@ The base command for the LangGraph CLI is `langgraph`.
|
||||
langgraph [OPTIONS] COMMAND [ARGS]
|
||||
```
|
||||
|
||||
### `dev`
|
||||
|
||||
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.
|
||||
|
||||
**Installation**
|
||||
|
||||
This command requires the "inmem" extra to be installed:
|
||||
|
||||
```bash
|
||||
pip install -U "langgraph-cli[inmem]"
|
||||
```
|
||||
|
||||
**Usage**
|
||||
|
||||
```
|
||||
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 |
|
||||
|
||||
### `build`
|
||||
|
||||
Build LangGraph Cloud API server Docker image.
|
||||
@@ -100,7 +131,7 @@ langgraph build [OPTIONS]
|
||||
|
||||
### `up`
|
||||
|
||||
Start langgraph API server. For local testing, requires a LangSmith API key with access to LangGraph Cloud closed beta. Requires a license key for production use.
|
||||
Start LangGraph API server. For local testing, requires a LangSmith API key with access to LangGraph Cloud closed beta. Requires a license key for production use.
|
||||
|
||||
**Usage**
|
||||
|
||||
@@ -120,8 +151,8 @@ langgraph up [OPTIONS]
|
||||
| `--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 test --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` |
|
||||
| `-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. |
|
||||
|
||||
@@ -148,9 +179,9 @@ Example:
|
||||
langgraph dockerfile -c langgraph.json Dockerfile
|
||||
```
|
||||
|
||||
Would generate something like the following:
|
||||
This generates a Dockerfile that looks similar to:
|
||||
|
||||
```text
|
||||
```dockerfile
|
||||
FROM langchain/langgraph-api:3.11
|
||||
|
||||
ADD ./pipconf.txt /pipconfig.txt
|
||||
@@ -170,6 +201,3 @@ RUN set -ex && \
|
||||
RUN PIP_CONFIG_FILE=/pipconfig.txt PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
|
||||
ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_graphs/src/agent.py:graph", "storm": "/deps/__outer_graphs/src/storm.py:graph"}'
|
||||
```
|
||||
|
||||
You can then customize, build images, push, and deploy from this file.
|
||||
|
||||
@@ -2,13 +2,13 @@
|
||||
|
||||
LLMs are extremely powerful, particularly when connected to other systems such as a retriever or APIs. This is why many LLM applications use a control flow of steps before and / or after LLM calls. As an example [RAG](https://github.com/langchain-ai/rag-from-scratch) performs retrieval of relevant documents to a question, and passes those documents to an LLM in order to ground the response. Often a control flow of steps before and / or after an LLM is called a "chain." Chains are a popular paradigm for programming with LLMs and offer a high degree of reliability; the same set of steps runs with each chain invocation.
|
||||
|
||||
However, we often want LLM systems that can pick their own control flow! This is one definition of an [agent](https://blog.langchain.dev/what-is-an-agent/): an agent is a system that uses an LLM to decide the control flow of an application. Unlike a chain, an agent given an LLM some degree of control over the sequence of steps in the application. Examples of using an LLM to decide the control of an application:
|
||||
However, we often want LLM systems that can pick their own control flow! This is one definition of an [agent](https://blog.langchain.dev/what-is-an-agent/): an agent is a system that uses an LLM to decide the control flow of an application. Unlike a chain, an agent gives an LLM some degree of control over the sequence of steps in the application. Examples of using an LLM to decide the control of an application:
|
||||
|
||||
- Using an LLM to route between two potential paths
|
||||
- Using an LLM to decide which of many tools to call
|
||||
- Using an LLM to decide whether the generated answer is sufficient or more work is need
|
||||
|
||||
There are many different types of [agent architectures](https://blog.langchain.dev/what-is-a-cognitive-architecture/) to consider, which given an LLM varying levels of control. On one extreme, a router allows an LLM to select a single step from a specified set of options and, on the other extreme, a fully autonomous long-running agent may have complete freedom to select any sequence of steps that it wants for a given problem.
|
||||
There are many different types of [agent architectures](https://blog.langchain.dev/what-is-a-cognitive-architecture/) to consider, which give an LLM varying levels of control. On one extreme, a router allows an LLM to select a single step from a specified set of options and, on the other extreme, a fully autonomous long-running agent may have complete freedom to select any sequence of steps that it wants for a given problem.
|
||||
|
||||

|
||||
|
||||
|
||||
@@ -30,7 +30,7 @@ The conceptual guide does not cover step-by-step instructions or specific implem
|
||||
- [Streaming](streaming.md): Streaming is crucial for enhancing the responsiveness of applications built on LLMs. By displaying output progressively, even before a complete response is ready, streaming significantly improves user experience (UX), particularly when dealing with the latency of LLMs.
|
||||
- [FAQ](faq.md): Frequently asked questions about LangGraph.
|
||||
|
||||
## LangGraph Platform
|
||||
## LangGraph Platform
|
||||
|
||||
LangGraph Platform is a commercial solution for deploying agentic applications in production, built on the open-source LangGraph framework.
|
||||
|
||||
|
||||
@@ -28,9 +28,28 @@ The CLI provides the following core functionality:
|
||||
|
||||
The `langgraph build` command builds a Docker image for the [LangGraph API server](./langgraph_server.md) that can be directly deployed.
|
||||
|
||||
### `dev`
|
||||
|
||||
!!! note "New in version 0.1.55"
|
||||
The `langgraph dev` command was introduced in langgraph-cli version 0.1.55.
|
||||
|
||||
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
|
||||
- Debugger support: Attach your IDE's debugger for line-by-line debugging
|
||||
- In-memory state with local persistence: Server state is stored in memory for speed but persisted locally between restarts
|
||||
|
||||
To use this command, you need to install the CLI with the "inmem" extra:
|
||||
|
||||
```bash
|
||||
pip install -U "langgraph-cli[inmem]"
|
||||
```
|
||||
|
||||
**Note**: This command is intended for local development and testing only. It is not recommended for production use. Since it does not use Docker, we recommend using virtual environments to manage your project's dependencies.
|
||||
|
||||
### `up`
|
||||
|
||||
The `langgraph up` command starts an instance of the [LangGraph API server](./langgraph_server.md) locally. This requires docker to be installed and running locally. It also requires a LangSmith API key for local development or a license key for production use.
|
||||
The `langgraph up` command starts an instance of the [LangGraph API server](./langgraph_server.md) locally in a docker container. This requires thedocker server to be running locally. It also requires a LangSmith API key for local development or a license key for production use.
|
||||
|
||||
The server includes all API endpoints for your graph's runs, threads, assistants, etc. as well as the other services required to run your agent, including a managed database for checkpointing and storage.
|
||||
|
||||
|
||||
@@ -14,6 +14,13 @@ A **deployment** is an instance of a LangGraph API. A single deployment can have
|
||||
|
||||
See the [how-to guide](../cloud/deployment/cloud.md#create-new-deployment) for creating a new deployment.
|
||||
|
||||
## Resource Allocation
|
||||
|
||||
| **Deployment Type** | **CPU** | **Memory** | **Scaling** |
|
||||
|---------------------|---------|------------|---------------------|
|
||||
| Development | 1 CPU | 1 GB | Up to 1 container |
|
||||
| Production | 1 CPU | 2 GB | Up to 10 containers |
|
||||
|
||||
## Revision
|
||||
|
||||
A revision is an iteration of a [deployment](#deployment). When a new deployment is created, an initial revision is automatically created. To deploy new code changes or update environment variable configurations for a deployment, a new revision must be created. When a revision is created, a new container image is built automatically.
|
||||
@@ -33,6 +40,7 @@ A high-level diagram of a Cloud SaaS deployment.
|
||||
|
||||

|
||||
|
||||
|
||||
## Related
|
||||
|
||||
- [Deployment Options](./deployment_options.md)
|
||||
|
||||
@@ -35,6 +35,16 @@ While in Beta, LangGraph Studio is available for free to all [LangSmith](https:/
|
||||
|
||||
If you have deployed your LangGraph application on LangGraph Platform (Cloud), you can access the studio as part of that
|
||||
|
||||
### Development server
|
||||
|
||||
LangGraph CLI also contains a command for running an in-memory development server that can be used to connect a local LangGraph app with the studio.
|
||||
See [instructions here](../cloud/reference/cli.md#dev) for more information.
|
||||
|
||||
The way this works is that it runs inside your local environment.
|
||||
It will spin up an in-memory, development server to deploy the graph.
|
||||
You can then connect to the studio via the Cloud hosted version of LangGraph Platform.
|
||||
To be clear, the web studio will connect to your locally running server - your agent is still running locally and never leaves your device.
|
||||
|
||||
## Studio FAQs
|
||||
|
||||
### Why is my project failing to start?
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
|
||||
## Versions
|
||||
|
||||
There are two versions of the self hosted deployment: [Self-Hosted Enterprise](./deployment_options.md#self-hosted-enterprise) and [Self-Hosted Lite](./deployment_options.md#self-hosted-lite).
|
||||
There are two versions of the self-hosted deployment: [Self-Hosted Enterprise](./deployment_options.md#self-hosted-enterprise) and [Self-Hosted Lite](./deployment_options.md#self-hosted-lite).
|
||||
|
||||
### Self-Hosted Lite
|
||||
|
||||
@@ -34,6 +34,10 @@ To use the Self-Hosted Enterprise version, you must acquire a license key that y
|
||||
|
||||
For step-by-step instructions, see [How to set up a self-hosted deployment of LangGraph](../how-tos/deploy-self-hosted.md).
|
||||
|
||||
## Helm Chart
|
||||
|
||||
If you would like to deploy LangGraph Cloud on Kubernetes, you can use this [Helm chart](https://github.com/langchain-ai/helm/blob/main/charts/langgraph-cloud/README.md).
|
||||
|
||||
## Related
|
||||
|
||||
- [How to set up a self-hosted deployment of LangGraph](../how-tos/deploy-self-hosted.md).
|
||||
|
||||
@@ -6,22 +6,14 @@
|
||||
|
||||
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.
|
||||
|
||||
Templates can be accessed via [LangGraph Studio](langgraph_studio.md), or cloned directly from Github. You can download LangGraph Studio and see available templates [here](https://studio.langchain.com/).
|
||||
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/).
|
||||
|
||||
## Available templates
|
||||
|
||||
- **New LangGraph Project**: A simple, minimal chatbot with memory.
|
||||
- [Python](https://github.com/langchain-ai/new-langgraph-project)
|
||||
- [JS/TS](https://github.com/langchain-ai/new-langgraphjs-project)
|
||||
- **ReAct Agent**: A simple agent that can be flexibly extended to many tools.
|
||||
- [Python](https://github.com/langchain-ai/react-agent)
|
||||
- [JS/TS](https://github.com/langchain-ai/react-agent-js)
|
||||
- **Memory Agent**: A ReAct-style agent with an additional tool to store memories for use across conversational threads.
|
||||
- [Python](https://github.com/langchain-ai/memory-agent)
|
||||
- [JS/TS](https://github.com/langchain-ai/memory-agent-js)
|
||||
- **Retrieval Agent**: An agent that includes a retrieval-based question-answering system.
|
||||
- [Python](https://github.com/langchain-ai/retrieval-agent-template)
|
||||
- [JS/TS](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.
|
||||
- [Python](https://github.com/langchain-ai/data-enrichment)
|
||||
- [JS/TS](https://github.com/langchain-ai/data-enrichment-js)
|
||||
| Template | Description | Python | JS/TS |
|
||||
|---------------------------|------------------------------------------------------------------------------------------|------------------------------------------------------------------|---------------------------------------------------------------------|
|
||||
| **New LangGraph Project** | A simple, minimal chatbot with memory. | [Repo](https://github.com/langchain-ai/new-langgraph-project) | [Repo](https://github.com/langchain-ai/new-langgraphjs-project) |
|
||||
| **ReAct Agent** | A simple agent that can be flexibly extended to many tools. | [Repo](https://github.com/langchain-ai/react-agent) | [Repo](https://github.com/langchain-ai/react-agent-js) |
|
||||
| **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) |
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,171 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "8381b6e0-29a6-48c5-b451-5d2549351249",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# How to use LangGraph Platform to deploy CrewAI, AutoGen, and other frameworks\n",
|
||||
"\n",
|
||||
"[LangGraph Platform](https://langchain-ai.github.io/langgraph/concepts/langgraph_platform/) provides infrastructure for deploying agents. This integrates seamlessly with LangGraph, but can also work with other frameworks. The way to make this work is to wrap the agent in a single LangGraph node, and have that be the entire graph.\n",
|
||||
"\n",
|
||||
"Doing so will allow you to deploy to LangGraph Platform, and allows you to get a lot of the [benefits](https://langchain-ai.github.io/langgraph/concepts/langgraph_platform/). You get horizontally scalable infrastructure, a task queue to handle bursty operations, a persistence layer to power short term memory, and long term memory support.\n",
|
||||
"\n",
|
||||
"In this guide we show how to do this with an AutoGen agent, but this method should work for agents defined in other frameworks like CrewAI, LlamaIndex, and others as well."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1113cb16-b538-448c-924c-85731ce96ebd",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Setup"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"id": "f05993fa-9d03-4f45-bc13-0a8d87260d86",
|
||||
"metadata": {
|
||||
"scrolled": true
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%pip install autogen langgraph"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "f4e0ca12-1714-4776-a30a-9527e519799b",
|
||||
"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",
|
||||
"id": "1926bbc3-6b06-41e0-9604-860a2bbf8fa3",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Define autogen agent\n",
|
||||
"\n",
|
||||
"Here we define our AutoGen agent. From https://github.com/microsoft/autogen/blob/0.2/notebook/agentchat_web_info.ipynb"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "d4a14dc7-d565-4207-8788-525f85b9fb27",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import autogen\n",
|
||||
"import os\n",
|
||||
"\n",
|
||||
"config_list = [{\"model\": \"gpt-4o\", \"api_key\": os.environ[\"OPENAI_API_KEY\"]}]\n",
|
||||
"\n",
|
||||
"llm_config = {\n",
|
||||
" \"timeout\": 600,\n",
|
||||
" \"cache_seed\": 42,\n",
|
||||
" \"config_list\": config_list,\n",
|
||||
" \"temperature\": 0,\n",
|
||||
"}\n",
|
||||
"\n",
|
||||
"autogen_agent = autogen.AssistantAgent(\n",
|
||||
" name=\"assistant\",\n",
|
||||
" llm_config=llm_config,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"user_proxy = autogen.UserProxyAgent(\n",
|
||||
" name=\"user_proxy\",\n",
|
||||
" human_input_mode=\"NEVER\",\n",
|
||||
" max_consecutive_auto_reply=10,\n",
|
||||
" is_termination_msg=lambda x: x.get(\"content\", \"\").rstrip().endswith(\"TERMINATE\"),\n",
|
||||
" code_execution_config={\n",
|
||||
" \"work_dir\": \"web\",\n",
|
||||
" \"use_docker\": False,\n",
|
||||
" }, # Please set use_docker=True if docker is available to run the generated code. Using docker is safer than running the generated code directly.\n",
|
||||
" llm_config=llm_config,\n",
|
||||
" system_message=\"Reply TERMINATE if the task has been solved at full satisfaction. Otherwise, reply CONTINUE, or the reason why the task is not solved yet.\",\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b1170836-f23e-4e4c-ab83-ce791cd7fbd2",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Wrap in LangGraph\n",
|
||||
"\n",
|
||||
"We now wrap the AutoGen agent in a single LangGraph node, and make that the entire graph.\n",
|
||||
"The main thing this involves is defining an Input and Output schema for the node, which you would need to do if deploying this manually, so it's no extra work"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"id": "7b417c16-ff4e-4d5c-a9a9-0aaeeef6ede5",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from langgraph.graph import StateGraph, MessagesState\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def call_autogen_agent(state: MessagesState):\n",
|
||||
" last_message = state[\"messages\"][-1]\n",
|
||||
" response = user_proxy.initiate_chat(autogen_agent, message=last_message.content)\n",
|
||||
" # get the final response from the agent\n",
|
||||
" content = response.chat_history[-1][\"content\"]\n",
|
||||
" return {\"messages\": {\"role\": \"assistant\", \"content\": content}}\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"graph = StateGraph(MessagesState)\n",
|
||||
"graph.add_node(call_autogen_agent)\n",
|
||||
"graph.set_entry_point(\"call_autogen_agent\")\n",
|
||||
"graph = graph.compile()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f6a18377-ac29-478f-a76a-b213f1a3c85d",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Deploy with LangGraph Platform\n",
|
||||
"\n",
|
||||
"You can now deploy this as you normally would with LangGraph Platform. See [these instructions](https://langchain-ai.github.io/langgraph/concepts/deployment_options/) for more details."
|
||||
]
|
||||
}
|
||||
],
|
||||
"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
|
||||
}
|
||||
@@ -345,7 +345,7 @@
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.4"
|
||||
"version": "3.11.1"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
|
||||
@@ -17,14 +17,18 @@ You will need to do the following:
|
||||
2. Build a docker image with the [LangGraph Server](../concepts/langgraph_server.md) using the [LangGraph CLI](../concepts/langgraph_cli.md).
|
||||
3. Deploy a web server that will run the docker image and pass in the necessary environment variables.
|
||||
|
||||
## Helm Chart
|
||||
|
||||
If you would like to deploy LangGraph Cloud on Kubernetes, you can use this [Helm chart](https://github.com/langchain-ai/helm/blob/main/charts/langgraph-cloud/README.md).
|
||||
|
||||
## Environment Variables
|
||||
|
||||
You will eventually need to pass in the following environment variables to the LangGraph Deploy server:
|
||||
|
||||
- `REDIS_URI`: Connection details to a Redis instance. Redis will be used as a pub-sub broker to enable streaming real time output from background runs.
|
||||
- `DATABASE_URI`: Postgres connection details. Postgres will be used to store assistants, threads, runs, persist thread state and long term memory, and to manage the state of the background task queue with 'exactly once' semantics.
|
||||
- `LANGSMITH_API_KEY`: (If using [Self-Hosted Lite]) LangSmith API key. This will be used to authenticate ONCE at server start up.
|
||||
- `LANGGRAPH_CLOUD_LICENSE_KEY`: (If using Self-Hosted Enterprise) LangGraph Platform license key. This will be used to authenticate ONCE at server start up.
|
||||
- `LANGSMITH_API_KEY`: (If using [Self-Hosted Lite](../concepts/deployment_options.md#self-hosted-lite)) LangSmith API key. This will be used to authenticate ONCE at server start up.
|
||||
- `LANGGRAPH_CLOUD_LICENSE_KEY`: (If using [Self-Hosted Enterprise](../concepts/deployment_options.md#self-hosted-enterprise)) LangGraph Platform license key. This will be used to authenticate ONCE at server start up.
|
||||
|
||||
|
||||
## Build the Docker Image
|
||||
@@ -70,7 +74,7 @@ If you want to run this quickly without setting up a separate Redis and Postgres
|
||||
* You need to replace `my-image` with the name of the image you built in the previous step (from `langgraph build`).
|
||||
and you should provide appropriate values for `REDIS_URI`, `DATABASE_URI`, and `LANGSMITH_API_KEY`.
|
||||
* If your application requires additional environment variables, you can pass them in a similar way.
|
||||
* If using Self-Hosted Enterprise, you must provide `LANGGRAPH_CLOUD_LICENSE_KEY` as an additional environment variable.
|
||||
* If using [Self-Hosted Enterprise](../concepts/deployment_options.md#self-hosted-enterprise), you must provide `LANGGRAPH_CLOUD_LICENSE_KEY` as an additional environment variable.
|
||||
|
||||
|
||||
### Using Docker Compose
|
||||
|
||||
@@ -103,6 +103,7 @@ These how-to guides show common patterns for tool calling with LangGraph:
|
||||
- [How to force function calling agent to structure output](react-agent-structured-output.ipynb)
|
||||
- [How to pass custom LangSmith run ID for graph runs](run-id-langsmith.ipynb)
|
||||
- [How to return state before hitting recursion limit](return-when-recursion-limit-hits.ipynb)
|
||||
- [How to integrate LangGraph with AutoGen, CrewAI, and other frameworks](autogen-integration.ipynb)
|
||||
|
||||
### Prebuilt ReAct Agent
|
||||
|
||||
@@ -141,6 +142,7 @@ Learn how to set up your app for deployment to LangGraph Platform:
|
||||
- [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)
|
||||
- [How to use LangGraph Platform to deploy CrewAI, AutoGen, and other frameworks](autogen-langgraph-platform.ipynb)
|
||||
|
||||
### Deployment
|
||||
|
||||
@@ -149,6 +151,7 @@ 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)
|
||||
|
||||
### Assistants
|
||||
|
||||
[Assistants](../concepts/assistants.md) is a configured instance of a template.
|
||||
@@ -163,7 +166,7 @@ LangGraph applications can be deployed using LangGraph Cloud, which provides a r
|
||||
|
||||
### Runs
|
||||
|
||||
LangGraph Cloud supports multiple types of runs besides streaming runs.
|
||||
LangGraph Platform supports multiple types of runs besides streaming runs.
|
||||
|
||||
- [How to run an agent in the background](../cloud/how-tos/background_run.md)
|
||||
- [How to run multiple agents in the same thread](../cloud/how-tos/same-thread.md)
|
||||
@@ -183,7 +186,7 @@ Streaming the results of your LLM application is vital for ensuring a good user
|
||||
|
||||
### Human-in-the-loop
|
||||
|
||||
When creating complex graphs, leaving every decision up to the LLM can be dangerous, especially when the decisions involve invoking certain tools or accessing specific documents. To remedy this, LangGraph allows you to insert human-in-the-loop behavior to ensure your graph does not have undesired outcomes. Read more about the different ways you can add human-in-the-loop capabilities to your LangGraph Cloud projects in these how-to guides:
|
||||
When designing complex graphs, relying entirely on the LLM for decision-making can be risky, particularly when it involves tools that interact with files, APIs, or databases. These interactions may lead to unintended data access or modifications, depending on the use case. To mitigate these risks, LangGraph allows you to integrate human-in-the-loop behavior, ensuring your LLM applications operate as intended without undesirable outcomes.
|
||||
|
||||
- [How to add a breakpoint](../cloud/how-tos/human_in_the_loop_breakpoint.md)
|
||||
- [How to wait for user input](../cloud/how-tos/human_in_the_loop_user_input.md)
|
||||
@@ -193,7 +196,7 @@ When creating complex graphs, leaving every decision up to the LLM can be danger
|
||||
|
||||
### 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. The following how-to guides provide information on the various options LangGraph Cloud gives you for dealing with 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.
|
||||
|
||||
- [How to use the interrupt option](../cloud/how-tos/interrupt_concurrent.md)
|
||||
- [How to use the rollback option](../cloud/how-tos/rollback_concurrent.md)
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
# How to connect a local agent to LangGraph Studio
|
||||
|
||||
This guide shows you how to connect your local agent to [LangGraph Studio](../concepts/langgraph_studio.md) for visualization, interaction, and debugging.
|
||||
|
||||
## Connection Options
|
||||
|
||||
There are two ways to connect your local agent to LangGraph Studio:
|
||||
|
||||
- [LangGraph Desktop](../concepts/langgraph_studio.md#desktop-app): Application, Mac only, requires Docker
|
||||
- [Development Server](../concepts/langgraph_studio.md#dev-server): Python package, all platforms, no Docker
|
||||
|
||||
In this guide we will cover how to use the development server as that is generally an easier and better experience.
|
||||
|
||||
## Setup your application
|
||||
|
||||
First, you will need to setup your application in the proper format.
|
||||
This means defining a `langgraph.json` file which contains paths to your agent(s).
|
||||
See [this guide](../concepts/application_structure.md) for information on how to do so.
|
||||
|
||||
## Install langgraph-cli
|
||||
|
||||
You will need to install [`langgraph-cli`](../cloud/reference/cli.md#langgraph-cli) (version `0.1.55` or higher).
|
||||
You will need to make sure to install the `inmem` extras.
|
||||
|
||||
```shell
|
||||
pip install "langgraph-cli[inmem]==0.1.55"
|
||||
```
|
||||
|
||||
## Run the development server
|
||||
|
||||
1. Navigate to your project directory (where `langgraph.json` is located)
|
||||
|
||||
2. Start the server:
|
||||
```bash
|
||||
langgraph dev
|
||||
```
|
||||
|
||||
This will look for the `langgraph.json` file in your current directory.
|
||||
In there, it will find the paths to the graph(s), and start those up.
|
||||
It will then automatically connect to the cloud-hosted studio.
|
||||
|
||||
## Use the studio
|
||||
|
||||
After connecting to the studio, a browser window should automatically pop up.
|
||||
This will use the cloud hosted studio UI to connect to your local development server.
|
||||
Your graph is still running locally, the UI is connecting to visualizing the agent and threads that are defined locally.
|
||||
|
||||
The graph will always use the most up-to-date code, so you will be able to change the underlying code and have it automatically reflected in the studio.
|
||||
This is useful for debugging workflows.
|
||||
You can run your graph in the UI until it messes up, go in and change your code, and then rerun from the node that failed.
|
||||
|
||||
# (Optional) Attach a debugger
|
||||
|
||||
For step-by-step debugging with breakpoints and variable inspection:
|
||||
|
||||
```bash
|
||||
# Install debugpy package
|
||||
pip install debugpy
|
||||
|
||||
# Start server with debugging enabled
|
||||
langgraph dev --debug-port 5678
|
||||
```
|
||||
|
||||
Then attach your preferred debugger:
|
||||
|
||||
=== "VS Code"
|
||||
Add this configuration to `launch.json`:
|
||||
```json
|
||||
{
|
||||
"name": "Attach to LangGraph",
|
||||
"type": "debugpy",
|
||||
"request": "attach",
|
||||
"connect": {
|
||||
"host": "0.0.0.0",
|
||||
"port": 5678
|
||||
}
|
||||
}
|
||||
```
|
||||
Specify the port number you chose in the previous step.
|
||||
|
||||
=== "PyCharm"
|
||||
1. Go to Run → Edit Configurations
|
||||
2. Click + and select "Python Debug Server"
|
||||
3. Set IDE host name: `localhost`
|
||||
4. Set port: `5678` (or the port number you chose in the previous step)
|
||||
5. Click "OK" and start debugging
|
||||
@@ -6,25 +6,23 @@ title: Tutorials
|
||||
|
||||
# Tutorials
|
||||
|
||||
Welcome to the LangGraph Tutorials! These notebooks introduce LangGraph through building various language agents and applications.
|
||||
New to LangGraph or LLM app development? Read this material to get up and running building your first applications.
|
||||
|
||||
## Quick Start
|
||||
## Get Started 🚀 {#quick-start}
|
||||
|
||||
Learn the basics of LangGraph through a comprehensive quick start in which you will build an agent from scratch.
|
||||
- [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.
|
||||
|
||||
- [Quick Start](introduction.ipynb): In this tutorial, you will build a support chatbot using LangGraph.
|
||||
- [LangGraph Cloud Quick Start](../cloud/quick_start.md): In this tutorial, you will build and deploy an agent to LangGraph Cloud.
|
||||
## Use cases 🛠️
|
||||
|
||||
## Use cases
|
||||
|
||||
Learn from example implementations of graphs designed for specific scenarios and that implement common design patterns.
|
||||
Explore practical implementations tailored for specific scenarios:
|
||||
|
||||
### Chatbots
|
||||
|
||||
- [Customer Support](customer-support/customer-support.ipynb): Build a customer support chatbot to manage flights, hotel reservations, car rentals, and other tasks
|
||||
- [Prompt Generation from User Requirements](chatbots/information-gather-prompting.ipynb): Build an information gathering chatbot
|
||||
- [Code Assistant](code_assistant/langgraph_code_assistant.ipynb): Build a code analysis and generation assistant
|
||||
|
||||
- [Customer Support](customer-support/customer-support.ipynb): Build a multi-functional support bot for flights, hotels, and car rentals.
|
||||
- [Prompt Generation from User Requirements](chatbots/information-gather-prompting.ipynb): Build an information gathering chatbot.
|
||||
- [Code Assistant](code_assistant/langgraph_code_assistant.ipynb): Build a code analysis and generation assistant.
|
||||
|
||||
### RAG
|
||||
|
||||
|
||||
@@ -5,21 +5,21 @@
|
||||
"id": "4a1aae78-88a6-4133-b905-7e46c8e3772f",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# LangGraph Quick Start\n",
|
||||
"# 🚀 LangGraph Quick Start\n",
|
||||
"\n",
|
||||
"In this comprehensive quick start, we will build a support chatbot in LangGraph that can:\n",
|
||||
"In this tutorial, we will build a support chatbot in LangGraph that can:\n",
|
||||
"\n",
|
||||
"- Answer common questions by searching the web\n",
|
||||
"- Maintain conversation state across calls\n",
|
||||
"- Route complex queries to a human for review\n",
|
||||
"- Use custom state to control its behavior\n",
|
||||
"- Rewind and explore alternative conversation paths\n",
|
||||
"✅ **Answer common questions** by searching the web \n",
|
||||
"✅ **Maintain conversation state** across calls \n",
|
||||
"✅ **Route complex queries** to a human for review \n",
|
||||
"✅ **Use custom state** to control its behavior \n",
|
||||
"✅ **Rewind and explore** alternative conversation paths \n",
|
||||
"\n",
|
||||
"We'll start with a basic chatbot and progressively add more sophisticated capabilities, introducing key LangGraph concepts along the way.\n",
|
||||
"We'll start with a **basic chatbot** and progressively add more sophisticated capabilities, introducing key LangGraph concepts along the way. Let’s dive in! 🌟\n",
|
||||
"\n",
|
||||
"## Setup\n",
|
||||
"\n",
|
||||
"First, install the required packages:"
|
||||
"First, install the required packages and configure your environment:"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -33,14 +33,6 @@
|
||||
"%pip install -U langgraph langsmith langchain_anthropic"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a6d1e870-1bc0-4d44-86c0-96681ccf6113",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Next, set your API keys:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
@@ -120,27 +112,24 @@
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "31c755cd-8994-4867-bdff-96a55d7beae7",
|
||||
"id": "c08c41da-0855-49d3-9a3d-b7eb94413367",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"<div class=\"admonition tip\">\n",
|
||||
" <p class=\"admonition-title\">Note</p>\n",
|
||||
" <p>\n",
|
||||
" The first thing you do when you define a graph is define the <code>State</code> of the graph. The <code>State</code> consists of the schema of the graph as well as <a href=\"https://langchain-ai.github.io/langgraph/concepts/low_level/#reducers\">reducer functions</a> which specify how to apply updates to the state. In our example <code>State</code> is a <code>TypedDict</code> with a single key: <code>messages</code>. The <code>messages</code> key is annotated with the <a href=\"https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.message.add_messages\"><code>add_messages</code></a> reducer function, which tells LangGraph to append new messages to the existing list, rather than overwriting it. State keys without an annotation will be overwritten by each update, storing the most recent value. Check out <a href=\"https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.message.add_messages\">this conceptual guide</a> to learn more about state, reducers and other low-level concepts.\n",
|
||||
" </p>\n",
|
||||
"</div>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "4137feed-746e-4c72-a34a-f7a699ad5dcf",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"So now our graph knows two things:\n",
|
||||
"Our graph can now handle two key tasks:\n",
|
||||
"\n",
|
||||
"1. Each `node` can receive the current `State` as input and output an update to the state.\n",
|
||||
"2. Updates to `messages` will be appended to the existing list rather than overwriting it, thanks to the prebuilt [`add_messages`](https://langchain-ai.github.io/langgraph/reference/graphs/?h=add+messages#add_messages) function used with the `Annotated` syntax.\n",
|
||||
"\n",
|
||||
"------\n",
|
||||
"\n",
|
||||
"!!! tip \"Concept\"\n",
|
||||
"\n",
|
||||
" When defining a graph, the first step is to define its `State`. The `State` includes the graph's schema and [reducer functions](https://langchain-ai.github.io/langgraph/concepts/low_level/#reducers) that handle state updates. In our example, `State` is a `TypedDict` with one key: `messages`. The [`add_messages`](https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.message.add_messages) reducer function is used to append new messages to the list instead of overwriting it. Keys without a reducer annotation will overwrite previous values. Learn more about state, reducers, and related concepts in [this guide](https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.message.add_messages).\n",
|
||||
"\n",
|
||||
"---------\n",
|
||||
"\n",
|
||||
"1. Every `node` we define will receive the current `State` as input and return a value that updates that state.\n",
|
||||
"2. `messages` will be _appended_ to the current list, rather than directly overwritten. This is communicated via the prebuilt [`add_messages`](https://langchain-ai.github.io/langgraph/reference/graphs/?h=add+messages#add_messages) function in the `Annotated` syntax.\n",
|
||||
"\n",
|
||||
"Next, add a \"`chatbot`\" node. Nodes represent units of work. They are typically regular python functions."
|
||||
]
|
||||
@@ -365,7 +354,7 @@
|
||||
"id": "f22c5d4a-3134-413c-81fe-dd9752fbeb66",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Part 2: Enhancing the Chatbot with Tools\n",
|
||||
"## Part 2: 🛠️ Enhancing the Chatbot with Tools\n",
|
||||
"\n",
|
||||
"To handle queries our chatbot can't answer \"from memory\", we'll integrate a web search tool. Our bot can use this tool to find relevant information and provide better responses.\n",
|
||||
"\n",
|
||||
@@ -2046,7 +2035,7 @@
|
||||
"\n",
|
||||
"So far, we've relied on a simple state (it's just a list of messages!). You can go far with this simple state, but if you want to define complex behavior without relying on the message list, you can add additional fields to the state. In this section, we will extend our chat bot with a new node to illustrate this.\n",
|
||||
"\n",
|
||||
"In the examples above, we involved a human deterministically: the graph __always__ interrupted whenever an tool was invoked. Suppose we wanted our chat bot to have the choice of relying on a human.\n",
|
||||
"In the examples above, we involved a human deterministically: the graph __always__ interrupted whenever a tool was invoked. Suppose we wanted our chat bot to have the choice of relying on a human.\n",
|
||||
"\n",
|
||||
"One way to do this is to create a passthrough \"human\" node, before which the graph will always stop. We will only execute this node if the LLM invokes a \"human\" tool. For our convenience, we will include an \"ask_human\" flag in our graph state that we will flip if the LLM calls this tool.\n",
|
||||
"\n",
|
||||
@@ -3136,11 +3125,29 @@
|
||||
"id": "e584d57f-5aad-4507-815f-0b2e4b64b791",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Conclusion\n",
|
||||
"## Next Steps\n",
|
||||
"\n",
|
||||
"Congrats! You've completed the intro tutorial and built a chat bot in LangGraph that supports tool calling, persistent memory, human-in-the-loop interactivity, and even time-travel!\n",
|
||||
"Take your journey further by exploring deployment and advanced features:\n",
|
||||
"\n",
|
||||
"The [LangGraph documentation](https://langchain-ai.github.io/langgraph/) is a great resource for diving deeper into the library's capabilities."
|
||||
"### Server Quickstart\n",
|
||||
"\n",
|
||||
"- **[LangGraph Server Quickstart](../langgraph-platform/local-server)**: Launch a LangGraph server locally and interact with it using the REST API and LangGraph Studio Web UI.\n",
|
||||
"\n",
|
||||
"### LangGraph Cloud\n",
|
||||
"\n",
|
||||
"- **[LangGraph Cloud QuickStart](../../cloud/quick_start)**: Deploy your LangGraph app using LangGraph Cloud.\n",
|
||||
"\n",
|
||||
"### LangGraph Framework\n",
|
||||
"\n",
|
||||
"- **[LangGraph Concepts](../../concepts)**: Learn the foundational concepts of LangGraph. \n",
|
||||
"- **[LangGraph How-to Guides](../../how-tos)**: Guides for common tasks with LangGraph.\n",
|
||||
"\n",
|
||||
"### LangGraph Platform\n",
|
||||
"\n",
|
||||
"Expand your knowledge with these resources:\n",
|
||||
"\n",
|
||||
"- **[LangGraph Platform Concepts](../../concepts#langgraph-platform)**: Understand the foundational concepts of the LangGraph Platform. \n",
|
||||
"- **[LangGraph Platform How-to Guides](../../how-tos#langgraph-platform)**: Guides for common tasks with LangGraph Platform. "
|
||||
]
|
||||
}
|
||||
],
|
||||
@@ -3160,7 +3167,7 @@
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.9"
|
||||
"version": "3.11.4"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
|
||||
@@ -0,0 +1,253 @@
|
||||
# Quick Start: Launch Local LangGraph Server
|
||||
|
||||
This is a quick start guide to help you get a LangGraph app up and running locally.
|
||||
|
||||
!!! 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
|
||||
```
|
||||
|
||||
## 🌱 Create a LangGraph App
|
||||
|
||||
Create a new app from the `react-agent` template. This template is a simple agent that can be flexibly extended to many tools.
|
||||
|
||||
=== "Python Server"
|
||||
|
||||
```shell
|
||||
langgraph new path/to/your/app --template react-agent-python
|
||||
```
|
||||
|
||||
=== "Node Server"
|
||||
|
||||
```shell
|
||||
langgraph new path/to/your/app --template react-agent-js
|
||||
```
|
||||
|
||||
!!! tip "Additional Templates"
|
||||
|
||||
If you use `langgraph new` without specifying a template, you will be presented with an interactive menu that will allow you to choose from a list of available templates.
|
||||
|
||||
## Install Dependencies
|
||||
|
||||
In the root of your new LangGraph app, install the dependencies:
|
||||
|
||||
```shell
|
||||
pip install .
|
||||
```
|
||||
|
||||
## Create a `.env` file
|
||||
|
||||
You will find a `.env.example` in the root of your new LangGraph app. Create
|
||||
a `.env` file in the root of your new LangGraph app and copy the contents of the `.env.example` file into it, filling in the necessary API keys:
|
||||
|
||||
```bash
|
||||
LANGSMITH_API_KEY=lsv2...
|
||||
TAVILY_API_KEY=tvly-...
|
||||
ANTHROPIC_API_KEY=sk-
|
||||
OPENAI_API_KEY=sk-...
|
||||
```
|
||||
|
||||
<details><summary>Get API Keys</summary>
|
||||
<ul>
|
||||
<li> <b>LANGSMITH_API_KEY</b>: Go to the <a href="https://smith.langchain.com/settings">LangSmith Settings page</a>. Then clck <b>Create API Key</b>.
|
||||
</li>
|
||||
<li>
|
||||
<b>ANTHROPIC_API_KEY</b>: Get an API key from <a href="https://console.anthropic.com/">Anthropic</a>.
|
||||
</li>
|
||||
<li>
|
||||
<b>OPENAI_API_KEY</b>: Get an API key from <a href="https://openai.com/">OpenAI</a>.
|
||||
</li>
|
||||
<li>
|
||||
<b>TAVILY_API_KEY</b>: Get an API key on the <a href="https://app.tavily.com/">Tavily website</a>.
|
||||
</li>
|
||||
</ul>
|
||||
</details>
|
||||
|
||||
## 🚀 Launch LangGraph Server
|
||||
|
||||
```shell
|
||||
langgraph dev
|
||||
```
|
||||
|
||||
This will start up the LangGraph API server locally. If this runs successfully, you should see something like:
|
||||
|
||||
> Ready!
|
||||
>
|
||||
> - API: [http://localhost:8123](http://localhost:8123/)
|
||||
>
|
||||
> - Docs: http://localhost:8123/docs
|
||||
>
|
||||
> - LangGraph Studio Web UI: https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:8123
|
||||
|
||||
|
||||
!!! note "In-Memory Mode"
|
||||
|
||||
The `langgraph dev` command starts LangGraph Server in an in-memory mode. This mode is suitable for development and testing purposes. For production use, you should deploy LangGraph Server with access to a persistent storage backend.
|
||||
|
||||
If you want to test your application with a persistent storage backend, you can use the `langgraph up` command instead of `langgraph dev`. You will
|
||||
need to have `docker` installed on your machine to use this command.
|
||||
|
||||
## LangGraph Studio Web UI
|
||||
|
||||
Test your graph in the LangGraph Studio Web UI by visiting the URL provided in the output of the `langgraph up` command.
|
||||
|
||||
> - LangGraph Studio Web UI: https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:8123
|
||||
|
||||
!!! warning "Safari Compatibility"
|
||||
|
||||
Currently, LangGraph Studio Web does not support Safari when running a server locally.
|
||||
|
||||
## Test the API
|
||||
|
||||
=== "Python SDK (Async)"
|
||||
|
||||
**Install the LangGraph Python SDK**
|
||||
|
||||
```shell
|
||||
pip install langgraph-sdk
|
||||
```
|
||||
|
||||
**Send a message to the assistant (threadless run)**
|
||||
|
||||
```python
|
||||
from langgraph_sdk import get_client
|
||||
|
||||
client = get_client(url="http://localhost:8123")
|
||||
|
||||
async for chunk in client.runs.stream(
|
||||
None, # Threadless run
|
||||
"agent", # Name of assistant. Defined in langgraph.json.
|
||||
input={
|
||||
"messages": [{
|
||||
"role": "human",
|
||||
"content": "What is LangGraph?",
|
||||
}],
|
||||
},
|
||||
stream_mode="updates",
|
||||
):
|
||||
print(f"Receiving new event of type: {chunk.event}...")
|
||||
print(chunk.data)
|
||||
print("\n\n")
|
||||
```
|
||||
|
||||
=== "Python SDK (Sync)"
|
||||
|
||||
**Install the LangGraph Python SDK**
|
||||
|
||||
```shell
|
||||
pip install langgraph-sdk
|
||||
```
|
||||
|
||||
**Send a message to the assistant (threadless run)**
|
||||
|
||||
```python
|
||||
from langgraph_sdk import get_sync_client
|
||||
|
||||
client = get_sync_client(url="http://localhost:8123")
|
||||
|
||||
for chunk in client.runs.stream(
|
||||
None, # Threadless run
|
||||
"agent", # Name of assistant. Defined in langgraph.json.
|
||||
input={
|
||||
"messages": [{
|
||||
"role": "human",
|
||||
"content": "What is LangGraph?",
|
||||
}],
|
||||
},
|
||||
stream_mode="updates",
|
||||
):
|
||||
print(f"Receiving new event of type: {chunk.event}...")
|
||||
print(chunk.data)
|
||||
print("\n\n")
|
||||
```
|
||||
|
||||
=== "Javascript SDK"
|
||||
|
||||
**Install the LangGraph JS SDK**
|
||||
|
||||
```shell
|
||||
npm install @langchain/langgraph-sdk
|
||||
```
|
||||
|
||||
**Send a message to the assistant (threadless run)**
|
||||
|
||||
```js
|
||||
const { Client } = await import("@langchain/langgraph-sdk");
|
||||
|
||||
// only set the apiUrl if you changed the default port when calling langgraph up
|
||||
const client = new Client({ apiUrl: "http://localhost:8123"});
|
||||
|
||||
const streamResponse = client.runs.stream(
|
||||
null, // Threadless run
|
||||
"agent", // Assistant ID
|
||||
{
|
||||
input: {
|
||||
"messages": [
|
||||
{ "role": "user", "content": "What is LangGraph?"}
|
||||
]
|
||||
},
|
||||
streamMode: "messages",
|
||||
}
|
||||
);
|
||||
|
||||
for await (const chunk of streamResponse) {
|
||||
console.log(`Receiving new event of type: ${chunk.event}...`);
|
||||
console.log(JSON.stringify(chunk.data));
|
||||
console.log("\n\n");
|
||||
}
|
||||
```
|
||||
|
||||
=== "Rest API"
|
||||
|
||||
```bash
|
||||
curl -s --request POST \
|
||||
--url "http://localhost:8123/runs/stream" \
|
||||
--header 'Content-Type: application/json' \
|
||||
--data "{
|
||||
\"assistant_id\": \"agent\",
|
||||
\"input\": {
|
||||
\"messages\": [
|
||||
{
|
||||
\"role\": \"human\",
|
||||
\"content\": \"What is LangGraph?\"
|
||||
}
|
||||
]
|
||||
},
|
||||
\"stream_mode\": \"updates\"
|
||||
}"
|
||||
```
|
||||
|
||||
!!! tip "Auth"
|
||||
|
||||
If you're connecting to a remote server, you will need to provide a LangSmith
|
||||
API Key for authorization. Please see the API Reference for the clients
|
||||
for more information.
|
||||
|
||||
## Next Steps
|
||||
|
||||
Now that you have a LangGraph app running locally, take your journey further by exploring deployment and advanced features:
|
||||
|
||||
### 🌐 Deploy to LangGraph Cloud
|
||||
|
||||
- **[LangGraph Cloud QuickStart](../../cloud/quick_start.md)**: Deploy your LangGraph app using LangGraph Cloud.
|
||||
|
||||
### 📚 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.
|
||||
|
||||
### 🛠️ Developer References
|
||||
|
||||
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.
|
||||
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
@@ -934,7 +934,7 @@
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.9"
|
||||
"version": "3.11.4"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
|
||||
@@ -112,7 +112,7 @@
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"_set_env(\"LANGCHAIN_API_KEY\")\n",
|
||||
"_set_env(\"LANGSMITH_API_KEY\")\n",
|
||||
"os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
|
||||
"os.environ[\"LANGCHAIN_PROJECT\"] = \"local-llama32-rag\""
|
||||
]
|
||||
|
||||
+2
-1
@@ -94,6 +94,7 @@ nav:
|
||||
- Quick Start:
|
||||
- Quick Start: tutorials#quick-start
|
||||
- tutorials/introduction.ipynb
|
||||
- tutorials/langgraph-platform/local-server.md
|
||||
- cloud/quick_start.md
|
||||
- Chatbots:
|
||||
- Chatbots: tutorials#chatbots
|
||||
@@ -438,4 +439,4 @@ validation:
|
||||
# and those anchors are not available in the actual doc
|
||||
anchors: info
|
||||
# this is needed to handle headers with anchors for nav
|
||||
not_found: info
|
||||
not_found: info
|
||||
|
||||
@@ -23,6 +23,7 @@ from langgraph.store.base import (
|
||||
Op,
|
||||
PutOp,
|
||||
Result,
|
||||
SearchItem,
|
||||
SearchOp,
|
||||
)
|
||||
|
||||
@@ -283,7 +284,7 @@ class DuckDBStore(BaseStore, BaseDuckDBStore[duckdb.DuckDBPyConnection]):
|
||||
|
||||
for cur, idx in cursors:
|
||||
rows = cur.fetchall()
|
||||
items = [_row_to_item(_convert_ns(row[0]), row) for row in rows]
|
||||
items = [_row_to_search_item(_convert_ns(row[0]), row) for row in rows]
|
||||
results[idx] = items
|
||||
|
||||
def _batch_list_namespaces_ops(
|
||||
@@ -376,6 +377,22 @@ def _row_to_item(
|
||||
)
|
||||
|
||||
|
||||
def _row_to_search_item(
|
||||
namespace: tuple[str, ...],
|
||||
row: tuple,
|
||||
) -> SearchItem:
|
||||
"""Convert a row from the database into an SearchItem."""
|
||||
# TODO: Add support for search
|
||||
_, key, val, created_at, updated_at = row
|
||||
return SearchItem(
|
||||
value=val if isinstance(val, dict) else json.loads(val),
|
||||
key=key,
|
||||
namespace=namespace,
|
||||
created_at=created_at,
|
||||
updated_at=updated_at,
|
||||
)
|
||||
|
||||
|
||||
def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int]:
|
||||
grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list)
|
||||
tot = 0
|
||||
|
||||
@@ -5,7 +5,11 @@
|
||||
######################
|
||||
|
||||
start-postgres:
|
||||
POSTGRES_VERSION=${POSTGRES_VERSION:-16} docker compose -f tests/compose-postgres.yml up -V --force-recreate --wait
|
||||
POSTGRES_VERSION=${POSTGRES_VERSION:-16} docker compose -f tests/compose-postgres.yml up -V --force-recreate --wait || ( \
|
||||
echo "Failed to start PostgreSQL, printing logs..."; \
|
||||
docker compose -f tests/compose-postgres.yml logs; \
|
||||
exit 1 \
|
||||
)
|
||||
|
||||
stop-postgres:
|
||||
docker compose -f tests/compose-postgres.yml down
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import threading
|
||||
from collections.abc import Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Iterator, Optional, Sequence, Union
|
||||
from typing import Any, Optional
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from psycopg import Capabilities, Connection, Cursor, Pipeline
|
||||
@@ -17,21 +18,11 @@ from langgraph.checkpoint.base import (
|
||||
CheckpointTuple,
|
||||
get_checkpoint_id,
|
||||
)
|
||||
from langgraph.checkpoint.postgres import _internal
|
||||
from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
|
||||
Conn = Union[Connection[DictRow], ConnectionPool[Connection[DictRow]]]
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _get_connection(conn: Conn) -> Iterator[Connection[DictRow]]:
|
||||
if isinstance(conn, Connection):
|
||||
yield conn
|
||||
elif isinstance(conn, ConnectionPool):
|
||||
with conn.connection() as conn:
|
||||
yield conn
|
||||
else:
|
||||
raise TypeError(f"Invalid connection type: {type(conn)}")
|
||||
Conn = _internal.Conn # For backward compatibility
|
||||
|
||||
|
||||
class PostgresSaver(BasePostgresSaver):
|
||||
@@ -39,7 +30,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: Conn,
|
||||
conn: _internal.Conn,
|
||||
pipe: Optional[Pipeline] = None,
|
||||
serde: Optional[SerializerProtocol] = None,
|
||||
) -> None:
|
||||
@@ -73,9 +64,9 @@ class PostgresSaver(BasePostgresSaver):
|
||||
) as conn:
|
||||
if pipeline:
|
||||
with conn.pipeline() as pipe:
|
||||
yield PostgresSaver(conn, pipe)
|
||||
yield cls(conn, pipe)
|
||||
else:
|
||||
yield PostgresSaver(conn)
|
||||
yield cls(conn)
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Set up the checkpoint database asynchronously.
|
||||
@@ -373,7 +364,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
Will be applied regardless of whether the PostgresSaver instance was initialized with a pipeline.
|
||||
If pipeline mode is not supported, will fall back to using transaction context manager.
|
||||
"""
|
||||
with _get_connection(self.conn) as conn:
|
||||
with _internal.get_connection(self.conn) as conn:
|
||||
if self.pipe:
|
||||
# a connection in pipeline mode can be used concurrently
|
||||
# in multiple threads/coroutines, but only one cursor can be
|
||||
@@ -388,19 +379,23 @@ class PostgresSaver(BasePostgresSaver):
|
||||
# a connection not in pipeline mode can only be used by one
|
||||
# thread/coroutine at a time, so we acquire a lock
|
||||
if self.supports_pipeline:
|
||||
with self.lock, conn.pipeline(), conn.cursor(
|
||||
binary=True, row_factory=dict_row
|
||||
) as cur:
|
||||
with (
|
||||
self.lock,
|
||||
conn.pipeline(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
# Use connection's transaction context manager when pipeline mode not supported
|
||||
with self.lock, conn.transaction(), conn.cursor(
|
||||
binary=True, row_factory=dict_row
|
||||
) as cur:
|
||||
with (
|
||||
self.lock,
|
||||
conn.transaction(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
|
||||
__all__ = ["PostgresSaver", "Conn"]
|
||||
__all__ = ["PostgresSaver", "BasePostgresSaver", "Conn"]
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
"""Shared async utility functions for the Postgres checkpoint & storage classes."""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Union
|
||||
|
||||
from psycopg import AsyncConnection
|
||||
from psycopg.rows import DictRow
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
|
||||
Conn = Union[AsyncConnection[DictRow], AsyncConnectionPool[AsyncConnection[DictRow]]]
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def get_connection(
|
||||
conn: Conn,
|
||||
) -> AsyncIterator[AsyncConnection[DictRow]]:
|
||||
if isinstance(conn, AsyncConnection):
|
||||
yield conn
|
||||
elif isinstance(conn, AsyncConnectionPool):
|
||||
async with conn.connection() as conn:
|
||||
yield conn
|
||||
else:
|
||||
raise TypeError(f"Invalid connection type: {type(conn)}")
|
||||
@@ -0,0 +1,22 @@
|
||||
"""Shared utility functions for the Postgres checkpoint & storage classes."""
|
||||
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from typing import Union
|
||||
|
||||
from psycopg import Connection
|
||||
from psycopg.rows import DictRow
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
Conn = Union[Connection[DictRow], ConnectionPool[Connection[DictRow]]]
|
||||
|
||||
|
||||
@contextmanager
|
||||
def get_connection(conn: Conn) -> Iterator[Connection[DictRow]]:
|
||||
if isinstance(conn, Connection):
|
||||
yield conn
|
||||
elif isinstance(conn, ConnectionPool):
|
||||
with conn.connection() as conn:
|
||||
yield conn
|
||||
else:
|
||||
raise TypeError(f"Invalid connection type: {type(conn)}")
|
||||
@@ -1,6 +1,7 @@
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, AsyncIterator, Iterator, Optional, Sequence, Union
|
||||
from typing import Any, Optional
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
|
||||
@@ -17,23 +18,11 @@ from langgraph.checkpoint.base import (
|
||||
CheckpointTuple,
|
||||
get_checkpoint_id,
|
||||
)
|
||||
from langgraph.checkpoint.postgres import _ainternal
|
||||
from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
|
||||
Conn = Union[AsyncConnection[DictRow], AsyncConnectionPool[AsyncConnection[DictRow]]]
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _get_connection(
|
||||
conn: Conn,
|
||||
) -> AsyncIterator[AsyncConnection[DictRow]]:
|
||||
if isinstance(conn, AsyncConnection):
|
||||
yield conn
|
||||
elif isinstance(conn, AsyncConnectionPool):
|
||||
async with conn.connection() as conn:
|
||||
yield conn
|
||||
else:
|
||||
raise TypeError(f"Invalid connection type: {type(conn)}")
|
||||
Conn = _ainternal.Conn # For backward compatibility
|
||||
|
||||
|
||||
class AsyncPostgresSaver(BasePostgresSaver):
|
||||
@@ -41,7 +30,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: Conn,
|
||||
conn: _ainternal.Conn,
|
||||
pipe: Optional[AsyncPipeline] = None,
|
||||
serde: Optional[SerializerProtocol] = None,
|
||||
) -> None:
|
||||
@@ -80,9 +69,9 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
) as conn:
|
||||
if pipeline:
|
||||
async with conn.pipeline() as pipe:
|
||||
yield AsyncPostgresSaver(conn=conn, pipe=pipe, serde=serde)
|
||||
yield cls(conn=conn, pipe=pipe, serde=serde)
|
||||
else:
|
||||
yield AsyncPostgresSaver(conn=conn, serde=serde)
|
||||
yield cls(conn=conn, serde=serde)
|
||||
|
||||
async def setup(self) -> None:
|
||||
"""Set up the checkpoint database asynchronously.
|
||||
@@ -157,15 +146,17 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
value["pending_sends"],
|
||||
),
|
||||
self._load_metadata(value["metadata"]),
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
}
|
||||
}
|
||||
}
|
||||
if value["parent_checkpoint_id"]
|
||||
else None,
|
||||
if value["parent_checkpoint_id"]
|
||||
else None
|
||||
),
|
||||
await asyncio.to_thread(self._load_writes, value["pending_writes"]),
|
||||
)
|
||||
|
||||
@@ -216,15 +207,17 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
value["pending_sends"],
|
||||
),
|
||||
self._load_metadata(value["metadata"]),
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
}
|
||||
}
|
||||
}
|
||||
if value["parent_checkpoint_id"]
|
||||
else None,
|
||||
if value["parent_checkpoint_id"]
|
||||
else None
|
||||
),
|
||||
await asyncio.to_thread(self._load_writes, value["pending_writes"]),
|
||||
)
|
||||
|
||||
@@ -331,7 +324,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
Will be applied regardless of whether the AsyncPostgresSaver instance was initialized with a pipeline.
|
||||
If pipeline mode is not supported, will fall back to using transaction context manager.
|
||||
"""
|
||||
async with _get_connection(self.conn) as conn:
|
||||
async with _ainternal.get_connection(self.conn) as conn:
|
||||
if self.pipe:
|
||||
# a connection in pipeline mode can be used concurrently
|
||||
# in multiple threads/coroutines, but only one cursor can be
|
||||
@@ -346,20 +339,25 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
# a connection not in pipeline mode can only be used by one
|
||||
# thread/coroutine at a time, so we acquire a lock
|
||||
if self.supports_pipeline:
|
||||
async with self.lock, conn.pipeline(), conn.cursor(
|
||||
binary=True, row_factory=dict_row
|
||||
) as cur:
|
||||
async with (
|
||||
self.lock,
|
||||
conn.pipeline(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
# Use connection's transaction context manager when pipeline mode not supported
|
||||
async with self.lock, conn.transaction(), conn.cursor(
|
||||
binary=True, row_factory=dict_row
|
||||
) as cur:
|
||||
async with (
|
||||
self.lock,
|
||||
conn.transaction(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
async with self.lock, conn.cursor(
|
||||
binary=True, row_factory=dict_row
|
||||
) as cur:
|
||||
async with (
|
||||
self.lock,
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
|
||||
def list(
|
||||
@@ -388,7 +386,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
while True:
|
||||
try:
|
||||
yield asyncio.run_coroutine_threadsafe(
|
||||
anext(aiter_),
|
||||
anext(aiter_), # noqa: F821
|
||||
self.loop,
|
||||
).result()
|
||||
except StopAsyncIteration:
|
||||
@@ -467,3 +465,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.aput_writes(config, writes, task_id), self.loop
|
||||
).result()
|
||||
|
||||
|
||||
__all__ = ["AsyncPostgresSaver", "Conn"]
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import random
|
||||
from typing import Any, List, Optional, Sequence, Tuple, cast
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, Optional, cast
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from psycopg.types.json import Jsonb
|
||||
@@ -249,7 +250,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
config: Optional[RunnableConfig],
|
||||
filter: MetadataInput,
|
||||
before: Optional[RunnableConfig] = None,
|
||||
) -> Tuple[str, List[Any]]:
|
||||
) -> tuple[str, list[Any]]:
|
||||
"""Return WHERE clause predicates for alist() given config, filter, before.
|
||||
|
||||
This method returns a tuple of a string and a tuple of values. The string
|
||||
|
||||
@@ -1,110 +1,241 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import AsyncIterator, Iterable, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncIterator,
|
||||
Callable,
|
||||
Iterable,
|
||||
Optional,
|
||||
Sequence,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
from typing import Any, Callable, Optional, Union, cast
|
||||
|
||||
import orjson
|
||||
from psycopg import AsyncConnection, AsyncCursor
|
||||
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
|
||||
from psycopg.errors import UndefinedTable
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
|
||||
from langgraph.store.base import GetOp, ListNamespacesOp, Op, PutOp, Result, SearchOp
|
||||
from langgraph.checkpoint.postgres import _ainternal
|
||||
from langgraph.store.base import (
|
||||
GetOp,
|
||||
ListNamespacesOp,
|
||||
Op,
|
||||
PutOp,
|
||||
Result,
|
||||
SearchOp,
|
||||
)
|
||||
from langgraph.store.base.batch import AsyncBatchedBaseStore
|
||||
from langgraph.store.postgres.base import (
|
||||
_PLACEHOLDER,
|
||||
BasePostgresStore,
|
||||
PoolConfig,
|
||||
PostgresIndexConfig,
|
||||
Row,
|
||||
_decode_ns_bytes,
|
||||
_ensure_index_config,
|
||||
_group_ops,
|
||||
_row_to_item,
|
||||
_row_to_search_item,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnection]):
|
||||
__slots__ = ("_deserializer",)
|
||||
class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
|
||||
__slots__ = (
|
||||
"_deserializer",
|
||||
"pipe",
|
||||
"lock",
|
||||
"supports_pipeline",
|
||||
"index_config",
|
||||
"embeddings",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: AsyncConnection[Any],
|
||||
conn: _ainternal.Conn,
|
||||
*,
|
||||
pipe: Optional[AsyncPipeline] = None,
|
||||
deserializer: Optional[
|
||||
Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
|
||||
] = None,
|
||||
index: Optional[PostgresIndexConfig] = None,
|
||||
) -> None:
|
||||
if isinstance(conn, AsyncConnectionPool) and pipe is not None:
|
||||
raise ValueError(
|
||||
"Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool."
|
||||
)
|
||||
super().__init__()
|
||||
self._deserializer = deserializer
|
||||
self.conn = conn
|
||||
self.pipe = pipe
|
||||
self.lock = asyncio.Lock()
|
||||
self.loop = asyncio.get_running_loop()
|
||||
self.supports_pipeline = Capabilities().has_pipeline()
|
||||
self.index_config = index
|
||||
if self.index_config:
|
||||
self.embeddings, self.index_config = _ensure_index_config(self.index_config)
|
||||
|
||||
else:
|
||||
self.embeddings = None
|
||||
|
||||
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
grouped_ops, num_ops = _group_ops(ops)
|
||||
results: list[Result] = [None] * num_ops
|
||||
|
||||
async with self.conn.pipeline():
|
||||
tasks = []
|
||||
|
||||
if GetOp in grouped_ops:
|
||||
tasks.append(
|
||||
self._batch_get_ops(
|
||||
cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results
|
||||
)
|
||||
)
|
||||
|
||||
if PutOp in grouped_ops:
|
||||
tasks.append(
|
||||
self._batch_put_ops(
|
||||
cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp])
|
||||
)
|
||||
)
|
||||
|
||||
if SearchOp in grouped_ops:
|
||||
tasks.append(
|
||||
self._batch_search_ops(
|
||||
cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
|
||||
results,
|
||||
)
|
||||
)
|
||||
|
||||
if ListNamespacesOp in grouped_ops:
|
||||
tasks.append(
|
||||
self._batch_list_namespaces_ops(
|
||||
cast(
|
||||
Sequence[tuple[int, ListNamespacesOp]],
|
||||
grouped_ops[ListNamespacesOp],
|
||||
),
|
||||
results,
|
||||
)
|
||||
)
|
||||
|
||||
await asyncio.gather(*tasks)
|
||||
async with _ainternal.get_connection(self.conn) as conn:
|
||||
if self.pipe:
|
||||
async with self.pipe:
|
||||
await self._execute_batch(grouped_ops, results, conn)
|
||||
else:
|
||||
await self._execute_batch(grouped_ops, results, conn)
|
||||
|
||||
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(
|
||||
cls,
|
||||
conn_string: str,
|
||||
*,
|
||||
pipeline: bool = False,
|
||||
pool_config: Optional[PoolConfig] = None,
|
||||
index: Optional[PostgresIndexConfig] = None,
|
||||
) -> AsyncIterator["AsyncPostgresStore"]:
|
||||
"""Create a new AsyncPostgresStore instance from a connection string.
|
||||
|
||||
Args:
|
||||
conn_string (str): The Postgres connection info string.
|
||||
pipeline (bool): Whether to use AsyncPipeline (only for single connections)
|
||||
pool_config (Optional[PoolConfig]): 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 embedding config.
|
||||
|
||||
Returns:
|
||||
AsyncPostgresStore: A new AsyncPostgresStore instance.
|
||||
"""
|
||||
if pool_config is not None:
|
||||
pc = pool_config.copy()
|
||||
async with cast(
|
||||
AsyncConnectionPool[AsyncConnection[DictRow]],
|
||||
AsyncConnectionPool(
|
||||
conn_string,
|
||||
min_size=pc.pop("min_size", 1),
|
||||
max_size=pc.pop("max_size", None),
|
||||
kwargs={
|
||||
"autocommit": True,
|
||||
"prepare_threshold": 0,
|
||||
"row_factory": dict_row,
|
||||
**(pc.pop("kwargs", None) or {}),
|
||||
},
|
||||
**cast(dict, pc),
|
||||
),
|
||||
) as pool:
|
||||
yield cls(conn=pool, index=index)
|
||||
else:
|
||||
async with await AsyncConnection.connect(
|
||||
conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
|
||||
) as conn:
|
||||
if pipeline:
|
||||
async with conn.pipeline() as pipe:
|
||||
yield cls(conn=conn, pipe=pipe, index=index)
|
||||
else:
|
||||
yield cls(conn=conn, index=index)
|
||||
|
||||
async def setup(self) -> None:
|
||||
"""Set up the store database asynchronously.
|
||||
|
||||
This method creates the necessary tables in the Postgres database if they don't
|
||||
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:
|
||||
try:
|
||||
await cur.execute(f"SELECT v FROM {table} 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(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {table} (
|
||||
v INTEGER PRIMARY KEY
|
||||
)
|
||||
"""
|
||||
)
|
||||
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,))
|
||||
|
||||
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 = {
|
||||
k: v(self) if v is not None and callable(v) else v
|
||||
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,)
|
||||
)
|
||||
|
||||
async def _execute_batch(
|
||||
self,
|
||||
grouped_ops: dict,
|
||||
results: list[Result],
|
||||
conn: AsyncConnection[DictRow],
|
||||
) -> None:
|
||||
async with self._cursor(pipeline=True) as cur:
|
||||
if GetOp in grouped_ops:
|
||||
await self._batch_get_ops(
|
||||
cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]),
|
||||
results,
|
||||
cur,
|
||||
)
|
||||
|
||||
if SearchOp in grouped_ops:
|
||||
await self._batch_search_ops(
|
||||
cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
|
||||
results,
|
||||
cur,
|
||||
)
|
||||
|
||||
if ListNamespacesOp in grouped_ops:
|
||||
await self._batch_list_namespaces_ops(
|
||||
cast(
|
||||
Sequence[tuple[int, ListNamespacesOp]],
|
||||
grouped_ops[ListNamespacesOp],
|
||||
),
|
||||
results,
|
||||
cur,
|
||||
)
|
||||
|
||||
if PutOp in grouped_ops:
|
||||
await self._batch_put_ops(
|
||||
cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]),
|
||||
cur,
|
||||
)
|
||||
|
||||
async def _batch_get_ops(
|
||||
self,
|
||||
get_ops: Sequence[tuple[int, GetOp]],
|
||||
results: list[Result],
|
||||
cur: AsyncCursor[DictRow],
|
||||
) -> None:
|
||||
cursors = []
|
||||
for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
|
||||
cur = self.conn.cursor(binary=True)
|
||||
await cur.execute(query, params)
|
||||
cursors.append((cur, namespace, items))
|
||||
|
||||
for cur, namespace, items in cursors:
|
||||
rows = cast(list[Row], await cur.fetchall())
|
||||
key_to_row = {row["key"]: row for row in rows}
|
||||
for idx, key in items:
|
||||
@@ -119,29 +250,59 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnectio
|
||||
async def _batch_put_ops(
|
||||
self,
|
||||
put_ops: Sequence[tuple[int, PutOp]],
|
||||
cur: AsyncCursor[DictRow],
|
||||
) -> None:
|
||||
queries = self._get_batch_PUT_queries(put_ops)
|
||||
queries, embedding_request = self._prepare_batch_PUT_queries(put_ops)
|
||||
if embedding_request:
|
||||
if self.embeddings is None:
|
||||
# Should not get here since the embedding config is required
|
||||
# to return an embedding_request above
|
||||
raise ValueError(
|
||||
"Embedding configuration is required for vector operations "
|
||||
f"(for semantic search). "
|
||||
f"Please provide an EmbeddingConfig when initializing the {self.__class__.__name__}."
|
||||
)
|
||||
query, txt_params = embedding_request
|
||||
vectors = await self.embeddings.aembed_documents(
|
||||
[param[-1] for param in txt_params]
|
||||
)
|
||||
queries.append(
|
||||
(
|
||||
query,
|
||||
[
|
||||
p
|
||||
for (ns, k, pathname, _), vector in zip(txt_params, vectors)
|
||||
for p in (ns, k, pathname, vector)
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
for query, params in queries:
|
||||
cur = self.conn.cursor(binary=True)
|
||||
await cur.execute(query, params)
|
||||
|
||||
async def _batch_search_ops(
|
||||
self,
|
||||
search_ops: Sequence[tuple[int, SearchOp]],
|
||||
results: list[Result],
|
||||
cur: AsyncCursor[DictRow],
|
||||
) -> None:
|
||||
queries = self._get_batch_search_queries(search_ops)
|
||||
cursors: list[tuple[AsyncCursor[Any], int]] = []
|
||||
queries, embedding_requests = self._prepare_batch_search_queries(search_ops)
|
||||
|
||||
for (query, params), (idx, _) in zip(queries, search_ops):
|
||||
cur = self.conn.cursor(binary=True)
|
||||
if embedding_requests and self.embeddings:
|
||||
vectors = await self.embeddings.aembed_documents(
|
||||
[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
|
||||
|
||||
for (idx, _), (query, params) in zip(search_ops, queries):
|
||||
await cur.execute(query, params)
|
||||
cursors.append((cur, idx))
|
||||
|
||||
for cur, idx in cursors:
|
||||
rows = cast(list[Row], await cur.fetchall())
|
||||
items = [
|
||||
_row_to_item(
|
||||
_row_to_search_item(
|
||||
_decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
|
||||
)
|
||||
for row in rows
|
||||
@@ -152,67 +313,57 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnectio
|
||||
self,
|
||||
list_ops: Sequence[tuple[int, ListNamespacesOp]],
|
||||
results: list[Result],
|
||||
cur: AsyncCursor[DictRow],
|
||||
) -> None:
|
||||
queries = self._get_batch_list_namespaces_queries(list_ops)
|
||||
cursors: list[tuple[AsyncCursor[Any], int]] = []
|
||||
for (query, params), (idx, _) in zip(queries, list_ops):
|
||||
cur = self.conn.cursor(binary=True)
|
||||
await cur.execute(query, params)
|
||||
cursors.append((cur, idx))
|
||||
|
||||
for cur, idx in cursors:
|
||||
rows = cast(list[dict], await cur.fetchall())
|
||||
namespaces = [_decode_ns_bytes(row["truncated_prefix"]) for row in rows]
|
||||
results[idx] = namespaces
|
||||
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
async def from_conn_string(
|
||||
cls,
|
||||
conn_string: str,
|
||||
) -> AsyncIterator["AsyncPostgresStore"]:
|
||||
"""Create a new AsyncPostgresStore instance from a connection string.
|
||||
async def _cursor(
|
||||
self, *, pipeline: bool = False
|
||||
) -> AsyncIterator[AsyncCursor[DictRow]]:
|
||||
"""Create a database cursor as a context manager.
|
||||
|
||||
Args:
|
||||
conn_string (str): The Postgres connection info string.
|
||||
|
||||
Returns:
|
||||
AsyncPostgresStore: A new AsyncPostgresStore instance.
|
||||
pipeline: whether to use pipeline for the DB operations inside the context manager.
|
||||
Will be applied regardless of whether the PostgresStore instance was initialized with a pipeline.
|
||||
If pipeline mode is not supported, will fall back to using transaction context manager.
|
||||
"""
|
||||
async with await AsyncConnection.connect(
|
||||
conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
|
||||
) as conn:
|
||||
yield cls(conn=conn)
|
||||
|
||||
async def setup(self) -> None:
|
||||
"""Set up the store database asynchronously.
|
||||
|
||||
This method creates the necessary tables in the Postgres database if they don't
|
||||
already exist and runs database migrations. It MUST be called directly by the user
|
||||
the first time the store is used.
|
||||
"""
|
||||
async with self.conn.cursor() as cur:
|
||||
try:
|
||||
await cur.execute(
|
||||
"SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
|
||||
)
|
||||
row = cast(dict, await cur.fetchone())
|
||||
if row is None:
|
||||
version = -1
|
||||
async with _ainternal.get_connection(self.conn) as conn:
|
||||
if self.pipe:
|
||||
# a connection in pipeline mode can be used concurrently
|
||||
# in multiple threads/coroutines, but only one cursor can be
|
||||
# used at a time
|
||||
try:
|
||||
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
finally:
|
||||
if pipeline:
|
||||
await self.pipe.sync()
|
||||
elif pipeline:
|
||||
# a connection not in pipeline mode can only be used by one
|
||||
# thread/coroutine at a time, so we acquire a lock
|
||||
if self.supports_pipeline:
|
||||
async with (
|
||||
self.lock,
|
||||
conn.pipeline(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
version = row["v"]
|
||||
except UndefinedTable:
|
||||
version = -1
|
||||
# Create store_migrations table if it doesn't exist
|
||||
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
|
||||
):
|
||||
await cur.execute(migration)
|
||||
await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
|
||||
async with (
|
||||
self.lock,
|
||||
conn.transaction(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
async with (
|
||||
self.lock,
|
||||
conn.cursor(binary=True) as cur,
|
||||
):
|
||||
yield cur
|
||||
|
||||
@@ -1,44 +1,64 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
from collections import defaultdict
|
||||
from collections.abc import Iterable, Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Callable,
|
||||
Generic,
|
||||
Iterable,
|
||||
Iterator,
|
||||
Literal,
|
||||
NamedTuple,
|
||||
Optional,
|
||||
Sequence,
|
||||
TypeVar,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
|
||||
import orjson
|
||||
from psycopg import BaseConnection, Connection, Cursor
|
||||
from psycopg import Capabilities, Connection, Cursor, Pipeline
|
||||
from psycopg.errors import UndefinedTable
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import ConnectionPool
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.checkpoint.postgres import _ainternal as _ainternal
|
||||
from langgraph.checkpoint.postgres import _internal as _pg_internal
|
||||
from langgraph.store.base import (
|
||||
BaseStore,
|
||||
GetOp,
|
||||
IndexConfig,
|
||||
Item,
|
||||
ListNamespacesOp,
|
||||
Op,
|
||||
PutOp,
|
||||
Result,
|
||||
SearchItem,
|
||||
SearchOp,
|
||||
ensure_embeddings,
|
||||
get_text_at_path,
|
||||
tokenize_path,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain_core.embeddings import Embeddings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
MIGRATIONS = [
|
||||
class Migration(NamedTuple):
|
||||
"""A database migration with optional conditions and parameters."""
|
||||
|
||||
sql: str
|
||||
params: Optional[dict[str, Any]] = None
|
||||
|
||||
|
||||
MIGRATIONS: Sequence[str] = [
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS store (
|
||||
-- 'prefix' represents the doc's 'namespace'
|
||||
@@ -56,13 +76,100 @@ CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pa
|
||||
""",
|
||||
]
|
||||
|
||||
C = TypeVar("C", bound=BaseConnection)
|
||||
VECTOR_MIGRATIONS: Sequence[Migration] = [
|
||||
Migration(
|
||||
"""
|
||||
CREATE EXTENSION IF NOT EXISTS vector;
|
||||
""",
|
||||
),
|
||||
Migration(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS store_vectors (
|
||||
prefix text NOT NULL,
|
||||
key text NOT NULL,
|
||||
field_name text NOT NULL,
|
||||
embedding %(vector_type)s(%(dims)s),
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
||||
PRIMARY KEY (prefix, key, field_name),
|
||||
FOREIGN KEY (prefix, key) REFERENCES store(prefix, key) ON DELETE CASCADE
|
||||
);
|
||||
""",
|
||||
params={
|
||||
"dims": lambda store: store.index_config["dims"],
|
||||
"vector_type": lambda store: (
|
||||
cast(PostgresIndexConfig, store.index_config)
|
||||
.get("ann_index_config", {})
|
||||
.get("vector_type", "vector")
|
||||
),
|
||||
},
|
||||
),
|
||||
# TODO: Add an HNSW or IVFFlat index depending on config
|
||||
# First must improve the search query when filtering by
|
||||
# namespace
|
||||
]
|
||||
|
||||
C = TypeVar("C", bound=Union[_pg_internal.Conn, _ainternal.Conn])
|
||||
|
||||
|
||||
class PoolConfig(TypedDict, total=False):
|
||||
"""Connection pool settings for PostgreSQL connections.
|
||||
|
||||
Controls connection lifecycle and resource utilization:
|
||||
- Small pools (1-5) suit low-concurrency workloads
|
||||
- Larger pools handle concurrent requests but consume more resources
|
||||
- Setting max_size prevents resource exhaustion under load
|
||||
"""
|
||||
|
||||
min_size: int
|
||||
"""Minimum number of connections maintained in the pool. Defaults to 1."""
|
||||
|
||||
max_size: Optional[int]
|
||||
"""Maximum number of connections allowed in the pool. None means unlimited."""
|
||||
|
||||
kwargs: dict
|
||||
"""Additional connection arguments passed to each connection in the pool.
|
||||
|
||||
Default kwargs set automatically:
|
||||
- autocommit: True
|
||||
- prepare_threshold: 0
|
||||
- row_factory: dict_row
|
||||
"""
|
||||
|
||||
|
||||
class ANNIndexConfig(TypedDict, total=False):
|
||||
"""Configuration for vector index in PostgreSQL store."""
|
||||
|
||||
vector_type: Literal["vector", "halfvec"]
|
||||
"""Type of vector storage to use.
|
||||
Options:
|
||||
- 'vector': Regular vectors (default)
|
||||
- 'halfvec': Half-precision vectors for reduced memory usage
|
||||
"""
|
||||
|
||||
|
||||
class PostgresIndexConfig(IndexConfig, total=False):
|
||||
"""Configuration for vector embeddings in PostgreSQL store with pgvector-specific options.
|
||||
|
||||
Extends EmbeddingConfig with additional configuration for pgvector index and vector types.
|
||||
"""
|
||||
|
||||
ann_index_config: ANNIndexConfig
|
||||
"""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:
|
||||
- 'l2': Euclidean distance
|
||||
- 'inner_product': Dot product
|
||||
- 'cosine': Cosine similarity
|
||||
"""
|
||||
|
||||
|
||||
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]
|
||||
|
||||
def _get_batch_GET_ops_queries(
|
||||
self,
|
||||
@@ -84,13 +191,21 @@ class BasePostgresStore(Generic[C]):
|
||||
results.append((query, params, namespace, items))
|
||||
return results
|
||||
|
||||
def _get_batch_PUT_queries(
|
||||
def _prepare_batch_PUT_queries(
|
||||
self,
|
||||
put_ops: Sequence[tuple[int, PutOp]],
|
||||
) -> list[tuple[str, Sequence]]:
|
||||
) -> tuple[
|
||||
list[tuple[str, Sequence]],
|
||||
Optional[tuple[str, Sequence[tuple[str, str, str, str]]]],
|
||||
]:
|
||||
# Last-write wins
|
||||
dedupped_ops: dict[tuple[tuple[str, ...], str], PutOp] = {}
|
||||
for _, op in put_ops:
|
||||
dedupped_ops[(op.namespace, op.key)] = op
|
||||
|
||||
inserts: list[PutOp] = []
|
||||
deletes: list[PutOp] = []
|
||||
for _, op in put_ops:
|
||||
for op in dedupped_ops.values():
|
||||
if op.value is None:
|
||||
deletes.append(op)
|
||||
else:
|
||||
@@ -109,60 +224,182 @@ class BasePostgresStore(Generic[C]):
|
||||
)
|
||||
params = (_namespace_to_text(namespace), *keys)
|
||||
queries.append((query, params))
|
||||
embedding_request: Optional[tuple[str, Sequence[tuple[str, str, str, str]]]] = (
|
||||
None
|
||||
)
|
||||
if inserts:
|
||||
values = []
|
||||
insertion_params = []
|
||||
vector_values = []
|
||||
embedding_request_params = []
|
||||
|
||||
# First handle main store insertions
|
||||
for op in inserts:
|
||||
values.append("(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)")
|
||||
insertion_params.extend(
|
||||
[
|
||||
_namespace_to_text(op.namespace),
|
||||
op.key,
|
||||
Jsonb(op.value),
|
||||
Jsonb(cast(dict, op.value)),
|
||||
]
|
||||
)
|
||||
|
||||
# Then handle embeddings if configured
|
||||
if self.index_config:
|
||||
for op in inserts:
|
||||
if op.index is False:
|
||||
continue
|
||||
value = op.value
|
||||
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):
|
||||
pathname = f"{path}.{i}" if len(texts) > 1 else path
|
||||
vector_values.append(
|
||||
"(%s, %s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)"
|
||||
)
|
||||
embedding_request_params.append((ns, k, pathname, text))
|
||||
|
||||
values_str = ",".join(values)
|
||||
query = f"""
|
||||
INSERT INTO store (prefix, key, value, created_at, updated_at)
|
||||
VALUES {values_str}
|
||||
ON CONFLICT (prefix, key) DO UPDATE
|
||||
SET value = EXCLUDED.value, updated_at = CURRENT_TIMESTAMP
|
||||
SET value = EXCLUDED.value,
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
"""
|
||||
queries.append((query, insertion_params))
|
||||
|
||||
return queries
|
||||
if vector_values:
|
||||
values_str = ",".join(vector_values)
|
||||
query = f"""
|
||||
INSERT INTO store_vectors (prefix, key, field_name, embedding, created_at, updated_at)
|
||||
VALUES {values_str}
|
||||
ON CONFLICT (prefix, key, field_name) DO UPDATE
|
||||
SET embedding = EXCLUDED.embedding,
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
"""
|
||||
embedding_request = (query, embedding_request_params)
|
||||
|
||||
def _get_batch_search_queries(
|
||||
return queries, embedding_request
|
||||
|
||||
def _prepare_batch_search_queries(
|
||||
self,
|
||||
search_ops: Sequence[tuple[int, SearchOp]],
|
||||
) -> list[tuple[str, Sequence]]:
|
||||
queries: list[tuple[str, Sequence]] = []
|
||||
for _, op in search_ops:
|
||||
query = """
|
||||
SELECT prefix, key, value, created_at, updated_at
|
||||
FROM store
|
||||
WHERE prefix LIKE %s
|
||||
"""
|
||||
params: list = [f"{_namespace_to_text(op.namespace_prefix)}%"]
|
||||
) -> tuple[
|
||||
list[tuple[str, list[Union[None, str, list[float]]]]], # queries, params
|
||||
list[tuple[int, str]], # idx, query_text pairs to embed
|
||||
]:
|
||||
queries = []
|
||||
embedding_requests = []
|
||||
|
||||
for idx, (_, op) in enumerate(search_ops):
|
||||
# Build filter conditions first
|
||||
filter_params = []
|
||||
filter_conditions = []
|
||||
if op.filter:
|
||||
filter_conditions = []
|
||||
for key, value in op.filter.items():
|
||||
if isinstance(value, list):
|
||||
filter_conditions.append("value->%s @> %s::jsonb")
|
||||
params.extend([key, json.dumps(value)])
|
||||
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")
|
||||
params.extend([key, json.dumps(value)])
|
||||
query += " AND " + " AND ".join(filter_conditions)
|
||||
filter_params.extend([key, json.dumps(value)])
|
||||
|
||||
# Note: we will need to not do this if sim/keyword search
|
||||
# is used
|
||||
query += " ORDER BY updated_at DESC LIMIT %s OFFSET %s"
|
||||
params.extend([op.limit, op.offset])
|
||||
# Vector search branch
|
||||
if op.query and self.index_config:
|
||||
embedding_requests.append((idx, op.query))
|
||||
|
||||
queries.append((query, params))
|
||||
return queries
|
||||
score_operator = _get_distance_operator(self)
|
||||
vector_type = (
|
||||
cast(PostgresIndexConfig, self.index_config)
|
||||
.get("ann_index_config", {})
|
||||
.get("vector_type", "vector")
|
||||
)
|
||||
|
||||
if (
|
||||
vector_type == "bit"
|
||||
and self.index_config.get("distance_type") == "hamming"
|
||||
):
|
||||
score_operator = score_operator % (
|
||||
"%s",
|
||||
self.index_config["dims"],
|
||||
)
|
||||
else:
|
||||
score_operator = score_operator % (
|
||||
"%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)
|
||||
)
|
||||
base_query = f"""
|
||||
WITH scored AS (
|
||||
SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, {score_operator} AS score
|
||||
FROM store s
|
||||
JOIN store_vectors sv ON s.prefix = sv.prefix AND s.key = sv.key
|
||||
WHERE s.prefix LIKE %s {filter_str}
|
||||
ORDER BY {score_operator} DESC
|
||||
LIMIT %s
|
||||
)
|
||||
SELECT * FROM (
|
||||
SELECT DISTINCT ON (prefix, key)
|
||||
prefix, key, value, created_at, updated_at, score
|
||||
FROM scored
|
||||
ORDER BY prefix, key, score DESC
|
||||
) AS unique_docs
|
||||
ORDER BY score DESC
|
||||
LIMIT %s
|
||||
OFFSET %s
|
||||
"""
|
||||
params = [
|
||||
_PLACEHOLDER, # Vector placeholder
|
||||
f"{_namespace_to_text(op.namespace_prefix)}%",
|
||||
*filter_params,
|
||||
_PLACEHOLDER,
|
||||
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 filter_conditions:
|
||||
params.extend(filter_params)
|
||||
base_query += " AND " + " AND ".join(filter_conditions)
|
||||
|
||||
base_query += " ORDER BY updated_at DESC"
|
||||
base_query += " LIMIT %s OFFSET %s"
|
||||
params.extend([op.limit, op.offset])
|
||||
|
||||
queries.append((base_query, params))
|
||||
|
||||
return queries, embedding_requests
|
||||
|
||||
def _get_batch_list_namespaces_queries(
|
||||
self,
|
||||
@@ -214,45 +451,167 @@ class BasePostgresStore(Generic[C]):
|
||||
|
||||
query += " ORDER BY truncated_prefix LIMIT %s OFFSET %s"
|
||||
params.extend([op.limit, op.offset])
|
||||
queries.append((query, params))
|
||||
queries.append((query, tuple(params)))
|
||||
|
||||
return queries
|
||||
|
||||
def _get_filter_condition(self, key: str, op: str, value: Any) -> tuple[str, list]:
|
||||
"""Helper to generate filter conditions."""
|
||||
if op == "$eq":
|
||||
return "value->%s = %s::jsonb", [key, json.dumps(value)]
|
||||
elif op == "$gt":
|
||||
return "value->>%s > %s", [key, str(value)]
|
||||
elif op == "$gte":
|
||||
return "value->>%s >= %s", [key, str(value)]
|
||||
elif op == "$lt":
|
||||
return "value->>%s < %s", [key, str(value)]
|
||||
elif op == "$lte":
|
||||
return "value->>%s <= %s", [key, str(value)]
|
||||
elif op == "$ne":
|
||||
return "value->%s != %s::jsonb", [key, json.dumps(value)]
|
||||
else:
|
||||
raise ValueError(f"Unsupported operator: {op}")
|
||||
|
||||
class PostgresStore(BaseStore, BasePostgresStore[Connection]):
|
||||
__slots__ = ("_deserializer",)
|
||||
|
||||
class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
__slots__ = (
|
||||
"_deserializer",
|
||||
"pipe",
|
||||
"lock",
|
||||
"supports_pipeline",
|
||||
"index_config",
|
||||
"embeddings",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: Connection[Any],
|
||||
conn: _pg_internal.Conn,
|
||||
*,
|
||||
pipe: Optional[Pipeline] = None,
|
||||
deserializer: Optional[
|
||||
Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
|
||||
] = None,
|
||||
index: Optional[PostgresIndexConfig] = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self._deserializer = deserializer
|
||||
self.conn = conn
|
||||
self.pipe = pipe
|
||||
self.supports_pipeline = Capabilities().has_pipeline()
|
||||
self.lock = threading.Lock()
|
||||
self.index_config = index
|
||||
if self.index_config:
|
||||
self.embeddings, self.index_config = _ensure_index_config(self.index_config)
|
||||
else:
|
||||
self.embeddings = None
|
||||
|
||||
@classmethod
|
||||
@contextmanager
|
||||
def from_conn_string(
|
||||
cls,
|
||||
conn_string: str,
|
||||
*,
|
||||
pipeline: bool = False,
|
||||
pool_config: Optional[PoolConfig] = None,
|
||||
index: Optional[PostgresIndexConfig] = None,
|
||||
) -> Iterator["PostgresStore"]:
|
||||
"""Create a new PostgresStore instance from a connection string.
|
||||
|
||||
Args:
|
||||
conn_string (str): The Postgres connection info string.
|
||||
pipeline (bool): whether to use Pipeline
|
||||
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.
|
||||
|
||||
Returns:
|
||||
PostgresStore: A new PostgresStore instance.
|
||||
"""
|
||||
if pool_config is not None:
|
||||
pc = pool_config.copy()
|
||||
with cast(
|
||||
ConnectionPool[Connection[DictRow]],
|
||||
ConnectionPool(
|
||||
conn_string,
|
||||
min_size=pc.pop("min_size", 1),
|
||||
max_size=pc.pop("max_size", None),
|
||||
kwargs={
|
||||
"autocommit": True,
|
||||
"prepare_threshold": 0,
|
||||
"row_factory": dict_row,
|
||||
**(pc.pop("kwargs", None) or {}),
|
||||
},
|
||||
**cast(dict, pc),
|
||||
),
|
||||
) as pool:
|
||||
yield cls(conn=pool, index=index)
|
||||
else:
|
||||
with Connection.connect(
|
||||
conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
|
||||
) as conn:
|
||||
if pipeline:
|
||||
with conn.pipeline() as pipe:
|
||||
yield cls(conn, pipe=pipe, index=index)
|
||||
else:
|
||||
yield cls(conn, index=index)
|
||||
|
||||
@contextmanager
|
||||
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
|
||||
"""Create a database cursor as a context manager.
|
||||
|
||||
Args:
|
||||
pipeline (bool): whether to use pipeline for the DB operations inside the context manager.
|
||||
Will be applied regardless of whether the PostgresStore instance was initialized with a pipeline.
|
||||
If pipeline mode is not supported, will fall back to using transaction context manager.
|
||||
"""
|
||||
with _pg_internal.get_connection(self.conn) as conn:
|
||||
if self.pipe:
|
||||
# a connection in pipeline mode can be used concurrently
|
||||
# in multiple threads/coroutines, but only one cursor can be
|
||||
# used at a time
|
||||
try:
|
||||
with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
finally:
|
||||
if pipeline:
|
||||
self.pipe.sync()
|
||||
elif pipeline:
|
||||
# a connection not in pipeline mode can only be used by one
|
||||
# thread/coroutine at a time, so we acquire a lock
|
||||
if self.supports_pipeline:
|
||||
with (
|
||||
self.lock,
|
||||
conn.pipeline(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
with (
|
||||
self.lock,
|
||||
conn.transaction(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
def batch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
grouped_ops, num_ops = _group_ops(ops)
|
||||
results: list[Result] = [None] * num_ops
|
||||
|
||||
with self.conn.pipeline():
|
||||
with self._cursor(pipeline=True) as cur:
|
||||
if GetOp in grouped_ops:
|
||||
self._batch_get_ops(
|
||||
cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results
|
||||
)
|
||||
|
||||
if PutOp in grouped_ops:
|
||||
self._batch_put_ops(
|
||||
cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp])
|
||||
cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results, cur
|
||||
)
|
||||
|
||||
if SearchOp in grouped_ops:
|
||||
self._batch_search_ops(
|
||||
cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
|
||||
results,
|
||||
cur,
|
||||
)
|
||||
|
||||
if ListNamespacesOp in grouped_ops:
|
||||
@@ -262,25 +621,23 @@ class PostgresStore(BaseStore, BasePostgresStore[Connection]):
|
||||
grouped_ops[ListNamespacesOp],
|
||||
),
|
||||
results,
|
||||
cur,
|
||||
)
|
||||
if PutOp in grouped_ops:
|
||||
self._batch_put_ops(
|
||||
cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]), cur
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
return await asyncio.get_running_loop().run_in_executor(None, self.batch, ops)
|
||||
|
||||
def _batch_get_ops(
|
||||
self,
|
||||
get_ops: Sequence[tuple[int, GetOp]],
|
||||
results: list[Result],
|
||||
cur: Cursor[DictRow],
|
||||
) -> None:
|
||||
cursors = []
|
||||
for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
|
||||
cur = self.conn.cursor(binary=True)
|
||||
cur.execute(query, params)
|
||||
cursors.append((cur, namespace, items))
|
||||
|
||||
for cur, namespace, items in cursors:
|
||||
rows = cast(list[Row], cur.fetchall())
|
||||
key_to_row = {row["key"]: row for row in rows}
|
||||
for idx, key in items:
|
||||
@@ -295,70 +652,80 @@ class PostgresStore(BaseStore, BasePostgresStore[Connection]):
|
||||
def _batch_put_ops(
|
||||
self,
|
||||
put_ops: Sequence[tuple[int, PutOp]],
|
||||
cur: Cursor[DictRow],
|
||||
) -> None:
|
||||
queries = self._get_batch_PUT_queries(put_ops)
|
||||
queries, embedding_request = self._prepare_batch_PUT_queries(put_ops)
|
||||
if embedding_request:
|
||||
if self.embeddings is None:
|
||||
# Should not get here since the embedding config is required
|
||||
# to return an embedding_request above
|
||||
raise ValueError(
|
||||
"Embedding configuration is required for vector operations "
|
||||
f"(for semantic search). "
|
||||
f"Please provide an Embeddings when initializing the {self.__class__.__name__}."
|
||||
)
|
||||
query, txt_params = embedding_request
|
||||
# Update the params to replace the raw text with the vectors
|
||||
vectors = self.embeddings.embed_documents(
|
||||
[param[-1] for param in txt_params]
|
||||
)
|
||||
queries.append(
|
||||
(
|
||||
query,
|
||||
[
|
||||
p
|
||||
for (ns, k, pathname, _), vector in zip(txt_params, vectors)
|
||||
for p in (ns, k, pathname, vector)
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
for query, params in queries:
|
||||
cur = self.conn.cursor(binary=True)
|
||||
cur.execute(query, params)
|
||||
|
||||
def _batch_search_ops(
|
||||
self,
|
||||
search_ops: Sequence[tuple[int, SearchOp]],
|
||||
results: list[Result],
|
||||
cur: Cursor[DictRow],
|
||||
) -> None:
|
||||
queries = self._get_batch_search_queries(search_ops)
|
||||
cursors: list[tuple[Cursor[Any], int]] = []
|
||||
queries, embedding_requests = self._prepare_batch_search_queries(search_ops)
|
||||
|
||||
for (query, params), (idx, _) in zip(queries, search_ops):
|
||||
cur = self.conn.cursor(binary=True)
|
||||
if embedding_requests and self.embeddings:
|
||||
embeddings = self.embeddings.embed_documents(
|
||||
[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
|
||||
|
||||
for (idx, _), (query, params) in zip(search_ops, queries):
|
||||
# Execute the actual query
|
||||
cur.execute(query, params)
|
||||
cursors.append((cur, idx))
|
||||
|
||||
for cur, idx in cursors:
|
||||
rows = cast(list[Row], cur.fetchall())
|
||||
items = [
|
||||
_row_to_item(
|
||||
results[idx] = [
|
||||
_row_to_search_item(
|
||||
_decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
results[idx] = items
|
||||
|
||||
def _batch_list_namespaces_ops(
|
||||
self,
|
||||
list_ops: Sequence[tuple[int, ListNamespacesOp]],
|
||||
results: list[Result],
|
||||
cur: Cursor[DictRow],
|
||||
) -> None:
|
||||
queries = self._get_batch_list_namespaces_queries(list_ops)
|
||||
cursors: list[tuple[Cursor[Any], int]] = []
|
||||
for (query, params), (idx, _) in zip(queries, list_ops):
|
||||
cur = self.conn.cursor(binary=True)
|
||||
for (query, params), (idx, _) in zip(
|
||||
self._get_batch_list_namespaces_queries(list_ops), list_ops
|
||||
):
|
||||
cur.execute(query, params)
|
||||
cursors.append((cur, idx))
|
||||
results[idx] = [_decode_ns_bytes(row["truncated_prefix"]) for row in cur]
|
||||
|
||||
for cur, idx in cursors:
|
||||
rows = cast(list[dict], cur.fetchall())
|
||||
namespaces = [_decode_ns_bytes(row["truncated_prefix"]) for row in rows]
|
||||
results[idx] = namespaces
|
||||
|
||||
@classmethod
|
||||
@contextmanager
|
||||
def from_conn_string(
|
||||
cls,
|
||||
conn_string: str,
|
||||
) -> Iterator["PostgresStore"]:
|
||||
"""Create a new BasePostgresStore instance from a connection string.
|
||||
|
||||
Args:
|
||||
conn_string (str): The Postgres connection info string.
|
||||
|
||||
Returns:
|
||||
BasePostgresStore: A new BasePostgresStore instance.
|
||||
"""
|
||||
with Connection.connect(
|
||||
conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
|
||||
) as conn:
|
||||
yield cls(conn=conn)
|
||||
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
return await asyncio.get_running_loop().run_in_executor(None, self.batch, ops)
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Set up the store database.
|
||||
@@ -367,31 +734,47 @@ class PostgresStore(BaseStore, BasePostgresStore[Connection]):
|
||||
already exist and runs database migrations. It MUST be called directly by the user
|
||||
the first time the store is used.
|
||||
"""
|
||||
with self.conn.cursor(binary=True) as cur:
|
||||
|
||||
def _get_version(cur: Cursor[dict[str, Any]], table: str) -> int:
|
||||
try:
|
||||
cur.execute("SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1")
|
||||
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"]
|
||||
except UndefinedTable:
|
||||
self.conn.rollback()
|
||||
version = -1
|
||||
# Create store_migrations table if it doesn't exist
|
||||
cur.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS store_migrations (
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {table} (
|
||||
v INTEGER PRIMARY KEY
|
||||
)
|
||||
"""
|
||||
)
|
||||
for v, migration in enumerate(
|
||||
self.MIGRATIONS[version + 1 :], start=version + 1
|
||||
):
|
||||
cur.execute(migration)
|
||||
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
|
||||
):
|
||||
sql = migration.sql
|
||||
if migration.params:
|
||||
params = {
|
||||
k: v(self) if v is not None and callable(v) else v
|
||||
for k, v in migration.params.items()
|
||||
}
|
||||
sql = sql % params
|
||||
cur.execute(sql)
|
||||
cur.execute("INSERT INTO vector_migrations (v) VALUES (%s)", (v,))
|
||||
|
||||
|
||||
class Row(TypedDict):
|
||||
key: str
|
||||
@@ -401,6 +784,45 @@ class Row(TypedDict):
|
||||
updated_at: datetime
|
||||
|
||||
|
||||
# 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."""
|
||||
if not store.index_config:
|
||||
return "vector_cosine_ops"
|
||||
|
||||
config = cast(PostgresIndexConfig, store.index_config)
|
||||
index_config = config.get("ann_index_config", _DEFAULT_ANN_CONFIG).copy()
|
||||
vector_type = cast(str, index_config.get("vector_type", "vector"))
|
||||
if vector_type not in ("vector", "halfvec"):
|
||||
raise ValueError(
|
||||
f"Vector type must be 'vector' or 'halfvec', got {vector_type}"
|
||||
)
|
||||
|
||||
distance_type = config.get("distance_type", "cosine")
|
||||
|
||||
# For regular vectors
|
||||
type_prefix = {"vector": "vector", "halfvec": "halfvec"}[vector_type]
|
||||
|
||||
if distance_type not in ("l2", "inner_product", "cosine"):
|
||||
raise ValueError(
|
||||
f"Vector type {vector_type} only supports 'l2', 'inner_product', or 'cosine' distance, got {distance_type}"
|
||||
)
|
||||
|
||||
distance_suffix = {
|
||||
"l2": "l2_ops",
|
||||
"inner_product": "ip_ops",
|
||||
"cosine": "cosine_ops",
|
||||
}[distance_type]
|
||||
|
||||
return f"{type_prefix}_{distance_suffix}"
|
||||
|
||||
|
||||
def _namespace_to_text(
|
||||
namespace: tuple[str, ...], handle_wildcards: bool = False
|
||||
) -> str:
|
||||
@@ -416,15 +838,51 @@ def _row_to_item(
|
||||
*,
|
||||
loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None,
|
||||
) -> Item:
|
||||
"""Convert a row from the database into an Item.
|
||||
|
||||
Args:
|
||||
namespace: Item namespace
|
||||
row: Database row
|
||||
loader: Optional value loader for non-dict values
|
||||
"""
|
||||
val = row["value"]
|
||||
if not isinstance(val, dict):
|
||||
val = (loader or _json_loads)(val)
|
||||
|
||||
kwargs = {
|
||||
"key": row["key"],
|
||||
"namespace": namespace,
|
||||
"value": val,
|
||||
"created_at": row["created_at"],
|
||||
"updated_at": row["updated_at"],
|
||||
}
|
||||
|
||||
return Item(**kwargs)
|
||||
|
||||
|
||||
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"]
|
||||
return Item(
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@@ -455,3 +913,62 @@ def _decode_ns_bytes(namespace: Union[str, bytes, list]) -> tuple[str, ...]:
|
||||
if isinstance(namespace, bytes):
|
||||
namespace = namespace.decode()[1:]
|
||||
return tuple(namespace.split("."))
|
||||
|
||||
|
||||
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 "
|
||||
f"(for semantic search). "
|
||||
f"Please provide an Embeddings when initializing the {store.__class__.__name__}."
|
||||
)
|
||||
|
||||
config = cast(PostgresIndexConfig, store.index_config)
|
||||
distance_type = config.get("distance_type", "cosine")
|
||||
|
||||
if distance_type == "l2":
|
||||
return "1 - (sv.embedding <-> %s::%s)"
|
||||
elif distance_type == "inner_product":
|
||||
return "-(sv.embedding <#> %s::%s)"
|
||||
else: # cosine
|
||||
return "1 - (sv.embedding <=> %s::%s)"
|
||||
|
||||
|
||||
def _ensure_index_config(
|
||||
index_config: PostgresIndexConfig,
|
||||
) -> tuple[Optional["Embeddings"], PostgresIndexConfig]:
|
||||
index_config = index_config.copy()
|
||||
tokenized: list[tuple[str, Union[Literal["$"], list[str]]]] = []
|
||||
tot = 0
|
||||
text_fields = index_config.get("text_fields") or ["$"]
|
||||
if isinstance(text_fields, str):
|
||||
text_fields = [text_fields]
|
||||
if not isinstance(text_fields, list):
|
||||
raise ValueError(f"Text fields must be a list or a string. Got {text_fields}")
|
||||
for p in text_fields:
|
||||
if p == "$":
|
||||
tokenized.append((p, "$"))
|
||||
tot += 1
|
||||
else:
|
||||
toks = tokenize_path(p)
|
||||
tokenized.append((p, toks))
|
||||
tot += len(toks)
|
||||
index_config["__tokenized_fields"] = tokenized
|
||||
index_config["__estimated_num_vectors"] = tot
|
||||
embeddings = ensure_embeddings(
|
||||
index_config.get("embed"),
|
||||
)
|
||||
return embeddings, index_config
|
||||
|
||||
|
||||
_PLACEHOLDER = object()
|
||||
|
||||
Generated
+529
-426
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "langgraph-checkpoint-postgres"
|
||||
version = "2.0.3"
|
||||
version = "2.0.5"
|
||||
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
|
||||
authors = []
|
||||
license = "MIT"
|
||||
@@ -10,7 +10,7 @@ packages = [{ include = "langgraph" }]
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
python = "^3.9.0,<4.0"
|
||||
langgraph-checkpoint = "^2.0.2"
|
||||
langgraph-checkpoint = "^2.0.7"
|
||||
orjson = ">=3.10.1"
|
||||
psycopg = "^3.0.0"
|
||||
psycopg-pool = "^3.0.0"
|
||||
|
||||
@@ -1,12 +1,13 @@
|
||||
services:
|
||||
postgres-test:
|
||||
image: postgres:${POSTGRES_VERSION:-16}
|
||||
image: pgvector/pgvector:pg${POSTGRES_VERSION:-16}
|
||||
ports:
|
||||
- "5441:5432"
|
||||
environment:
|
||||
POSTGRES_DB: postgres
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
command: ["postgres", "-c", "shared_preload_libraries=vector"]
|
||||
healthcheck:
|
||||
test: pg_isready -U postgres
|
||||
start_period: 10s
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
from typing import AsyncIterator
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
import pytest
|
||||
from psycopg import AsyncConnection
|
||||
from psycopg.errors import UndefinedTable
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
|
||||
from tests.embed_test_utils import CharacterEmbeddings
|
||||
|
||||
DEFAULT_URI = "postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable"
|
||||
|
||||
|
||||
@@ -24,6 +26,18 @@ async def clear_test_db(conn: AsyncConnection[DictRow]) -> None:
|
||||
await conn.execute("DELETE FROM checkpoint_blobs")
|
||||
await conn.execute("DELETE FROM checkpoint_writes")
|
||||
await conn.execute("DELETE FROM checkpoint_migrations")
|
||||
await conn.execute("DELETE FROM store_migrations")
|
||||
except UndefinedTable:
|
||||
pass
|
||||
try:
|
||||
await conn.execute("DELETE FROM store_migrations")
|
||||
await conn.execute("DELETE FROM store")
|
||||
except UndefinedTable:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_embeddings() -> CharacterEmbeddings:
|
||||
return CharacterEmbeddings(dims=500)
|
||||
|
||||
|
||||
VECTOR_TYPES = ["vector", "halfvec"]
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
"""Embedding utilities for testing."""
|
||||
|
||||
import math
|
||||
import random
|
||||
from collections import Counter, defaultdict
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.embeddings import Embeddings
|
||||
|
||||
|
||||
class CharacterEmbeddings(Embeddings):
|
||||
"""Simple character-frequency based embeddings using random projections."""
|
||||
|
||||
def __init__(self, dims: int = 50, seed: int = 42):
|
||||
"""Initialize with embedding dimensions and random seed."""
|
||||
self._rng = random.Random(seed)
|
||||
self.dims = dims
|
||||
# Create projection vector for each character lazily
|
||||
self._char_projections: defaultdict[str, list[float]] = defaultdict(
|
||||
lambda: [
|
||||
self._rng.gauss(0, 1 / math.sqrt(self.dims)) for _ in range(self.dims)
|
||||
]
|
||||
)
|
||||
|
||||
def _embed_one(self, text: str) -> list[float]:
|
||||
"""Embed a single text."""
|
||||
counts = Counter(text)
|
||||
total = sum(counts.values())
|
||||
|
||||
if total == 0:
|
||||
return [0.0] * self.dims
|
||||
|
||||
embedding = [0.0] * self.dims
|
||||
for char, count in counts.items():
|
||||
weight = count / total
|
||||
char_proj = self._char_projections[char]
|
||||
for i, proj in enumerate(char_proj):
|
||||
embedding[i] += weight * proj
|
||||
|
||||
norm = math.sqrt(sum(x * x for x in embedding))
|
||||
if norm > 0:
|
||||
embedding = [x / norm for x in embedding]
|
||||
|
||||
return embedding
|
||||
|
||||
def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Embed a list of documents."""
|
||||
return [self._embed_one(text) for text in texts]
|
||||
|
||||
def embed_query(self, text: str) -> list[float]:
|
||||
"""Embed a query string."""
|
||||
return self._embed_one(text)
|
||||
|
||||
def __eq__(self, other: Any) -> bool:
|
||||
return isinstance(other, CharacterEmbeddings) and self.dims == other.dims
|
||||
@@ -1,7 +1,6 @@
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from conftest import DEFAULT_URI # type: ignore
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
@@ -11,6 +10,7 @@ from langgraph.checkpoint.base import (
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
from tests.conftest import DEFAULT_URI
|
||||
|
||||
|
||||
class TestAsyncPostgresSaver:
|
||||
|
||||
@@ -1,114 +1,84 @@
|
||||
# type: ignore
|
||||
import itertools
|
||||
import sys
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, Optional
|
||||
|
||||
import pytest
|
||||
from conftest import DEFAULT_URI # type: ignore
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from psycopg import AsyncConnection
|
||||
|
||||
from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
|
||||
from langgraph.store.postgres import AsyncPostgresStore
|
||||
from tests.conftest import (
|
||||
DEFAULT_URI,
|
||||
VECTOR_TYPES,
|
||||
CharacterEmbeddings,
|
||||
)
|
||||
|
||||
|
||||
class MockAsyncCursor:
|
||||
def __init__(self, fetch_result: Any) -> None:
|
||||
self.fetch_result = fetch_result
|
||||
self.execute = AsyncMock()
|
||||
self.fetchall = AsyncMock(return_value=self.fetch_result)
|
||||
@pytest.fixture(scope="function", params=["default", "pipe", "pool"])
|
||||
async def store(request) -> AsyncIterator[AsyncPostgresStore]:
|
||||
if sys.version_info < (3, 10):
|
||||
pytest.skip("Async Postgres tests require Python 3.10+")
|
||||
|
||||
database = f"test_{uuid.uuid4().hex[:16]}"
|
||||
uri_parts = DEFAULT_URI.split("/")
|
||||
uri_base = "/".join(uri_parts[:-1])
|
||||
query_params = ""
|
||||
if "?" in uri_parts[-1]:
|
||||
db_name, query_params = uri_parts[-1].split("?", 1)
|
||||
query_params = "?" + query_params
|
||||
|
||||
class MockAsyncConnection:
|
||||
def __init__(self) -> None:
|
||||
self.cursor = MagicMock()
|
||||
self.pipeline = MagicMock(
|
||||
return_value=AsyncMock(__aenter__=AsyncMock(), __aexit__=AsyncMock())
|
||||
)
|
||||
conn_string = f"{uri_base}/{database}{query_params}"
|
||||
admin_conn_string = DEFAULT_URI
|
||||
|
||||
async with await AsyncConnection.connect(
|
||||
admin_conn_string, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
async with AsyncPostgresStore.from_conn_string(conn_string) as store:
|
||||
await store.setup()
|
||||
|
||||
@pytest.fixture
|
||||
def mock_connection() -> MockAsyncConnection:
|
||||
return MockAsyncConnection()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def store(mock_connection: MockAsyncConnection) -> AsyncPostgresStore:
|
||||
return AsyncPostgresStore(mock_connection)
|
||||
if request.param == "pipe":
|
||||
async with AsyncPostgresStore.from_conn_string(
|
||||
conn_string, pipeline=True
|
||||
) as store:
|
||||
yield store
|
||||
elif request.param == "pool":
|
||||
async with AsyncPostgresStore.from_conn_string(
|
||||
conn_string, pool_config={"min_size": 1, "max_size": 10}
|
||||
) as store:
|
||||
yield store
|
||||
else: # default
|
||||
async with AsyncPostgresStore.from_conn_string(conn_string) as store:
|
||||
yield store
|
||||
finally:
|
||||
async with await AsyncConnection.connect(
|
||||
admin_conn_string, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
async def test_abatch_order(store: AsyncPostgresStore) -> None:
|
||||
mock_connection = store.conn
|
||||
mock_get_cursor = MockAsyncCursor(
|
||||
[
|
||||
{
|
||||
"key": "key1",
|
||||
"value": '{"data": "value1"}',
|
||||
"created_at": datetime.now(),
|
||||
"updated_at": datetime.now(),
|
||||
"prefix": "test.foo",
|
||||
},
|
||||
{
|
||||
"key": "key2",
|
||||
"value": '{"data": "value2"}',
|
||||
"created_at": datetime.now(),
|
||||
"updated_at": datetime.now(),
|
||||
"prefix": "test.bar",
|
||||
},
|
||||
]
|
||||
)
|
||||
mock_search_cursor = MockAsyncCursor(
|
||||
[
|
||||
{
|
||||
"key": "key1",
|
||||
"value": '{"data": "value1"}',
|
||||
"created_at": datetime.now(),
|
||||
"updated_at": datetime.now(),
|
||||
"prefix": "test.foo",
|
||||
},
|
||||
]
|
||||
)
|
||||
mock_list_namespaces_cursor = MockAsyncCursor(
|
||||
[
|
||||
{"truncated_prefix": b"\x01test"},
|
||||
]
|
||||
)
|
||||
|
||||
failures = []
|
||||
|
||||
def cursor_side_effect(binary: bool = False) -> Any:
|
||||
cursor = MagicMock()
|
||||
|
||||
async def execute_side_effect(query: str, *params: Any) -> None:
|
||||
# My super sophisticated database.
|
||||
if "SELECT prefix, key," in query:
|
||||
cursor.fetchall = mock_search_cursor.fetchall
|
||||
elif "SELECT DISTINCT ON (truncated_prefix)" in query:
|
||||
cursor.fetchall = mock_list_namespaces_cursor.fetchall
|
||||
elif "WHERE prefix = %s AND key" in query:
|
||||
cursor.fetchall = mock_get_cursor.fetchall
|
||||
elif "INSERT INTO " in query:
|
||||
pass
|
||||
else:
|
||||
e = ValueError(f"Unmatched query: {query}")
|
||||
failures.append(e)
|
||||
raise e
|
||||
|
||||
cursor.execute = AsyncMock(side_effect=execute_side_effect)
|
||||
return cursor
|
||||
|
||||
mock_connection.cursor.side_effect = cursor_side_effect # type: ignore
|
||||
# Setup test data
|
||||
await store.aput(("test", "foo"), "key1", {"data": "value1"})
|
||||
await store.aput(("test", "bar"), "key2", {"data": "value2"})
|
||||
|
||||
ops = [
|
||||
GetOp(namespace=("test",), key="key1"),
|
||||
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
|
||||
GetOp(namespace=("test", "foo"), key="key1"),
|
||||
PutOp(namespace=("test", "bar"), key="key2", value={"data": "value2"}),
|
||||
SearchOp(
|
||||
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
|
||||
),
|
||||
ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0),
|
||||
GetOp(namespace=("test",), key="key3"),
|
||||
]
|
||||
|
||||
results = await store.abatch(ops)
|
||||
assert not failures
|
||||
assert len(results) == 5
|
||||
assert isinstance(results[0], Item)
|
||||
assert isinstance(results[0].value, dict)
|
||||
@@ -118,27 +88,29 @@ async def test_abatch_order(store: AsyncPostgresStore) -> None:
|
||||
assert isinstance(results[2], list)
|
||||
assert len(results[2]) == 1
|
||||
assert isinstance(results[3], list)
|
||||
assert results[3] == [("test",)]
|
||||
assert ("test", "foo") in results[3] and ("test", "bar") in results[3]
|
||||
assert results[4] is None
|
||||
|
||||
ops_reordered = [
|
||||
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
|
||||
GetOp(namespace=("test",), key="key2"),
|
||||
GetOp(namespace=("test", "bar"), key="key2"),
|
||||
ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0),
|
||||
PutOp(namespace=("test",), key="key3", value={"data": "value3"}),
|
||||
GetOp(namespace=("test",), key="key1"),
|
||||
GetOp(namespace=("test", "foo"), key="key1"),
|
||||
]
|
||||
|
||||
results_reordered = await store.abatch(ops_reordered)
|
||||
assert not failures
|
||||
assert len(results_reordered) == 5
|
||||
assert isinstance(results_reordered[0], list)
|
||||
assert len(results_reordered[0]) == 1
|
||||
assert len(results_reordered[0]) == 2
|
||||
assert isinstance(results_reordered[1], Item)
|
||||
assert results_reordered[1].value == {"data": "value2"}
|
||||
assert results_reordered[1].key == "key2"
|
||||
assert isinstance(results_reordered[2], list)
|
||||
assert results_reordered[2] == [("test",)]
|
||||
assert ("test", "foo") in results_reordered[2] and (
|
||||
"test",
|
||||
"bar",
|
||||
) in results_reordered[2]
|
||||
assert results_reordered[3] is None
|
||||
assert isinstance(results_reordered[4], Item)
|
||||
assert results_reordered[4].value == {"data": "value1"}
|
||||
@@ -146,26 +118,9 @@ async def test_abatch_order(store: AsyncPostgresStore) -> None:
|
||||
|
||||
|
||||
async def test_batch_get_ops(store: AsyncPostgresStore) -> None:
|
||||
mock_connection = store.conn
|
||||
mock_cursor = MockAsyncCursor(
|
||||
[
|
||||
{
|
||||
"key": "key1",
|
||||
"value": '{"data": "value1"}',
|
||||
"created_at": datetime.now(),
|
||||
"updated_at": datetime.now(),
|
||||
"prefix": "test.foo",
|
||||
},
|
||||
{
|
||||
"key": "key2",
|
||||
"value": '{"data": "value2"}',
|
||||
"created_at": datetime.now(),
|
||||
"updated_at": datetime.now(),
|
||||
"prefix": "test.bar",
|
||||
},
|
||||
]
|
||||
)
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
# Setup test data
|
||||
await store.aput(("test",), "key1", {"data": "value1"})
|
||||
await store.aput(("test",), "key2", {"data": "value2"})
|
||||
|
||||
ops = [
|
||||
GetOp(namespace=("test",), key="key1"),
|
||||
@@ -184,10 +139,6 @@ async def test_batch_get_ops(store: AsyncPostgresStore) -> None:
|
||||
|
||||
|
||||
async def test_batch_put_ops(store: AsyncPostgresStore) -> None:
|
||||
mock_connection = store.conn
|
||||
mock_cursor = MockAsyncCursor([])
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
|
||||
ops = [
|
||||
PutOp(namespace=("test",), key="key1", value={"data": "value1"}),
|
||||
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
|
||||
@@ -198,30 +149,16 @@ async def test_batch_put_ops(store: AsyncPostgresStore) -> None:
|
||||
|
||||
assert len(results) == 3
|
||||
assert all(result is None for result in results)
|
||||
assert mock_cursor.execute.call_count == 2
|
||||
|
||||
# Verify the puts worked
|
||||
items = await store.asearch(["test"], limit=10)
|
||||
assert len(items) == 2 # key3 had None value so wasn't stored
|
||||
|
||||
|
||||
async def test_batch_search_ops(store: AsyncPostgresStore) -> None:
|
||||
mock_connection = store.conn
|
||||
mock_cursor = MockAsyncCursor(
|
||||
[
|
||||
{
|
||||
"key": "key1",
|
||||
"value": '{"data": "value1"}',
|
||||
"created_at": datetime.now(),
|
||||
"updated_at": datetime.now(),
|
||||
"prefix": "test.foo",
|
||||
},
|
||||
{
|
||||
"key": "key2",
|
||||
"value": '{"data": "value2"}',
|
||||
"created_at": datetime.now(),
|
||||
"updated_at": datetime.now(),
|
||||
"prefix": "test.bar",
|
||||
},
|
||||
]
|
||||
)
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
# Setup test data
|
||||
await store.aput(("test", "foo"), "key1", {"data": "value1"})
|
||||
await store.aput(("test", "bar"), "key2", {"data": "value2"})
|
||||
|
||||
ops = [
|
||||
SearchOp(
|
||||
@@ -233,297 +170,338 @@ async def test_batch_search_ops(store: AsyncPostgresStore) -> None:
|
||||
results = await store.abatch(ops)
|
||||
|
||||
assert len(results) == 2
|
||||
assert len(results[0]) == 2
|
||||
assert len(results[1]) == 2
|
||||
assert len(results[0]) == 1 # Filtered results
|
||||
assert len(results[1]) == 2 # All results
|
||||
|
||||
|
||||
async def test_batch_list_namespaces_ops(store: AsyncPostgresStore) -> None:
|
||||
mock_connection = store.conn
|
||||
mock_cursor = MockAsyncCursor(
|
||||
[
|
||||
{"truncated_prefix": b"\x01test.namespace1"},
|
||||
{"truncated_prefix": b"\x01test.namespace2"},
|
||||
]
|
||||
)
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
# Setup test data
|
||||
await store.aput(("test", "namespace1"), "key1", {"data": "value1"})
|
||||
await store.aput(("test", "namespace2"), "key2", {"data": "value2"})
|
||||
|
||||
ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)]
|
||||
|
||||
results = await store.abatch(ops)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0] == [("test", "namespace1"), ("test", "namespace2")]
|
||||
assert len(results[0]) == 2
|
||||
assert ("test", "namespace1") in results[0]
|
||||
assert ("test", "namespace2") in results[0]
|
||||
|
||||
|
||||
# The following use the actual DB connection
|
||||
@asynccontextmanager
|
||||
async def _create_vector_store(
|
||||
vector_type: str,
|
||||
distance_type: str,
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
text_fields: Optional[list[str]] = None,
|
||||
) -> AsyncIterator[AsyncPostgresStore]:
|
||||
"""Create a store with vector search enabled."""
|
||||
if sys.version_info < (3, 10):
|
||||
pytest.skip("Async Postgres tests require Python 3.10+")
|
||||
|
||||
database = f"test_{uuid.uuid4().hex[:16]}"
|
||||
uri_parts = DEFAULT_URI.split("/")
|
||||
uri_base = "/".join(uri_parts[:-1])
|
||||
query_params = ""
|
||||
if "?" in uri_parts[-1]:
|
||||
db_name, query_params = uri_parts[-1].split("?", 1)
|
||||
query_params = "?" + query_params
|
||||
|
||||
class TestAsyncPostgresStore:
|
||||
@pytest.fixture(autouse=True)
|
||||
async def setup(self) -> None:
|
||||
async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
|
||||
conn_string = f"{uri_base}/{database}{query_params}"
|
||||
admin_conn_string = DEFAULT_URI
|
||||
|
||||
index_config = {
|
||||
"dims": fake_embeddings.dims,
|
||||
"embed": fake_embeddings,
|
||||
"ann_index_config": {
|
||||
"vector_type": vector_type,
|
||||
},
|
||||
"distance_type": distance_type,
|
||||
"text_fields": text_fields,
|
||||
}
|
||||
|
||||
async with await AsyncConnection.connect(
|
||||
admin_conn_string, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
async with AsyncPostgresStore.from_conn_string(
|
||||
conn_string,
|
||||
index=index_config,
|
||||
) as store:
|
||||
await store.setup()
|
||||
yield store
|
||||
finally:
|
||||
async with await AsyncConnection.connect(
|
||||
admin_conn_string, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
async def test_basic_store_ops(self) -> None:
|
||||
async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
|
||||
namespace = ("test", "documents")
|
||||
item_id = "doc1"
|
||||
item_value = {"title": "Test Document", "content": "Hello, World!"}
|
||||
|
||||
await store.aput(namespace, item_id, item_value)
|
||||
item = await store.aget(namespace, item_id)
|
||||
@pytest.fixture(
|
||||
scope="function",
|
||||
params=[
|
||||
(vector_type, distance_type)
|
||||
for vector_type in VECTOR_TYPES
|
||||
for distance_type in (
|
||||
["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
|
||||
)
|
||||
],
|
||||
ids=lambda p: f"{p[0]}_{p[1]}",
|
||||
)
|
||||
async def vector_store(
|
||||
request,
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
) -> AsyncIterator[AsyncPostgresStore]:
|
||||
"""Create a store with vector search enabled."""
|
||||
vector_type, distance_type = request.param
|
||||
async with _create_vector_store(
|
||||
vector_type, distance_type, fake_embeddings
|
||||
) as store:
|
||||
yield store
|
||||
|
||||
assert item
|
||||
assert item.namespace == namespace
|
||||
assert item.key == item_id
|
||||
assert item.value == item_value
|
||||
|
||||
updated_value = {
|
||||
"title": "Updated Test Document",
|
||||
"content": "Hello, LangGraph!",
|
||||
}
|
||||
await store.aput(namespace, item_id, updated_value)
|
||||
updated_item = await store.aget(namespace, item_id)
|
||||
async def test_vector_store_initialization(
|
||||
vector_store: AsyncPostgresStore, fake_embeddings: CharacterEmbeddings
|
||||
) -> None:
|
||||
"""Test store initialization with embedding config."""
|
||||
assert vector_store.index_config is not None
|
||||
assert vector_store.index_config["dims"] == fake_embeddings.dims
|
||||
if isinstance(vector_store.index_config["embed"], Embeddings):
|
||||
assert vector_store.index_config["embed"] == fake_embeddings
|
||||
|
||||
assert updated_item.value == updated_value
|
||||
assert updated_item.updated_at > item.updated_at
|
||||
different_namespace = ("test", "other_documents")
|
||||
item_in_different_namespace = await store.aget(different_namespace, item_id)
|
||||
assert item_in_different_namespace is None
|
||||
|
||||
new_item_id = "doc2"
|
||||
new_item_value = {"title": "Another Document", "content": "Greetings!"}
|
||||
await store.aput(namespace, new_item_id, new_item_value)
|
||||
async def test_vector_insert_with_auto_embedding(
|
||||
vector_store: AsyncPostgresStore,
|
||||
) -> None:
|
||||
"""Test inserting items that get auto-embedded."""
|
||||
docs = [
|
||||
("doc1", {"text": "short text"}),
|
||||
("doc2", {"text": "longer text document"}),
|
||||
("doc3", {"text": "longest text document here"}),
|
||||
("doc4", {"description": "text in description field"}),
|
||||
("doc5", {"content": "text in content field"}),
|
||||
("doc6", {"body": "text in body field"}),
|
||||
]
|
||||
|
||||
search_results = await store.asearch(["test"], limit=10)
|
||||
items = search_results
|
||||
assert len(items) == 2
|
||||
assert any(item.key == item_id for item in items)
|
||||
assert any(item.key == new_item_id for item in items)
|
||||
for key, value in docs:
|
||||
await vector_store.aput(("test",), key, value)
|
||||
|
||||
namespaces = await store.alist_namespaces(prefix=["test"])
|
||||
assert ("test", "documents") in namespaces
|
||||
results = await vector_store.asearch(("test",), query="long text")
|
||||
assert len(results) > 0
|
||||
|
||||
await store.adelete(namespace, item_id)
|
||||
await store.adelete(namespace, new_item_id)
|
||||
deleted_item = await store.aget(namespace, item_id)
|
||||
assert deleted_item is None
|
||||
doc_order = [r.key for r in results]
|
||||
assert "doc2" in doc_order
|
||||
assert "doc3" in doc_order
|
||||
|
||||
deleted_item = await store.aget(namespace, new_item_id)
|
||||
assert deleted_item is None
|
||||
|
||||
empty_search_results = await store.asearch(["test"], limit=10)
|
||||
assert len(empty_search_results) == 0
|
||||
async def test_vector_update_with_embedding(vector_store: AsyncPostgresStore) -> None:
|
||||
"""Test that updating items properly updates their embeddings."""
|
||||
await vector_store.aput(("test",), "doc1", {"text": "zany zebra Xerxes"})
|
||||
await vector_store.aput(("test",), "doc2", {"text": "something about dogs"})
|
||||
await vector_store.aput(("test",), "doc3", {"text": "text about birds"})
|
||||
|
||||
async def test_list_namespaces(self) -> None:
|
||||
async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
|
||||
test_pref = str(uuid.uuid4())
|
||||
test_namespaces = [
|
||||
(test_pref, "test", "documents", "public", test_pref),
|
||||
(test_pref, "test", "documents", "private", test_pref),
|
||||
(test_pref, "test", "images", "public", test_pref),
|
||||
(test_pref, "test", "images", "private", test_pref),
|
||||
(test_pref, "prod", "documents", "public", test_pref),
|
||||
(
|
||||
test_pref,
|
||||
"prod",
|
||||
"documents",
|
||||
"some",
|
||||
"nesting",
|
||||
"public",
|
||||
test_pref,
|
||||
),
|
||||
(test_pref, "prod", "documents", "private", test_pref),
|
||||
]
|
||||
results_initial = await vector_store.asearch(("test",), query="Zany Xerxes")
|
||||
assert len(results_initial) > 0
|
||||
assert results_initial[0].key == "doc1"
|
||||
initial_score = results_initial[0].score
|
||||
|
||||
for namespace in test_namespaces:
|
||||
await store.aput(namespace, "dummy", {"content": "dummy"})
|
||||
await vector_store.aput(("test",), "doc1", {"text": "new text about dogs"})
|
||||
|
||||
prefix_result = await store.alist_namespaces(prefix=[test_pref, "test"])
|
||||
assert len(prefix_result) == 4
|
||||
assert all([ns[1] == "test" for ns in prefix_result])
|
||||
results_after = await vector_store.asearch(("test",), query="Zany Xerxes")
|
||||
after_score = next((r.score for r in results_after if r.key == "doc1"), 0.0)
|
||||
assert after_score < initial_score
|
||||
|
||||
specific_prefix_result = await store.alist_namespaces(
|
||||
prefix=[test_pref, "test", "documents"]
|
||||
)
|
||||
assert len(specific_prefix_result) == 2
|
||||
assert all(
|
||||
[ns[1:3] == ("test", "documents") for ns in specific_prefix_result]
|
||||
)
|
||||
results_new = await vector_store.asearch(("test",), query="new text about dogs")
|
||||
for r in results_new:
|
||||
if r.key == "doc1":
|
||||
assert r.score > after_score
|
||||
|
||||
suffix_result = await store.alist_namespaces(suffix=["public", test_pref])
|
||||
assert len(suffix_result) == 4
|
||||
assert all(ns[-2] == "public" for ns in suffix_result)
|
||||
# Don't index this one
|
||||
await vector_store.aput(
|
||||
("test",), "doc4", {"text": "new text about dogs"}, index=False
|
||||
)
|
||||
results_new = await vector_store.asearch(
|
||||
("test",), query="new text about dogs", limit=3
|
||||
)
|
||||
assert not any(r.key == "doc4" for r in results_new)
|
||||
|
||||
prefix_suffix_result = await store.alist_namespaces(
|
||||
prefix=[test_pref, "test"], suffix=["public", test_pref]
|
||||
)
|
||||
assert len(prefix_suffix_result) == 2
|
||||
assert all(
|
||||
ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result
|
||||
)
|
||||
|
||||
wildcard_prefix_result = await store.alist_namespaces(
|
||||
prefix=[test_pref, "*", "documents"]
|
||||
)
|
||||
assert len(wildcard_prefix_result) == 5
|
||||
assert all(ns[2] == "documents" for ns in wildcard_prefix_result)
|
||||
async def test_vector_search_with_filters(vector_store: AsyncPostgresStore) -> None:
|
||||
"""Test combining vector search with filters."""
|
||||
docs = [
|
||||
("doc1", {"text": "red apple", "color": "red", "score": 4.5}),
|
||||
("doc2", {"text": "red car", "color": "red", "score": 3.0}),
|
||||
("doc3", {"text": "green apple", "color": "green", "score": 4.0}),
|
||||
("doc4", {"text": "blue car", "color": "blue", "score": 3.5}),
|
||||
]
|
||||
|
||||
wildcard_suffix_result = await store.alist_namespaces(
|
||||
suffix=["*", "public", test_pref]
|
||||
)
|
||||
assert len(wildcard_suffix_result) == 4
|
||||
assert all(ns[-2] == "public" for ns in wildcard_suffix_result)
|
||||
wildcard_single = await store.alist_namespaces(
|
||||
suffix=["some", "*", "public", test_pref]
|
||||
)
|
||||
assert len(wildcard_single) == 1
|
||||
assert wildcard_single[0] == (
|
||||
test_pref,
|
||||
"prod",
|
||||
"documents",
|
||||
"some",
|
||||
"nesting",
|
||||
"public",
|
||||
test_pref,
|
||||
)
|
||||
for key, value in docs:
|
||||
await vector_store.aput(("test",), key, value)
|
||||
|
||||
max_depth_result = await store.alist_namespaces(max_depth=3)
|
||||
assert all([len(ns) <= 3 for ns in max_depth_result])
|
||||
max_depth_result = await store.alist_namespaces(
|
||||
max_depth=4, prefix=[test_pref, "*", "documents"]
|
||||
)
|
||||
assert (
|
||||
len(set(tuple(res) for res in max_depth_result))
|
||||
== len(max_depth_result)
|
||||
== 5
|
||||
)
|
||||
results = await vector_store.asearch(
|
||||
("test",), query="apple", filter={"color": "red"}
|
||||
)
|
||||
assert len(results) == 2
|
||||
assert results[0].key == "doc1"
|
||||
|
||||
limit_result = await store.alist_namespaces(prefix=[test_pref], limit=3)
|
||||
assert len(limit_result) == 3
|
||||
results = await vector_store.asearch(
|
||||
("test",), query="car", filter={"color": "red"}
|
||||
)
|
||||
assert len(results) == 2
|
||||
assert results[0].key == "doc2"
|
||||
|
||||
offset_result = await store.alist_namespaces(prefix=[test_pref], offset=3)
|
||||
assert len(offset_result) == len(test_namespaces) - 3
|
||||
results = await vector_store.asearch(
|
||||
("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
|
||||
)
|
||||
assert len(results) == 3
|
||||
assert results[0].key == "doc4"
|
||||
|
||||
empty_prefix_result = await store.alist_namespaces(prefix=[test_pref])
|
||||
assert len(empty_prefix_result) == len(test_namespaces)
|
||||
assert set(tuple(ns) for ns in empty_prefix_result) == set(
|
||||
tuple(ns) for ns in test_namespaces
|
||||
)
|
||||
results = await vector_store.asearch(
|
||||
("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"}
|
||||
)
|
||||
assert len(results) == 1
|
||||
assert results[0].key == "doc3"
|
||||
|
||||
for namespace in test_namespaces:
|
||||
await store.adelete(namespace, "dummy")
|
||||
|
||||
async def test_search(self):
|
||||
async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
|
||||
test_namespaces = [
|
||||
("test_search", "documents", "user1"),
|
||||
("test_search", "documents", "user2"),
|
||||
("test_search", "reports", "department1"),
|
||||
("test_search", "reports", "department2"),
|
||||
]
|
||||
test_items = [
|
||||
{"title": "Doc 1", "author": "John Doe", "tags": ["important"]},
|
||||
{"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]},
|
||||
{"title": "Report A", "author": "John Doe", "tags": ["final"]},
|
||||
{"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]},
|
||||
]
|
||||
empty = await store.asearch(
|
||||
(
|
||||
"scoped",
|
||||
"assistant_id",
|
||||
"shared",
|
||||
"6c5356f6-63ab-4158-868d-cd9fd14c736e",
|
||||
),
|
||||
limit=10,
|
||||
offset=0,
|
||||
)
|
||||
assert len(empty) == 0
|
||||
async def test_vector_search_pagination(vector_store: AsyncPostgresStore) -> None:
|
||||
"""Test pagination with vector search."""
|
||||
for i in range(5):
|
||||
await vector_store.aput(
|
||||
("test",), f"doc{i}", {"text": f"test document number {i}"}
|
||||
)
|
||||
|
||||
for namespace, item in zip(test_namespaces, test_items):
|
||||
await store.aput(namespace, f"item_{namespace[-1]}", item)
|
||||
results_page1 = await vector_store.asearch(("test",), query="test", limit=2)
|
||||
results_page2 = await vector_store.asearch(
|
||||
("test",), query="test", limit=2, offset=2
|
||||
)
|
||||
|
||||
docs_result = await store.asearch(["test_search", "documents"])
|
||||
assert len(docs_result) == 2
|
||||
assert all([item.namespace[1] == "documents" for item in docs_result]), [
|
||||
item.namespace for item in docs_result
|
||||
]
|
||||
assert len(results_page1) == 2
|
||||
assert len(results_page2) == 2
|
||||
assert results_page1[0].key != results_page2[0].key
|
||||
|
||||
reports_result = await store.asearch(["test_search", "reports"])
|
||||
assert len(reports_result) == 2
|
||||
assert all(item.namespace[1] == "reports" for item in reports_result)
|
||||
all_results = await vector_store.asearch(("test",), query="test", limit=10)
|
||||
assert len(all_results) == 5
|
||||
|
||||
limited_result = await store.asearch(["test_search"], limit=2)
|
||||
assert len(limited_result) == 2
|
||||
offset_result = await store.asearch(["test_search"])
|
||||
assert len(offset_result) == 4
|
||||
|
||||
offset_result = await store.asearch(["test_search"], offset=2)
|
||||
assert len(offset_result) == 2
|
||||
assert all(item not in limited_result for item in offset_result)
|
||||
async def test_vector_search_edge_cases(vector_store: AsyncPostgresStore) -> None:
|
||||
"""Test edge cases in vector search."""
|
||||
await vector_store.aput(("test",), "doc1", {"text": "test document"})
|
||||
|
||||
john_doe_result = await store.asearch(
|
||||
["test_search"], filter={"author": "John Doe"}
|
||||
)
|
||||
assert len(john_doe_result) == 2
|
||||
assert all(item.value["author"] == "John Doe" for item in john_doe_result)
|
||||
perfect_match = await vector_store.asearch(("test",), query="text test document")
|
||||
perfect_score = perfect_match[0].score
|
||||
|
||||
draft_result = await store.asearch(
|
||||
["test_search"], filter={"tags": ["draft"]}
|
||||
)
|
||||
assert len(draft_result) == 2
|
||||
assert all("draft" in item.value["tags"] for item in draft_result)
|
||||
results = await vector_store.asearch(("test",), query="")
|
||||
assert len(results) == 1
|
||||
assert results[0].score is None
|
||||
|
||||
page1 = await store.asearch(["test_search"], limit=2, offset=0)
|
||||
page2 = await store.asearch(["test_search"], limit=2, offset=2)
|
||||
all_items = page1 + page2
|
||||
assert len(all_items) == 4
|
||||
assert len(set(item.key for item in all_items)) == 4
|
||||
empty = await store.asearch(
|
||||
(
|
||||
"scoped",
|
||||
"assistant_id",
|
||||
"shared",
|
||||
"again",
|
||||
"maybe",
|
||||
"some-long",
|
||||
"6be5cb0e-2eb4-42e6-bb6b-fba3c269db25",
|
||||
),
|
||||
limit=10,
|
||||
offset=0,
|
||||
)
|
||||
assert len(empty) == 0
|
||||
results = await vector_store.asearch(("test",), query=None)
|
||||
assert len(results) == 1
|
||||
assert results[0].score is None
|
||||
|
||||
# Test with a namespace beginning with a number (like a UUID)
|
||||
uuid_namespace = (str(uuid.uuid4()), "documents")
|
||||
uuid_item_id = "uuid_doc"
|
||||
uuid_item_value = {
|
||||
"title": "UUID Document",
|
||||
"content": "This document has a UUID namespace.",
|
||||
}
|
||||
long_query = "foo " * 100
|
||||
results = await vector_store.asearch(("test",), query=long_query)
|
||||
assert len(results) == 1
|
||||
assert results[0].score < perfect_score
|
||||
|
||||
# Insert the item with the UUID namespace
|
||||
await store.aput(uuid_namespace, uuid_item_id, uuid_item_value)
|
||||
special_query = "test!@#$%^&*()"
|
||||
results = await vector_store.asearch(("test",), query=special_query)
|
||||
assert len(results) == 1
|
||||
assert results[0].score < perfect_score
|
||||
|
||||
# Retrieve the item to verify it was stored correctly
|
||||
retrieved_item = await store.aget(uuid_namespace, uuid_item_id)
|
||||
assert retrieved_item is not None
|
||||
assert retrieved_item.namespace == uuid_namespace
|
||||
assert retrieved_item.key == uuid_item_id
|
||||
assert retrieved_item.value == uuid_item_value
|
||||
|
||||
# Search for the item using the UUID namespace
|
||||
search_result = await store.asearch([uuid_namespace[0]])
|
||||
assert len(search_result) == 1
|
||||
assert search_result[0].key == uuid_item_id
|
||||
assert search_result[0].value == uuid_item_value
|
||||
@pytest.mark.parametrize(
|
||||
"vector_type,distance_type",
|
||||
[
|
||||
*itertools.product(["vector", "halfvec"], ["cosine", "inner_product", "l2"]),
|
||||
],
|
||||
)
|
||||
async def test_embed_with_path(
|
||||
request: Any,
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
vector_type: str,
|
||||
distance_type: str,
|
||||
) -> None:
|
||||
"""Test vector search with specific text fields in Postgres store."""
|
||||
async with _create_vector_store(
|
||||
vector_type,
|
||||
distance_type,
|
||||
fake_embeddings,
|
||||
text_fields=["key0", "key1", "key3"],
|
||||
) as store:
|
||||
# This will have 2 vectors representing it
|
||||
doc1 = {
|
||||
# Omit key0 - check it doesn't raise an error
|
||||
"key1": "xxx",
|
||||
"key2": "yyy",
|
||||
"key3": "zzz",
|
||||
}
|
||||
# This will have 3 vectors representing it
|
||||
doc2 = {
|
||||
"key0": "uuu",
|
||||
"key1": "vvv",
|
||||
"key2": "www",
|
||||
"key3": "xxx",
|
||||
}
|
||||
await store.aput(("test",), "doc1", doc1)
|
||||
await store.aput(("test",), "doc2", doc2)
|
||||
|
||||
# Clean up: delete the item with the UUID namespace
|
||||
await store.adelete(uuid_namespace, uuid_item_id)
|
||||
# doc2.key3 and doc1.key1 both would have the highest score
|
||||
results = await store.asearch(("test",), query="xxx")
|
||||
assert len(results) == 2
|
||||
assert results[0].key != results[1].key
|
||||
ascore = results[0].score
|
||||
bscore = results[1].score
|
||||
assert ascore == pytest.approx(bscore, abs=1e-3)
|
||||
|
||||
# Verify the item was deleted
|
||||
deleted_item = await store.aget(uuid_namespace, uuid_item_id)
|
||||
assert deleted_item is None
|
||||
results = await store.asearch(("test",), query="uuu")
|
||||
assert len(results) == 2
|
||||
assert results[0].key != results[1].key
|
||||
assert results[0].key == "doc2"
|
||||
assert results[0].score > results[1].score
|
||||
assert ascore == pytest.approx(results[0].score, abs=1e-3)
|
||||
|
||||
for namespace in test_namespaces:
|
||||
await store.adelete(namespace, f"item_{namespace[-1]}")
|
||||
# Un-indexed - will have low results for both. Not zero (because we're projecting)
|
||||
# but less than the above.
|
||||
results = await store.asearch(("test",), query="www")
|
||||
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
|
||||
|
||||
@@ -1,174 +1,125 @@
|
||||
# type: ignore
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Optional
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from conftest import DEFAULT_URI # type: ignore
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from psycopg import Connection
|
||||
|
||||
from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
|
||||
from langgraph.store.base import (
|
||||
GetOp,
|
||||
Item,
|
||||
ListNamespacesOp,
|
||||
MatchCondition,
|
||||
PutOp,
|
||||
SearchOp,
|
||||
)
|
||||
from langgraph.store.postgres import PostgresStore
|
||||
from tests.conftest import (
|
||||
DEFAULT_URI,
|
||||
VECTOR_TYPES,
|
||||
CharacterEmbeddings,
|
||||
)
|
||||
|
||||
|
||||
class MockCursor:
|
||||
def __init__(self, fetch_result: Any) -> None:
|
||||
self.fetch_result = fetch_result
|
||||
self.execute = MagicMock()
|
||||
self.fetchall = MagicMock(return_value=self.fetch_result)
|
||||
@pytest.fixture(scope="function", params=["default", "pipe", "pool"])
|
||||
def store(request) -> PostgresStore:
|
||||
database = f"test_{uuid4().hex[:16]}"
|
||||
uri_parts = DEFAULT_URI.split("/")
|
||||
uri_base = "/".join(uri_parts[:-1])
|
||||
query_params = ""
|
||||
if "?" in uri_parts[-1]:
|
||||
db_name, query_params = uri_parts[-1].split("?", 1)
|
||||
query_params = "?" + query_params
|
||||
|
||||
conn_string = f"{uri_base}/{database}{query_params}"
|
||||
admin_conn_string = DEFAULT_URI
|
||||
|
||||
class MockConnection:
|
||||
def __init__(self) -> None:
|
||||
self.cursor = MagicMock()
|
||||
self.pipeline = MagicMock()
|
||||
with Connection.connect(admin_conn_string, autocommit=True) as conn:
|
||||
conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
with PostgresStore.from_conn_string(conn_string) as store:
|
||||
store.setup()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_connection() -> MockConnection:
|
||||
return MockConnection()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(mock_connection: MockConnection) -> PostgresStore:
|
||||
return PostgresStore(mock_connection)
|
||||
if request.param == "pipe":
|
||||
with PostgresStore.from_conn_string(conn_string, pipeline=True) as store:
|
||||
yield store
|
||||
elif request.param == "pool":
|
||||
with PostgresStore.from_conn_string(
|
||||
conn_string, pool_config={"min_size": 1, "max_size": 10}
|
||||
) as store:
|
||||
yield store
|
||||
else: # default
|
||||
with PostgresStore.from_conn_string(conn_string) as store:
|
||||
yield store
|
||||
finally:
|
||||
with Connection.connect(admin_conn_string, autocommit=True) as conn:
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
def test_batch_order(store: PostgresStore) -> None:
|
||||
mock_connection = store.conn
|
||||
mock_get_cursor = MockCursor(
|
||||
[
|
||||
{
|
||||
"key": "key1",
|
||||
"value": '{"data": "value1"}',
|
||||
"created_at": datetime.now(),
|
||||
"updated_at": datetime.now(),
|
||||
"prefix": "test.foo",
|
||||
},
|
||||
{
|
||||
"key": "key2",
|
||||
"value": '{"data": "value2"}',
|
||||
"created_at": datetime.now(),
|
||||
"updated_at": datetime.now(),
|
||||
"prefix": "test.bar",
|
||||
},
|
||||
]
|
||||
)
|
||||
mock_search_cursor = MockCursor(
|
||||
[
|
||||
{
|
||||
"key": "key1",
|
||||
"value": '{"data": "value1"}',
|
||||
"created_at": datetime.now(),
|
||||
"updated_at": datetime.now(),
|
||||
"prefix": "test.foo",
|
||||
},
|
||||
]
|
||||
)
|
||||
mock_list_namespaces_cursor = MockCursor(
|
||||
[
|
||||
{"truncated_prefix": b"\x01test"},
|
||||
]
|
||||
)
|
||||
|
||||
failures = []
|
||||
|
||||
def cursor_side_effect(binary: bool = False) -> Any:
|
||||
cursor = MagicMock()
|
||||
|
||||
def execute_side_effect(query: str, *params: Any) -> None:
|
||||
# My super sophisticated database.
|
||||
if "SELECT prefix, key, value" in query:
|
||||
cursor.fetchall = mock_search_cursor.fetchall
|
||||
elif "SELECT DISTINCT ON (truncated_prefix)" in query:
|
||||
cursor.fetchall = mock_list_namespaces_cursor.fetchall
|
||||
elif "WHERE prefix = %s AND key" in query:
|
||||
cursor.fetchall = mock_get_cursor.fetchall
|
||||
elif "INSERT INTO " in query:
|
||||
pass
|
||||
else:
|
||||
e = ValueError(f"Unmatched query: {query}")
|
||||
failures.append(e)
|
||||
raise e
|
||||
|
||||
cursor.execute = MagicMock(side_effect=execute_side_effect)
|
||||
return cursor
|
||||
|
||||
mock_connection.cursor.side_effect = cursor_side_effect
|
||||
# Setup test data
|
||||
store.put(("test", "foo"), "key1", {"data": "value1"})
|
||||
store.put(("test", "bar"), "key2", {"data": "value2"})
|
||||
|
||||
ops = [
|
||||
GetOp(namespace=("test",), key="key1"),
|
||||
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
|
||||
GetOp(namespace=("test", "foo"), key="key1"),
|
||||
PutOp(namespace=("test", "bar"), key="key2", value={"data": "value2"}),
|
||||
SearchOp(
|
||||
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
|
||||
),
|
||||
ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0),
|
||||
GetOp(namespace=("test",), key="key3"),
|
||||
]
|
||||
|
||||
results = store.batch(ops)
|
||||
assert not failures
|
||||
assert len(results) == 5
|
||||
assert isinstance(results[0], Item)
|
||||
assert isinstance(results[0].value, dict)
|
||||
assert results[0].value == {"data": "value1"}
|
||||
assert results[0].key == "key1"
|
||||
assert results[1] is None
|
||||
assert results[1] is None # Put operation returns None
|
||||
assert isinstance(results[2], list)
|
||||
assert len(results[2]) == 1
|
||||
assert isinstance(results[3], list)
|
||||
assert results[3] == [("test",)]
|
||||
assert results[4] is None
|
||||
assert len(results[3]) > 0 # Should contain at least our test namespaces
|
||||
assert results[4] is None # Non-existent key returns None
|
||||
|
||||
# Test reordered operations
|
||||
ops_reordered = [
|
||||
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
|
||||
GetOp(namespace=("test",), key="key2"),
|
||||
GetOp(namespace=("test", "bar"), key="key2"),
|
||||
ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0),
|
||||
PutOp(namespace=("test",), key="key3", value={"data": "value3"}),
|
||||
GetOp(namespace=("test",), key="key1"),
|
||||
GetOp(namespace=("test", "foo"), key="key1"),
|
||||
]
|
||||
|
||||
results_reordered = store.batch(ops_reordered)
|
||||
assert not failures
|
||||
assert len(results_reordered) == 5
|
||||
assert isinstance(results_reordered[0], list)
|
||||
assert len(results_reordered[0]) == 1
|
||||
assert len(results_reordered[0]) >= 2 # Should find at least our two test items
|
||||
assert isinstance(results_reordered[1], Item)
|
||||
assert results_reordered[1].value == {"data": "value2"}
|
||||
assert results_reordered[1].key == "key2"
|
||||
assert isinstance(results_reordered[2], list)
|
||||
assert results_reordered[2] == [("test",)]
|
||||
assert results_reordered[3] is None
|
||||
assert len(results_reordered[2]) > 0
|
||||
assert results_reordered[3] is None # Put operation returns None
|
||||
assert isinstance(results_reordered[4], Item)
|
||||
assert results_reordered[4].value == {"data": "value1"}
|
||||
assert results_reordered[4].key == "key1"
|
||||
|
||||
|
||||
def test_batch_get_ops(store: PostgresStore) -> None:
|
||||
mock_connection = store.conn
|
||||
mock_cursor = MockCursor(
|
||||
[
|
||||
{
|
||||
"key": "key1",
|
||||
"value": '{"data": "value1"}',
|
||||
"created_at": datetime.now(),
|
||||
"updated_at": datetime.now(),
|
||||
"prefix": "test.foo",
|
||||
},
|
||||
{
|
||||
"key": "key2",
|
||||
"value": '{"data": "value2"}',
|
||||
"created_at": datetime.now(),
|
||||
"updated_at": datetime.now(),
|
||||
"prefix": "test.bar",
|
||||
},
|
||||
]
|
||||
)
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
# Setup test data
|
||||
store.put(("test",), "key1", {"data": "value1"})
|
||||
store.put(("test",), "key2", {"data": "value2"})
|
||||
|
||||
ops = [
|
||||
GetOp(namespace=("test",), key="key1"),
|
||||
GetOp(namespace=("test",), key="key2"),
|
||||
GetOp(namespace=("test",), key="key3"),
|
||||
GetOp(namespace=("test",), key="key3"), # Non-existent key
|
||||
]
|
||||
|
||||
results = store.batch(ops)
|
||||
@@ -182,75 +133,90 @@ def test_batch_get_ops(store: PostgresStore) -> None:
|
||||
|
||||
|
||||
def test_batch_put_ops(store: PostgresStore) -> None:
|
||||
mock_connection = store.conn
|
||||
mock_cursor = MockCursor([])
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
|
||||
ops = [
|
||||
PutOp(namespace=("test",), key="key1", value={"data": "value1"}),
|
||||
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
|
||||
PutOp(namespace=("test",), key="key3", value=None),
|
||||
PutOp(namespace=("test",), key="key3", value=None), # Delete operation
|
||||
]
|
||||
|
||||
results = store.batch(ops)
|
||||
|
||||
assert len(results) == 3
|
||||
assert all(result is None for result in results)
|
||||
assert mock_cursor.execute.call_count == 2
|
||||
|
||||
# Verify the puts worked
|
||||
item1 = store.get(("test",), "key1")
|
||||
item2 = store.get(("test",), "key2")
|
||||
item3 = store.get(("test",), "key3")
|
||||
|
||||
assert item1 and item1.value == {"data": "value1"}
|
||||
assert item2 and item2.value == {"data": "value2"}
|
||||
assert item3 is None
|
||||
|
||||
|
||||
def test_batch_search_ops(store: PostgresStore) -> None:
|
||||
mock_connection = store.conn
|
||||
mock_cursor = MockCursor(
|
||||
[
|
||||
{
|
||||
"key": "key1",
|
||||
"value": '{"data": "value1"}',
|
||||
"created_at": datetime.now(),
|
||||
"updated_at": datetime.now(),
|
||||
"prefix": "test.foo",
|
||||
},
|
||||
{
|
||||
"key": "key2",
|
||||
"value": '{"data": "value2"}',
|
||||
"created_at": datetime.now(),
|
||||
"updated_at": datetime.now(),
|
||||
"prefix": "test.bar",
|
||||
},
|
||||
]
|
||||
)
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
# Setup test data
|
||||
test_data = [
|
||||
(("test", "foo"), "key1", {"data": "value1", "tag": "a"}),
|
||||
(("test", "bar"), "key2", {"data": "value2", "tag": "a"}),
|
||||
(("test", "baz"), "key3", {"data": "value3", "tag": "b"}),
|
||||
]
|
||||
for namespace, key, value in test_data:
|
||||
store.put(namespace, key, value)
|
||||
|
||||
ops = [
|
||||
SearchOp(
|
||||
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
|
||||
),
|
||||
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
|
||||
SearchOp(namespace_prefix=("test",), filter={"tag": "a"}, limit=10, offset=0),
|
||||
SearchOp(namespace_prefix=("test",), filter=None, limit=2, offset=0),
|
||||
SearchOp(namespace_prefix=("test", "foo"), filter=None, limit=10, offset=0),
|
||||
]
|
||||
|
||||
results = store.batch(ops)
|
||||
assert len(results) == 3
|
||||
|
||||
assert len(results) == 2
|
||||
# First search should find items with tag "a"
|
||||
assert len(results[0]) == 2
|
||||
assert all(item.value["tag"] == "a" for item in results[0])
|
||||
|
||||
# Second search should return first 2 items
|
||||
assert len(results[1]) == 2
|
||||
|
||||
# Third search should only find items in test/foo namespace
|
||||
assert len(results[2]) == 1
|
||||
assert results[2][0].namespace == ("test", "foo")
|
||||
|
||||
|
||||
def test_batch_list_namespaces_ops(store: PostgresStore) -> None:
|
||||
mock_connection = store.conn
|
||||
mock_cursor = MockCursor(
|
||||
[
|
||||
{"truncated_prefix": b"\x01test.namespace1"},
|
||||
{"truncated_prefix": b"\x01test.namespace2"},
|
||||
]
|
||||
)
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
# Setup test data with various namespaces
|
||||
test_data = [
|
||||
(("test", "documents", "public"), "doc1", {"content": "public doc"}),
|
||||
(("test", "documents", "private"), "doc2", {"content": "private doc"}),
|
||||
(("test", "images", "public"), "img1", {"content": "public image"}),
|
||||
(("prod", "documents", "public"), "doc3", {"content": "prod doc"}),
|
||||
]
|
||||
for namespace, key, value in test_data:
|
||||
store.put(namespace, key, value)
|
||||
|
||||
ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)]
|
||||
ops = [
|
||||
ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0),
|
||||
ListNamespacesOp(match_conditions=None, max_depth=2, limit=10, offset=0),
|
||||
ListNamespacesOp(
|
||||
match_conditions=[MatchCondition("suffix", "public")],
|
||||
max_depth=None,
|
||||
limit=10,
|
||||
offset=0,
|
||||
),
|
||||
]
|
||||
|
||||
results = store.batch(ops)
|
||||
assert len(results) == 3
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0] == [("test", "namespace1"), ("test", "namespace2")]
|
||||
# First operation should list all namespaces
|
||||
assert len(results[0]) == len(test_data)
|
||||
|
||||
# Second operation should only return namespaces up to depth 2
|
||||
assert all(len(ns) <= 2 for ns in results[1])
|
||||
|
||||
# Third operation should only return namespaces ending with "public"
|
||||
assert all(ns[-1] == "public" for ns in results[2])
|
||||
|
||||
|
||||
class TestPostgresStore:
|
||||
@@ -273,195 +239,459 @@ class TestPostgresStore:
|
||||
assert item.key == item_id
|
||||
assert item.value == item_value
|
||||
|
||||
updated_value = {
|
||||
"title": "Updated Test Document",
|
||||
"content": "Hello, LangGraph!",
|
||||
}
|
||||
# Test update
|
||||
updated_value = {"title": "Updated Document", "content": "Hello, Updated!"}
|
||||
store.put(namespace, item_id, updated_value)
|
||||
updated_item = store.get(namespace, item_id)
|
||||
|
||||
assert updated_item.value == updated_value
|
||||
assert updated_item.updated_at > item.updated_at
|
||||
|
||||
# Test get from non-existent namespace
|
||||
different_namespace = ("test", "other_documents")
|
||||
item_in_different_namespace = store.get(different_namespace, item_id)
|
||||
assert item_in_different_namespace is None
|
||||
|
||||
new_item_id = "doc2"
|
||||
new_item_value = {"title": "Another Document", "content": "Greetings!"}
|
||||
store.put(namespace, new_item_id, new_item_value)
|
||||
|
||||
search_results = store.search(["test"], limit=10)
|
||||
items = search_results
|
||||
assert len(items) == 2
|
||||
assert any(item.key == item_id for item in items)
|
||||
assert any(item.key == new_item_id for item in items)
|
||||
|
||||
namespaces = store.list_namespaces(prefix=["test"])
|
||||
assert ("test", "documents") in namespaces
|
||||
|
||||
# Test delete
|
||||
store.delete(namespace, item_id)
|
||||
store.delete(namespace, new_item_id)
|
||||
deleted_item = store.get(namespace, item_id)
|
||||
assert deleted_item is None
|
||||
|
||||
deleted_item = store.get(namespace, new_item_id)
|
||||
assert deleted_item is None
|
||||
|
||||
empty_search_results = store.search(["test"], limit=10)
|
||||
assert len(empty_search_results) == 0
|
||||
|
||||
def test_list_namespaces(self) -> None:
|
||||
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
|
||||
test_pref = str(uuid.uuid4())
|
||||
# Create test data with various namespaces
|
||||
test_namespaces = [
|
||||
(test_pref, "test", "documents", "public", test_pref),
|
||||
(test_pref, "test", "documents", "private", test_pref),
|
||||
(test_pref, "test", "images", "public", test_pref),
|
||||
(test_pref, "test", "images", "private", test_pref),
|
||||
(test_pref, "prod", "documents", "public", test_pref),
|
||||
(
|
||||
test_pref,
|
||||
"prod",
|
||||
"documents",
|
||||
"some",
|
||||
"nesting",
|
||||
"public",
|
||||
test_pref,
|
||||
),
|
||||
(test_pref, "prod", "documents", "private", test_pref),
|
||||
("test", "documents", "public"),
|
||||
("test", "documents", "private"),
|
||||
("test", "images", "public"),
|
||||
("test", "images", "private"),
|
||||
("prod", "documents", "public"),
|
||||
("prod", "documents", "private"),
|
||||
]
|
||||
|
||||
# Insert test data
|
||||
for namespace in test_namespaces:
|
||||
store.put(namespace, "dummy", {"content": "dummy"})
|
||||
|
||||
prefix_result = store.list_namespaces(prefix=[test_pref, "test"])
|
||||
assert len(prefix_result) == 4
|
||||
assert all([ns[1] == "test" for ns in prefix_result])
|
||||
# Test listing with various filters
|
||||
all_namespaces = store.list_namespaces()
|
||||
assert len(all_namespaces) == len(test_namespaces)
|
||||
|
||||
specific_prefix_result = store.list_namespaces(
|
||||
prefix=[test_pref, "test", "documents"]
|
||||
)
|
||||
assert len(specific_prefix_result) == 2
|
||||
assert all(
|
||||
[ns[1:3] == ("test", "documents") for ns in specific_prefix_result]
|
||||
)
|
||||
# Test prefix filtering
|
||||
test_prefix_namespaces = store.list_namespaces(prefix=["test"])
|
||||
assert len(test_prefix_namespaces) == 4
|
||||
assert all(ns[0] == "test" for ns in test_prefix_namespaces)
|
||||
|
||||
suffix_result = store.list_namespaces(suffix=["public", test_pref])
|
||||
assert len(suffix_result) == 4
|
||||
assert all(ns[-2] == "public" for ns in suffix_result)
|
||||
# Test suffix filtering
|
||||
public_namespaces = store.list_namespaces(suffix=["public"])
|
||||
assert len(public_namespaces) == 3
|
||||
assert all(ns[-1] == "public" for ns in public_namespaces)
|
||||
|
||||
prefix_suffix_result = store.list_namespaces(
|
||||
prefix=[test_pref, "test"], suffix=["public", test_pref]
|
||||
)
|
||||
assert len(prefix_suffix_result) == 2
|
||||
assert all(
|
||||
ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result
|
||||
)
|
||||
# Test max depth
|
||||
depth_2_namespaces = store.list_namespaces(max_depth=2)
|
||||
assert all(len(ns) <= 2 for ns in depth_2_namespaces)
|
||||
|
||||
wildcard_prefix_result = store.list_namespaces(
|
||||
prefix=[test_pref, "*", "documents"]
|
||||
)
|
||||
assert len(wildcard_prefix_result) == 5
|
||||
assert all(ns[2] == "documents" for ns in wildcard_prefix_result)
|
||||
|
||||
wildcard_suffix_result = store.list_namespaces(
|
||||
suffix=["*", "public", test_pref]
|
||||
)
|
||||
assert len(wildcard_suffix_result) == 4
|
||||
assert all(ns[-2] == "public" for ns in wildcard_suffix_result)
|
||||
wildcard_single = store.list_namespaces(
|
||||
suffix=["some", "*", "public", test_pref]
|
||||
)
|
||||
assert len(wildcard_single) == 1
|
||||
assert wildcard_single[0] == (
|
||||
test_pref,
|
||||
"prod",
|
||||
"documents",
|
||||
"some",
|
||||
"nesting",
|
||||
"public",
|
||||
test_pref,
|
||||
)
|
||||
|
||||
max_depth_result = store.list_namespaces(max_depth=3)
|
||||
assert all([len(ns) <= 3 for ns in max_depth_result])
|
||||
|
||||
max_depth_result = store.list_namespaces(
|
||||
max_depth=4, prefix=[test_pref, "*", "documents"]
|
||||
)
|
||||
assert (
|
||||
len(set(tuple(res) for res in max_depth_result))
|
||||
== len(max_depth_result)
|
||||
== 5
|
||||
)
|
||||
|
||||
limit_result = store.list_namespaces(prefix=[test_pref], limit=3)
|
||||
assert len(limit_result) == 3
|
||||
|
||||
offset_result = store.list_namespaces(prefix=[test_pref], offset=3)
|
||||
assert len(offset_result) == len(test_namespaces) - 3
|
||||
|
||||
empty_prefix_result = store.list_namespaces(prefix=[test_pref])
|
||||
assert len(empty_prefix_result) == len(test_namespaces)
|
||||
assert set(tuple(ns) for ns in empty_prefix_result) == set(
|
||||
tuple(ns) for ns in test_namespaces
|
||||
)
|
||||
# Test pagination
|
||||
paginated_namespaces = store.list_namespaces(limit=3)
|
||||
assert len(paginated_namespaces) == 3
|
||||
|
||||
# Cleanup
|
||||
for namespace in test_namespaces:
|
||||
store.delete(namespace, "dummy")
|
||||
|
||||
def test_search(self):
|
||||
def test_search(self) -> None:
|
||||
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
|
||||
test_namespaces = [
|
||||
("test_search", "documents", "user1"),
|
||||
("test_search", "documents", "user2"),
|
||||
("test_search", "reports", "department1"),
|
||||
("test_search", "reports", "department2"),
|
||||
]
|
||||
test_items = [
|
||||
{"title": "Doc 1", "author": "John Doe", "tags": ["important"]},
|
||||
{"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]},
|
||||
{"title": "Report A", "author": "John Doe", "tags": ["final"]},
|
||||
{"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]},
|
||||
# Create test data
|
||||
test_data = [
|
||||
(
|
||||
("test", "docs"),
|
||||
"doc1",
|
||||
{"title": "First Doc", "author": "Alice", "tags": ["important"]},
|
||||
),
|
||||
(
|
||||
("test", "docs"),
|
||||
"doc2",
|
||||
{"title": "Second Doc", "author": "Bob", "tags": ["draft"]},
|
||||
),
|
||||
(
|
||||
("test", "images"),
|
||||
"img1",
|
||||
{"title": "Image 1", "author": "Alice", "tags": ["final"]},
|
||||
),
|
||||
]
|
||||
|
||||
for namespace, item in zip(test_namespaces, test_items):
|
||||
store.put(namespace, f"item_{namespace[-1]}", item)
|
||||
for namespace, key, value in test_data:
|
||||
store.put(namespace, key, value)
|
||||
|
||||
docs_result = store.search(["test_search", "documents"])
|
||||
assert len(docs_result) == 2
|
||||
assert all(
|
||||
[item.namespace[1] == "documents" for item in docs_result]
|
||||
), docs_result
|
||||
# Test basic search
|
||||
all_items = store.search(["test"])
|
||||
assert len(all_items) == 3
|
||||
|
||||
reports_result = store.search(["test_search", "reports"])
|
||||
assert len(reports_result) == 2
|
||||
assert all(item.namespace[1] == "reports" for item in reports_result)
|
||||
# Test namespace filtering
|
||||
docs_items = store.search(["test", "docs"])
|
||||
assert len(docs_items) == 2
|
||||
assert all(item.namespace == ("test", "docs") for item in docs_items)
|
||||
|
||||
limited_result = store.search(["test_search"], limit=2)
|
||||
assert len(limited_result) == 2
|
||||
offset_result = store.search(["test_search"])
|
||||
assert len(offset_result) == 4
|
||||
# Test value filtering
|
||||
alice_items = store.search(["test"], filter={"author": "Alice"})
|
||||
assert len(alice_items) == 2
|
||||
assert all(item.value["author"] == "Alice" for item in alice_items)
|
||||
|
||||
offset_result = store.search(["test_search"], offset=2)
|
||||
assert len(offset_result) == 2
|
||||
assert all(item not in limited_result for item in offset_result)
|
||||
# Test pagination
|
||||
paginated_items = store.search(["test"], limit=2)
|
||||
assert len(paginated_items) == 2
|
||||
|
||||
john_doe_result = store.search(
|
||||
["test_search"], filter={"author": "John Doe"}
|
||||
)
|
||||
assert len(john_doe_result) == 2
|
||||
assert all(item.value["author"] == "John Doe" for item in john_doe_result)
|
||||
offset_items = store.search(["test"], offset=2)
|
||||
assert len(offset_items) == 1
|
||||
|
||||
draft_result = store.search(["test_search"], filter={"tags": ["draft"]})
|
||||
assert len(draft_result) == 2
|
||||
assert all("draft" in item.value["tags"] for item in draft_result)
|
||||
# Cleanup
|
||||
for namespace, key, _ in test_data:
|
||||
store.delete(namespace, key)
|
||||
|
||||
page1 = store.search(["test_search"], limit=2, offset=0)
|
||||
page2 = store.search(["test_search"], limit=2, offset=2)
|
||||
all_items = page1 + page2
|
||||
assert len(all_items) == 4
|
||||
assert len(set(item.key for item in all_items)) == 4
|
||||
|
||||
for namespace in test_namespaces:
|
||||
store.delete(namespace, f"item_{namespace[-1]}")
|
||||
@contextmanager
|
||||
def _create_vector_store(
|
||||
vector_type: str,
|
||||
distance_type: str,
|
||||
fake_embeddings: Embeddings,
|
||||
text_fields: Optional[list[str]] = None,
|
||||
) -> PostgresStore:
|
||||
"""Create a store with vector search enabled."""
|
||||
database = f"test_{uuid4().hex[:16]}"
|
||||
uri_parts = DEFAULT_URI.split("/")
|
||||
uri_base = "/".join(uri_parts[:-1])
|
||||
query_params = ""
|
||||
if "?" in uri_parts[-1]:
|
||||
db_name, query_params = uri_parts[-1].split("?", 1)
|
||||
query_params = "?" + query_params
|
||||
|
||||
conn_string = f"{uri_base}/{database}{query_params}"
|
||||
admin_conn_string = DEFAULT_URI
|
||||
|
||||
index_config = {
|
||||
"dims": fake_embeddings.dims,
|
||||
"embed": fake_embeddings,
|
||||
"ann_index_config": {
|
||||
"vector_type": vector_type,
|
||||
},
|
||||
"distance_type": distance_type,
|
||||
"text_fields": text_fields,
|
||||
}
|
||||
|
||||
with Connection.connect(admin_conn_string, autocommit=True) as conn:
|
||||
conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
with PostgresStore.from_conn_string(
|
||||
conn_string,
|
||||
index=index_config,
|
||||
) as store:
|
||||
store.setup()
|
||||
yield store
|
||||
finally:
|
||||
with Connection.connect(admin_conn_string, autocommit=True) as conn:
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
scope="function",
|
||||
params=[
|
||||
(vector_type, distance_type)
|
||||
for vector_type in VECTOR_TYPES
|
||||
for distance_type in (
|
||||
["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
|
||||
)
|
||||
],
|
||||
ids=lambda p: f"{p[0]}_{p[1]}",
|
||||
)
|
||||
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:
|
||||
yield store
|
||||
|
||||
|
||||
def test_vector_store_initialization(
|
||||
vector_store: PostgresStore, fake_embeddings: CharacterEmbeddings
|
||||
) -> None:
|
||||
"""Test store initialization with embedding config."""
|
||||
# Store should be initialized with embedding config
|
||||
assert vector_store.index_config is not None
|
||||
assert vector_store.index_config["dims"] == fake_embeddings.dims
|
||||
assert vector_store.index_config["embed"] == fake_embeddings
|
||||
|
||||
|
||||
def test_vector_insert_with_auto_embedding(vector_store: PostgresStore) -> None:
|
||||
"""Test inserting items that get auto-embedded."""
|
||||
docs = [
|
||||
("doc1", {"text": "short text"}),
|
||||
("doc2", {"text": "longer text document"}),
|
||||
("doc3", {"text": "longest text document here"}),
|
||||
("doc4", {"description": "text in description field"}),
|
||||
("doc5", {"content": "text in content field"}),
|
||||
("doc6", {"body": "text in body field"}),
|
||||
]
|
||||
|
||||
for key, value in docs:
|
||||
vector_store.put(("test",), key, value)
|
||||
|
||||
results = vector_store.search(("test",), query="long text")
|
||||
assert len(results) > 0
|
||||
|
||||
doc_order = [r.key for r in results]
|
||||
assert "doc2" in doc_order
|
||||
assert "doc3" in doc_order
|
||||
|
||||
|
||||
def test_vector_update_with_embedding(vector_store: PostgresStore) -> None:
|
||||
"""Test that updating items properly updates their embeddings."""
|
||||
vector_store.put(("test",), "doc1", {"text": "zany zebra Xerxes"})
|
||||
vector_store.put(("test",), "doc2", {"text": "something about dogs"})
|
||||
vector_store.put(("test",), "doc3", {"text": "text about birds"})
|
||||
|
||||
results_initial = vector_store.search(("test",), query="Zany Xerxes")
|
||||
assert len(results_initial) > 0
|
||||
assert results_initial[0].key == "doc1"
|
||||
initial_score = results_initial[0].score
|
||||
|
||||
vector_store.put(("test",), "doc1", {"text": "new text about dogs"})
|
||||
|
||||
results_after = vector_store.search(("test",), query="Zany Xerxes")
|
||||
after_score = next((r.score for r in results_after if r.key == "doc1"), 0.0)
|
||||
assert after_score < initial_score
|
||||
|
||||
results_new = vector_store.search(("test",), query="new text about dogs")
|
||||
for r in results_new:
|
||||
if r.key == "doc1":
|
||||
assert r.score > after_score
|
||||
|
||||
# Don't index this one
|
||||
vector_store.put(("test",), "doc4", {"text": "new text about dogs"}, index=False)
|
||||
results_new = vector_store.search(("test",), query="new text about dogs", limit=3)
|
||||
assert not any(r.key == "doc4" for r in results_new)
|
||||
|
||||
|
||||
def test_vector_search_with_filters(vector_store: PostgresStore) -> None:
|
||||
"""Test combining vector search with filters."""
|
||||
# Insert test documents
|
||||
docs = [
|
||||
("doc1", {"text": "red apple", "color": "red", "score": 4.5}),
|
||||
("doc2", {"text": "red car", "color": "red", "score": 3.0}),
|
||||
("doc3", {"text": "green apple", "color": "green", "score": 4.0}),
|
||||
("doc4", {"text": "blue car", "color": "blue", "score": 3.5}),
|
||||
]
|
||||
|
||||
for key, value in docs:
|
||||
vector_store.put(("test",), key, value)
|
||||
|
||||
results = vector_store.search(("test",), query="apple", filter={"color": "red"})
|
||||
assert len(results) == 2
|
||||
assert results[0].key == "doc1"
|
||||
|
||||
results = vector_store.search(("test",), query="car", filter={"color": "red"})
|
||||
assert len(results) == 2
|
||||
assert results[0].key == "doc2"
|
||||
|
||||
results = vector_store.search(
|
||||
("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
|
||||
)
|
||||
assert len(results) == 3
|
||||
assert results[0].key == "doc4"
|
||||
|
||||
# Multiple filters
|
||||
results = vector_store.search(
|
||||
("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"}
|
||||
)
|
||||
assert len(results) == 1
|
||||
assert results[0].key == "doc3"
|
||||
|
||||
|
||||
def test_vector_search_pagination(vector_store: PostgresStore) -> None:
|
||||
"""Test pagination with vector search."""
|
||||
# Insert multiple similar documents
|
||||
for i in range(5):
|
||||
vector_store.put(("test",), f"doc{i}", {"text": f"test document number {i}"})
|
||||
|
||||
# Test with different page sizes
|
||||
results_page1 = vector_store.search(("test",), query="test", limit=2)
|
||||
results_page2 = vector_store.search(("test",), query="test", limit=2, offset=2)
|
||||
|
||||
assert len(results_page1) == 2
|
||||
assert len(results_page2) == 2
|
||||
assert results_page1[0].key != results_page2[0].key
|
||||
|
||||
# Get all results
|
||||
all_results = vector_store.search(("test",), query="test", limit=10)
|
||||
assert len(all_results) == 5
|
||||
|
||||
|
||||
def test_vector_search_edge_cases(vector_store: PostgresStore) -> None:
|
||||
"""Test edge cases in vector search."""
|
||||
vector_store.put(("test",), "doc1", {"text": "test document"})
|
||||
|
||||
results = vector_store.search(("test",), query="")
|
||||
assert len(results) == 1
|
||||
|
||||
results = vector_store.search(("test",), query=None)
|
||||
assert len(results) == 1
|
||||
|
||||
long_query = "test " * 100
|
||||
results = vector_store.search(("test",), query=long_query)
|
||||
assert len(results) == 1
|
||||
|
||||
special_query = "test!@#$%^&*()"
|
||||
results = vector_store.search(("test",), query=special_query)
|
||||
assert len(results) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"vector_type,distance_type",
|
||||
[
|
||||
("vector", "cosine"),
|
||||
("vector", "inner_product"),
|
||||
("halfvec", "cosine"),
|
||||
("halfvec", "inner_product"),
|
||||
],
|
||||
)
|
||||
def test_embed_with_path_sync(
|
||||
request: Any,
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
vector_type: str,
|
||||
distance_type: str,
|
||||
) -> None:
|
||||
"""Test vector search with specific text fields in Postgres store."""
|
||||
with _create_vector_store(
|
||||
vector_type,
|
||||
distance_type,
|
||||
fake_embeddings,
|
||||
text_fields=["key0", "key1", "key3"],
|
||||
) as store:
|
||||
# This will have 2 vectors representing it
|
||||
doc1 = {
|
||||
# Omit key0 - check it doesn't raise an error
|
||||
"key1": "xxx",
|
||||
"key2": "yyy",
|
||||
"key3": "zzz",
|
||||
}
|
||||
# This will have 3 vectors representing it
|
||||
doc2 = {
|
||||
"key0": "uuu",
|
||||
"key1": "vvv",
|
||||
"key2": "www",
|
||||
"key3": "xxx",
|
||||
}
|
||||
store.put(("test",), "doc1", doc1)
|
||||
store.put(("test",), "doc2", doc2)
|
||||
|
||||
# doc2.key3 and doc1.key1 both would have the highest score
|
||||
results = store.search(("test",), query="xxx")
|
||||
assert len(results) == 2
|
||||
assert results[0].key != results[1].key
|
||||
ascore = results[0].score
|
||||
bscore = results[1].score
|
||||
assert ascore == pytest.approx(bscore, abs=1e-3)
|
||||
|
||||
# ~Only match doc2
|
||||
results = store.search(("test",), query="uuu")
|
||||
assert len(results) == 2
|
||||
assert results[0].key != results[1].key
|
||||
assert results[0].key == "doc2"
|
||||
assert results[0].score > results[1].score
|
||||
assert ascore == pytest.approx(results[0].score, abs=1e-3)
|
||||
|
||||
# ~Only match doc1
|
||||
results = store.search(("test",), query="zzz")
|
||||
assert len(results) == 2
|
||||
assert results[0].key != results[1].key
|
||||
assert results[0].key == "doc1"
|
||||
assert results[0].score > results[1].score
|
||||
assert ascore == pytest.approx(results[0].score, abs=1e-3)
|
||||
|
||||
# Un-indexed - will have low results for both. Not zero (because we're projecting)
|
||||
# but less than the above.
|
||||
results = store.search(("test",), query="www")
|
||||
assert len(results) == 2
|
||||
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
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from conftest import DEFAULT_URI # type: ignore
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
@@ -11,6 +10,7 @@ from langgraph.checkpoint.base import (
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
from tests.conftest import DEFAULT_URI
|
||||
|
||||
|
||||
class TestPostgresSaver:
|
||||
|
||||
@@ -4,11 +4,13 @@
|
||||
# TESTING AND COVERAGE
|
||||
######################
|
||||
|
||||
TEST ?= .
|
||||
|
||||
test:
|
||||
poetry run pytest tests
|
||||
poetry run pytest $(TEST)
|
||||
|
||||
test_watch:
|
||||
poetry run ptw .
|
||||
poetry run ptw $(TEST)
|
||||
|
||||
######################
|
||||
# LINTING AND FORMATTING
|
||||
|
||||
@@ -372,7 +372,7 @@ class MemorySaver(
|
||||
RunnableConfig: The updated config containing the saved writes' timestamp.
|
||||
"""
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"]["checkpoint_ns"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = config["configurable"]["checkpoint_id"]
|
||||
outer_key = (thread_id, checkpoint_ns, checkpoint_id)
|
||||
outer_writes_ = self.writes.get(outer_key)
|
||||
|
||||
@@ -1,12 +1,27 @@
|
||||
"""Base classes and types for persistent key-value stores.
|
||||
|
||||
Stores enable persistence and memory that can be shared across threads,
|
||||
scoped to user IDs, assistant IDs, or other arbitrary namespaces.
|
||||
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
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime
|
||||
from typing import Any, Iterable, Literal, NamedTuple, Optional, Union, cast
|
||||
from typing import Any, Iterable, Literal, NamedTuple, Optional, TypedDict, Union, cast
|
||||
|
||||
from langchain_core.embeddings import Embeddings
|
||||
|
||||
from langgraph.store.base.embed import (
|
||||
AEmbeddingsFunc,
|
||||
EmbeddingsFunc,
|
||||
ensure_embeddings,
|
||||
get_text_at_path,
|
||||
tokenize_path,
|
||||
)
|
||||
|
||||
|
||||
class Item:
|
||||
@@ -73,112 +88,415 @@ class Item:
|
||||
}
|
||||
|
||||
|
||||
class SearchItem(Item):
|
||||
"""Represents a result item with additional response metadata."""
|
||||
|
||||
__slots__ = ("score",)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
namespace: tuple[str, ...],
|
||||
key: str,
|
||||
value: dict[str, Any],
|
||||
created_at: datetime,
|
||||
updated_at: datetime,
|
||||
score: Optional[float] = None,
|
||||
) -> None:
|
||||
"""Initialize a result item.
|
||||
|
||||
Args:
|
||||
namespace: Hierarchical path to the item.
|
||||
key: Unique identifier within the namespace.
|
||||
value: The stored value.
|
||||
created_at: When the item was first created.
|
||||
updated_at: When the item was last updated.
|
||||
score: Relevance/similarity score if from a ranked operation.
|
||||
"""
|
||||
super().__init__(
|
||||
value=value,
|
||||
key=key,
|
||||
namespace=namespace,
|
||||
created_at=created_at,
|
||||
updated_at=updated_at,
|
||||
)
|
||||
self.score = score
|
||||
|
||||
def dict(self) -> dict:
|
||||
result = super().dict()
|
||||
result["score"] = self.score
|
||||
return result
|
||||
|
||||
|
||||
class GetOp(NamedTuple):
|
||||
"""Operation to retrieve an item by namespace and key."""
|
||||
"""Operation to retrieve a specific item by its namespace and key.
|
||||
|
||||
This operation allows precise retrieval of stored items using their full path
|
||||
(namespace) and unique identifier (key) combination.
|
||||
|
||||
??? example "Examples"
|
||||
|
||||
Basic item retrieval:
|
||||
```python
|
||||
GetOp(namespace=("users", "profiles"), key="user123")
|
||||
GetOp(namespace=("cache", "embeddings"), key="doc456")
|
||||
```
|
||||
"""
|
||||
|
||||
namespace: tuple[str, ...]
|
||||
"""Hierarchical path for the item."""
|
||||
"""Hierarchical path that uniquely identifies the item's location.
|
||||
|
||||
??? example "Examples"
|
||||
|
||||
```python
|
||||
("users",) # Root level users namespace
|
||||
("users", "profiles") # Profiles within users namespace
|
||||
```
|
||||
"""
|
||||
|
||||
key: str
|
||||
"""Unique identifier within the namespace."""
|
||||
"""Unique identifier for the item within its specific namespace.
|
||||
|
||||
??? example "Examples"
|
||||
|
||||
```python
|
||||
"user123" # For a user profile
|
||||
"doc456" # For a document
|
||||
```
|
||||
"""
|
||||
|
||||
|
||||
class SearchOp(NamedTuple):
|
||||
"""Operation to search for items within a namespace prefix."""
|
||||
"""Operation to search for items within a specified namespace hierarchy.
|
||||
|
||||
This operation supports both structured filtering and natural language search
|
||||
within a given namespace prefix. It provides pagination through limit and offset
|
||||
parameters.
|
||||
|
||||
Note:
|
||||
Natural language search support depends on your store implementation.
|
||||
|
||||
??? example "Examples"
|
||||
Search with filters and pagination:
|
||||
```python
|
||||
SearchOp(
|
||||
namespace_prefix=("documents",),
|
||||
filter={"type": "report", "status": "active"},
|
||||
limit=5,
|
||||
offset=10
|
||||
)
|
||||
```
|
||||
|
||||
Natural language search:
|
||||
```python
|
||||
SearchOp(
|
||||
namespace_prefix=("users", "content"),
|
||||
query="technical documentation about APIs",
|
||||
limit=20
|
||||
)
|
||||
```
|
||||
"""
|
||||
|
||||
namespace_prefix: tuple[str, ...]
|
||||
"""Hierarchical path prefix to search within."""
|
||||
"""Hierarchical path prefix defining the search scope.
|
||||
|
||||
??? example "Examples"
|
||||
|
||||
```python
|
||||
() # Search entire store
|
||||
("documents",) # Search all documents
|
||||
("users", "content") # Search within user content
|
||||
```
|
||||
"""
|
||||
|
||||
filter: Optional[dict[str, Any]] = None
|
||||
"""Key-value pairs to filter results."""
|
||||
"""Key-value pairs for filtering results based on exact matches or comparison operators.
|
||||
|
||||
The filter supports both exact matches and operator-based comparisons.
|
||||
|
||||
Supported Operators:
|
||||
- $eq: Equal to (same as direct value comparison)
|
||||
- $ne: Not equal to
|
||||
- $gt: Greater than
|
||||
- $gte: Greater than or equal to
|
||||
- $lt: Less than
|
||||
- $lte: Less than or equal to
|
||||
|
||||
??? example "Examples"
|
||||
|
||||
Simple exact match:
|
||||
|
||||
```python
|
||||
{"status": "active"}
|
||||
```
|
||||
|
||||
Comparison operators:
|
||||
|
||||
```python
|
||||
{"score": {"$gt": 4.99}} # Score greater than 4.99
|
||||
```
|
||||
|
||||
Multiple conditions:
|
||||
|
||||
```python
|
||||
{
|
||||
"score": {"$gte": 3.0},
|
||||
"color": "red"
|
||||
}
|
||||
```
|
||||
|
||||
Note:
|
||||
Comparison operator support depends on your store implementation.
|
||||
"""
|
||||
|
||||
limit: int = 10
|
||||
"""Maximum number of items to return."""
|
||||
"""Maximum number of items to return in the search results."""
|
||||
|
||||
offset: int = 0
|
||||
"""Number of items to skip before returning results."""
|
||||
"""Number of matching items to skip for pagination."""
|
||||
|
||||
query: Optional[str] = None
|
||||
"""Natural language search query for semantic search capabilities.
|
||||
|
||||
class PutOp(NamedTuple):
|
||||
"""Operation to store, update, or delete an item."""
|
||||
|
||||
namespace: tuple[str, ...]
|
||||
"""Hierarchical path for the item.
|
||||
|
||||
Represented as a tuple of strings, allowing for nested categorization.
|
||||
For example: ("documents", "user123")
|
||||
"""
|
||||
|
||||
key: str
|
||||
"""Unique identifier for the document.
|
||||
|
||||
Should be distinct within its namespace.
|
||||
"""
|
||||
|
||||
value: Optional[dict[str, Any]]
|
||||
"""Data to be stored, or None to delete the item.
|
||||
|
||||
Schema:
|
||||
- Should be a dictionary where:
|
||||
- Keys are strings representing field names
|
||||
- Values can be of any serializable type
|
||||
- If None, it indicates that the item should be deleted
|
||||
??? example "Examples"
|
||||
- "technical documentation about REST APIs"
|
||||
- "machine learning papers from 2023"
|
||||
"""
|
||||
|
||||
|
||||
NameSpacePath = tuple[Union[str, Literal["*"]], ...]
|
||||
# Type representing a namespace path that can include wildcards
|
||||
NamespacePath = tuple[Union[str, Literal["*"]], ...]
|
||||
"""A tuple representing a namespace path that can include wildcards.
|
||||
|
||||
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
|
||||
NamespaceMatchType = Literal["prefix", "suffix"]
|
||||
"""Specifies how to match namespace paths.
|
||||
|
||||
Values:
|
||||
"prefix": Match from the start of the namespace
|
||||
"suffix": Match from the end of the namespace
|
||||
"""
|
||||
|
||||
|
||||
class MatchCondition(NamedTuple):
|
||||
"""Represents a single match condition."""
|
||||
"""Represents a pattern for matching namespaces in the store.
|
||||
|
||||
This class combines a match type (prefix or suffix) with a namespace path
|
||||
pattern that can include wildcards to flexibly match different namespace
|
||||
hierarchies.
|
||||
|
||||
??? example "Examples"
|
||||
Prefix matching:
|
||||
```python
|
||||
MatchCondition(match_type="prefix", path=("users", "profiles"))
|
||||
```
|
||||
|
||||
Suffix matching with wildcard:
|
||||
```python
|
||||
MatchCondition(match_type="suffix", path=("cache", "*"))
|
||||
```
|
||||
|
||||
Simple suffix matching:
|
||||
```python
|
||||
MatchCondition(match_type="suffix", path=("v1",))
|
||||
```
|
||||
"""
|
||||
|
||||
match_type: NamespaceMatchType
|
||||
path: NameSpacePath
|
||||
"""Type of namespace matching to perform."""
|
||||
|
||||
path: NamespacePath
|
||||
"""Namespace path pattern that can include wildcards."""
|
||||
|
||||
|
||||
class ListNamespacesOp(NamedTuple):
|
||||
"""Operation to list namespaces with optional match conditions."""
|
||||
"""Operation to list and filter namespaces in the store.
|
||||
|
||||
This operation allows exploring the organization of data, finding specific
|
||||
collections, and navigating the namespace hierarchy.
|
||||
|
||||
??? example "Examples"
|
||||
|
||||
List all namespaces under the "documents" path:
|
||||
```python
|
||||
ListNamespacesOp(
|
||||
match_conditions=(MatchCondition(match_type="prefix", path=("documents",)),),
|
||||
max_depth=2
|
||||
)
|
||||
```
|
||||
|
||||
List all namespaces that end with "v1":
|
||||
```python
|
||||
ListNamespacesOp(
|
||||
match_conditions=(MatchCondition(match_type="suffix", path=("v1",)),),
|
||||
limit=50
|
||||
)
|
||||
```
|
||||
|
||||
"""
|
||||
|
||||
match_conditions: Optional[tuple[MatchCondition, ...]] = None
|
||||
"""A tuple of match conditions to apply to namespaces."""
|
||||
"""Optional conditions for filtering namespaces.
|
||||
|
||||
??? example "Examples"
|
||||
All user namespaces:
|
||||
```python
|
||||
(MatchCondition(match_type="prefix", path=("users",)),)
|
||||
```
|
||||
|
||||
All namespaces that start with "docs" and end with "draft":
|
||||
```python
|
||||
(
|
||||
MatchCondition(match_type="prefix", path=("docs",)),
|
||||
MatchCondition(match_type="suffix", path=("draft",))
|
||||
)
|
||||
```
|
||||
"""
|
||||
|
||||
max_depth: Optional[int] = None
|
||||
"""Return namespaces up to this depth in the hierarchy."""
|
||||
"""Maximum depth of namespace hierarchy to return.
|
||||
|
||||
Note:
|
||||
Namespaces deeper than this level will be truncated.
|
||||
"""
|
||||
|
||||
limit: int = 100
|
||||
"""Maximum number of namespaces to return."""
|
||||
|
||||
offset: int = 0
|
||||
"""Number of namespaces to skip before returning results."""
|
||||
"""Number of namespaces to skip for pagination."""
|
||||
|
||||
|
||||
class PutOp(NamedTuple):
|
||||
"""Operation to store, update, or delete an item in the store.
|
||||
|
||||
This class represents a single operation to modify the store's contents,
|
||||
whether adding new items, updating existing ones, or removing them.
|
||||
"""
|
||||
|
||||
namespace: tuple[str, ...]
|
||||
"""Hierarchical path that identifies the location of the item.
|
||||
|
||||
The namespace acts as a folder-like structure to organize items.
|
||||
Each element in the tuple represents one level in the hierarchy.
|
||||
|
||||
??? example "Examples"
|
||||
Root level documents
|
||||
```python
|
||||
("documents",)
|
||||
```
|
||||
|
||||
User-specific documents
|
||||
```python
|
||||
("documents", "user123")
|
||||
```
|
||||
|
||||
Nested cache structure
|
||||
```python
|
||||
("cache", "embeddings", "v1")
|
||||
```
|
||||
"""
|
||||
|
||||
key: str
|
||||
"""Unique identifier for the item within its namespace.
|
||||
|
||||
The key must be unique within the specific namespace to avoid conflicts.
|
||||
Together with the namespace, it forms a complete path to the item.
|
||||
|
||||
Example:
|
||||
If namespace is ("documents", "user123") and key is "report1",
|
||||
the full path would effectively be "documents/user123/report1"
|
||||
"""
|
||||
|
||||
value: Optional[dict[str, Any]]
|
||||
"""The data to store, or None to mark the item for deletion.
|
||||
|
||||
The value must be a dictionary with string keys and JSON-serializable values.
|
||||
Setting this to None signals that the item should be deleted.
|
||||
|
||||
Example:
|
||||
{
|
||||
"field1": "string value",
|
||||
"field2": 123,
|
||||
"nested": {"can": "contain", "any": "serializable data"}
|
||||
}
|
||||
"""
|
||||
|
||||
index: Optional[Union[Literal[False], list[str]]] = None # type: ignore[assignment]
|
||||
"""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
|
||||
|
||||
The item remains accessible through direct get() operations regardless of indexing.
|
||||
When indexed, fields can be searched using natural language queries through
|
||||
vector similarity search (if supported by the store implementation).
|
||||
|
||||
Path Syntax:
|
||||
- Simple field access: "field"
|
||||
- Nested fields: "parent.child.grandchild"
|
||||
- Array indexing:
|
||||
- Specific index: "array[0]"
|
||||
- Last element: "array[-1]"
|
||||
- All elements (each individually): "array[*]"
|
||||
|
||||
??? 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
|
||||
"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
|
||||
"metadata.tags[*]", # All tags in metadata
|
||||
]
|
||||
```
|
||||
"""
|
||||
|
||||
|
||||
Op = Union[GetOp, SearchOp, PutOp, ListNamespacesOp]
|
||||
Result = Union[Item, list[Item], list[tuple[str, ...]], None]
|
||||
Result = Union[Item, list[Item], list[SearchItem], list[tuple[str, ...]], None]
|
||||
|
||||
|
||||
class InvalidNamespaceError(ValueError):
|
||||
"""Provided namespace is invalid."""
|
||||
|
||||
|
||||
def _validate_namespace(namespace: tuple[str, ...]) -> None:
|
||||
if not namespace:
|
||||
raise InvalidNamespaceError("Namespace cannot be empty.")
|
||||
for label in namespace:
|
||||
if not isinstance(label, str):
|
||||
raise InvalidNamespaceError(
|
||||
f"Invalid namespace label '{label}' found in {namespace}. Namespace labels"
|
||||
f" must be strings, but got {type(label).__name__}."
|
||||
)
|
||||
if "." in label:
|
||||
raise InvalidNamespaceError(
|
||||
f"Invalid namespace label '{label}' found in {namespace}. Namespace labels cannot contain periods ('.')."
|
||||
)
|
||||
elif not label:
|
||||
raise InvalidNamespaceError(
|
||||
f"Namespace labels cannot be empty strings. Got {label} in {namespace}"
|
||||
)
|
||||
if namespace[0] == "langgraph":
|
||||
raise InvalidNamespaceError(
|
||||
f'Root label for namespace cannot be "langgraph". Got: {namespace}'
|
||||
)
|
||||
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: 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."""
|
||||
|
||||
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 BaseStore(ABC):
|
||||
@@ -231,14 +549,16 @@ class BaseStore(ABC):
|
||||
namespace_prefix: tuple[str, ...],
|
||||
/,
|
||||
*,
|
||||
query: Optional[str] = None,
|
||||
filter: Optional[dict[str, Any]] = None,
|
||||
limit: int = 10,
|
||||
offset: int = 0,
|
||||
) -> list[Item]:
|
||||
) -> list[SearchItem]:
|
||||
"""Search for items within a namespace prefix.
|
||||
|
||||
Args:
|
||||
namespace_prefix: Hierarchical path prefix to search within.
|
||||
query: Optional query for natural language search.
|
||||
filter: Key-value pairs to filter results.
|
||||
limit: Maximum number of items to return.
|
||||
offset: Number of items to skip before returning results.
|
||||
@@ -246,18 +566,54 @@ class BaseStore(ABC):
|
||||
Returns:
|
||||
List of items matching the search criteria.
|
||||
"""
|
||||
return self.batch([SearchOp(namespace_prefix, filter, limit, offset)])[0]
|
||||
return self.batch([SearchOp(namespace_prefix, filter, limit, offset, query)])[0]
|
||||
|
||||
def put(self, namespace: tuple[str, ...], key: str, value: dict[str, Any]) -> None:
|
||||
"""Store or update an item.
|
||||
def put(
|
||||
self,
|
||||
namespace: tuple[str, ...],
|
||||
key: str,
|
||||
value: dict[str, Any],
|
||||
index: Optional[Union[Literal[False], list[str]]] = None,
|
||||
) -> None:
|
||||
"""Store or update an item in the store.
|
||||
|
||||
Args:
|
||||
namespace: Hierarchical path for the item.
|
||||
key: Unique identifier within the namespace.
|
||||
value: Dictionary containing the item's data.
|
||||
namespace: Hierarchical path for the item, represented as a tuple of strings.
|
||||
Example: ("documents", "user123")
|
||||
key: Unique identifier within the namespace. Together with namespace forms
|
||||
the complete path to the item.
|
||||
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 store's default indexing configuration
|
||||
- False: Disable indexing for this item
|
||||
- list[str]: List of field paths to index, supporting:
|
||||
- Nested fields: "metadata.title"
|
||||
- Array access: "chapters[*].content" (each indexed separately)
|
||||
- Specific indices: "authors[0].name"
|
||||
|
||||
Note:
|
||||
Indexing capabilities depend on your store implementation.
|
||||
Some implementations may support only a subset of indexing features.
|
||||
|
||||
??? example "Examples"
|
||||
Simple storage without special indexing (respects store defaults)
|
||||
```python
|
||||
store.put(("docs",), "report", {"title": "Annual Report"})
|
||||
```
|
||||
|
||||
Index specific fields for search
|
||||
```python
|
||||
store.put(("docs",), "report", {"title": "Annual Report"}, index=["title"])
|
||||
```
|
||||
|
||||
Do not index for semantic search
|
||||
```python
|
||||
store.put(("docs",), "report", {"title": "Annual Report"}, index=False)
|
||||
```
|
||||
"""
|
||||
_validate_namespace(namespace)
|
||||
self.batch([PutOp(namespace, key, value)])
|
||||
self.batch([PutOp(namespace, key, value, index=index)])
|
||||
|
||||
def delete(self, namespace: tuple[str, ...], key: str) -> None:
|
||||
"""Delete an item.
|
||||
@@ -271,8 +627,8 @@ class BaseStore(ABC):
|
||||
def list_namespaces(
|
||||
self,
|
||||
*,
|
||||
prefix: Optional[NameSpacePath] = None,
|
||||
suffix: Optional[NameSpacePath] = None,
|
||||
prefix: Optional[NamespacePath] = None,
|
||||
suffix: Optional[NamespacePath] = None,
|
||||
max_depth: Optional[int] = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
@@ -286,7 +642,7 @@ class BaseStore(ABC):
|
||||
prefix (Optional[Tuple[str, ...]]): Filter namespaces that start with this path.
|
||||
suffix (Optional[Tuple[str, ...]]): Filter namespaces that end with this path.
|
||||
max_depth (Optional[int]): Return namespaces up to this depth in the hierarchy.
|
||||
Namespaces deeper than this level will be truncated to this depth.
|
||||
Namespaces deeper than this level will be truncated.
|
||||
limit (int): Maximum number of namespaces to return (default 100).
|
||||
offset (int): Number of namespaces to skip for pagination (default 0).
|
||||
|
||||
@@ -294,16 +650,18 @@ 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`.
|
||||
|
||||
Examples:
|
||||
|
||||
??? example "Examples":
|
||||
Setting max_depth=3. Given the namespaces:
|
||||
# ("a", "b", "c")
|
||||
# ("a", "b", "d", "e")
|
||||
# ("a", "b", "d", "i")
|
||||
# ("a", "b", "f")
|
||||
# ("a", "c", "f")
|
||||
store.list_namespaces(prefix=("a", "b"), max_depth=3)
|
||||
# [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
|
||||
```python
|
||||
# Example if you have the following namespaces:
|
||||
# ("a", "b", "c")
|
||||
# ("a", "b", "d", "e")
|
||||
# ("a", "b", "d", "i")
|
||||
# ("a", "b", "f")
|
||||
# ("a", "c", "f")
|
||||
store.list_namespaces(prefix=("a", "b"), max_depth=3)
|
||||
# [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
|
||||
```
|
||||
"""
|
||||
match_conditions = []
|
||||
if prefix:
|
||||
@@ -336,14 +694,16 @@ class BaseStore(ABC):
|
||||
namespace_prefix: tuple[str, ...],
|
||||
/,
|
||||
*,
|
||||
query: Optional[str] = None,
|
||||
filter: Optional[dict[str, Any]] = None,
|
||||
limit: int = 10,
|
||||
offset: int = 0,
|
||||
) -> list[Item]:
|
||||
) -> list[SearchItem]:
|
||||
"""Asynchronously search for items within a namespace prefix.
|
||||
|
||||
Args:
|
||||
namespace_prefix: Hierarchical path prefix to search within.
|
||||
query: Optional query for natural language search.
|
||||
filter: Key-value pairs to filter results.
|
||||
limit: Maximum number of items to return.
|
||||
offset: Number of items to skip before returning results.
|
||||
@@ -351,22 +711,61 @@ class BaseStore(ABC):
|
||||
Returns:
|
||||
List of items matching the search criteria.
|
||||
"""
|
||||
return (await self.abatch([SearchOp(namespace_prefix, filter, limit, offset)]))[
|
||||
0
|
||||
]
|
||||
return (
|
||||
await self.abatch(
|
||||
[SearchOp(namespace_prefix, filter, limit, offset, query)]
|
||||
)
|
||||
)[0]
|
||||
|
||||
async def aput(
|
||||
self, namespace: tuple[str, ...], key: str, value: dict[str, Any]
|
||||
self,
|
||||
namespace: tuple[str, ...],
|
||||
key: str,
|
||||
value: dict[str, Any],
|
||||
index: Optional[Union[Literal[False], list[str]]] = None,
|
||||
) -> None:
|
||||
"""Asynchronously store or update an item.
|
||||
"""Asynchronously store or update an item in the store.
|
||||
|
||||
Args:
|
||||
namespace: Hierarchical path for the item.
|
||||
key: Unique identifier within the namespace.
|
||||
value: Dictionary containing the item's data.
|
||||
namespace: Hierarchical path for the item, represented as a tuple of strings.
|
||||
Example: ("documents", "user123")
|
||||
key: Unique identifier within the namespace. Together with namespace forms
|
||||
the complete path to the item.
|
||||
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 store's default indexing configuration
|
||||
- False: Disable indexing for this item
|
||||
- list[str]: List of field paths to index, supporting:
|
||||
- Nested fields: "metadata.title"
|
||||
- Array access: "chapters[*].content" (each indexed separately)
|
||||
- Specific indices: "authors[0].name"
|
||||
|
||||
Note:
|
||||
Indexing capabilities depend on your store implementation.
|
||||
Some implementations may support only a subset of indexing features.
|
||||
|
||||
??? example "Examples"
|
||||
Simple storage without special indexing:
|
||||
```python
|
||||
await store.aput(("docs",), "report", {"title": "Annual Report"})
|
||||
```
|
||||
|
||||
Index specific fields for search:
|
||||
```python
|
||||
await store.aput(
|
||||
("docs",),
|
||||
"report",
|
||||
{
|
||||
"title": "Q4 Report",
|
||||
"chapters": [{"content": "..."}, {"content": "..."}]
|
||||
},
|
||||
index=["title", "chapters[*].content"]
|
||||
)
|
||||
```
|
||||
"""
|
||||
_validate_namespace(namespace)
|
||||
await self.abatch([PutOp(namespace, key, value)])
|
||||
await self.abatch([PutOp(namespace, key, value, index=index)])
|
||||
|
||||
async def adelete(self, namespace: tuple[str, ...], key: str) -> None:
|
||||
"""Asynchronously delete an item.
|
||||
@@ -380,8 +779,8 @@ class BaseStore(ABC):
|
||||
async def alist_namespaces(
|
||||
self,
|
||||
*,
|
||||
prefix: Optional[NameSpacePath] = None,
|
||||
suffix: Optional[NameSpacePath] = None,
|
||||
prefix: Optional[NamespacePath] = None,
|
||||
suffix: Optional[NamespacePath] = None,
|
||||
max_depth: Optional[int] = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
@@ -403,16 +802,19 @@ 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`.
|
||||
|
||||
Examples:
|
||||
??? example "Examples"
|
||||
Setting max_depth=3 with existing namespaces:
|
||||
```python
|
||||
# Given the following namespaces:
|
||||
# ("a", "b", "c")
|
||||
# ("a", "b", "d", "e")
|
||||
# ("a", "b", "d", "i")
|
||||
# ("a", "b", "f")
|
||||
# ("a", "c", "f")
|
||||
|
||||
Setting max_depth=3. Given the namespaces:
|
||||
# ("a", "b", "c")
|
||||
# ("a", "b", "d", "e")
|
||||
# ("a", "b", "d", "i")
|
||||
# ("a", "b", "f")
|
||||
# ("a", "c", "f")
|
||||
await store.alist_namespaces(prefix=("a", "b"), max_depth=3)
|
||||
# [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
|
||||
await store.alist_namespaces(prefix=("a", "b"), max_depth=3)
|
||||
# Returns: [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
|
||||
```
|
||||
"""
|
||||
match_conditions = []
|
||||
if prefix:
|
||||
@@ -427,3 +829,44 @@ class BaseStore(ABC):
|
||||
offset=offset,
|
||||
)
|
||||
return (await self.abatch([op]))[0]
|
||||
|
||||
|
||||
def _validate_namespace(namespace: tuple[str, ...]) -> None:
|
||||
if not namespace:
|
||||
raise InvalidNamespaceError("Namespace cannot be empty.")
|
||||
for label in namespace:
|
||||
if not isinstance(label, str):
|
||||
raise InvalidNamespaceError(
|
||||
f"Invalid namespace label '{label}' found in {namespace}. Namespace labels"
|
||||
f" must be strings, but got {type(label).__name__}."
|
||||
)
|
||||
if "." in label:
|
||||
raise InvalidNamespaceError(
|
||||
f"Invalid namespace label '{label}' found in {namespace}. Namespace labels cannot contain periods ('.')."
|
||||
)
|
||||
elif not label:
|
||||
raise InvalidNamespaceError(
|
||||
f"Namespace labels cannot be empty strings. Got {label} in {namespace}"
|
||||
)
|
||||
if namespace[0] == "langgraph":
|
||||
raise InvalidNamespaceError(
|
||||
f'Root label for namespace cannot be "langgraph". Got: {namespace}'
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BaseStore",
|
||||
"Item",
|
||||
"Op",
|
||||
"PutOp",
|
||||
"GetOp",
|
||||
"SearchOp",
|
||||
"ListNamespacesOp",
|
||||
"MatchCondition",
|
||||
"NamespacePath",
|
||||
"NamespaceMatchType",
|
||||
"Embeddings",
|
||||
"ensure_embeddings",
|
||||
"tokenize_path",
|
||||
"get_text_at_path",
|
||||
]
|
||||
|
||||
@@ -1,13 +1,17 @@
|
||||
import asyncio
|
||||
import weakref
|
||||
from typing import Any, Optional
|
||||
from typing import Any, Literal, Optional, Union
|
||||
|
||||
from langgraph.store.base import (
|
||||
BaseStore,
|
||||
GetOp,
|
||||
Item,
|
||||
ListNamespacesOp,
|
||||
MatchCondition,
|
||||
NamespacePath,
|
||||
Op,
|
||||
PutOp,
|
||||
SearchItem,
|
||||
SearchOp,
|
||||
_validate_namespace,
|
||||
)
|
||||
@@ -40,12 +44,13 @@ class AsyncBatchedBaseStore(BaseStore):
|
||||
namespace_prefix: tuple[str, ...],
|
||||
/,
|
||||
*,
|
||||
query: Optional[str] = None,
|
||||
filter: Optional[dict[str, Any]] = None,
|
||||
limit: int = 10,
|
||||
offset: int = 0,
|
||||
) -> list[Item]:
|
||||
) -> list[SearchItem]:
|
||||
fut = self._loop.create_future()
|
||||
self._aqueue[fut] = SearchOp(namespace_prefix, filter, limit, offset)
|
||||
self._aqueue[fut] = SearchOp(namespace_prefix, filter, limit, offset, query)
|
||||
return await fut
|
||||
|
||||
async def aput(
|
||||
@@ -53,10 +58,11 @@ class AsyncBatchedBaseStore(BaseStore):
|
||||
namespace: tuple[str, ...],
|
||||
key: str,
|
||||
value: dict[str, Any],
|
||||
index: Optional[Union[Literal[False], list[str]]] = None,
|
||||
) -> None:
|
||||
_validate_namespace(namespace)
|
||||
fut = self._loop.create_future()
|
||||
self._aqueue[fut] = PutOp(namespace, key, value)
|
||||
self._aqueue[fut] = PutOp(namespace, key, value, index)
|
||||
return await fut
|
||||
|
||||
async def adelete(
|
||||
@@ -68,6 +74,74 @@ class AsyncBatchedBaseStore(BaseStore):
|
||||
self._aqueue[fut] = PutOp(namespace, key, None)
|
||||
return await fut
|
||||
|
||||
async def alist_namespaces(
|
||||
self,
|
||||
*,
|
||||
prefix: Optional[NamespacePath] = None,
|
||||
suffix: Optional[NamespacePath] = None,
|
||||
max_depth: Optional[int] = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> list[tuple[str, ...]]:
|
||||
fut = self._loop.create_future()
|
||||
match_conditions = []
|
||||
if prefix:
|
||||
match_conditions.append(MatchCondition(match_type="prefix", path=prefix))
|
||||
if suffix:
|
||||
match_conditions.append(MatchCondition(match_type="suffix", path=suffix))
|
||||
|
||||
op = ListNamespacesOp(
|
||||
match_conditions=tuple(match_conditions),
|
||||
max_depth=max_depth,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
self._aqueue[fut] = op
|
||||
return await fut
|
||||
|
||||
|
||||
def _dedupe_ops(values: list[Op]) -> tuple[Optional[list[int]], list[Op]]:
|
||||
"""Dedupe operations while preserving order for results.
|
||||
|
||||
Args:
|
||||
values: List of operations to dedupe
|
||||
|
||||
Returns:
|
||||
Tuple of (listen indices, deduped operations)
|
||||
where listen indices map deduped operation results back to original positions
|
||||
"""
|
||||
if len(values) <= 1:
|
||||
return None, list(values)
|
||||
|
||||
dedupped: list[Op] = []
|
||||
listen: list[int] = []
|
||||
puts: dict[tuple[tuple[str, ...], str], int] = {}
|
||||
|
||||
for op in values:
|
||||
if isinstance(op, (GetOp, SearchOp, ListNamespacesOp)):
|
||||
try:
|
||||
listen.append(dedupped.index(op))
|
||||
except ValueError:
|
||||
listen.append(len(dedupped))
|
||||
dedupped.append(op)
|
||||
elif isinstance(op, PutOp):
|
||||
putkey = (op.namespace, op.key)
|
||||
if putkey in puts:
|
||||
# Overwrite previous put
|
||||
ix = puts[putkey]
|
||||
dedupped[ix] = op
|
||||
listen.append(ix)
|
||||
else:
|
||||
puts[putkey] = len(dedupped)
|
||||
listen.append(len(dedupped))
|
||||
dedupped.append(op)
|
||||
|
||||
else: # Any new ops will be treated regularly
|
||||
listen.append(len(dedupped))
|
||||
dedupped.append(op)
|
||||
|
||||
return listen, dedupped
|
||||
|
||||
|
||||
async def _run(
|
||||
aqueue: dict[asyncio.Future, Op], store: weakref.ReferenceType[BaseStore]
|
||||
@@ -81,7 +155,12 @@ async def _run(
|
||||
taken = aqueue.copy()
|
||||
# action each operation
|
||||
try:
|
||||
results = await s.abatch(taken.values())
|
||||
values = list(taken.values())
|
||||
listen, dedupped = _dedupe_ops(values)
|
||||
results = await s.abatch(dedupped)
|
||||
if listen is not None:
|
||||
results = [results[ix] for ix in listen]
|
||||
|
||||
# set the results of each operation
|
||||
for fut, result in zip(taken, results):
|
||||
fut.set_result(result)
|
||||
|
||||
@@ -0,0 +1,380 @@
|
||||
"""Utilities for working with embedding functions and LangChain's Embeddings interface.
|
||||
|
||||
This module provides tools to wrap arbitrary embedding functions (both sync and async)
|
||||
into LangChain's Embeddings interface. This enables using custom embedding functions
|
||||
with LangChain-compatible tools while maintaining support for both synchronous and
|
||||
asynchronous operations.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any, Awaitable, Callable, Optional, Sequence, Union
|
||||
|
||||
from langchain_core.embeddings import Embeddings
|
||||
|
||||
EmbeddingsFunc = Callable[[Sequence[str]], list[list[float]]]
|
||||
"""Type for synchronous embedding functions.
|
||||
|
||||
The function should take a sequence of strings and return a list of embeddings,
|
||||
where each embedding is a list of floats. The dimensionality of the embeddings
|
||||
should be consistent for all inputs.
|
||||
"""
|
||||
|
||||
AEmbeddingsFunc = Callable[[Sequence[str]], Awaitable[list[list[float]]]]
|
||||
"""Type for asynchronous embedding functions.
|
||||
|
||||
Similar to EmbeddingsFunc, but returns an awaitable that resolves to the embeddings.
|
||||
"""
|
||||
|
||||
|
||||
def ensure_embeddings(
|
||||
embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc, None],
|
||||
) -> Embeddings:
|
||||
"""Ensure that an embedding function conforms to LangChain's Embeddings interface.
|
||||
|
||||
This function wraps arbitrary embedding functions to make them compatible with
|
||||
LangChain's Embeddings interface. It handles both synchronous and asynchronous
|
||||
functions.
|
||||
|
||||
Args:
|
||||
embed: Either an existing Embeddings instance, or a function that converts
|
||||
text to embeddings. If the function is async, it will be used for both
|
||||
sync and async operations.
|
||||
|
||||
Returns:
|
||||
An Embeddings instance that wraps the provided function(s).
|
||||
|
||||
??? example "Examples"
|
||||
Wrap a synchronous embedding function:
|
||||
```python
|
||||
def my_embed_fn(texts):
|
||||
return [[0.1, 0.2] for _ in texts]
|
||||
|
||||
embeddings = ensure_embeddings(my_embed_fn)
|
||||
result = embeddings.embed_query("hello") # Returns [0.1, 0.2]
|
||||
```
|
||||
|
||||
Wrap an asynchronous embedding function:
|
||||
```python
|
||||
async def my_async_fn(texts):
|
||||
return [[0.1, 0.2] for _ in texts]
|
||||
|
||||
embeddings = ensure_embeddings(my_async_fn)
|
||||
result = await embeddings.aembed_query("hello") # Returns [0.1, 0.2]
|
||||
```
|
||||
"""
|
||||
if embed is None:
|
||||
raise ValueError("embed must be provided")
|
||||
if isinstance(embed, Embeddings):
|
||||
return embed
|
||||
return EmbeddingsLambda(embed)
|
||||
|
||||
|
||||
class EmbeddingsLambda(Embeddings):
|
||||
"""Wrapper to convert embedding functions into LangChain's Embeddings interface.
|
||||
|
||||
This class allows arbitrary embedding functions to be used with LangChain-compatible
|
||||
tools. It supports both synchronous and asynchronous operations, and can handle:
|
||||
1. A synchronous function for sync operations (async operations will use sync function)
|
||||
2. An async function for both sync/async operations (sync operations will raise an error)
|
||||
|
||||
The embedding functions should convert text into fixed-dimensional vectors that
|
||||
capture the semantic meaning of the text.
|
||||
|
||||
Args:
|
||||
func: Function that converts text to embeddings. Can be sync or async.
|
||||
If async, it will be used for async operations, but sync operations
|
||||
will raise an error. If sync, it will be used for both sync and async operations.
|
||||
|
||||
??? example "Examples"
|
||||
With a sync function:
|
||||
```python
|
||||
def my_embed_fn(texts):
|
||||
# Return 2D embeddings for each text
|
||||
return [[0.1, 0.2] for _ in texts]
|
||||
|
||||
embeddings = EmbeddingsLambda(my_embed_fn)
|
||||
result = embeddings.embed_query("hello") # Returns [0.1, 0.2]
|
||||
await embeddings.aembed_query("hello") # Also returns [0.1, 0.2]
|
||||
```
|
||||
|
||||
With an async function:
|
||||
```python
|
||||
async def my_async_fn(texts):
|
||||
return [[0.1, 0.2] for _ in texts]
|
||||
|
||||
embeddings = EmbeddingsLambda(my_async_fn)
|
||||
await embeddings.aembed_query("hello") # Returns [0.1, 0.2]
|
||||
# Note: embed_query() would raise an error
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
func: Union[EmbeddingsFunc, AEmbeddingsFunc],
|
||||
) -> None:
|
||||
if func is None:
|
||||
raise ValueError("func must be provided")
|
||||
if _is_async_callable(func):
|
||||
self.afunc = func
|
||||
else:
|
||||
self.func = func
|
||||
|
||||
def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Embed a list of texts into vectors.
|
||||
|
||||
Args:
|
||||
texts: list of texts to convert to embeddings.
|
||||
|
||||
Returns:
|
||||
list of embeddings, one per input text. Each embedding is a list of floats.
|
||||
|
||||
Raises:
|
||||
ValueError: If the instance was initialized with only an async function.
|
||||
"""
|
||||
func = getattr(self, "func", None)
|
||||
if func is None:
|
||||
raise ValueError(
|
||||
"EmbeddingsLambda was initialized with an async function but no sync function. "
|
||||
"Use aembed_documents for async operation or provide a sync function."
|
||||
)
|
||||
return func(texts)
|
||||
|
||||
def embed_query(self, text: str) -> list[float]:
|
||||
"""Embed a single piece of text.
|
||||
|
||||
Args:
|
||||
text: Text to convert to an embedding.
|
||||
|
||||
Returns:
|
||||
Embedding vector as a list of floats.
|
||||
|
||||
Note:
|
||||
This is equivalent to calling embed_documents with a single text
|
||||
and taking the first result.
|
||||
"""
|
||||
return self.embed_documents([text])[0]
|
||||
|
||||
async def aembed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Asynchronously embed a list of texts into vectors.
|
||||
|
||||
Args:
|
||||
texts: list of texts to convert to embeddings.
|
||||
|
||||
Returns:
|
||||
list of embeddings, one per input text. Each embedding is a list of floats.
|
||||
|
||||
Note:
|
||||
If no async function was provided, this falls back to the sync implementation.
|
||||
"""
|
||||
afunc = getattr(self, "afunc", None)
|
||||
if afunc is None:
|
||||
return await super().aembed_documents(texts)
|
||||
return await afunc(texts)
|
||||
|
||||
async def aembed_query(self, text: str) -> list[float]:
|
||||
"""Asynchronously embed a single piece of text.
|
||||
|
||||
Args:
|
||||
text: Text to convert to an embedding.
|
||||
|
||||
Returns:
|
||||
Embedding vector as a list of floats.
|
||||
|
||||
Note:
|
||||
This is equivalent to calling aembed_documents with a single text
|
||||
and taking the first result.
|
||||
"""
|
||||
afunc = getattr(self, "afunc", None)
|
||||
if afunc is None:
|
||||
return await super().aembed_query(text)
|
||||
return (await afunc([text]))[0]
|
||||
|
||||
|
||||
def get_text_at_path(obj: Any, path: Union[str, list[str]]) -> list[str]:
|
||||
"""Extract text from an object using a path expression or pre-tokenized path.
|
||||
|
||||
Args:
|
||||
obj: The object to extract text from
|
||||
path: Either a path string or pre-tokenized path list.
|
||||
|
||||
!!! info "Path types handled"
|
||||
- Simple paths: "field1.field2"
|
||||
- Array indexing: "[0]", "[*]", "[-1]"
|
||||
- Wildcards: "*"
|
||||
- Multi-field selection: "{field1,field2}"
|
||||
- Nested paths in multi-field: "{field1,nested.field2}"
|
||||
"""
|
||||
if not path or path == "$":
|
||||
return [json.dumps(obj, sort_keys=True)]
|
||||
|
||||
tokens = tokenize_path(path) if isinstance(path, str) else path
|
||||
|
||||
def _extract_from_obj(obj: Any, tokens: list[str], pos: int) -> list[str]:
|
||||
if pos >= len(tokens):
|
||||
if isinstance(obj, (str, int, float, bool)):
|
||||
return [str(obj)]
|
||||
elif obj is None:
|
||||
return []
|
||||
elif isinstance(obj, (list, dict)):
|
||||
return [json.dumps(obj, sort_keys=True)]
|
||||
return []
|
||||
|
||||
token = tokens[pos]
|
||||
results = []
|
||||
|
||||
if token.startswith("[") and token.endswith("]"):
|
||||
if not isinstance(obj, list):
|
||||
return []
|
||||
|
||||
index = token[1:-1]
|
||||
if index == "*":
|
||||
for item in obj:
|
||||
results.extend(_extract_from_obj(item, tokens, pos + 1))
|
||||
else:
|
||||
try:
|
||||
idx = int(index)
|
||||
if idx < 0:
|
||||
idx = len(obj) + idx
|
||||
if 0 <= idx < len(obj):
|
||||
results.extend(_extract_from_obj(obj[idx], tokens, pos + 1))
|
||||
except (ValueError, IndexError):
|
||||
return []
|
||||
|
||||
elif token.startswith("{") and token.endswith("}"):
|
||||
if not isinstance(obj, dict):
|
||||
return []
|
||||
|
||||
fields = [f.strip() for f in token[1:-1].split(",")]
|
||||
for field in fields:
|
||||
nested_tokens = tokenize_path(field)
|
||||
if nested_tokens:
|
||||
current_obj: Optional[dict] = obj
|
||||
for nested_token in nested_tokens:
|
||||
if (
|
||||
isinstance(current_obj, dict)
|
||||
and nested_token in current_obj
|
||||
):
|
||||
current_obj = current_obj[nested_token]
|
||||
else:
|
||||
current_obj = None
|
||||
break
|
||||
if current_obj is not None:
|
||||
if isinstance(current_obj, (str, int, float, bool)):
|
||||
results.append(str(current_obj))
|
||||
elif isinstance(current_obj, (list, dict)):
|
||||
results.append(json.dumps(current_obj, sort_keys=True))
|
||||
|
||||
# Handle wildcard
|
||||
elif token == "*":
|
||||
if isinstance(obj, dict):
|
||||
for value in obj.values():
|
||||
results.extend(_extract_from_obj(value, tokens, pos + 1))
|
||||
elif isinstance(obj, list):
|
||||
for item in obj:
|
||||
results.extend(_extract_from_obj(item, tokens, pos + 1))
|
||||
|
||||
# Handle regular field
|
||||
else:
|
||||
if isinstance(obj, dict) and token in obj:
|
||||
results.extend(_extract_from_obj(obj[token], tokens, pos + 1))
|
||||
|
||||
return results
|
||||
|
||||
return _extract_from_obj(obj, tokens, 0)
|
||||
|
||||
|
||||
# Private utility functions
|
||||
|
||||
|
||||
def tokenize_path(path: str) -> list[str]:
|
||||
"""Tokenize a path into components.
|
||||
|
||||
!!! info "Types handled"
|
||||
- Simple paths: "field1.field2"
|
||||
- Array indexing: "[0]", "[*]", "[-1]"
|
||||
- Wildcards: "*"
|
||||
- Multi-field selection: "{field1,field2}"
|
||||
"""
|
||||
if not path:
|
||||
return []
|
||||
|
||||
tokens = []
|
||||
current: list[str] = []
|
||||
i = 0
|
||||
while i < len(path):
|
||||
char = path[i]
|
||||
|
||||
if char == "[": # Handle array index
|
||||
if current:
|
||||
tokens.append("".join(current))
|
||||
current = []
|
||||
bracket_count = 1
|
||||
index_chars = ["["]
|
||||
i += 1
|
||||
while i < len(path) and bracket_count > 0:
|
||||
if path[i] == "[":
|
||||
bracket_count += 1
|
||||
elif path[i] == "]":
|
||||
bracket_count -= 1
|
||||
index_chars.append(path[i])
|
||||
i += 1
|
||||
tokens.append("".join(index_chars))
|
||||
continue
|
||||
|
||||
elif char == "{": # Handle multi-field selection
|
||||
if current:
|
||||
tokens.append("".join(current))
|
||||
current = []
|
||||
brace_count = 1
|
||||
field_chars = ["{"]
|
||||
i += 1
|
||||
while i < len(path) and brace_count > 0:
|
||||
if path[i] == "{":
|
||||
brace_count += 1
|
||||
elif path[i] == "}":
|
||||
brace_count -= 1
|
||||
field_chars.append(path[i])
|
||||
i += 1
|
||||
tokens.append("".join(field_chars))
|
||||
continue
|
||||
|
||||
elif char == ".": # Handle regular field
|
||||
if current:
|
||||
tokens.append("".join(current))
|
||||
current = []
|
||||
else:
|
||||
current.append(char)
|
||||
i += 1
|
||||
|
||||
if current:
|
||||
tokens.append("".join(current))
|
||||
|
||||
return tokens
|
||||
|
||||
|
||||
def _is_async_callable(
|
||||
func: Any,
|
||||
) -> bool:
|
||||
"""Check if a function is async.
|
||||
|
||||
This includes both async def functions and classes with async __call__ methods.
|
||||
|
||||
Args:
|
||||
func: Function or callable object to check.
|
||||
|
||||
Returns:
|
||||
True if the function is async, False otherwise.
|
||||
"""
|
||||
return (
|
||||
asyncio.iscoroutinefunction(func)
|
||||
or hasattr(func, "__call__") # noqa: B004
|
||||
and asyncio.iscoroutinefunction(func.__call__)
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ensure_embeddings",
|
||||
"EmbeddingsFunc",
|
||||
"AEmbeddingsFunc",
|
||||
]
|
||||
@@ -1,79 +1,379 @@
|
||||
"""In-memory key-value store.
|
||||
|
||||
A lightweight store implementation using Python dictionaries. Supports basic
|
||||
key-value operations and vector search when configured with embeddings.
|
||||
|
||||
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_openai import OpenAIEmbeddings
|
||||
store = InMemoryStore(index={
|
||||
"dims": 1536,
|
||||
"embed": OpenAIEmbeddings(model="text-embedding-3-small"),
|
||||
})
|
||||
|
||||
# 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")
|
||||
|
||||
|
||||
Note:
|
||||
For production use cases requiring persistence, use a database-backed store instead.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import concurrent.futures as cf
|
||||
import functools
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from datetime import datetime, timezone
|
||||
from typing import Iterable
|
||||
from importlib import util
|
||||
from typing import Any, Iterable, Optional
|
||||
|
||||
from langchain_core.embeddings import Embeddings
|
||||
|
||||
from langgraph.store.base import (
|
||||
BaseStore,
|
||||
GetOp,
|
||||
IndexConfig,
|
||||
Item,
|
||||
ListNamespacesOp,
|
||||
MatchCondition,
|
||||
Op,
|
||||
PutOp,
|
||||
Result,
|
||||
SearchItem,
|
||||
SearchOp,
|
||||
ensure_embeddings,
|
||||
get_text_at_path,
|
||||
tokenize_path,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class InMemoryStore(BaseStore):
|
||||
"""A KV store backed by an in-memory python dictionary.
|
||||
"""In-memory dictionary-backed store with optional vector search.
|
||||
|
||||
Useful for testing/experimentation and lightweight PoC's.
|
||||
For actual persistence, use a Store backed by a proper database.
|
||||
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_openai import OpenAIEmbeddings
|
||||
store = InMemoryStore(index={
|
||||
"dims": 1536,
|
||||
"embed": OpenAIEmbeddings(model="text-embedding-3-small"),
|
||||
})
|
||||
|
||||
# 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")
|
||||
|
||||
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
|
||||
```
|
||||
"""
|
||||
|
||||
__slots__ = ("_data",)
|
||||
__slots__ = (
|
||||
"_data",
|
||||
"_vectors",
|
||||
"index_config",
|
||||
"embeddings",
|
||||
)
|
||||
|
||||
def __init__(self) -> None:
|
||||
def __init__(self, *, index: Optional[IndexConfig] = None) -> None:
|
||||
# Both _data and _vectors are wrapped in the In-memory API
|
||||
# Do not change their names
|
||||
self._data: dict[tuple[str, ...], dict[str, Item]] = defaultdict(dict)
|
||||
# [ns][key][path]
|
||||
self._vectors: dict[tuple[str, ...], dict[str, dict[str, list[float]]]] = (
|
||||
defaultdict(lambda: defaultdict(dict))
|
||||
)
|
||||
self.index_config = index
|
||||
if self.index_config:
|
||||
self.index_config = self.index_config.copy()
|
||||
self.embeddings: Optional[Embeddings] = ensure_embeddings(
|
||||
self.index_config.get("embed"),
|
||||
)
|
||||
self.index_config["__tokenized_fields"] = [
|
||||
(p, tokenize_path(p)) if p != "$" else (p, p)
|
||||
for p in (self.index_config.get("fields") or ["$"])
|
||||
]
|
||||
|
||||
else:
|
||||
self.index_config = None
|
||||
self.embeddings = None
|
||||
|
||||
def batch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
# The batch/abatch methods are treated as internal.
|
||||
# Users should access via put/search/get/list_namespaces/etc.
|
||||
results, put_ops, search_ops = self._prepare_ops(ops)
|
||||
if search_ops:
|
||||
queryinmem_store = self._embed_search_queries(search_ops)
|
||||
self._batch_search(search_ops, queryinmem_store, results)
|
||||
|
||||
to_embed = self._extract_texts(put_ops)
|
||||
if to_embed and self.index_config and self.embeddings:
|
||||
embeddings = self.embeddings.embed_documents(list(to_embed))
|
||||
self._insertinmem_store(to_embed, embeddings)
|
||||
self._apply_put_ops(put_ops)
|
||||
return results
|
||||
|
||||
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
# The batch/abatch methods are treated as internal.
|
||||
# Users should access via put/search/get/list_namespaces/etc.
|
||||
results, put_ops, search_ops = self._prepare_ops(ops)
|
||||
if search_ops:
|
||||
queryinmem_store = await self._aembed_search_queries(search_ops)
|
||||
self._batch_search(search_ops, queryinmem_store, results)
|
||||
|
||||
to_embed = self._extract_texts(put_ops)
|
||||
if to_embed and self.index_config and self.embeddings:
|
||||
embeddings = await self.embeddings.aembed_documents(list(to_embed))
|
||||
self._insertinmem_store(to_embed, embeddings)
|
||||
self._apply_put_ops(put_ops)
|
||||
return results
|
||||
|
||||
# Helpers
|
||||
|
||||
def _filter_items(self, op: SearchOp) -> list[tuple[Item, list[list[float]]]]:
|
||||
"""Filter items by namespace and filter function, return items with their embeddings."""
|
||||
namespace_prefix = op.namespace_prefix
|
||||
|
||||
def filter_func(item: Item) -> bool:
|
||||
if not op.filter:
|
||||
return True
|
||||
|
||||
return all(
|
||||
_compare_values(item.value.get(key), filter_value)
|
||||
for key, filter_value in op.filter.items()
|
||||
)
|
||||
|
||||
filtered = []
|
||||
for namespace in self._data:
|
||||
if not (
|
||||
namespace[: len(namespace_prefix)] == namespace_prefix
|
||||
if len(namespace) >= len(namespace_prefix)
|
||||
else False
|
||||
):
|
||||
continue
|
||||
|
||||
for key, item in self._data[namespace].items():
|
||||
if filter_func(item):
|
||||
if op.query and (embeddings := self._vectors[namespace].get(key)):
|
||||
filtered.append((item, list(embeddings.values())))
|
||||
else:
|
||||
filtered.append((item, []))
|
||||
return filtered
|
||||
|
||||
def _embed_search_queries(
|
||||
self,
|
||||
search_ops: dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]],
|
||||
) -> dict[str, list[float]]:
|
||||
queryinmem_store = {}
|
||||
if self.index_config and self.embeddings and search_ops:
|
||||
queries = {op.query for (op, _) in search_ops.values() if op.query}
|
||||
|
||||
if queries:
|
||||
with cf.ThreadPoolExecutor() as executor:
|
||||
futures = {
|
||||
q: executor.submit(self.embeddings.embed_query, q)
|
||||
for q in list(queries)
|
||||
}
|
||||
for query, future in futures.items():
|
||||
queryinmem_store[query] = future.result()
|
||||
|
||||
return queryinmem_store
|
||||
|
||||
async def _aembed_search_queries(
|
||||
self,
|
||||
search_ops: dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]],
|
||||
) -> dict[str, list[float]]:
|
||||
queryinmem_store = {}
|
||||
if self.index_config and self.embeddings and search_ops:
|
||||
queries = {op.query for (op, _) in search_ops.values() if op.query}
|
||||
|
||||
if queries:
|
||||
coros = [self.embeddings.aembed_query(q) for q in list(queries)]
|
||||
results = await asyncio.gather(*coros)
|
||||
queryinmem_store = dict(zip(queries, results))
|
||||
|
||||
return queryinmem_store
|
||||
|
||||
def _batch_search(
|
||||
self,
|
||||
ops: dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]],
|
||||
queryinmem_store: dict[str, list[float]],
|
||||
results: list[Result],
|
||||
) -> None:
|
||||
"""Perform batch similarity search for multiple queries."""
|
||||
for i, (op, candidates) in ops.items():
|
||||
if not candidates:
|
||||
results[i] = []
|
||||
continue
|
||||
if op.query and queryinmem_store:
|
||||
query_embedding = queryinmem_store[op.query]
|
||||
flat_items, flat_vectors = [], []
|
||||
scoreless = []
|
||||
for item, vectors in candidates:
|
||||
for vector in vectors:
|
||||
flat_items.append(item)
|
||||
flat_vectors.append(vector)
|
||||
if not vectors:
|
||||
scoreless.append(item)
|
||||
|
||||
scores = _cosine_similarity(query_embedding, flat_vectors)
|
||||
sorted_results = sorted(
|
||||
zip(scores, flat_items), key=lambda x: x[0], reverse=True
|
||||
)
|
||||
# max pooling
|
||||
seen: set[tuple[tuple[str, ...], str]] = set()
|
||||
kept: list[tuple[Optional[float], Item]] = []
|
||||
for score, item in sorted_results:
|
||||
key = (item.namespace, item.key)
|
||||
if key in seen:
|
||||
continue
|
||||
ix = len(seen)
|
||||
seen.add(key)
|
||||
if ix >= op.offset + op.limit:
|
||||
break
|
||||
if ix < op.offset:
|
||||
continue
|
||||
|
||||
kept.append((score, item))
|
||||
if scoreless and len(kept) < op.limit:
|
||||
# Corner case: if we request more items than what we have embedded,
|
||||
# fill the rest with non-scored items
|
||||
kept.extend(
|
||||
(None, item) for item in scoreless[: op.limit - len(kept)]
|
||||
)
|
||||
|
||||
results[i] = [
|
||||
SearchItem(
|
||||
namespace=item.namespace,
|
||||
key=item.key,
|
||||
value=item.value,
|
||||
created_at=item.created_at,
|
||||
updated_at=item.updated_at,
|
||||
score=float(score) if score is not None else None,
|
||||
)
|
||||
for score, item in kept
|
||||
]
|
||||
else:
|
||||
results[i] = [
|
||||
SearchItem(
|
||||
namespace=item.namespace,
|
||||
key=item.key,
|
||||
value=item.value,
|
||||
created_at=item.created_at,
|
||||
updated_at=item.updated_at,
|
||||
)
|
||||
for (item, _) in candidates[op.offset : op.offset + op.limit]
|
||||
]
|
||||
|
||||
def _prepare_ops(
|
||||
self, ops: Iterable[Op]
|
||||
) -> tuple[
|
||||
list[Result],
|
||||
dict[tuple[tuple[str, ...], str], PutOp],
|
||||
dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]],
|
||||
]:
|
||||
results: list[Result] = []
|
||||
for op in ops:
|
||||
put_ops: dict[tuple[tuple[str, ...], str], PutOp] = {}
|
||||
search_ops: dict[
|
||||
int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]
|
||||
] = {}
|
||||
for i, op in enumerate(ops):
|
||||
if isinstance(op, GetOp):
|
||||
item = self._data[op.namespace].get(op.key)
|
||||
results.append(item)
|
||||
elif isinstance(op, SearchOp):
|
||||
candidates = [
|
||||
item
|
||||
for namespace, items in self._data.items()
|
||||
if (
|
||||
namespace[: len(op.namespace_prefix)] == op.namespace_prefix
|
||||
if len(namespace) >= len(op.namespace_prefix)
|
||||
else False
|
||||
)
|
||||
for item in items.values()
|
||||
]
|
||||
if op.filter:
|
||||
candidates = [
|
||||
item
|
||||
for item in candidates
|
||||
if item.value.items() >= op.filter.items()
|
||||
]
|
||||
results.append(candidates[op.offset : op.offset + op.limit])
|
||||
elif isinstance(op, PutOp):
|
||||
if op.value is None:
|
||||
self._data[op.namespace].pop(op.key, None)
|
||||
elif op.key in self._data[op.namespace]:
|
||||
self._data[op.namespace][op.key].value = op.value
|
||||
self._data[op.namespace][op.key].updated_at = datetime.now(
|
||||
timezone.utc
|
||||
)
|
||||
else:
|
||||
self._data[op.namespace][op.key] = Item(
|
||||
value=op.value,
|
||||
key=op.key,
|
||||
namespace=op.namespace,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
search_ops[i] = (op, self._filter_items(op))
|
||||
results.append(None)
|
||||
elif isinstance(op, ListNamespacesOp):
|
||||
results.append(self._handle_list_namespaces(op))
|
||||
return results
|
||||
elif isinstance(op, PutOp):
|
||||
put_ops[(op.namespace, op.key)] = op
|
||||
results.append(None)
|
||||
else:
|
||||
raise ValueError(f"Unknown operation type: {type(op)}")
|
||||
|
||||
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
return self.batch(ops)
|
||||
return results, put_ops, search_ops
|
||||
|
||||
def _apply_put_ops(self, put_ops: dict[tuple[tuple[str, ...], str], PutOp]) -> None:
|
||||
for (namespace, key), op in put_ops.items():
|
||||
if op.value is None:
|
||||
self._data[namespace].pop(key, None)
|
||||
self._vectors[namespace].pop(key, None)
|
||||
else:
|
||||
self._data[namespace][key] = Item(
|
||||
value=op.value,
|
||||
key=key,
|
||||
namespace=namespace,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
def _extract_texts(
|
||||
self, put_ops: dict[tuple[tuple[str, ...], str], PutOp]
|
||||
) -> dict[str, list[tuple[tuple[str, ...], str, str]]]:
|
||||
if put_ops and self.index_config and self.embeddings:
|
||||
to_embed = defaultdict(list)
|
||||
|
||||
for op in put_ops.values():
|
||||
if op.value is not None and op.index is not False:
|
||||
if op.index is None:
|
||||
paths = self.index_config["__tokenized_fields"]
|
||||
else:
|
||||
paths = [(ix, tokenize_path(ix)) for ix in op.index]
|
||||
for path, field in paths:
|
||||
texts = get_text_at_path(op.value, field)
|
||||
if texts:
|
||||
if len(texts) > 1:
|
||||
for i, text in enumerate(texts):
|
||||
to_embed[text].append(
|
||||
(op.namespace, op.key, f"{path}.{i}")
|
||||
)
|
||||
|
||||
else:
|
||||
to_embed[texts[0]].append((op.namespace, op.key, path))
|
||||
|
||||
return to_embed
|
||||
|
||||
return {}
|
||||
|
||||
def _insertinmem_store(
|
||||
self,
|
||||
to_embed: dict[str, list[tuple[tuple[str, ...], str, str]]],
|
||||
embeddings: list[list[float]],
|
||||
) -> None:
|
||||
indices = [index for indices in to_embed.values() for index in indices]
|
||||
if len(indices) != len(embeddings):
|
||||
raise ValueError(
|
||||
f"Number of embeddings ({len(embeddings)}) does not"
|
||||
f" match number of indices ({len(indices)})"
|
||||
)
|
||||
for embedding, (ns, key, path) in zip(embeddings, indices):
|
||||
self._vectors[ns][key][path] = embedding
|
||||
|
||||
def _handle_list_namespaces(self, op: ListNamespacesOp) -> list[tuple[str, ...]]:
|
||||
all_namespaces = list(
|
||||
@@ -94,7 +394,52 @@ class InMemoryStore(BaseStore):
|
||||
return namespaces[op.offset : op.offset + op.limit]
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _check_numpy() -> bool:
|
||||
if bool(util.find_spec("numpy")):
|
||||
return True
|
||||
logger.warning(
|
||||
"NumPy not found in the current Python environment. "
|
||||
"The InMemoryStore will use a pure Python implementation for vector operations, "
|
||||
"which may significantly impact performance, especially for large datasets or frequent searches. "
|
||||
"For optimal speed and efficiency, consider installing NumPy: "
|
||||
"pip install numpy"
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
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 _check_numpy():
|
||||
import numpy as np # type: ignore
|
||||
|
||||
X_arr = np.array(X) if not isinstance(X, np.ndarray) else X
|
||||
Y_arr = np.array(Y) if not isinstance(Y, np.ndarray) else Y
|
||||
X_norm = np.linalg.norm(X_arr)
|
||||
Y_norm = np.linalg.norm(Y_arr, axis=1)
|
||||
|
||||
# Avoid division by zero
|
||||
mask = Y_norm != 0
|
||||
similarities = np.zeros_like(Y_norm)
|
||||
similarities[mask] = np.dot(Y_arr[mask], X_arr) / (Y_norm[mask] * X_norm)
|
||||
return similarities.tolist()
|
||||
|
||||
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 _does_match(match_condition: MatchCondition, key: tuple[str, ...]) -> bool:
|
||||
"""Whether a namespace key matches a match condition."""
|
||||
match_type = match_condition.match_type
|
||||
path = match_condition.path
|
||||
|
||||
@@ -117,3 +462,44 @@ def _does_match(match_condition: MatchCondition, key: tuple[str, ...]) -> bool:
|
||||
return True
|
||||
else:
|
||||
raise ValueError(f"Unsupported match type: {match_type}")
|
||||
|
||||
|
||||
def _compare_values(item_value: Any, filter_value: Any) -> bool:
|
||||
"""Compare values in a JSONB-like way, handling nested objects."""
|
||||
if isinstance(filter_value, dict):
|
||||
if any(k.startswith("$") for k in filter_value):
|
||||
return all(
|
||||
_apply_operator(item_value, op_key, op_value)
|
||||
for op_key, op_value in filter_value.items()
|
||||
)
|
||||
if not isinstance(item_value, dict):
|
||||
return False
|
||||
return all(
|
||||
_compare_values(item_value.get(k), v) for k, v in filter_value.items()
|
||||
)
|
||||
elif isinstance(filter_value, (list, tuple)):
|
||||
return (
|
||||
isinstance(item_value, (list, tuple))
|
||||
and len(item_value) == len(filter_value)
|
||||
and all(_compare_values(iv, fv) for iv, fv in zip(item_value, filter_value))
|
||||
)
|
||||
else:
|
||||
return item_value == filter_value
|
||||
|
||||
|
||||
def _apply_operator(value: Any, operator: str, op_value: Any) -> bool:
|
||||
"""Apply a comparison operator, matching PostgreSQL's JSONB behavior."""
|
||||
if operator == "$eq":
|
||||
return value == op_value
|
||||
elif operator == "$gt":
|
||||
return float(value) > float(op_value)
|
||||
elif operator == "$gte":
|
||||
return float(value) >= float(op_value)
|
||||
elif operator == "$lt":
|
||||
return float(value) < float(op_value)
|
||||
elif operator == "$lte":
|
||||
return float(value) <= float(op_value)
|
||||
elif operator == "$ne":
|
||||
return value != op_value
|
||||
else:
|
||||
raise ValueError(f"Unsupported operator: {operator}")
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "2.0.5"
|
||||
version = "2.0.7"
|
||||
description = "Library with base interfaces for LangGraph checkpoint savers."
|
||||
authors = []
|
||||
license = "MIT"
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
"""Embedding utilities for testing."""
|
||||
|
||||
import math
|
||||
import random
|
||||
from collections import Counter, defaultdict
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.embeddings import Embeddings
|
||||
|
||||
|
||||
class CharacterEmbeddings(Embeddings):
|
||||
"""Simple character-frequency based embeddings using random projections."""
|
||||
|
||||
def __init__(self, dims: int = 50, seed: int = 42):
|
||||
"""Initialize with embedding dimensions and random seed."""
|
||||
self._rng = random.Random(seed)
|
||||
self.dims = dims
|
||||
# Create projection vector for each character lazily
|
||||
self._char_projections: defaultdict[str, list[float]] = defaultdict(
|
||||
lambda: [
|
||||
self._rng.gauss(0, 1 / math.sqrt(self.dims)) for _ in range(self.dims)
|
||||
]
|
||||
)
|
||||
|
||||
def _embed_one(self, text: str) -> list[float]:
|
||||
"""Embed a single text."""
|
||||
counts = Counter(text)
|
||||
total = sum(counts.values())
|
||||
|
||||
if total == 0:
|
||||
return [0.0] * self.dims
|
||||
|
||||
embedding = [0.0] * self.dims
|
||||
for char, count in counts.items():
|
||||
weight = count / total
|
||||
char_proj = self._char_projections[char]
|
||||
for i, proj in enumerate(char_proj):
|
||||
embedding[i] += weight * proj
|
||||
|
||||
norm = math.sqrt(sum(x * x for x in embedding))
|
||||
if norm > 0:
|
||||
embedding = [x / norm for x in embedding]
|
||||
|
||||
return embedding
|
||||
|
||||
def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Embed a list of documents."""
|
||||
return [self._embed_one(text) for text in texts]
|
||||
|
||||
def embed_query(self, text: str) -> list[float]:
|
||||
"""Embed a query string."""
|
||||
return self._embed_one(text)
|
||||
|
||||
def __eq__(self, other: Any) -> bool:
|
||||
return isinstance(other, CharacterEmbeddings) and self.dims == other.dims
|
||||
@@ -1,13 +1,104 @@
|
||||
# mypy: disable-error-code="operator"
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import datetime
|
||||
from typing import Iterable
|
||||
from typing import Any, Iterable
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from langgraph.store.base import GetOp, InvalidNamespaceError, Item, Op, PutOp, Result
|
||||
from langgraph.store.base import (
|
||||
GetOp,
|
||||
InvalidNamespaceError,
|
||||
Item,
|
||||
Op,
|
||||
PutOp,
|
||||
Result,
|
||||
get_text_at_path,
|
||||
)
|
||||
from langgraph.store.base.batch import AsyncBatchedBaseStore
|
||||
from langgraph.store.memory import InMemoryStore
|
||||
from tests.embed_test_utils import CharacterEmbeddings
|
||||
|
||||
|
||||
class MockAsyncBatchedStore(AsyncBatchedBaseStore):
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
super().__init__()
|
||||
self._store = InMemoryStore(**kwargs)
|
||||
|
||||
def batch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
return self._store.batch(ops)
|
||||
|
||||
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
return self._store.batch(ops)
|
||||
|
||||
|
||||
def test_get_text_at_path() -> None:
|
||||
nested_data = {
|
||||
"name": "test",
|
||||
"info": {
|
||||
"age": 25,
|
||||
"tags": ["a", "b", "c"],
|
||||
"metadata": {"created": "2024-01-01", "updated": "2024-01-02"},
|
||||
},
|
||||
"items": [
|
||||
{"id": 1, "value": "first", "tags": ["x", "y"]},
|
||||
{"id": 2, "value": "second", "tags": ["y", "z"]},
|
||||
{"id": 3, "value": "third", "tags": ["z", "w"]},
|
||||
],
|
||||
"empty": None,
|
||||
"zeros": [0, 0.0, "0"],
|
||||
"empty_list": [],
|
||||
"empty_dict": {},
|
||||
}
|
||||
|
||||
assert get_text_at_path(nested_data, "$") == [
|
||||
json.dumps(nested_data, sort_keys=True)
|
||||
]
|
||||
|
||||
assert get_text_at_path(nested_data, "name") == ["test"]
|
||||
assert get_text_at_path(nested_data, "info.age") == ["25"]
|
||||
|
||||
assert get_text_at_path(nested_data, "info.metadata.created") == ["2024-01-01"]
|
||||
|
||||
assert get_text_at_path(nested_data, "items[0].value") == ["first"]
|
||||
assert get_text_at_path(nested_data, "items[-1].value") == ["third"]
|
||||
assert get_text_at_path(nested_data, "items[1].tags[0]") == ["y"]
|
||||
|
||||
values = get_text_at_path(nested_data, "items[*].value")
|
||||
assert set(values) == {"first", "second", "third"}
|
||||
|
||||
metadata_dates = get_text_at_path(nested_data, "info.metadata.*")
|
||||
assert set(metadata_dates) == {"2024-01-01", "2024-01-02"}
|
||||
name_and_age = get_text_at_path(nested_data, "{name,info.age}")
|
||||
assert set(name_and_age) == {"test", "25"}
|
||||
|
||||
item_fields = get_text_at_path(nested_data, "items[*].{id,value}")
|
||||
assert set(item_fields) == {"1", "2", "3", "first", "second", "third"}
|
||||
|
||||
all_tags = get_text_at_path(nested_data, "items[*].tags[*]")
|
||||
assert set(all_tags) == {"x", "y", "z", "w"}
|
||||
|
||||
assert get_text_at_path(None, "any.path") == []
|
||||
assert get_text_at_path({}, "any.path") == []
|
||||
assert get_text_at_path(nested_data, "") == [
|
||||
json.dumps(nested_data, sort_keys=True)
|
||||
]
|
||||
assert get_text_at_path(nested_data, "nonexistent") == []
|
||||
assert get_text_at_path(nested_data, "items[99].value") == []
|
||||
assert get_text_at_path(nested_data, "items[*].nonexistent") == []
|
||||
|
||||
assert get_text_at_path(nested_data, "empty") == []
|
||||
assert get_text_at_path(nested_data, "empty_list") == ["[]"]
|
||||
assert get_text_at_path(nested_data, "empty_dict") == ["{}"]
|
||||
|
||||
zeros = get_text_at_path(nested_data, "zeros[*]")
|
||||
assert set(zeros) == {"0", "0.0"}
|
||||
|
||||
assert get_text_at_path(nested_data, "items[].value") == []
|
||||
assert get_text_at_path(nested_data, "items[abc].value") == []
|
||||
assert get_text_at_path(nested_data, "{unclosed") == []
|
||||
assert get_text_at_path(nested_data, "nested[{invalid}]") == []
|
||||
|
||||
|
||||
async def test_async_batch_store(mocker: MockerFixture) -> None:
|
||||
@@ -292,12 +383,14 @@ async def test_cannot_put_empty_namespace() -> None:
|
||||
|
||||
await store.aput(("foo", "langgraph", "foo"), "bar", doc)
|
||||
assert (await store.aget(("foo", "langgraph", "foo"), "bar")).value == doc # type: ignore[union-attr]
|
||||
assert (await store.asearch(("foo", "langgraph", "foo")))[0].value == doc
|
||||
assert (await store.asearch(("foo", "langgraph", "foo"), query="bar"))[
|
||||
0
|
||||
].value == doc
|
||||
await store.adelete(("foo", "langgraph", "foo"), "bar")
|
||||
assert (await store.aget(("foo", "langgraph", "foo"), "bar")) is None
|
||||
store.put(("foo", "langgraph", "foo"), "bar", doc)
|
||||
assert store.get(("foo", "langgraph", "foo"), "bar").value == doc # type: ignore[union-attr]
|
||||
assert store.search(("foo", "langgraph", "foo"))[0].value == doc
|
||||
assert store.search(("foo", "langgraph", "foo"), query="bar")[0].value == doc
|
||||
store.delete(("foo", "langgraph", "foo"), "bar")
|
||||
assert store.get(("foo", "langgraph", "foo"), "bar") is None
|
||||
|
||||
@@ -313,17 +406,6 @@ async def test_cannot_put_empty_namespace() -> None:
|
||||
store.delete(("langgraph", "foo"), "bar")
|
||||
assert store.get(("langgraph", "foo"), "bar") is None
|
||||
|
||||
class MockAsyncBatchedStore(AsyncBatchedBaseStore):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self._store = InMemoryStore()
|
||||
|
||||
def batch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
return self._store.batch(ops)
|
||||
|
||||
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
return self._store.batch(ops)
|
||||
|
||||
async_store = MockAsyncBatchedStore()
|
||||
doc = {"foo": "bar"}
|
||||
|
||||
@@ -344,6 +426,9 @@ async def test_cannot_put_empty_namespace() -> None:
|
||||
assert val is not None
|
||||
assert val.value == doc
|
||||
assert (await async_store.asearch(("foo", "langgraph", "foo")))[0].value == doc
|
||||
assert (await async_store.asearch(("foo", "langgraph", "foo"), query="bar"))[
|
||||
0
|
||||
].value == doc
|
||||
await async_store.adelete(("foo", "langgraph", "foo"), "bar")
|
||||
assert (await async_store.aget(("foo", "langgraph", "foo"), "bar")) is None
|
||||
|
||||
@@ -354,3 +439,511 @@ async def test_cannot_put_empty_namespace() -> None:
|
||||
assert (await async_store.asearch(("valid", "namespace")))[0].value == doc
|
||||
await async_store.adelete(("valid", "namespace"), "key")
|
||||
assert (await async_store.aget(("valid", "namespace"), "key")) is None
|
||||
|
||||
|
||||
async def test_async_batch_store_deduplication(mocker: MockerFixture) -> None:
|
||||
abatch = mocker.spy(InMemoryStore, "batch")
|
||||
store = MockAsyncBatchedStore()
|
||||
|
||||
same_doc = {"value": "same"}
|
||||
diff_doc = {"value": "different"}
|
||||
await asyncio.gather(
|
||||
store.aput(namespace=("test",), key="same", value=same_doc),
|
||||
store.aput(namespace=("test",), key="different", value=diff_doc),
|
||||
)
|
||||
abatch.reset_mock()
|
||||
|
||||
results = await asyncio.gather(
|
||||
store.aget(namespace=("test",), key="same"),
|
||||
store.aget(namespace=("test",), key="same"),
|
||||
store.aget(namespace=("test",), key="different"),
|
||||
)
|
||||
|
||||
assert len(results) == 3
|
||||
assert results[0] == results[1]
|
||||
assert results[0] != results[2]
|
||||
assert results[0].value == same_doc # type: ignore
|
||||
assert results[2].value == diff_doc # type: ignore
|
||||
assert len(abatch.call_args_list) == 1
|
||||
ops = list(abatch.call_args_list[0].args[1])
|
||||
assert len(ops) == 2
|
||||
assert GetOp(("test",), "same") in ops
|
||||
assert GetOp(("test",), "different") in ops
|
||||
|
||||
abatch.reset_mock()
|
||||
|
||||
doc1 = {"value": 1}
|
||||
doc2 = {"value": 2}
|
||||
results = await asyncio.gather(
|
||||
store.aput(namespace=("test",), key="key", value=doc1),
|
||||
store.aput(namespace=("test",), key="key", value=doc2),
|
||||
)
|
||||
assert len(abatch.call_args_list) == 1
|
||||
ops = list(abatch.call_args_list[0].args[1])
|
||||
assert len(ops) == 1
|
||||
assert ops[0] == PutOp(("test",), "key", doc2)
|
||||
assert len(results) == 2
|
||||
assert all(result is None for result in results)
|
||||
|
||||
result = await store.aget(namespace=("test",), key="key")
|
||||
assert result is not None
|
||||
assert result.value == doc2
|
||||
|
||||
abatch.reset_mock()
|
||||
|
||||
results = await asyncio.gather(
|
||||
store.asearch(("test",), filter={"value": 2}),
|
||||
store.asearch(("test",), filter={"value": 2}),
|
||||
)
|
||||
assert len(abatch.call_args_list) == 1
|
||||
ops = list(abatch.call_args_list[0].args[1])
|
||||
assert len(ops) == 1
|
||||
assert len(results) == 2
|
||||
assert results[0] == results[1]
|
||||
assert len(results[0]) == 1
|
||||
assert results[0][0].value == doc2
|
||||
|
||||
abatch.reset_mock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_embeddings() -> CharacterEmbeddings:
|
||||
return CharacterEmbeddings(dims=500)
|
||||
|
||||
|
||||
def test_vector_store_initialization(fake_embeddings: CharacterEmbeddings) -> None:
|
||||
"""Test store initialization with embedding config."""
|
||||
store = InMemoryStore(
|
||||
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
|
||||
)
|
||||
assert store.index_config is not None
|
||||
assert store.index_config["dims"] == fake_embeddings.dims
|
||||
assert store.index_config["embed"] == fake_embeddings
|
||||
|
||||
|
||||
def test_vector_insert_with_auto_embedding(
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
) -> None:
|
||||
"""Test inserting items that get auto-embedded."""
|
||||
store = InMemoryStore(
|
||||
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
|
||||
)
|
||||
docs = [
|
||||
("doc1", {"text": "short text"}),
|
||||
("doc2", {"text": "longer text document"}),
|
||||
("doc3", {"text": "longest text document here"}),
|
||||
("doc4", {"description": "text in description field"}),
|
||||
("doc5", {"content": "text in content field"}),
|
||||
("doc6", {"body": "text in body field"}),
|
||||
]
|
||||
|
||||
for key, value in docs:
|
||||
store.put(("test",), key, value)
|
||||
|
||||
results = store.search(("test",), query="long text")
|
||||
assert len(results) > 0
|
||||
|
||||
doc_order = [r.key for r in results]
|
||||
assert "doc2" in doc_order
|
||||
assert "doc3" in doc_order
|
||||
|
||||
|
||||
async def test_async_vector_insert_with_auto_embedding(
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
) -> None:
|
||||
"""Test inserting items that get auto-embedded using async methods."""
|
||||
store = InMemoryStore(
|
||||
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
|
||||
)
|
||||
docs = [
|
||||
("doc1", {"text": "short text"}),
|
||||
("doc2", {"text": "longer text document"}),
|
||||
("doc3", {"text": "longest text document here"}),
|
||||
("doc4", {"description": "text in description field"}),
|
||||
("doc5", {"content": "text in content field"}),
|
||||
("doc6", {"body": "text in body field"}),
|
||||
]
|
||||
|
||||
for key, value in docs:
|
||||
await store.aput(("test",), key, value)
|
||||
|
||||
results = await store.asearch(("test",), query="long text")
|
||||
assert len(results) > 0
|
||||
|
||||
doc_order = [r.key for r in results]
|
||||
assert "doc2" in doc_order
|
||||
assert "doc3" in doc_order
|
||||
|
||||
|
||||
def test_vector_update_with_embedding(fake_embeddings: CharacterEmbeddings) -> None:
|
||||
"""Test that updating items properly updates their embeddings."""
|
||||
store = InMemoryStore(
|
||||
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
|
||||
)
|
||||
store.put(("test",), "doc1", {"text": "zany zebra Xerxes"})
|
||||
store.put(("test",), "doc2", {"text": "something about dogs"})
|
||||
store.put(("test",), "doc3", {"text": "text about birds"})
|
||||
|
||||
results_initial = store.search(("test",), query="Zany Xerxes")
|
||||
assert len(results_initial) > 0
|
||||
assert results_initial[0].key == "doc1"
|
||||
initial_score = results_initial[0].score
|
||||
assert initial_score is not None
|
||||
|
||||
store.put(("test",), "doc1", {"text": "new text about dogs"})
|
||||
|
||||
results_after = store.search(("test",), query="Zany Xerxes")
|
||||
after_score = next((r.score for r in results_after if r.key == "doc1"), 0.0)
|
||||
assert after_score is not None
|
||||
assert after_score < initial_score
|
||||
|
||||
results_new = store.search(("test",), query="new text about dogs")
|
||||
for r in results_new:
|
||||
if r.key == "doc1":
|
||||
assert r.score > after_score
|
||||
|
||||
# Don't index this one
|
||||
store.put(("test",), "doc4", {"text": "new text about dogs"}, index=False)
|
||||
results_new = store.search(("test",), query="new text about dogs", limit=3)
|
||||
assert not any(r.key == "doc4" for r in results_new)
|
||||
|
||||
|
||||
async def test_async_vector_update_with_embedding(
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
) -> None:
|
||||
"""Test that updating items properly updates their embeddings using async methods."""
|
||||
store = InMemoryStore(
|
||||
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
|
||||
)
|
||||
await store.aput(("test",), "doc1", {"text": "zany zebra Xerxes"})
|
||||
await store.aput(("test",), "doc2", {"text": "something about dogs"})
|
||||
await store.aput(("test",), "doc3", {"text": "text about birds"})
|
||||
|
||||
results_initial = await store.asearch(("test",), query="Zany Xerxes")
|
||||
assert len(results_initial) > 0
|
||||
assert results_initial[0].key == "doc1"
|
||||
initial_score = results_initial[0].score
|
||||
|
||||
await store.aput(("test",), "doc1", {"text": "new text about dogs"})
|
||||
|
||||
results_after = await store.asearch(("test",), query="Zany Xerxes")
|
||||
after_score = next((r.score for r in results_after if r.key == "doc1"), 0.0)
|
||||
assert after_score is not None
|
||||
assert after_score < initial_score
|
||||
|
||||
results_new = await store.asearch(("test",), query="new text about dogs")
|
||||
for r in results_new:
|
||||
if r.key == "doc1":
|
||||
assert r.score is not None
|
||||
assert r.score > after_score
|
||||
|
||||
# Don't index this one
|
||||
await store.aput(("test",), "doc4", {"text": "new text about dogs"}, index=False)
|
||||
results_new = await store.asearch(("test",), query="new text about dogs", limit=3)
|
||||
assert not any(r.key == "doc4" for r in results_new)
|
||||
|
||||
|
||||
def test_vector_search_with_filters(fake_embeddings: CharacterEmbeddings) -> None:
|
||||
"""Test combining vector search with filters."""
|
||||
inmem_store = InMemoryStore(
|
||||
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
|
||||
)
|
||||
# Insert test documents
|
||||
docs = [
|
||||
("doc1", {"text": "red apple", "color": "red", "score": 4.5}),
|
||||
("doc2", {"text": "red car", "color": "red", "score": 3.0}),
|
||||
("doc3", {"text": "green apple", "color": "green", "score": 4.0}),
|
||||
("doc4", {"text": "blue car", "color": "blue", "score": 3.5}),
|
||||
]
|
||||
|
||||
for key, value in docs:
|
||||
inmem_store.put(("test",), key, value)
|
||||
|
||||
results = inmem_store.search(("test",), query="apple", filter={"color": "red"})
|
||||
assert len(results) == 2
|
||||
assert results[0].key == "doc1"
|
||||
|
||||
results = inmem_store.search(("test",), query="car", filter={"color": "red"})
|
||||
assert len(results) == 2
|
||||
assert results[0].key == "doc2"
|
||||
|
||||
results = inmem_store.search(
|
||||
("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
|
||||
)
|
||||
assert len(results) == 3
|
||||
assert results[0].key == "doc4"
|
||||
|
||||
# Multiple filters
|
||||
results = inmem_store.search(
|
||||
("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"}
|
||||
)
|
||||
assert len(results) == 1
|
||||
assert results[0].key == "doc3"
|
||||
|
||||
|
||||
async def test_async_vector_search_with_filters(
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
) -> None:
|
||||
"""Test combining vector search with filters using async methods."""
|
||||
store = InMemoryStore(
|
||||
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
|
||||
)
|
||||
# Insert test documents
|
||||
docs = [
|
||||
("doc1", {"text": "red apple", "color": "red", "score": 4.5}),
|
||||
("doc2", {"text": "red car", "color": "red", "score": 3.0}),
|
||||
("doc3", {"text": "green apple", "color": "green", "score": 4.0}),
|
||||
("doc4", {"text": "blue car", "color": "blue", "score": 3.5}),
|
||||
]
|
||||
|
||||
for key, value in docs:
|
||||
await store.aput(("test",), key, value)
|
||||
|
||||
results = await store.asearch(("test",), query="apple", filter={"color": "red"})
|
||||
assert len(results) == 2
|
||||
assert results[0].key == "doc1"
|
||||
|
||||
results = await store.asearch(("test",), query="car", filter={"color": "red"})
|
||||
assert len(results) == 2
|
||||
assert results[0].key == "doc2"
|
||||
|
||||
results = await store.asearch(
|
||||
("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
|
||||
)
|
||||
assert len(results) == 3
|
||||
assert results[0].key == "doc4"
|
||||
|
||||
# Multiple filters
|
||||
results = await store.asearch(
|
||||
("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"}
|
||||
)
|
||||
assert len(results) == 1
|
||||
assert results[0].key == "doc3"
|
||||
|
||||
|
||||
async def test_async_batched_vector_search_concurrent(
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
) -> None:
|
||||
"""Test concurrent vector search operations using async batched store."""
|
||||
store = MockAsyncBatchedStore(
|
||||
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
|
||||
)
|
||||
|
||||
colors = ["red", "blue", "green", "yellow", "purple"]
|
||||
items = ["apple", "car", "house", "book", "phone"]
|
||||
scores = [3.0, 3.5, 4.0, 4.5, 5.0]
|
||||
|
||||
docs = []
|
||||
for i in range(50):
|
||||
color = colors[i % len(colors)]
|
||||
item = items[i % len(items)]
|
||||
score = scores[i % len(scores)]
|
||||
docs.append(
|
||||
(
|
||||
f"doc{i}",
|
||||
{"text": f"{color} {item}", "color": color, "score": score, "index": i},
|
||||
)
|
||||
)
|
||||
coros = [
|
||||
*[store.aput(("test",), key, value) for key, value in docs],
|
||||
*[store.adelete(("test",), key) for key, value in docs],
|
||||
*[store.aput(("test",), key, value) for key, value in docs],
|
||||
]
|
||||
await asyncio.gather(*coros)
|
||||
|
||||
# Prepare multiple search queries with different filters
|
||||
search_queries: list[tuple[str, dict[str, Any]]] = [
|
||||
("apple", {"color": "red"}),
|
||||
("car", {"color": "blue"}),
|
||||
("house", {"color": "green"}),
|
||||
("phone", {"score": {"$gt": 4.99}}),
|
||||
("book", {"score": {"$lte": 3.5}}),
|
||||
("apple", {"score": {"$gte": 3.0}, "color": "red"}),
|
||||
("car", {"score": {"$lt": 5.1}, "color": "blue"}),
|
||||
("house", {"index": {"$gt": 25}}),
|
||||
("phone", {"index": {"$lte": 10}}),
|
||||
]
|
||||
|
||||
all_results = await asyncio.gather(
|
||||
*[
|
||||
store.asearch(("test",), query=query, filter=filter_)
|
||||
for query, filter_ in search_queries
|
||||
]
|
||||
)
|
||||
|
||||
for results, (query, filter_) in zip(all_results, search_queries):
|
||||
assert len(results) > 0, f"No results for query '{query}' with filter {filter_}"
|
||||
|
||||
for result in results:
|
||||
if "color" in filter_:
|
||||
assert result.value["color"] == filter_["color"]
|
||||
|
||||
if "score" in filter_:
|
||||
score = result.value["score"]
|
||||
for op, value in filter_["score"].items():
|
||||
if op == "$gt":
|
||||
assert score > value
|
||||
elif op == "$gte":
|
||||
assert score >= value
|
||||
elif op == "$lt":
|
||||
assert score < value
|
||||
elif op == "$lte":
|
||||
assert score <= value
|
||||
|
||||
if "index" in filter_:
|
||||
index = result.value["index"]
|
||||
for op, value in filter_["index"].items():
|
||||
if op == "$gt":
|
||||
assert index > value
|
||||
elif op == "$gte":
|
||||
assert index >= value
|
||||
elif op == "$lt":
|
||||
assert index < value
|
||||
elif op == "$lte":
|
||||
assert index <= value
|
||||
|
||||
|
||||
def test_vector_search_pagination(fake_embeddings: CharacterEmbeddings) -> None:
|
||||
"""Test pagination with vector search."""
|
||||
store = InMemoryStore(
|
||||
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
|
||||
)
|
||||
for i in range(5):
|
||||
store.put(("test",), f"doc{i}", {"text": f"test document number {i}"})
|
||||
|
||||
results_page1 = store.search(("test",), query="test", limit=2)
|
||||
results_page2 = store.search(("test",), query="test", limit=2, offset=2)
|
||||
|
||||
assert len(results_page1) == 2
|
||||
assert len(results_page2) == 2
|
||||
assert results_page1[0].key != results_page2[0].key
|
||||
|
||||
all_results = store.search(("test",), query="test", limit=10)
|
||||
assert len(all_results) == 5
|
||||
|
||||
|
||||
async def test_async_vector_search_pagination(
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
) -> None:
|
||||
"""Test pagination with vector search using async methods."""
|
||||
store = InMemoryStore(
|
||||
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
|
||||
)
|
||||
for i in range(5):
|
||||
await store.aput(("test",), f"doc{i}", {"text": f"test document number {i}"})
|
||||
|
||||
results_page1 = await store.asearch(("test",), query="test", limit=2)
|
||||
results_page2 = await store.asearch(("test",), query="test", limit=2, offset=2)
|
||||
|
||||
assert len(results_page1) == 2
|
||||
assert len(results_page2) == 2
|
||||
assert results_page1[0].key != results_page2[0].key
|
||||
|
||||
all_results = await store.asearch(("test",), query="test", limit=10)
|
||||
assert len(all_results) == 5
|
||||
|
||||
|
||||
async def test_embed_with_path(fake_embeddings: CharacterEmbeddings) -> None:
|
||||
# Test store-level field configuration
|
||||
store = InMemoryStore(
|
||||
index={
|
||||
"dims": fake_embeddings.dims,
|
||||
"embed": fake_embeddings,
|
||||
# Key 2 isn't included. Don't index it.
|
||||
"fields": ["key0", "key1", "key3"],
|
||||
}
|
||||
)
|
||||
# This will have 2 vectors representing it
|
||||
doc1 = {
|
||||
# Omit key0 - check it doesn't raise an error
|
||||
"key1": "xxx",
|
||||
"key2": "yyy",
|
||||
"key3": "zzz",
|
||||
}
|
||||
# This will have 3 vectors representing it
|
||||
doc2 = {
|
||||
"key0": "uuu",
|
||||
"key1": "vvv",
|
||||
"key2": "www",
|
||||
"key3": "xxx",
|
||||
}
|
||||
await store.aput(("test",), "doc1", doc1)
|
||||
await store.aput(("test",), "doc2", doc2)
|
||||
|
||||
# doc2.key3 and doc1.key1 both would have the highest score
|
||||
results = await store.asearch(("test",), query="xxx")
|
||||
assert len(results) == 2
|
||||
assert results[0].key != results[1].key
|
||||
ascore = results[0].score
|
||||
bscore = results[1].score
|
||||
assert ascore == bscore
|
||||
assert ascore is not None and bscore is not None
|
||||
|
||||
results = await store.asearch(("test",), query="uuu")
|
||||
assert len(results) == 2
|
||||
assert results[0].key != results[1].key
|
||||
assert results[0].key == "doc2"
|
||||
assert results[0].score is not None and results[0].score > results[1].score
|
||||
assert ascore == pytest.approx(results[0].score, abs=1e-5)
|
||||
|
||||
# Un-indexed - will have low results for both. Not zero (because we're projecting)
|
||||
# but less than the above.
|
||||
results = await store.asearch(("test",), query="www")
|
||||
assert len(results) == 2
|
||||
assert results[0].score < ascore
|
||||
assert results[1].score < ascore
|
||||
|
||||
# Test operation-level field configuration
|
||||
store_no_defaults = InMemoryStore(
|
||||
index={
|
||||
"dims": fake_embeddings.dims,
|
||||
"embed": fake_embeddings,
|
||||
"fields": ["key17"],
|
||||
}
|
||||
)
|
||||
|
||||
doc3 = {
|
||||
"key0": "aaa",
|
||||
"key1": "bbb",
|
||||
"key2": "ccc",
|
||||
"key3": "ddd",
|
||||
}
|
||||
doc4 = {
|
||||
"key0": "eee",
|
||||
"key1": "bbb", # Same as doc3.key1
|
||||
"key2": "fff",
|
||||
"key3": "ggg",
|
||||
}
|
||||
|
||||
await store_no_defaults.aput(("test",), "doc3", doc3, index=["key0", "key1"])
|
||||
await store_no_defaults.aput(("test",), "doc4", doc4, index=["key1", "key3"])
|
||||
|
||||
results = await store_no_defaults.asearch(("test",), query="aaa")
|
||||
assert len(results) == 2
|
||||
assert results[0].key == "doc3"
|
||||
assert results[0].score is not None and results[0].score > results[1].score
|
||||
|
||||
results = await store_no_defaults.asearch(("test",), query="ggg")
|
||||
assert len(results) == 2
|
||||
assert results[0].key == "doc4"
|
||||
assert results[0].score is not None and results[0].score > results[1].score
|
||||
|
||||
results = await store_no_defaults.asearch(("test",), query="bbb")
|
||||
assert len(results) == 2
|
||||
assert results[0].key != results[1].key
|
||||
assert results[0].score == results[1].score
|
||||
|
||||
results = await store_no_defaults.asearch(("test",), query="ccc")
|
||||
assert len(results) == 2
|
||||
assert all(r.score < ascore for r in results)
|
||||
|
||||
doc5 = {
|
||||
"key0": "hhh",
|
||||
"key1": "iii",
|
||||
}
|
||||
await store_no_defaults.aput(("test",), "doc5", doc5, index=False)
|
||||
|
||||
results = await store_no_defaults.asearch(("test",), query="hhh")
|
||||
assert len(results) == 3
|
||||
doc5_result = next(r for r in results if r.key == "doc5")
|
||||
assert doc5_result.score is None
|
||||
|
||||
+103
-8
@@ -1,10 +1,105 @@
|
||||
# langchain-cli
|
||||
# LangGraph CLI
|
||||
|
||||
This package implements the official CLI for LangGraph API.
|
||||
The official command-line interface for LangGraph, providing tools to create, develop, and deploy LangGraph applications.
|
||||
|
||||
## How to Test CLI Changes Locally
|
||||
These instructions are for CLI development and testing. Use the CLI examples to test CLI changes locally.
|
||||
1. Make changes to the CLI code.
|
||||
1. Navigate to the `libs/cli/examples`: `cd libs/cli/examples`
|
||||
1. Install CLI examples dependencies: `poetry install`
|
||||
1. Run/test CLI command (e.g. `langgraph build`).
|
||||
## Installation
|
||||
|
||||
Install via pip:
|
||||
```bash
|
||||
pip install langgraph-cli
|
||||
```
|
||||
|
||||
For development mode with hot reloading:
|
||||
```bash
|
||||
pip install "langgraph-cli[inmem]"
|
||||
```
|
||||
|
||||
## Commands
|
||||
|
||||
### `langgraph new` 🌱
|
||||
Create a new LangGraph project from a template
|
||||
```bash
|
||||
langgraph new [PATH] --template TEMPLATE_NAME
|
||||
```
|
||||
|
||||
### `langgraph dev` 🏃♀️
|
||||
Run LangGraph API server in development mode with hot reloading
|
||||
```bash
|
||||
langgraph dev [OPTIONS]
|
||||
--host TEXT Host to bind to (default: 127.0.0.1)
|
||||
--port INTEGER Port to bind to (default: 2024)
|
||||
--no-reload Disable auto-reload
|
||||
--debug-port INTEGER Enable remote debugging
|
||||
--no-browser Skip opening browser window
|
||||
-c, --config FILE Config file path (default: langgraph.json)
|
||||
```
|
||||
|
||||
### `langgraph up` 🚀
|
||||
Launch LangGraph API server in Docker
|
||||
```bash
|
||||
langgraph up [OPTIONS]
|
||||
-p, --port INTEGER Port to expose (default: 8123)
|
||||
--wait Wait for services to start
|
||||
--watch Restart on file changes
|
||||
--verbose Show detailed logs
|
||||
-c, --config FILE Config file path
|
||||
-d, --docker-compose Additional services file
|
||||
```
|
||||
|
||||
### `langgraph build`
|
||||
Build a Docker image for your LangGraph application
|
||||
```bash
|
||||
langgraph build -t IMAGE_TAG [OPTIONS]
|
||||
--platform TEXT Target platforms (e.g., linux/amd64,linux/arm64)
|
||||
--pull / --no-pull Use latest/local base image
|
||||
-c, --config FILE Config file path
|
||||
```
|
||||
|
||||
### `langgraph dockerfile`
|
||||
Generate a Dockerfile for custom deployments
|
||||
```bash
|
||||
langgraph dockerfile SAVE_PATH [OPTIONS]
|
||||
-c, --config FILE Config file path
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
The CLI uses a `langgraph.json` configuration file with these key settings:
|
||||
|
||||
```json
|
||||
{
|
||||
"dependencies": ["langchain_openai", "./your_package"], // Required: Package dependencies
|
||||
"graphs": {
|
||||
"my_graph": "./your_package/file.py:graph" // Required: Graph definitions
|
||||
},
|
||||
"env": "./.env", // Optional: Environment variables
|
||||
"python_version": "3.11", // Optional: Python version (3.11/3.12)
|
||||
"pip_config_file": "./pip.conf", // Optional: pip configuration
|
||||
"dockerfile_lines": [] // Optional: Additional Dockerfile commands
|
||||
}
|
||||
```
|
||||
|
||||
See the [full documentation](https://langchain-ai.github.io/langgraph/docs/cloud/reference/cli.html) for detailed configuration options.
|
||||
|
||||
## Development
|
||||
|
||||
To develop the CLI itself:
|
||||
|
||||
1. Clone the repository
|
||||
2. Navigate to the CLI directory: `cd libs/cli`
|
||||
3. Install development dependencies: `poetry install`
|
||||
4. Make your changes to the CLI code
|
||||
5. Test your changes:
|
||||
```bash
|
||||
# Run CLI commands directly
|
||||
poetry run langgraph --help
|
||||
|
||||
# Or use the examples
|
||||
cd examples
|
||||
poetry install
|
||||
poetry run langgraph dev # or other commands
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
This project is licensed under the terms specified in the repository's LICENSE file.
|
||||
|
||||
@@ -1299,9 +1299,9 @@ create-jest@^29.7.0:
|
||||
prompts "^2.0.1"
|
||||
|
||||
cross-spawn@^7.0.2, cross-spawn@^7.0.3:
|
||||
version "7.0.3"
|
||||
resolved "https://registry.yarnpkg.com/cross-spawn/-/cross-spawn-7.0.3.tgz#f73a85b9d5d41d045551c177e2882d4ac85728a6"
|
||||
integrity sha512-iRDPJKUPVEND7dHPO8rkbOnPpyDygcDFtWjpeWNCgy8WP2rXcxXL8TskReQl6OrB2G7+UJrags1q15Fudc7G6w==
|
||||
version "7.0.6"
|
||||
resolved "https://registry.yarnpkg.com/cross-spawn/-/cross-spawn-7.0.6.tgz#8a58fe78f00dcd70c370451759dfbfaf03e8ee9f"
|
||||
integrity sha512-uV2QOWP2nWzsy2aMp8aRibhi9dlzF5Hgh5SHaB9OiTGEyDTiJJyx0uy51QXdyWbtAHNua4XJzUKca3OzKUd3vA==
|
||||
dependencies:
|
||||
path-key "^3.1.0"
|
||||
shebang-command "^2.0.0"
|
||||
|
||||
+127
-22
@@ -1,4 +1,4 @@
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
import shutil
|
||||
import sys
|
||||
@@ -45,7 +45,7 @@ OPT_CONFIG = click.option(
|
||||
- "graphs": mapping from graph ID to path where the compiled graph is defined, i.e. ./your_package/your_file.py:variable, where
|
||||
"variable" is an instance of langgraph.graph.graph.CompiledGraph
|
||||
- "env": (optional) path to .env file or a mapping from environment variable to its value
|
||||
- "python_version": (optional) 3.11 or 3.12. Defaults to 3.11
|
||||
- "python_version": (optional) 3.11, 3.12, or 3.13. Defaults to 3.11
|
||||
- "pip_config_file": (optional) path to pip config file
|
||||
- "dockerfile_lines": (optional) array of additional lines to add to Dockerfile following the import from parent image
|
||||
|
||||
@@ -190,7 +190,6 @@ def up(
|
||||
click.secho(
|
||||
"""For local dev, requires env var LANGSMITH_API_KEY with access to LangGraph Cloud closed beta.
|
||||
For production use, requires a license key in env var LANGGRAPH_CLOUD_LICENSE_KEY.""",
|
||||
fg="red",
|
||||
)
|
||||
with Runner() as runner, Progress(message="Pulling...") as set:
|
||||
capabilities = langgraph_cli.docker.check_capabilities(runner)
|
||||
@@ -285,9 +284,11 @@ def _build(
|
||||
subp_exec(
|
||||
"docker",
|
||||
"pull",
|
||||
f"{base_image}:{config_json['node_version']}"
|
||||
if config_json.get("node_version")
|
||||
else f"{base_image}:{config_json['python_version']}",
|
||||
(
|
||||
f"{base_image}:{config_json['node_version']}"
|
||||
if config_json.get("node_version")
|
||||
else f"{base_image}:{config_json['python_version']}"
|
||||
),
|
||||
verbose=True,
|
||||
)
|
||||
)
|
||||
@@ -352,8 +353,7 @@ def build(
|
||||
with Runner() as runner, Progress(message="Pulling...") as set:
|
||||
if shutil.which("docker") is None:
|
||||
raise click.UsageError("Docker not installed") from None
|
||||
with open(config) as f:
|
||||
config_json = langgraph_cli.config.validate_config(json.load(f))
|
||||
config_json = langgraph_cli.config.validate_config_file(config)
|
||||
_build(
|
||||
runner, set, config, config_json, base_image, pull, tag, docker_build_args
|
||||
)
|
||||
@@ -433,8 +433,7 @@ tests
|
||||
def dockerfile(save_path: str, config: pathlib.Path, add_docker_compose: bool) -> None:
|
||||
save_path = pathlib.Path(save_path).absolute()
|
||||
secho(f"🔍 Validating configuration at path: {config}", fg="yellow")
|
||||
with open(config, encoding="utf-8") as f:
|
||||
config_json = langgraph_cli.config.validate_config(json.load(f))
|
||||
config_json = langgraph_cli.config.validate_config_file(config)
|
||||
secho("✅ Configuration validated!", fg="green")
|
||||
|
||||
secho(f"📝 Generating Dockerfile at {save_path}", fg="yellow")
|
||||
@@ -443,9 +442,11 @@ def dockerfile(save_path: str, config: pathlib.Path, add_docker_compose: bool) -
|
||||
langgraph_cli.config.config_to_docker(
|
||||
config,
|
||||
config_json,
|
||||
"langchain/langgraphjs-api"
|
||||
if config_json.get("node_version")
|
||||
else "langchain/langgraph-api",
|
||||
(
|
||||
"langchain/langgraphjs-api"
|
||||
if config_json.get("node_version")
|
||||
else "langchain/langgraph-api"
|
||||
),
|
||||
)
|
||||
)
|
||||
secho("✅ Created: Dockerfile", fg="green")
|
||||
@@ -510,6 +511,107 @@ def dockerfile(save_path: str, config: pathlib.Path, add_docker_compose: bool) -
|
||||
)
|
||||
|
||||
|
||||
@click.option(
|
||||
"--host",
|
||||
default="127.0.0.1",
|
||||
help="Network interface to bind the development server to. Default 127.0.0.1 is recommended for security. Only use 0.0.0.0 in trusted networks",
|
||||
)
|
||||
@click.option(
|
||||
"--port",
|
||||
default=2024,
|
||||
type=int,
|
||||
help="Port number to bind the development server to. Example: langgraph dev --port 8000",
|
||||
)
|
||||
@click.option(
|
||||
"--no-reload",
|
||||
is_flag=True,
|
||||
help="Disable automatic reloading when code changes are detected",
|
||||
)
|
||||
@click.option(
|
||||
"--config",
|
||||
type=click.Path(exists=True),
|
||||
default="langgraph.json",
|
||||
help="Path to configuration file declaring dependencies, graphs and environment variables",
|
||||
)
|
||||
@click.option(
|
||||
"--n-jobs-per-worker",
|
||||
default=None,
|
||||
type=int,
|
||||
help="Maximum number of concurrent jobs each worker process can handle. Default: 10",
|
||||
)
|
||||
@click.option(
|
||||
"--no-browser",
|
||||
is_flag=True,
|
||||
help="Skip automatically opening the browser when the server starts",
|
||||
)
|
||||
@click.option(
|
||||
"--debug-port",
|
||||
default=None,
|
||||
type=int,
|
||||
help="Enable remote debugging by listening on specified port. Requires debugpy to be installed",
|
||||
)
|
||||
@cli.command(
|
||||
"dev",
|
||||
help="🏃♀️➡️ Run LangGraph API server in development mode with hot reloading and debugging support",
|
||||
)
|
||||
@log_command
|
||||
def dev(
|
||||
host: str,
|
||||
port: int,
|
||||
no_reload: bool,
|
||||
config: pathlib.Path,
|
||||
n_jobs_per_worker: Optional[int],
|
||||
no_browser: bool,
|
||||
debug_port: Optional[int],
|
||||
):
|
||||
"""CLI entrypoint for running the LangGraph API server."""
|
||||
try:
|
||||
from langgraph_api.cli import run_server
|
||||
except ImportError:
|
||||
try:
|
||||
import pkg_resources
|
||||
|
||||
pkg_resources.require("langgraph-api-inmem")
|
||||
except (ImportError, pkg_resources.DistributionNotFound):
|
||||
raise click.UsageError(
|
||||
"Required package 'langgraph-api-inmem' is not installed.\n"
|
||||
"Please install it with:\n\n"
|
||||
' pip install -U "langgraph-cli[inmem]"\n\n'
|
||||
"If you're developing the langgraph-cli package locally, you can install in development mode:\n"
|
||||
" pip install -e ."
|
||||
) from None
|
||||
raise click.UsageError(
|
||||
"Could not import run_server. This likely means your installation is incomplete.\n"
|
||||
"Please ensure langgraph-cli is installed with the 'inmem' extra: pip install -U \"langgraph-cli[inmem]\""
|
||||
) from None
|
||||
|
||||
config_json = langgraph_cli.config.validate_config_file(config)
|
||||
cwd = os.getcwd()
|
||||
sys.path.append(cwd)
|
||||
dependencies = config_json.get("dependencies", [])
|
||||
for dep in dependencies:
|
||||
dep_path = pathlib.Path(cwd) / dep
|
||||
if dep_path.is_dir() and dep_path.exists():
|
||||
sys.path.append(str(dep_path))
|
||||
|
||||
graphs = config_json.get("graphs", {})
|
||||
additional_config = {}
|
||||
if config_json.get("store"):
|
||||
additional_config["store"] = config_json["store"]
|
||||
|
||||
run_server(
|
||||
host,
|
||||
port,
|
||||
not no_reload,
|
||||
graphs,
|
||||
n_jobs_per_worker=n_jobs_per_worker,
|
||||
open_browser=not no_browser,
|
||||
debug_port=debug_port,
|
||||
env=config_json.get("env", None),
|
||||
config=additional_config,
|
||||
)
|
||||
|
||||
|
||||
@click.argument("path", required=False)
|
||||
@click.option(
|
||||
"--template",
|
||||
@@ -556,9 +658,11 @@ def prepare_args_and_stdin(
|
||||
config_path,
|
||||
config,
|
||||
watch=watch,
|
||||
base_image="langchain/langgraphjs-api"
|
||||
if config.get("node_version")
|
||||
else "langchain/langgraph-api",
|
||||
base_image=(
|
||||
"langchain/langgraphjs-api"
|
||||
if config.get("node_version")
|
||||
else "langchain/langgraph-api"
|
||||
),
|
||||
)
|
||||
return args, stdin
|
||||
|
||||
@@ -577,17 +681,18 @@ def prepare(
|
||||
debugger_base_url: Optional[str] = None,
|
||||
postgres_uri: Optional[str] = None,
|
||||
):
|
||||
with open(config_path) as f:
|
||||
config = langgraph_cli.config.validate_config(json.load(f))
|
||||
config_json = langgraph_cli.config.validate_config_file(config_path)
|
||||
# pull latest images
|
||||
if pull:
|
||||
runner.run(
|
||||
subp_exec(
|
||||
"docker",
|
||||
"pull",
|
||||
f"langchain/langgraphjs-api:{config['node_version']}"
|
||||
if config.get("node_version")
|
||||
else f"langchain/langgraph-api:{config['python_version']}",
|
||||
(
|
||||
f"langchain/langgraphjs-api:{config_json['node_version']}"
|
||||
if config_json.get("node_version")
|
||||
else f"langchain/langgraph-api:{config_json['python_version']}"
|
||||
),
|
||||
verbose=verbose,
|
||||
)
|
||||
)
|
||||
@@ -595,7 +700,7 @@ def prepare(
|
||||
args, stdin = prepare_args_and_stdin(
|
||||
capabilities=capabilities,
|
||||
config_path=config_path,
|
||||
config=config,
|
||||
config=config_json,
|
||||
docker_compose=docker_compose,
|
||||
port=port,
|
||||
watch=watch,
|
||||
|
||||
@@ -6,8 +6,48 @@ from typing import NamedTuple, Optional, TypedDict, Union
|
||||
|
||||
import click
|
||||
|
||||
MIN_NODE_VERSION = "20"
|
||||
MIN_PYTHON_VERSION = "3.11"
|
||||
|
||||
class Config(TypedDict):
|
||||
|
||||
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: 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: 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):
|
||||
python_version: str
|
||||
node_version: Optional[str]
|
||||
pip_config_file: Optional[str]
|
||||
@@ -15,6 +55,29 @@ class Config(TypedDict):
|
||||
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]:
|
||||
"""Parse a version string into a tuple of (major, minor)."""
|
||||
try:
|
||||
major, minor = map(int, version_str.split("."))
|
||||
return (major, minor)
|
||||
except ValueError:
|
||||
raise click.UsageError(f"Invalid version format: {version_str}") from None
|
||||
|
||||
|
||||
def _parse_node_version(version_str: str) -> int:
|
||||
"""Parse a Node.js version string into a major version number."""
|
||||
try:
|
||||
if "." in version_str:
|
||||
raise ValueError("Node.js version must be major version only")
|
||||
return int(version_str)
|
||||
except ValueError:
|
||||
raise click.UsageError(
|
||||
f"Invalid Node.js version format: {version_str}. "
|
||||
"Use major version only (e.g., '20')."
|
||||
) from None
|
||||
|
||||
|
||||
def validate_config(config: Config) -> Config:
|
||||
@@ -24,6 +87,7 @@ 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 {
|
||||
@@ -33,25 +97,39 @@ def validate_config(config: Config) -> Config:
|
||||
"dependencies": config.get("dependencies", []),
|
||||
"graphs": config.get("graphs", {}),
|
||||
"env": config.get("env", {}),
|
||||
"store": config.get("store"),
|
||||
}
|
||||
)
|
||||
|
||||
if config.get("node_version"):
|
||||
if config["node_version"] not in ("20",):
|
||||
raise click.UsageError(
|
||||
f"Unsupported Node.js version: {config['node_version']}. "
|
||||
"Currently only `node_version: \"20\"` is supported."
|
||||
)
|
||||
node_version = config["node_version"]
|
||||
try:
|
||||
major = _parse_node_version(node_version)
|
||||
min_major = _parse_node_version(MIN_NODE_VERSION)
|
||||
if major < min_major:
|
||||
raise click.UsageError(
|
||||
f"Node.js version {node_version} is not supported. "
|
||||
f"Minimum required version is {MIN_NODE_VERSION}."
|
||||
)
|
||||
except ValueError as e:
|
||||
raise click.UsageError(str(e)) from None
|
||||
|
||||
if config.get("python_version"):
|
||||
if config["python_version"] not in (
|
||||
"3.11",
|
||||
"3.12",
|
||||
pyversion = config["python_version"]
|
||||
if not pyversion.count(".") == 1 or not all(
|
||||
part.isdigit() for part in pyversion.split(".")
|
||||
):
|
||||
raise click.UsageError(
|
||||
f"Unsupported Python version: {config['python_version']}. "
|
||||
"Supported versions are 3.11 and 3.12."
|
||||
f"Invalid Python version format: {pyversion}. "
|
||||
"Use 'major.minor' format (e.g., '3.11'). "
|
||||
"Patch version cannot be specified."
|
||||
)
|
||||
if _parse_version(pyversion) < _parse_version(MIN_PYTHON_VERSION):
|
||||
raise click.UsageError(
|
||||
f"Python version {pyversion} is not supported. "
|
||||
f"Minimum required version is {MIN_PYTHON_VERSION}."
|
||||
)
|
||||
|
||||
if not config["dependencies"]:
|
||||
raise click.UsageError(
|
||||
"No dependencies found in config. "
|
||||
@@ -66,6 +144,48 @@ def validate_config(config: Config) -> Config:
|
||||
return config
|
||||
|
||||
|
||||
def validate_config_file(config_path: pathlib.Path) -> Config:
|
||||
with open(config_path) as f:
|
||||
config = json.load(f)
|
||||
validated = validate_config(config)
|
||||
# Enforce the package.json doesn't enforce an
|
||||
# incompatible Node.js version
|
||||
if validated.get("node_version"):
|
||||
package_json_path = config_path.parent / "package.json"
|
||||
if package_json_path.is_file():
|
||||
try:
|
||||
with open(package_json_path) as f:
|
||||
package_json = json.load(f)
|
||||
if "engines" in package_json:
|
||||
engines = package_json["engines"]
|
||||
if any(engine != "node" for engine in engines.keys()):
|
||||
raise click.UsageError(
|
||||
"Only 'node' engine is supported in package.json engines."
|
||||
f" Got engines: {list(engines.keys())}"
|
||||
)
|
||||
if engines:
|
||||
node_version = engines["node"]
|
||||
try:
|
||||
major = _parse_node_version(node_version)
|
||||
min_major = _parse_node_version(MIN_NODE_VERSION)
|
||||
if major < min_major:
|
||||
raise click.UsageError(
|
||||
f"Node.js version in package.json engines must be >= {MIN_NODE_VERSION} "
|
||||
f"(major version only), got '{node_version}'. Minor/patch versions "
|
||||
"(like '20.x.y') are not supported to prevent deployment issues "
|
||||
"when new Node.js versions are released."
|
||||
)
|
||||
except ValueError as e:
|
||||
raise click.UsageError(str(e)) from None
|
||||
|
||||
except json.JSONDecodeError:
|
||||
raise click.UsageError(
|
||||
"Invalid package.json found in langgraph "
|
||||
f"config directory {package_json_path}: file is not valid JSON"
|
||||
) from None
|
||||
return validated
|
||||
|
||||
|
||||
class LocalDeps(NamedTuple):
|
||||
pip_reqs: list[tuple[pathlib.Path, str]]
|
||||
real_pkgs: dict[pathlib.Path, str]
|
||||
@@ -272,7 +392,16 @@ RUN set -ex && \\
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
additional_config = {}
|
||||
if config.get("store"):
|
||||
additional_config["store"] = config["store"]
|
||||
env_additional_config = (
|
||||
""
|
||||
if not additional_config
|
||||
else f"""
|
||||
ENV LANGGRAPH_CONFIG='{json.dumps(additional_config)}'
|
||||
"""
|
||||
)
|
||||
return f"""FROM {base_image}:{config['python_version']}
|
||||
|
||||
{os.linesep.join(config["dockerfile_lines"])}
|
||||
@@ -280,7 +409,7 @@ RUN set -ex && \\
|
||||
{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 ""}"""
|
||||
@@ -310,7 +439,16 @@ def node_config_to_docker(config_path: pathlib.Path, config: Config, base_image:
|
||||
install_cmd = "npm ci"
|
||||
else:
|
||||
install_cmd = "npm i"
|
||||
|
||||
additional_config = {}
|
||||
if config.get("store"):
|
||||
additional_config["store"] = config["store"]
|
||||
env_additional_config = (
|
||||
""
|
||||
if not additional_config
|
||||
else f"""
|
||||
ENV LANGGRAPH_CONFIG='{json.dumps(additional_config)}'
|
||||
"""
|
||||
)
|
||||
return f"""FROM {base_image}:{config['node_version']}
|
||||
|
||||
{os.linesep.join(config["dockerfile_lines"])}
|
||||
@@ -318,7 +456,7 @@ def node_config_to_docker(config_path: pathlib.Path, config: Config, base_image:
|
||||
ADD . {faux_path}
|
||||
|
||||
RUN cd {faux_path} && {install_cmd}
|
||||
|
||||
{env_additional_config}
|
||||
ENV LANGSERVE_GRAPHS='{json.dumps(config["graphs"])}'
|
||||
|
||||
WORKDIR {faux_path}
|
||||
|
||||
Generated
+1347
-111
File diff suppressed because it is too large
Load Diff
@@ -1,12 +1,12 @@
|
||||
[tool.poetry]
|
||||
name = "langgraph-cli"
|
||||
version = "0.1.54"
|
||||
version = "0.1.60"
|
||||
description = "CLI for interacting with LangGraph API"
|
||||
authors = []
|
||||
license = "MIT"
|
||||
readme = "README.md"
|
||||
repository = "https://www.github.com/langchain-ai/langgraph"
|
||||
packages = [{include = "langgraph_cli"}]
|
||||
packages = [{ include = "langgraph_cli" }]
|
||||
|
||||
[tool.poetry.scripts]
|
||||
langgraph = "langgraph_cli.cli:cli"
|
||||
@@ -14,6 +14,8 @@ langgraph = "langgraph_cli.cli:cli"
|
||||
[tool.poetry.dependencies]
|
||||
python = "^3.9.0,<4.0"
|
||||
click = "^8.1.7"
|
||||
langgraph-api = { version = ">=0.0.5,<0.1.0", optional = true, python = ">=3.11,<4.0" }
|
||||
python-dotenv = { version = ">=0.8.0", optional = true }
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
ruff = "^0.6.2"
|
||||
@@ -24,6 +26,9 @@ pytest-mock = "^3.11.1"
|
||||
pytest-watch = "^4.2.0"
|
||||
mypy = "^1.10.0"
|
||||
|
||||
[tool.poetry.extras]
|
||||
inmem = ["langgraph-api", "python-dotenv"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
# --strict-markers will raise errors on unknown marks.
|
||||
# https://docs.pytest.org/en/7.1.x/how-to/mark.html#raising-errors-on-unknown-marks
|
||||
@@ -52,4 +57,4 @@ lint.select = [
|
||||
# isort
|
||||
"I",
|
||||
]
|
||||
lint.ignore = [ "E501", "B008" ]
|
||||
lint.ignore = ["E501", "B008"]
|
||||
|
||||
@@ -1,10 +1,17 @@
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
import tempfile
|
||||
|
||||
import click
|
||||
import pytest
|
||||
|
||||
from langgraph_cli.config import config_to_compose, config_to_docker, validate_config
|
||||
from langgraph_cli.config import (
|
||||
config_to_compose,
|
||||
config_to_docker,
|
||||
validate_config,
|
||||
validate_config_file,
|
||||
)
|
||||
from langgraph_cli.util import clean_empty_lines
|
||||
|
||||
PATH_TO_CONFIG = pathlib.Path(__file__).parent / "test_config.json"
|
||||
@@ -23,6 +30,7 @@ def test_validate_config():
|
||||
"pip_config_file": None,
|
||||
"dockerfile_lines": [],
|
||||
"env": {},
|
||||
"store": None,
|
||||
**expected_config,
|
||||
}
|
||||
actual_config = validate_config(expected_config)
|
||||
@@ -39,9 +47,13 @@ def test_validate_config():
|
||||
"agent": "./agent.py:graph",
|
||||
},
|
||||
"env": env,
|
||||
"store": None,
|
||||
}
|
||||
actual_config = validate_config(expected_config)
|
||||
assert actual_config == expected_config
|
||||
expected_config["python_version"] = "3.13"
|
||||
actual_config = validate_config(expected_config)
|
||||
assert actual_config == expected_config
|
||||
|
||||
# check wrong python version raises
|
||||
with pytest.raises(click.UsageError):
|
||||
@@ -61,6 +73,86 @@ def test_validate_config():
|
||||
with pytest.raises(click.UsageError):
|
||||
validate_config({"python_version": "3.9", "dependencies": ["."]})
|
||||
|
||||
with pytest.raises(click.UsageError) as exc_info:
|
||||
validate_config({"python_version": "3.11.0"})
|
||||
assert "Invalid Python version format" in str(exc_info.value)
|
||||
|
||||
with pytest.raises(click.UsageError) as exc_info:
|
||||
validate_config({"python_version": "3"})
|
||||
assert "Invalid Python version format" in str(exc_info.value)
|
||||
|
||||
with pytest.raises(click.UsageError) as exc_info:
|
||||
validate_config({"python_version": "abc.def"})
|
||||
assert "Invalid Python version format" in str(exc_info.value)
|
||||
|
||||
with pytest.raises(click.UsageError) as exc_info:
|
||||
validate_config({"python_version": "3.10"})
|
||||
assert "Minimum required version" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_validate_config_file():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmpdir_path = pathlib.Path(tmpdir)
|
||||
|
||||
config_path = tmpdir_path / "langgraph.json"
|
||||
|
||||
node_config = {"node_version": "20", "graphs": {"agent": "./agent.js:graph"}}
|
||||
with open(config_path, "w") as f:
|
||||
json.dump(node_config, f)
|
||||
|
||||
validate_config_file(config_path)
|
||||
|
||||
package_json = {"name": "test", "engines": {"node": "20"}}
|
||||
with open(tmpdir_path / "package.json", "w") as f:
|
||||
json.dump(package_json, f)
|
||||
validate_config_file(config_path)
|
||||
|
||||
package_json["engines"]["node"] = "20.18"
|
||||
with open(tmpdir_path / "package.json", "w") as f:
|
||||
json.dump(package_json, f)
|
||||
with pytest.raises(click.UsageError, match="Use major version only"):
|
||||
validate_config_file(config_path)
|
||||
|
||||
package_json["engines"] = {"node": "18"}
|
||||
with open(tmpdir_path / "package.json", "w") as f:
|
||||
json.dump(package_json, f)
|
||||
with pytest.raises(click.UsageError, match="must be >= 20"):
|
||||
validate_config_file(config_path)
|
||||
|
||||
package_json["engines"] = {"node": "20", "deno": "1.0"}
|
||||
with open(tmpdir_path / "package.json", "w") as f:
|
||||
json.dump(package_json, f)
|
||||
with pytest.raises(click.UsageError, match="Only 'node' engine is supported"):
|
||||
validate_config_file(config_path)
|
||||
|
||||
with open(tmpdir_path / "package.json", "w") as f:
|
||||
f.write("{invalid json")
|
||||
with pytest.raises(click.UsageError, match="Invalid package.json"):
|
||||
validate_config_file(config_path)
|
||||
|
||||
python_config = {
|
||||
"python_version": "3.11",
|
||||
"dependencies": ["."],
|
||||
"graphs": {"agent": "./agent.py:graph"},
|
||||
}
|
||||
with open(config_path, "w") as f:
|
||||
json.dump(python_config, f)
|
||||
|
||||
validate_config_file(config_path)
|
||||
|
||||
for package_content in [
|
||||
{"name": "test"},
|
||||
{"engines": {"node": "18"}},
|
||||
{"engines": {"node": "20", "deno": "1.0"}},
|
||||
"{invalid json",
|
||||
]:
|
||||
with open(tmpdir_path / "package.json", "w") as f:
|
||||
if isinstance(package_content, dict):
|
||||
json.dump(package_content, f)
|
||||
else:
|
||||
f.write(package_content)
|
||||
validate_config_file(config_path)
|
||||
|
||||
|
||||
# config_to_docker
|
||||
def test_config_to_docker_simple():
|
||||
|
||||
@@ -48,8 +48,13 @@ test:
|
||||
make stop-postgres; \
|
||||
exit $$EXIT_CODE
|
||||
|
||||
WORKERS ?= auto
|
||||
XDIST_ARGS := $(if $(WORKERS),-n $(WORKERS) --dist worksteal,)
|
||||
MAXFAIL ?=
|
||||
MAXFAIL_ARGS := $(if $(MAXFAIL),--maxfail $(MAXFAIL),)
|
||||
|
||||
test_watch:
|
||||
make start-postgres && poetry run ptw . -- --ff -vv -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \
|
||||
make start-postgres && poetry run ptw . -- --ff -vv -x $(XDIST_ARGS) $(MAXFAIL_ARGS) --snapshot-update --tb short $(TEST); \
|
||||
EXIT_CODE=$$?; \
|
||||
make stop-postgres; \
|
||||
exit $$EXIT_CODE
|
||||
|
||||
@@ -238,7 +238,7 @@ final_state["messages"][-1].content
|
||||
* [How-to Guides](https://langchain-ai.github.io/langgraph/how-tos/): Accomplish specific things within LangGraph, from streaming, to adding memory & persistence, to common design patterns (branching, subgraphs, etc.), these are the place to go if you want to copy and run a specific code snippet.
|
||||
* [Conceptual Guides](https://langchain-ai.github.io/langgraph/concepts/high_level/): In-depth explanations of the key concepts and principles behind LangGraph, such as nodes, edges, state and more.
|
||||
* [API Reference](https://langchain-ai.github.io/langgraph/reference/graphs/): Review important classes and methods, simple examples of how to use the graph and checkpointing APIs, higher-level prebuilt components and more.
|
||||
* [Cloud (beta)](https://langchain-ai.github.io/langgraph/cloud/): With one click, deploy LangGraph applications to LangGraph Cloud.
|
||||
* [LangGraph Platform](https://langchain-ai.github.io/langgraph/concepts/#langgraph-platform): LangGraph Platform is a commercial solution for deploying agentic applications in production, built on the open-source LangGraph framework.
|
||||
|
||||
## Contributing
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@ from enum import Enum
|
||||
from typing import Any, Sequence
|
||||
|
||||
from langgraph.checkpoint.base import EmptyChannelError # noqa: F401
|
||||
from langgraph.types import Interrupt
|
||||
from langgraph.types import Command, Interrupt
|
||||
|
||||
# EmptyChannelError re-exported for backwards compatibility
|
||||
|
||||
@@ -58,7 +58,11 @@ class InvalidUpdateError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class GraphInterrupt(Exception):
|
||||
class GraphBubbleUp(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class GraphInterrupt(GraphBubbleUp):
|
||||
"""Raised when a subgraph is interrupted, suppressed by the root graph.
|
||||
Never raised directly, or surfaced to the user."""
|
||||
|
||||
@@ -73,13 +77,20 @@ class NodeInterrupt(GraphInterrupt):
|
||||
super().__init__([Interrupt(value=value)])
|
||||
|
||||
|
||||
class GraphDelegate(Exception):
|
||||
class GraphDelegate(GraphBubbleUp):
|
||||
"""Raised when a graph is delegated (for distributed mode)."""
|
||||
|
||||
def __init__(self, *args: dict[str, Any]) -> None:
|
||||
super().__init__(*args)
|
||||
|
||||
|
||||
class ParentCommand(GraphBubbleUp):
|
||||
args: tuple[Command]
|
||||
|
||||
def __init__(self, command: Command) -> None:
|
||||
super().__init__(command)
|
||||
|
||||
|
||||
class EmptyInputError(Exception):
|
||||
"""Raised when graph receives an empty input."""
|
||||
|
||||
|
||||
@@ -37,7 +37,12 @@ from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.channels.named_barrier_value import NamedBarrierValue
|
||||
from langgraph.constants import EMPTY_SEQ, NS_END, NS_SEP, SELF, TAG_HIDDEN
|
||||
from langgraph.errors import ErrorCode, InvalidUpdateError, create_error_message
|
||||
from langgraph.errors import (
|
||||
ErrorCode,
|
||||
InvalidUpdateError,
|
||||
ParentCommand,
|
||||
create_error_message,
|
||||
)
|
||||
from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send
|
||||
from langgraph.managed.base import (
|
||||
ChannelKeyPlaceholder,
|
||||
@@ -623,20 +628,27 @@ class CompiledStateGraph(CompiledGraph):
|
||||
|
||||
def _get_root(input: Any) -> Any:
|
||||
if isinstance(input, Command):
|
||||
if input.graph == Command.PARENT:
|
||||
return SKIP_WRITE
|
||||
return input.update
|
||||
else:
|
||||
return input
|
||||
|
||||
# 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 SKIP_WRITE
|
||||
elif isinstance(input, dict):
|
||||
if all(k not in output_keys for k in input):
|
||||
raise InvalidUpdateError(
|
||||
f"Expected node {key} to update at least one of {output_keys}, got {input}"
|
||||
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 SKIP_WRITE
|
||||
return _get_state_key(input.update, key=key)
|
||||
elif get_type_hints(type(input)):
|
||||
value = getattr(input, key, SKIP_WRITE)
|
||||
@@ -819,6 +831,8 @@ def _control_branch(value: Any) -> Sequence[Union[str, Send]]:
|
||||
return [value]
|
||||
if not isinstance(value, GraphCommand):
|
||||
return EMPTY_SEQ
|
||||
if value.graph == Command.PARENT:
|
||||
raise ParentCommand(value)
|
||||
rtn: list[Union[str, Send]] = []
|
||||
if isinstance(value.goto, str):
|
||||
rtn.append(value.goto)
|
||||
@@ -836,6 +850,8 @@ async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]:
|
||||
return [value]
|
||||
if not isinstance(value, GraphCommand):
|
||||
return EMPTY_SEQ
|
||||
if value.graph == Command.PARENT:
|
||||
raise ParentCommand(value)
|
||||
rtn: list[Union[str, Send]] = []
|
||||
if isinstance(value.goto, str):
|
||||
rtn.append(value.goto)
|
||||
|
||||
@@ -212,6 +212,7 @@ def create_react_agent(
|
||||
Args:
|
||||
model: The `LangChain` chat model that supports tool calling.
|
||||
tools: A list of tools, a ToolExecutor, or a ToolNode instance.
|
||||
If an empty list is provided, the agent will consist of a single LLM node without tool calling.
|
||||
state_schema: An optional state schema that defines graph state.
|
||||
Must have `messages` and `is_last_step` keys.
|
||||
Defaults to `AgentState` that defines those two keys.
|
||||
@@ -540,19 +541,10 @@ def create_react_agent(
|
||||
# get the tool functions wrapped in a tool class from the ToolNode
|
||||
tool_classes = list(tool_node.tools_by_name.values())
|
||||
|
||||
if _should_bind_tools(model, tool_classes):
|
||||
model = cast(BaseChatModel, model).bind_tools(tool_classes)
|
||||
tool_calling_enabled = len(tool_classes) > 0
|
||||
|
||||
# Define the function that determines whether to continue or not
|
||||
def should_continue(state: AgentState) -> Literal["tools", "__end__"]:
|
||||
messages = state["messages"]
|
||||
last_message = messages[-1]
|
||||
# If there is no function call, then we finish
|
||||
if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
|
||||
return "__end__"
|
||||
# Otherwise if there is, we continue
|
||||
else:
|
||||
return "tools"
|
||||
if _should_bind_tools(model, tool_classes) and tool_calling_enabled:
|
||||
model = cast(BaseChatModel, model).bind_tools(tool_classes)
|
||||
|
||||
# we're passing store here for validation
|
||||
preprocessor = _get_model_preprocessing_runnable(
|
||||
@@ -635,6 +627,30 @@ def create_react_agent(
|
||||
# We return a list, because this will get added to the existing list
|
||||
return {"messages": [response]}
|
||||
|
||||
if not tool_calling_enabled:
|
||||
# Define a new graph
|
||||
workflow = StateGraph(state_schema or AgentState)
|
||||
workflow.add_node("agent", RunnableCallable(call_model, acall_model))
|
||||
workflow.set_entry_point("agent")
|
||||
return workflow.compile(
|
||||
checkpointer=checkpointer,
|
||||
store=store,
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
debug=debug,
|
||||
)
|
||||
|
||||
# Define the function that determines whether to continue or not
|
||||
def should_continue(state: AgentState) -> Literal["tools", "__end__"]:
|
||||
messages = state["messages"]
|
||||
last_message = messages[-1]
|
||||
# If there is no function call, then we finish
|
||||
if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
|
||||
return "__end__"
|
||||
# Otherwise if there is, we continue
|
||||
else:
|
||||
return "tools"
|
||||
|
||||
# Define a new graph
|
||||
workflow = StateGraph(state_schema or AgentState)
|
||||
|
||||
|
||||
@@ -37,7 +37,7 @@ from langchain_core.tools import tool as create_tool
|
||||
from langchain_core.tools.base import get_all_basemodel_annotations
|
||||
from typing_extensions import Annotated, get_args, get_origin
|
||||
|
||||
from langgraph.errors import GraphInterrupt
|
||||
from langgraph.errors import GraphBubbleUp
|
||||
from langgraph.store.base import BaseStore
|
||||
from langgraph.utils.runnable import RunnableCallable
|
||||
|
||||
@@ -275,7 +275,7 @@ class ToolNode(RunnableCallable):
|
||||
# (2) a NodeInterrupt is raised inside a graph node for a graph called as a tool
|
||||
# (3) a GraphInterrupt is raised when a subgraph is interrupted inside a graph called as a tool
|
||||
# (2 and 3 can happen in a "supervisor w/ tools" multi-agent architecture)
|
||||
except GraphInterrupt as e:
|
||||
except GraphBubbleUp as e:
|
||||
raise e
|
||||
except Exception as e:
|
||||
if isinstance(self.handle_tool_errors, tuple):
|
||||
@@ -316,7 +316,7 @@ class ToolNode(RunnableCallable):
|
||||
# (2) a NodeInterrupt is raised inside a graph node for a graph called as a tool
|
||||
# (3) a GraphInterrupt is raised when a subgraph is interrupted inside a graph called as a tool
|
||||
# (2 and 3 can happen in a "supervisor w/ tools" multi-agent architecture)
|
||||
except GraphInterrupt as e:
|
||||
except GraphBubbleUp as e:
|
||||
raise e
|
||||
except Exception as e:
|
||||
if isinstance(self.handle_tool_errors, tuple):
|
||||
|
||||
@@ -107,6 +107,7 @@ from langgraph.types import (
|
||||
Checkpointer,
|
||||
LoopProtocol,
|
||||
StateSnapshot,
|
||||
StreamChunk,
|
||||
StreamMode,
|
||||
)
|
||||
from langgraph.utils.config import (
|
||||
@@ -1752,6 +1753,10 @@ class Pregel(PregelProtocol):
|
||||
|
||||
stream = AsyncQueue()
|
||||
aioloop = asyncio.get_running_loop()
|
||||
stream_put = cast(
|
||||
Callable[[StreamChunk], None],
|
||||
partial(aioloop.call_soon_threadsafe, stream.put_nowait),
|
||||
)
|
||||
|
||||
def output() -> Iterator:
|
||||
while True:
|
||||
@@ -1806,12 +1811,14 @@ class Pregel(PregelProtocol):
|
||||
# set up messages stream mode
|
||||
if "messages" in stream_modes:
|
||||
run_manager.inheritable_handlers.append(
|
||||
StreamMessagesHandler(stream.put_nowait)
|
||||
StreamMessagesHandler(stream_put)
|
||||
)
|
||||
# set up custom stream mode
|
||||
if "custom" in stream_modes:
|
||||
config[CONF][CONFIG_KEY_STREAM_WRITER] = lambda c: stream.put_nowait(
|
||||
((), "custom", c)
|
||||
config[CONF][CONFIG_KEY_STREAM_WRITER] = (
|
||||
lambda c: aioloop.call_soon_threadsafe(
|
||||
stream.put_nowait, ((), "custom", c)
|
||||
)
|
||||
)
|
||||
async with AsyncPregelLoop(
|
||||
input,
|
||||
@@ -1838,7 +1845,9 @@ class Pregel(PregelProtocol):
|
||||
)
|
||||
# enable subgraph streaming
|
||||
if subgraphs:
|
||||
loop.config[CONF][CONFIG_KEY_STREAM] = loop.stream
|
||||
loop.config[CONF][CONFIG_KEY_STREAM] = StreamProtocol(
|
||||
stream_put, stream_modes
|
||||
)
|
||||
# enable concurrent streaming
|
||||
if subgraphs or "messages" in stream_modes or "custom" in stream_modes:
|
||||
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import sys
|
||||
from collections import defaultdict, deque
|
||||
from functools import partial
|
||||
from hashlib import sha1
|
||||
@@ -66,6 +67,7 @@ from langgraph.types import All, LoopProtocol, PregelExecutableTask, PregelTask
|
||||
from langgraph.utils.config import merge_configs, patch_config
|
||||
|
||||
GetNextVersion = Callable[[Optional[V], BaseChannel], V]
|
||||
SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11)
|
||||
|
||||
|
||||
class WritesProtocol(Protocol):
|
||||
@@ -602,6 +604,7 @@ def prepare_single_task(
|
||||
None,
|
||||
task_id,
|
||||
task_path,
|
||||
writers=proc.flat_writers,
|
||||
)
|
||||
|
||||
else:
|
||||
@@ -633,6 +636,12 @@ def prepare_single_task(
|
||||
)
|
||||
except StopIteration:
|
||||
return
|
||||
except Exception as exc:
|
||||
if SUPPORTS_EXC_NOTES:
|
||||
exc.add_note(
|
||||
f"Before task with name '{name}' and path '{task_path[:3]}'"
|
||||
)
|
||||
raise
|
||||
|
||||
# create task id
|
||||
checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name
|
||||
@@ -720,6 +729,7 @@ def prepare_single_task(
|
||||
None,
|
||||
task_id,
|
||||
task_path,
|
||||
writers=proc.flat_writers,
|
||||
)
|
||||
else:
|
||||
return PregelTask(task_id, name, task_path)
|
||||
|
||||
@@ -191,6 +191,14 @@ def map_debug_checkpoint(
|
||||
"state": t.state,
|
||||
}
|
||||
if t.error
|
||||
else {
|
||||
"id": t.id,
|
||||
"name": t.name,
|
||||
"result": t.result,
|
||||
"interrupts": tuple(asdict(i) for i in t.interrupts),
|
||||
"state": t.state,
|
||||
}
|
||||
if t.result
|
||||
else {
|
||||
"id": t.id,
|
||||
"name": t.name,
|
||||
|
||||
@@ -20,7 +20,7 @@ from langchain_core.runnables import RunnableConfig
|
||||
from langchain_core.runnables.config import get_executor_for_config
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
from langgraph.errors import GraphInterrupt
|
||||
from langgraph.errors import GraphBubbleUp
|
||||
|
||||
P = ParamSpec("P")
|
||||
T = TypeVar("T")
|
||||
@@ -68,7 +68,7 @@ class BackgroundExecutor(ContextManager):
|
||||
def done(self, task: concurrent.futures.Future) -> None:
|
||||
try:
|
||||
task.result()
|
||||
except GraphInterrupt:
|
||||
except GraphBubbleUp:
|
||||
# This exception is an interruption signal, not an error
|
||||
# so we don't want to re-raise it on exit
|
||||
self.tasks.pop(task)
|
||||
@@ -155,7 +155,7 @@ class AsyncBackgroundExecutor(AsyncContextManager):
|
||||
if exc := task.exception():
|
||||
# This exception is an interruption signal, not an error
|
||||
# so we don't want to re-raise it on exit
|
||||
if isinstance(exc, GraphInterrupt):
|
||||
if isinstance(exc, GraphBubbleUp):
|
||||
self.tasks.pop(task)
|
||||
else:
|
||||
self.tasks.pop(task)
|
||||
|
||||
@@ -15,6 +15,7 @@ from langgraph.constants import (
|
||||
TAG_HIDDEN,
|
||||
TASKS,
|
||||
)
|
||||
from langgraph.errors import InvalidUpdateError
|
||||
from langgraph.pregel.log import logger
|
||||
from langgraph.types import Command, PregelExecutableTask, Send
|
||||
|
||||
@@ -68,6 +69,8 @@ def map_command(
|
||||
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.send:
|
||||
if isinstance(cmd.send, (tuple, list)):
|
||||
sends = cmd.send
|
||||
|
||||
@@ -110,13 +110,13 @@ from langgraph.types import (
|
||||
Command,
|
||||
LoopProtocol,
|
||||
PregelExecutableTask,
|
||||
StreamChunk,
|
||||
StreamProtocol,
|
||||
)
|
||||
from langgraph.utils.config import patch_configurable
|
||||
|
||||
V = TypeVar("V")
|
||||
P = ParamSpec("P")
|
||||
StreamChunk = tuple[tuple[str, ...], str, Any]
|
||||
|
||||
INPUT_DONE = object()
|
||||
INPUT_RESUMING = object()
|
||||
|
||||
@@ -18,7 +18,7 @@ from langchain_core.outputs import ChatGenerationChunk, LLMResult
|
||||
from langchain_core.tracers._streaming import T, _StreamingCallbackHandler
|
||||
|
||||
from langgraph.constants import NS_SEP, TAG_HIDDEN, TAG_NOSTREAM
|
||||
from langgraph.pregel.loop import StreamChunk
|
||||
from langgraph.types import StreamChunk
|
||||
|
||||
Meta = tuple[tuple[str, ...], dict[str, Any]]
|
||||
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import replace
|
||||
from functools import partial
|
||||
from typing import Any, Callable, Optional, Sequence
|
||||
|
||||
@@ -10,12 +12,14 @@ from langgraph.constants import (
|
||||
CONFIG_KEY_CHECKPOINT_NS,
|
||||
CONFIG_KEY_RESUMING,
|
||||
CONFIG_KEY_SEND,
|
||||
NS_SEP,
|
||||
)
|
||||
from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphInterrupt
|
||||
from langgraph.types import PregelExecutableTask, RetryPolicy
|
||||
from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphBubbleUp, ParentCommand
|
||||
from langgraph.types import Command, PregelExecutableTask, RetryPolicy
|
||||
from langgraph.utils.config import patch_configurable
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11)
|
||||
|
||||
|
||||
def run_with_retry(
|
||||
@@ -40,10 +44,26 @@ def run_with_retry(
|
||||
task.proc.invoke(task.input, config)
|
||||
# if successful, end
|
||||
break
|
||||
except GraphInterrupt:
|
||||
except ParentCommand as exc:
|
||||
ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
|
||||
cmd = exc.args[0]
|
||||
if cmd.graph == ns:
|
||||
# this command is for the current graph, handle it
|
||||
for w in task.writers:
|
||||
w.invoke(cmd, config)
|
||||
break
|
||||
elif cmd.graph == Command.PARENT:
|
||||
# this command is for the parent graph, assign it to the parent
|
||||
parent_ns = NS_SEP.join(ns.split(NS_SEP)[:-1])
|
||||
exc.args = (replace(cmd, graph=parent_ns),)
|
||||
# bubble up
|
||||
raise
|
||||
except GraphBubbleUp:
|
||||
# if interrupted, end
|
||||
raise
|
||||
except Exception as exc:
|
||||
if SUPPORTS_EXC_NOTES:
|
||||
exc.add_note(f"During task with name '{task.name}' and id '{task.id}'")
|
||||
if retry_policy is None:
|
||||
raise
|
||||
# increment attempts
|
||||
@@ -118,10 +138,26 @@ async def arun_with_retry(
|
||||
await task.proc.ainvoke(task.input, config)
|
||||
# if successful, end
|
||||
break
|
||||
except GraphInterrupt:
|
||||
except ParentCommand as exc:
|
||||
ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
|
||||
cmd = exc.args[0]
|
||||
if cmd.graph == ns:
|
||||
# this command is for the current graph, handle it
|
||||
for w in task.writers:
|
||||
w.invoke(cmd, config)
|
||||
break
|
||||
elif cmd.graph == Command.PARENT:
|
||||
# this command is for the parent graph, assign it to the parent
|
||||
parent_ns = NS_SEP.join(ns.split(NS_SEP)[:-1])
|
||||
exc.args = (replace(cmd, graph=parent_ns),)
|
||||
# bubble up
|
||||
raise
|
||||
except GraphBubbleUp:
|
||||
# if interrupted, end
|
||||
raise
|
||||
except Exception as exc:
|
||||
if SUPPORTS_EXC_NOTES:
|
||||
exc.add_note(f"During task with name '{task.name}' and id '{task.id}'")
|
||||
if retry_policy is None:
|
||||
raise
|
||||
# increment attempts
|
||||
|
||||
@@ -23,7 +23,7 @@ from langgraph.constants import (
|
||||
PUSH,
|
||||
TAG_HIDDEN,
|
||||
)
|
||||
from langgraph.errors import GraphDelegate, GraphInterrupt
|
||||
from langgraph.errors import GraphBubbleUp, GraphInterrupt
|
||||
from langgraph.pregel.executor import Submit
|
||||
from langgraph.pregel.retry import arun_with_retry, run_with_retry
|
||||
from langgraph.types import PregelExecutableTask, RetryPolicy
|
||||
@@ -298,7 +298,7 @@ class PregelRunner:
|
||||
# save interrupt to checkpointer
|
||||
if interrupts := [(INTERRUPT, i) for i in exception.args[0]]:
|
||||
self.put_writes(task.id, interrupts)
|
||||
elif isinstance(exception, GraphDelegate):
|
||||
elif isinstance(exception, GraphBubbleUp):
|
||||
raise exception
|
||||
else:
|
||||
# save error to checkpointer
|
||||
@@ -324,7 +324,7 @@ def _should_stop_others(
|
||||
if fut.cancelled():
|
||||
return True
|
||||
if exc := fut.exception():
|
||||
return not isinstance(exc, GraphInterrupt)
|
||||
return not isinstance(exc, GraphBubbleUp)
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
@@ -48,7 +48,7 @@ def find_subgraph_pregel(candidate: Runnable) -> Optional[Runnable]:
|
||||
nl.__self__ if hasattr(nl, "__self__") else nl
|
||||
for nl in get_function_nonlocals(c.func)
|
||||
)
|
||||
if c.afunc is not None:
|
||||
elif c.afunc is not None:
|
||||
candidates.extend(
|
||||
nl.__self__ if hasattr(nl, "__self__") else nl
|
||||
for nl in get_function_nonlocals(c.afunc)
|
||||
|
||||
@@ -5,6 +5,7 @@ from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Callable,
|
||||
ClassVar,
|
||||
Generic,
|
||||
Hashable,
|
||||
Literal,
|
||||
@@ -140,6 +141,7 @@ class PregelExecutableTask(NamedTuple):
|
||||
id: str
|
||||
path: tuple[Union[str, int, tuple], ...]
|
||||
scheduled: bool = False
|
||||
writers: Sequence[Runnable] = ()
|
||||
|
||||
|
||||
class StateSnapshot(NamedTuple):
|
||||
@@ -239,6 +241,7 @@ N = TypeVar("N", bound=Hashable)
|
||||
class Command(Generic[N]):
|
||||
"""One or more commands to update the graph's state and send messages to nodes."""
|
||||
|
||||
graph: Optional[str] = None
|
||||
update: Optional[dict[str, Any]] = None
|
||||
send: Union[Send, Sequence[Send]] = ()
|
||||
resume: Optional[Union[Any, dict[str, Any]]] = None
|
||||
@@ -252,6 +255,8 @@ class Command(Generic[N]):
|
||||
)
|
||||
return f"Command({contents})"
|
||||
|
||||
PARENT: ClassVar[Literal["__parent__"]] = "__parent__"
|
||||
|
||||
|
||||
StreamChunk = tuple[tuple[str, ...], str, Any]
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import asyncio
|
||||
import sys
|
||||
from collections import ChainMap
|
||||
from typing import Any, Optional, Sequence
|
||||
from typing import Any, Optional, Sequence, cast
|
||||
|
||||
from langchain_core.callbacks import (
|
||||
AsyncCallbackManager,
|
||||
@@ -280,7 +280,10 @@ def ensure_config(*configs: Optional[RunnableConfig]) -> RunnableConfig:
|
||||
continue
|
||||
for k, v in config.items():
|
||||
if v is not None and k in CONFIG_KEYS:
|
||||
empty[k] = v # type: ignore[literal-required]
|
||||
if k == CONF:
|
||||
empty[k] = cast(dict, v).copy()
|
||||
else:
|
||||
empty[k] = v # type: ignore[literal-required]
|
||||
for k, v in config.items():
|
||||
if v is not None and k not in CONFIG_KEYS:
|
||||
empty[CONF][k] = v
|
||||
|
||||
Generated
+13
-13
@@ -1,4 +1,4 @@
|
||||
# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
|
||||
# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
|
||||
|
||||
[[package]]
|
||||
name = "aiosqlite"
|
||||
@@ -1253,13 +1253,13 @@ test = ["jupyter-server (>=2.0.0)", "pytest (>=7.0)", "pytest-jupyter[server] (>
|
||||
|
||||
[[package]]
|
||||
name = "jupyterlab"
|
||||
version = "4.3.0"
|
||||
version = "4.2.5"
|
||||
description = "JupyterLab computational environment"
|
||||
optional = false
|
||||
python-versions = ">=3.8"
|
||||
files = [
|
||||
{file = "jupyterlab-4.3.0-py3-none-any.whl", hash = "sha256:f67e1095ad61ae04349024f0b40345062ab108a0c6998d9810fec6a3c1a70cd5"},
|
||||
{file = "jupyterlab-4.3.0.tar.gz", hash = "sha256:7c6835cbf8df0af0ec8a39332e85ff11693fb9a468205343b4fc0bfbc74817e5"},
|
||||
{file = "jupyterlab-4.2.5-py3-none-any.whl", hash = "sha256:73b6e0775d41a9fee7ee756c80f58a6bed4040869ccc21411dc559818874d321"},
|
||||
{file = "jupyterlab-4.2.5.tar.gz", hash = "sha256:ae7f3a1b8cb88b4f55009ce79fa7c06f99d70cd63601ee4aa91815d054f46f75"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
@@ -1280,9 +1280,9 @@ tornado = ">=6.2.0"
|
||||
traitlets = "*"
|
||||
|
||||
[package.extras]
|
||||
dev = ["build", "bump2version", "coverage", "hatch", "pre-commit", "pytest-cov", "ruff (==0.6.9)"]
|
||||
docs = ["jsx-lexer", "myst-parser", "pydata-sphinx-theme (>=0.13.0)", "pytest", "pytest-check-links", "pytest-jupyter", "sphinx (>=1.8,<8.1.0)", "sphinx-copybutton"]
|
||||
docs-screenshots = ["altair (==5.4.1)", "ipython (==8.16.1)", "ipywidgets (==8.1.5)", "jupyterlab-geojson (==3.4.0)", "jupyterlab-language-pack-zh-cn (==4.2.post3)", "matplotlib (==3.9.2)", "nbconvert (>=7.0.0)", "pandas (==2.2.3)", "scipy (==1.14.1)", "vega-datasets (==0.9.0)"]
|
||||
dev = ["build", "bump2version", "coverage", "hatch", "pre-commit", "pytest-cov", "ruff (==0.3.5)"]
|
||||
docs = ["jsx-lexer", "myst-parser", "pydata-sphinx-theme (>=0.13.0)", "pytest", "pytest-check-links", "pytest-jupyter", "sphinx (>=1.8,<7.3.0)", "sphinx-copybutton"]
|
||||
docs-screenshots = ["altair (==5.3.0)", "ipython (==8.16.1)", "ipywidgets (==8.1.2)", "jupyterlab-geojson (==3.4.0)", "jupyterlab-language-pack-zh-cn (==4.1.post2)", "matplotlib (==3.8.3)", "nbconvert (>=7.0.0)", "pandas (==2.2.1)", "scipy (==1.12.0)", "vega-datasets (==0.9.0)"]
|
||||
test = ["coverage", "pytest (>=7.0)", "pytest-check-links (>=0.7)", "pytest-console-scripts", "pytest-cov", "pytest-jupyter (>=0.5.3)", "pytest-timeout", "pytest-tornasync", "requests", "requests-cache", "virtualenv"]
|
||||
upgrade-extension = ["copier (>=9,<10)", "jinja2-time (<0.3)", "pydantic (<3.0)", "pyyaml-include (<3.0)", "tomli-w (<2.0)"]
|
||||
|
||||
@@ -1792,26 +1792,26 @@ files = [
|
||||
|
||||
[[package]]
|
||||
name = "notebook"
|
||||
version = "7.0.7"
|
||||
version = "7.2.2"
|
||||
description = "Jupyter Notebook - A web-based notebook environment for interactive computing"
|
||||
optional = false
|
||||
python-versions = ">=3.8"
|
||||
files = [
|
||||
{file = "notebook-7.0.7-py3-none-any.whl", hash = "sha256:289b606d7e173f75a18beb1406ef411b43f97f7a9c55ba03efa3622905a62346"},
|
||||
{file = "notebook-7.0.7.tar.gz", hash = "sha256:3bcff00c17b3ac142ef5f436d50637d936b274cfa0b41f6ac0175363de9b4e09"},
|
||||
{file = "notebook-7.2.2-py3-none-any.whl", hash = "sha256:c89264081f671bc02eec0ed470a627ed791b9156cad9285226b31611d3e9fe1c"},
|
||||
{file = "notebook-7.2.2.tar.gz", hash = "sha256:2ef07d4220421623ad3fe88118d687bc0450055570cdd160814a59cf3a1c516e"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
jupyter-server = ">=2.4.0,<3"
|
||||
jupyterlab = ">=4.0.2,<5"
|
||||
jupyterlab-server = ">=2.22.1,<3"
|
||||
jupyterlab = ">=4.2.0,<4.3"
|
||||
jupyterlab-server = ">=2.27.1,<3"
|
||||
notebook-shim = ">=0.2,<0.3"
|
||||
tornado = ">=6.2.0"
|
||||
|
||||
[package.extras]
|
||||
dev = ["hatch", "pre-commit"]
|
||||
docs = ["myst-parser", "nbsphinx", "pydata-sphinx-theme", "sphinx (>=1.3.6)", "sphinxcontrib-github-alt", "sphinxcontrib-spelling"]
|
||||
test = ["importlib-resources (>=5.0)", "ipykernel", "jupyter-server[test] (>=2.4.0,<3)", "jupyterlab-server[test] (>=2.22.1,<3)", "nbval", "pytest (>=7.0)", "pytest-console-scripts", "pytest-timeout", "pytest-tornasync", "requests"]
|
||||
test = ["importlib-resources (>=5.0)", "ipykernel", "jupyter-server[test] (>=2.4.0,<3)", "jupyterlab-server[test] (>=2.27.1,<3)", "nbval", "pytest (>=7.0)", "pytest-console-scripts", "pytest-timeout", "pytest-tornasync", "requests"]
|
||||
|
||||
[[package]]
|
||||
name = "notebook-shim"
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "langgraph"
|
||||
version = "0.2.50"
|
||||
version = "0.2.53"
|
||||
description = "Building stateful, multi-actor applications with LLMs"
|
||||
authors = []
|
||||
license = "MIT"
|
||||
|
||||
@@ -272,6 +272,54 @@ async def _store_postgres_aio():
|
||||
await conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _store_postgres_aio_pipe():
|
||||
if sys.version_info < (3, 10):
|
||||
pytest.skip("Async Postgres tests require Python 3.10+")
|
||||
database = f"test_{uuid4().hex[:16]}"
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
async with AsyncPostgresStore.from_conn_string(
|
||||
DEFAULT_POSTGRES_URI + database
|
||||
) as store:
|
||||
await store.setup() # Run in its own transaction
|
||||
async with AsyncPostgresStore.from_conn_string(
|
||||
DEFAULT_POSTGRES_URI + database, pipeline=True
|
||||
) as store:
|
||||
yield store
|
||||
finally:
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _store_postgres_aio_pool():
|
||||
if sys.version_info < (3, 10):
|
||||
pytest.skip("Async Postgres tests require Python 3.10+")
|
||||
database = f"test_{uuid4().hex[:16]}"
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
async with AsyncPostgresStore.from_conn_string(
|
||||
DEFAULT_POSTGRES_URI + database,
|
||||
pool_config={"max_size": 10},
|
||||
) as store:
|
||||
await store.setup()
|
||||
yield store
|
||||
finally:
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _store_duckdb_aio():
|
||||
async with AsyncDuckDBStore.from_conn_string(":memory:") as store:
|
||||
@@ -296,6 +344,45 @@ def store_postgres():
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def store_postgres_pipe():
|
||||
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 store
|
||||
with PostgresStore.from_conn_string(DEFAULT_POSTGRES_URI + database) as store:
|
||||
store.setup() # Run in its own transaction
|
||||
with PostgresStore.from_conn_string(
|
||||
DEFAULT_POSTGRES_URI + database, pipeline=True
|
||||
) as store:
|
||||
yield store
|
||||
finally:
|
||||
# drop unique db
|
||||
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def store_postgres_pool():
|
||||
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 store
|
||||
with PostgresStore.from_conn_string(
|
||||
DEFAULT_POSTGRES_URI + database, pool_config={"max_size": 10}
|
||||
) as store:
|
||||
store.setup()
|
||||
yield store
|
||||
finally:
|
||||
# drop unique db
|
||||
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def store_duckdb():
|
||||
with DuckDBStore.from_conn_string(":memory:") as store:
|
||||
@@ -317,6 +404,12 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
|
||||
elif store_name == "postgres_aio":
|
||||
async with _store_postgres_aio() as store:
|
||||
yield store
|
||||
elif store_name == "postgres_aio_pipe":
|
||||
async with _store_postgres_aio_pipe() as store:
|
||||
yield store
|
||||
elif store_name == "postgres_aio_pool":
|
||||
async with _store_postgres_aio_pool() as store:
|
||||
yield store
|
||||
elif store_name == "duckdb_aio":
|
||||
async with _store_duckdb_aio() as store:
|
||||
yield store
|
||||
@@ -342,5 +435,17 @@ ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [
|
||||
*ALL_CHECKPOINTERS_ASYNC,
|
||||
None,
|
||||
]
|
||||
ALL_STORES_SYNC = ["in_memory", "postgres", "duckdb"]
|
||||
ALL_STORES_ASYNC = ["in_memory", "postgres_aio", "duckdb_aio"]
|
||||
ALL_STORES_SYNC = [
|
||||
"in_memory",
|
||||
"postgres",
|
||||
"postgres_pipe",
|
||||
"postgres_pool",
|
||||
"duckdb",
|
||||
]
|
||||
ALL_STORES_ASYNC = [
|
||||
"in_memory",
|
||||
"postgres_aio",
|
||||
"postgres_aio_pipe",
|
||||
"postgres_aio_pool",
|
||||
"duckdb_aio",
|
||||
]
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
import re
|
||||
from typing import Any, Iterator, List, Optional, cast
|
||||
from typing import Any, AsyncIterator, Iterator, List, Optional, cast
|
||||
|
||||
from langchain_core.callbacks import CallbackManagerForLLMRun
|
||||
from langchain_core.callbacks import (
|
||||
AsyncCallbackManagerForLLMRun,
|
||||
CallbackManagerForLLMRun,
|
||||
)
|
||||
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
|
||||
from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage
|
||||
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
|
||||
@@ -84,3 +87,51 @@ class FakeChatModel(GenericFakeChatModel):
|
||||
if run_manager:
|
||||
run_manager.on_llm_new_token("", chunk=chunk)
|
||||
yield chunk
|
||||
|
||||
async def _astream(
|
||||
self,
|
||||
messages: List[BaseMessage],
|
||||
stop: Optional[List[str]] = None,
|
||||
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterator[ChatGenerationChunk]:
|
||||
"""Stream the output of the model."""
|
||||
chat_result = self._generate(
|
||||
messages, stop=stop, run_manager=run_manager, **kwargs
|
||||
)
|
||||
if not isinstance(chat_result, ChatResult):
|
||||
raise ValueError(
|
||||
f"Expected generate to return a ChatResult, "
|
||||
f"but got {type(chat_result)} instead."
|
||||
)
|
||||
|
||||
message = chat_result.generations[0].message
|
||||
|
||||
if not isinstance(message, AIMessage):
|
||||
raise ValueError(
|
||||
f"Expected invoke to return an AIMessage, "
|
||||
f"but got {type(message)} instead."
|
||||
)
|
||||
|
||||
content = message.content
|
||||
|
||||
if content:
|
||||
# Use a regular expression to split on whitespace with a capture group
|
||||
# so that we can preserve the whitespace in the output.
|
||||
assert isinstance(content, str)
|
||||
content_chunks = cast(list[str], re.split(r"(\s)", content))
|
||||
|
||||
for token in content_chunks:
|
||||
chunk = ChatGenerationChunk(
|
||||
message=AIMessageChunk(content=token, id=message.id)
|
||||
)
|
||||
if run_manager:
|
||||
run_manager.on_llm_new_token(token, chunk=chunk)
|
||||
yield chunk
|
||||
else:
|
||||
args = message.__dict__
|
||||
args.pop("type")
|
||||
chunk = ChatGenerationChunk(message=AIMessageChunk(**args))
|
||||
if run_manager:
|
||||
await run_manager.on_llm_new_token("", chunk=chunk)
|
||||
yield chunk
|
||||
|
||||
@@ -102,6 +102,9 @@ class FakeToolCallingModel(BaseChatModel):
|
||||
tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]],
|
||||
**kwargs: Any,
|
||||
) -> Runnable[LanguageModelInput, BaseMessage]:
|
||||
if len(tools) == 0:
|
||||
raise ValueError("Must provide at least one tool")
|
||||
|
||||
tool_dicts = []
|
||||
for tool in tools:
|
||||
if not isinstance(tool, BaseTool):
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import enum
|
||||
import json
|
||||
import logging
|
||||
import operator
|
||||
import re
|
||||
import time
|
||||
@@ -67,16 +68,9 @@ from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInt
|
||||
from langgraph.graph import END, Graph, GraphCommand, StateGraph
|
||||
from langgraph.graph.message import MessageGraph, MessagesState, add_messages
|
||||
from langgraph.managed.shared_value import SharedValue
|
||||
from langgraph.prebuilt.chat_agent_executor import (
|
||||
create_tool_calling_executor,
|
||||
)
|
||||
from langgraph.prebuilt.chat_agent_executor import create_tool_calling_executor
|
||||
from langgraph.prebuilt.tool_node import ToolNode
|
||||
from langgraph.pregel import (
|
||||
Channel,
|
||||
GraphRecursionError,
|
||||
Pregel,
|
||||
StateSnapshot,
|
||||
)
|
||||
from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
|
||||
from langgraph.pregel.retry import RetryPolicy
|
||||
from langgraph.store.base import BaseStore
|
||||
from langgraph.store.memory import InMemoryStore
|
||||
@@ -104,6 +98,8 @@ from tests.messages import (
|
||||
_AnyIdToolMessage,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# define these objects to avoid importing langchain_core.agents
|
||||
# and therefore avoid relying on core Pydantic version
|
||||
@@ -6628,11 +6624,7 @@ def test_message_graph(
|
||||
from langchain_core.language_models.fake_chat_models import (
|
||||
FakeMessagesListChatModel,
|
||||
)
|
||||
from langchain_core.messages import (
|
||||
AIMessage,
|
||||
BaseMessage,
|
||||
HumanMessage,
|
||||
)
|
||||
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
|
||||
from langchain_core.outputs import ChatGeneration, ChatResult
|
||||
from langchain_core.tools import tool
|
||||
|
||||
@@ -13937,50 +13929,75 @@ def test_store_injected(
|
||||
|
||||
doc_id = str(uuid.uuid4())
|
||||
doc = {"some-key": "this-is-a-val"}
|
||||
|
||||
def node(input: State, config: RunnableConfig, store: BaseStore):
|
||||
assert isinstance(store, BaseStore)
|
||||
store.put(
|
||||
("foo", "bar"),
|
||||
doc_id,
|
||||
{
|
||||
**doc,
|
||||
"from_thread": config["configurable"]["thread_id"],
|
||||
"some_val": input["count"],
|
||||
},
|
||||
)
|
||||
return {"count": 1}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("node", node)
|
||||
builder.add_edge("__start__", "node")
|
||||
graph = builder.compile(store=the_store, checkpointer=checkpointer)
|
||||
|
||||
uid = uuid.uuid4().hex
|
||||
namespace = (f"foo-{uid}", "bar")
|
||||
thread_1 = str(uuid.uuid4())
|
||||
result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_1}})
|
||||
assert result == {"count": 1}
|
||||
returned_doc = the_store.get(("foo", "bar"), doc_id).value
|
||||
assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0}
|
||||
assert len(the_store.search(("foo", "bar"))) == 1
|
||||
|
||||
# Check update on existing thread
|
||||
result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_1}})
|
||||
assert result == {"count": 2}
|
||||
returned_doc = the_store.get(("foo", "bar"), doc_id).value
|
||||
assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 1}
|
||||
assert len(the_store.search(("foo", "bar"))) == 1
|
||||
|
||||
thread_2 = str(uuid.uuid4())
|
||||
|
||||
class Node:
|
||||
def __init__(self, i: Optional[int] = None):
|
||||
self.i = i
|
||||
|
||||
def __call__(self, inputs: State, config: RunnableConfig, store: BaseStore):
|
||||
assert isinstance(store, BaseStore)
|
||||
store.put(
|
||||
namespace
|
||||
if self.i is not None
|
||||
and config["configurable"]["thread_id"] in (thread_1, thread_2)
|
||||
else (f"foo_{self.i}", "bar"),
|
||||
doc_id,
|
||||
{
|
||||
**doc,
|
||||
"from_thread": config["configurable"]["thread_id"],
|
||||
"some_val": inputs["count"],
|
||||
},
|
||||
)
|
||||
return {"count": 1}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("node", Node())
|
||||
builder.add_edge("__start__", "node")
|
||||
N = 500
|
||||
M = 1
|
||||
if "duckdb" in store_name:
|
||||
logger.warning(
|
||||
"DuckDB store implementation has a known issue that does not"
|
||||
" support concurrent writes, so we're reducing the test scope"
|
||||
)
|
||||
N = M = 1
|
||||
|
||||
for i in range(N):
|
||||
builder.add_node(f"node_{i}", Node(i))
|
||||
builder.add_edge("__start__", f"node_{i}")
|
||||
|
||||
graph = builder.compile(store=the_store, checkpointer=checkpointer)
|
||||
|
||||
results = graph.batch(
|
||||
[{"count": 0}] * M,
|
||||
([{"configurable": {"thread_id": str(uuid.uuid4())}}] * (M - 1))
|
||||
+ [{"configurable": {"thread_id": thread_1}}],
|
||||
)
|
||||
result = results[-1]
|
||||
assert result == {"count": N + 1}
|
||||
returned_doc = the_store.get(namespace, doc_id).value
|
||||
assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0}
|
||||
assert len(the_store.search(namespace)) == 1
|
||||
# Check results after another turn of the same thread
|
||||
result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_1}})
|
||||
assert result == {"count": (N + 1) * 2}
|
||||
returned_doc = the_store.get(namespace, doc_id).value
|
||||
assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1}
|
||||
assert len(the_store.search(namespace)) == 1
|
||||
|
||||
result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_2}})
|
||||
assert result == {"count": 1}
|
||||
returned_doc = the_store.get(("foo", "bar"), doc_id).value
|
||||
assert result == {"count": N + 1}
|
||||
returned_doc = the_store.get(namespace, doc_id).value
|
||||
assert returned_doc == {
|
||||
**doc,
|
||||
"from_thread": thread_2,
|
||||
"some_val": 0,
|
||||
} # Overwrites the whole doc
|
||||
assert len(the_store.search(("foo", "bar"))) == 1 # still overwriting the same one
|
||||
assert len(the_store.search(namespace)) == 1 # still overwriting the same one
|
||||
|
||||
|
||||
def test_enum_node_names():
|
||||
@@ -14378,3 +14395,79 @@ def test_runnable_passthrough_node_graph() -> None:
|
||||
graph = graph_builder.compile()
|
||||
|
||||
assert graph.get_graph(xray=True).to_json() == graph.get_graph(xray=False).to_json()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str) -> None:
|
||||
from langchain_core.messages import BaseMessage
|
||||
from langchain_core.tools import tool
|
||||
|
||||
@tool(return_direct=True)
|
||||
def get_user_name() -> GraphCommand:
|
||||
"""Retrieve user name"""
|
||||
return GraphCommand(update={"user_name": "Meow"}, graph=GraphCommand.PARENT)
|
||||
|
||||
subgraph_builder = StateGraph(MessagesState)
|
||||
subgraph_builder.add_node("tool", get_user_name)
|
||||
subgraph_builder.add_edge(START, "tool")
|
||||
subgraph = subgraph_builder.compile()
|
||||
|
||||
class CustomParentState(TypedDict):
|
||||
messages: Annotated[list[BaseMessage], add_messages]
|
||||
# this key is not available to the child graph
|
||||
user_name: str
|
||||
|
||||
builder = StateGraph(CustomParentState)
|
||||
builder.add_node("alice", subgraph)
|
||||
builder.add_edge(START, "alice")
|
||||
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
|
||||
graph = builder.compile(checkpointer=checkpointer)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
assert graph.invoke({"messages": [("user", "get user name")]}, config) == {
|
||||
"messages": [
|
||||
_AnyIdHumanMessage(
|
||||
content="get user name", additional_kwargs={}, response_metadata={}
|
||||
),
|
||||
],
|
||||
"user_name": "Meow",
|
||||
}
|
||||
assert graph.get_state(config) == StateSnapshot(
|
||||
values={
|
||||
"messages": [
|
||||
_AnyIdHumanMessage(
|
||||
content="get user name", additional_kwargs={}, response_metadata={}
|
||||
),
|
||||
],
|
||||
"user_name": "Meow",
|
||||
},
|
||||
next=(),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"source": "loop",
|
||||
"writes": {
|
||||
"alice": {
|
||||
"user_name": "Meow",
|
||||
}
|
||||
},
|
||||
"thread_id": "1",
|
||||
"step": 1,
|
||||
"parents": {},
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
tasks=(),
|
||||
)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import operator
|
||||
import random
|
||||
import re
|
||||
@@ -100,6 +101,8 @@ from tests.messages import (
|
||||
_AnyIdToolMessage,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
@@ -12272,60 +12275,89 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) ->
|
||||
|
||||
doc_id = str(uuid.uuid4())
|
||||
doc = {"some-key": "this-is-a-val"}
|
||||
uid = uuid.uuid4().hex
|
||||
namespace = (f"foo-{uid}", "bar")
|
||||
thread_1 = str(uuid.uuid4())
|
||||
thread_2 = str(uuid.uuid4())
|
||||
|
||||
async def node(input: State, config: RunnableConfig, store: BaseStore):
|
||||
assert isinstance(store, BaseStore)
|
||||
await store.aput(
|
||||
("foo", "bar"),
|
||||
doc_id,
|
||||
{
|
||||
**doc,
|
||||
"from_thread": config["configurable"]["thread_id"],
|
||||
"some_val": input["count"],
|
||||
},
|
||||
)
|
||||
return {"count": 1}
|
||||
class Node:
|
||||
def __init__(self, i: Optional[int] = None):
|
||||
self.i = i
|
||||
|
||||
async def __call__(
|
||||
self, inputs: State, config: RunnableConfig, store: BaseStore
|
||||
):
|
||||
assert isinstance(store, BaseStore)
|
||||
await store.aput(
|
||||
namespace
|
||||
if self.i is not None
|
||||
and config["configurable"]["thread_id"] in (thread_1, thread_2)
|
||||
else (f"foo_{self.i}", "bar"),
|
||||
doc_id,
|
||||
{
|
||||
**doc,
|
||||
"from_thread": config["configurable"]["thread_id"],
|
||||
"some_val": inputs["count"],
|
||||
},
|
||||
)
|
||||
return {"count": 1}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("node", node)
|
||||
builder.add_node("node", Node())
|
||||
builder.add_edge("__start__", "node")
|
||||
|
||||
N = 500
|
||||
M = 1
|
||||
if "duckdb" in store_name:
|
||||
logger.warning(
|
||||
"DuckDB store implementation has a known issue that does not"
|
||||
" support concurrent writes, so we're reducing the test scope"
|
||||
)
|
||||
N = M = 1
|
||||
|
||||
for i in range(N):
|
||||
builder.add_node(f"node_{i}", Node(i))
|
||||
builder.add_edge("__start__", f"node_{i}")
|
||||
|
||||
async with awith_checkpointer(checkpointer_name) as checkpointer, awith_store(
|
||||
store_name
|
||||
) as the_store:
|
||||
graph = builder.compile(store=the_store, checkpointer=checkpointer)
|
||||
|
||||
thread_1 = str(uuid.uuid4())
|
||||
result = await graph.ainvoke(
|
||||
{"count": 0}, {"configurable": {"thread_id": thread_1}}
|
||||
# Test batch operations with multiple threads
|
||||
results = await graph.abatch(
|
||||
[{"count": 0}] * M,
|
||||
([{"configurable": {"thread_id": str(uuid.uuid4())}}] * (M - 1))
|
||||
+ [{"configurable": {"thread_id": thread_1}}],
|
||||
)
|
||||
assert result == {"count": 1}
|
||||
returned_doc = (await the_store.aget(("foo", "bar"), doc_id)).value
|
||||
result = results[-1]
|
||||
assert result == {"count": N + 1}
|
||||
returned_doc = (await the_store.aget(namespace, doc_id)).value
|
||||
assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0}
|
||||
assert len((await the_store.asearch(("foo", "bar")))) == 1
|
||||
assert len((await the_store.asearch(namespace))) == 1
|
||||
|
||||
# Check update on existing thread
|
||||
# Check results after another turn of the same thread
|
||||
result = await graph.ainvoke(
|
||||
{"count": 0}, {"configurable": {"thread_id": thread_1}}
|
||||
)
|
||||
assert result == {"count": 2}
|
||||
returned_doc = (await the_store.aget(("foo", "bar"), doc_id)).value
|
||||
assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 1}
|
||||
assert len((await the_store.asearch(("foo", "bar")))) == 1
|
||||
|
||||
thread_2 = str(uuid.uuid4())
|
||||
assert result == {"count": (N + 1) * 2}
|
||||
returned_doc = (await the_store.aget(namespace, doc_id)).value
|
||||
assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1}
|
||||
assert len((await the_store.asearch(namespace))) == 1
|
||||
|
||||
# Test with a different thread
|
||||
result = await graph.ainvoke(
|
||||
{"count": 0}, {"configurable": {"thread_id": thread_2}}
|
||||
)
|
||||
assert result == {"count": 1}
|
||||
returned_doc = (await the_store.aget(("foo", "bar"), doc_id)).value
|
||||
assert result == {"count": N + 1}
|
||||
returned_doc = (await the_store.aget(namespace, doc_id)).value
|
||||
assert returned_doc == {
|
||||
**doc,
|
||||
"from_thread": thread_2,
|
||||
"some_val": 0,
|
||||
} # Overwrites the whole doc
|
||||
assert (
|
||||
len((await the_store.asearch(("foo", "bar")))) == 1
|
||||
len((await the_store.asearch(namespace))) == 1
|
||||
) # still overwriting the same one
|
||||
|
||||
|
||||
@@ -12565,3 +12597,83 @@ async def test_debug_nested_subgraphs():
|
||||
assert stream_task["interrupts"] == history_task.interrupts
|
||||
assert stream_task.get("error") == history_task.error
|
||||
assert stream_task.get("state") == history_task.state
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
async def test_parent_command(checkpointer_name: str) -> None:
|
||||
from langchain_core.messages import BaseMessage
|
||||
from langchain_core.tools import tool
|
||||
|
||||
@tool(return_direct=True)
|
||||
def get_user_name() -> GraphCommand:
|
||||
"""Retrieve user name"""
|
||||
return GraphCommand(update={"user_name": "Meow"}, graph=GraphCommand.PARENT)
|
||||
|
||||
subgraph_builder = StateGraph(MessagesState)
|
||||
subgraph_builder.add_node("tool", get_user_name)
|
||||
subgraph_builder.add_edge(START, "tool")
|
||||
subgraph = subgraph_builder.compile()
|
||||
|
||||
class CustomParentState(TypedDict):
|
||||
messages: Annotated[list[BaseMessage], add_messages]
|
||||
# this key is not available to the child graph
|
||||
user_name: str
|
||||
|
||||
builder = StateGraph(CustomParentState)
|
||||
builder.add_node("alice", subgraph)
|
||||
builder.add_edge(START, "alice")
|
||||
async with awith_checkpointer(checkpointer_name) as checkpointer:
|
||||
graph = builder.compile(checkpointer=checkpointer)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
assert await graph.ainvoke(
|
||||
{"messages": [("user", "get user name")]}, config
|
||||
) == {
|
||||
"messages": [
|
||||
_AnyIdHumanMessage(
|
||||
content="get user name", additional_kwargs={}, response_metadata={}
|
||||
),
|
||||
],
|
||||
"user_name": "Meow",
|
||||
}
|
||||
assert await graph.aget_state(config) == StateSnapshot(
|
||||
values={
|
||||
"messages": [
|
||||
_AnyIdHumanMessage(
|
||||
content="get user name",
|
||||
additional_kwargs={},
|
||||
response_metadata={},
|
||||
),
|
||||
],
|
||||
"user_name": "Meow",
|
||||
},
|
||||
next=(),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"source": "loop",
|
||||
"writes": {
|
||||
"alice": {
|
||||
"user_name": "Meow",
|
||||
}
|
||||
},
|
||||
"thread_id": "1",
|
||||
"step": 1,
|
||||
"parents": {},
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
tasks=(),
|
||||
)
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user