mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-05 17:27:47 +02:00
Compare commits
10
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1b9b0a686e | ||
|
|
9e02a23682 | ||
|
|
5e70e6f307 | ||
|
|
eb57c06896 | ||
|
|
9284b57ba0 | ||
|
|
25fea591b5 | ||
|
|
b9fe53777f | ||
|
|
35c2e8a679 | ||
|
|
e0fb56c6a3 | ||
|
|
3458a3cecb |
@@ -1,60 +1,44 @@
|
||||
name: "\U0001F41B Bug Report"
|
||||
description: Report a bug in LangGraph. To report a security issue, please instead use the security option (below). For questions, please use the LangChain forum (below).
|
||||
labels: ["bug"]
|
||||
type: bug
|
||||
description: Report a bug in LangGraph. To report a security issue, please instead use the security option below. For questions, please use the GitHub Discussions.
|
||||
labels: ["02 Bug Report"]
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
Thank you for taking the time to file a bug report.
|
||||
|
||||
For usage questions, feature requests and general design questions, please use the [LangChain Forum](https://forum.langchain.com/).
|
||||
|
||||
Check these before submitting to see if your issue has already been reported, fixed or if there's another way to solve your problem:
|
||||
|
||||
* [Documentation](https://docs.langchain.com/oss/python/langgraph/overview),
|
||||
* [API Reference Documentation](https://reference.langchain.com/python/),
|
||||
* [LangChain ChatBot](https://chat.langchain.com/)
|
||||
* [GitHub search](https://github.com/langchain-ai/langgraph),
|
||||
* [LangChain Forum](https://forum.langchain.com/),
|
||||
value: >
|
||||
Thank you for taking the time to file a bug report.
|
||||
|
||||
Use this to report BUGS in LangGraph. For usage questions, feature requests and general design questions, please use [GitHub Discussions](https://github.com/langchain-ai/langgraph/discussions).
|
||||
|
||||
Relevant links to check before filing a bug report to see if your issue has already been reported, fixed or
|
||||
if there's another way to solve your problem:
|
||||
|
||||
[LangGraph Github Discussions](https://github.com/langchain-ai/langgraph/discussions),
|
||||
[LangGraph Github Issues](https://github.com/langchain-ai/langgraph/issues),
|
||||
[LangGraph how-to guides](https://langchain-ai.github.io/langgraph/how-tos/).
|
||||
[LangChain documentation with the integrated search](https://python.langchain.com/docs/get_started/introduction),
|
||||
[GitHub search](https://github.com/langchain-ai/langgraph),
|
||||
- type: checkboxes
|
||||
id: checks
|
||||
attributes:
|
||||
label: Checked other resources
|
||||
description: Please confirm and check all the following options.
|
||||
description: Before submitting this issue, please confirm that you have completed all the steps below by checking each option. These steps help ensure your issue is well-defined, relevant, and actionable.
|
||||
options:
|
||||
- label: This is a bug, not a usage question.
|
||||
- label: This is a bug, not a usage question. For questions, please use GitHub Discussions.
|
||||
required: true
|
||||
- label: I added a clear and descriptive title that summarizes this issue.
|
||||
- label: I added a clear and detailed title that summarizes the issue.
|
||||
required: true
|
||||
- label: I used the GitHub search to find a similar question and didn't find it.
|
||||
- label: I read what a minimal reproducible example is (https://stackoverflow.com/help/minimal-reproducible-example).
|
||||
required: true
|
||||
- label: I am sure that this is a bug in LangGraph rather than my code.
|
||||
required: true
|
||||
- label: The bug is not resolved by updating to the latest stable version of LangGraph (or the specific integration package).
|
||||
required: true
|
||||
- label: This is not related to the langchain-community package.
|
||||
required: true
|
||||
- label: I posted a self-contained, minimal, reproducible example. A maintainer can copy it and run it AS IS.
|
||||
- label: I included a self-contained, minimal example that demonstrates the issue INCLUDING all the relevant imports. The code run AS IS to reproduce the issue.
|
||||
required: true
|
||||
- type: textarea
|
||||
id: reproduction
|
||||
validations:
|
||||
required: true
|
||||
attributes:
|
||||
label: Reproduction Steps / Example Code (Python)
|
||||
label: Example Code
|
||||
description: |
|
||||
Please add a self-contained, [minimal, reproducible, example](https://stackoverflow.com/help/minimal-reproducible-example) with your use case.
|
||||
|
||||
If a maintainer can copy it, run it, and see it right away, there's a much higher chance that you'll be able to get help.
|
||||
|
||||
**Important!**
|
||||
|
||||
* Avoid screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.
|
||||
* Reduce your code to the minimum required to reproduce the issue if possible.
|
||||
|
||||
(This will be automatically formatted into code, so no need for backticks.)
|
||||
render: python
|
||||
placeholder: |
|
||||
from langgraph.graph import StateGraph
|
||||
|
||||
@@ -63,13 +47,17 @@ body:
|
||||
|
||||
chain = StateGraph(list)
|
||||
chain.invoke('Hello!')
|
||||
render: python
|
||||
- type: textarea
|
||||
id: error
|
||||
validations:
|
||||
required: false
|
||||
attributes:
|
||||
label: Error Message and Stack Trace (if applicable)
|
||||
description: |
|
||||
If you are reporting an error, please copy and paste the full error message and
|
||||
stack trace.
|
||||
(This will be automatically formatted into code, so no need for backticks.)
|
||||
If you are reporting an error, please include the full error message and stack trace.
|
||||
placeholder: |
|
||||
Exception + full stack trace
|
||||
render: shell
|
||||
- type: textarea
|
||||
id: description
|
||||
@@ -90,18 +78,7 @@ body:
|
||||
attributes:
|
||||
label: System Info
|
||||
description: |
|
||||
Please share your system info with us.
|
||||
|
||||
Run the following command in your terminal and paste the output here:
|
||||
|
||||
`python -m langchain_core.sys_info`
|
||||
|
||||
or if you have an existing python interpreter running:
|
||||
|
||||
```python
|
||||
from langchain_core import sys_info
|
||||
sys_info.print_sys_info()
|
||||
```
|
||||
python -m langchain_core.sys_info
|
||||
placeholder: |
|
||||
python -m langchain_core.sys_info
|
||||
validations:
|
||||
|
||||
@@ -1,15 +1,15 @@
|
||||
blank_issues_enabled: false
|
||||
version: 2.1
|
||||
contact_links:
|
||||
- name: 💬 LangChain Forum
|
||||
url: https://forum.langchain.com/
|
||||
about: General community discussions and support
|
||||
- name: 📚 LangGraph Documentation
|
||||
url: https://docs.langchain.com/oss/python/langgraph/overview
|
||||
about: View the official LangGraph documentation
|
||||
- name: 📚 API Reference Documentation
|
||||
url: https://reference.langchain.com/python/
|
||||
about: View the official LangGraph API reference documentation
|
||||
- name: 📚 Documentation issue
|
||||
url: https://github.com/langchain-ai/docs/issues/new?template=02-langgraph.yml
|
||||
about: Report an issue related to the LangGraph documentation
|
||||
- name: 🤔 Question or Problem
|
||||
about: Ask a question or ask about a problem in GitHub Discussions.
|
||||
url: https://github.com/langchain-ai/langgraph/discussions/categories/q-a
|
||||
- name: Feature Request
|
||||
url: https://github.com/langchain-ai/langgraph/discussions/categories/ideas
|
||||
about: Suggest a feature or an idea
|
||||
- name: Show and tell
|
||||
about: Show what you built with LangChain
|
||||
url: https://github.com/langchain-ai/langgraph/discussions/categories/show-and-tell
|
||||
- name: Slack
|
||||
url: https://www.langchain.com/join-community
|
||||
about: General community discussions
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
name: Documentation
|
||||
description: Report an issue related to the LangGraph documentation.
|
||||
title: "DOC: <Please write a comprehensive title after the 'DOC: ' prefix>"
|
||||
labels: [03 - Documentation]
|
||||
|
||||
body:
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: "Issue with current documentation:"
|
||||
description: >
|
||||
Please make sure to leave a reference to the document/code you're
|
||||
referring to.
|
||||
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: "Idea or request for content:"
|
||||
description: >
|
||||
Please describe as clearly as possible what topics you think are missing
|
||||
from the current documentation.
|
||||
@@ -1,29 +1,25 @@
|
||||
name: 🔒 Privileged
|
||||
description: You are a LangGraph maintainer, or was asked directly by a maintainer to create an issue here. If not, check the other options.
|
||||
description: You are a LangChain maintainer, or was asked directly by a maintainer to create an issue here. If not, check the other options.
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
Thanks for your interest in LangGraph! 🚀
|
||||
|
||||
If you are not a LangGraph maintainer or were not asked directly by a maintainer to create an issue, then please start the conversation on the [LangChain Forum](https://forum.langchain.com/) instead.
|
||||
|
||||
You are a LangGraph maintainer if you maintain any of the packages inside of the LangGraph repository
|
||||
or are a regular contributor to LangGraph with previous merged merged pull requests.
|
||||
Thanks for your interest in LangChain! 🚀
|
||||
|
||||
If you are not a LangChain maintainer or were not asked directly by a maintainer to create an issue, then please start the conversation in a [Question in GitHub Discussions](https://github.com/langchain-ai/langchain/discussions/categories/q-a) instead.
|
||||
|
||||
You are a LangChain maintainer if you maintain any of the packages inside of the LangChain repository
|
||||
or are a regular contributor to LangChain with previous merged merged pull requests.
|
||||
- type: checkboxes
|
||||
id: privileged
|
||||
attributes:
|
||||
label: Privileged issue
|
||||
description: Confirm that you are allowed to create an issue here.
|
||||
options:
|
||||
- label: I am a LangGraph maintainer, or was asked directly by a LangGraph maintainer to create an issue here.
|
||||
- label: I am a LangChain maintainer, or was asked directly by a LangChain maintainer to create an issue here.
|
||||
required: true
|
||||
- type: textarea
|
||||
id: content
|
||||
attributes:
|
||||
label: Issue Content
|
||||
description: Add the content of the issue here.
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
Community members should **NOT** work on Privileged issues unless these issues have been explicitly marked with a "help-wanted" tag.
|
||||
|
||||
@@ -1,31 +0,0 @@
|
||||
Thank you for contributing to LangGraph! Follow these steps to mark your pull request as ready for review. **If any of these steps are not completed, your PR will not be considered for review.**
|
||||
|
||||
- [ ] **PR title**: Follows the format: {TYPE}({SCOPE}): {DESCRIPTION}
|
||||
- Examples:
|
||||
- feat(core): add multi-tenant support
|
||||
- fix(cli): resolve flag parsing error
|
||||
- docs(openai): update API usage examples
|
||||
- Allowed `{TYPE}` values:
|
||||
- feat, fix, docs, style, refactor, perf, test, build, ci, chore, revert, release
|
||||
- Allowed `{SCOPE}` values (optional):
|
||||
- langgraph, docs, cli, checkpoint, checkpoint-postgres, checkpoint-sqlite, prebuilt, scheduler-kafka, sdk-py
|
||||
- Once you've written the title, please delete this checklist item; do not include it in the PR.
|
||||
|
||||
- [ ] **PR message**: ***Delete this entire checklist*** and replace with
|
||||
- **Description:** a description of the change. Include a [closing keyword](https://docs.github.com/en/issues/tracking-your-work-with-issues/using-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword) if applicable.
|
||||
- **Issue:** the issue # it fixes, if applicable
|
||||
- **Dependencies:** any dependencies required for this change
|
||||
- **Twitter handle:** if your PR gets announced, and you'd like a mention, we'll gladly shout you out!
|
||||
|
||||
- [ ] **Add tests and docs**: If you're adding a new integration, you must include:
|
||||
1. A test for the integration, preferably unit tests that do not rely on network access,
|
||||
2. An example notebook showing its use. It lives in `docs/docs/integrations` directory.
|
||||
|
||||
- [ ] **Lint and test**: Run `make format`, `make lint` and `make test` from the root of the package(s) you've modified. We will not consider a PR unless these three are passing in CI. See [contribution guidelines](https://docs.langchain.com/oss/python/contributing/overview) for more.
|
||||
|
||||
Additional guidelines:
|
||||
|
||||
- Make sure optional dependencies are imported within a function.
|
||||
- Please do not add dependencies to `pyproject.toml` files (even optional ones) unless they are **required** for unit tests.
|
||||
- Most PRs should not touch more than one package.
|
||||
- Changes should be backwards compatible.
|
||||
@@ -0,0 +1,88 @@
|
||||
# An action for setting up poetry install with caching.
|
||||
# Using a custom action since the default action does not
|
||||
# take poetry install groups into account.
|
||||
# Action code from:
|
||||
# https://github.com/actions/setup-python/issues/505#issuecomment-1273013236
|
||||
name: poetry-install-with-caching
|
||||
description: Poetry install with support for caching of dependency groups.
|
||||
|
||||
inputs:
|
||||
python-version:
|
||||
description: Python version, supporting MAJOR.MINOR only
|
||||
required: true
|
||||
|
||||
poetry-version:
|
||||
description: Poetry version
|
||||
required: true
|
||||
|
||||
cache-key:
|
||||
description: Cache key to use for manual handling of caching
|
||||
required: true
|
||||
|
||||
runs:
|
||||
using: composite
|
||||
steps:
|
||||
- uses: actions/setup-python@v5
|
||||
name: Setup python ${{ inputs.python-version }}
|
||||
id: setup-python
|
||||
with:
|
||||
python-version: ${{ inputs.python-version }}
|
||||
|
||||
- uses: actions/cache@v3
|
||||
id: cache-bin-poetry
|
||||
name: Cache Poetry binary - Python ${{ inputs.python-version }}
|
||||
env:
|
||||
SEGMENT_DOWNLOAD_TIMEOUT_MIN: "1"
|
||||
with:
|
||||
path: |
|
||||
/opt/pipx/venvs/poetry
|
||||
# This step caches the poetry installation, so make sure it's keyed on the poetry version as well.
|
||||
key: bin-poetry-${{ runner.os }}-${{ runner.arch }}-py-${{ inputs.python-version }}-${{ inputs.poetry-version }}
|
||||
|
||||
- name: Refresh shell hashtable and fixup softlinks
|
||||
if: steps.cache-bin-poetry.outputs.cache-hit == 'true'
|
||||
shell: bash
|
||||
env:
|
||||
POETRY_VERSION: ${{ inputs.poetry-version }}
|
||||
PYTHON_VERSION: ${{ inputs.python-version }}
|
||||
run: |
|
||||
set -eux
|
||||
|
||||
# Refresh the shell hashtable, to ensure correct `which` output.
|
||||
hash -r
|
||||
|
||||
# `actions/cache@v3` doesn't always seem able to correctly unpack softlinks.
|
||||
# Delete and recreate the softlinks pipx expects to have.
|
||||
rm /opt/pipx/venvs/poetry/bin/python
|
||||
cd /opt/pipx/venvs/poetry/bin
|
||||
ln -s "$(which "python$PYTHON_VERSION")" python
|
||||
chmod +x python
|
||||
cd /opt/pipx_bin/
|
||||
ln -s /opt/pipx/venvs/poetry/bin/poetry poetry
|
||||
chmod +x poetry
|
||||
|
||||
# Ensure everything got set up correctly.
|
||||
/opt/pipx/venvs/poetry/bin/python --version
|
||||
/opt/pipx_bin/poetry --version
|
||||
|
||||
- name: Install poetry
|
||||
if: steps.cache-bin-poetry.outputs.cache-hit != 'true'
|
||||
shell: bash
|
||||
env:
|
||||
POETRY_VERSION: ${{ inputs.poetry-version }}
|
||||
PYTHON_VERSION: ${{ inputs.python-version }}
|
||||
# Install poetry using the python version installed by setup-python step.
|
||||
run: pipx install "poetry==$POETRY_VERSION" --python '${{ steps.setup-python.outputs.python-path }}' --verbose
|
||||
|
||||
- name: Restore pip and poetry cached dependencies
|
||||
uses: actions/cache@v3
|
||||
env:
|
||||
SEGMENT_DOWNLOAD_TIMEOUT_MIN: "4"
|
||||
with:
|
||||
path: |
|
||||
~/.cache/pip
|
||||
~/.cache/pypoetry/virtualenvs
|
||||
~/.cache/pypoetry/cache
|
||||
~/.cache/pypoetry/artifacts
|
||||
./.venv
|
||||
key: py-deps-${{ runner.os }}-${{ runner.arch }}-py-${{ inputs.python-version }}-poetry-${{ inputs.poetry-version }}-${{ inputs.cache-key }}-${{ hashFiles('./poetry.lock') }}
|
||||
@@ -1,18 +0,0 @@
|
||||
version: 2
|
||||
updates:
|
||||
- package-ecosystem: "github-actions"
|
||||
directory: "/"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
|
||||
- package-ecosystem: "pip"
|
||||
directories:
|
||||
- "libs/checkpoint"
|
||||
- "libs/checkpoint-postgres"
|
||||
- "libs/checkpoint-sqlite"
|
||||
- "libs/cli"
|
||||
- "libs/langgraph"
|
||||
- "libs/prebuilt"
|
||||
- "libs/sdk-py"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
@@ -1,15 +1,10 @@
|
||||
import ast
|
||||
import os
|
||||
from itertools import filterfalse
|
||||
from typing import Dict, List, Tuple
|
||||
from typing import List, Tuple
|
||||
|
||||
ROOT_PATH = os.path.abspath(os.path.join(__file__, "..", "..", ".."))
|
||||
CLIENT_PATH = os.path.join(ROOT_PATH, "libs", "sdk-py", "langgraph_sdk", "client.py")
|
||||
ASYNC_TO_SYNC_METHOD_MAP: Dict[str, str] = {
|
||||
"aclose": "close",
|
||||
"__aenter__": "__enter__",
|
||||
"__aexit__": "__exit__",
|
||||
}
|
||||
|
||||
|
||||
def get_class_methods(node: ast.ClassDef) -> List[str]:
|
||||
@@ -27,7 +22,7 @@ def find_classes(tree: ast.AST) -> List[Tuple[str, List[str]]]:
|
||||
|
||||
def compare_sync_async_methods(sync_methods: List[str], async_methods: List[str]) -> List[str]:
|
||||
sync_set = set(sync_methods)
|
||||
async_set = {ASYNC_TO_SYNC_METHOD_MAP.get(async_method, async_method) for async_method in async_methods}
|
||||
async_set = set(async_methods)
|
||||
missing_in_sync = list(async_set - sync_set)
|
||||
missing_in_async = list(sync_set - async_set)
|
||||
return missing_in_sync + missing_in_async
|
||||
@@ -38,7 +33,7 @@ def main():
|
||||
tree = ast.parse(file.read())
|
||||
|
||||
classes = find_classes(tree)
|
||||
|
||||
|
||||
def is_sync(class_spec: Tuple[str, List[str]]) -> bool:
|
||||
return class_spec[0].startswith("Sync")
|
||||
|
||||
|
||||
@@ -1,164 +1,108 @@
|
||||
import logging
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
import sys
|
||||
import time
|
||||
from urllib import error, request
|
||||
|
||||
import langgraph_cli
|
||||
import langgraph_cli.config
|
||||
import langgraph_cli.docker
|
||||
from langgraph_cli.cli import prepare_args_and_stdin
|
||||
from langgraph_cli.constants import DEFAULT_PORT
|
||||
import langgraph_cli.config
|
||||
|
||||
from langgraph_cli.exec import Runner, subp_exec
|
||||
from langgraph_cli.progress import Progress
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
from langgraph_cli.constants import DEFAULT_PORT
|
||||
|
||||
|
||||
def test(config: pathlib.Path, port: int, tag: str, verbose: bool):
|
||||
"""Spin up API with Postgres/Redis via docker compose and wait until ready."""
|
||||
logger.info("Starting test...")
|
||||
def test(
|
||||
config: pathlib.Path,
|
||||
port: int,
|
||||
tag: str,
|
||||
verbose: bool,
|
||||
):
|
||||
with Runner() as runner, Progress(message="Pulling...") as set:
|
||||
# Detect docker/compose capabilities
|
||||
# check docker available
|
||||
capabilities = langgraph_cli.docker.check_capabilities(runner)
|
||||
|
||||
# Validate config and prepare compose stdin/args using built image
|
||||
# open config
|
||||
config_json = langgraph_cli.config.validate_config_file(config)
|
||||
args, stdin = prepare_args_and_stdin(
|
||||
capabilities=capabilities,
|
||||
config_path=config,
|
||||
config=config_json,
|
||||
docker_compose=None,
|
||||
port=port,
|
||||
watch=False,
|
||||
debugger_port=None,
|
||||
debugger_base_url=f"http://127.0.0.1:{port}",
|
||||
postgres_uri=None,
|
||||
api_version=None,
|
||||
image=tag,
|
||||
base_image=None,
|
||||
)
|
||||
|
||||
# Compose up with wait (implies detach), similar to `langgraph up --wait`
|
||||
args_up = [*args, "up", "--remove-orphans", "--wait"]
|
||||
|
||||
compose_cmd = ["docker", "compose"]
|
||||
if capabilities.compose_type == "standalone":
|
||||
compose_cmd = ["docker-compose"]
|
||||
|
||||
set("Starting...")
|
||||
try:
|
||||
runner.run(
|
||||
subp_exec(
|
||||
*compose_cmd,
|
||||
*args_up,
|
||||
input=stdin,
|
||||
verbose=verbose,
|
||||
)
|
||||
set("Running...")
|
||||
args = [
|
||||
"run",
|
||||
"--rm",
|
||||
"-p",
|
||||
f"{port}:8000",
|
||||
]
|
||||
if isinstance(config_json["env"], str):
|
||||
args.extend(
|
||||
[
|
||||
"--env-file",
|
||||
str(config.parent / config_json["env"]),
|
||||
]
|
||||
)
|
||||
except Exception as e: # noqa: BLE001
|
||||
# On failure, show diagnostics then ensure clean teardown
|
||||
sys.stderr.write(f"docker compose up failed: {e}\n")
|
||||
try:
|
||||
sys.stderr.write("\n== docker compose ps ==\n")
|
||||
runner.run(
|
||||
subp_exec(*compose_cmd, *args, "ps", input=stdin, verbose=False)
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
sys.stderr.write("\n== docker compose logs (api) ==\n")
|
||||
runner.run(
|
||||
subp_exec(
|
||||
*compose_cmd,
|
||||
*args,
|
||||
"logs",
|
||||
"langgraph-api",
|
||||
input=stdin,
|
||||
verbose=False,
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
try:
|
||||
runner.run(
|
||||
subp_exec(
|
||||
*compose_cmd,
|
||||
*args,
|
||||
"down",
|
||||
"-v",
|
||||
"--remove-orphans",
|
||||
input=stdin,
|
||||
verbose=False,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
raise
|
||||
|
||||
set("")
|
||||
base_url = f"http://localhost:{port}"
|
||||
ok_url = f"{base_url}/ok"
|
||||
logger.info(f"Waiting for {ok_url} to respond with 200...")
|
||||
deadline = time.time() + 30
|
||||
last_err: Exception | None = None
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
with request.urlopen(ok_url, timeout=2) as resp:
|
||||
if resp.status == 200:
|
||||
sys.stdout.write(
|
||||
f"""Ready!\n- API: {base_url}\n- /ok: 200 OK\n"""
|
||||
)
|
||||
sys.stdout.flush()
|
||||
break
|
||||
else:
|
||||
last_err = RuntimeError(f"Unexpected status: {resp.status}")
|
||||
logger.error(f"Unexpected status: {resp.status}")
|
||||
except error.URLError as e:
|
||||
logger.error(f"URLError: {e}")
|
||||
last_err = e
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.error(f"Exception: {e}")
|
||||
last_err = e
|
||||
time.sleep(0.5)
|
||||
else:
|
||||
logger.error("Timeout waiting for /ok to return 200")
|
||||
# Bring stack down before raising
|
||||
args_down = [*args, "down", "-v", "--remove-orphans"]
|
||||
try:
|
||||
runner.run(
|
||||
subp_exec(
|
||||
*compose_cmd,
|
||||
*args_down,
|
||||
input=stdin,
|
||||
verbose=verbose,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
raise SystemExit(
|
||||
f"/ok did not return 202 within timeout. Last error: {last_err}"
|
||||
for k, v in config_json["env"].items():
|
||||
args.extend(
|
||||
[
|
||||
"-e",
|
||||
f"{k}={v}",
|
||||
]
|
||||
)
|
||||
if capabilities.healthcheck_start_interval:
|
||||
args.extend(
|
||||
[
|
||||
"--health-interval",
|
||||
"5s",
|
||||
"--health-retries",
|
||||
"1",
|
||||
"--health-start-period",
|
||||
"10s",
|
||||
"--health-start-interval",
|
||||
"1s",
|
||||
]
|
||||
)
|
||||
else:
|
||||
args.extend(
|
||||
[
|
||||
"--health-interval",
|
||||
"5s",
|
||||
"--health-retries",
|
||||
"2",
|
||||
]
|
||||
)
|
||||
|
||||
_task = None
|
||||
|
||||
def on_stdout(line: str):
|
||||
nonlocal _task
|
||||
if "GET /ok" in line or "Uvicorn running on" in line:
|
||||
set("")
|
||||
sys.stdout.write(
|
||||
f"""Ready!
|
||||
- API: http://localhost:{port}
|
||||
"""
|
||||
)
|
||||
sys.stdout.flush()
|
||||
_task.cancel()
|
||||
return True
|
||||
return False
|
||||
|
||||
async def subp_exec_task(*args, **kwargs):
|
||||
nonlocal _task
|
||||
_task = asyncio.create_task(subp_exec(*args, **kwargs))
|
||||
await _task
|
||||
|
||||
# Clean up: bring compose stack down to free ports for next test
|
||||
logger.info("Test succeeded. Bringing down compose stack...")
|
||||
try:
|
||||
args_down = [*args, "down", "-v", "--remove-orphans"]
|
||||
runner.run(
|
||||
subp_exec(
|
||||
*compose_cmd,
|
||||
*args_down,
|
||||
input=stdin,
|
||||
subp_exec_task(
|
||||
"docker",
|
||||
*args,
|
||||
tag,
|
||||
verbose=verbose,
|
||||
on_stdout=on_stdout,
|
||||
)
|
||||
)
|
||||
logger.info("Compose stack down. Finishing...")
|
||||
except Exception:
|
||||
logger.exception("Failed to bring down compose stack")
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
logger.info("Test finished")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
@@ -166,12 +110,6 @@ if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-t", "--tag", type=str)
|
||||
parser.add_argument("-c", "--config", type=str, default="./langgraph.json")
|
||||
parser.add_argument("-p", "--port", type=int, default=DEFAULT_PORT)
|
||||
parser.add_argument("-p", "--port", default=DEFAULT_PORT)
|
||||
args = parser.parse_args()
|
||||
try:
|
||||
test(pathlib.Path(args.config), args.port, args.tag, verbose=True)
|
||||
except BaseException:
|
||||
logger.exception("Test failed")
|
||||
raise
|
||||
|
||||
logger.info("Test execution finished")
|
||||
test(pathlib.Path(args.config), args.port, args.tag, verbose=True)
|
||||
|
||||
@@ -3,8 +3,8 @@ name: CLI integration test
|
||||
on:
|
||||
workflow_call:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
env:
|
||||
POETRY_VERSION: "1.7.1"
|
||||
|
||||
jobs:
|
||||
build:
|
||||
@@ -13,106 +13,62 @@ jobs:
|
||||
matrix:
|
||||
python-version:
|
||||
- "3.10"
|
||||
- "3.14"
|
||||
example:
|
||||
- name: A
|
||||
workdir: libs/cli/examples
|
||||
tag: langgraph-test-a
|
||||
- name: B
|
||||
workdir: libs/cli/examples/graphs
|
||||
tag: langgraph-test-b
|
||||
- name: C
|
||||
workdir: libs/cli/examples/graphs_reqs_a
|
||||
tag: langgraph-test-c
|
||||
- name: D
|
||||
workdir: libs/cli/examples/graphs_reqs_b
|
||||
tag: langgraph-test-d
|
||||
- "3.11"
|
||||
name: "CLI integration test"
|
||||
defaults:
|
||||
run:
|
||||
working-directory: libs/cli
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- name: Get changed files
|
||||
id: changed-files
|
||||
uses: Ana06/get-changed-files@v2.3.0
|
||||
with:
|
||||
filter: "libs/cli/**"
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
- name: Set up Python ${{ matrix.python-version }} + Poetry ${{ env.POETRY_VERSION }}
|
||||
if: steps.changed-files.outputs.all
|
||||
uses: astral-sh/setup-uv@v7
|
||||
uses: "./.github/actions/poetry_setup"
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
enable-cache: true
|
||||
cache-suffix: "cli-integration-test"
|
||||
ignore-nothing-to-cache: true
|
||||
poetry-version: ${{ env.POETRY_VERSION }}
|
||||
cache-key: integration-test-cli
|
||||
- name: Setup env
|
||||
if: steps.changed-files.outputs.all
|
||||
working-directory: libs/cli/examples
|
||||
run: cat .env.example > .env
|
||||
- name: Install cli globally
|
||||
if: steps.changed-files.outputs.all
|
||||
run: pip install -e .
|
||||
- name: Build and test service ${{ matrix.example.name }}
|
||||
- name: Build and test service A
|
||||
if: steps.changed-files.outputs.all
|
||||
working-directory: ${{ matrix.example.workdir }}
|
||||
env:
|
||||
LANGSMITH_API_KEY: ${{ secrets.LANGSMITH_API_KEY }}
|
||||
working-directory: libs/cli/examples
|
||||
run: |
|
||||
# Build the image for this example
|
||||
langgraph build -t ${{ matrix.example.tag }}
|
||||
# Prepare environment file from local or parent example directory
|
||||
if [ -f .env.example ]; then cp .env.example .env; elif [ -f ../.env.example ]; then cp ../.env.example .env && cp ../.env.example ../.env; fi
|
||||
if [ -n "${{ secrets.LANGSMITH_API_KEY }}" ]; then echo "LANGSMITH_API_KEY=${{ secrets.LANGSMITH_API_KEY }}" >> .env; if [ -f ../.env ]; then echo "LANGSMITH_API_KEY=${{ secrets.LANGSMITH_API_KEY }}" >> ../.env; fi; fi
|
||||
# Run the integration test using the built tag
|
||||
# Compute repo root to reference the shared script robustly
|
||||
REPO_ROOT=$(git rev-parse --show-toplevel)
|
||||
timeout 60 python "$REPO_ROOT/.github/scripts/run_langgraph_cli_test.py" -t ${{ matrix.example.tag }}
|
||||
# The build-arg isn't used; just testing that we accept other args
|
||||
langgraph build -t langgraph-test-a --base-image "langchain/langgraph-trial"
|
||||
cp .env.example .envg
|
||||
timeout 60 python ../../../.github/scripts/run_langgraph_cli_test.py -c langgraph.json -t langgraph-test-a
|
||||
- name: Build and test service B
|
||||
if: steps.changed-files.outputs.all
|
||||
working-directory: libs/cli/examples/graphs
|
||||
run: |
|
||||
langgraph build -t langgraph-test-b --base-image "langchain/langgraph-trial"
|
||||
timeout 60 python ../../../../.github/scripts/run_langgraph_cli_test.py -t langgraph-test-b
|
||||
- name: Build and test service C
|
||||
if: steps.changed-files.outputs.all
|
||||
working-directory: libs/cli/examples/graphs_reqs_a
|
||||
run: |
|
||||
langgraph build -t langgraph-test-c --base-image "langchain/langgraph-trial"
|
||||
timeout 60 python ../../../../.github/scripts/run_langgraph_cli_test.py -t langgraph-test-c
|
||||
- name: Build and test service D
|
||||
if: steps.changed-files.outputs.all
|
||||
working-directory: libs/cli/examples/graphs_reqs_b
|
||||
run: |
|
||||
langgraph build -t langgraph-test-d --base-image "langchain/langgraph-trial"
|
||||
timeout 60 python ../../../../.github/scripts/run_langgraph_cli_test.py -t langgraph-test-d
|
||||
|
||||
- name: Build JS service
|
||||
if: ${{ steps.changed-files.outputs.all && matrix.example.name == 'A' }}
|
||||
if: steps.changed-files.outputs.all
|
||||
working-directory: libs/cli/js-examples
|
||||
run: |
|
||||
langgraph build -t langgraph-test-e
|
||||
|
||||
- name: Build JS monorepo service
|
||||
if: ${{ steps.changed-files.outputs.all && matrix.example.name == 'A' }}
|
||||
working-directory: libs/cli/js-monorepo-example
|
||||
run: |
|
||||
langgraph build -t langgraph-test-f -c apps/agent/langgraph.json --build-command "yarn run turbo build" --install-command "yarn install"
|
||||
|
||||
- name: Build Python monorepo service
|
||||
if: ${{ steps.changed-files.outputs.all && matrix.example.name == 'A' }}
|
||||
working-directory: libs/cli/python-monorepo-example
|
||||
run: |
|
||||
langgraph build -t langgraph-test-g -c apps/agent/langgraph.json
|
||||
cp apps/agent/.env.example apps/agent/.env
|
||||
if [ -n "${{ secrets.LANGSMITH_API_KEY }}" ]; then echo "LANGSMITH_API_KEY=${{ secrets.LANGSMITH_API_KEY }}" >> apps/agent/.env; fi
|
||||
timeout 60 python ../../../.github/scripts/run_langgraph_cli_test.py -t langgraph-test-g -c apps/agent/langgraph.json
|
||||
|
||||
- name: Build and test prerelease reqs service
|
||||
if: ${{ steps.changed-files.outputs.all && matrix.example.name == 'A' }}
|
||||
working-directory: libs/cli/examples/graph_prerelease_reqs
|
||||
run: |
|
||||
langgraph build -t langgraph-test-h
|
||||
cp ../.env.example .env
|
||||
if [ -n "${{ secrets.LANGSMITH_API_KEY }}" ]; then echo "LANGSMITH_API_KEY=${{ secrets.LANGSMITH_API_KEY }}" >> .env; fi
|
||||
timeout 60 python ../../../../.github/scripts/run_langgraph_cli_test.py -t langgraph-test-h
|
||||
echo "Finished starting up langgraph-test-h"
|
||||
LANGGRAPH_VERSION=$(docker run --rm --entrypoint "" langgraph-test-h python -c "import sys; from importlib.metadata import version; v = version('langgraph'); print(v);")
|
||||
if [ "$LANGGRAPH_VERSION" != "1.0.2" ]; then
|
||||
echo "LANGGRAPH_VERSION != 1.0.2; $LANGGRAPH_VERSION"
|
||||
exit 1
|
||||
fi
|
||||
LANGCHAIN_OPENAI_VERSION=$(docker run --rm --entrypoint "" langgraph-test-h python -c "import sys; from importlib.metadata import version; v = version('langchain-openai'); print(v);")
|
||||
if [ "$LANGCHAIN_OPENAI_VERSION" != "1.0.1" ]; then
|
||||
echo "LANGCHAIN_OPENAI_VERSION != 1.0.1; $LANGCHAIN_OPENAI_VERSION"
|
||||
exit 1
|
||||
fi
|
||||
LANGCHAIN_ANTHROPIC_VERSION=$(docker run --rm --entrypoint "" langgraph-test-h python -c "import sys; from importlib.metadata import version; v = version('langchain-anthropic'); print(v);")
|
||||
if [ "$LANGCHAIN_ANTHROPIC_VERSION" != "1.0.0a5" ]; then
|
||||
echo "LANGCHAIN_ANTHROPIC_VERSION != 1.0.0a5; $LANGCHAIN_ANTHROPIC_VERSION"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- name: Build and test prerelease reqs fail service
|
||||
if: ${{ steps.changed-files.outputs.all && matrix.example.name == 'A' }}
|
||||
working-directory: libs/cli/examples/graph_prerelease_reqs_fail
|
||||
run: |
|
||||
langgraph build -t langgraph-test-i || [ $? -eq 1 ]
|
||||
|
||||
+42
-14
@@ -8,10 +8,9 @@ on:
|
||||
type: string
|
||||
description: "From which folder this pipeline executes"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
env:
|
||||
POETRY_VERSION: "1.7.1"
|
||||
|
||||
# This env var allows us to get inline annotations when ruff has complaints.
|
||||
RUFF_OUTPUT_FORMAT: github
|
||||
|
||||
@@ -31,34 +30,54 @@ jobs:
|
||||
- "3.12"
|
||||
name: "lint #${{ matrix.python-version }}"
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- name: Get changed files
|
||||
id: changed-files
|
||||
uses: Ana06/get-changed-files@v2.3.0
|
||||
with:
|
||||
filter: "${{ inputs.working-directory }}/**"
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
- name: Set up Python ${{ matrix.python-version }} + Poetry ${{ env.POETRY_VERSION }}
|
||||
if: steps.changed-files.outputs.all
|
||||
uses: astral-sh/setup-uv@v7
|
||||
uses: "./.github/actions/poetry_setup"
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
enable-cache: true
|
||||
cache-suffix: lint-${{ inputs.working-directory }}
|
||||
poetry-version: ${{ env.POETRY_VERSION }}
|
||||
cache-key: lint-${{ inputs.working-directory }}
|
||||
|
||||
- name: Check Poetry File
|
||||
if: steps.changed-files.outputs.all
|
||||
shell: bash
|
||||
working-directory: ${{ inputs.working-directory }}
|
||||
run: poetry check
|
||||
|
||||
- name: Check lock file
|
||||
if: steps.changed-files.outputs.all
|
||||
shell: bash
|
||||
working-directory: ${{ inputs.working-directory }}
|
||||
run: poetry lock --check
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changed-files.outputs.all
|
||||
# Also installs dev/lint/test/typing dependencies, to ensure we have
|
||||
# type hints for as many of our libraries as possible.
|
||||
# This helps catch errors that require dependencies to be spotted, for example:
|
||||
# https://github.com/langchain-ai/langchain/pull/10249/files#diff-935185cd488d015f026dcd9e19616ff62863e8cde8c0bee70318d3ccbca98341
|
||||
#
|
||||
# If you change this configuration, make sure to change the `cache-key`
|
||||
# in the `poetry_setup` action above to stop using the old cache.
|
||||
# It doesn't matter how you change it, any change will cause a cache-bust.
|
||||
working-directory: ${{ inputs.working-directory }}
|
||||
run: uv sync --frozen --group lint
|
||||
run: poetry install --with dev
|
||||
|
||||
- name: Get .mypy_cache to speed up mypy
|
||||
if: steps.changed-files.outputs.all
|
||||
uses: actions/cache@v5
|
||||
uses: actions/cache@v3
|
||||
env:
|
||||
SEGMENT_DOWNLOAD_TIMEOUT_MIN: "2"
|
||||
with:
|
||||
path: |
|
||||
${{ inputs.working-directory }}/.mypy_cache
|
||||
key: mypy-lint-${{ runner.os }}-${{ runner.arch }}-py${{ matrix.python-version }}-${{ inputs.working-directory }}-${{ hashFiles(format('{0}/uv.lock', inputs.working-directory)) }}
|
||||
key: mypy-lint-${{ runner.os }}-${{ runner.arch }}-py${{ matrix.python-version }}-${{ inputs.working-directory }}-${{ hashFiles(format('{0}/poetry.lock', inputs.working-directory)) }}
|
||||
|
||||
- name: Analysing package code with our lint
|
||||
if: steps.changed-files.outputs.all
|
||||
@@ -73,18 +92,27 @@ jobs:
|
||||
|
||||
- name: Install test dependencies
|
||||
if: steps.changed-files.outputs.all
|
||||
# Also installs dev/lint/test/typing dependencies, to ensure we have
|
||||
# type hints for as many of our libraries as possible.
|
||||
# This helps catch errors that require dependencies to be spotted, for example:
|
||||
# https://github.com/langchain-ai/langchain/pull/10249/files#diff-935185cd488d015f026dcd9e19616ff62863e8cde8c0bee70318d3ccbca98341
|
||||
#
|
||||
# If you change this configuration, make sure to change the `cache-key`
|
||||
# in the `poetry_setup` action above to stop using the old cache.
|
||||
# It doesn't matter how you change it, any change will cause a cache-bust.
|
||||
working-directory: ${{ inputs.working-directory }}
|
||||
run: uv sync --group lint
|
||||
run: |
|
||||
poetry install --with dev
|
||||
|
||||
- name: Get .mypy_cache_test to speed up mypy
|
||||
if: steps.changed-files.outputs.all
|
||||
uses: actions/cache@v5
|
||||
uses: actions/cache@v3
|
||||
env:
|
||||
SEGMENT_DOWNLOAD_TIMEOUT_MIN: "2"
|
||||
with:
|
||||
path: |
|
||||
${{ inputs.working-directory }}/.mypy_cache_test
|
||||
key: mypy-test-${{ runner.os }}-${{ runner.arch }}-py${{ matrix.python-version }}-${{ inputs.working-directory }}-${{ hashFiles(format('{0}/uv.lock', inputs.working-directory)) }}
|
||||
key: mypy-test-${{ runner.os }}-${{ runner.arch }}-py${{ matrix.python-version }}-${{ inputs.working-directory }}-${{ hashFiles(format('{0}/poetry.lock', inputs.working-directory)) }}
|
||||
|
||||
- name: Analysing tests with our lint
|
||||
if: steps.changed-files.outputs.all
|
||||
|
||||
+12
-10
@@ -8,8 +8,8 @@ on:
|
||||
type: string
|
||||
description: "From which folder this pipeline executes"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
env:
|
||||
POETRY_VERSION: "1.7.1"
|
||||
|
||||
jobs:
|
||||
build:
|
||||
@@ -17,21 +17,21 @@ jobs:
|
||||
strategy:
|
||||
matrix:
|
||||
python-version:
|
||||
- "3.9"
|
||||
- "3.10"
|
||||
- "3.11"
|
||||
- "3.12"
|
||||
- "3.13"
|
||||
- "3.14"
|
||||
|
||||
name: "test #${{ matrix.python-version }}"
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: astral-sh/setup-uv@v7
|
||||
- uses: actions/checkout@v4
|
||||
- name: Set up Python ${{ matrix.python-version }} + Poetry ${{ env.POETRY_VERSION }}
|
||||
uses: "./.github/actions/poetry_setup"
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
enable-cache: true
|
||||
cache-suffix: test-${{ inputs.working-directory }}
|
||||
poetry-version: ${{ env.POETRY_VERSION }}
|
||||
cache-key: test-${{ inputs.working-directory }}
|
||||
- name: Login to Docker Hub
|
||||
uses: docker/login-action@v3
|
||||
if: ${{ !github.event.pull_request.head.repo.fork }}
|
||||
@@ -42,12 +42,14 @@ jobs:
|
||||
- name: Install dependencies
|
||||
shell: bash
|
||||
working-directory: ${{ inputs.working-directory }}
|
||||
run: uv sync --frozen --group test --no-dev
|
||||
run: |
|
||||
poetry install --with dev
|
||||
|
||||
- name: Run tests
|
||||
shell: bash
|
||||
working-directory: ${{ inputs.working-directory }}
|
||||
run: make test
|
||||
run: |
|
||||
make test
|
||||
|
||||
- name: Ensure the tests did not create any additional files
|
||||
shell: bash
|
||||
|
||||
@@ -3,8 +3,8 @@ name: test
|
||||
on:
|
||||
workflow_call:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
env:
|
||||
POETRY_VERSION: "1.7.1"
|
||||
|
||||
jobs:
|
||||
build:
|
||||
@@ -12,24 +12,24 @@ jobs:
|
||||
strategy:
|
||||
matrix:
|
||||
python-version:
|
||||
- "3.9"
|
||||
- "3.10"
|
||||
- "3.11"
|
||||
- "3.12"
|
||||
- "3.13"
|
||||
- "3.14"
|
||||
|
||||
defaults:
|
||||
run:
|
||||
working-directory: libs/langgraph
|
||||
name: "test #${{ matrix.python-version }}"
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: astral-sh/setup-uv@v7
|
||||
- uses: actions/checkout@v4
|
||||
- name: Set up Python ${{ matrix.python-version }} + Poetry ${{ env.POETRY_VERSION }}
|
||||
uses: "./.github/actions/poetry_setup"
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
enable-cache: true
|
||||
cache-suffix: "test-langgraph"
|
||||
poetry-version: ${{ env.POETRY_VERSION }}
|
||||
cache-key: test-langgraph
|
||||
- name: Login to Docker Hub
|
||||
uses: docker/login-action@v3
|
||||
if: ${{ !github.event.pull_request.head.repo.fork }}
|
||||
@@ -39,11 +39,13 @@ jobs:
|
||||
|
||||
- name: Install dependencies
|
||||
shell: bash
|
||||
run: uv sync --frozen --group test --no-dev
|
||||
run: |
|
||||
poetry install --with dev
|
||||
|
||||
- name: Run tests
|
||||
shell: bash
|
||||
run: make test_parallel
|
||||
run: |
|
||||
make test_parallel
|
||||
|
||||
- name: Ensure the tests did not create any additional files
|
||||
shell: bash
|
||||
|
||||
@@ -9,13 +9,12 @@ on:
|
||||
description: "From which folder this pipeline executes"
|
||||
|
||||
env:
|
||||
POETRY_VERSION: "1.7.1"
|
||||
PYTHON_VERSION: "3.10"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
build:
|
||||
if: github.ref == 'refs/heads/main'
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
outputs:
|
||||
@@ -23,14 +22,14 @@ jobs:
|
||||
version: ${{ steps.check-version.outputs.version }}
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python $${ env.PYTHON_VERSION }}
|
||||
uses: astral-sh/setup-uv@v7
|
||||
- name: Set up Python + Poetry ${{ env.POETRY_VERSION }}
|
||||
uses: "./.github/actions/poetry_setup"
|
||||
with:
|
||||
python-version: ${{ env.PYTHON_VERSION }}
|
||||
enable-cache: true
|
||||
cache-suffix: "release"
|
||||
poetry-version: ${{ env.POETRY_VERSION }}
|
||||
cache-key: release
|
||||
|
||||
# We want to keep this build stage *separate* from the release stage,
|
||||
# so that there's no sharing of permissions between them.
|
||||
@@ -44,11 +43,11 @@ jobs:
|
||||
# > from the publish job.
|
||||
# https://github.com/pypa/gh-action-pypi-publish#non-goals
|
||||
- name: Build project for distribution
|
||||
run: uv build
|
||||
run: poetry build
|
||||
working-directory: ${{ inputs.working-directory }}
|
||||
|
||||
- name: Upload build
|
||||
uses: actions/upload-artifact@v6
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: test-dist
|
||||
path: ${{ inputs.working-directory }}/dist/
|
||||
@@ -58,8 +57,8 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: ${{ inputs.working-directory }}
|
||||
run: |
|
||||
echo pkg-name=$(grep -m 1 "^name = " pyproject.toml | cut -d '"' -f 2)
|
||||
echo version=$(grep -m 1 "^version = " pyproject.toml | cut -d '"' -f 2)
|
||||
echo pkg-name="$(poetry version | cut -d ' ' -f 1)" >> $GITHUB_OUTPUT
|
||||
echo version="$(poetry version --short)" >> $GITHUB_OUTPUT
|
||||
|
||||
publish:
|
||||
needs:
|
||||
@@ -74,9 +73,9 @@ jobs:
|
||||
id-token: write
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- uses: actions/download-artifact@v7
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: test-dist
|
||||
path: ${{ inputs.working-directory }}/dist/
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
name: test
|
||||
|
||||
on:
|
||||
workflow_call:
|
||||
|
||||
env:
|
||||
POETRY_VERSION: "1.7.1"
|
||||
|
||||
jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
python-version:
|
||||
- "3.11"
|
||||
- "3.12"
|
||||
|
||||
defaults:
|
||||
run:
|
||||
working-directory: libs/scheduler-kafka
|
||||
name: "test #${{ matrix.python-version }}"
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Set up Python ${{ matrix.python-version }} + Poetry ${{ env.POETRY_VERSION }}
|
||||
uses: "./.github/actions/poetry_setup"
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
poetry-version: ${{ env.POETRY_VERSION }}
|
||||
cache-key: test-scheduler-kafka
|
||||
- name: Login to Docker Hub
|
||||
uses: docker/login-action@v3
|
||||
if: ${{ !github.event.pull_request.head.repo.fork }}
|
||||
with:
|
||||
username: ${{ secrets.DOCKERHUB_USERNAME }}
|
||||
password: ${{ secrets.DOCKERHUB_RO_TOKEN }}
|
||||
|
||||
- name: Install dependencies
|
||||
shell: bash
|
||||
run: |
|
||||
poetry install --with dev
|
||||
|
||||
- name: Run tests
|
||||
shell: bash
|
||||
run: |
|
||||
make test
|
||||
|
||||
- name: Ensure the tests did not create any additional files
|
||||
shell: bash
|
||||
run: |
|
||||
set -eu
|
||||
|
||||
STATUS="$(git status)"
|
||||
echo "$STATUS"
|
||||
|
||||
# grep will exit non-zero if the target message isn't found,
|
||||
# and `set -e` above will cause the step to fail.
|
||||
echo "$STATUS" | grep 'nothing to commit, working tree clean'
|
||||
@@ -7,8 +7,8 @@ on:
|
||||
paths:
|
||||
- "libs/**"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
env:
|
||||
POETRY_VERSION: "1.7.1"
|
||||
|
||||
jobs:
|
||||
benchmark:
|
||||
@@ -17,20 +17,20 @@ jobs:
|
||||
run:
|
||||
working-directory: libs/langgraph
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- run: SHA=$(git rev-parse HEAD) && echo "SHA=$SHA" >> $GITHUB_ENV
|
||||
- name: Set up Python 3.11
|
||||
uses: astral-sh/setup-uv@v7
|
||||
- name: Set up Python 3.11 + Poetry ${{ env.POETRY_VERSION }}
|
||||
uses: "./.github/actions/poetry_setup"
|
||||
with:
|
||||
python-version: "3.11"
|
||||
enable-cache: true
|
||||
cache-suffix: "bench"
|
||||
poetry-version: ${{ env.POETRY_VERSION }}
|
||||
cache-key: bench
|
||||
- name: Install dependencies
|
||||
run: uv sync --group test
|
||||
run: poetry install --with dev
|
||||
- name: Run benchmarks
|
||||
run: OUTPUT=out/benchmark-baseline.json make -s benchmark
|
||||
- name: Save outputs
|
||||
uses: actions/cache/save@v5
|
||||
uses: actions/cache/save@v4
|
||||
with:
|
||||
key: ${{ runner.os }}-benchmark-baseline-${{ env.SHA }}
|
||||
path: |
|
||||
|
||||
+12
-12
@@ -5,8 +5,8 @@ on:
|
||||
paths:
|
||||
- "libs/**"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
env:
|
||||
POETRY_VERSION: "1.7.1"
|
||||
|
||||
jobs:
|
||||
benchmark:
|
||||
@@ -15,22 +15,22 @@ jobs:
|
||||
run:
|
||||
working-directory: libs/langgraph
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- id: files
|
||||
name: Get changed files
|
||||
uses: Ana06/get-changed-files@v2.3.0
|
||||
with:
|
||||
format: json
|
||||
- name: Set up Python 3.11
|
||||
uses: astral-sh/setup-uv@v7
|
||||
- name: Set up Python 3.11 + Poetry ${{ env.POETRY_VERSION }}
|
||||
uses: "./.github/actions/poetry_setup"
|
||||
with:
|
||||
python-version: "3.11"
|
||||
enable-cache: true
|
||||
cache-suffix: "bench"
|
||||
poetry-version: ${{ env.POETRY_VERSION }}
|
||||
cache-key: bench
|
||||
- name: Install dependencies
|
||||
run: uv sync --group test
|
||||
run: poetry install --with dev
|
||||
- name: Download baseline
|
||||
uses: actions/cache/restore@v5
|
||||
uses: actions/cache/restore@v4
|
||||
with:
|
||||
key: ${{ runner.os }}-benchmark-baseline
|
||||
restore-keys: |
|
||||
@@ -43,7 +43,7 @@ jobs:
|
||||
run: |
|
||||
{
|
||||
echo 'OUTPUT<<EOF'
|
||||
make -s benchmark-fast
|
||||
make -s benchmark
|
||||
echo EOF
|
||||
} >> "$GITHUB_OUTPUT"
|
||||
- name: Compare benchmarks
|
||||
@@ -53,11 +53,11 @@ jobs:
|
||||
echo 'OUTPUT<<EOF'
|
||||
mv out/benchmark-baseline.json out/main.json
|
||||
mv out/benchmark.json out/changes.json
|
||||
uv run pyperf compare_to out/main.json out/changes.json --table --group-by-speed
|
||||
poetry run pyperf compare_to out/main.json out/changes.json --table --group-by-speed
|
||||
echo EOF
|
||||
} >> "$GITHUB_OUTPUT"
|
||||
- name: Annotation
|
||||
uses: actions/github-script@v8
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
script: |
|
||||
const file = JSON.parse(`${{ steps.files.outputs.added_modified_renamed }}`)[0]
|
||||
|
||||
+77
-56
@@ -3,13 +3,9 @@ name: CI
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
branches: [main]
|
||||
pull_request:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
# If another push to the same PR or branch happens while this workflow is still running,
|
||||
# cancel the earlier run in favor of the next run.
|
||||
#
|
||||
@@ -20,14 +16,17 @@ concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
env:
|
||||
POETRY_VERSION: "1.7.1"
|
||||
|
||||
jobs:
|
||||
changes:
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
python: ${{ steps.filter.outputs.python }}
|
||||
deps: ${{ steps.filter.outputs.deps }}
|
||||
sdk-js: ${{ steps.filter.outputs.sdk-js }}
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- uses: dorny/paths-filter@v3
|
||||
id: filter
|
||||
with:
|
||||
@@ -39,10 +38,10 @@ jobs:
|
||||
- 'libs/checkpoint/**'
|
||||
- 'libs/checkpoint-sqlite/**'
|
||||
- 'libs/checkpoint-postgres/**'
|
||||
- 'libs/scheduler-kafka/**'
|
||||
- 'libs/prebuilt/**'
|
||||
deps:
|
||||
- '**/pyproject.toml'
|
||||
- '**/uv.lock'
|
||||
sdk-js:
|
||||
- 'libs/sdk-js/**'
|
||||
|
||||
lint:
|
||||
needs: changes
|
||||
@@ -57,10 +56,10 @@ jobs:
|
||||
"libs/checkpoint",
|
||||
"libs/checkpoint-sqlite",
|
||||
"libs/checkpoint-postgres",
|
||||
|
||||
"libs/scheduler-kafka",
|
||||
"libs/prebuilt",
|
||||
]
|
||||
if: needs.changes.outputs.python == 'true' || needs.changes.outputs.deps == 'true'
|
||||
if: needs.changes.outputs.python == 'true'
|
||||
uses: ./.github/workflows/_lint.yml
|
||||
with:
|
||||
working-directory: ${{ matrix.working-directory }}
|
||||
@@ -78,9 +77,8 @@ jobs:
|
||||
"libs/checkpoint-sqlite",
|
||||
"libs/checkpoint-postgres",
|
||||
"libs/prebuilt",
|
||||
"libs/sdk-py",
|
||||
]
|
||||
if: needs.changes.outputs.python == 'true' || needs.changes.outputs.deps == 'true'
|
||||
if: needs.changes.outputs.python == 'true'
|
||||
uses: ./.github/workflows/_test.yml
|
||||
with:
|
||||
working-directory: ${{ matrix.working-directory }}
|
||||
@@ -89,78 +87,101 @@ jobs:
|
||||
# NOTE: we're testing langgraph separately because it requires a different matrix
|
||||
test-langgraph:
|
||||
needs: changes
|
||||
if: needs.changes.outputs.python == 'true' || needs.changes.outputs.deps == 'true'
|
||||
if: needs.changes.outputs.python == 'true'
|
||||
name: "cd libs/langgraph"
|
||||
uses: ./.github/workflows/_test_langgraph.yml
|
||||
secrets: inherit
|
||||
|
||||
# NOTE: we're testing scheduler-kafka separately because it requires a different matrix
|
||||
test-scheduler-kafka:
|
||||
needs: changes
|
||||
if: needs.changes.outputs.python == 'true'
|
||||
name: "cd libs/scheduler-kafka"
|
||||
uses: ./.github/workflows/_test_scheduler_kafka.yml
|
||||
secrets: inherit
|
||||
|
||||
check-sdk-methods:
|
||||
needs: changes
|
||||
if: needs.changes.outputs.python == 'true'
|
||||
name: "Check SDK methods matching"
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.11"
|
||||
- name: Run check_sdk_methods script
|
||||
run: python .github/scripts/check_sdk_methods.py
|
||||
|
||||
check-schema:
|
||||
needs: changes
|
||||
if: needs.changes.outputs.python == 'true'
|
||||
name: "Check CLI schema hasn't changed #${{ matrix.python-version }}"
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
python-version:
|
||||
- "3.13"
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
python-version: "3.13"
|
||||
enable-cache: true
|
||||
cache-suffix: "schema-check-cli"
|
||||
- name: Install CLI dependencies
|
||||
run: |
|
||||
cd libs/cli
|
||||
uv sync
|
||||
- name: Generate schema and check for changes
|
||||
run: |
|
||||
cd libs/cli
|
||||
# Create a temporary copy of the current schema
|
||||
cp schemas/schema.json schemas/schema.current.json
|
||||
# Generate new schema
|
||||
uv run python generate_schema.py
|
||||
# Compare the new schema with the original
|
||||
if ! diff -q schemas/schema.json schemas/schema.current.json > /dev/null; then
|
||||
echo "Error: Langgraph.json configuration schema has changed. Please run 'uv run python generate_schema.py' in the libs/cli directory and commit the changes."
|
||||
diff schemas/schema.json schemas/schema.current.json
|
||||
exit 1
|
||||
fi
|
||||
echo "Schema check passed - no changes detected"
|
||||
|
||||
integration-test:
|
||||
needs: changes
|
||||
if: needs.changes.outputs.python == 'true' || needs.changes.outputs.deps == 'true'
|
||||
if: needs.changes.outputs.python == 'true'
|
||||
name: CLI integration test
|
||||
uses: ./.github/workflows/_integration_test.yml
|
||||
secrets: inherit
|
||||
|
||||
lint-js:
|
||||
needs: changes
|
||||
if: needs.changes.outputs.sdk-js == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
working-directory:
|
||||
- "libs/sdk-js"
|
||||
defaults:
|
||||
run:
|
||||
working-directory: ${{ matrix.working-directory }}
|
||||
steps:
|
||||
- uses: actions/checkout@v3
|
||||
- name: Setup Node.js (LTS)
|
||||
uses: actions/setup-node@v3
|
||||
with:
|
||||
node-version: "20"
|
||||
cache: "yarn"
|
||||
cache-dependency-path: ${{ matrix.working-directory }}/yarn.lock
|
||||
- name: Install dependencies
|
||||
run: yarn install
|
||||
- name: Run lint
|
||||
run: yarn lint
|
||||
- name: Build
|
||||
run: yarn build
|
||||
|
||||
test-js:
|
||||
needs: changes
|
||||
if: needs.changes.outputs.sdk-js == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
working-directory:
|
||||
- "libs/sdk-js"
|
||||
defaults:
|
||||
run:
|
||||
working-directory: ${{ matrix.working-directory }}
|
||||
steps:
|
||||
- uses: actions/checkout@v3
|
||||
- name: Setup Node.js (LTS)
|
||||
uses: actions/setup-node@v3
|
||||
with:
|
||||
node-version: "20"
|
||||
cache: "yarn"
|
||||
cache-dependency-path: ${{ matrix.working-directory }}/yarn.lock
|
||||
- name: Install dependencies
|
||||
run: yarn install
|
||||
- name: Run tests
|
||||
run: yarn test
|
||||
|
||||
ci_success:
|
||||
name: "CI Success"
|
||||
needs:
|
||||
[
|
||||
lint,
|
||||
lint-js,
|
||||
test,
|
||||
test-langgraph,
|
||||
check-sdk-methods,
|
||||
check-schema,
|
||||
test-scheduler-kafka,
|
||||
integration-test,
|
||||
test-js,
|
||||
]
|
||||
if: |
|
||||
always()
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
---
|
||||
name: CI / cd . / make spell_check
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
pull_request:
|
||||
branches: [main]
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
defaults:
|
||||
run:
|
||||
working-directory: docs
|
||||
|
||||
jobs:
|
||||
codespell:
|
||||
name: (Check for spelling errors)
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Install Dependencies
|
||||
run: |
|
||||
pip install toml codespell==2.3.0 jupytext
|
||||
|
||||
- name: Extract Ignore Words List
|
||||
run: |
|
||||
# Use a Python script to extract the ignore words list from pyproject.toml
|
||||
python ../.github/workflows/extract_ignored_words_list.py
|
||||
id: extract_ignore_words
|
||||
|
||||
- name: Codespell
|
||||
uses: codespell-project/actions-codespell@v2
|
||||
with:
|
||||
skip: '*.ambr,*.lock,*.ipynb,*.yaml,*.zlib,*.md'
|
||||
ignore_words_list: ${{ steps.extract_ignore_words.outputs.ignore_words_list }}
|
||||
# We do this to avoid spellchecking cell outputs
|
||||
- name: Codespell Notebooks
|
||||
run: make codespell
|
||||
@@ -0,0 +1,170 @@
|
||||
name: Deploy Docs
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
workflow_dispatch:
|
||||
|
||||
env:
|
||||
POETRY_VERSION: "1.7.1"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
pages: write
|
||||
id-token: write
|
||||
|
||||
concurrency:
|
||||
group: "pages"
|
||||
cancel-in-progress: false
|
||||
|
||||
defaults:
|
||||
run:
|
||||
working-directory: docs
|
||||
|
||||
jobs:
|
||||
get-changed-files:
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
changed-files: ${{ steps.changed-files.outputs.added_modified }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Get changed files
|
||||
id: changed-files
|
||||
uses: Ana06/get-changed-files@v2.3.0
|
||||
with:
|
||||
filter: "docs/docs/**"
|
||||
|
||||
run-changed-notebooks:
|
||||
needs: get-changed-files
|
||||
uses: ./.github/workflows/run_notebooks.yml
|
||||
secrets: inherit
|
||||
with:
|
||||
changed-files: ${{ needs.get-changed-files.outputs.changed-files }}
|
||||
|
||||
deploy:
|
||||
# needs: run-changed-notebooks
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10 # Job will be cancelled if it runs for more than 10 minutes
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.MKDOCS_GITHUB_TOKEN }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Set up Python + Poetry ${{ env.POETRY_VERSION }}
|
||||
uses: "./.github/actions/poetry_setup"
|
||||
with:
|
||||
python-version: "3.12"
|
||||
poetry-version: ${{ env.POETRY_VERSION }}
|
||||
cache-key: docs
|
||||
|
||||
- name: Use Node.js
|
||||
uses: actions/setup-node@v3
|
||||
with:
|
||||
node-version: "22"
|
||||
cache: "yarn"
|
||||
cache-dependency-path: docs/yarn.lock
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
yarn
|
||||
poetry install --with test --with docs --no-root
|
||||
poetry run pip install -U \
|
||||
pytest \
|
||||
pytest-check-links \
|
||||
GitPython \
|
||||
"git+https://github.com/benjamincburns/markdown-exec.git@cc0d39d737e5ffd4b83d23cd8729d7ea16e363c8"
|
||||
|
||||
# we run this installation only for internal PRs
|
||||
# as GITHUB_TOKEN is not available for PRs from outside contributors
|
||||
if [ -n "${GITHUB_TOKEN}" ]; then
|
||||
poetry run pip install "git+https://${GITHUB_TOKEN}@github.com/langchain-ai/mkdocs-material-insiders.git"
|
||||
fi
|
||||
|
||||
poetry run jupyter kernelspec list
|
||||
poetry run python3 -m ipykernel install --user --name=python3
|
||||
npm install -g tslab
|
||||
poetry run tslab install --python=python3
|
||||
poetry run jupyter kernelspec list
|
||||
|
||||
- name: Run unit tests
|
||||
# Run unit tests on the docs build pipeline
|
||||
run: make tests
|
||||
- name: Lint Docs
|
||||
# This step lints the docs using the existing linting set up.
|
||||
# It should be very fast and should not require any external services.
|
||||
run: make lint-docs
|
||||
- name: Build llms-text
|
||||
run: make llms-text
|
||||
- name: Build site
|
||||
run: make build-docs
|
||||
env:
|
||||
MKDOCS_GIT_COMMITTERS_APIKEY: ${{ secrets.MKDOCS_GIT_COMMITTERS_APIKEY }}
|
||||
OPENAI_API_KEY: sf-proj-1234567890 # fake placeholder, shouldn't actually be used
|
||||
ANTHROPIC_API_KEY: sk-ant-api03-1234567890 # fake placeholder, shouldn't actually be used
|
||||
- name: Check links in notebooks
|
||||
env:
|
||||
LANGCHAIN_API_KEY: test
|
||||
run: |
|
||||
if [ "${{ github.event_name }}" == "schedule" ] || [ "${{ github.event_name }}" == "workflow_dispatch" ] || ([ "${{ github.event_name }}" == "push" ] && [ "${{ github.ref }}" == "refs/heads/main" ]); then
|
||||
echo "Running link check on all HTML files matching notebooks in docs directory..."
|
||||
poetry run pytest -v \
|
||||
--check-links-ignore "https://(api|web|docs)\.smith\.langchain\.com/.*" \
|
||||
--check-links-ignore "https://academy\.langchain\.com/.*" \
|
||||
--check-links-ignore "https://x.com/.*" \
|
||||
--check-links-ignore "https://twitter.com/.*" \
|
||||
--check-links-ignore "https://github\.com/.*" \
|
||||
--check-links-ignore "http://localhost:8123/.*" \
|
||||
--check-links-ignore "http://localhost:2024.*" \
|
||||
--check-links-ignore "http://127.0.0.1:.*" \
|
||||
--check-links-ignore "/.*\.(ipynb|html)$" \
|
||||
--check-links-ignore "https://python\.langchain\.com/.*" \
|
||||
--check-links-ignore "https://openai\.com/.*" \
|
||||
--check-links-ignore "https://www\.uber\.com/.*" \
|
||||
--check-links-ignore "https://pepy\.tech/.*" \
|
||||
--check-links $(find site -name "index.html" | grep -v 'storm/index.html')
|
||||
|
||||
else
|
||||
echo "Fetching changes from origin/main..."
|
||||
git fetch origin main
|
||||
echo "Checking for changed notebook files..."
|
||||
CHANGED_FILES=$(git diff --name-only --diff-filter=d origin/main | grep 'docs/docs/.*\.ipynb$' | grep -v 'storm.ipynb' | sed -E 's|^docs/docs/|site/|; s/\.ipynb$/\/index.html/' || true)
|
||||
echo "Changed files: ${CHANGED_FILES}"
|
||||
if [ -n "${CHANGED_FILES}" ]; then
|
||||
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 "https://academy\.langchain\.com/.*" \
|
||||
--check-links-ignore "http://localhost:8123/.*" \
|
||||
--check-links-ignore "http://localhost:2024.*" \
|
||||
--check-links-ignore "http://127.0.0.1:.*" \
|
||||
--check-links-ignore "https://x.com/.*" \
|
||||
--check-links-ignore "https://twitter.com/.*" \
|
||||
--check-links-ignore "https://github\.com/.*" \
|
||||
--check-links-ignore "/.*\.(ipynb|html)$" \
|
||||
--check-links ${CHANGED_FILES} \
|
||||
|| ([ $? = 5 ] && exit 0 || exit $?)
|
||||
else
|
||||
echo "No notebook files changed."
|
||||
fi
|
||||
fi
|
||||
|
||||
- name: Configure GitHub Pages
|
||||
if: github.ref == 'refs/heads/main'
|
||||
uses: actions/configure-pages@v4
|
||||
|
||||
- name: Upload Pages Artifact
|
||||
# if: github.ref == 'refs/heads/main'
|
||||
uses: actions/upload-pages-artifact@v3
|
||||
with:
|
||||
path: ./docs/site/
|
||||
|
||||
- name: Deploy to GitHub Pages
|
||||
if: github.ref == 'refs/heads/main'
|
||||
id: deployment
|
||||
uses: actions/deploy-pages@v4
|
||||
@@ -0,0 +1,10 @@
|
||||
import toml
|
||||
|
||||
pyproject_toml = toml.load("pyproject.toml")
|
||||
|
||||
# Extract the ignore words list (adjust the key as per your TOML structure)
|
||||
ignore_words_list = (
|
||||
pyproject_toml.get("tool", {}).get("codespell", {}).get("ignore-words-list")
|
||||
)
|
||||
|
||||
print(f"::set-output name=ignore_words_list::{ignore_words_list}") # noqa: T201
|
||||
@@ -0,0 +1,49 @@
|
||||
name: Check Docs & Links
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
schedule:
|
||||
- cron: "0 5 * * *"
|
||||
workflow_dispatch:
|
||||
|
||||
env:
|
||||
POETRY_VERSION: "1.7.1"
|
||||
|
||||
jobs:
|
||||
markdown-link-check:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Check links in Markdown files
|
||||
uses: gaurav-nelson/github-action-markdown-link-check@v1
|
||||
with:
|
||||
folder-path: "docs/"
|
||||
check-modified-files-only: ${{ github.event_name != 'schedule' }}
|
||||
file-path: "./README.md"
|
||||
config-file: "./.markdown-link-check.config.json"
|
||||
|
||||
check-readmes-synced:
|
||||
# This checks that the repo README.md is identical to the libs/langgraph/README.md
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 1
|
||||
|
||||
- name: Check README.md is in sync
|
||||
run: |
|
||||
if ! diff -q README.md libs/langgraph/README.md >/dev/null; then
|
||||
echo "README.md is out of sync with libs/langgraph/README.md"
|
||||
diff -C 3 README.md libs/langgraph/README.md
|
||||
exit 1
|
||||
fi
|
||||
@@ -1,46 +0,0 @@
|
||||
name: PR Title Lint
|
||||
|
||||
permissions:
|
||||
pull-requests: read
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [opened, edited, synchronize]
|
||||
|
||||
jobs:
|
||||
lint-pr-title:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Validate PR Title
|
||||
uses: amannn/action-semantic-pull-request@v6
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
with:
|
||||
types: |
|
||||
feat
|
||||
fix
|
||||
docs
|
||||
style
|
||||
refactor
|
||||
perf
|
||||
test
|
||||
build
|
||||
ci
|
||||
chore
|
||||
revert
|
||||
release
|
||||
scopes: |
|
||||
checkpoint
|
||||
checkpoint-postgres
|
||||
checkpoint-sqlite
|
||||
cli
|
||||
langgraph
|
||||
prebuilt
|
||||
scheduler-kafka
|
||||
sdk-py
|
||||
docs
|
||||
ci
|
||||
deps
|
||||
requireScope: false
|
||||
ignoreLabels: |
|
||||
ignore-lint-pr-title
|
||||
@@ -8,14 +8,13 @@ on:
|
||||
type: string
|
||||
default: "libs/langgraph"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
env:
|
||||
PYTHON_VERSION: "3.11"
|
||||
POETRY_VERSION: "1.7.1"
|
||||
|
||||
jobs:
|
||||
build:
|
||||
if: github.ref == 'refs/heads/main'
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
outputs:
|
||||
@@ -25,14 +24,14 @@ jobs:
|
||||
tag: ${{ steps.check-version.outputs.tag }}
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python
|
||||
uses: astral-sh/setup-uv@v7
|
||||
- name: Set up Python + Poetry ${{ env.POETRY_VERSION }}
|
||||
uses: "./.github/actions/poetry_setup"
|
||||
with:
|
||||
python-version: ${{ env.PYTHON_VERSION }}
|
||||
enable-cache: true
|
||||
cache-suffix: "release"
|
||||
poetry-version: ${{ env.POETRY_VERSION }}
|
||||
cache-key: release
|
||||
|
||||
# We want to keep this build stage *separate* from the release stage,
|
||||
# so that there's no sharing of permissions between them.
|
||||
@@ -46,11 +45,11 @@ jobs:
|
||||
# > from the publish job.
|
||||
# https://github.com/pypa/gh-action-pypi-publish#non-goals
|
||||
- name: Build project for distribution
|
||||
run: uv build
|
||||
run: poetry build
|
||||
working-directory: ${{ inputs.working-directory }}
|
||||
|
||||
- name: Upload build
|
||||
uses: actions/upload-artifact@v6
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: dist
|
||||
path: ${{ inputs.working-directory }}/dist/
|
||||
@@ -60,14 +59,8 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: ${{ inputs.working-directory }}
|
||||
run: |
|
||||
PKG_NAME=$(grep -m 1 "^name = " pyproject.toml | cut -d '"' -f 2)
|
||||
if grep -q 'dynamic.*=.*\[.*"version".*\]' pyproject.toml; then
|
||||
# handle dynamic versioning
|
||||
DIR_NAME=$(echo "$PKG_NAME" | tr '-' '_')
|
||||
VERSION=$(grep -m 1 '^__version__' "${DIR_NAME}/__init__.py" | cut -d '"' -f 2)
|
||||
else
|
||||
VERSION=$(grep -m 1 "^version = " pyproject.toml | cut -d '"' -f 2)
|
||||
fi
|
||||
PKG_NAME="$(poetry version | cut -d ' ' -f 1)"
|
||||
VERSION="$(poetry version --short)"
|
||||
SHORT_PKG_NAME="$(echo "$PKG_NAME" | sed -e 's/langgraph//g' -e 's/-//g')"
|
||||
if [ -z $SHORT_PKG_NAME ]; then
|
||||
TAG="$VERSION"
|
||||
@@ -86,7 +79,7 @@ jobs:
|
||||
outputs:
|
||||
release-body: ${{ steps.generate-release-body.outputs.release-body }}
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
repository: langchain-ai/langgraph
|
||||
path: langgraph
|
||||
@@ -142,9 +135,7 @@ jobs:
|
||||
needs:
|
||||
- build
|
||||
- release-notes
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
permissions: write-all
|
||||
uses: ./.github/workflows/_test_release.yml
|
||||
with:
|
||||
working-directory: ${{ inputs.working-directory }}
|
||||
@@ -157,7 +148,7 @@ jobs:
|
||||
- test-pypi-publish
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
# We explicitly *don't* set up caching here. This ensures our tests are
|
||||
# maximally sensitive to catching breakage.
|
||||
@@ -172,11 +163,11 @@ jobs:
|
||||
# - The package is published, and it breaks on the missing dependency when
|
||||
# used in the real world.
|
||||
|
||||
- name: Set up Python
|
||||
uses: astral-sh/setup-uv@v7
|
||||
- name: Set up Python + Poetry ${{ env.POETRY_VERSION }}
|
||||
uses: "./.github/actions/poetry_setup"
|
||||
with:
|
||||
python-version: ${{ env.PYTHON_VERSION }}
|
||||
enable-cache: true
|
||||
poetry-version: ${{ env.POETRY_VERSION }}
|
||||
|
||||
- name: Import published package
|
||||
shell: bash
|
||||
@@ -194,18 +185,18 @@ jobs:
|
||||
# - attempt install again after 5 seconds if it fails because there is
|
||||
# sometimes a delay in availability on test pypi
|
||||
run: |
|
||||
uv run pip install \
|
||||
poetry run pip install \
|
||||
--extra-index-url https://test.pypi.org/simple/ \
|
||||
"$PKG_NAME==$VERSION" || \
|
||||
( \
|
||||
sleep 5 && \
|
||||
uv run pip install \
|
||||
poetry run pip install \
|
||||
--extra-index-url https://test.pypi.org/simple/ \
|
||||
"$PKG_NAME==$VERSION" \
|
||||
)
|
||||
|
||||
if [[ "$PKG_NAME" == *prebuilt* ]]; then
|
||||
uv run pip install langgraph
|
||||
poetry run pip install langgraph
|
||||
fi
|
||||
|
||||
if [[ "$PKG_NAME" == *checkpoint* || "$PKG_NAME" == *prebuilt* ]]; then
|
||||
@@ -218,10 +209,10 @@ jobs:
|
||||
IMPORT_NAME="$(echo "$PKG_NAME" | sed s/-/_/g)"
|
||||
fi
|
||||
|
||||
uv run python -c "import $IMPORT_NAME; print(dir($IMPORT_NAME))"
|
||||
poetry run python -c "import $IMPORT_NAME; print(dir($IMPORT_NAME))"
|
||||
|
||||
- name: Import test dependencies
|
||||
run: uv sync --group test
|
||||
run: poetry install --with dev
|
||||
working-directory: ${{ inputs.working-directory }}
|
||||
|
||||
# Overwrite the local version of the package with the test PyPI version.
|
||||
@@ -232,7 +223,7 @@ jobs:
|
||||
PKG_NAME: ${{ needs.build.outputs.pkg-name }}
|
||||
VERSION: ${{ needs.build.outputs.version }}
|
||||
run: |
|
||||
uv run pip install \
|
||||
poetry run pip install \
|
||||
--extra-index-url https://test.pypi.org/simple/ \
|
||||
"$PKG_NAME==$VERSION"
|
||||
|
||||
@@ -260,16 +251,16 @@ jobs:
|
||||
working-directory: ${{ inputs.working-directory }}
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python
|
||||
uses: astral-sh/setup-uv@v7
|
||||
- name: Set up Python + Poetry ${{ env.POETRY_VERSION }}
|
||||
uses: "./.github/actions/poetry_setup"
|
||||
with:
|
||||
python-version: ${{ env.PYTHON_VERSION }}
|
||||
enable-cache: true
|
||||
cache-suffix: "release"
|
||||
poetry-version: ${{ env.POETRY_VERSION }}
|
||||
cache-key: release
|
||||
|
||||
- uses: actions/download-artifact@v7
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: dist
|
||||
path: ${{ inputs.working-directory }}/dist/
|
||||
@@ -301,16 +292,16 @@ jobs:
|
||||
working-directory: ${{ inputs.working-directory }}
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python
|
||||
uses: astral-sh/setup-uv@v7
|
||||
- name: Set up Python + Poetry ${{ env.POETRY_VERSION }}
|
||||
uses: "./.github/actions/poetry_setup"
|
||||
with:
|
||||
python-version: ${{ env.PYTHON_VERSION }}
|
||||
enable-cache: true
|
||||
cache-suffix: "release"
|
||||
poetry-version: ${{ env.POETRY_VERSION }}
|
||||
cache-key: release
|
||||
|
||||
- uses: actions/download-artifact@v7
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: dist
|
||||
path: ${{ inputs.working-directory }}/dist/
|
||||
@@ -322,6 +313,5 @@ jobs:
|
||||
token: ${{ secrets.GITHUB_TOKEN }}
|
||||
generateReleaseNotes: false
|
||||
tag: ${{needs.build.outputs.tag}}
|
||||
name: ${{ needs.build.outputs.pkg-name }}==${{ needs.build.outputs.version }}
|
||||
body: ${{ needs.release-notes.outputs.release-body }}
|
||||
commit: ${{ github.sha }}
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
name: JS Release
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
publish:
|
||||
# Disallow publishing from branches that aren't `main`.
|
||||
if: github.ref == 'refs/heads/main'
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
strategy:
|
||||
matrix:
|
||||
working-directory:
|
||||
- "libs/sdk-js"
|
||||
|
||||
defaults:
|
||||
run:
|
||||
working-directory: ${{ matrix.working-directory }}
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
# JS Build
|
||||
- name: Use Node.js
|
||||
uses: actions/setup-node@v3
|
||||
with:
|
||||
node-version: "20"
|
||||
cache: "yarn"
|
||||
cache-dependency-path: ${{ matrix.working-directory }}/yarn.lock
|
||||
|
||||
- name: Install dependencies
|
||||
run: yarn install
|
||||
- name: Build
|
||||
run: yarn build
|
||||
- name: Publish package to NPM
|
||||
run: |
|
||||
echo "//registry.npmjs.org/:_authToken=${{ secrets.NPM_TOKEN }}" > .npmrc
|
||||
npm publish
|
||||
@@ -0,0 +1,82 @@
|
||||
name: Run notebooks
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
workflow_call:
|
||||
inputs:
|
||||
changed-files:
|
||||
required: false
|
||||
type: string
|
||||
description: "JSON string of changed files"
|
||||
schedule:
|
||||
- cron: '0 13 * * *'
|
||||
|
||||
defaults:
|
||||
run:
|
||||
working-directory: docs
|
||||
|
||||
jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
lib-version:
|
||||
- "development"
|
||||
- "latest"
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Set up Python + Poetry
|
||||
uses: "./.github/actions/poetry_setup"
|
||||
with:
|
||||
python-version: 3.11
|
||||
poetry-version: 1.7.1
|
||||
cache-key: test-langgraph-notebooks
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
poetry install --with test
|
||||
poetry run pip install jupyter
|
||||
|
||||
- name: Start services
|
||||
run: make start-services
|
||||
|
||||
- name: Pre-download tiktoken files
|
||||
run: |
|
||||
poetry run python _scripts/download_tiktoken.py
|
||||
|
||||
- name: Prepare notebooks
|
||||
run: |
|
||||
if [ "${{ matrix.lib-version }}" = "development" ]; then
|
||||
poetry run python _scripts/prepare_notebooks_for_ci.py --comment-install-cells
|
||||
else
|
||||
poetry run python _scripts/prepare_notebooks_for_ci.py
|
||||
fi
|
||||
|
||||
- name: Run notebooks
|
||||
env:
|
||||
# these won't actually be used because of the VCR cassettes
|
||||
# but need to set them to avoid triggering getpass()
|
||||
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
|
||||
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
|
||||
TAVILY_API_KEY: ${{ secrets.TAVILY_API_KEY }}
|
||||
LANGSMITH_API_KEY: ${{ secrets.LANGSMITH_API_KEY }}
|
||||
NOMIC_API_KEY: ${{ secrets.NOMIC_API_KEY }}
|
||||
COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
|
||||
FIREWORKS_API_KEY: ${{ secrets.FIREWORKS_API_KEY }}
|
||||
run: |
|
||||
if [ "${{ github.event_name }}" = "workflow_dispatch" ] || [ "${{ github.event_name }}" = "schedule" ]; then
|
||||
echo "Running all notebooks"
|
||||
./_scripts/execute_notebooks.sh
|
||||
else
|
||||
CHANGED_FILES=$(echo '${{ inputs.changed-files }}' | tr ' ' '\n' | sed 's|^docs/docs/|docs/|' | grep '\.ipynb$' || true)
|
||||
if [ -n "$CHANGED_FILES" ]; then
|
||||
echo "Running changed notebooks: $CHANGED_FILES"
|
||||
./_scripts/execute_notebooks.sh $CHANGED_FILES
|
||||
else
|
||||
echo "No notebook files changed, skipping execution"
|
||||
fi
|
||||
fi
|
||||
|
||||
- name: Stop services
|
||||
run: make stop-services
|
||||
@@ -0,0 +1,29 @@
|
||||
name: Check File Size
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
file-size-check:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Get changed files
|
||||
id: changed-files
|
||||
uses: tj-actions/changed-files@v44
|
||||
- name: Filter by size
|
||||
# TODO: roll back the web voyager hack
|
||||
run: |
|
||||
large_added_files=$(find ${{ steps.changed-files.outputs.added_files }} -maxdepth 0 -size +1M | grep -v "web_voyager" || true)
|
||||
if [ -n "$large_added_files" ]; then
|
||||
echo "Large files added: $large_added_files"
|
||||
echo "# Large files added:" >> $GITHUB_STEP_SUMMARY
|
||||
echo "$large_added_files" >> $GITHUB_STEP_SUMMARY
|
||||
exit 1
|
||||
fi
|
||||
@@ -1,45 +0,0 @@
|
||||
name: UV Lock Upgrade
|
||||
|
||||
on:
|
||||
schedule:
|
||||
# run at midnight every Sunday
|
||||
- cron: '0 0 * * 0'
|
||||
# allow manual triggering
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
pull-requests: write
|
||||
|
||||
jobs:
|
||||
upgrade-dependencies:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
# use minimum supported Python version
|
||||
python-version: "3.10"
|
||||
enable-cache: true
|
||||
cache-suffix: "uv-lock-upgrade"
|
||||
|
||||
- name: Run uv lock --upgrade in all Python packages
|
||||
run: make lock-upgrade
|
||||
|
||||
- name: Create Pull Request
|
||||
uses: peter-evans/create-pull-request@v8
|
||||
with:
|
||||
token: ${{ secrets.GITHUB_TOKEN }}
|
||||
commit-message: "chore(deps): upgrade dependencies with `uv lock --upgrade`"
|
||||
title: "chore(deps): upgrade dependencies with `uv lock --upgrade`"
|
||||
body: |
|
||||
This PR updates the dependencies in all Python packages using `uv lock --upgrade`.
|
||||
|
||||
This is an automated PR created by the UV Lock Upgrade workflow.
|
||||
branch: deps/uv-lock-upgrade
|
||||
delete-branch: true
|
||||
labels: |
|
||||
dependencies
|
||||
+82
-2
@@ -6,6 +6,9 @@ __pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
@@ -51,12 +54,27 @@ coverage.xml
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
docs/docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
target/
|
||||
|
||||
@@ -71,9 +89,23 @@ ipython_config.py
|
||||
# pyenv
|
||||
.python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.envrc
|
||||
@@ -85,6 +117,16 @@ ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
@@ -96,7 +138,45 @@ dmypy.json
|
||||
# macOS display setting files
|
||||
.DS_Store
|
||||
|
||||
# Wandb directory
|
||||
wandb/
|
||||
|
||||
# asdf tool versions
|
||||
.tool-versions
|
||||
/.ruff_cache/
|
||||
|
||||
*.pkl
|
||||
*.bin
|
||||
|
||||
# integration test artifacts
|
||||
data_map*
|
||||
\[('_type', 'fake'), ('stop', None)]
|
||||
|
||||
# Replit files
|
||||
*replit*
|
||||
|
||||
node_modules
|
||||
docs/.yarn/
|
||||
docs/node_modules/
|
||||
docs/.docusaurus/
|
||||
docs/.cache-loader/
|
||||
docs/_dist
|
||||
docs/api_reference/api_reference.rst
|
||||
docs/api_reference/experimental_api_reference.rst
|
||||
docs/api_reference/_build
|
||||
docs/api_reference/*/
|
||||
!docs/api_reference/_static/
|
||||
!docs/api_reference/templates/
|
||||
!docs/api_reference/themes/
|
||||
docs/docs_skeleton/build
|
||||
docs/docs_skeleton/node_modules
|
||||
docs/docs_skeleton/yarn.lock
|
||||
|
||||
# Any new jupyter notebooks
|
||||
# not intended for the repo
|
||||
Untitled*.ipynb
|
||||
|
||||
Chinook.db
|
||||
|
||||
.vercel
|
||||
.turbo
|
||||
.editorconfig
|
||||
.scratch
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
{
|
||||
"aliveStatusCodes": [200, 206, 402],
|
||||
"ignorePatterns": ["*dcbadge.vercel.app*"]
|
||||
}
|
||||
@@ -1,55 +0,0 @@
|
||||
# AGENTS Instructions
|
||||
|
||||
This repository is a monorepo. Each library lives in a subdirectory under `libs/`.
|
||||
|
||||
When you modify code in any library, run the following commands in that library's directory before creating a pull request:
|
||||
|
||||
- `make format` – run code formatters
|
||||
- `make lint` – run the linter
|
||||
- `make test` – execute the test suite
|
||||
|
||||
To run a particular test file or to pass additional pytest options you can specify the `TEST` variable:
|
||||
|
||||
```txt
|
||||
TEST=path/to/test.py make test
|
||||
```
|
||||
|
||||
Other pytest arguments can also be supplied inside the `TEST` variable.
|
||||
|
||||
## Libraries
|
||||
|
||||
The repository contains several Python and JavaScript/TypeScript libraries.
|
||||
Below is a high-level overview:
|
||||
|
||||
- **checkpoint** – base interfaces for LangGraph checkpointers.
|
||||
- **checkpoint-postgres** – Postgres implementation of the checkpoint saver.
|
||||
- **checkpoint-sqlite** – SQLite implementation of the checkpoint saver.
|
||||
- **cli** – official command-line interface for LangGraph.
|
||||
- **langgraph** – core framework for building stateful, multi-actor agents.
|
||||
- **prebuilt** – high-level APIs for creating and running agents and tools.
|
||||
- **sdk-js** – JS/TS SDK for interacting with the LangGraph REST API.
|
||||
- **sdk-py** – Python SDK for the LangGraph Server API.
|
||||
|
||||
### Dependency map
|
||||
|
||||
The diagram below lists downstream libraries for each production dependency as
|
||||
declared in that library's `pyproject.toml` (or `package.json`).
|
||||
|
||||
```text
|
||||
checkpoint
|
||||
├── checkpoint-postgres
|
||||
├── checkpoint-sqlite
|
||||
├── prebuilt
|
||||
└── langgraph
|
||||
|
||||
prebuilt
|
||||
└── langgraph
|
||||
|
||||
sdk-py
|
||||
├── langgraph
|
||||
└── cli
|
||||
|
||||
sdk-js (standalone)
|
||||
```
|
||||
|
||||
Changes to a library may impact all of its dependents shown above.
|
||||
@@ -1,55 +0,0 @@
|
||||
# AGENTS Instructions
|
||||
|
||||
This repository is a monorepo. Each library lives in a subdirectory under `libs/`.
|
||||
|
||||
When you modify code in any library, run the following commands in that library's directory before creating a pull request:
|
||||
|
||||
- `make format` – run code formatters
|
||||
- `make lint` – run the linter
|
||||
- `make test` – execute the test suite
|
||||
|
||||
To run a particular test file or to pass additional pytest options you can specify the `TEST` variable:
|
||||
|
||||
```
|
||||
TEST=path/to/test.py make test
|
||||
```
|
||||
|
||||
Other pytest arguments can also be supplied inside the `TEST` variable.
|
||||
|
||||
## Libraries
|
||||
|
||||
The repository contains several Python and JavaScript/TypeScript libraries.
|
||||
Below is a high-level overview:
|
||||
|
||||
- **checkpoint** – base interfaces for LangGraph checkpointers.
|
||||
- **checkpoint-postgres** – Postgres implementation of the checkpoint saver.
|
||||
- **checkpoint-sqlite** – SQLite implementation of the checkpoint saver.
|
||||
- **cli** – official command-line interface for LangGraph.
|
||||
- **langgraph** – core framework for building stateful, multi-actor agents.
|
||||
- **prebuilt** – high-level APIs for creating and running agents and tools.
|
||||
- **sdk-js** – JS/TS SDK for interacting with the LangGraph REST API.
|
||||
- **sdk-py** – Python SDK for the LangGraph Server API.
|
||||
|
||||
### Dependency map
|
||||
|
||||
The diagram below lists downstream libraries for each production dependency as
|
||||
declared in that library's `pyproject.toml` (or `package.json`).
|
||||
|
||||
```text
|
||||
checkpoint
|
||||
├── checkpoint-postgres
|
||||
├── checkpoint-sqlite
|
||||
├── prebuilt
|
||||
└── langgraph
|
||||
|
||||
prebuilt
|
||||
└── langgraph
|
||||
|
||||
sdk-py
|
||||
├── langgraph
|
||||
└── cli
|
||||
|
||||
sdk-js (standalone)
|
||||
```
|
||||
|
||||
Changes to a library may impact all of its dependents shown above.
|
||||
+293
@@ -0,0 +1,293 @@
|
||||
# Contributing to LangGraph
|
||||
|
||||
Thank you for being interested in contributing to LangGraph!
|
||||
|
||||
## General guidelines
|
||||
|
||||
Here are some things to keep in mind for all types of contributions:
|
||||
|
||||
- Follow the ["fork and pull request"](https://docs.github.com/en/get-started/exploring-projects-on-github/contributing-to-a-project) workflow.
|
||||
- Fill out the checked-in pull request template when opening pull requests. Note related issues and tag relevant maintainers.
|
||||
- Ensure your PR passes formatting, linting, and testing checks before requesting a review.
|
||||
- If you would like comments or feedback, please open an issue or discussion and tag a maintainer.
|
||||
- Backwards compatibility is key. Your changes must not be breaking, except in case of critical bug and security fixes.
|
||||
- Look for duplicate PRs or issues that have already been opened before opening a new one.
|
||||
- Keep scope as isolated as possible. As a general rule, your changes should not affect more than one package at a time.
|
||||
|
||||
### Bugfixes
|
||||
|
||||
For bug fixes, please open up an issue before proposing a fix to ensure the proposal properly addresses the underlying problem. In general, bug fixes should all have an accompanying unit test that fails before the fix.
|
||||
|
||||
### New features
|
||||
|
||||
For new features, please start a new [discussion](https://github.com/langchain-ai/langgraph/discussions), where the maintainers will help with scoping out the necessary changes.
|
||||
|
||||
## Contribute Documentation
|
||||
|
||||
Documentation is a vital part of LangGraph. We welcome both new documentation for new features and
|
||||
community improvements to our current documentation. Please read the resources below before getting started:
|
||||
|
||||
- [Documentation style guide](#documentation-style-guide)
|
||||
- [Documentation setup](#setup)
|
||||
|
||||
## Documentation Style Guide
|
||||
|
||||
As LangGraph continues to grow, the surface area of documentation required to cover it continues to grow too.
|
||||
This page provides guidelines for anyone writing documentation for LangGraph, as well as some of our philosophies around organization and structure.
|
||||
|
||||
## Philosophy
|
||||
|
||||
LangGraph's documentation follows the [Diataxis framework](https://diataxis.fr).
|
||||
Under this framework, all documentation falls under one of four categories: [Tutorials](#tutorials),
|
||||
[How-to guides](#how-to-guides),
|
||||
[References](#references), and [Explanations (aka conceptual guides)](#conceptual-guide).
|
||||
|
||||
### Tutorials
|
||||
|
||||
Tutorials are lessons that take the reader through a practical activity. Their purpose is to help the user
|
||||
gain understanding of concepts and how they interact by showing one way to achieve some goal in a hands-on way.
|
||||
|
||||
They should **avoid** giving
|
||||
multiple permutations of ways to achieve that goal in-depth. Choice is burdensome. Instead, they should guide a new user through a recommended path to accomplishing a concrete goal. While the end result of a tutorial does not necessarily need to
|
||||
be completely production-ready, it should be useful and practically satisfy the goal that you clearly stated in the tutorial's introduction.
|
||||
|
||||
To quote the Diataxis website:
|
||||
|
||||
> A tutorial serves the user’s *acquisition* of skills and knowledge - their study. Its purpose is not to help the user get something done, but to help them learn.
|
||||
|
||||
In LangGraph, these are often higher level guides that show off end-to-end use cases.
|
||||
|
||||
Some examples include:
|
||||
|
||||
- [Build a Customer Support Bot](https://langchain-ai.github.io/langgraph/tutorials/customer-support/customer-support/)
|
||||
- [Build a SQL Agent](https://langchain-ai.github.io/langgraph/tutorials/sql-agent/)
|
||||
|
||||
Here are some high-level tips on writing a good tutorial:
|
||||
|
||||
- Focus on guiding the user to get something done, but keep in mind the end-goal is more to impart principles than to create a perfect production system.
|
||||
- Be specific, not abstract and follow one path.
|
||||
- No need to go deeply into alternative approaches, but it’s ok to reference them, ideally with a link to an appropriate how-to guide.
|
||||
- Get "a point on the board" as soon as possible - something the user can run that outputs something.
|
||||
- You can iterate and expand afterwards.
|
||||
- Try to frequently checkpoint at given steps where the user can run code and see progress.
|
||||
- Focus on results, not technical explanation.
|
||||
- Crosslink heavily to appropriate conceptual/reference pages
|
||||
- The first time you mention a LangGraph concept, use its full name (e.g. "human-in-the-loop"), and link to its conceptual/other documentation page.
|
||||
- It's also helpful to add a prerequisite callout that links to any pages with necessary background information.
|
||||
- End with a recap/next steps section summarizing what the tutorial covered and future reading, such as related how-to guides.
|
||||
- Use phrases like "Next we can run X & Y. We will expect Z.". Then afterwards, use language like "Notice Z" that recalls our expectations and directs the reader's attention to the topic we are trying to teach.
|
||||
- Do not shy away from repetition.
|
||||
|
||||
### How-to guides
|
||||
|
||||
A how-to guide, as the name implies, demonstrates how to do something discrete and specific.
|
||||
It should assume that the user is already familiar with underlying concepts, and is trying to solve an immediate problem, but
|
||||
should still give some background or list the scenarios where the information contained within can be relevant.
|
||||
They can and should discuss alternatives if one approach may be better than another in certain cases.
|
||||
|
||||
To quote the Diataxis website:
|
||||
|
||||
> A how-to guide serves the work of the already-competent user, whom you can assume to know what they want to do, and to be able to follow your instructions correctly.
|
||||
|
||||
Some examples include:
|
||||
|
||||
- [How to add persistence to your graph](https://langchain-ai.github.io/langgraph/how-tos/persistence/)
|
||||
- [How to view and update past graph state](https://langchain-ai.github.io/langgraph/how-tos/human_in_the_loop/time-travel/)
|
||||
|
||||
Here are some high-level tips on writing a good how-to guide:
|
||||
|
||||
- Clearly explain what you are guiding the user through at the start
|
||||
- Assume higher intent than a tutorial and show what the user needs to do to get that task done
|
||||
- Assume familiarity of concepts, but explain why suggested actions are helpful
|
||||
- Crosslink heavily to conceptual/reference pages
|
||||
- Discuss alternatives and responses to real-world tradeoffs that may arise when solving a problem
|
||||
- Use lots of example code, ideally within complete code blocks that the reader can copy and run.
|
||||
- End with a recap/next steps section summarizing what the tutorial covered and future reading, such as other related how-to guides
|
||||
|
||||
### Conceptual guides
|
||||
|
||||
LangGraph's conceptual guides fall under the **Explanation** quadrant of Diataxis. They should cover LangChain terms and concepts
|
||||
in a more abstract way than how-to guides or tutorials, and should be geared towards curious users interested in
|
||||
gaining a deeper understanding of the framework. Try to avoid excessively large code examples. The goal here is to
|
||||
impart perspective to the user rather than to finish a practical project. These guides should cover **why** things work they way they do.
|
||||
|
||||
|
||||
To quote the Diataxis website:
|
||||
|
||||
> The perspective of explanation is higher and wider than that of the other types. It does not take the user’s eye-level view, as in a how-to guide, or a close-up view of the machinery, like reference material. Its scope in each case is a topic - “an area of knowledge”, that somehow has to be bounded in a reasonable, meaningful way.
|
||||
|
||||
Some examples include:
|
||||
|
||||
- [What does it mean to be agentic?](https://langchain-ai.github.io/langgraph/concepts/high_level/)
|
||||
- [Tool calling](https://langchain-ai.github.io/langgraph/concepts/agentic_concepts/#tool-calling)
|
||||
|
||||
Here are some high-level tips on writing a good conceptual guide:
|
||||
|
||||
- Explain design decisions. Why does concept X exist and why was it designed this way?
|
||||
- Use analogies and reference other concepts and alternatives
|
||||
- Avoid blending in too much reference content
|
||||
- You can and should reference content covered in other guides, but make sure to link to them
|
||||
|
||||
### References
|
||||
|
||||
References contain detailed, low-level information that describes exactly what functionality exists and how to use it.
|
||||
In LangGraph, this is mainly our API reference pages, which are populated from docstrings within code.
|
||||
References pages are generally not read end-to-end, but are consulted as necessary when a user needs to know
|
||||
how to use something specific.
|
||||
|
||||
To quote the Diataxis website:
|
||||
|
||||
> The only purpose of a reference guide is to describe, as succinctly as possible, and in an orderly way. Whereas the content of tutorials and how-to guides are led by needs of the user, reference material is led by the product it describes.
|
||||
|
||||
Many of the reference pages in LangChain are automatically generated from code,
|
||||
but here are some high-level tips on writing a good docstring:
|
||||
|
||||
- Be concise
|
||||
- Discuss special cases and deviations from a user's expectations
|
||||
- Go into detail on required inputs and outputs
|
||||
- Light details on when one might use the feature are fine, but in-depth details belong in other sections.
|
||||
|
||||
Each category serves a distinct purpose and requires a specific approach to writing and structuring the content.
|
||||
|
||||
## General guidelines
|
||||
|
||||
Here are some other guidelines you should think about when writing and organizing documentation.
|
||||
|
||||
We generally do not merge new tutorials from outside contributors without an actue need.
|
||||
We welcome updates as well as new integration docs, how-tos, and references.
|
||||
|
||||
### Avoid duplication
|
||||
|
||||
Multiple pages that cover the same material in depth are difficult to maintain and cause confusion. There should
|
||||
be only one (very rarely two), canonical pages for a given concept or feature. Instead, you should link to other guides.
|
||||
|
||||
### Link to other sections
|
||||
|
||||
Because sections of the docs do not exist in a vacuum, it is important to link to other sections as often as possible
|
||||
to allow a developer to learn more about an unfamiliar topic inline.
|
||||
|
||||
This includes linking to the API references as well as conceptual sections!
|
||||
|
||||
### Be concise
|
||||
|
||||
In general, take a less-is-more approach. If a section with a good explanation of a concept already exists, you should link to it rather than
|
||||
re-explain it, unless the concept you are documenting presents some new wrinkle.
|
||||
|
||||
Be concise, including in code samples.
|
||||
|
||||
### General style
|
||||
|
||||
- Use active voice and present tense whenever possible
|
||||
- Use examples and code snippets to illustrate concepts and usage
|
||||
- Use appropriate header levels (`#`, `##`, `###`, etc.) to organize the content hierarchically
|
||||
- Use fewer cells with more code to make copy/paste easier
|
||||
- Use bullet points and numbered lists to break down information into easily digestible chunks
|
||||
- Use tables (especially for **Reference** sections) and diagrams often to present information visually
|
||||
- Include the table of contents for longer documentation pages to help readers navigate the content, but hide it for shorter pages
|
||||
|
||||
## Setup
|
||||
|
||||
LangChain documentation consists of two components:
|
||||
|
||||
1. Main Documentation: Hosted at [https://langchain-ai.github.io](https://langchain-ai.github.io/langgraph/),
|
||||
this comprehensive resource serves as the primary user-facing documentation.
|
||||
It covers a wide array of topics, including tutorials, use cases, integrations,
|
||||
and more, offering extensive guidance on building with LangGraph.
|
||||
The content for this documentation lives in the `/docs` directory of the monorepo.
|
||||
2. In-code Documentation: This is documentation of the codebase itself, which is also
|
||||
used to generate the externally facing [API Reference](https://langchain-ai.github.io/langgraph/reference/graphs/).
|
||||
The content for the API reference is autogenerated by scanning the docstrings in the codebase. For this reason we ask that developers document their code well.
|
||||
|
||||
We appreciate all contributions to the documentation, whether it be fixing a typo,
|
||||
adding a new tutorial or example and whether it be in the main documentation or the API Reference.
|
||||
|
||||
### 📜 Main Documentation
|
||||
|
||||
The content for the main documentation is located in the `/docs` directory of the monorepo.
|
||||
|
||||
The documentation is written using a combination of ipython notebooks (`.ipynb` files)
|
||||
and markdown (`.md` files). The notebooks are converted to markdown
|
||||
and then built using [MkDocs](https://www.mkdocs.org/).
|
||||
|
||||
Feel free to make contributions to the main documentation! 🥰
|
||||
|
||||
After modifying the documentation:
|
||||
|
||||
1. Run the linting and formatting commands (see below) to ensure that the documentation is well-formatted and free of errors.
|
||||
2. Optionally build the documentation locally to verify that the changes look good.
|
||||
3. Make a pull request with the changes.
|
||||
|
||||
### ⚒️ Linting and Building Documentation Locally
|
||||
|
||||
After writing up the documentation, you may want to lint and build the documentation
|
||||
locally to ensure that it looks good and is free of errors.
|
||||
|
||||
If you're unable to build it locally that's okay as well, as you will be able to
|
||||
see a preview of the documentation on the pull request page.
|
||||
|
||||
From the **monorepo root**, run the following command to install the dependencies:
|
||||
|
||||
```bash
|
||||
poetry install --with docs --no-root
|
||||
```
|
||||
|
||||
#### Building
|
||||
|
||||
The code that builds the documentation is located in the `/docs` directory of the monorepo.
|
||||
|
||||
Before building the documentation, it is always a good idea to clean the build directory:
|
||||
|
||||
```bash
|
||||
make clean-docs
|
||||
```
|
||||
|
||||
You can build and preview the documentation as outlined below:
|
||||
|
||||
```bash
|
||||
make serve-docs
|
||||
```
|
||||
|
||||
#### Linting
|
||||
|
||||
The documentation is linted from the **monorepo root**. To lint it, run the following from there:
|
||||
|
||||
```bash
|
||||
make spellcheck
|
||||
```
|
||||
|
||||
### ️In-code Documentation
|
||||
|
||||
The in-code documentation is autogenerated from docstrings.
|
||||
|
||||
For the API reference to be useful, the codebase must be well-documented. This means that all functions, classes, and methods should have a docstring that explains what they do, what the arguments are, and what the return value is. This is a good practice in general, but it is especially important for LangChain because the API reference is the primary resource for developers to understand how to use the codebase.
|
||||
|
||||
We generally follow the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) for docstrings.
|
||||
|
||||
Here is an example of a well-documented function:
|
||||
|
||||
```python
|
||||
|
||||
def my_function(arg1: int, arg2: str) -> float:
|
||||
"""This is a short description of the function. (It should be a single sentence.)
|
||||
|
||||
This is a longer description of the function. It should explain what
|
||||
the function does, what the arguments are, and what the return value is.
|
||||
It should wrap at 88 characters.
|
||||
|
||||
Examples:
|
||||
This is a section for examples of how to use the function.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
my_function(1, "hello")
|
||||
|
||||
Args:
|
||||
arg1: This is a description of arg1. We do not need to specify the type since
|
||||
it is already specified in the function signature.
|
||||
arg2: This is a description of arg2.
|
||||
|
||||
Returns:
|
||||
This is a description of the return value.
|
||||
"""
|
||||
return 3.14
|
||||
```
|
||||
@@ -1,68 +0,0 @@
|
||||
# Define the directories containing projects
|
||||
LIBS_DIRS := $(wildcard libs/*)
|
||||
|
||||
# Default target
|
||||
.PHONY: all
|
||||
all: lint format lock test
|
||||
|
||||
# Install dependencies for all projects
|
||||
.PHONY: install
|
||||
install:
|
||||
@echo "Creating virtual environment..."
|
||||
@uv venv
|
||||
@for dir in $(LIBS_DIRS); do \
|
||||
if [ -f $$dir/pyproject.toml ]; then \
|
||||
echo "Installing dependencies for $$dir"; \
|
||||
uv pip install -e $$dir; \
|
||||
fi; \
|
||||
done
|
||||
|
||||
# Lint all projects
|
||||
.PHONY: lint
|
||||
lint:
|
||||
@for dir in $(LIBS_DIRS); do \
|
||||
if [ -f $$dir/Makefile ]; then \
|
||||
echo "Running lint in $$dir"; \
|
||||
$(MAKE) -C $$dir lint; \
|
||||
fi; \
|
||||
done
|
||||
|
||||
# Format all projects
|
||||
.PHONY: format
|
||||
format:
|
||||
@for dir in $(LIBS_DIRS); do \
|
||||
if [ -f $$dir/Makefile ]; then \
|
||||
echo "Running format in $$dir"; \
|
||||
$(MAKE) -C $$dir format; \
|
||||
fi; \
|
||||
done
|
||||
|
||||
# Lock all projects
|
||||
.PHONY: lock
|
||||
lock:
|
||||
@for dir in $(LIBS_DIRS); do \
|
||||
if [ -f $$dir/Makefile ]; then \
|
||||
echo "Running lock in $$dir"; \
|
||||
(cd $$dir && uv lock); \
|
||||
fi; \
|
||||
done
|
||||
|
||||
# Lock all projects and upgrade dependencies
|
||||
.PHONY: lock-upgrade
|
||||
lock-upgrade:
|
||||
@for dir in $(LIBS_DIRS); do \
|
||||
if [ -f $$dir/Makefile ]; then \
|
||||
echo "Running lock-upgrade in $$dir"; \
|
||||
(cd $$dir && uv lock --upgrade); \
|
||||
fi; \
|
||||
done
|
||||
|
||||
# Test all projects
|
||||
.PHONY: test
|
||||
test:
|
||||
@for dir in $(LIBS_DIRS); do \
|
||||
if [ -f $$dir/Makefile ]; then \
|
||||
echo "Running test in $$dir"; \
|
||||
$(MAKE) -C $$dir test; \
|
||||
fi; \
|
||||
done
|
||||
@@ -1,91 +1,339 @@
|
||||
<picture class="github-only">
|
||||
<source media="(prefers-color-scheme: light)" srcset="https://langchain-ai.github.io/langgraph/static/wordmark_dark.svg">
|
||||
<source media="(prefers-color-scheme: dark)" srcset="https://langchain-ai.github.io/langgraph/static/wordmark_light.svg">
|
||||
<img alt="LangGraph Logo" src="https://langchain-ai.github.io/langgraph/static/wordmark_dark.svg" width="80%">
|
||||
</picture>
|
||||
# 🦜🕸️LangGraph
|
||||
|
||||
<div>
|
||||
<br>
|
||||
</div>
|
||||
|
||||
[](https://pypi.org/project/langgraph/)
|
||||

|
||||
[](https://pepy.tech/project/langgraph)
|
||||
[](https://github.com/langchain-ai/langgraph/issues)
|
||||
[](https://docs.langchain.com/oss/python/langgraph/overview)
|
||||
[](https://langchain-ai.github.io/langgraph/)
|
||||
|
||||
Trusted by companies shaping the future of agents – including Klarna, Replit, Elastic, and more – LangGraph is a low-level orchestration framework for building, managing, and deploying long-running, stateful agents.
|
||||
⚡ Building language agents as graphs ⚡
|
||||
|
||||
## Get started
|
||||
> [!NOTE]
|
||||
> Looking for the JS version? See the [JS repo](https://github.com/langchain-ai/langgraphjs) and the [JS docs](https://langchain-ai.github.io/langgraphjs/).
|
||||
|
||||
Install LangGraph:
|
||||
## Overview
|
||||
|
||||
```
|
||||
[LangGraph](https://langchain-ai.github.io/langgraph/) is a library for building
|
||||
stateful, multi-actor applications with LLMs, used to create agent and multi-agent
|
||||
workflows. Check out an introductory tutorial [here](https://langchain-ai.github.io/langgraph/tutorials/introduction/).
|
||||
|
||||
|
||||
LangGraph is inspired by [Pregel](https://research.google/pubs/pub37252/) and [Apache Beam](https://beam.apache.org/). The public interface draws inspiration from [NetworkX](https://networkx.org/documentation/latest/). LangGraph is built by LangChain Inc, the creators of LangChain, but can be used without LangChain.
|
||||
|
||||
### Why use LangGraph?
|
||||
|
||||
LangGraph powers [production-grade agents](https://www.langchain.com/built-with-langgraph), trusted by Linkedin, Uber, Klarna, GitLab, and many more. LangGraph provides fine-grained control over both the flow and state of your agent applications. It implements a central [persistence layer](https://langchain-ai.github.io/langgraph/concepts/persistence/), enabling features that are common to most agent architectures:
|
||||
|
||||
- **Memory**: LangGraph persists arbitrary aspects of your application's state,
|
||||
supporting memory of conversations and other updates within and across user
|
||||
interactions;
|
||||
- **Human-in-the-loop**: Because state is checkpointed, execution can be interrupted
|
||||
and resumed, allowing for decisions, validation, and corrections at key stages via
|
||||
human input.
|
||||
|
||||
Standardizing these components allows individuals and teams to focus on the behavior
|
||||
of their agent, instead of its supporting infrastructure.
|
||||
|
||||
Through [LangGraph Platform](#langgraph-platform), LangGraph also provides tooling for
|
||||
the development, deployment, debugging, and monitoring of your applications.
|
||||
|
||||
LangGraph integrates seamlessly with
|
||||
[LangChain](https://python.langchain.com/docs/introduction/) and
|
||||
[LangSmith](https://docs.smith.langchain.com/) (but does not require them).
|
||||
|
||||
To learn more about LangGraph, check out our first LangChain Academy
|
||||
course, *Introduction to LangGraph*, available for free
|
||||
[here](https://academy.langchain.com/courses/intro-to-langgraph).
|
||||
|
||||
### LangGraph Platform
|
||||
|
||||
[LangGraph Platform](https://langchain-ai.github.io/langgraph/concepts/langgraph_platform) is infrastructure for deploying LangGraph agents. It is a commercial solution for deploying agentic applications to production, built on the open-source LangGraph framework. The LangGraph Platform consists of several components that work together to support the development, deployment, debugging, and monitoring of LangGraph applications: [LangGraph Server](https://langchain-ai.github.io/langgraph/concepts/langgraph_server) (APIs), [LangGraph SDKs](https://langchain-ai.github.io/langgraph/concepts/sdk) (clients for the APIs), [LangGraph CLI](https://langchain-ai.github.io/langgraph/concepts/langgraph_cli) (command line tool for building the server), and [LangGraph Studio](https://langchain-ai.github.io/langgraph/concepts/langgraph_studio) (UI/debugger).
|
||||
|
||||
See deployment options [here](https://langchain-ai.github.io/langgraph/concepts/deployment_options/)
|
||||
(includes a free tier).
|
||||
|
||||
Here are some common issues that arise in complex deployments, which LangGraph Platform addresses:
|
||||
|
||||
- **Streaming support**: LangGraph Server provides [multiple streaming modes](https://langchain-ai.github.io/langgraph/concepts/streaming) optimized for various application needs
|
||||
- **Background runs**: Runs agents asynchronously in the background
|
||||
- **Support for long running agents**: Infrastructure that can handle long running processes
|
||||
- **[Double texting](https://langchain-ai.github.io/langgraph/concepts/double_texting)**: Handle the case where you get two messages from the user before the agent can respond
|
||||
- **Handle burstiness**: Task queue for ensuring requests are handled consistently without loss, even under heavy loads
|
||||
|
||||
## Installation
|
||||
|
||||
```shell
|
||||
pip install -U langgraph
|
||||
```
|
||||
|
||||
Create a simple workflow:
|
||||
## Example
|
||||
|
||||
```python
|
||||
from langgraph.graph import START, StateGraph
|
||||
from typing_extensions import TypedDict
|
||||
Let's build a tool-calling [ReAct-style](https://langchain-ai.github.io/langgraph/concepts/agentic_concepts/#react-implementation) agent that uses a search tool!
|
||||
|
||||
|
||||
class State(TypedDict):
|
||||
text: str
|
||||
|
||||
|
||||
def node_a(state: State) -> dict:
|
||||
return {"text": state["text"] + "a"}
|
||||
|
||||
|
||||
def node_b(state: State) -> dict:
|
||||
return {"text": state["text"] + "b"}
|
||||
|
||||
|
||||
graph = StateGraph(State)
|
||||
graph.add_node("node_a", node_a)
|
||||
graph.add_node("node_b", node_b)
|
||||
graph.add_edge(START, "node_a")
|
||||
graph.add_edge("node_a", "node_b")
|
||||
|
||||
print(graph.compile().invoke({"text": ""}))
|
||||
# {'text': 'ab'}
|
||||
```shell
|
||||
pip install langchain-anthropic
|
||||
```
|
||||
|
||||
Get started with the [LangGraph Quickstart](https://docs.langchain.com/oss/python/langgraph/quickstart).
|
||||
```shell
|
||||
export ANTHROPIC_API_KEY=sk-...
|
||||
```
|
||||
|
||||
To quickly build agents with LangChain's `create_agent` (built on LangGraph), see the [LangChain Agents documentation](https://docs.langchain.com/oss/python/langchain/agents).
|
||||
Optionally, we can set up [LangSmith](https://docs.smith.langchain.com/) for best-in-class observability.
|
||||
|
||||
## Core benefits
|
||||
```shell
|
||||
export LANGSMITH_TRACING=true
|
||||
export LANGSMITH_API_KEY=lsv2_sk_...
|
||||
```
|
||||
|
||||
LangGraph provides low-level supporting infrastructure for *any* long-running, stateful workflow or agent. LangGraph does not abstract prompts or architecture, and provides the following central benefits:
|
||||
The simplest way to create a tool-calling agent in LangGraph is to use `create_react_agent`:
|
||||
|
||||
- [Durable execution](https://docs.langchain.com/oss/python/langgraph/durable-execution): Build agents that persist through failures and can run for extended periods, automatically resuming from exactly where they left off.
|
||||
- [Human-in-the-loop](https://docs.langchain.com/oss/python/langgraph/interrupts): Seamlessly incorporate human oversight by inspecting and modifying agent state at any point during execution.
|
||||
- [Comprehensive memory](https://docs.langchain.com/oss/python/langgraph/memory): Create truly stateful agents with both short-term working memory for ongoing reasoning and long-term persistent memory across sessions.
|
||||
- [Debugging with LangSmith](http://www.langchain.com/langsmith): Gain deep visibility into complex agent behavior with visualization tools that trace execution paths, capture state transitions, and provide detailed runtime metrics.
|
||||
- [Production-ready deployment](https://docs.langchain.com/langsmith/app-development): Deploy sophisticated agent systems confidently with scalable infrastructure designed to handle the unique challenges of stateful, long-running workflows.
|
||||
<details open>
|
||||
<summary>High-level implementation</summary>
|
||||
|
||||
## LangGraph’s ecosystem
|
||||
```python
|
||||
from langgraph.prebuilt import create_react_agent
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
from langchain_anthropic import ChatAnthropic
|
||||
from langchain_core.tools import tool
|
||||
|
||||
While LangGraph can be used standalone, it also integrates seamlessly with any LangChain product, giving developers a full suite of tools for building agents. To improve your LLM application development, pair LangGraph with:
|
||||
# Define the tools for the agent to use
|
||||
@tool
|
||||
def search(query: str):
|
||||
"""Call to surf the web."""
|
||||
# This is a placeholder, but don't tell the LLM that...
|
||||
if "sf" in query.lower() or "san francisco" in query.lower():
|
||||
return "It's 60 degrees and foggy."
|
||||
return "It's 90 degrees and sunny."
|
||||
|
||||
- [LangSmith](http://www.langchain.com/langsmith) — Helpful for agent evals and observability. Debug poor-performing LLM app runs, evaluate agent trajectories, gain visibility in production, and improve performance over time.
|
||||
- [LangSmith Deployment](https://docs.langchain.com/langsmith/deployments) — Deploy and scale agents effortlessly with a purpose-built deployment platform for long running, stateful workflows. Discover, reuse, configure, and share agents across teams — and iterate quickly with visual prototyping in [LangGraph Studio](https://docs.langchain.com/oss/python/langgraph/studio).
|
||||
- [LangChain](https://docs.langchain.com/oss/python/langchain/overview) – Provides integrations and composable components to streamline LLM application development.
|
||||
|
||||
> [!NOTE]
|
||||
> Looking for the JS version of LangGraph? See the [JS repo](https://github.com/langchain-ai/langgraphjs) and the [JS docs](https://docs.langchain.com/oss/javascript/langgraph/overview).
|
||||
tools = [search]
|
||||
model = ChatAnthropic(model="claude-3-5-sonnet-latest", temperature=0)
|
||||
|
||||
## Additional resources
|
||||
# Initialize memory to persist state between graph runs
|
||||
checkpointer = MemorySaver()
|
||||
|
||||
- [Guides](https://docs.langchain.com/oss/python/langgraph/overview): Quick, actionable code snippets for topics such as streaming, adding memory & persistence, and design patterns (e.g. branching, subgraphs, etc.).
|
||||
- [Reference](https://reference.langchain.com/python/langgraph/): Detailed reference on core classes, methods, how to use the graph and checkpointing APIs, and higher-level prebuilt components.
|
||||
- [Examples](https://docs.langchain.com/oss/python/langgraph/agentic-rag): Guided examples on getting started with LangGraph.
|
||||
- [LangChain Forum](https://forum.langchain.com/): Connect with the community and share all of your technical questions, ideas, and feedback.
|
||||
- [LangChain Academy](https://academy.langchain.com/courses/intro-to-langgraph): Learn the basics of LangGraph in our free, structured course.
|
||||
- [Case studies](https://www.langchain.com/built-with-langgraph): Hear how industry leaders use LangGraph to ship AI applications at scale.
|
||||
app = create_react_agent(model, tools, checkpointer=checkpointer)
|
||||
|
||||
## Acknowledgements
|
||||
# Use the agent
|
||||
final_state = app.invoke(
|
||||
{"messages": [{"role": "user", "content": "what is the weather in sf"}]},
|
||||
config={"configurable": {"thread_id": 42}}
|
||||
)
|
||||
final_state["messages"][-1].content
|
||||
```
|
||||
```
|
||||
"Based on the search results, I can tell you that the current weather in San Francisco is:\n\nTemperature: 60 degrees Fahrenheit\nConditions: Foggy\n\nSan Francisco is known for its microclimates and frequent fog, especially during the summer months. The temperature of 60°F (about 15.5°C) is quite typical for the city, which tends to have mild temperatures year-round. The fog, often referred to as "Karl the Fog" by locals, is a characteristic feature of San Francisco\'s weather, particularly in the mornings and evenings.\n\nIs there anything else you\'d like to know about the weather in San Francisco or any other location?"
|
||||
```
|
||||
|
||||
LangGraph is inspired by [Pregel](https://research.google/pubs/pub37252/) and [Apache Beam](https://beam.apache.org/). The public interface draws inspiration from [NetworkX](https://networkx.org/documentation/latest/). LangGraph is built by LangChain Inc, the creators of LangChain, but can be used without LangChain.
|
||||
Now when we pass the same <code>"thread_id"</code>, the conversation context is retained via the saved state (i.e. stored list of messages)
|
||||
|
||||
```python
|
||||
final_state = app.invoke(
|
||||
{"messages": [{"role": "user", "content": "what about ny"}]},
|
||||
config={"configurable": {"thread_id": 42}}
|
||||
)
|
||||
final_state["messages"][-1].content
|
||||
```
|
||||
|
||||
```
|
||||
"Based on the search results, I can tell you that the current weather in New York City is:\n\nTemperature: 90 degrees Fahrenheit (approximately 32.2 degrees Celsius)\nConditions: Sunny\n\nThis weather is quite different from what we just saw in San Francisco. New York is experiencing much warmer temperatures right now. Here are a few points to note:\n\n1. The temperature of 90°F is quite hot, typical of summer weather in New York City.\n2. The sunny conditions suggest clear skies, which is great for outdoor activities but also means it might feel even hotter due to direct sunlight.\n3. This kind of weather in New York often comes with high humidity, which can make it feel even warmer than the actual temperature suggests.\n\nIt's interesting to see the stark contrast between San Francisco's mild, foggy weather and New York's hot, sunny conditions. This difference illustrates how varied weather can be across different parts of the United States, even on the same day.\n\nIs there anything else you'd like to know about the weather in New York or any other location?"
|
||||
```
|
||||
</details>
|
||||
|
||||
> [!TIP]
|
||||
> LangGraph is a **low-level** framework that allows you to implement any custom agent
|
||||
architectures. Click on the low-level implementation below to see how to implement a
|
||||
tool-calling agent from scratch.
|
||||
|
||||
<details>
|
||||
<summary>Low-level implementation</summary>
|
||||
|
||||
```python
|
||||
from typing import Literal
|
||||
|
||||
from langchain_anthropic import ChatAnthropic
|
||||
from langchain_core.tools import tool
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
from langgraph.graph import END, START, StateGraph, MessagesState
|
||||
from langgraph.prebuilt import ToolNode
|
||||
|
||||
|
||||
# Define the tools for the agent to use
|
||||
@tool
|
||||
def search(query: str):
|
||||
"""Call to surf the web."""
|
||||
# This is a placeholder, but don't tell the LLM that...
|
||||
if "sf" in query.lower() or "san francisco" in query.lower():
|
||||
return "It's 60 degrees and foggy."
|
||||
return "It's 90 degrees and sunny."
|
||||
|
||||
|
||||
tools = [search]
|
||||
|
||||
tool_node = ToolNode(tools)
|
||||
|
||||
model = ChatAnthropic(model="claude-3-5-sonnet-latest", temperature=0).bind_tools(tools)
|
||||
|
||||
# Define the function that determines whether to continue or not
|
||||
def should_continue(state: MessagesState) -> Literal["tools", END]:
|
||||
messages = state['messages']
|
||||
last_message = messages[-1]
|
||||
# If the LLM makes a tool call, then we route to the "tools" node
|
||||
if last_message.tool_calls:
|
||||
return "tools"
|
||||
# Otherwise, we stop (reply to the user)
|
||||
return END
|
||||
|
||||
|
||||
# Define the function that calls the model
|
||||
def call_model(state: MessagesState):
|
||||
messages = state['messages']
|
||||
response = model.invoke(messages)
|
||||
# We return a list, because this will get added to the existing list
|
||||
return {"messages": [response]}
|
||||
|
||||
|
||||
# Define a new graph
|
||||
workflow = StateGraph(MessagesState)
|
||||
|
||||
# Define the two nodes we will cycle between
|
||||
workflow.add_node("agent", call_model)
|
||||
workflow.add_node("tools", tool_node)
|
||||
|
||||
# Set the entrypoint as `agent`
|
||||
# This means that this node is the first one called
|
||||
workflow.add_edge(START, "agent")
|
||||
|
||||
# We now add a conditional edge
|
||||
workflow.add_conditional_edges(
|
||||
# First, we define the start node. We use `agent`.
|
||||
# This means these are the edges taken after the `agent` node is called.
|
||||
"agent",
|
||||
# Next, we pass in the function that will determine which node is called next.
|
||||
should_continue,
|
||||
)
|
||||
|
||||
# We now add a normal edge from `tools` to `agent`.
|
||||
# This means that after `tools` is called, `agent` node is called next.
|
||||
workflow.add_edge("tools", 'agent')
|
||||
|
||||
# Initialize memory to persist state between graph runs
|
||||
checkpointer = MemorySaver()
|
||||
|
||||
# Finally, we compile it!
|
||||
# This compiles it into a LangChain Runnable,
|
||||
# meaning you can use it as you would any other runnable.
|
||||
# Note that we're (optionally) passing the memory when compiling the graph
|
||||
app = workflow.compile(checkpointer=checkpointer)
|
||||
|
||||
# Use the agent
|
||||
final_state = app.invoke(
|
||||
{"messages": [{"role": "user", "content": "what is the weather in sf"}]},
|
||||
config={"configurable": {"thread_id": 42}}
|
||||
)
|
||||
final_state["messages"][-1].content
|
||||
```
|
||||
|
||||
<b>Step-by-step Breakdown</b>:
|
||||
|
||||
<details>
|
||||
<summary>Initialize the model and tools.</summary>
|
||||
<ul>
|
||||
<li>
|
||||
We use <code>ChatAnthropic</code> as our LLM. <strong>NOTE:</strong> we need to make sure the model knows that it has these tools available to call. We can do this by converting the LangChain tools into the format for OpenAI tool calling using the <code>.bind_tools()</code> method.
|
||||
</li>
|
||||
<li>
|
||||
We define the tools we want to use - a search tool in our case. It is really easy to create your own tools - see documentation here on how to do that <a href="https://python.langchain.com/docs/how_to/custom_tools/">here</a>.
|
||||
</li>
|
||||
</ul>
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Initialize graph with state.</summary>
|
||||
|
||||
<ul>
|
||||
<li>We initialize graph (<code>StateGraph</code>) by passing state schema (in our case <code>MessagesState</code>)</li>
|
||||
<li><code>MessagesState</code> is a prebuilt state schema that has one attribute -- a list of LangChain <code>Message</code> objects, as well as logic for merging the updates from each node into the state.</li>
|
||||
</ul>
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Define graph nodes.</summary>
|
||||
|
||||
There are two main nodes we need:
|
||||
|
||||
<ul>
|
||||
<li>The <code>agent</code> node: responsible for deciding what (if any) actions to take.</li>
|
||||
<li>The <code>tools</code> node that invokes tools: if the agent decides to take an action, this node will then execute that action.</li>
|
||||
</ul>
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Define entry point and graph edges.</summary>
|
||||
|
||||
First, we need to set the entry point for graph execution - <code>agent</code> node.
|
||||
|
||||
Then we define one normal and one conditional edge. Conditional edge means that the destination depends on the contents of the graph's state (<code>MessagesState</code>). In our case, the destination is not known until the agent (LLM) decides.
|
||||
|
||||
<ul>
|
||||
<li>Conditional edge: after the agent is called, we should either:
|
||||
<ul>
|
||||
<li>a. Run tools if the agent said to take an action, OR</li>
|
||||
<li>b. Finish (respond to the user) if the agent did not ask to run tools</li>
|
||||
</ul>
|
||||
</li>
|
||||
<li>Normal edge: after the tools are invoked, the graph should always return to the agent to decide what to do next</li>
|
||||
</ul>
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Compile the graph.</summary>
|
||||
|
||||
<ul>
|
||||
<li>
|
||||
When we compile the graph, we turn it into a LangChain
|
||||
<a href="https://python.langchain.com/docs/concepts/runnables/">Runnable</a>,
|
||||
which automatically enables calling <code>.invoke()</code>, <code>.stream()</code> and <code>.batch()</code>
|
||||
with your inputs
|
||||
</li>
|
||||
<li>
|
||||
We can also optionally pass checkpointer object for persisting state between graph runs, and enabling memory,
|
||||
human-in-the-loop workflows, time travel and more. In our case we use <code>MemorySaver</code> -
|
||||
a simple in-memory checkpointer
|
||||
</li>
|
||||
</ul>
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Execute the graph.</summary>
|
||||
|
||||
<ol>
|
||||
<li>LangGraph adds the input message to the internal state, then passes the state to the entrypoint node, <code>"agent"</code>.</li>
|
||||
<li>The <code>"agent"</code> node executes, invoking the chat model.</li>
|
||||
<li>The chat model returns an <code>AIMessage</code>. LangGraph adds this to the state.</li>
|
||||
<li>Graph cycles the following steps until there are no more <code>tool_calls</code> on <code>AIMessage</code>:
|
||||
<ul>
|
||||
<li>If <code>AIMessage</code> has <code>tool_calls</code>, <code>"tools"</code> node executes</li>
|
||||
<li>The <code>"agent"</code> node executes again and returns <code>AIMessage</code></li>
|
||||
</ul>
|
||||
</li>
|
||||
<li>Execution progresses to the special <code>END</code> value and outputs the final state. And as a result, we get a list of all our chat messages as output.</li>
|
||||
</ol>
|
||||
</details>
|
||||
|
||||
</details>
|
||||
|
||||
## Documentation
|
||||
|
||||
* [Tutorials](https://langchain-ai.github.io/langgraph/tutorials/): Learn to build with LangGraph through guided examples.
|
||||
* [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.
|
||||
* [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.
|
||||
|
||||
## Resources
|
||||
|
||||
* [Built with LangGraph](https://www.langchain.com/built-with-langgraph): Hear how industry leaders use LangGraph to ship powerful, production-ready AI applications.
|
||||
|
||||
## Contributing
|
||||
|
||||
For more information on how to contribute, see [here](https://github.com/langchain-ai/langgraph/blob/main/CONTRIBUTING.md).
|
||||
|
||||
@@ -1,3 +0,0 @@
|
||||
# LangGraph examples
|
||||
|
||||
This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview). Please refer to the LangChain docs for the most up-to-date examples and usage guidelines for LangGraph.
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "10251c1c",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/chatbot-simulation-evaluation/agent-simulation-evaluation.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c5fc63df",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.1"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a4351a24",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/chatbot-simulation-evaluation/langsmith-agent-simulation-evaluation.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "4cc9af1e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,203 +0,0 @@
|
||||
import functools
|
||||
from typing import Annotated, Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
from langchain_community.adapters.openai import convert_message_to_dict
|
||||
from langchain_core.messages import AIMessage, AnyMessage, BaseMessage, HumanMessage
|
||||
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
|
||||
from langchain_core.runnables import Runnable, RunnableLambda
|
||||
from langchain_core.runnables import chain as as_runnable
|
||||
from langchain_openai import ChatOpenAI
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.graph import END, StateGraph, START
|
||||
|
||||
|
||||
def langchain_to_openai_messages(messages: List[BaseMessage]):
|
||||
"""
|
||||
Convert a list of langchain base messages to a list of openai messages.
|
||||
|
||||
Parameters:
|
||||
messages (List[BaseMessage]): A list of langchain base messages.
|
||||
|
||||
Returns:
|
||||
List[dict]: A list of openai messages.
|
||||
"""
|
||||
|
||||
return [
|
||||
convert_message_to_dict(m) if isinstance(m, BaseMessage) else m
|
||||
for m in messages
|
||||
]
|
||||
|
||||
|
||||
def create_simulated_user(
|
||||
system_prompt: str, llm: Runnable | None = None
|
||||
) -> Runnable[Dict, AIMessage]:
|
||||
"""
|
||||
Creates a simulated user for chatbot simulation.
|
||||
|
||||
Args:
|
||||
system_prompt (str): The system prompt to be used by the simulated user.
|
||||
llm (Runnable | None, optional): The language model to be used for the simulation.
|
||||
Defaults to gpt-3.5-turbo.
|
||||
|
||||
Returns:
|
||||
Runnable[Dict, AIMessage]: The simulated user for chatbot simulation.
|
||||
"""
|
||||
return ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", system_prompt),
|
||||
MessagesPlaceholder(variable_name="messages"),
|
||||
]
|
||||
) | (llm or ChatOpenAI(model="gpt-3.5-turbo")).with_config(
|
||||
run_name="simulated_user"
|
||||
)
|
||||
|
||||
|
||||
Messages = Union[list[AnyMessage], AnyMessage]
|
||||
|
||||
|
||||
def add_messages(left: Messages, right: Messages) -> Messages:
|
||||
if not isinstance(left, list):
|
||||
left = [left]
|
||||
if not isinstance(right, list):
|
||||
right = [right]
|
||||
return left + right
|
||||
|
||||
|
||||
class SimulationState(TypedDict):
|
||||
"""
|
||||
Represents the state of a simulation.
|
||||
|
||||
Attributes:
|
||||
messages (List[AnyMessage]): A list of messages in the simulation.
|
||||
inputs (Optional[dict[str, Any]]): Optional inputs for the simulation.
|
||||
"""
|
||||
|
||||
messages: Annotated[List[AnyMessage], add_messages]
|
||||
inputs: Optional[dict[str, Any]]
|
||||
|
||||
|
||||
def create_chat_simulator(
|
||||
assistant: (
|
||||
Callable[[List[AnyMessage]], str | AIMessage]
|
||||
| Runnable[List[AnyMessage], str | AIMessage]
|
||||
),
|
||||
simulated_user: Runnable[Dict, AIMessage],
|
||||
*,
|
||||
input_key: str,
|
||||
max_turns: int = 6,
|
||||
should_continue: Optional[Callable[[SimulationState], str]] = None,
|
||||
):
|
||||
"""Creates a chat simulator for evaluating a chatbot.
|
||||
|
||||
Args:
|
||||
assistant: The chatbot assistant function or runnable object.
|
||||
simulated_user: The simulated user object.
|
||||
input_key: The key for the input to the chat simulation.
|
||||
max_turns: The maximum number of turns in the chat simulation. Default is 6.
|
||||
should_continue: Optional function to determine if the simulation should continue.
|
||||
If not provided, a default function will be used.
|
||||
|
||||
Returns:
|
||||
The compiled chat simulation graph.
|
||||
|
||||
"""
|
||||
graph_builder = StateGraph(SimulationState)
|
||||
graph_builder.add_node(
|
||||
"user",
|
||||
_create_simulated_user_node(simulated_user),
|
||||
)
|
||||
graph_builder.add_node(
|
||||
"assistant", _fetch_messages | assistant | _coerce_to_message
|
||||
)
|
||||
graph_builder.add_edge("assistant", "user")
|
||||
graph_builder.add_conditional_edges(
|
||||
"user",
|
||||
should_continue or functools.partial(_should_continue, max_turns=max_turns),
|
||||
)
|
||||
# If your dataset has a 'leading question/input', then we route first to the assistant, otherwise, we let the user take the lead.
|
||||
graph_builder.add_edge(START, "assistant" if input_key is not None else "user")
|
||||
|
||||
return (
|
||||
RunnableLambda(_prepare_example).bind(input_key=input_key)
|
||||
| graph_builder.compile()
|
||||
)
|
||||
|
||||
|
||||
## Private methods
|
||||
|
||||
|
||||
def _prepare_example(inputs: dict[str, Any], input_key: Optional[str] = None):
|
||||
if input_key is not None:
|
||||
if input_key not in inputs:
|
||||
raise ValueError(
|
||||
f"Dataset's example input must contain the provided input key: '{input_key}'.\nFound: {list(inputs.keys())}"
|
||||
)
|
||||
messages = [HumanMessage(content=inputs[input_key])]
|
||||
return {
|
||||
"inputs": {k: v for k, v in inputs.items() if k != input_key},
|
||||
"messages": messages,
|
||||
}
|
||||
return {"inputs": inputs, "messages": []}
|
||||
|
||||
|
||||
def _invoke_simulated_user(state: SimulationState, simulated_user: Runnable):
|
||||
"""Invoke the simulated user node."""
|
||||
runnable = (
|
||||
simulated_user
|
||||
if isinstance(simulated_user, Runnable)
|
||||
else RunnableLambda(simulated_user)
|
||||
)
|
||||
inputs = state.get("inputs", {})
|
||||
inputs["messages"] = state["messages"]
|
||||
return runnable.invoke(inputs)
|
||||
|
||||
|
||||
def _swap_roles(state: SimulationState):
|
||||
new_messages = []
|
||||
for m in state["messages"]:
|
||||
if isinstance(m, AIMessage):
|
||||
new_messages.append(HumanMessage(content=m.content))
|
||||
else:
|
||||
new_messages.append(AIMessage(content=m.content))
|
||||
return {
|
||||
"inputs": state.get("inputs", {}),
|
||||
"messages": new_messages,
|
||||
}
|
||||
|
||||
|
||||
@as_runnable
|
||||
def _fetch_messages(state: SimulationState):
|
||||
"""Invoke the simulated user node."""
|
||||
return state["messages"]
|
||||
|
||||
|
||||
def _convert_to_human_message(message: BaseMessage):
|
||||
return {"messages": [HumanMessage(content=message.content)]}
|
||||
|
||||
|
||||
def _create_simulated_user_node(simulated_user: Runnable):
|
||||
"""Simulated user accepts a {"messages": [...]} argument and returns a single message."""
|
||||
return (
|
||||
_swap_roles
|
||||
| RunnableLambda(_invoke_simulated_user).bind(simulated_user=simulated_user)
|
||||
| _convert_to_human_message
|
||||
)
|
||||
|
||||
|
||||
def _coerce_to_message(assistant_output: str | BaseMessage):
|
||||
if isinstance(assistant_output, str):
|
||||
return {"messages": [AIMessage(content=assistant_output)]}
|
||||
else:
|
||||
return {"messages": [assistant_output]}
|
||||
|
||||
|
||||
def _should_continue(state: SimulationState, max_turns: int = 6):
|
||||
messages = state["messages"]
|
||||
# TODO support other stop criteria
|
||||
if len(messages) > max_turns:
|
||||
return END
|
||||
elif messages[-1].content.strip() == "FINISHED":
|
||||
return END
|
||||
else:
|
||||
return "assistant"
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a9014f94",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/chatbots/information-gather-prompting.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f47ce992",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.9"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1f2f13ca",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/code_assistant/langgraph_code_assistant.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "5e4c9bfe",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a8232bc9",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/customer-support/customer-support.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "63da8671",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "8dbdba5b",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/extraction/retries.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1d444b7f",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "3ecab357",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/how-tos/human_in_the_loop/wait-user-input.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "3f2866bd",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.9"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "09038b53",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/lats/lats.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b1669748",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "85205e97",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/llm-compiler/LLMCompiler.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "2fdab366",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.9"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "5cc8a2ad",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/multi_agent/hierarchical_agent_teams.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b9f3508a",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.9"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "d2b507b9",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/multi_agent/multi-agent-collaboration.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "41a8f10a",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "9138f92e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/plan-and-execute/plan-and-execute.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "093678ba",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -1,40 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "294995c4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/how-tos/react-agent-from-scratch.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.9"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
@@ -1,40 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "40f0d107",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/how-tos/react-agent-structured-output.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.9"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "658773a2",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/reflection/reflection.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1cb60657",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.9"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "caf07859",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/reflexion/reflexion.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "cd1df0e0",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "961f43ec",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/rewoo/rewoo.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7f00c427",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,40 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "bbd6e9b8",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/how-tos/run-id-langsmith.md)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.9"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f6db1873",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/self-discover/self-discover.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "219a78f9",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,40 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f49876e1",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/how-tos/subgraph.md)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.9"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
@@ -1,40 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7fd8bd65",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/how-tos/tool-calling.md)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.9"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "83c2223f",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/sql/sql-agent.md)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "57f924b1",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.9"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "11140167",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/tnt-llm/tnt-llm.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1a2ba3e6",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "9dffdb54",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/usaco/usaco.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "579c9959",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "007ea2e9",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"[This file has been moved](https://github.com/langchain-ai/langgraph/blob/23961cff61a42b52525f3b20b4094d8d2fba1744/docs/docs/tutorials/web-navigation/web_voyager.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f0d7b895",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This directory is retained purely for archival purposes and is no longer updated. The examples previously found here have been moved to the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview)."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,21 +0,0 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2024 LangChain, Inc.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -18,7 +18,7 @@ POSTGRES_VERSIONS ?= 15 16
|
||||
test_pg_version:
|
||||
@echo "Testing PostgreSQL $(POSTGRES_VERSION)"
|
||||
@POSTGRES_VERSION=$(POSTGRES_VERSION) make start-postgres
|
||||
@uv run pytest $(TEST)
|
||||
@poetry run pytest $(TEST)
|
||||
@EXIT_CODE=$$?; \
|
||||
make stop-postgres; \
|
||||
echo "Finished testing PostgreSQL $(POSTGRES_VERSION); Exit code: $$EXIT_CODE"; \
|
||||
@@ -36,7 +36,7 @@ test:
|
||||
TEST ?= .
|
||||
test_watch:
|
||||
POSTGRES_VERSION=${POSTGRES_VERSION:-16} make start-postgres; \
|
||||
uv run ptw $(TEST); \
|
||||
poetry run ptw $(TEST); \
|
||||
EXIT_CODE=$$?; \
|
||||
make stop-postgres; \
|
||||
exit $$EXIT_CODE
|
||||
@@ -55,12 +55,12 @@ lint_tests: PYTHON_FILES=tests
|
||||
lint_tests: MYPY_CACHE=.mypy_cache_test
|
||||
|
||||
lint lint_diff lint_package lint_tests:
|
||||
uv run ruff check .
|
||||
[ "$(PYTHON_FILES)" = "" ] || uv run ruff format $(PYTHON_FILES) --diff
|
||||
[ "$(PYTHON_FILES)" = "" ] || uv run ruff check --select I $(PYTHON_FILES)
|
||||
poetry run ruff check .
|
||||
[ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff
|
||||
[ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES)
|
||||
[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE)
|
||||
[ "$(PYTHON_FILES)" = "" ] || uv run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
|
||||
[ "$(PYTHON_FILES)" = "" ] || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
|
||||
|
||||
format format_diff:
|
||||
uv run ruff format $(PYTHON_FILES)
|
||||
uv run ruff check --select I --fix $(PYTHON_FILES)
|
||||
poetry run ruff format $(PYTHON_FILES)
|
||||
poetry run ruff check --select I --fix $(PYTHON_FILES)
|
||||
|
||||
@@ -13,20 +13,6 @@ By default `langgraph-checkpoint-postgres` installs `psycopg` (Psycopg 3) withou
|
||||
|
||||
> [!IMPORTANT]
|
||||
> When manually creating Postgres connections and passing them to `PostgresSaver` or `AsyncPostgresSaver`, make sure to include `autocommit=True` and `row_factory=dict_row` (`from psycopg.rows import dict_row`). See a full example in this [how-to guide](https://langchain-ai.github.io/langgraph/how-tos/persistence_postgres/).
|
||||
>
|
||||
> **Why these parameters are required:**
|
||||
> - `autocommit=True`: Required for the `.setup()` method to properly commit the checkpoint tables to the database. Without this, table creation may not be persisted.
|
||||
> - `row_factory=dict_row`: Required because the PostgresSaver implementation accesses database rows using dictionary-style syntax (e.g., `row["column_name"]`). The default `tuple_row` factory returns tuples that only support index-based access (e.g., `row[0]`), which will cause `TypeError` exceptions when the checkpointer tries to access columns by name.
|
||||
>
|
||||
> **Example of incorrect usage:**
|
||||
> ```python
|
||||
> # ❌ This will fail with TypeError during checkpointer operations
|
||||
> with psycopg.connect(DB_URI) as conn: # Missing autocommit=True and row_factory=dict_row
|
||||
> checkpointer = PostgresSaver(conn)
|
||||
> checkpointer.setup() # May not persist tables properly
|
||||
> # Any operation that reads from database will fail with:
|
||||
> # TypeError: tuple indices must be integers or slices, not str
|
||||
> ```
|
||||
|
||||
```python
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
@@ -39,7 +25,7 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
|
||||
# call .setup() the first time you're using the checkpointer
|
||||
checkpointer.setup()
|
||||
checkpoint = {
|
||||
"v": 4,
|
||||
"v": 1,
|
||||
"ts": "2024-07-31T20:14:19.804150+00:00",
|
||||
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
|
||||
"channel_values": {
|
||||
@@ -61,6 +47,7 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
|
||||
"start:node": 2
|
||||
}
|
||||
},
|
||||
"pending_sends": [],
|
||||
}
|
||||
|
||||
# store checkpoint
|
||||
@@ -72,44 +59,3 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
|
||||
# list checkpoints
|
||||
list(checkpointer.list(read_config))
|
||||
```
|
||||
|
||||
### Async
|
||||
|
||||
```python
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
|
||||
async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:
|
||||
checkpoint = {
|
||||
"v": 4,
|
||||
"ts": "2024-07-31T20:14:19.804150+00:00",
|
||||
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
|
||||
"channel_values": {
|
||||
"my_key": "meow",
|
||||
"node": "node"
|
||||
},
|
||||
"channel_versions": {
|
||||
"__start__": 2,
|
||||
"my_key": 3,
|
||||
"start:node": 3,
|
||||
"node": 3
|
||||
},
|
||||
"versions_seen": {
|
||||
"__input__": {},
|
||||
"__start__": {
|
||||
"__start__": 1
|
||||
},
|
||||
"node": {
|
||||
"start:node": 2
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
# store checkpoint
|
||||
await checkpointer.aput(write_config, checkpoint, {}, {})
|
||||
|
||||
# load checkpoint
|
||||
await checkpointer.aget(read_config)
|
||||
|
||||
# list checkpoints
|
||||
[c async for c in checkpointer.alist(read_config)]
|
||||
```
|
||||
|
||||
@@ -1,44 +1,38 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from collections import defaultdict
|
||||
from collections.abc import Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
from typing import Any, Optional
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
WRITES_IDX_MAP,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
get_checkpoint_id,
|
||||
get_serializable_checkpoint_metadata,
|
||||
)
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
from psycopg import Capabilities, Connection, Cursor, Pipeline
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
WRITES_IDX_MAP,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointConfig,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
get_checkpoint_id,
|
||||
get_checkpoint_metadata,
|
||||
)
|
||||
from langgraph.checkpoint.postgres import _internal
|
||||
from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
|
||||
Conn = _internal.Conn # For backward compatibility
|
||||
|
||||
|
||||
class PostgresSaver(BasePostgresSaver):
|
||||
"""Checkpointer that stores checkpoints in a Postgres database."""
|
||||
|
||||
lock: threading.Lock
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: _internal.Conn,
|
||||
pipe: Pipeline | None = None,
|
||||
serde: SerializerProtocol | None = None,
|
||||
pipe: Optional[Pipeline] = None,
|
||||
serde: Optional[SerializerProtocol] = None,
|
||||
) -> None:
|
||||
super().__init__(serde=serde)
|
||||
if isinstance(conn, ConnectionPool) and pipe is not None:
|
||||
@@ -55,12 +49,12 @@ class PostgresSaver(BasePostgresSaver):
|
||||
@contextmanager
|
||||
def from_conn_string(
|
||||
cls, conn_string: str, *, pipeline: bool = False
|
||||
) -> Iterator[PostgresSaver]:
|
||||
) -> Iterator["PostgresSaver"]:
|
||||
"""Create a new PostgresSaver instance from a connection string.
|
||||
|
||||
Args:
|
||||
conn_string: The Postgres connection info string.
|
||||
pipeline: whether to use Pipeline
|
||||
conn_string (str): The Postgres connection info string.
|
||||
pipeline (bool): whether to use Pipeline
|
||||
|
||||
Returns:
|
||||
PostgresSaver: A new PostgresSaver instance.
|
||||
@@ -94,20 +88,19 @@ class PostgresSaver(BasePostgresSaver):
|
||||
for v, migration in zip(
|
||||
range(version + 1, len(self.MIGRATIONS)),
|
||||
self.MIGRATIONS[version + 1 :],
|
||||
strict=False,
|
||||
):
|
||||
cur.execute(migration)
|
||||
cur.execute("INSERT INTO checkpoint_migrations (v) VALUES (%s)", (v,))
|
||||
cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})")
|
||||
if self.pipe:
|
||||
self.pipe.sync()
|
||||
|
||||
def list(
|
||||
self,
|
||||
config: RunnableConfig | None,
|
||||
config: Optional[CheckpointConfig],
|
||||
*,
|
||||
filter: dict[str, Any] | None = None,
|
||||
before: RunnableConfig | None = None,
|
||||
limit: int | None = None,
|
||||
filter: Optional[dict[str, Any]] = None,
|
||||
before: Optional[CheckpointConfig] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> Iterator[CheckpointTuple]:
|
||||
"""List checkpoints from the database.
|
||||
|
||||
@@ -115,13 +108,13 @@ class PostgresSaver(BasePostgresSaver):
|
||||
on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
|
||||
|
||||
Args:
|
||||
config: The config to use for listing the checkpoints.
|
||||
filter: Additional filtering criteria for metadata.
|
||||
before: If provided, only checkpoints before the specified checkpoint ID are returned.
|
||||
limit: The maximum number of checkpoints to return.
|
||||
config (CheckpointConfig): The config to use for listing the checkpoints.
|
||||
filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. Defaults to None.
|
||||
before (Optional[CheckpointConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
|
||||
limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
|
||||
|
||||
Yields:
|
||||
An iterator of checkpoint tuples.
|
||||
Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
|
||||
|
||||
Examples:
|
||||
>>> from langgraph.checkpoint.postgres import PostgresSaver
|
||||
@@ -143,57 +136,53 @@ class PostgresSaver(BasePostgresSaver):
|
||||
"""
|
||||
where, args = self._search_where(config, filter, before)
|
||||
query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
|
||||
params = list(args)
|
||||
if limit is not None:
|
||||
query += " LIMIT %s"
|
||||
params.append(int(limit))
|
||||
if limit:
|
||||
query += f" LIMIT {limit}"
|
||||
# if we change this to use .stream() we need to make sure to close the cursor
|
||||
with self._cursor() as cur:
|
||||
cur.execute(query, params)
|
||||
values = cur.fetchall()
|
||||
if not values:
|
||||
return
|
||||
# migrate pending sends if necessary
|
||||
if to_migrate := [
|
||||
v
|
||||
for v in values
|
||||
if v["checkpoint"]["v"] < 4 and v["parent_checkpoint_id"]
|
||||
]:
|
||||
cur.execute(
|
||||
self.SELECT_PENDING_SENDS_SQL,
|
||||
(
|
||||
values[0]["thread_id"],
|
||||
[v["parent_checkpoint_id"] for v in to_migrate],
|
||||
cur.execute(query, args, binary=True)
|
||||
for value in cur:
|
||||
yield CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
self._load_checkpoint(
|
||||
value["checkpoint"],
|
||||
value["channel_values"],
|
||||
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"],
|
||||
}
|
||||
}
|
||||
if value["parent_checkpoint_id"]
|
||||
else None
|
||||
),
|
||||
self._load_writes(value["pending_writes"]),
|
||||
)
|
||||
grouped_by_parent = defaultdict(list)
|
||||
for value in to_migrate:
|
||||
grouped_by_parent[value["parent_checkpoint_id"]].append(value)
|
||||
for sends in cur:
|
||||
for value in grouped_by_parent[sends["checkpoint_id"]]:
|
||||
if value["channel_values"] is None:
|
||||
value["channel_values"] = []
|
||||
self._migrate_pending_sends(
|
||||
sends["sends"],
|
||||
value["checkpoint"],
|
||||
value["channel_values"],
|
||||
)
|
||||
for value in values:
|
||||
yield self._load_checkpoint_tuple(value)
|
||||
|
||||
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
def get_tuple(self, config: CheckpointConfig) -> Optional[CheckpointTuple]:
|
||||
"""Get a checkpoint tuple from the database.
|
||||
|
||||
This method retrieves a checkpoint tuple from the Postgres database based on the
|
||||
provided config. If the config contains a `checkpoint_id` key, the checkpoint with
|
||||
provided config. If the config contains a "checkpoint_id" key, the checkpoint with
|
||||
the matching thread ID and timestamp is retrieved. Otherwise, the latest checkpoint
|
||||
for the given thread ID is retrieved.
|
||||
|
||||
Args:
|
||||
config: The config to use for retrieving the checkpoint.
|
||||
config (CheckpointConfig): The config to use for retrieving the checkpoint.
|
||||
|
||||
Returns:
|
||||
The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
||||
Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
||||
|
||||
Examples:
|
||||
|
||||
@@ -230,48 +219,58 @@ class PostgresSaver(BasePostgresSaver):
|
||||
cur.execute(
|
||||
self.SELECT_SQL + where,
|
||||
args,
|
||||
binary=True,
|
||||
)
|
||||
value = cur.fetchone()
|
||||
if value is None:
|
||||
return None
|
||||
|
||||
# migrate pending sends if necessary
|
||||
if value["checkpoint"]["v"] < 4 and value["parent_checkpoint_id"]:
|
||||
cur.execute(
|
||||
self.SELECT_PENDING_SENDS_SQL,
|
||||
(thread_id, [value["parent_checkpoint_id"]]),
|
||||
)
|
||||
if sends := cur.fetchone():
|
||||
if value["channel_values"] is None:
|
||||
value["channel_values"] = []
|
||||
self._migrate_pending_sends(
|
||||
sends["sends"],
|
||||
for value in cur:
|
||||
return CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
self._load_checkpoint(
|
||||
value["checkpoint"],
|
||||
value["channel_values"],
|
||||
)
|
||||
|
||||
return self._load_checkpoint_tuple(value)
|
||||
value["pending_sends"],
|
||||
),
|
||||
self._load_metadata(value["metadata"]),
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
}
|
||||
}
|
||||
if value["parent_checkpoint_id"]
|
||||
else None
|
||||
),
|
||||
self._load_writes(value["pending_writes"]),
|
||||
)
|
||||
|
||||
def put(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
config: CheckpointConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
) -> CheckpointConfig:
|
||||
"""Save a checkpoint to the database.
|
||||
|
||||
This method saves a checkpoint to the Postgres database. The checkpoint is associated
|
||||
with the provided config and its parent config (if any).
|
||||
|
||||
Args:
|
||||
config: The config to associate with the checkpoint.
|
||||
checkpoint: The checkpoint to save.
|
||||
metadata: Additional metadata to save with the checkpoint.
|
||||
new_versions: New channel versions as of this write.
|
||||
config (CheckpointConfig): The config to associate with the checkpoint.
|
||||
checkpoint (Checkpoint): The checkpoint to save.
|
||||
metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
|
||||
new_versions (ChannelVersions): New channel versions as of this write.
|
||||
|
||||
Returns:
|
||||
RunnableConfig: Updated configuration after storing the checkpoint.
|
||||
CheckpointConfig: Updated configuration after storing the checkpoint.
|
||||
|
||||
Examples:
|
||||
|
||||
@@ -287,9 +286,11 @@ class PostgresSaver(BasePostgresSaver):
|
||||
configurable = config["configurable"].copy()
|
||||
thread_id = configurable.pop("thread_id")
|
||||
checkpoint_ns = configurable.pop("checkpoint_ns")
|
||||
checkpoint_id = configurable.pop("checkpoint_id", None)
|
||||
checkpoint_id = configurable.pop(
|
||||
"checkpoint_id", configurable.pop("thread_ts", None)
|
||||
)
|
||||
|
||||
copy = checkpoint.copy()
|
||||
copy["channel_values"] = copy["channel_values"].copy()
|
||||
next_config = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
@@ -298,28 +299,16 @@ class PostgresSaver(BasePostgresSaver):
|
||||
}
|
||||
}
|
||||
|
||||
# inline primitive values in checkpoint table
|
||||
# others are stored in blobs table
|
||||
blob_values = {}
|
||||
for k, v in checkpoint["channel_values"].items():
|
||||
if v is None or isinstance(v, (str, int, float, bool)):
|
||||
pass
|
||||
else:
|
||||
blob_values[k] = copy["channel_values"].pop(k)
|
||||
|
||||
with self._cursor(pipeline=True) as cur:
|
||||
if blob_versions := {
|
||||
k: v for k, v in new_versions.items() if k in blob_values
|
||||
}:
|
||||
cur.executemany(
|
||||
self.UPSERT_CHECKPOINT_BLOBS_SQL,
|
||||
self._dump_blobs(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
blob_values,
|
||||
blob_versions,
|
||||
),
|
||||
)
|
||||
cur.executemany(
|
||||
self.UPSERT_CHECKPOINT_BLOBS_SQL,
|
||||
self._dump_blobs(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
copy.pop("channel_values"), # type: ignore[misc]
|
||||
new_versions,
|
||||
),
|
||||
)
|
||||
cur.execute(
|
||||
self.UPSERT_CHECKPOINTS_SQL,
|
||||
(
|
||||
@@ -327,15 +316,15 @@ class PostgresSaver(BasePostgresSaver):
|
||||
checkpoint_ns,
|
||||
checkpoint["id"],
|
||||
checkpoint_id,
|
||||
Jsonb(copy),
|
||||
Jsonb(get_serializable_checkpoint_metadata(config, metadata)),
|
||||
Jsonb(self._dump_checkpoint(copy)),
|
||||
self._dump_metadata(get_checkpoint_metadata(config, metadata)),
|
||||
),
|
||||
)
|
||||
return next_config
|
||||
|
||||
def put_writes(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
config: CheckpointConfig,
|
||||
writes: Sequence[tuple[str, Any]],
|
||||
task_id: str,
|
||||
task_path: str = "",
|
||||
@@ -345,9 +334,9 @@ class PostgresSaver(BasePostgresSaver):
|
||||
This method saves intermediate writes associated with a checkpoint to the Postgres database.
|
||||
|
||||
Args:
|
||||
config: Configuration of the related checkpoint.
|
||||
writes: List of writes to store.
|
||||
task_id: Identifier for the task creating the writes.
|
||||
config (CheckpointConfig): Configuration of the related checkpoint.
|
||||
writes (List[Tuple[str, Any]]): List of writes to store.
|
||||
task_id (str): Identifier for the task creating the writes.
|
||||
"""
|
||||
query = (
|
||||
self.UPSERT_CHECKPOINT_WRITES_SQL
|
||||
@@ -367,39 +356,16 @@ class PostgresSaver(BasePostgresSaver):
|
||||
),
|
||||
)
|
||||
|
||||
def delete_thread(self, thread_id: str) -> None:
|
||||
"""Delete all checkpoints and writes associated with a thread ID.
|
||||
|
||||
Args:
|
||||
thread_id: The thread ID to delete.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
with self._cursor(pipeline=True) as cur:
|
||||
cur.execute(
|
||||
"DELETE FROM checkpoints WHERE thread_id = %s",
|
||||
(str(thread_id),),
|
||||
)
|
||||
cur.execute(
|
||||
"DELETE FROM checkpoint_blobs WHERE thread_id = %s",
|
||||
(str(thread_id),),
|
||||
)
|
||||
cur.execute(
|
||||
"DELETE FROM checkpoint_writes WHERE thread_id = %s",
|
||||
(str(thread_id),),
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
|
||||
"""Create a database cursor as a context manager.
|
||||
|
||||
Args:
|
||||
pipeline: whether to use pipeline for the DB operations inside the context manager.
|
||||
pipeline (bool): whether to use pipeline for the DB operations inside the context manager.
|
||||
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 self.lock, _internal.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
|
||||
@@ -415,6 +381,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
# 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,
|
||||
):
|
||||
@@ -422,55 +389,14 @@ class PostgresSaver(BasePostgresSaver):
|
||||
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,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
||||
"""
|
||||
Convert a database row into a CheckpointTuple object.
|
||||
|
||||
Args:
|
||||
value: A row from the database containing checkpoint data.
|
||||
|
||||
Returns:
|
||||
CheckpointTuple: A structured representation of the checkpoint,
|
||||
including its configuration, metadata, parent checkpoint (if any),
|
||||
and pending writes.
|
||||
"""
|
||||
return CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
{
|
||||
**value["checkpoint"],
|
||||
"channel_values": {
|
||||
**(value["checkpoint"].get("channel_values") or {}),
|
||||
**self._load_blobs(value["channel_values"]),
|
||||
},
|
||||
},
|
||||
value["metadata"],
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
}
|
||||
}
|
||||
if value["parent_checkpoint_id"]
|
||||
else None
|
||||
),
|
||||
self._load_writes(value["pending_writes"]),
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["PostgresSaver", "BasePostgresSaver", "ShallowPostgresSaver", "Conn"]
|
||||
__all__ = ["PostgresSaver", "BasePostgresSaver", "Conn"]
|
||||
|
||||
@@ -1,23 +0,0 @@
|
||||
"""Shared async utility functions for the Postgres checkpoint & storage classes."""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from psycopg import AsyncConnection
|
||||
from psycopg.rows import DictRow
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
|
||||
Conn = 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)}")
|
||||
@@ -2,12 +2,13 @@
|
||||
|
||||
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 = Connection[DictRow] | ConnectionPool[Connection[DictRow]]
|
||||
Conn = Union[Connection[DictRow], ConnectionPool[Connection[DictRow]]]
|
||||
|
||||
|
||||
@contextmanager
|
||||
|
||||
@@ -1,582 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections import defaultdict
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
WRITES_IDX_MAP,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
get_checkpoint_id,
|
||||
get_serializable_checkpoint_metadata,
|
||||
)
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
|
||||
from langgraph.checkpoint.postgres import _ainternal
|
||||
from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
|
||||
|
||||
Conn = _ainternal.Conn # For backward compatibility
|
||||
|
||||
|
||||
class AsyncPostgresSaver(BasePostgresSaver):
|
||||
"""Asynchronous checkpointer that stores checkpoints in a Postgres database."""
|
||||
|
||||
lock: asyncio.Lock
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: _ainternal.Conn,
|
||||
pipe: AsyncPipeline | None = None,
|
||||
serde: SerializerProtocol | None = None,
|
||||
) -> None:
|
||||
super().__init__(serde=serde)
|
||||
if isinstance(conn, AsyncConnectionPool) and pipe is not None:
|
||||
raise ValueError(
|
||||
"Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool."
|
||||
)
|
||||
|
||||
self.conn = conn
|
||||
self.pipe = pipe
|
||||
self.lock = asyncio.Lock()
|
||||
self.loop = asyncio.get_running_loop()
|
||||
self.supports_pipeline = Capabilities().has_pipeline()
|
||||
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
async def from_conn_string(
|
||||
cls,
|
||||
conn_string: str,
|
||||
*,
|
||||
pipeline: bool = False,
|
||||
serde: SerializerProtocol | None = None,
|
||||
) -> AsyncIterator[AsyncPostgresSaver]:
|
||||
"""Create a new AsyncPostgresSaver instance from a connection string.
|
||||
|
||||
Args:
|
||||
conn_string: The Postgres connection info string.
|
||||
pipeline: whether to use AsyncPipeline
|
||||
|
||||
Returns:
|
||||
AsyncPostgresSaver: A new AsyncPostgresSaver instance.
|
||||
"""
|
||||
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, serde=serde)
|
||||
else:
|
||||
yield cls(conn=conn, serde=serde)
|
||||
|
||||
async def setup(self) -> None:
|
||||
"""Set up the checkpoint 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 checkpointer is used.
|
||||
"""
|
||||
async with self._cursor() as cur:
|
||||
await cur.execute(self.MIGRATIONS[0])
|
||||
results = await cur.execute(
|
||||
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
|
||||
)
|
||||
row = await results.fetchone()
|
||||
if row is None:
|
||||
version = -1
|
||||
else:
|
||||
version = row["v"]
|
||||
for v, migration in zip(
|
||||
range(version + 1, len(self.MIGRATIONS)),
|
||||
self.MIGRATIONS[version + 1 :],
|
||||
strict=False,
|
||||
):
|
||||
await cur.execute(migration)
|
||||
await cur.execute(
|
||||
"INSERT INTO checkpoint_migrations (v) VALUES (%s)", (v,)
|
||||
)
|
||||
if self.pipe:
|
||||
await self.pipe.sync()
|
||||
|
||||
async def alist(
|
||||
self,
|
||||
config: RunnableConfig | None,
|
||||
*,
|
||||
filter: dict[str, Any] | None = None,
|
||||
before: RunnableConfig | None = None,
|
||||
limit: int | None = None,
|
||||
) -> AsyncIterator[CheckpointTuple]:
|
||||
"""List checkpoints from the database asynchronously.
|
||||
|
||||
This method retrieves a list of checkpoint tuples from the Postgres database based
|
||||
on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
|
||||
|
||||
Args:
|
||||
config: Base configuration for filtering checkpoints.
|
||||
filter: Additional filtering criteria for metadata.
|
||||
before: If provided, only checkpoints before the specified checkpoint ID are returned.
|
||||
limit: Maximum number of checkpoints to return.
|
||||
|
||||
Yields:
|
||||
An asynchronous iterator of matching checkpoint tuples.
|
||||
"""
|
||||
where, args = self._search_where(config, filter, before)
|
||||
query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
|
||||
params = list(args)
|
||||
if limit is not None:
|
||||
query += " LIMIT %s"
|
||||
params.append(int(limit))
|
||||
# if we change this to use .stream() we need to make sure to close the cursor
|
||||
async with self._cursor() as cur:
|
||||
await cur.execute(query, params, binary=True)
|
||||
values = await cur.fetchall()
|
||||
if not values:
|
||||
return
|
||||
# migrate pending sends if necessary
|
||||
if to_migrate := [
|
||||
v
|
||||
for v in values
|
||||
if v["checkpoint"]["v"] < 4 and v["parent_checkpoint_id"]
|
||||
]:
|
||||
await cur.execute(
|
||||
self.SELECT_PENDING_SENDS_SQL,
|
||||
(
|
||||
values[0]["thread_id"],
|
||||
[v["parent_checkpoint_id"] for v in to_migrate],
|
||||
),
|
||||
)
|
||||
grouped_by_parent = defaultdict(list)
|
||||
for value in to_migrate:
|
||||
grouped_by_parent[value["parent_checkpoint_id"]].append(value)
|
||||
async for sends in cur:
|
||||
for value in grouped_by_parent[sends["checkpoint_id"]]:
|
||||
if value["channel_values"] is None:
|
||||
value["channel_values"] = []
|
||||
self._migrate_pending_sends(
|
||||
sends["sends"],
|
||||
value["checkpoint"],
|
||||
value["channel_values"],
|
||||
)
|
||||
for value in values:
|
||||
yield await self._load_checkpoint_tuple(value)
|
||||
|
||||
async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Get a checkpoint tuple from the database asynchronously.
|
||||
|
||||
This method retrieves a checkpoint tuple from the Postgres database based on the
|
||||
provided config. If the config contains a `checkpoint_id` key, the checkpoint with
|
||||
the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
|
||||
for the given thread ID is retrieved.
|
||||
|
||||
Args:
|
||||
config: The config to use for retrieving the checkpoint.
|
||||
|
||||
Returns:
|
||||
The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
||||
"""
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_id = get_checkpoint_id(config)
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
if checkpoint_id:
|
||||
args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
|
||||
where = "WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s"
|
||||
else:
|
||||
args = (thread_id, checkpoint_ns)
|
||||
where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1"
|
||||
|
||||
async with self._cursor() as cur:
|
||||
await cur.execute(
|
||||
self.SELECT_SQL + where,
|
||||
args,
|
||||
binary=True,
|
||||
)
|
||||
value = await cur.fetchone()
|
||||
if value is None:
|
||||
return None
|
||||
|
||||
# migrate pending sends if necessary
|
||||
if value["checkpoint"]["v"] < 4 and value["parent_checkpoint_id"]:
|
||||
await cur.execute(
|
||||
self.SELECT_PENDING_SENDS_SQL,
|
||||
(thread_id, [value["parent_checkpoint_id"]]),
|
||||
)
|
||||
if sends := await cur.fetchone():
|
||||
if value["channel_values"] is None:
|
||||
value["channel_values"] = []
|
||||
self._migrate_pending_sends(
|
||||
sends["sends"],
|
||||
value["checkpoint"],
|
||||
value["channel_values"],
|
||||
)
|
||||
|
||||
return await self._load_checkpoint_tuple(value)
|
||||
|
||||
async def aput(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
"""Save a checkpoint to the database asynchronously.
|
||||
|
||||
This method saves a checkpoint to the Postgres database. The checkpoint is associated
|
||||
with the provided config and its parent config (if any).
|
||||
|
||||
Args:
|
||||
config: The config to associate with the checkpoint.
|
||||
checkpoint: The checkpoint to save.
|
||||
metadata: Additional metadata to save with the checkpoint.
|
||||
new_versions: New channel versions as of this write.
|
||||
|
||||
Returns:
|
||||
RunnableConfig: Updated configuration after storing the checkpoint.
|
||||
"""
|
||||
configurable = config["configurable"].copy()
|
||||
thread_id = configurable.pop("thread_id")
|
||||
checkpoint_ns = configurable.pop("checkpoint_ns")
|
||||
checkpoint_id = configurable.pop("checkpoint_id", None)
|
||||
|
||||
copy = checkpoint.copy()
|
||||
copy["channel_values"] = copy["channel_values"].copy()
|
||||
next_config = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
}
|
||||
|
||||
# inline primitive values in checkpoint table
|
||||
# others are stored in blobs table
|
||||
blob_values = {}
|
||||
for k, v in checkpoint["channel_values"].items():
|
||||
if v is None or isinstance(v, (str, int, float, bool)):
|
||||
pass
|
||||
else:
|
||||
blob_values[k] = copy["channel_values"].pop(k)
|
||||
|
||||
async with self._cursor(pipeline=True) as cur:
|
||||
if blob_versions := {
|
||||
k: v for k, v in new_versions.items() if k in blob_values
|
||||
}:
|
||||
await cur.executemany(
|
||||
self.UPSERT_CHECKPOINT_BLOBS_SQL,
|
||||
await asyncio.to_thread(
|
||||
self._dump_blobs,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
blob_values,
|
||||
blob_versions,
|
||||
),
|
||||
)
|
||||
await cur.execute(
|
||||
self.UPSERT_CHECKPOINTS_SQL,
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint["id"],
|
||||
checkpoint_id,
|
||||
Jsonb(copy),
|
||||
Jsonb(get_serializable_checkpoint_metadata(config, metadata)),
|
||||
),
|
||||
)
|
||||
return next_config
|
||||
|
||||
async def aput_writes(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
writes: Sequence[tuple[str, Any]],
|
||||
task_id: str,
|
||||
task_path: str = "",
|
||||
) -> None:
|
||||
"""Store intermediate writes linked to a checkpoint asynchronously.
|
||||
|
||||
This method saves intermediate writes associated with a checkpoint to the database.
|
||||
|
||||
Args:
|
||||
config: Configuration of the related checkpoint.
|
||||
writes: List of writes to store, each as (channel, value) pair.
|
||||
task_id: Identifier for the task creating the writes.
|
||||
"""
|
||||
query = (
|
||||
self.UPSERT_CHECKPOINT_WRITES_SQL
|
||||
if all(w[0] in WRITES_IDX_MAP for w in writes)
|
||||
else self.INSERT_CHECKPOINT_WRITES_SQL
|
||||
)
|
||||
params = await asyncio.to_thread(
|
||||
self._dump_writes,
|
||||
config["configurable"]["thread_id"],
|
||||
config["configurable"]["checkpoint_ns"],
|
||||
config["configurable"]["checkpoint_id"],
|
||||
task_id,
|
||||
task_path,
|
||||
writes,
|
||||
)
|
||||
async with self._cursor(pipeline=True) as cur:
|
||||
await cur.executemany(query, params)
|
||||
|
||||
async def adelete_thread(self, thread_id: str) -> None:
|
||||
"""Delete all checkpoints and writes associated with a thread ID.
|
||||
|
||||
Args:
|
||||
thread_id: The thread ID to delete.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
async with self._cursor(pipeline=True) as cur:
|
||||
await cur.execute(
|
||||
"DELETE FROM checkpoints WHERE thread_id = %s",
|
||||
(str(thread_id),),
|
||||
)
|
||||
await cur.execute(
|
||||
"DELETE FROM checkpoint_blobs WHERE thread_id = %s",
|
||||
(str(thread_id),),
|
||||
)
|
||||
await cur.execute(
|
||||
"DELETE FROM checkpoint_writes WHERE thread_id = %s",
|
||||
(str(thread_id),),
|
||||
)
|
||||
|
||||
@asynccontextmanager
|
||||
async def _cursor(
|
||||
self, *, pipeline: bool = False
|
||||
) -> AsyncIterator[AsyncCursor[DictRow]]:
|
||||
"""Create a database cursor as a context manager.
|
||||
|
||||
Args:
|
||||
pipeline: whether to use pipeline for the DB operations inside the context manager.
|
||||
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 self.lock, _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 (
|
||||
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 (
|
||||
conn.transaction(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
||||
"""
|
||||
Convert a database row into a CheckpointTuple object.
|
||||
|
||||
Args:
|
||||
value: A row from the database containing checkpoint data.
|
||||
|
||||
Returns:
|
||||
CheckpointTuple: A structured representation of the checkpoint,
|
||||
including its configuration, metadata, parent checkpoint (if any),
|
||||
and pending writes.
|
||||
"""
|
||||
return CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
{
|
||||
**value["checkpoint"],
|
||||
"channel_values": {
|
||||
**(value["checkpoint"].get("channel_values") or {}),
|
||||
**self._load_blobs(value["channel_values"]),
|
||||
},
|
||||
},
|
||||
value["metadata"],
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
}
|
||||
}
|
||||
if value["parent_checkpoint_id"]
|
||||
else None
|
||||
),
|
||||
await asyncio.to_thread(self._load_writes, value["pending_writes"]),
|
||||
)
|
||||
|
||||
def list(
|
||||
self,
|
||||
config: RunnableConfig | None,
|
||||
*,
|
||||
filter: dict[str, Any] | None = None,
|
||||
before: RunnableConfig | None = None,
|
||||
limit: int | None = None,
|
||||
) -> Iterator[CheckpointTuple]:
|
||||
"""List checkpoints from the database.
|
||||
|
||||
This method retrieves a list of checkpoint tuples from the Postgres database based
|
||||
on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
|
||||
|
||||
Args:
|
||||
config: Base configuration for filtering checkpoints.
|
||||
filter: Additional filtering criteria for metadata.
|
||||
before: If provided, only checkpoints before the specified checkpoint ID are returned.
|
||||
limit: Maximum number of checkpoints to return.
|
||||
|
||||
Yields:
|
||||
An iterator of matching checkpoint tuples.
|
||||
"""
|
||||
try:
|
||||
# check if we are in the main thread, only bg threads can block
|
||||
# we don't check in other methods to avoid the overhead
|
||||
if asyncio.get_running_loop() is self.loop:
|
||||
raise asyncio.InvalidStateError(
|
||||
"Synchronous calls to AsyncPostgresSaver are only allowed from a "
|
||||
"different thread. From the main thread, use the async interface. "
|
||||
"For example, use `checkpointer.alist(...)` or `await "
|
||||
"graph.ainvoke(...)`."
|
||||
)
|
||||
except RuntimeError:
|
||||
pass
|
||||
aiter_ = self.alist(config, filter=filter, before=before, limit=limit)
|
||||
while True:
|
||||
try:
|
||||
yield asyncio.run_coroutine_threadsafe(
|
||||
anext(aiter_), # type: ignore[arg-type] # noqa: F821
|
||||
self.loop,
|
||||
).result()
|
||||
except StopAsyncIteration:
|
||||
break
|
||||
|
||||
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Get a checkpoint tuple from the database.
|
||||
|
||||
This method retrieves a checkpoint tuple from the Postgres database based on the
|
||||
provided config. If the config contains a `checkpoint_id` key, the checkpoint with
|
||||
the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
|
||||
for the given thread ID is retrieved.
|
||||
|
||||
Args:
|
||||
config: The config to use for retrieving the checkpoint.
|
||||
|
||||
Returns:
|
||||
The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
||||
"""
|
||||
try:
|
||||
# check if we are in the main thread, only bg threads can block
|
||||
# we don't check in other methods to avoid the overhead
|
||||
if asyncio.get_running_loop() is self.loop:
|
||||
raise asyncio.InvalidStateError(
|
||||
"Synchronous calls to AsyncPostgresSaver are only allowed from a "
|
||||
"different thread. From the main thread, use the async interface. "
|
||||
"For example, use `await checkpointer.aget_tuple(...)` or `await "
|
||||
"graph.ainvoke(...)`."
|
||||
)
|
||||
except RuntimeError:
|
||||
pass
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.aget_tuple(config), self.loop
|
||||
).result()
|
||||
|
||||
def put(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
"""Save a checkpoint to the database.
|
||||
|
||||
This method saves a checkpoint to the Postgres database. The checkpoint is associated
|
||||
with the provided config and its parent config (if any).
|
||||
|
||||
Args:
|
||||
config: The config to associate with the checkpoint.
|
||||
checkpoint: The checkpoint to save.
|
||||
metadata: Additional metadata to save with the checkpoint.
|
||||
new_versions: New channel versions as of this write.
|
||||
|
||||
Returns:
|
||||
RunnableConfig: Updated configuration after storing the checkpoint.
|
||||
"""
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.aput(config, checkpoint, metadata, new_versions), self.loop
|
||||
).result()
|
||||
|
||||
def put_writes(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
writes: Sequence[tuple[str, Any]],
|
||||
task_id: str,
|
||||
task_path: str = "",
|
||||
) -> None:
|
||||
"""Store intermediate writes linked to a checkpoint.
|
||||
|
||||
This method saves intermediate writes associated with a checkpoint to the database.
|
||||
|
||||
Args:
|
||||
config: Configuration of the related checkpoint.
|
||||
writes: List of writes to store, each as (channel, value) pair.
|
||||
task_id: Identifier for the task creating the writes.
|
||||
task_path: Path of the task creating the writes.
|
||||
"""
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.aput_writes(config, writes, task_id, task_path), self.loop
|
||||
).result()
|
||||
|
||||
def delete_thread(self, thread_id: str) -> None:
|
||||
"""Delete all checkpoints and writes associated with a thread ID.
|
||||
|
||||
Args:
|
||||
thread_id: The thread ID to delete.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
try:
|
||||
# check if we are in the main thread, only bg threads can block
|
||||
# we don't check in other methods to avoid the overhead
|
||||
if asyncio.get_running_loop() is self.loop:
|
||||
raise asyncio.InvalidStateError(
|
||||
"Synchronous calls to AsyncPostgresSaver are only allowed from a "
|
||||
"different thread. From the main thread, use the async interface. "
|
||||
"For example, use `await checkpointer.aget_tuple(...)` or `await "
|
||||
"graph.ainvoke(...)`."
|
||||
)
|
||||
except RuntimeError:
|
||||
pass
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.adelete_thread(thread_id), self.loop
|
||||
).result()
|
||||
|
||||
|
||||
__all__ = ["AsyncPostgresSaver", "AsyncShallowPostgresSaver", "Conn"]
|
||||
@@ -1,34 +1,22 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import warnings
|
||||
from collections.abc import Sequence
|
||||
from importlib.metadata import version as get_version
|
||||
from typing import Any, cast
|
||||
from typing import Any, Optional, cast
|
||||
|
||||
from psycopg.types.json import Jsonb
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
WRITES_IDX_MAP,
|
||||
BaseCheckpointSaver,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointConfig,
|
||||
CheckpointMetadata,
|
||||
get_checkpoint_id,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import TASKS
|
||||
from psycopg.types.json import Jsonb
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
from langgraph.checkpoint.serde.types import TASKS, ChannelProtocol
|
||||
|
||||
MetadataInput = dict[str, Any] | None
|
||||
|
||||
try:
|
||||
major, minor = get_version("langgraph").split(".")[:2]
|
||||
if int(major) == 0 and int(minor) < 5:
|
||||
warnings.warn(
|
||||
"You're using incompatible versions of langgraph and checkpoint-postgres. Please upgrade langgraph to avoid unexpected behavior.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
except Exception:
|
||||
# skip version check if running from source
|
||||
pass
|
||||
MetadataInput = Optional[dict[str, Any]]
|
||||
|
||||
"""
|
||||
To add a new migration, add a new string to the MIGRATIONS list.
|
||||
@@ -81,10 +69,10 @@ MIGRATIONS = [
|
||||
"""
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id);
|
||||
""",
|
||||
"""ALTER TABLE checkpoint_writes ADD COLUMN IF NOT EXISTS task_path TEXT NOT NULL DEFAULT '';""",
|
||||
"""ALTER TABLE checkpoint_writes ADD COLUMN task_path TEXT NOT NULL DEFAULT '';""",
|
||||
]
|
||||
|
||||
SELECT_SQL = """
|
||||
SELECT_SQL = f"""
|
||||
select
|
||||
thread_id,
|
||||
checkpoint,
|
||||
@@ -108,20 +96,17 @@ select
|
||||
where cw.thread_id = checkpoints.thread_id
|
||||
and cw.checkpoint_ns = checkpoints.checkpoint_ns
|
||||
and cw.checkpoint_id = checkpoints.checkpoint_id
|
||||
) as pending_writes
|
||||
) as pending_writes,
|
||||
(
|
||||
select array_agg(array[cw.type::bytea, cw.blob] order by cw.task_path, cw.task_id, cw.idx)
|
||||
from checkpoint_writes cw
|
||||
where cw.thread_id = checkpoints.thread_id
|
||||
and cw.checkpoint_ns = checkpoints.checkpoint_ns
|
||||
and cw.checkpoint_id = checkpoints.parent_checkpoint_id
|
||||
and cw.channel = '{TASKS}'
|
||||
) as pending_sends
|
||||
from checkpoints """
|
||||
|
||||
SELECT_PENDING_SENDS_SQL = f"""
|
||||
select
|
||||
checkpoint_id,
|
||||
array_agg(array[type::bytea, blob] order by task_path, task_id, idx) as sends
|
||||
from checkpoint_writes
|
||||
where thread_id = %s
|
||||
and checkpoint_id = any(%s)
|
||||
and channel = '{TASKS}'
|
||||
group by checkpoint_id
|
||||
"""
|
||||
|
||||
UPSERT_CHECKPOINT_BLOBS_SQL = """
|
||||
INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob)
|
||||
VALUES (%s, %s, %s, %s, %s, %s)
|
||||
@@ -155,34 +140,31 @@ INSERT_CHECKPOINT_WRITES_SQL = """
|
||||
|
||||
class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
SELECT_SQL = SELECT_SQL
|
||||
SELECT_PENDING_SENDS_SQL = SELECT_PENDING_SENDS_SQL
|
||||
MIGRATIONS = MIGRATIONS
|
||||
UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
|
||||
UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
|
||||
UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
|
||||
INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL
|
||||
|
||||
jsonplus_serde = JsonPlusSerializer()
|
||||
supports_pipeline: bool
|
||||
|
||||
def _migrate_pending_sends(
|
||||
def _load_checkpoint(
|
||||
self,
|
||||
pending_sends: list[tuple[bytes, bytes]],
|
||||
checkpoint: dict[str, Any],
|
||||
channel_values: list[tuple[bytes, bytes, bytes]],
|
||||
) -> None:
|
||||
if not pending_sends:
|
||||
return
|
||||
# add to values
|
||||
enc, blob = self.serde.dumps_typed(
|
||||
[self.serde.loads_typed((c.decode(), b)) for c, b in pending_sends],
|
||||
)
|
||||
channel_values.append((TASKS.encode(), enc.encode(), blob))
|
||||
# add to versions
|
||||
checkpoint["channel_versions"][TASKS] = (
|
||||
max(checkpoint["channel_versions"].values())
|
||||
if checkpoint["channel_versions"]
|
||||
else self.get_next_version(None, None)
|
||||
)
|
||||
pending_sends: list[tuple[bytes, bytes]],
|
||||
) -> Checkpoint:
|
||||
return {
|
||||
**checkpoint,
|
||||
"pending_sends": [
|
||||
self.serde.loads_typed((c.decode(), b)) for c, b in pending_sends or []
|
||||
],
|
||||
"channel_values": self._load_blobs(channel_values),
|
||||
}
|
||||
|
||||
def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]:
|
||||
return {**checkpoint, "pending_sends": []}
|
||||
|
||||
def _load_blobs(
|
||||
self, blob_values: list[tuple[bytes, bytes, bytes]]
|
||||
@@ -201,7 +183,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
checkpoint_ns: str,
|
||||
values: dict[str, Any],
|
||||
versions: ChannelVersions,
|
||||
) -> list[tuple[str, str, str, str, str, bytes | None]]:
|
||||
) -> list[tuple[str, str, str, str, str, Optional[bytes]]]:
|
||||
if not versions:
|
||||
return []
|
||||
|
||||
@@ -259,7 +241,15 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
for idx, (channel, value) in enumerate(writes)
|
||||
]
|
||||
|
||||
def get_next_version(self, current: str | None, channel: None) -> str:
|
||||
def _load_metadata(self, metadata: dict[str, Any]) -> CheckpointMetadata:
|
||||
return self.jsonplus_serde.loads(self.jsonplus_serde.dumps(metadata))
|
||||
|
||||
def _dump_metadata(self, metadata: CheckpointMetadata) -> str:
|
||||
serialized_metadata = self.jsonplus_serde.dumps(metadata)
|
||||
# NOTE: we're using JSON serializer (not msgpack), so we need to remove null characters before writing
|
||||
return serialized_metadata.decode().replace("\\u0000", "")
|
||||
|
||||
def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
|
||||
if current is None:
|
||||
current_v = 0
|
||||
elif isinstance(current, int):
|
||||
@@ -272,9 +262,9 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
|
||||
def _search_where(
|
||||
self,
|
||||
config: RunnableConfig | None,
|
||||
config: Optional[CheckpointConfig],
|
||||
filter: MetadataInput,
|
||||
before: RunnableConfig | None = None,
|
||||
before: Optional[CheckpointConfig] = None,
|
||||
) -> tuple[str, list[Any]]:
|
||||
"""Return WHERE clause predicates for alist() given config, filter, before.
|
||||
|
||||
|
||||
@@ -1,967 +0,0 @@
|
||||
import asyncio
|
||||
import threading
|
||||
import warnings
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
WRITES_IDX_MAP,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
get_serializable_checkpoint_metadata,
|
||||
)
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
from langgraph.checkpoint.serde.types import TASKS
|
||||
from psycopg import (
|
||||
AsyncConnection,
|
||||
AsyncCursor,
|
||||
AsyncPipeline,
|
||||
Capabilities,
|
||||
Connection,
|
||||
Cursor,
|
||||
Pipeline,
|
||||
)
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import AsyncConnectionPool, ConnectionPool
|
||||
|
||||
from langgraph.checkpoint.postgres import _ainternal, _internal
|
||||
from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
|
||||
"""
|
||||
To add a new migration, add a new string to the MIGRATIONS list.
|
||||
The position of the migration in the list is the version number.
|
||||
"""
|
||||
MIGRATIONS = [
|
||||
"""CREATE TABLE IF NOT EXISTS checkpoint_migrations (
|
||||
v INTEGER PRIMARY KEY
|
||||
);""",
|
||||
"""CREATE TABLE IF NOT EXISTS checkpoints (
|
||||
thread_id TEXT NOT NULL,
|
||||
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
||||
type TEXT,
|
||||
checkpoint JSONB NOT NULL,
|
||||
metadata JSONB NOT NULL DEFAULT '{}',
|
||||
PRIMARY KEY (thread_id, checkpoint_ns)
|
||||
);""",
|
||||
"""CREATE TABLE IF NOT EXISTS checkpoint_blobs (
|
||||
thread_id TEXT NOT NULL,
|
||||
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
||||
channel TEXT NOT NULL,
|
||||
type TEXT NOT NULL,
|
||||
blob BYTEA,
|
||||
PRIMARY KEY (thread_id, checkpoint_ns, channel)
|
||||
);""",
|
||||
"""CREATE TABLE IF NOT EXISTS checkpoint_writes (
|
||||
thread_id TEXT NOT NULL,
|
||||
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
||||
checkpoint_id TEXT NOT NULL,
|
||||
task_id TEXT NOT NULL,
|
||||
idx INTEGER NOT NULL,
|
||||
channel TEXT NOT NULL,
|
||||
type TEXT,
|
||||
blob BYTEA NOT NULL,
|
||||
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
|
||||
);""",
|
||||
"""
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoints_thread_id_idx ON checkpoints(thread_id);
|
||||
""",
|
||||
"""
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_blobs_thread_id_idx ON checkpoint_blobs(thread_id);
|
||||
""",
|
||||
"""
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id);
|
||||
""",
|
||||
"""
|
||||
ALTER TABLE checkpoint_writes ADD COLUMN IF NOT EXISTS task_path TEXT NOT NULL DEFAULT '';
|
||||
""",
|
||||
]
|
||||
|
||||
SELECT_SQL = f"""
|
||||
select
|
||||
thread_id,
|
||||
checkpoint,
|
||||
checkpoint_ns,
|
||||
metadata,
|
||||
(
|
||||
select array_agg(array[bl.channel::bytea, bl.type::bytea, bl.blob])
|
||||
from jsonb_each_text(checkpoint -> 'channel_versions')
|
||||
inner join checkpoint_blobs bl
|
||||
on bl.thread_id = checkpoints.thread_id
|
||||
and bl.checkpoint_ns = checkpoints.checkpoint_ns
|
||||
and bl.channel = jsonb_each_text.key
|
||||
) as channel_values,
|
||||
(
|
||||
select
|
||||
array_agg(array[cw.task_id::text::bytea, cw.channel::bytea, cw.type::bytea, cw.blob] order by cw.task_id, cw.idx)
|
||||
from checkpoint_writes cw
|
||||
where cw.thread_id = checkpoints.thread_id
|
||||
and cw.checkpoint_ns = checkpoints.checkpoint_ns
|
||||
and cw.checkpoint_id = (checkpoint->>'id')
|
||||
) as pending_writes,
|
||||
(
|
||||
select array_agg(array[cw.type::bytea, cw.blob] order by cw.task_path, cw.task_id, cw.idx)
|
||||
from checkpoint_writes cw
|
||||
where cw.thread_id = checkpoints.thread_id
|
||||
and cw.checkpoint_ns = checkpoints.checkpoint_ns
|
||||
and cw.channel = '{TASKS}'
|
||||
) as pending_sends
|
||||
from checkpoints """
|
||||
|
||||
UPSERT_CHECKPOINT_BLOBS_SQL = """
|
||||
INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, type, blob)
|
||||
VALUES (%s, %s, %s, %s, %s)
|
||||
ON CONFLICT (thread_id, checkpoint_ns, channel) DO UPDATE SET
|
||||
type = EXCLUDED.type,
|
||||
blob = EXCLUDED.blob;
|
||||
"""
|
||||
|
||||
UPSERT_CHECKPOINTS_SQL = """
|
||||
INSERT INTO checkpoints (thread_id, checkpoint_ns, checkpoint, metadata)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
ON CONFLICT (thread_id, checkpoint_ns)
|
||||
DO UPDATE SET
|
||||
checkpoint = EXCLUDED.checkpoint,
|
||||
metadata = EXCLUDED.metadata;
|
||||
"""
|
||||
|
||||
UPSERT_CHECKPOINT_WRITES_SQL = """
|
||||
INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, blob)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||||
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO UPDATE SET
|
||||
channel = EXCLUDED.channel,
|
||||
type = EXCLUDED.type,
|
||||
blob = EXCLUDED.blob;
|
||||
"""
|
||||
|
||||
INSERT_CHECKPOINT_WRITES_SQL = """
|
||||
INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, blob)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||||
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
|
||||
"""
|
||||
|
||||
|
||||
def _dump_blobs(
|
||||
serde: SerializerProtocol,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
values: dict[str, Any],
|
||||
versions: ChannelVersions,
|
||||
) -> list[tuple[str, str, str, str, bytes | None]]:
|
||||
if not versions:
|
||||
return []
|
||||
|
||||
return [
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
k,
|
||||
*(serde.dumps_typed(values[k]) if k in values else ("empty", None)),
|
||||
)
|
||||
for k in versions
|
||||
]
|
||||
|
||||
|
||||
class ShallowPostgresSaver(BasePostgresSaver):
|
||||
"""A checkpoint saver that uses Postgres to store checkpoints.
|
||||
|
||||
This checkpointer ONLY stores the most recent checkpoint and does NOT retain any history.
|
||||
It is meant to be a light-weight drop-in replacement for the PostgresSaver that
|
||||
supports most of the LangGraph persistence functionality with the exception of time travel.
|
||||
"""
|
||||
|
||||
SELECT_SQL = SELECT_SQL
|
||||
MIGRATIONS = MIGRATIONS
|
||||
UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
|
||||
UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
|
||||
UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
|
||||
INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL
|
||||
|
||||
lock: threading.Lock
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: _internal.Conn,
|
||||
pipe: Pipeline | None = None,
|
||||
serde: SerializerProtocol | None = None,
|
||||
) -> None:
|
||||
warnings.warn(
|
||||
"ShallowPostgresSaver is deprecated as of version 2.0.20 and will be removed in 3.0.0. "
|
||||
"Use PostgresSaver instead, and invoke the graph with `graph.invoke(..., durability='exit')`.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
super().__init__(serde=serde)
|
||||
if isinstance(conn, ConnectionPool) and pipe is not None:
|
||||
raise ValueError(
|
||||
"Pipeline should be used only with a single Connection, not ConnectionPool."
|
||||
)
|
||||
|
||||
self.conn = conn
|
||||
self.pipe = pipe
|
||||
self.lock = threading.Lock()
|
||||
self.supports_pipeline = Capabilities().has_pipeline()
|
||||
|
||||
@classmethod
|
||||
@contextmanager
|
||||
def from_conn_string(
|
||||
cls, conn_string: str, *, pipeline: bool = False
|
||||
) -> Iterator["ShallowPostgresSaver"]:
|
||||
"""Create a new ShallowPostgresSaver instance from a connection string.
|
||||
|
||||
Args:
|
||||
conn_string: The Postgres connection info string.
|
||||
pipeline: whether to use Pipeline
|
||||
|
||||
Returns:
|
||||
ShallowPostgresSaver: A new ShallowPostgresSaver instance.
|
||||
"""
|
||||
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)
|
||||
else:
|
||||
yield cls(conn)
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Set up the checkpoint 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 checkpointer is used.
|
||||
"""
|
||||
with self._cursor() as cur:
|
||||
cur.execute(self.MIGRATIONS[0])
|
||||
results = cur.execute(
|
||||
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
|
||||
)
|
||||
row = results.fetchone()
|
||||
if row is None:
|
||||
version = -1
|
||||
else:
|
||||
version = row["v"]
|
||||
for v, migration in zip(
|
||||
range(version + 1, len(self.MIGRATIONS)),
|
||||
self.MIGRATIONS[version + 1 :],
|
||||
strict=False,
|
||||
):
|
||||
cur.execute(migration)
|
||||
cur.execute("INSERT INTO checkpoint_migrations (v) VALUES (%s)", (v,))
|
||||
if self.pipe:
|
||||
self.pipe.sync()
|
||||
|
||||
def list(
|
||||
self,
|
||||
config: RunnableConfig | None,
|
||||
*,
|
||||
filter: dict[str, Any] | None = None,
|
||||
before: RunnableConfig | None = None,
|
||||
limit: int | None = None,
|
||||
) -> Iterator[CheckpointTuple]:
|
||||
"""List checkpoints from the database.
|
||||
|
||||
This method retrieves a list of checkpoint tuples from the Postgres database based
|
||||
on the provided config. For ShallowPostgresSaver, this method returns a list with
|
||||
ONLY the most recent checkpoint.
|
||||
"""
|
||||
where, args = self._search_where(config, filter, before)
|
||||
query = self.SELECT_SQL + where
|
||||
params = list(args)
|
||||
if limit is not None:
|
||||
query += " LIMIT %s"
|
||||
params.append(int(limit))
|
||||
with self._cursor() as cur:
|
||||
cur.execute(query, params, binary=True)
|
||||
for value in cur:
|
||||
checkpoint: Checkpoint = {
|
||||
**value["checkpoint"],
|
||||
"channel_values": self._load_blobs(value["channel_values"]),
|
||||
"pending_sends": [
|
||||
self.serde.loads_typed((t.decode(), v))
|
||||
for t, v in value["pending_sends"]
|
||||
]
|
||||
if value["pending_sends"]
|
||||
else [],
|
||||
}
|
||||
yield CheckpointTuple(
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
},
|
||||
checkpoint=checkpoint,
|
||||
metadata=value["metadata"],
|
||||
pending_writes=self._load_writes(value["pending_writes"]),
|
||||
)
|
||||
|
||||
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Get a checkpoint tuple from the database.
|
||||
|
||||
This method retrieves a checkpoint tuple from the Postgres database based on the
|
||||
provided config (matching the thread ID in the config).
|
||||
|
||||
Args:
|
||||
config: The config to use for retrieving the checkpoint.
|
||||
|
||||
Returns:
|
||||
The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
||||
|
||||
Examples:
|
||||
|
||||
Basic:
|
||||
>>> config = {"configurable": {"thread_id": "1"}}
|
||||
>>> checkpoint_tuple = memory.get_tuple(config)
|
||||
>>> print(checkpoint_tuple)
|
||||
CheckpointTuple(...)
|
||||
|
||||
With timestamp:
|
||||
|
||||
>>> config = {
|
||||
... "configurable": {
|
||||
... "thread_id": "1",
|
||||
... "checkpoint_ns": "",
|
||||
... "checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875",
|
||||
... }
|
||||
... }
|
||||
>>> checkpoint_tuple = memory.get_tuple(config)
|
||||
>>> print(checkpoint_tuple)
|
||||
CheckpointTuple(...)
|
||||
""" # noqa
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
args = (thread_id, checkpoint_ns)
|
||||
where = "WHERE thread_id = %s AND checkpoint_ns = %s"
|
||||
|
||||
with self._cursor() as cur:
|
||||
cur.execute(
|
||||
self.SELECT_SQL + where,
|
||||
args,
|
||||
binary=True,
|
||||
)
|
||||
|
||||
for value in cur:
|
||||
checkpoint: Checkpoint = {
|
||||
**value["checkpoint"],
|
||||
"channel_values": self._load_blobs(value["channel_values"]),
|
||||
"pending_sends": [
|
||||
self.serde.loads_typed((t.decode(), v))
|
||||
for t, v in value["pending_sends"]
|
||||
]
|
||||
if value["pending_sends"]
|
||||
else [],
|
||||
}
|
||||
return CheckpointTuple(
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
},
|
||||
checkpoint=checkpoint,
|
||||
metadata=value["metadata"],
|
||||
pending_writes=self._load_writes(value["pending_writes"]),
|
||||
)
|
||||
|
||||
def put(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
"""Save a checkpoint to the database.
|
||||
|
||||
This method saves a checkpoint to the Postgres database. The checkpoint is associated
|
||||
with the provided config. For ShallowPostgresSaver, this method saves ONLY the most recent
|
||||
checkpoint and overwrites a previous checkpoint, if it exists.
|
||||
|
||||
Args:
|
||||
config: The config to associate with the checkpoint.
|
||||
checkpoint: The checkpoint to save.
|
||||
metadata: Additional metadata to save with the checkpoint.
|
||||
new_versions: New channel versions as of this write.
|
||||
|
||||
Returns:
|
||||
RunnableConfig: Updated configuration after storing the checkpoint.
|
||||
|
||||
Examples:
|
||||
|
||||
>>> from langgraph.checkpoint.postgres import ShallowPostgresSaver
|
||||
>>> DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
|
||||
>>> with ShallowPostgresSaver.from_conn_string(DB_URI) as memory:
|
||||
>>> config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
|
||||
>>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}}
|
||||
>>> saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {})
|
||||
>>> print(saved_config)
|
||||
{'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}}
|
||||
"""
|
||||
configurable = config["configurable"].copy()
|
||||
thread_id = configurable.pop("thread_id")
|
||||
checkpoint_ns = configurable.pop("checkpoint_ns")
|
||||
|
||||
copy = checkpoint.copy()
|
||||
next_config = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
}
|
||||
|
||||
with self._cursor(pipeline=True) as cur:
|
||||
cur.execute(
|
||||
"""DELETE FROM checkpoint_writes
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id NOT IN (%s, %s)""",
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint["id"],
|
||||
configurable.get("checkpoint_id", ""),
|
||||
),
|
||||
)
|
||||
cur.executemany(
|
||||
self.UPSERT_CHECKPOINT_BLOBS_SQL,
|
||||
_dump_blobs(
|
||||
self.serde,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
copy.pop("channel_values"), # type: ignore[misc]
|
||||
new_versions,
|
||||
),
|
||||
)
|
||||
cur.execute(
|
||||
self.UPSERT_CHECKPOINTS_SQL,
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
Jsonb(copy),
|
||||
Jsonb(get_serializable_checkpoint_metadata(config, metadata)),
|
||||
),
|
||||
)
|
||||
return next_config
|
||||
|
||||
def put_writes(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
writes: Sequence[tuple[str, Any]],
|
||||
task_id: str,
|
||||
task_path: str = "",
|
||||
) -> None:
|
||||
"""Store intermediate writes linked to a checkpoint.
|
||||
|
||||
This method saves intermediate writes associated with a checkpoint to the Postgres database.
|
||||
|
||||
Args:
|
||||
config: Configuration of the related checkpoint.
|
||||
writes: List of writes to store.
|
||||
task_id: Identifier for the task creating the writes.
|
||||
"""
|
||||
query = (
|
||||
self.UPSERT_CHECKPOINT_WRITES_SQL
|
||||
if all(w[0] in WRITES_IDX_MAP for w in writes)
|
||||
else self.INSERT_CHECKPOINT_WRITES_SQL
|
||||
)
|
||||
with self._cursor(pipeline=True) as cur:
|
||||
cur.executemany(
|
||||
query,
|
||||
self._dump_writes(
|
||||
config["configurable"]["thread_id"],
|
||||
config["configurable"]["checkpoint_ns"],
|
||||
config["configurable"]["checkpoint_id"],
|
||||
task_id,
|
||||
task_path,
|
||||
writes,
|
||||
),
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
|
||||
"""Create a database cursor as a context manager.
|
||||
|
||||
Args:
|
||||
pipeline: whether to use pipeline for the DB operations inside the context manager.
|
||||
Will be applied regardless of whether the ShallowPostgresSaver instance was initialized with a pipeline.
|
||||
If pipeline mode is not supported, will fall back to using transaction context manager.
|
||||
"""
|
||||
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
|
||||
# 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:
|
||||
# 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,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
|
||||
class AsyncShallowPostgresSaver(BasePostgresSaver):
|
||||
"""A checkpoint saver that uses Postgres to store checkpoints asynchronously.
|
||||
|
||||
This checkpointer ONLY stores the most recent checkpoint and does NOT retain any history.
|
||||
It is meant to be a light-weight drop-in replacement for the AsyncPostgresSaver that
|
||||
supports most of the LangGraph persistence functionality with the exception of time travel.
|
||||
"""
|
||||
|
||||
SELECT_SQL = SELECT_SQL
|
||||
MIGRATIONS = MIGRATIONS
|
||||
UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
|
||||
UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
|
||||
UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
|
||||
INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL
|
||||
lock: asyncio.Lock
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: _ainternal.Conn,
|
||||
pipe: AsyncPipeline | None = None,
|
||||
serde: SerializerProtocol | None = None,
|
||||
) -> None:
|
||||
warnings.warn(
|
||||
"AsyncShallowPostgresSaver is deprecated as of version 2.0.20 and will be removed in 3.0.0. "
|
||||
"Use AsyncPostgresSaver instead, and invoke the graph with `await graph.ainvoke(..., durability='exit')`.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
super().__init__(serde=serde)
|
||||
if isinstance(conn, AsyncConnectionPool) and pipe is not None:
|
||||
raise ValueError(
|
||||
"Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool."
|
||||
)
|
||||
|
||||
self.conn = conn
|
||||
self.pipe = pipe
|
||||
self.lock = asyncio.Lock()
|
||||
self.loop = asyncio.get_running_loop()
|
||||
self.supports_pipeline = Capabilities().has_pipeline()
|
||||
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
async def from_conn_string(
|
||||
cls,
|
||||
conn_string: str,
|
||||
*,
|
||||
pipeline: bool = False,
|
||||
serde: SerializerProtocol | None = None,
|
||||
) -> AsyncIterator["AsyncShallowPostgresSaver"]:
|
||||
"""Create a new AsyncShallowPostgresSaver instance from a connection string.
|
||||
|
||||
Args:
|
||||
conn_string: The Postgres connection info string.
|
||||
pipeline: whether to use AsyncPipeline
|
||||
|
||||
Returns:
|
||||
AsyncShallowPostgresSaver: A new AsyncShallowPostgresSaver instance.
|
||||
"""
|
||||
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, serde=serde)
|
||||
else:
|
||||
yield cls(conn=conn, serde=serde)
|
||||
|
||||
async def setup(self) -> None:
|
||||
"""Set up the checkpoint 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 checkpointer is used.
|
||||
"""
|
||||
async with self._cursor() as cur:
|
||||
await cur.execute(self.MIGRATIONS[0])
|
||||
results = await cur.execute(
|
||||
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
|
||||
)
|
||||
row = await results.fetchone()
|
||||
if row is None:
|
||||
version = -1
|
||||
else:
|
||||
version = row["v"]
|
||||
for v, migration in zip(
|
||||
range(version + 1, len(self.MIGRATIONS)),
|
||||
self.MIGRATIONS[version + 1 :],
|
||||
strict=False,
|
||||
):
|
||||
await cur.execute(migration)
|
||||
await cur.execute(
|
||||
"INSERT INTO checkpoint_migrations (v) VALUES (%s)", (v,)
|
||||
)
|
||||
if self.pipe:
|
||||
await self.pipe.sync()
|
||||
|
||||
async def alist(
|
||||
self,
|
||||
config: RunnableConfig | None,
|
||||
*,
|
||||
filter: dict[str, Any] | None = None,
|
||||
before: RunnableConfig | None = None,
|
||||
limit: int | None = None,
|
||||
) -> AsyncIterator[CheckpointTuple]:
|
||||
"""List checkpoints from the database asynchronously.
|
||||
|
||||
This method retrieves a list of checkpoint tuples from the Postgres database based
|
||||
on the provided config. For ShallowPostgresSaver, this method returns a list with
|
||||
ONLY the most recent checkpoint.
|
||||
"""
|
||||
where, args = self._search_where(config, filter, before)
|
||||
query = self.SELECT_SQL + where
|
||||
params = list(args)
|
||||
if limit is not None:
|
||||
query += " LIMIT %s"
|
||||
params.append(int(limit))
|
||||
async with self._cursor() as cur:
|
||||
await cur.execute(query, params, binary=True)
|
||||
async for value in cur:
|
||||
checkpoint: Checkpoint = {
|
||||
**value["checkpoint"],
|
||||
"channel_values": self._load_blobs(value["channel_values"]),
|
||||
"pending_sends": [
|
||||
self.serde.loads_typed((t.decode(), v))
|
||||
for t, v in value["pending_sends"]
|
||||
]
|
||||
if value["pending_sends"]
|
||||
else [],
|
||||
}
|
||||
yield CheckpointTuple(
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
},
|
||||
checkpoint=checkpoint,
|
||||
metadata=value["metadata"],
|
||||
pending_writes=await asyncio.to_thread(
|
||||
self._load_writes, value["pending_writes"]
|
||||
),
|
||||
)
|
||||
|
||||
async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Get a checkpoint tuple from the database asynchronously.
|
||||
|
||||
This method retrieves a checkpoint tuple from the Postgres database based on the
|
||||
provided config (matching the thread ID in the config).
|
||||
|
||||
Args:
|
||||
config: The config to use for retrieving the checkpoint.
|
||||
|
||||
Returns:
|
||||
The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
||||
"""
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
args = (thread_id, checkpoint_ns)
|
||||
where = "WHERE thread_id = %s AND checkpoint_ns = %s"
|
||||
|
||||
async with self._cursor() as cur:
|
||||
await cur.execute(
|
||||
self.SELECT_SQL + where,
|
||||
args,
|
||||
binary=True,
|
||||
)
|
||||
|
||||
async for value in cur:
|
||||
checkpoint: Checkpoint = {
|
||||
**value["checkpoint"],
|
||||
"channel_values": self._load_blobs(value["channel_values"]),
|
||||
"pending_sends": [
|
||||
self.serde.loads_typed((t.decode(), v))
|
||||
for t, v in value["pending_sends"]
|
||||
]
|
||||
if value["pending_sends"]
|
||||
else [],
|
||||
}
|
||||
return CheckpointTuple(
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
},
|
||||
checkpoint=checkpoint,
|
||||
metadata=value["metadata"],
|
||||
pending_writes=await asyncio.to_thread(
|
||||
self._load_writes, value["pending_writes"]
|
||||
),
|
||||
)
|
||||
|
||||
async def aput(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
"""Save a checkpoint to the database asynchronously.
|
||||
|
||||
This method saves a checkpoint to the Postgres database. The checkpoint is associated
|
||||
with the provided config. For AsyncShallowPostgresSaver, this method saves ONLY the most recent
|
||||
checkpoint and overwrites a previous checkpoint, if it exists.
|
||||
|
||||
Args:
|
||||
config: The config to associate with the checkpoint.
|
||||
checkpoint: The checkpoint to save.
|
||||
metadata: Additional metadata to save with the checkpoint.
|
||||
new_versions: New channel versions as of this write.
|
||||
|
||||
Returns:
|
||||
RunnableConfig: Updated configuration after storing the checkpoint.
|
||||
"""
|
||||
configurable = config["configurable"].copy()
|
||||
thread_id = configurable.pop("thread_id")
|
||||
checkpoint_ns = configurable.pop("checkpoint_ns")
|
||||
|
||||
copy = checkpoint.copy()
|
||||
next_config = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
}
|
||||
|
||||
async with self._cursor(pipeline=True) as cur:
|
||||
await cur.execute(
|
||||
"""DELETE FROM checkpoint_writes
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id NOT IN (%s, %s)""",
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint["id"],
|
||||
configurable.get("checkpoint_id", ""),
|
||||
),
|
||||
)
|
||||
await cur.executemany(
|
||||
self.UPSERT_CHECKPOINT_BLOBS_SQL,
|
||||
_dump_blobs(
|
||||
self.serde,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
copy.pop("channel_values"), # type: ignore[misc]
|
||||
new_versions,
|
||||
),
|
||||
)
|
||||
await cur.execute(
|
||||
self.UPSERT_CHECKPOINTS_SQL,
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
Jsonb(copy),
|
||||
Jsonb(get_serializable_checkpoint_metadata(config, metadata)),
|
||||
),
|
||||
)
|
||||
return next_config
|
||||
|
||||
async def aput_writes(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
writes: Sequence[tuple[str, Any]],
|
||||
task_id: str,
|
||||
task_path: str = "",
|
||||
) -> None:
|
||||
"""Store intermediate writes linked to a checkpoint asynchronously.
|
||||
|
||||
This method saves intermediate writes associated with a checkpoint to the database.
|
||||
|
||||
Args:
|
||||
config: Configuration of the related checkpoint.
|
||||
writes: List of writes to store, each as (channel, value) pair.
|
||||
task_id: Identifier for the task creating the writes.
|
||||
"""
|
||||
query = (
|
||||
self.UPSERT_CHECKPOINT_WRITES_SQL
|
||||
if all(w[0] in WRITES_IDX_MAP for w in writes)
|
||||
else self.INSERT_CHECKPOINT_WRITES_SQL
|
||||
)
|
||||
params = await asyncio.to_thread(
|
||||
self._dump_writes,
|
||||
config["configurable"]["thread_id"],
|
||||
config["configurable"]["checkpoint_ns"],
|
||||
config["configurable"]["checkpoint_id"],
|
||||
task_id,
|
||||
task_path,
|
||||
writes,
|
||||
)
|
||||
async with self._cursor(pipeline=True) as cur:
|
||||
await cur.executemany(query, params)
|
||||
|
||||
@asynccontextmanager
|
||||
async def _cursor(
|
||||
self, *, pipeline: bool = False
|
||||
) -> AsyncIterator[AsyncCursor[DictRow]]:
|
||||
"""Create a database cursor as a context manager.
|
||||
|
||||
Args:
|
||||
pipeline: whether to use pipeline for the DB operations inside the context manager.
|
||||
Will be applied regardless of whether the AsyncShallowPostgresSaver instance was initialized with a pipeline.
|
||||
If pipeline mode is not supported, will fall back to using transaction context manager.
|
||||
"""
|
||||
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:
|
||||
# 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,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
async with (
|
||||
self.lock,
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
|
||||
def list(
|
||||
self,
|
||||
config: RunnableConfig | None,
|
||||
*,
|
||||
filter: dict[str, Any] | None = None,
|
||||
before: RunnableConfig | None = None,
|
||||
limit: int | None = None,
|
||||
) -> Iterator[CheckpointTuple]:
|
||||
"""List checkpoints from the database.
|
||||
|
||||
This method retrieves a list of checkpoint tuples from the Postgres database based
|
||||
on the provided config. For ShallowPostgresSaver, this method returns a list with
|
||||
ONLY the most recent checkpoint.
|
||||
"""
|
||||
aiter_ = self.alist(config, filter=filter, before=before, limit=limit)
|
||||
while True:
|
||||
try:
|
||||
yield asyncio.run_coroutine_threadsafe(
|
||||
anext(aiter_), # type: ignore[arg-type] # noqa: F821
|
||||
self.loop,
|
||||
).result()
|
||||
except StopAsyncIteration:
|
||||
break
|
||||
|
||||
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Get a checkpoint tuple from the database.
|
||||
|
||||
This method retrieves a checkpoint tuple from the Postgres database based on the
|
||||
provided config (matching the thread ID in the config).
|
||||
|
||||
Args:
|
||||
config: The config to use for retrieving the checkpoint.
|
||||
|
||||
Returns:
|
||||
The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
||||
"""
|
||||
try:
|
||||
# check if we are in the main thread, only bg threads can block
|
||||
# we don't check in other methods to avoid the overhead
|
||||
if asyncio.get_running_loop() is self.loop:
|
||||
raise asyncio.InvalidStateError(
|
||||
"Synchronous calls to AsyncShallowPostgresSaver are only allowed from a "
|
||||
"different thread. From the main thread, use the async interface."
|
||||
"For example, use `await checkpointer.aget_tuple(...)` or `await "
|
||||
"graph.ainvoke(...)`."
|
||||
)
|
||||
except RuntimeError:
|
||||
pass
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.aget_tuple(config), self.loop
|
||||
).result()
|
||||
|
||||
def put(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
"""Save a checkpoint to the database.
|
||||
|
||||
This method saves a checkpoint to the Postgres database. The checkpoint is associated
|
||||
with the provided config. For AsyncShallowPostgresSaver, this method saves ONLY the most recent
|
||||
checkpoint and overwrites a previous checkpoint, if it exists.
|
||||
|
||||
Args:
|
||||
config: The config to associate with the checkpoint.
|
||||
checkpoint: The checkpoint to save.
|
||||
metadata: Additional metadata to save with the checkpoint.
|
||||
new_versions: New channel versions as of this write.
|
||||
|
||||
Returns:
|
||||
RunnableConfig: Updated configuration after storing the checkpoint.
|
||||
"""
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.aput(config, checkpoint, metadata, new_versions), self.loop
|
||||
).result()
|
||||
|
||||
def put_writes(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
writes: Sequence[tuple[str, Any]],
|
||||
task_id: str,
|
||||
task_path: str = "",
|
||||
) -> None:
|
||||
"""Store intermediate writes linked to a checkpoint.
|
||||
|
||||
This method saves intermediate writes associated with a checkpoint to the database.
|
||||
|
||||
Args:
|
||||
config: Configuration of the related checkpoint.
|
||||
writes: List of writes to store, each as (channel, value) pair.
|
||||
task_id: Identifier for the task creating the writes.
|
||||
task_path: Path of the task creating the writes.
|
||||
"""
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.aput_writes(config, writes, task_id, task_path), self.loop
|
||||
).result()
|
||||
@@ -1,4 +1,3 @@
|
||||
from langgraph.store.postgres.aio import AsyncPostgresStore
|
||||
from langgraph.store.postgres.base import PoolConfig, PostgresStore
|
||||
from langgraph.store.postgres.base import PostgresStore
|
||||
|
||||
__all__ = ["AsyncPostgresStore", "PoolConfig", "PostgresStore"]
|
||||
__all__ = ["PostgresStore"]
|
||||
|
||||
@@ -1,591 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import AsyncIterator, Callable, Iterable, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from types import TracebackType
|
||||
from typing import Any, cast
|
||||
|
||||
import orjson
|
||||
from langgraph.store.base import (
|
||||
GetOp,
|
||||
ListNamespacesOp,
|
||||
Op,
|
||||
PutOp,
|
||||
Result,
|
||||
SearchOp,
|
||||
)
|
||||
from langgraph.store.base.batch import AsyncBatchedBaseStore
|
||||
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
|
||||
from langgraph.checkpoint.postgres import _ainternal
|
||||
from langgraph.store.postgres.base import (
|
||||
PLACEHOLDER,
|
||||
BasePostgresStore,
|
||||
PoolConfig,
|
||||
PostgresIndexConfig,
|
||||
Row,
|
||||
TTLConfig,
|
||||
_decode_ns_bytes,
|
||||
_ensure_index_config,
|
||||
_group_ops,
|
||||
_row_to_item,
|
||||
_row_to_search_item,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
|
||||
"""Asynchronous Postgres-backed store with optional vector search using pgvector.
|
||||
|
||||
!!! example "Examples"
|
||||
Basic setup and usage:
|
||||
```python
|
||||
from langgraph.store.postgres import AsyncPostgresStore
|
||||
|
||||
conn_string = "postgresql://user:pass@localhost:5432/dbname"
|
||||
|
||||
async with AsyncPostgresStore.from_conn_string(conn_string) as store:
|
||||
await store.setup() # Run migrations. Done once
|
||||
|
||||
# Store and retrieve data
|
||||
await store.aput(("users", "123"), "prefs", {"theme": "dark"})
|
||||
item = await store.aget(("users", "123"), "prefs")
|
||||
```
|
||||
|
||||
Vector search using LangChain embeddings:
|
||||
```python
|
||||
from langchain.embeddings import init_embeddings
|
||||
from langgraph.store.postgres import AsyncPostgresStore
|
||||
|
||||
conn_string = "postgresql://user:pass@localhost:5432/dbname"
|
||||
|
||||
async with AsyncPostgresStore.from_conn_string(
|
||||
conn_string,
|
||||
index={
|
||||
"dims": 1536,
|
||||
"embed": init_embeddings("openai:text-embedding-3-small"),
|
||||
"fields": ["text"] # specify which fields to embed. Default is the whole serialized value
|
||||
}
|
||||
) as store:
|
||||
await store.setup() # Run migrations. Done once
|
||||
|
||||
# Store documents
|
||||
await store.aput(("docs",), "doc1", {"text": "Python tutorial"})
|
||||
await store.aput(("docs",), "doc2", {"text": "TypeScript guide"})
|
||||
await store.aput(("docs",), "doc3", {"text": "Other guide"}, index=False) # don't index
|
||||
|
||||
# Search by similarity
|
||||
results = await store.asearch(("docs",), query="programming guides", limit=2)
|
||||
```
|
||||
|
||||
Using connection pooling for better performance:
|
||||
```python
|
||||
from langgraph.store.postgres import AsyncPostgresStore, PoolConfig
|
||||
|
||||
conn_string = "postgresql://user:pass@localhost:5432/dbname"
|
||||
|
||||
async with AsyncPostgresStore.from_conn_string(
|
||||
conn_string,
|
||||
pool_config=PoolConfig(
|
||||
min_size=5,
|
||||
max_size=20
|
||||
)
|
||||
) as store:
|
||||
await store.setup() # Run migrations. Done once
|
||||
# Use store with connection pooling...
|
||||
```
|
||||
|
||||
Warning:
|
||||
Make sure to:
|
||||
1. Call `setup()` before first use to create necessary tables and indexes
|
||||
2. Have the pgvector extension available to use vector search
|
||||
3. Use Python 3.10+ for async functionality
|
||||
|
||||
Note:
|
||||
Semantic search is disabled by default. You can enable it by providing an `index` configuration
|
||||
when creating the store. Without this configuration, all `index` arguments passed to
|
||||
`put` or `aput` will have no effect.
|
||||
|
||||
Note:
|
||||
If you provide a TTL configuration, you must explicitly call `start_ttl_sweeper()` to begin
|
||||
the background task that removes expired items. Call `stop_ttl_sweeper()` to properly
|
||||
clean up resources when you're done with the store.
|
||||
"""
|
||||
|
||||
__slots__ = (
|
||||
"_deserializer",
|
||||
"pipe",
|
||||
"lock",
|
||||
"supports_pipeline",
|
||||
"index_config",
|
||||
"embeddings",
|
||||
"ttl_config",
|
||||
"_ttl_sweeper_task",
|
||||
"_ttl_stop_event",
|
||||
)
|
||||
supports_ttl: bool = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: _ainternal.Conn,
|
||||
*,
|
||||
pipe: AsyncPipeline | None = None,
|
||||
deserializer: Callable[[bytes | orjson.Fragment], dict[str, Any]] | None = None,
|
||||
index: PostgresIndexConfig | None = None,
|
||||
ttl: TTLConfig | None = 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
|
||||
|
||||
self.ttl_config = ttl
|
||||
self._ttl_sweeper_task: asyncio.Task[None] | None = None
|
||||
self._ttl_stop_event = asyncio.Event()
|
||||
|
||||
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
grouped_ops, num_ops = _group_ops(ops)
|
||||
results: list[Result] = [None] * num_ops
|
||||
|
||||
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
|
||||
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
async def from_conn_string(
|
||||
cls,
|
||||
conn_string: str,
|
||||
*,
|
||||
pipeline: bool = False,
|
||||
pool_config: PoolConfig | None = None,
|
||||
index: PostgresIndexConfig | None = None,
|
||||
ttl: TTLConfig | None = None,
|
||||
) -> AsyncIterator[AsyncPostgresStore]:
|
||||
"""Create a new AsyncPostgresStore instance from a connection string.
|
||||
|
||||
Args:
|
||||
conn_string: The Postgres connection info string.
|
||||
pipeline: Whether to use AsyncPipeline (only for single connections)
|
||||
pool_config: 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: 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, ttl=ttl)
|
||||
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, ttl=ttl)
|
||||
else:
|
||||
yield cls(conn=conn, index=index, ttl=ttl)
|
||||
|
||||
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:
|
||||
await cur.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {table} (
|
||||
v INTEGER PRIMARY KEY
|
||||
)
|
||||
"""
|
||||
)
|
||||
await cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
|
||||
row = cast(dict, await cur.fetchone())
|
||||
if row is None:
|
||||
version = -1
|
||||
else:
|
||||
version = row["v"]
|
||||
return version
|
||||
|
||||
async with self._cursor() as cur:
|
||||
version = await _get_version(cur, table="store_migrations")
|
||||
for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1):
|
||||
await cur.execute(sql)
|
||||
await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
|
||||
|
||||
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()
|
||||
}
|
||||
if "dims" in params:
|
||||
try:
|
||||
params["dims"] = int(params["dims"])
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Invalid dims for vector index: {params['dims']}"
|
||||
) from e
|
||||
if "vector_type" in params:
|
||||
vt = str(params["vector_type"])
|
||||
if vt not in ("vector", "halfvec"):
|
||||
raise ValueError(
|
||||
f"Invalid vector_type for pgvector: {vt}"
|
||||
)
|
||||
params["vector_type"] = vt
|
||||
if "index_type" in params:
|
||||
it = str(params["index_type"])
|
||||
if it not in ("hnsw", "ivfflat"):
|
||||
raise ValueError(
|
||||
f"Invalid index_type for pgvector: {it}"
|
||||
)
|
||||
params["index_type"] = it
|
||||
sql = sql % params
|
||||
await cur.execute(sql)
|
||||
await cur.execute(
|
||||
"INSERT INTO vector_migrations (v) VALUES (%s)", (v,)
|
||||
)
|
||||
|
||||
async def sweep_ttl(self) -> int:
|
||||
"""Delete expired store items based on TTL.
|
||||
|
||||
Returns:
|
||||
int: The number of deleted items.
|
||||
"""
|
||||
async with self._cursor() as cur:
|
||||
await cur.execute(
|
||||
"""
|
||||
DELETE FROM store
|
||||
WHERE expires_at IS NOT NULL AND expires_at < NOW()
|
||||
"""
|
||||
)
|
||||
deleted_count = cur.rowcount
|
||||
return deleted_count
|
||||
|
||||
async def start_ttl_sweeper(
|
||||
self, sweep_interval_minutes: int | None = None
|
||||
) -> asyncio.Task[None]:
|
||||
"""Periodically delete expired store items based on TTL.
|
||||
|
||||
Returns:
|
||||
Task that can be awaited or cancelled.
|
||||
"""
|
||||
if not self.ttl_config:
|
||||
return asyncio.create_task(asyncio.sleep(0))
|
||||
|
||||
if self._ttl_sweeper_task is not None and not self._ttl_sweeper_task.done():
|
||||
return self._ttl_sweeper_task
|
||||
|
||||
self._ttl_stop_event.clear()
|
||||
|
||||
interval = float(
|
||||
sweep_interval_minutes or self.ttl_config.get("sweep_interval_minutes") or 5
|
||||
)
|
||||
logger.info(f"Starting store TTL sweeper with interval {interval} minutes")
|
||||
|
||||
async def _sweep_loop() -> None:
|
||||
while not self._ttl_stop_event.is_set():
|
||||
try:
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._ttl_stop_event.wait(),
|
||||
timeout=interval * 60,
|
||||
)
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
expired_items = await self.sweep_ttl()
|
||||
if expired_items > 0:
|
||||
logger.info(f"Store swept {expired_items} expired items")
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as exc:
|
||||
logger.exception("Store TTL sweep iteration failed", exc_info=exc)
|
||||
|
||||
task = asyncio.create_task(_sweep_loop())
|
||||
task.set_name("ttl_sweeper")
|
||||
self._ttl_sweeper_task = task
|
||||
return task
|
||||
|
||||
async def stop_ttl_sweeper(self, timeout: float | None = None) -> bool:
|
||||
"""Stop the TTL sweeper task if it's running.
|
||||
|
||||
Args:
|
||||
timeout: Maximum time to wait for the task to stop, in seconds.
|
||||
If `None`, wait indefinitely.
|
||||
|
||||
Returns:
|
||||
bool: True if the task was successfully stopped or wasn't running,
|
||||
False if the timeout was reached before the task stopped.
|
||||
"""
|
||||
if self._ttl_sweeper_task is None or self._ttl_sweeper_task.done():
|
||||
return True
|
||||
|
||||
logger.info("Stopping TTL sweeper task")
|
||||
self._ttl_stop_event.set()
|
||||
|
||||
if timeout is not None:
|
||||
try:
|
||||
await asyncio.wait_for(self._ttl_sweeper_task, timeout=timeout)
|
||||
success = True
|
||||
except asyncio.TimeoutError:
|
||||
success = False
|
||||
else:
|
||||
await self._ttl_sweeper_task
|
||||
success = True
|
||||
|
||||
if success:
|
||||
self._ttl_sweeper_task = None
|
||||
logger.info("TTL sweeper task stopped")
|
||||
else:
|
||||
logger.warning("Timed out waiting for TTL sweeper task to stop")
|
||||
|
||||
return success
|
||||
|
||||
async def __aenter__(self) -> AsyncPostgresStore:
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
# Ensure the TTL sweeper task is stopped when exiting the context
|
||||
if hasattr(self, "_ttl_sweeper_task") and self._ttl_sweeper_task is not None:
|
||||
# Set the event to signal the task to stop
|
||||
self._ttl_stop_event.set()
|
||||
# We don't wait for the task to complete here to avoid blocking
|
||||
# The task will clean up itself gracefully
|
||||
|
||||
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:
|
||||
for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
|
||||
await cur.execute(query, params)
|
||||
rows = cast(list[Row], await cur.fetchall())
|
||||
key_to_row = {row["key"]: row for row in rows}
|
||||
for idx, key in items:
|
||||
row = key_to_row.get(key)
|
||||
if row:
|
||||
results[idx] = _row_to_item(
|
||||
namespace, row, loader=self._deserializer
|
||||
)
|
||||
else:
|
||||
results[idx] = None
|
||||
|
||||
async def _batch_put_ops(
|
||||
self,
|
||||
put_ops: Sequence[tuple[int, PutOp]],
|
||||
cur: AsyncCursor[DictRow],
|
||||
) -> None:
|
||||
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, strict=False
|
||||
)
|
||||
for p in (ns, k, pathname, vector)
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
for query, params in queries:
|
||||
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, embedding_requests = self._prepare_batch_search_queries(search_ops)
|
||||
|
||||
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, strict=False):
|
||||
_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, strict=False):
|
||||
await cur.execute(query, params)
|
||||
rows = cast(list[Row], await cur.fetchall())
|
||||
items = [
|
||||
_row_to_search_item(
|
||||
_decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
results[idx] = items
|
||||
|
||||
async def _batch_list_namespaces_ops(
|
||||
self,
|
||||
list_ops: Sequence[tuple[int, ListNamespacesOp]],
|
||||
results: list[Result],
|
||||
cur: AsyncCursor[DictRow],
|
||||
) -> None:
|
||||
queries = self._get_batch_list_namespaces_queries(list_ops)
|
||||
for (query, params), (idx, _) in zip(queries, list_ops, strict=False):
|
||||
await cur.execute(query, params)
|
||||
rows = cast(list[dict], await cur.fetchall())
|
||||
namespaces = [_decode_ns_bytes(row["truncated_prefix"]) for row in rows]
|
||||
results[idx] = namespaces
|
||||
|
||||
@asynccontextmanager
|
||||
async def _cursor(
|
||||
self, *, pipeline: bool = False
|
||||
) -> AsyncIterator[AsyncCursor[DictRow]]:
|
||||
"""Create a database cursor as a context manager.
|
||||
|
||||
Args:
|
||||
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.
|
||||
"""
|
||||
is_pooled_conn = isinstance(self.conn, AsyncConnectionPool)
|
||||
# With AsyncConnectionPool, each _cursor() call checks out its own connection.
|
||||
# The pool does not hand out the same connection concurrently, so a shared lock
|
||||
# across calls is unnecessary here.
|
||||
lock = asyncio.Lock() if is_pooled_conn else self.lock
|
||||
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 (
|
||||
lock,
|
||||
conn.pipeline(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
async with (
|
||||
lock,
|
||||
conn.transaction(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
async with (
|
||||
lock,
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
@@ -1,27 +1,34 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable, Iterable, Iterator, Sequence
|
||||
from collections.abc import Iterable, Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Callable,
|
||||
Generic,
|
||||
Literal,
|
||||
NamedTuple,
|
||||
Optional,
|
||||
TypeVar,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
|
||||
import orjson
|
||||
from psycopg import Capabilities, Connection, Cursor, Pipeline
|
||||
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 _internal as _pg_internal
|
||||
from langgraph.store.base import (
|
||||
BaseStore,
|
||||
Embeddings,
|
||||
GetOp,
|
||||
IndexConfig,
|
||||
Item,
|
||||
@@ -31,22 +38,10 @@ from langgraph.store.base import (
|
||||
Result,
|
||||
SearchItem,
|
||||
SearchOp,
|
||||
TTLConfig,
|
||||
ensure_embeddings,
|
||||
get_text_at_path,
|
||||
tokenize_path,
|
||||
)
|
||||
from psycopg import Capabilities, Connection, Cursor, Pipeline
|
||||
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
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain_core.embeddings import Embeddings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -55,8 +50,8 @@ class Migration(NamedTuple):
|
||||
"""A database migration with optional conditions and parameters."""
|
||||
|
||||
sql: str
|
||||
params: dict[str, Any] | None = None
|
||||
condition: Callable[[BasePostgresStore], bool] | None = None
|
||||
params: Optional[dict[str, Any]] = None
|
||||
condition: Optional[Callable[["BasePostgresStore"], bool]] = None
|
||||
|
||||
|
||||
MIGRATIONS: Sequence[str] = [
|
||||
@@ -74,29 +69,13 @@ CREATE TABLE IF NOT EXISTS store (
|
||||
"""
|
||||
-- For faster lookups by prefix
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
|
||||
""",
|
||||
"""
|
||||
-- Add expires_at column to store table
|
||||
ALTER TABLE store
|
||||
ADD COLUMN IF NOT EXISTS expires_at TIMESTAMP WITH TIME ZONE,
|
||||
ADD COLUMN IF NOT EXISTS ttl_minutes INT;
|
||||
""",
|
||||
"""
|
||||
-- Add indexes for efficient TTL sweeping
|
||||
CREATE INDEX IF NOT EXISTS idx_store_expires_at ON store (expires_at)
|
||||
WHERE expires_at IS NOT NULL;
|
||||
""",
|
||||
]
|
||||
|
||||
VECTOR_MIGRATIONS: Sequence[Migration] = [
|
||||
Migration(
|
||||
"""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'vector') THEN
|
||||
CREATE EXTENSION vector;
|
||||
END IF;
|
||||
END $$;
|
||||
CREATE EXTENSION IF NOT EXISTS vector;
|
||||
""",
|
||||
),
|
||||
Migration(
|
||||
@@ -144,14 +123,13 @@ CREATE INDEX CONCURRENTLY IF NOT EXISTS store_vectors_embedding_idx ON store_vec
|
||||
]
|
||||
|
||||
|
||||
C = TypeVar("C", bound=_pg_internal.Conn | _ainternal.Conn)
|
||||
C = TypeVar("C", bound=_pg_internal.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
|
||||
@@ -160,14 +138,13 @@ class PoolConfig(TypedDict, total=False):
|
||||
min_size: int
|
||||
"""Minimum number of connections maintained in the pool. Defaults to 1."""
|
||||
|
||||
max_size: int | None
|
||||
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
|
||||
@@ -236,62 +213,27 @@ class BasePostgresStore(Generic[C]):
|
||||
MIGRATIONS = MIGRATIONS
|
||||
VECTOR_MIGRATIONS = VECTOR_MIGRATIONS
|
||||
conn: C
|
||||
_deserializer: Callable[[bytes | orjson.Fragment], dict[str, Any]] | None
|
||||
index_config: PostgresIndexConfig | None
|
||||
_deserializer: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]]
|
||||
index_config: Optional[PostgresIndexConfig]
|
||||
|
||||
def _get_batch_GET_ops_queries(
|
||||
self,
|
||||
get_ops: Sequence[tuple[int, GetOp]],
|
||||
) -> list[tuple[str, tuple, tuple[str, ...], list]]:
|
||||
"""
|
||||
Build queries to fetch (and optionally refresh the TTL of) multiple keys per namespace.
|
||||
|
||||
Each returned element is a tuple of:
|
||||
(sql_query_string, sql_params, namespace, items_for_this_namespace)
|
||||
|
||||
where items_for_this_namespace is the original list of (idx, key, refresh_ttl).
|
||||
"""
|
||||
|
||||
namespace_groups = defaultdict(list)
|
||||
refresh_ttls = defaultdict(list)
|
||||
for idx, op in get_ops:
|
||||
namespace_groups[op.namespace].append((idx, op.key))
|
||||
refresh_ttls[op.namespace].append(op.refresh_ttl)
|
||||
|
||||
results = []
|
||||
for namespace, items in namespace_groups.items():
|
||||
_, keys = zip(*items, strict=False)
|
||||
this_refresh_ttls = refresh_ttls[namespace]
|
||||
|
||||
query = """
|
||||
WITH passed_in AS (
|
||||
SELECT unnest(%s::text[]) AS key,
|
||||
unnest(%s::bool[]) AS do_refresh
|
||||
),
|
||||
updated AS (
|
||||
UPDATE store s
|
||||
SET expires_at = NOW() + (s.ttl_minutes || ' minutes')::interval
|
||||
FROM passed_in p
|
||||
WHERE s.prefix = %s
|
||||
AND s.key = p.key
|
||||
AND p.do_refresh = TRUE
|
||||
AND s.ttl_minutes IS NOT NULL
|
||||
RETURNING s.key
|
||||
)
|
||||
SELECT s.key, s.value, s.created_at, s.updated_at
|
||||
FROM store s
|
||||
JOIN passed_in p ON s.key = p.key
|
||||
WHERE s.prefix = %s
|
||||
_, keys = zip(*items)
|
||||
keys_to_query = ",".join(["%s"] * len(keys))
|
||||
query = f"""
|
||||
SELECT key, value, created_at, updated_at
|
||||
FROM store
|
||||
WHERE prefix = %s AND key IN ({keys_to_query})
|
||||
"""
|
||||
ns_text = _namespace_to_text(namespace)
|
||||
params = (
|
||||
list(keys), # -> unnest(%s::text[])
|
||||
list(this_refresh_ttls), # -> unnest(%s::bool[])
|
||||
ns_text, # -> prefix = %s (for UPDATE)
|
||||
ns_text, # -> prefix = %s (for final SELECT)
|
||||
)
|
||||
params = (_namespace_to_text(namespace), *keys)
|
||||
results.append((query, params, namespace, items))
|
||||
|
||||
return results
|
||||
|
||||
def _prepare_batch_PUT_queries(
|
||||
@@ -299,8 +241,9 @@ class BasePostgresStore(Generic[C]):
|
||||
put_ops: Sequence[tuple[int, PutOp]],
|
||||
) -> tuple[
|
||||
list[tuple[str, Sequence]],
|
||||
tuple[str, Sequence[tuple[str, str, str, str]]] | None,
|
||||
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
|
||||
@@ -326,39 +269,25 @@ class BasePostgresStore(Generic[C]):
|
||||
)
|
||||
params = (_namespace_to_text(namespace), *keys)
|
||||
queries.append((query, params))
|
||||
embedding_request: tuple[str, Sequence[tuple[str, str, str, str]]] | None = None
|
||||
embedding_request: Optional[tuple[str, Sequence[tuple[str, str, str, str]]]] = (
|
||||
None
|
||||
)
|
||||
if inserts:
|
||||
values = []
|
||||
insertion_params: list[Any] = []
|
||||
insertion_params = []
|
||||
vector_values = []
|
||||
embedding_request_params = []
|
||||
# Handle TTL expiration
|
||||
|
||||
# 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(cast(dict, op.value)),
|
||||
)
|
||||
]
|
||||
)
|
||||
if op.ttl is not None:
|
||||
values.append(
|
||||
"(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, NOW() + %s::interval, %s)"
|
||||
)
|
||||
ttl_minutes = float(op.ttl)
|
||||
insertion_params.extend(
|
||||
(
|
||||
f"{ttl_minutes * 60} seconds",
|
||||
ttl_minutes,
|
||||
)
|
||||
)
|
||||
else:
|
||||
values.append(
|
||||
"(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, NULL, %s)"
|
||||
)
|
||||
insertion_params.append(None)
|
||||
|
||||
# Then handle embeddings if configured
|
||||
if self.index_config:
|
||||
@@ -370,7 +299,7 @@ class BasePostgresStore(Generic[C]):
|
||||
k = op.key
|
||||
|
||||
if op.index is None:
|
||||
paths = cast(dict, self.index_config)["__tokenized_fields"]
|
||||
paths = self.index_config["__tokenized_fields"]
|
||||
else:
|
||||
paths = [(ix, tokenize_path(ix)) for ix in op.index]
|
||||
|
||||
@@ -385,13 +314,11 @@ class BasePostgresStore(Generic[C]):
|
||||
|
||||
values_str = ",".join(values)
|
||||
query = f"""
|
||||
INSERT INTO store (prefix, key, value, created_at, updated_at, expires_at, ttl_minutes)
|
||||
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,
|
||||
expires_at = EXCLUDED.expires_at,
|
||||
ttl_minutes = EXCLUDED.ttl_minutes
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
"""
|
||||
queries.append((query, insertion_params))
|
||||
|
||||
@@ -412,156 +339,120 @@ class BasePostgresStore(Generic[C]):
|
||||
self,
|
||||
search_ops: Sequence[tuple[int, SearchOp]],
|
||||
) -> tuple[
|
||||
list[tuple[str, list[None | str | list[float]]]], # queries, params
|
||||
list[tuple[str, list[Union[None, str, list[float]]]]], # queries, params
|
||||
list[tuple[int, str]], # idx, query_text pairs to embed
|
||||
]:
|
||||
"""
|
||||
Build per-SearchOp SQL queries (with optional TTL refresh) plus embedding requests.
|
||||
Returns:
|
||||
- queries: list of (SQL, param_list)
|
||||
- embedding_requests: list of (original_index_in_search_ops, text_query)
|
||||
"""
|
||||
|
||||
queries = []
|
||||
embedding_requests = []
|
||||
|
||||
for idx, (_, op) in enumerate(search_ops):
|
||||
# Build filter conditions first
|
||||
filter_params = []
|
||||
filter_clauses = []
|
||||
filter_conditions = []
|
||||
if op.filter:
|
||||
for key, value in op.filter.items():
|
||||
if isinstance(value, dict):
|
||||
for op_name, val in value.items():
|
||||
condition, params_ = self._get_filter_condition(
|
||||
condition, filter_params_ = self._get_filter_condition(
|
||||
key, op_name, val
|
||||
)
|
||||
filter_clauses.append(condition)
|
||||
filter_params.extend(params_)
|
||||
filter_conditions.append(condition)
|
||||
filter_params.extend(filter_params_)
|
||||
else:
|
||||
filter_clauses.append("value->%s = %s::jsonb")
|
||||
filter_params.extend([key, orjson.dumps(value).decode("utf-8")])
|
||||
|
||||
ns_condition = "TRUE"
|
||||
ns_param: Sequence[str] | None = None
|
||||
if op.namespace_prefix:
|
||||
ns_condition = "store.prefix LIKE %s"
|
||||
ns_param = (f"{_namespace_to_text(op.namespace_prefix)}%",)
|
||||
else:
|
||||
ns_param = ()
|
||||
|
||||
extra_filters = (
|
||||
" AND " + " AND ".join(filter_clauses) if filter_clauses else ""
|
||||
)
|
||||
filter_conditions.append("value->%s = %s::jsonb")
|
||||
filter_params.extend([key, json.dumps(value)])
|
||||
|
||||
# Vector search branch
|
||||
if op.query and self.index_config:
|
||||
# We'll embed the text later, so record the request.
|
||||
embedding_requests.append((idx, op.query))
|
||||
|
||||
score_operator, post_operator = get_distance_operator(self)
|
||||
post_operator = post_operator.replace("scored", "uniq")
|
||||
vector_type = self.index_config.get("ann_index_config", {}).get(
|
||||
"vector_type", "vector"
|
||||
score_operator, post_operator = _get_distance_operator(self)
|
||||
vector_type = (
|
||||
cast(PostgresIndexConfig, self.index_config)
|
||||
.get("ann_index_config", {})
|
||||
.get("vector_type", "vector")
|
||||
)
|
||||
|
||||
# For hamming bit vectors, or “regular” vectors
|
||||
if (
|
||||
vector_type == "bit"
|
||||
and cast(dict, self.index_config).get("distance_type") == "hamming"
|
||||
and self.index_config.get("distance_type") == "hamming"
|
||||
):
|
||||
score_operator = score_operator % (
|
||||
"%s",
|
||||
cast(dict, self.index_config)["dims"],
|
||||
self.index_config["dims"],
|
||||
)
|
||||
else:
|
||||
if vector_type not in ("vector", "halfvec"):
|
||||
raise ValueError(
|
||||
f"Invalid vector_type for pgvector: {vector_type}"
|
||||
)
|
||||
score_operator = score_operator % ("%s", vector_type)
|
||||
score_operator = score_operator % (
|
||||
"%s",
|
||||
vector_type,
|
||||
)
|
||||
|
||||
vectors_per_doc_estimate = cast(dict, self.index_config)[
|
||||
"__estimated_num_vectors"
|
||||
]
|
||||
vectors_per_doc_estimate = self.index_config["__estimated_num_vectors"]
|
||||
expanded_limit = (op.limit * vectors_per_doc_estimate * 2) + 1
|
||||
|
||||
# “sub_scored” does the main vector search
|
||||
# Then we do DISTINCT ON to drop duplicates if your store can have them
|
||||
# Finally we limit & offset
|
||||
vector_search_cte = f"""
|
||||
SELECT store.prefix, store.key, store.value, store.created_at, store.updated_at,
|
||||
{score_operator} AS neg_score
|
||||
FROM store
|
||||
JOIN store_vectors sv ON store.prefix = sv.prefix AND store.key = sv.key
|
||||
WHERE {ns_condition} {extra_filters}
|
||||
ORDER BY {score_operator} ASC
|
||||
LIMIT %s
|
||||
"""
|
||||
# Vector search with CTE for proper score handling
|
||||
filter_str = (
|
||||
""
|
||||
if not filter_conditions
|
||||
else " AND " + " AND ".join(filter_conditions)
|
||||
)
|
||||
if op.namespace_prefix:
|
||||
prefix_filter_str = f"WHERE s.prefix LIKE %s {filter_str} "
|
||||
ns_args: Sequence = (f"{_namespace_to_text(op.namespace_prefix)}%",)
|
||||
else:
|
||||
ns_args = ()
|
||||
if filter_str:
|
||||
prefix_filter_str = f"WHERE {filter_str} "
|
||||
else:
|
||||
prefix_filter_str = ""
|
||||
|
||||
search_results_sql = f"""
|
||||
WITH scored AS (
|
||||
{vector_search_cte}
|
||||
)
|
||||
SELECT uniq.prefix, uniq.key, uniq.value, uniq.created_at, uniq.updated_at,
|
||||
{post_operator} AS score
|
||||
FROM (
|
||||
SELECT DISTINCT ON (scored.prefix, scored.key)
|
||||
scored.prefix, scored.key, scored.value, scored.created_at, scored.updated_at, scored.neg_score
|
||||
FROM scored
|
||||
ORDER BY scored.prefix, scored.key, scored.neg_score ASC
|
||||
) uniq
|
||||
ORDER BY score DESC
|
||||
base_query = f"""
|
||||
WITH scored AS (
|
||||
SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, {score_operator} AS neg_score
|
||||
FROM store s
|
||||
JOIN store_vectors sv ON s.prefix = sv.prefix AND s.key = sv.key
|
||||
{prefix_filter_str}
|
||||
ORDER BY {score_operator} ASC
|
||||
LIMIT %s
|
||||
OFFSET %s
|
||||
"""
|
||||
|
||||
search_results_params = [
|
||||
PLACEHOLDER,
|
||||
*ns_param,
|
||||
)
|
||||
SELECT * FROM (
|
||||
SELECT DISTINCT ON (prefix, key)
|
||||
prefix, key, value, created_at, updated_at, {post_operator} as score
|
||||
FROM scored
|
||||
ORDER BY prefix, key, score DESC
|
||||
) AS unique_docs
|
||||
ORDER BY score DESC
|
||||
LIMIT %s
|
||||
OFFSET %s
|
||||
"""
|
||||
params = [
|
||||
_PLACEHOLDER, # Vector placeholder
|
||||
*ns_args,
|
||||
*filter_params,
|
||||
PLACEHOLDER,
|
||||
_PLACEHOLDER,
|
||||
expanded_limit,
|
||||
op.limit,
|
||||
op.offset,
|
||||
]
|
||||
|
||||
# Regular search branch
|
||||
else:
|
||||
base_query = f"""
|
||||
SELECT store.prefix, store.key, store.value, store.created_at, store.updated_at, NULL AS score
|
||||
FROM store
|
||||
WHERE {ns_condition} {extra_filters}
|
||||
ORDER BY store.updated_at DESC
|
||||
LIMIT %s
|
||||
OFFSET %s
|
||||
"""
|
||||
search_results_sql = base_query
|
||||
search_results_params = [
|
||||
*ns_param,
|
||||
*filter_params,
|
||||
op.limit,
|
||||
op.offset,
|
||||
]
|
||||
base_query = """
|
||||
SELECT prefix, key, value, created_at, updated_at
|
||||
FROM store
|
||||
WHERE prefix LIKE %s
|
||||
"""
|
||||
params = [f"{_namespace_to_text(op.namespace_prefix)}%"]
|
||||
|
||||
if op.refresh_ttl:
|
||||
# Wrap entire primary query in a CTE, then perform "update_at"
|
||||
final_sql = f"""
|
||||
WITH search_results AS (
|
||||
{search_results_sql}
|
||||
),
|
||||
updated AS (
|
||||
UPDATE store s
|
||||
SET expires_at = NOW() + (s.ttl_minutes || ' minutes')::interval
|
||||
FROM search_results sr
|
||||
WHERE s.prefix = sr.prefix
|
||||
AND s.key = sr.key
|
||||
AND s.ttl_minutes IS NOT NULL
|
||||
)
|
||||
SELECT sr.prefix, sr.key, sr.value, sr.created_at, sr.updated_at, sr.score
|
||||
FROM search_results sr
|
||||
"""
|
||||
final_params = search_results_params[:] # copy
|
||||
else:
|
||||
final_sql = search_results_sql
|
||||
final_params = search_results_params
|
||||
queries.append((final_sql, final_params))
|
||||
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
|
||||
|
||||
@@ -571,7 +462,7 @@ class BasePostgresStore(Generic[C]):
|
||||
) -> list[tuple[str, Sequence]]:
|
||||
queries: list[tuple[str, Sequence]] = []
|
||||
for _, op in list_ops:
|
||||
query = r"""
|
||||
query = """
|
||||
SELECT DISTINCT ON (truncated_prefix) truncated_prefix, prefix
|
||||
FROM (
|
||||
SELECT
|
||||
@@ -658,8 +549,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
item = store.get(("users", "123"), "prefs")
|
||||
```
|
||||
|
||||
Or using the convenient `from_conn_string` helper:
|
||||
|
||||
Or using the convenient from_conn_string helper:
|
||||
```python
|
||||
from langgraph.store.postgres import PostgresStore
|
||||
|
||||
@@ -696,7 +586,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
store.put(("docs",), "doc2", {"text": "Other guide"}, index=False) # don't index
|
||||
|
||||
# Search by similarity
|
||||
results = store.search(("docs",), query="programming guides", limit=2)
|
||||
results = store.search(("docs",), "programming guides", limit=2)
|
||||
```
|
||||
|
||||
Note:
|
||||
@@ -708,11 +598,6 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
Make sure to call `setup()` before first use to create necessary tables and indexes.
|
||||
The pgvector extension must be available to use vector search.
|
||||
|
||||
Note:
|
||||
If you provide a TTL configuration, you must explicitly call `start_ttl_sweeper()` to begin
|
||||
the background thread that removes expired items. Call `stop_ttl_sweeper()` to properly
|
||||
clean up resources when you're done with the store.
|
||||
|
||||
"""
|
||||
|
||||
__slots__ = (
|
||||
@@ -722,19 +607,17 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
"supports_pipeline",
|
||||
"index_config",
|
||||
"embeddings",
|
||||
"_ttl_sweeper_thread",
|
||||
"_ttl_stop_event",
|
||||
)
|
||||
supports_ttl: bool = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: _pg_internal.Conn,
|
||||
*,
|
||||
pipe: Pipeline | None = None,
|
||||
deserializer: Callable[[bytes | orjson.Fragment], dict[str, Any]] | None = None,
|
||||
index: PostgresIndexConfig | None = None,
|
||||
ttl: TTLConfig | None = None,
|
||||
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
|
||||
@@ -747,9 +630,6 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
self.embeddings, self.index_config = _ensure_index_config(self.index_config)
|
||||
else:
|
||||
self.embeddings = None
|
||||
self.ttl_config = ttl
|
||||
self._ttl_sweeper_thread: threading.Thread | None = None
|
||||
self._ttl_stop_event = threading.Event()
|
||||
|
||||
@classmethod
|
||||
@contextmanager
|
||||
@@ -758,20 +638,18 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
conn_string: str,
|
||||
*,
|
||||
pipeline: bool = False,
|
||||
pool_config: PoolConfig | None = None,
|
||||
index: PostgresIndexConfig | None = None,
|
||||
ttl: TTLConfig | None = None,
|
||||
) -> Iterator[PostgresStore]:
|
||||
pool_config: Optional[PoolConfig] = None,
|
||||
index: Optional[PostgresIndexConfig] = None,
|
||||
) -> Iterator["PostgresStore"]:
|
||||
"""Create a new PostgresStore instance from a connection string.
|
||||
|
||||
Args:
|
||||
conn_string: The Postgres connection info string.
|
||||
pipeline: whether to use Pipeline
|
||||
pool_config: Configuration for the connection pool.
|
||||
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: The index configuration for the store.
|
||||
ttl: The TTL configuration for the store.
|
||||
index (Optional[PostgresIndexConfig]): The index configuration for the store.
|
||||
|
||||
Returns:
|
||||
PostgresStore: A new PostgresStore instance.
|
||||
@@ -793,130 +671,23 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
**cast(dict, pc),
|
||||
),
|
||||
) as pool:
|
||||
yield cls(conn=pool, index=index, ttl=ttl)
|
||||
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, ttl=ttl)
|
||||
yield cls(conn, pipe=pipe, index=index)
|
||||
else:
|
||||
yield cls(conn, index=index, ttl=ttl)
|
||||
|
||||
def sweep_ttl(self) -> int:
|
||||
"""Delete expired store items based on TTL.
|
||||
|
||||
Returns:
|
||||
int: The number of deleted items.
|
||||
"""
|
||||
with self._cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
DELETE FROM store
|
||||
WHERE expires_at IS NOT NULL AND expires_at < NOW()
|
||||
"""
|
||||
)
|
||||
deleted_count = cur.rowcount
|
||||
return deleted_count
|
||||
|
||||
def start_ttl_sweeper(
|
||||
self, sweep_interval_minutes: int | None = None
|
||||
) -> concurrent.futures.Future[None]:
|
||||
"""Periodically delete expired store items based on TTL.
|
||||
|
||||
Returns:
|
||||
Future that can be waited on or cancelled.
|
||||
"""
|
||||
if not self.ttl_config:
|
||||
future: concurrent.futures.Future[None] = concurrent.futures.Future()
|
||||
future.set_result(None)
|
||||
return future
|
||||
|
||||
if self._ttl_sweeper_thread and self._ttl_sweeper_thread.is_alive():
|
||||
logger.info("TTL sweeper thread is already running")
|
||||
# Return a future that can be used to cancel the existing thread
|
||||
future = concurrent.futures.Future()
|
||||
future.add_done_callback(
|
||||
lambda f: self._ttl_stop_event.set() if f.cancelled() else None
|
||||
)
|
||||
return future
|
||||
|
||||
self._ttl_stop_event.clear()
|
||||
|
||||
interval = float(
|
||||
sweep_interval_minutes or self.ttl_config.get("sweep_interval_minutes") or 5
|
||||
)
|
||||
logger.info(f"Starting store TTL sweeper with interval {interval} minutes")
|
||||
|
||||
future = concurrent.futures.Future()
|
||||
|
||||
def _sweep_loop() -> None:
|
||||
try:
|
||||
while not self._ttl_stop_event.is_set():
|
||||
if self._ttl_stop_event.wait(interval * 60):
|
||||
break
|
||||
|
||||
try:
|
||||
expired_items = self.sweep_ttl()
|
||||
if expired_items > 0:
|
||||
logger.info(f"Store swept {expired_items} expired items")
|
||||
except Exception as exc:
|
||||
logger.exception(
|
||||
"Store TTL sweep iteration failed", exc_info=exc
|
||||
)
|
||||
future.set_result(None)
|
||||
except Exception as exc:
|
||||
future.set_exception(exc)
|
||||
|
||||
thread = threading.Thread(target=_sweep_loop, daemon=True, name="ttl-sweeper")
|
||||
self._ttl_sweeper_thread = thread
|
||||
thread.start()
|
||||
|
||||
future.add_done_callback(
|
||||
lambda f: self._ttl_stop_event.set() if f.cancelled() else None
|
||||
)
|
||||
return future
|
||||
|
||||
def stop_ttl_sweeper(self, timeout: float | None = None) -> bool:
|
||||
"""Stop the TTL sweeper thread if it's running.
|
||||
|
||||
Args:
|
||||
timeout: Maximum time to wait for the thread to stop, in seconds.
|
||||
If `None`, wait indefinitely.
|
||||
|
||||
Returns:
|
||||
bool: True if the thread was successfully stopped or wasn't running,
|
||||
False if the timeout was reached before the thread stopped.
|
||||
"""
|
||||
if not self._ttl_sweeper_thread or not self._ttl_sweeper_thread.is_alive():
|
||||
return True
|
||||
|
||||
logger.info("Stopping TTL sweeper thread")
|
||||
self._ttl_stop_event.set()
|
||||
|
||||
self._ttl_sweeper_thread.join(timeout)
|
||||
success = not self._ttl_sweeper_thread.is_alive()
|
||||
|
||||
if success:
|
||||
self._ttl_sweeper_thread = None
|
||||
logger.info("TTL sweeper thread stopped")
|
||||
else:
|
||||
logger.warning("Timed out waiting for TTL sweeper thread to stop")
|
||||
|
||||
return success
|
||||
|
||||
def __del__(self) -> None:
|
||||
"""Ensure the TTL sweeper thread is stopped when the object is garbage collected."""
|
||||
if hasattr(self, "_ttl_stop_event") and hasattr(self, "_ttl_sweeper_thread"):
|
||||
self.stop_ttl_sweeper(timeout=0.1)
|
||||
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: whether to use pipeline for the DB operations inside the context manager.
|
||||
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.
|
||||
"""
|
||||
@@ -1029,9 +800,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
query,
|
||||
[
|
||||
p
|
||||
for (ns, k, pathname, _), vector in zip(
|
||||
txt_params, vectors, strict=False
|
||||
)
|
||||
for (ns, k, pathname, _), vector in zip(txt_params, vectors)
|
||||
for p in (ns, k, pathname, vector)
|
||||
],
|
||||
)
|
||||
@@ -1052,15 +821,13 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
embeddings = self.embeddings.embed_documents(
|
||||
[query for _, query in embedding_requests]
|
||||
)
|
||||
for (idx, _), embedding in zip(
|
||||
embedding_requests, embeddings, strict=False
|
||||
):
|
||||
for (idx, _), embedding in zip(embedding_requests, embeddings):
|
||||
_paramslist = queries[idx][1]
|
||||
for i in range(len(_paramslist)):
|
||||
if _paramslist[i] is PLACEHOLDER:
|
||||
if _paramslist[i] is _PLACEHOLDER:
|
||||
_paramslist[i] = embedding
|
||||
|
||||
for (idx, _), (query, params) in zip(search_ops, queries, strict=False):
|
||||
for (idx, _), (query, params) in zip(search_ops, queries):
|
||||
cur.execute(query, params)
|
||||
rows = cast(list[Row], cur.fetchall())
|
||||
results[idx] = [
|
||||
@@ -1077,14 +844,11 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
cur: Cursor[DictRow],
|
||||
) -> None:
|
||||
for (query, params), (idx, _) in zip(
|
||||
self._get_batch_list_namespaces_queries(list_ops), list_ops, strict=False
|
||||
self._get_batch_list_namespaces_queries(list_ops), list_ops
|
||||
):
|
||||
cur.execute(query, params)
|
||||
results[idx] = [_decode_ns_bytes(row["truncated_prefix"]) for row in cur]
|
||||
|
||||
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.
|
||||
|
||||
@@ -1112,14 +876,8 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
with self._cursor() as cur:
|
||||
version = _get_version(cur, table="store_migrations")
|
||||
for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1):
|
||||
try:
|
||||
cur.execute(sql)
|
||||
cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to apply migration {v}.\nSql={sql}\nError={e}"
|
||||
)
|
||||
raise
|
||||
cur.execute(sql)
|
||||
cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
|
||||
|
||||
if self.index_config:
|
||||
version = _get_version(cur, table="vector_migrations")
|
||||
@@ -1134,27 +892,6 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
k: v(self) if v is not None and callable(v) else v
|
||||
for k, v in migration.params.items()
|
||||
}
|
||||
if "dims" in params:
|
||||
try:
|
||||
params["dims"] = int(params["dims"])
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Invalid dims for vector index: {params['dims']}"
|
||||
) from e
|
||||
if "vector_type" in params:
|
||||
vt = str(params["vector_type"])
|
||||
if vt not in ("vector", "halfvec"):
|
||||
raise ValueError(
|
||||
f"Invalid vector_type for pgvector: {vt}"
|
||||
)
|
||||
params["vector_type"] = vt
|
||||
if "index_type" in params:
|
||||
it = str(params["index_type"])
|
||||
if it not in ("hnsw", "ivfflat"):
|
||||
raise ValueError(
|
||||
f"Invalid index_type for pgvector: {it}"
|
||||
)
|
||||
params["index_type"] = it
|
||||
sql = sql % params
|
||||
cur.execute(sql)
|
||||
cur.execute("INSERT INTO vector_migrations (v) VALUES (%s)", (v,))
|
||||
@@ -1180,7 +917,7 @@ def _get_vector_type_ops(store: BasePostgresStore) -> str:
|
||||
if not store.index_config:
|
||||
return "vector_cosine_ops"
|
||||
|
||||
config = store.index_config
|
||||
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"):
|
||||
@@ -1208,44 +945,15 @@ def _get_vector_type_ops(store: BasePostgresStore) -> str:
|
||||
|
||||
|
||||
def _get_index_params(store: Any) -> tuple[str, dict[str, Any]]:
|
||||
"""Get a sanitized index type and configuration based on config.
|
||||
|
||||
Only allow known-safe kinds and integer parameters to avoid SQL injection
|
||||
when constructing DDL strings for index creation.
|
||||
"""
|
||||
"""Get the index type and configuration based on config."""
|
||||
if not store.index_config:
|
||||
return "hnsw", {}
|
||||
|
||||
config = cast(PostgresIndexConfig, store.index_config)
|
||||
raw = config.get("ann_index_config", _DEFAULT_ANN_CONFIG).copy()
|
||||
|
||||
kind = str(raw.pop("kind", "hnsw"))
|
||||
if kind not in ("hnsw", "ivfflat", "flat"):
|
||||
raise ValueError(
|
||||
f"Invalid index kind for pgvector: {kind}. Expected 'hnsw', 'ivfflat', or 'flat'."
|
||||
)
|
||||
|
||||
raw.pop("vector_type", None)
|
||||
|
||||
if kind == "hnsw":
|
||||
allowed_keys = {"m", "ef_construction"}
|
||||
else: # ivfflat/flat
|
||||
allowed_keys = {"lists", "nlist"}
|
||||
|
||||
sanitized: dict[str, int] = {}
|
||||
for k, v in list(raw.items()):
|
||||
if k not in allowed_keys:
|
||||
continue
|
||||
key = "lists" if k == "nlist" else k
|
||||
try:
|
||||
ivalue = int(v) # type: ignore[call-overload]
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid index parameter value for {k}: {v}") from e
|
||||
if ivalue <= 0:
|
||||
continue
|
||||
sanitized[key] = ivalue
|
||||
|
||||
return kind, sanitized
|
||||
index_config = config.get("ann_index_config", _DEFAULT_ANN_CONFIG).copy()
|
||||
kind = index_config.pop("kind", "hnsw")
|
||||
index_config.pop("vector_type", None)
|
||||
return kind, index_config
|
||||
|
||||
|
||||
def _namespace_to_text(
|
||||
@@ -1261,7 +969,7 @@ def _row_to_item(
|
||||
namespace: tuple[str, ...],
|
||||
row: Row,
|
||||
*,
|
||||
loader: Callable[[bytes | orjson.Fragment], dict[str, Any]] | None = None,
|
||||
loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None,
|
||||
) -> Item:
|
||||
"""Convert a row from the database into an Item.
|
||||
|
||||
@@ -1289,7 +997,7 @@ def _row_to_search_item(
|
||||
namespace: tuple[str, ...],
|
||||
row: Row,
|
||||
*,
|
||||
loader: Callable[[bytes | orjson.Fragment], dict[str, Any]] | None = None,
|
||||
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
|
||||
@@ -1320,7 +1028,7 @@ def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int
|
||||
return grouped_ops, tot
|
||||
|
||||
|
||||
def _json_loads(content: bytes | orjson.Fragment) -> Any:
|
||||
def _json_loads(content: Union[bytes, orjson.Fragment]) -> Any:
|
||||
if isinstance(content, orjson.Fragment):
|
||||
if hasattr(content, "buf"):
|
||||
content = content.buf
|
||||
@@ -1332,7 +1040,7 @@ def _json_loads(content: bytes | orjson.Fragment) -> Any:
|
||||
return orjson.loads(cast(bytes, content))
|
||||
|
||||
|
||||
def _decode_ns_bytes(namespace: str | bytes | list) -> tuple[str, ...]:
|
||||
def _decode_ns_bytes(namespace: Union[str, bytes, list]) -> tuple[str, ...]:
|
||||
if isinstance(namespace, list):
|
||||
return tuple(namespace)
|
||||
if isinstance(namespace, bytes):
|
||||
@@ -1340,7 +1048,7 @@ def _decode_ns_bytes(namespace: str | bytes | list) -> tuple[str, ...]:
|
||||
return tuple(namespace.split("."))
|
||||
|
||||
|
||||
def get_distance_operator(store: Any) -> tuple[str, str]:
|
||||
def _get_distance_operator(store: Any) -> tuple[str, 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
|
||||
@@ -1381,16 +1089,16 @@ def get_distance_operator(store: Any) -> tuple[str, str]:
|
||||
|
||||
def _ensure_index_config(
|
||||
index_config: PostgresIndexConfig,
|
||||
) -> tuple[Embeddings | None, PostgresIndexConfig]:
|
||||
) -> tuple[Optional["Embeddings"], PostgresIndexConfig]:
|
||||
index_config = index_config.copy()
|
||||
tokenized: list[tuple[str, Literal["$"] | list[str]]] = []
|
||||
tokenized: list[tuple[str, Union[Literal["$"], list[str]]]] = []
|
||||
tot = 0
|
||||
fields = index_config.get("fields") or ["$"]
|
||||
if isinstance(fields, str):
|
||||
fields = [fields]
|
||||
if not isinstance(fields, list):
|
||||
raise ValueError(f"Text fields must be a list or a string. Got {fields}")
|
||||
for p in fields:
|
||||
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
|
||||
@@ -1406,4 +1114,4 @@ def _ensure_index_config(
|
||||
return embeddings, index_config
|
||||
|
||||
|
||||
PLACEHOLDER = object()
|
||||
_PLACEHOLDER = object()
|
||||
|
||||
Generated
+1258
File diff suppressed because it is too large
Load Diff
@@ -1,62 +1,47 @@
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
[tool.poetry]
|
||||
name = "langgraph-checkpoint-postgres"
|
||||
version = "3.0.4"
|
||||
version = "2.0.15"
|
||||
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
license-files = ['LICENSE']
|
||||
dependencies = [
|
||||
"langgraph-checkpoint>=2.1.2,<5.0.0",
|
||||
"orjson>=3.10.1",
|
||||
"psycopg>=3.2.0",
|
||||
"psycopg-pool>=3.2.0",
|
||||
]
|
||||
readme = "README.md"
|
||||
repository = "https://www.github.com/langchain-ai/langgraph"
|
||||
packages = [{ include = "langgraph" }]
|
||||
|
||||
[project.urls]
|
||||
Source = "https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint-postgres"
|
||||
Twitter = "https://x.com/LangChain"
|
||||
Slack = "https://www.langchain.com/join-community"
|
||||
Reddit = "https://www.reddit.com/r/LangChain/"
|
||||
[tool.poetry.dependencies]
|
||||
python = "^3.9.0,<4.0"
|
||||
langgraph-checkpoint = "^2.0.15"
|
||||
orjson = ">=3.10.1"
|
||||
psycopg = "^3.2.0"
|
||||
psycopg-pool = "^3.2.0"
|
||||
|
||||
[dependency-groups]
|
||||
test = [
|
||||
"pytest",
|
||||
"anyio",
|
||||
"pytest-asyncio",
|
||||
"pytest-mock",
|
||||
"psycopg[binary]",
|
||||
"langgraph-checkpoint",
|
||||
"pytest-watcher",
|
||||
]
|
||||
lint = [
|
||||
"ruff",
|
||||
"codespell",
|
||||
"mypy",
|
||||
]
|
||||
dev = [
|
||||
{include-group = "test"},
|
||||
{include-group = "lint"},
|
||||
]
|
||||
|
||||
[tool.uv]
|
||||
default-groups = ['dev']
|
||||
|
||||
[tool.uv.sources]
|
||||
langgraph-checkpoint = { path = "../checkpoint", editable = true }
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
include = ["langgraph"]
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
ruff = "^0.6.2"
|
||||
codespell = "^2.2.0"
|
||||
pytest = "^7.2.1"
|
||||
anyio = "^4.4.0"
|
||||
pytest-asyncio = "^0.21.1"
|
||||
pytest-mock = "^3.11.1"
|
||||
mypy = "^1.10.0"
|
||||
psycopg = {extras = ["binary"], version = ">=3.0.0"}
|
||||
langgraph-checkpoint = {path = "../checkpoint", develop = true}
|
||||
pytest-watcher = "^0.4.3"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
# --strict-markers will raise errors on unknown marks.
|
||||
# https://docs.pytest.org/en/7.1.x/how-to/mark.html#raising-errors-on-unknown-marks
|
||||
#
|
||||
# https://docs.pytest.org/en/7.1.x/reference/reference.html
|
||||
# --strict-config any warnings encountered while parsing the `pytest`
|
||||
# section of the configuration file raise errors.
|
||||
addopts = "--strict-markers --strict-config --durations=5 -vv"
|
||||
asyncio_mode = "auto"
|
||||
|
||||
|
||||
[build-system]
|
||||
requires = ["poetry-core"]
|
||||
build-backend = "poetry.core.masonry.api"
|
||||
|
||||
[tool.ruff]
|
||||
lint.select = [
|
||||
"E", # pycodestyle
|
||||
@@ -64,10 +49,8 @@ lint.select = [
|
||||
"UP", # pyupgrade
|
||||
"B", # flake8-bugbear
|
||||
"I", # isort
|
||||
"UP", # pyupgrade
|
||||
]
|
||||
lint.ignore = ["E501", "B008"]
|
||||
target-version = "py310"
|
||||
lint.ignore = ["E501", "B008", "UP007", "UP006"]
|
||||
|
||||
[tool.mypy]
|
||||
# https://mypy.readthedocs.io/en/stable/config_file.html
|
||||
|
||||
@@ -1,373 +0,0 @@
|
||||
# type: ignore
|
||||
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
EXCLUDED_METADATA_KEYS,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import TASKS
|
||||
from psycopg import AsyncConnection
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
|
||||
from langgraph.checkpoint.postgres.aio import (
|
||||
AsyncPostgresSaver,
|
||||
AsyncShallowPostgresSaver,
|
||||
)
|
||||
from tests.conftest import DEFAULT_POSTGRES_URI
|
||||
|
||||
|
||||
def _exclude_keys(config: dict[str, Any]) -> dict[str, Any]:
|
||||
return {k: v for k, v in config.items() if k not in EXCLUDED_METADATA_KEYS}
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _pool_saver():
|
||||
"""Fixture for pool mode testing."""
|
||||
database = f"test_{uuid4().hex[:16]}"
|
||||
# create unique db
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
# yield checkpointer
|
||||
async with AsyncConnectionPool(
|
||||
DEFAULT_POSTGRES_URI + database,
|
||||
max_size=10,
|
||||
kwargs={"autocommit": True, "row_factory": dict_row},
|
||||
) as pool:
|
||||
checkpointer = AsyncPostgresSaver(pool)
|
||||
await checkpointer.setup()
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _pipe_saver():
|
||||
"""Fixture for pipeline mode testing."""
|
||||
database = f"test_{uuid4().hex[:16]}"
|
||||
# create unique db
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI + database,
|
||||
autocommit=True,
|
||||
prepare_threshold=0,
|
||||
row_factory=dict_row,
|
||||
) as conn:
|
||||
checkpointer = AsyncPostgresSaver(conn)
|
||||
await checkpointer.setup()
|
||||
async with conn.pipeline() as pipe:
|
||||
checkpointer = AsyncPostgresSaver(conn, pipe=pipe)
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _base_saver():
|
||||
"""Fixture for regular connection mode testing."""
|
||||
database = f"test_{uuid4().hex[:16]}"
|
||||
# create unique db
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI + database,
|
||||
autocommit=True,
|
||||
prepare_threshold=0,
|
||||
row_factory=dict_row,
|
||||
) as conn:
|
||||
checkpointer = AsyncPostgresSaver(conn)
|
||||
await checkpointer.setup()
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _shallow_saver():
|
||||
"""Fixture for shallow connection mode testing."""
|
||||
database = f"test_{uuid4().hex[:16]}"
|
||||
# create unique db
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI + database,
|
||||
autocommit=True,
|
||||
prepare_threshold=0,
|
||||
row_factory=dict_row,
|
||||
) as conn:
|
||||
checkpointer = AsyncShallowPostgresSaver(conn)
|
||||
await checkpointer.setup()
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _saver(name: str):
|
||||
if name == "base":
|
||||
async with _base_saver() as saver:
|
||||
yield saver
|
||||
elif name == "shallow":
|
||||
async with _shallow_saver() as saver:
|
||||
yield saver
|
||||
elif name == "pool":
|
||||
async with _pool_saver() as saver:
|
||||
yield saver
|
||||
elif name == "pipe":
|
||||
async with _pipe_saver() as saver:
|
||||
yield saver
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_data():
|
||||
"""Fixture providing test data for checkpoint tests."""
|
||||
config_1: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-1",
|
||||
"checkpoint_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
config_2: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_id": "2",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
config_3: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_id": "2-inner",
|
||||
"checkpoint_ns": "inner",
|
||||
}
|
||||
}
|
||||
|
||||
chkpnt_1: Checkpoint = empty_checkpoint()
|
||||
chkpnt_2: Checkpoint = create_checkpoint(chkpnt_1, {}, 1)
|
||||
chkpnt_3: Checkpoint = empty_checkpoint()
|
||||
|
||||
metadata_1: CheckpointMetadata = {
|
||||
"source": "input",
|
||||
"step": 2,
|
||||
"score": 1,
|
||||
}
|
||||
metadata_2: CheckpointMetadata = {
|
||||
"source": "loop",
|
||||
"step": 1,
|
||||
"score": None,
|
||||
}
|
||||
metadata_3: CheckpointMetadata = {}
|
||||
|
||||
return {
|
||||
"configs": [config_1, config_2, config_3],
|
||||
"checkpoints": [chkpnt_1, chkpnt_2, chkpnt_3],
|
||||
"metadata": [metadata_1, metadata_2, metadata_3],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
||||
async def test_combined_metadata(saver_name: str, test_data) -> None:
|
||||
async with _saver(saver_name) as saver:
|
||||
config = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_ns": "",
|
||||
"__super_private_key": "super_private_value",
|
||||
},
|
||||
"metadata": {"run_id": "my_run_id"},
|
||||
}
|
||||
chkpnt: Checkpoint = create_checkpoint(empty_checkpoint(), {}, 1)
|
||||
metadata: CheckpointMetadata = {
|
||||
"source": "loop",
|
||||
"step": 1,
|
||||
"score": None,
|
||||
}
|
||||
await saver.aput(config, chkpnt, metadata, {})
|
||||
checkpoint = await saver.aget_tuple(config)
|
||||
assert checkpoint.metadata == {
|
||||
**metadata,
|
||||
"run_id": "my_run_id",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
||||
async def test_asearch(saver_name: str, test_data) -> None:
|
||||
async with _saver(saver_name) as saver:
|
||||
configs = test_data["configs"]
|
||||
checkpoints = test_data["checkpoints"]
|
||||
metadata = test_data["metadata"]
|
||||
|
||||
await saver.aput(configs[0], checkpoints[0], metadata[0], {})
|
||||
await saver.aput(configs[1], checkpoints[1], metadata[1], {})
|
||||
await saver.aput(configs[2], checkpoints[2], metadata[2], {})
|
||||
|
||||
# call method / assertions
|
||||
query_1 = {"source": "input"} # search by 1 key
|
||||
query_2 = {
|
||||
"step": 1,
|
||||
} # search by multiple keys
|
||||
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
|
||||
query_4 = {"source": "update", "step": 1} # no match
|
||||
|
||||
search_results_1 = [c async for c in saver.alist(None, filter=query_1)]
|
||||
assert len(search_results_1) == 1
|
||||
assert search_results_1[0].metadata == {
|
||||
**_exclude_keys(configs[0]["configurable"]),
|
||||
**metadata[0],
|
||||
}
|
||||
|
||||
search_results_2 = [c async for c in saver.alist(None, filter=query_2)]
|
||||
assert len(search_results_2) == 1
|
||||
assert search_results_2[0].metadata == {
|
||||
**_exclude_keys(configs[1]["configurable"]),
|
||||
**metadata[1],
|
||||
}
|
||||
|
||||
search_results_3 = [c async for c in saver.alist(None, filter=query_3)]
|
||||
assert len(search_results_3) == 3
|
||||
|
||||
search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
|
||||
assert len(search_results_4) == 0
|
||||
|
||||
# search by config (defaults to checkpoints across all namespaces)
|
||||
search_results_5 = [
|
||||
c async for c in saver.alist({"configurable": {"thread_id": "thread-2"}})
|
||||
]
|
||||
assert len(search_results_5) == 2
|
||||
assert {
|
||||
search_results_5[0].config["configurable"]["checkpoint_ns"],
|
||||
search_results_5[1].config["configurable"]["checkpoint_ns"],
|
||||
} == {"", "inner"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
||||
async def test_null_chars(saver_name: str, test_data) -> None:
|
||||
async with _saver(saver_name) as saver:
|
||||
config = await saver.aput(
|
||||
test_data["configs"][0],
|
||||
test_data["checkpoints"][0],
|
||||
{"my_key": "\x00abc"},
|
||||
{},
|
||||
)
|
||||
assert (await saver.aget_tuple(config)).metadata["my_key"] == "abc" # type: ignore
|
||||
assert [c async for c in saver.alist(None, filter={"my_key": "abc"})][
|
||||
0
|
||||
].metadata["my_key"] == "abc"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
async def test_pending_sends_migration(saver_name: str) -> None:
|
||||
async with _saver(saver_name) as saver:
|
||||
config = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-1",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
|
||||
# create the first checkpoint
|
||||
# and put some pending sends
|
||||
checkpoint_0 = empty_checkpoint()
|
||||
config = await saver.aput(config, checkpoint_0, {}, {})
|
||||
await saver.aput_writes(
|
||||
config, [(TASKS, "send-1"), (TASKS, "send-2")], task_id="task-1"
|
||||
)
|
||||
await saver.aput_writes(config, [(TASKS, "send-3")], task_id="task-2")
|
||||
|
||||
# check that fetching checkpoint_0 doesn't attach pending sends
|
||||
# (they should be attached to the next checkpoint)
|
||||
tuple_0 = await saver.aget_tuple(config)
|
||||
assert tuple_0.checkpoint["channel_values"] == {}
|
||||
assert tuple_0.checkpoint["channel_versions"] == {}
|
||||
|
||||
# create the second checkpoint
|
||||
checkpoint_1 = create_checkpoint(checkpoint_0, {}, 1)
|
||||
config = await saver.aput(config, checkpoint_1, {}, {})
|
||||
|
||||
# check that pending sends are attached to checkpoint_1
|
||||
tuple_1 = await saver.aget_tuple(config)
|
||||
assert tuple_1.checkpoint["channel_values"] == {
|
||||
TASKS: ["send-1", "send-2", "send-3"]
|
||||
}
|
||||
assert TASKS in tuple_1.checkpoint["channel_versions"]
|
||||
|
||||
# check that list also applies the migration
|
||||
search_results = [
|
||||
c async for c in saver.alist({"configurable": {"thread_id": "thread-1"}})
|
||||
]
|
||||
assert len(search_results) == 2
|
||||
assert search_results[-1].checkpoint["channel_values"] == {}
|
||||
assert search_results[-1].checkpoint["channel_versions"] == {}
|
||||
assert search_results[0].checkpoint["channel_values"] == {
|
||||
TASKS: ["send-1", "send-2", "send-3"]
|
||||
}
|
||||
assert TASKS in search_results[0].checkpoint["channel_versions"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
async def test_get_checkpoint_no_channel_values(
|
||||
monkeypatch, saver_name: str, test_data
|
||||
) -> None:
|
||||
"""Backwards compatibility test that verifies a checkpoint with no channel_values key can be retrieved without throwing an error."""
|
||||
async with _saver(saver_name) as saver:
|
||||
config = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_ns": "",
|
||||
"__super_private_key": "super_private_value",
|
||||
},
|
||||
"metadata": {"run_id": "my_run_id"},
|
||||
}
|
||||
chkpnt: Checkpoint = create_checkpoint(empty_checkpoint(), {}, 1)
|
||||
await saver.aput(config, chkpnt, {}, {})
|
||||
|
||||
load_checkpoint_tuple = saver._load_checkpoint_tuple
|
||||
|
||||
def patched_load_checkpoint_tuple(value):
|
||||
value["checkpoint"].pop("channel_values", None)
|
||||
return load_checkpoint_tuple(value)
|
||||
|
||||
monkeypatch.setattr(
|
||||
saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple
|
||||
)
|
||||
|
||||
checkpoint = await saver.aget_tuple(config)
|
||||
assert checkpoint.checkpoint["channel_values"] == {}
|
||||
@@ -1,687 +0,0 @@
|
||||
# type: ignore
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import itertools
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from langgraph.store.base import (
|
||||
GetOp,
|
||||
Item,
|
||||
ListNamespacesOp,
|
||||
PutOp,
|
||||
SearchOp,
|
||||
)
|
||||
from psycopg import AsyncConnection
|
||||
|
||||
from langgraph.store.postgres import AsyncPostgresStore
|
||||
from tests.conftest import (
|
||||
DEFAULT_URI,
|
||||
VECTOR_TYPES,
|
||||
CharacterEmbeddings,
|
||||
)
|
||||
|
||||
TTL_SECONDS = 6
|
||||
TTL_MINUTES = TTL_SECONDS / 60
|
||||
|
||||
|
||||
@pytest.fixture(scope="function", params=["default", "pipe", "pool"])
|
||||
async def store(request) -> AsyncIterator[AsyncPostgresStore]:
|
||||
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
|
||||
|
||||
conn_string = f"{uri_base}/{database}{query_params}"
|
||||
admin_conn_string = DEFAULT_URI
|
||||
ttl_config = {
|
||||
"default_ttl": TTL_MINUTES,
|
||||
"refresh_on_read": True,
|
||||
"sweep_interval_minutes": TTL_MINUTES / 2,
|
||||
}
|
||||
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, ttl=ttl_config
|
||||
) as store:
|
||||
store.MIGRATIONS = [
|
||||
(
|
||||
mig.replace("ttl_minutes INT;", "ttl_minutes FLOAT;")
|
||||
if isinstance(mig, str)
|
||||
else mig
|
||||
)
|
||||
for mig in store.MIGRATIONS
|
||||
]
|
||||
await store.setup()
|
||||
async with store._cursor() as cur:
|
||||
# drop the migration index
|
||||
await cur.execute("DROP TABLE IF EXISTS store_migrations")
|
||||
await store.setup() # Will fail if migrations aren't idempotent
|
||||
|
||||
if request.param == "pipe":
|
||||
async with AsyncPostgresStore.from_conn_string(
|
||||
conn_string, pipeline=True, ttl=ttl_config
|
||||
) as store:
|
||||
await store.start_ttl_sweeper()
|
||||
yield store
|
||||
await store.stop_ttl_sweeper()
|
||||
elif request.param == "pool":
|
||||
async with AsyncPostgresStore.from_conn_string(
|
||||
conn_string, pool_config={"min_size": 1, "max_size": 10}, ttl=ttl_config
|
||||
) as store:
|
||||
await store.start_ttl_sweeper()
|
||||
yield store
|
||||
await store.stop_ttl_sweeper()
|
||||
else: # default
|
||||
async with AsyncPostgresStore.from_conn_string(
|
||||
conn_string, ttl=ttl_config
|
||||
) as store:
|
||||
await store.start_ttl_sweeper()
|
||||
yield store
|
||||
await store.stop_ttl_sweeper()
|
||||
finally:
|
||||
async with await AsyncConnection.connect(
|
||||
admin_conn_string, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
async def test_no_running_loop(store: AsyncPostgresStore) -> None:
|
||||
with pytest.raises(asyncio.InvalidStateError):
|
||||
store.put(("foo", "bar"), "baz", {"val": "baz"})
|
||||
with pytest.raises(asyncio.InvalidStateError):
|
||||
store.get(("foo", "bar"), "baz")
|
||||
with pytest.raises(asyncio.InvalidStateError):
|
||||
store.delete(("foo", "bar"), "baz")
|
||||
with pytest.raises(asyncio.InvalidStateError):
|
||||
store.search(("foo", "bar"))
|
||||
with pytest.raises(asyncio.InvalidStateError):
|
||||
store.list_namespaces(prefix=("foo",))
|
||||
with pytest.raises(asyncio.InvalidStateError):
|
||||
store.batch([PutOp(namespace=("foo", "bar"), key="baz", value={"val": "baz"})])
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(store.put, ("foo", "bar"), "baz", {"val": "baz"})
|
||||
result = await asyncio.wrap_future(future)
|
||||
assert result is None
|
||||
future = executor.submit(store.get, ("foo", "bar"), "baz")
|
||||
result = await asyncio.wrap_future(future)
|
||||
assert result.value == {"val": "baz"}
|
||||
result = await asyncio.wrap_future(
|
||||
executor.submit(store.list_namespaces, prefix=("foo",))
|
||||
)
|
||||
|
||||
|
||||
async def test_large_batches(request: Any, store: AsyncPostgresStore) -> None:
|
||||
N = 100 # less important that we are performant here
|
||||
M = 10
|
||||
|
||||
with ThreadPoolExecutor(max_workers=10) as executor:
|
||||
futures = []
|
||||
for m in range(M):
|
||||
for i in range(N):
|
||||
futures += [
|
||||
executor.submit(
|
||||
store.put,
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
value={"foo": "bar" + str(i)},
|
||||
),
|
||||
executor.submit(
|
||||
store.get,
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
),
|
||||
executor.submit(
|
||||
store.list_namespaces,
|
||||
prefix=None,
|
||||
max_depth=m + 1,
|
||||
),
|
||||
executor.submit(
|
||||
store.search,
|
||||
("test",),
|
||||
),
|
||||
executor.submit(
|
||||
store.put,
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
value={"foo": "bar" + str(i)},
|
||||
),
|
||||
executor.submit(
|
||||
store.put,
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
None,
|
||||
),
|
||||
]
|
||||
|
||||
results = await asyncio.gather(
|
||||
*(asyncio.wrap_future(future) for future in futures)
|
||||
)
|
||||
assert len(results) == M * N * 6
|
||||
|
||||
|
||||
async def test_large_batches_async(store: AsyncPostgresStore) -> None:
|
||||
N = 1000
|
||||
M = 10
|
||||
coros = []
|
||||
for m in range(M):
|
||||
for i in range(N):
|
||||
coros.append(
|
||||
store.aput(
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
value={"foo": "bar" + str(i)},
|
||||
)
|
||||
)
|
||||
coros.append(
|
||||
store.aget(
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
)
|
||||
)
|
||||
coros.append(
|
||||
store.alist_namespaces(
|
||||
prefix=None,
|
||||
max_depth=m + 1,
|
||||
)
|
||||
)
|
||||
coros.append(
|
||||
store.asearch(
|
||||
("test",),
|
||||
)
|
||||
)
|
||||
coros.append(
|
||||
store.aput(
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
value={"foo": "bar" + str(i)},
|
||||
)
|
||||
)
|
||||
coros.append(
|
||||
store.adelete(
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
)
|
||||
)
|
||||
|
||||
results = await asyncio.gather(*coros)
|
||||
assert len(results) == M * N * 6
|
||||
|
||||
|
||||
async def test_abatch_order(store: AsyncPostgresStore) -> None:
|
||||
# Setup test data
|
||||
await store.aput(("test", "foo"), "key1", {"data": "value1"})
|
||||
await store.aput(("test", "bar"), "key2", {"data": "value2"})
|
||||
|
||||
ops = [
|
||||
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 len(results) == 5
|
||||
assert isinstance(results[0], Item)
|
||||
assert isinstance(results[0].value, dict)
|
||||
assert results[0].value == {"data": "value1"}
|
||||
assert results[0].key == "key1"
|
||||
assert results[1] is None
|
||||
assert isinstance(results[2], list)
|
||||
assert len(results[2]) == 1
|
||||
assert isinstance(results[3], list)
|
||||
assert ("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", "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", "foo"), key="key1"),
|
||||
]
|
||||
|
||||
results_reordered = await store.abatch(ops_reordered)
|
||||
assert len(results_reordered) == 5
|
||||
assert isinstance(results_reordered[0], list)
|
||||
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 ("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"}
|
||||
assert results_reordered[4].key == "key1"
|
||||
|
||||
|
||||
async def test_batch_get_ops(store: AsyncPostgresStore) -> None:
|
||||
# Setup test data
|
||||
await store.aput(("test",), "key1", {"data": "value1"})
|
||||
await store.aput(("test",), "key2", {"data": "value2"})
|
||||
|
||||
ops = [
|
||||
GetOp(namespace=("test",), key="key1"),
|
||||
GetOp(namespace=("test",), key="key2"),
|
||||
GetOp(namespace=("test",), key="key3"),
|
||||
]
|
||||
|
||||
results = await store.abatch(ops)
|
||||
|
||||
assert len(results) == 3
|
||||
assert results[0] is not None
|
||||
assert results[1] is not None
|
||||
assert results[2] is None
|
||||
assert results[0].key == "key1"
|
||||
assert results[1].key == "key2"
|
||||
|
||||
|
||||
async def test_batch_put_ops(store: AsyncPostgresStore) -> None:
|
||||
ops = [
|
||||
PutOp(namespace=("test",), key="key1", value={"data": "value1"}),
|
||||
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
|
||||
PutOp(namespace=("test",), key="key3", value=None),
|
||||
]
|
||||
|
||||
results = await store.abatch(ops)
|
||||
|
||||
assert len(results) == 3
|
||||
assert all(result is None for result in results)
|
||||
|
||||
# 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:
|
||||
# Setup test data
|
||||
await store.aput(("test", "foo"), "key1", {"data": "value1"})
|
||||
await store.aput(("test", "bar"), "key2", {"data": "value2"})
|
||||
|
||||
ops = [
|
||||
SearchOp(
|
||||
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
|
||||
),
|
||||
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
|
||||
]
|
||||
|
||||
results = await store.abatch(ops)
|
||||
|
||||
assert len(results) == 2
|
||||
assert len(results[0]) == 1 # Filtered results
|
||||
assert len(results[1]) == 2 # All results
|
||||
|
||||
|
||||
async def test_batch_list_namespaces_ops(store: AsyncPostgresStore) -> None:
|
||||
# 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 len(results[0]) == 2
|
||||
assert ("test", "namespace1") in results[0]
|
||||
assert ("test", "namespace2") in results[0]
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _create_vector_store(
|
||||
vector_type: str,
|
||||
distance_type: str,
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
text_fields: list[str] | None = None,
|
||||
) -> AsyncIterator[AsyncPostgresStore]:
|
||||
"""Create a store with vector search enabled."""
|
||||
|
||||
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
|
||||
|
||||
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,
|
||||
"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}")
|
||||
|
||||
|
||||
@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
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
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"}),
|
||||
]
|
||||
|
||||
for key, value in docs:
|
||||
await vector_store.aput(("test",), key, value)
|
||||
|
||||
results = await vector_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
|
||||
|
||||
|
||||
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"})
|
||||
|
||||
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
|
||||
|
||||
await vector_store.aput(("test",), "doc1", {"text": "new text about dogs"})
|
||||
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
# 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)
|
||||
|
||||
|
||||
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}),
|
||||
]
|
||||
|
||||
for key, value in docs:
|
||||
await vector_store.aput(("test",), key, value)
|
||||
|
||||
results = await vector_store.asearch(
|
||||
("test",), query="apple", filter={"color": "red"}
|
||||
)
|
||||
assert len(results) == 2
|
||||
assert results[0].key == "doc1"
|
||||
|
||||
results = await vector_store.asearch(
|
||||
("test",), query="car", filter={"color": "red"}
|
||||
)
|
||||
assert len(results) == 2
|
||||
assert results[0].key == "doc2"
|
||||
|
||||
results = await vector_store.asearch(
|
||||
("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
|
||||
)
|
||||
assert len(results) == 3
|
||||
assert results[0].key == "doc4"
|
||||
|
||||
results = await vector_store.asearch(
|
||||
("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"}
|
||||
)
|
||||
assert len(results) == 1
|
||||
assert results[0].key == "doc3"
|
||||
|
||||
|
||||
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}"}
|
||||
)
|
||||
|
||||
results_page1 = await vector_store.asearch(("test",), query="test", limit=2)
|
||||
results_page2 = await vector_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 vector_store.asearch(("test",), query="test", limit=10)
|
||||
assert len(all_results) == 5
|
||||
|
||||
|
||||
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"})
|
||||
|
||||
perfect_match = await vector_store.asearch(("test",), query="text test document")
|
||||
perfect_score = perfect_match[0].score
|
||||
|
||||
results = await vector_store.asearch(("test",), query="")
|
||||
assert len(results) == 1
|
||||
assert results[0].score is None
|
||||
|
||||
results = await vector_store.asearch(("test",), query=None)
|
||||
assert len(results) == 1
|
||||
assert results[0].score is None
|
||||
|
||||
long_query = "foo " * 100
|
||||
results = await vector_store.asearch(("test",), query=long_query)
|
||||
assert len(results) == 1
|
||||
assert results[0].score < perfect_score
|
||||
|
||||
special_query = "test!@#$%^&*()"
|
||||
results = await vector_store.asearch(("test",), query=special_query)
|
||||
assert len(results) == 1
|
||||
assert results[0].score < perfect_score
|
||||
|
||||
|
||||
@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)
|
||||
|
||||
# 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)
|
||||
|
||||
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)
|
||||
|
||||
# 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
|
||||
|
||||
|
||||
async def test_store_ttl(store):
|
||||
# Assumes a TTL of 1 minute = 60 seconds
|
||||
ns = ("foo",)
|
||||
await store.start_ttl_sweeper()
|
||||
await store.aput(
|
||||
ns,
|
||||
key="item1",
|
||||
value={"foo": "bar"},
|
||||
ttl=TTL_MINUTES, # type: ignore
|
||||
)
|
||||
await asyncio.sleep(TTL_SECONDS - 2)
|
||||
res = await store.aget(ns, key="item1", refresh_ttl=True)
|
||||
assert res is not None
|
||||
await asyncio.sleep(TTL_SECONDS - 2)
|
||||
results = await store.asearch(ns, query="foo", refresh_ttl=True)
|
||||
assert len(results) == 1
|
||||
await asyncio.sleep(TTL_SECONDS - 2)
|
||||
res = await store.aget(ns, key="item1", refresh_ttl=False)
|
||||
assert res is not None
|
||||
await asyncio.sleep(TTL_SECONDS - 1)
|
||||
# Now has been (TTL_SECONDS-2)*2 > TTL_SECONDS + TTL_SECONDS/2
|
||||
results = await store.asearch(ns, query="bar", refresh_ttl=False)
|
||||
assert len(results) == 0
|
||||
@@ -1,14 +1,14 @@
|
||||
# type: ignore
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
from typing import Any, Optional
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from psycopg import Connection
|
||||
|
||||
from langgraph.store.base import (
|
||||
GetOp,
|
||||
Item,
|
||||
@@ -17,8 +17,6 @@ from langgraph.store.base import (
|
||||
PutOp,
|
||||
SearchOp,
|
||||
)
|
||||
from psycopg import Connection
|
||||
|
||||
from langgraph.store.postgres import PostgresStore
|
||||
from tests.conftest import (
|
||||
DEFAULT_URI,
|
||||
@@ -26,9 +24,6 @@ from tests.conftest import (
|
||||
CharacterEmbeddings,
|
||||
)
|
||||
|
||||
TTL_SECONDS = 6
|
||||
TTL_MINUTES = TTL_SECONDS / 60
|
||||
|
||||
|
||||
@pytest.fixture(scope="function", params=["default", "pipe", "pool"])
|
||||
def store(request) -> PostgresStore:
|
||||
@@ -37,56 +32,29 @@ def store(request) -> PostgresStore:
|
||||
uri_base = "/".join(uri_parts[:-1])
|
||||
query_params = ""
|
||||
if "?" in uri_parts[-1]:
|
||||
_, query_params = uri_parts[-1].split("?", 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
|
||||
ttl_config = {
|
||||
"default_ttl": TTL_MINUTES,
|
||||
"refresh_on_read": True,
|
||||
"sweep_interval_minutes": TTL_MINUTES / 2,
|
||||
}
|
||||
|
||||
with Connection.connect(admin_conn_string, autocommit=True) as conn:
|
||||
conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
with PostgresStore.from_conn_string(conn_string, ttl=ttl_config) as store:
|
||||
store.MIGRATIONS = [
|
||||
(
|
||||
mig.replace("ttl_minutes INT;", "ttl_minutes FLOAT;")
|
||||
if isinstance(mig, str)
|
||||
else mig
|
||||
)
|
||||
for mig in store.MIGRATIONS
|
||||
]
|
||||
with PostgresStore.from_conn_string(conn_string) as store:
|
||||
store.setup()
|
||||
|
||||
if request.param == "pipe":
|
||||
with PostgresStore.from_conn_string(
|
||||
conn_string,
|
||||
pipeline=True,
|
||||
ttl=ttl_config,
|
||||
) as store:
|
||||
store.start_ttl_sweeper()
|
||||
with PostgresStore.from_conn_string(conn_string, pipeline=True) as store:
|
||||
yield store
|
||||
|
||||
store.stop_ttl_sweeper()
|
||||
elif request.param == "pool":
|
||||
with PostgresStore.from_conn_string(
|
||||
conn_string,
|
||||
pool_config={"min_size": 1, "max_size": 10},
|
||||
ttl=ttl_config,
|
||||
conn_string, pool_config={"min_size": 1, "max_size": 10}
|
||||
) as store:
|
||||
store.start_ttl_sweeper()
|
||||
yield store
|
||||
|
||||
store.stop_ttl_sweeper()
|
||||
else: # default
|
||||
with PostgresStore.from_conn_string(conn_string, ttl=ttl_config) as store:
|
||||
store.start_ttl_sweeper()
|
||||
with PostgresStore.from_conn_string(conn_string) as store:
|
||||
yield store
|
||||
|
||||
store.stop_ttl_sweeper()
|
||||
finally:
|
||||
with Connection.connect(admin_conn_string, autocommit=True) as conn:
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
@@ -252,127 +220,134 @@ def test_batch_list_namespaces_ops(store: PostgresStore) -> None:
|
||||
assert all(ns[-1] == "public" for ns in results[2])
|
||||
|
||||
|
||||
def test_basic_store_ops(store) -> None:
|
||||
namespace = ("test", "documents")
|
||||
item_id = "doc1"
|
||||
item_value = {"title": "Test Document", "content": "Hello, World!"}
|
||||
class TestPostgresStore:
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup(self) -> None:
|
||||
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
|
||||
store.setup()
|
||||
|
||||
store.put(namespace, item_id, item_value)
|
||||
item = store.get(namespace, item_id)
|
||||
def test_basic_store_ops(self) -> None:
|
||||
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
|
||||
namespace = ("test", "documents")
|
||||
item_id = "doc1"
|
||||
item_value = {"title": "Test Document", "content": "Hello, World!"}
|
||||
|
||||
assert item
|
||||
assert item.namespace == namespace
|
||||
assert item.key == item_id
|
||||
assert item.value == item_value
|
||||
store.put(namespace, item_id, item_value)
|
||||
item = store.get(namespace, item_id)
|
||||
|
||||
# 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 item
|
||||
assert item.namespace == namespace
|
||||
assert item.key == item_id
|
||||
assert item.value == item_value
|
||||
|
||||
assert updated_item.value == updated_value
|
||||
assert updated_item.updated_at > item.updated_at
|
||||
# Test update
|
||||
updated_value = {"title": "Updated Document", "content": "Hello, Updated!"}
|
||||
store.put(namespace, item_id, updated_value)
|
||||
updated_item = store.get(namespace, item_id)
|
||||
|
||||
# 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
|
||||
assert updated_item.value == updated_value
|
||||
assert updated_item.updated_at > item.updated_at
|
||||
|
||||
# Test delete
|
||||
store.delete(namespace, item_id)
|
||||
deleted_item = store.get(namespace, item_id)
|
||||
assert deleted_item is None
|
||||
# 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
|
||||
|
||||
# Test delete
|
||||
store.delete(namespace, item_id)
|
||||
deleted_item = store.get(namespace, item_id)
|
||||
assert deleted_item is None
|
||||
|
||||
def test_list_namespaces(store) -> None:
|
||||
# Create test data with various namespaces
|
||||
test_namespaces = [
|
||||
("test", "documents", "public"),
|
||||
("test", "documents", "private"),
|
||||
("test", "images", "public"),
|
||||
("test", "images", "private"),
|
||||
("prod", "documents", "public"),
|
||||
("prod", "documents", "private"),
|
||||
]
|
||||
def test_list_namespaces(self) -> None:
|
||||
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
|
||||
# Create test data with various namespaces
|
||||
test_namespaces = [
|
||||
("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"})
|
||||
# Insert test data
|
||||
for namespace in test_namespaces:
|
||||
store.put(namespace, "dummy", {"content": "dummy"})
|
||||
|
||||
# Test listing with various filters
|
||||
all_namespaces = store.list_namespaces()
|
||||
assert len(all_namespaces) == len(test_namespaces)
|
||||
# Test listing with various filters
|
||||
all_namespaces = store.list_namespaces()
|
||||
assert len(all_namespaces) == len(test_namespaces)
|
||||
|
||||
# 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)
|
||||
# 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)
|
||||
|
||||
# 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)
|
||||
# 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)
|
||||
|
||||
# Test max depth
|
||||
depth_2_namespaces = store.list_namespaces(max_depth=2)
|
||||
assert all(len(ns) <= 2 for ns in depth_2_namespaces)
|
||||
# Test max depth
|
||||
depth_2_namespaces = store.list_namespaces(max_depth=2)
|
||||
assert all(len(ns) <= 2 for ns in depth_2_namespaces)
|
||||
|
||||
# Test pagination
|
||||
paginated_namespaces = store.list_namespaces(limit=3)
|
||||
assert len(paginated_namespaces) == 3
|
||||
# Test pagination
|
||||
paginated_namespaces = store.list_namespaces(limit=3)
|
||||
assert len(paginated_namespaces) == 3
|
||||
|
||||
# Cleanup
|
||||
for namespace in test_namespaces:
|
||||
store.delete(namespace, "dummy")
|
||||
# Cleanup
|
||||
for namespace in test_namespaces:
|
||||
store.delete(namespace, "dummy")
|
||||
|
||||
def test_search(self) -> None:
|
||||
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
|
||||
# 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"]},
|
||||
),
|
||||
]
|
||||
|
||||
def test_search(store) -> None:
|
||||
# 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, key, value in test_data:
|
||||
store.put(namespace, key, value)
|
||||
|
||||
for namespace, key, value in test_data:
|
||||
store.put(namespace, key, value)
|
||||
# Test basic search
|
||||
all_items = store.search(["test"])
|
||||
assert len(all_items) == 3
|
||||
|
||||
# Test basic search
|
||||
all_items = store.search(["test"])
|
||||
assert len(all_items) == 3
|
||||
# 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)
|
||||
|
||||
# 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)
|
||||
# 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)
|
||||
|
||||
# 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)
|
||||
# Test pagination
|
||||
paginated_items = store.search(["test"], limit=2)
|
||||
assert len(paginated_items) == 2
|
||||
|
||||
# Test pagination
|
||||
paginated_items = store.search(["test"], limit=2)
|
||||
assert len(paginated_items) == 2
|
||||
offset_items = store.search(["test"], offset=2)
|
||||
assert len(offset_items) == 1
|
||||
|
||||
offset_items = store.search(["test"], offset=2)
|
||||
assert len(offset_items) == 1
|
||||
|
||||
# Cleanup
|
||||
for namespace, key, _ in test_data:
|
||||
store.delete(namespace, key)
|
||||
# Cleanup
|
||||
for namespace, key, _ in test_data:
|
||||
store.delete(namespace, key)
|
||||
|
||||
|
||||
@contextmanager
|
||||
@@ -380,8 +355,7 @@ def _create_vector_store(
|
||||
vector_type: str,
|
||||
distance_type: str,
|
||||
fake_embeddings: Embeddings,
|
||||
text_fields: list[str] | None = None,
|
||||
enable_ttl: bool = True,
|
||||
text_fields: Optional[list[str]] = None,
|
||||
) -> PostgresStore:
|
||||
"""Create a store with vector search enabled."""
|
||||
database = f"test_{uuid4().hex[:16]}"
|
||||
@@ -402,7 +376,7 @@ def _create_vector_store(
|
||||
"vector_type": vector_type,
|
||||
},
|
||||
"distance_type": distance_type,
|
||||
"fields": text_fields,
|
||||
"text_fields": text_fields,
|
||||
}
|
||||
|
||||
with Connection.connect(admin_conn_string, autocommit=True) as conn:
|
||||
@@ -411,32 +385,23 @@ def _create_vector_store(
|
||||
with PostgresStore.from_conn_string(
|
||||
conn_string,
|
||||
index=index_config,
|
||||
ttl={"default_ttl": 2, "refresh_on_read": True} if enable_ttl else None,
|
||||
) as store:
|
||||
store.setup()
|
||||
with store._cursor() as cur:
|
||||
# drop the migration index
|
||||
cur.execute("DROP TABLE IF EXISTS store_migrations")
|
||||
store.setup() # Will fail if migrations aren't idempotent
|
||||
yield store
|
||||
finally:
|
||||
with Connection.connect(admin_conn_string, autocommit=True) as conn:
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
_vector_params = [
|
||||
(vector_type, distance_type, True)
|
||||
for vector_type in VECTOR_TYPES
|
||||
for distance_type in (
|
||||
["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
|
||||
)
|
||||
]
|
||||
_vector_params += [(*_vector_params[-1][:2], False)]
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
scope="function",
|
||||
params=_vector_params,
|
||||
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(
|
||||
@@ -444,10 +409,8 @@ def vector_store(
|
||||
fake_embeddings: Embeddings,
|
||||
) -> PostgresStore:
|
||||
"""Create a store with vector search enabled."""
|
||||
vector_type, distance_type, enable_ttl = request.param
|
||||
with _create_vector_store(
|
||||
vector_type, distance_type, fake_embeddings, enable_ttl=enable_ttl
|
||||
) as store:
|
||||
vector_type, distance_type = request.param
|
||||
with _create_vector_store(vector_type, distance_type, fake_embeddings) as store:
|
||||
yield store
|
||||
|
||||
|
||||
@@ -511,10 +474,7 @@ def test_vector_update_with_embedding(vector_store: PostgresStore) -> None:
|
||||
assert not any(r.key == "doc4" for r in results_new)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("refresh_ttl", [True, False])
|
||||
def test_vector_search_with_filters(
|
||||
vector_store: PostgresStore, refresh_ttl: bool
|
||||
) -> None:
|
||||
def test_vector_search_with_filters(vector_store: PostgresStore) -> None:
|
||||
"""Test combining vector search with filters."""
|
||||
# Insert test documents
|
||||
docs = [
|
||||
@@ -527,23 +487,16 @@ def test_vector_search_with_filters(
|
||||
for key, value in docs:
|
||||
vector_store.put(("test",), key, value)
|
||||
|
||||
results = vector_store.search(
|
||||
("test",), query="apple", filter={"color": "red"}, refresh_ttl=refresh_ttl
|
||||
)
|
||||
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"}, refresh_ttl=refresh_ttl
|
||||
)
|
||||
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}},
|
||||
refresh_ttl=refresh_ttl,
|
||||
("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
|
||||
)
|
||||
assert len(results) == 3
|
||||
assert results[0].key == "doc4"
|
||||
@@ -735,7 +688,7 @@ def test_embed_with_path_operation_config(
|
||||
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), f"{results}"
|
||||
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")
|
||||
@@ -754,7 +707,7 @@ def _cosine_similarity(X: list[float], Y: list[list[float]]) -> list[float]:
|
||||
|
||||
similarities = []
|
||||
for y in Y:
|
||||
dot_product = sum(a * b for a, b in zip(X, y, strict=False))
|
||||
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
|
||||
@@ -771,7 +724,7 @@ def _inner_product(X: list[float], Y: list[list[float]]) -> list[float]:
|
||||
|
||||
similarities = []
|
||||
for y in Y:
|
||||
similarity = sum(a * b for a, b in zip(X, y, strict=False))
|
||||
similarity = sum(a * b for a, b in zip(X, y))
|
||||
similarities.append(similarity)
|
||||
|
||||
return similarities
|
||||
@@ -785,7 +738,7 @@ def _neg_l2_distance(X: list[float], Y: list[list[float]]) -> list[float]:
|
||||
|
||||
similarities = []
|
||||
for y in Y:
|
||||
similarity = sum((a - b) ** 2 for a, b in zip(X, y, strict=False)) ** 0.5
|
||||
similarity = sum((a - b) ** 2 for a, b in zip(X, y)) ** 0.5
|
||||
similarities.append(-similarity)
|
||||
|
||||
return similarities
|
||||
@@ -837,65 +790,3 @@ def test_nonnull_migrations() -> None:
|
||||
for migration in PostgresStore.MIGRATIONS:
|
||||
statement = _leading_comment_remover.sub("", migration).split()[0]
|
||||
assert statement.strip()
|
||||
|
||||
|
||||
def test_store_ttl(store):
|
||||
# Assumes a TTL of 1 minute = 60 seconds
|
||||
ns = ("foo",)
|
||||
store.put(
|
||||
ns,
|
||||
key="item1",
|
||||
value={"foo": "bar"},
|
||||
ttl=TTL_MINUTES, # type: ignore
|
||||
)
|
||||
time.sleep(TTL_SECONDS - 2)
|
||||
res = store.get(ns, key="item1", refresh_ttl=True)
|
||||
assert res is not None
|
||||
time.sleep(TTL_SECONDS - 2)
|
||||
results = store.search(ns, query="foo", refresh_ttl=True)
|
||||
assert len(results) == 1
|
||||
time.sleep(TTL_SECONDS - 2)
|
||||
res = store.get(ns, key="item1", refresh_ttl=False)
|
||||
assert res is not None
|
||||
time.sleep(TTL_SECONDS - 1)
|
||||
# Now has been (TTL_SECONDS-2)*2 > TTL_SECONDS + TTL_SECONDS/2
|
||||
res = store.search(ns, query="bar", refresh_ttl=False)
|
||||
assert len(res) == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"vector_type,distance_type",
|
||||
[
|
||||
("vector", "cosine"),
|
||||
("vector", "inner_product"),
|
||||
("halfvec", "cosine"),
|
||||
("halfvec", "inner_product"),
|
||||
],
|
||||
)
|
||||
def test_non_ascii(
|
||||
request: Any,
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
vector_type: str,
|
||||
distance_type: str,
|
||||
) -> None:
|
||||
"""Test support for non-ascii characters"""
|
||||
with _create_vector_store(vector_type, distance_type, fake_embeddings) as store:
|
||||
store.put(("user_123", "memories"), "1", {"text": "这是中文"}) # Chinese
|
||||
store.put(
|
||||
("user_123", "memories"), "2", {"text": "これは日本語です"}
|
||||
) # Japanese
|
||||
store.put(("user_123", "memories"), "3", {"text": "이건 한국어야"}) # Korean
|
||||
store.put(("user_123", "memories"), "4", {"text": "Это русский"}) # Russian
|
||||
store.put(("user_123", "memories"), "5", {"text": "यह रूसी है"}) # Hindi
|
||||
|
||||
result1 = store.search(("user_123", "memories"), query="这是中文")
|
||||
result2 = store.search(("user_123", "memories"), query="これは日本語です")
|
||||
result3 = store.search(("user_123", "memories"), query="이건 한국어야")
|
||||
result4 = store.search(("user_123", "memories"), query="Это русский")
|
||||
result5 = store.search(("user_123", "memories"), query="यह रूसी है")
|
||||
|
||||
assert result1[0].key == "1"
|
||||
assert result2[0].key == "2"
|
||||
assert result3[0].key == "3"
|
||||
assert result4[0].key == "4"
|
||||
assert result5[0].key == "5"
|
||||
|
||||
@@ -7,6 +7,10 @@ from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from psycopg import Connection
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
EXCLUDED_METADATA_KEYS,
|
||||
Checkpoint,
|
||||
@@ -14,11 +18,6 @@ from langgraph.checkpoint.base import (
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import TASKS
|
||||
from psycopg import Connection
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver
|
||||
from tests.conftest import DEFAULT_POSTGRES_URI
|
||||
|
||||
@@ -143,7 +142,8 @@ def test_data():
|
||||
config_1: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-1",
|
||||
"checkpoint_id": "1",
|
||||
# for backwards compatibility testing
|
||||
"thread_ts": "1",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
@@ -169,11 +169,13 @@ def test_data():
|
||||
metadata_1: CheckpointMetadata = {
|
||||
"source": "input",
|
||||
"step": 2,
|
||||
"writes": {},
|
||||
"score": 1,
|
||||
}
|
||||
metadata_2: CheckpointMetadata = {
|
||||
"source": "loop",
|
||||
"step": 1,
|
||||
"writes": {"foo": "bar"},
|
||||
"score": None,
|
||||
}
|
||||
metadata_3: CheckpointMetadata = {}
|
||||
@@ -200,12 +202,14 @@ def test_combined_metadata(saver_name: str, test_data) -> None:
|
||||
metadata: CheckpointMetadata = {
|
||||
"source": "loop",
|
||||
"step": 1,
|
||||
"writes": {"foo": "bar"},
|
||||
"score": None,
|
||||
}
|
||||
saver.put(config, chkpnt, metadata, {})
|
||||
checkpoint = saver.get_tuple(config)
|
||||
assert checkpoint.metadata == {
|
||||
**metadata,
|
||||
"thread_id": "thread-2",
|
||||
"run_id": "my_run_id",
|
||||
}
|
||||
|
||||
@@ -225,6 +229,7 @@ def test_search(saver_name: str, test_data) -> None:
|
||||
query_1 = {"source": "input"} # search by 1 key
|
||||
query_2 = {
|
||||
"step": 1,
|
||||
"writes": {"foo": "bar"},
|
||||
} # search by multiple keys
|
||||
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
|
||||
query_4 = {"source": "update", "step": 1} # no match
|
||||
@@ -279,82 +284,3 @@ def test_nonnull_migrations() -> None:
|
||||
for migration in PostgresSaver.MIGRATIONS:
|
||||
statement = _leading_comment_remover.sub("", migration).split()[0]
|
||||
assert statement.strip()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
def test_pending_sends_migration(saver_name: str) -> None:
|
||||
with _saver(saver_name) as saver:
|
||||
config = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-1",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
|
||||
# create the first checkpoint
|
||||
# and put some pending sends
|
||||
checkpoint_0 = empty_checkpoint()
|
||||
config = saver.put(config, checkpoint_0, {}, {})
|
||||
saver.put_writes(
|
||||
config, [(TASKS, "send-1"), (TASKS, "send-2")], task_id="task-1"
|
||||
)
|
||||
saver.put_writes(config, [(TASKS, "send-3")], task_id="task-2")
|
||||
|
||||
# check that fetching checkpoint_0 doesn't attach pending sends
|
||||
# (they should be attached to the next checkpoint)
|
||||
tuple_0 = saver.get_tuple(config)
|
||||
assert tuple_0.checkpoint["channel_values"] == {}
|
||||
assert tuple_0.checkpoint["channel_versions"] == {}
|
||||
|
||||
# create the second checkpoint
|
||||
checkpoint_1 = create_checkpoint(checkpoint_0, {}, 1)
|
||||
config = saver.put(config, checkpoint_1, {}, {})
|
||||
|
||||
# check that pending sends are attached to checkpoint_1
|
||||
checkpoint_1 = saver.get_tuple(config)
|
||||
assert checkpoint_1.checkpoint["channel_values"] == {
|
||||
TASKS: ["send-1", "send-2", "send-3"]
|
||||
}
|
||||
assert TASKS in checkpoint_1.checkpoint["channel_versions"]
|
||||
|
||||
# check that list also applies the migration
|
||||
search_results = [
|
||||
c for c in saver.list({"configurable": {"thread_id": "thread-1"}})
|
||||
]
|
||||
assert len(search_results) == 2
|
||||
assert search_results[-1].checkpoint["channel_values"] == {}
|
||||
assert search_results[-1].checkpoint["channel_versions"] == {}
|
||||
assert search_results[0].checkpoint["channel_values"] == {
|
||||
TASKS: ["send-1", "send-2", "send-3"]
|
||||
}
|
||||
assert TASKS in search_results[0].checkpoint["channel_versions"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
def test_get_checkpoint_no_channel_values(
|
||||
monkeypatch, saver_name: str, test_data
|
||||
) -> None:
|
||||
"""Backwards compatibility test that verifies a checkpoint with no channel_values key can be retrieved without throwing an error."""
|
||||
with _saver(saver_name) as saver:
|
||||
config = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_ns": "",
|
||||
"__super_private_key": "super_private_value",
|
||||
},
|
||||
}
|
||||
chkpnt: Checkpoint = create_checkpoint(empty_checkpoint(), {}, 1)
|
||||
saver.put(config, chkpnt, {}, {})
|
||||
|
||||
load_checkpoint_tuple = saver._load_checkpoint_tuple
|
||||
|
||||
def patched_load_checkpoint_tuple(value):
|
||||
value["checkpoint"].pop("channel_values", None)
|
||||
return load_checkpoint_tuple(value)
|
||||
|
||||
monkeypatch.setattr(
|
||||
saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple
|
||||
)
|
||||
|
||||
checkpoint = saver.get_tuple(config)
|
||||
assert checkpoint.checkpoint["channel_values"] == {}
|
||||
|
||||
Generated
-1351
File diff suppressed because it is too large
Load Diff
@@ -1,21 +0,0 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2024 LangChain, Inc.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -1,37 +0,0 @@
|
||||
.PHONY: test test_watch lint format
|
||||
|
||||
######################
|
||||
# TESTING AND COVERAGE
|
||||
######################
|
||||
|
||||
TEST ?= .
|
||||
|
||||
test:
|
||||
uv run pytest $(TEST)
|
||||
|
||||
test_watch:
|
||||
uv run ptw $(TEST)
|
||||
|
||||
######################
|
||||
# LINTING AND FORMATTING
|
||||
######################
|
||||
|
||||
# Define a variable for Python and notebook files.
|
||||
PYTHON_FILES=.
|
||||
MYPY_CACHE=.mypy_cache
|
||||
lint format: PYTHON_FILES=.
|
||||
lint_diff format_diff: PYTHON_FILES=$(shell git diff --name-only --relative --diff-filter=d main . | grep -E '\.py$$|\.ipynb$$')
|
||||
lint_package: PYTHON_FILES=langgraph
|
||||
lint_tests: PYTHON_FILES=tests
|
||||
lint_tests: MYPY_CACHE=.mypy_cache_test
|
||||
|
||||
lint lint_diff lint_package lint_tests:
|
||||
uv run ruff check .
|
||||
[ "$(PYTHON_FILES)" = "" ] || uv run ruff format $(PYTHON_FILES) --diff
|
||||
[ "$(PYTHON_FILES)" = "" ] || uv run ruff check --select I $(PYTHON_FILES)
|
||||
[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE)
|
||||
[ "$(PYTHON_FILES)" = "" ] || uv run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
|
||||
|
||||
format format_diff:
|
||||
uv run ruff format $(PYTHON_FILES)
|
||||
uv run ruff check --select I --fix $(PYTHON_FILES)
|
||||
@@ -1,88 +0,0 @@
|
||||
# LangGraph SQLite Checkpoint
|
||||
|
||||
Implementation of LangGraph CheckpointSaver that uses SQLite DB (both sync and async, via `aiosqlite`)
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
|
||||
write_config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
|
||||
read_config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
|
||||
checkpoint = {
|
||||
"v": 4,
|
||||
"ts": "2024-07-31T20:14:19.804150+00:00",
|
||||
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
|
||||
"channel_values": {
|
||||
"my_key": "meow",
|
||||
"node": "node"
|
||||
},
|
||||
"channel_versions": {
|
||||
"__start__": 2,
|
||||
"my_key": 3,
|
||||
"start:node": 3,
|
||||
"node": 3
|
||||
},
|
||||
"versions_seen": {
|
||||
"__input__": {},
|
||||
"__start__": {
|
||||
"__start__": 1
|
||||
},
|
||||
"node": {
|
||||
"start:node": 2
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
# store checkpoint
|
||||
checkpointer.put(write_config, checkpoint, {}, {})
|
||||
|
||||
# load checkpoint
|
||||
checkpointer.get(read_config)
|
||||
|
||||
# list checkpoints
|
||||
list(checkpointer.list(read_config))
|
||||
```
|
||||
|
||||
### Async
|
||||
|
||||
```python
|
||||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||||
|
||||
async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
|
||||
checkpoint = {
|
||||
"v": 4,
|
||||
"ts": "2024-07-31T20:14:19.804150+00:00",
|
||||
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
|
||||
"channel_values": {
|
||||
"my_key": "meow",
|
||||
"node": "node"
|
||||
},
|
||||
"channel_versions": {
|
||||
"__start__": 2,
|
||||
"my_key": 3,
|
||||
"start:node": 3,
|
||||
"node": 3
|
||||
},
|
||||
"versions_seen": {
|
||||
"__input__": {},
|
||||
"__start__": {
|
||||
"__start__": 1
|
||||
},
|
||||
"node": {
|
||||
"start:node": 2
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
# store checkpoint
|
||||
await checkpointer.aput(write_config, checkpoint, {}, {})
|
||||
|
||||
# load checkpoint
|
||||
await checkpointer.aget(read_config)
|
||||
|
||||
# list checkpoints
|
||||
[c async for c in checkpointer.alist(read_config)]
|
||||
```
|
||||
@@ -1,120 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
import sqlite3
|
||||
import threading
|
||||
from collections.abc import Mapping, Sequence
|
||||
|
||||
from langgraph.cache.base import BaseCache, FullKey, Namespace, ValueT
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
|
||||
|
||||
class SqliteCache(BaseCache[ValueT]):
|
||||
"""File-based cache using SQLite."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
path: str,
|
||||
serde: SerializerProtocol | None = None,
|
||||
) -> None:
|
||||
"""Initialize the cache with a file path."""
|
||||
super().__init__(serde=serde)
|
||||
# SQLite backing store
|
||||
self._conn = sqlite3.connect(
|
||||
path,
|
||||
check_same_thread=False,
|
||||
)
|
||||
# Serialize access to the shared connection across threads
|
||||
self._lock = threading.RLock()
|
||||
# Better concurrency & atomicity
|
||||
self._conn.execute("PRAGMA journal_mode=WAL;")
|
||||
# Schema: key -> (expiry, encoding, value)
|
||||
self._conn.execute(
|
||||
"""CREATE TABLE IF NOT EXISTS cache (
|
||||
ns TEXT,
|
||||
key TEXT,
|
||||
expiry REAL,
|
||||
encoding TEXT NOT NULL,
|
||||
val BLOB NOT NULL,
|
||||
PRIMARY KEY (ns, key)
|
||||
)"""
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
|
||||
"""Get the cached values for the given keys."""
|
||||
with self._lock, self._conn:
|
||||
now = datetime.datetime.now(datetime.timezone.utc).timestamp()
|
||||
if not keys:
|
||||
return {}
|
||||
placeholders = ",".join("(?, ?)" for _ in keys)
|
||||
params: list[str] = []
|
||||
for ns_tuple, key in keys:
|
||||
params.extend((",".join(ns_tuple), key))
|
||||
cursor = self._conn.execute(
|
||||
f"SELECT ns, key, expiry, encoding, val FROM cache WHERE (ns, key) IN ({placeholders})",
|
||||
tuple(params),
|
||||
)
|
||||
values: dict[FullKey, ValueT] = {}
|
||||
rows = cursor.fetchall()
|
||||
for ns, key, expiry, encoding, raw in rows:
|
||||
if expiry is not None and now > expiry:
|
||||
# purge expired entry
|
||||
self._conn.execute(
|
||||
"DELETE FROM cache WHERE (ns, key) = (?, ?)", (ns, key)
|
||||
)
|
||||
continue
|
||||
values[(tuple(ns.split(",")), key)] = self.serde.loads_typed(
|
||||
(encoding, raw)
|
||||
)
|
||||
return values
|
||||
|
||||
async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
|
||||
"""Asynchronously get the cached values for the given keys."""
|
||||
return await asyncio.to_thread(self.get, keys)
|
||||
|
||||
def set(self, mapping: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
|
||||
"""Set the cached values for the given keys and TTLs."""
|
||||
with self._lock, self._conn:
|
||||
now = datetime.datetime.now(datetime.timezone.utc)
|
||||
for key, (value, ttl) in mapping.items():
|
||||
if ttl is not None:
|
||||
delta = datetime.timedelta(seconds=ttl)
|
||||
expiry: float | None = (now + delta).timestamp()
|
||||
else:
|
||||
expiry = None
|
||||
encoding, raw = self.serde.dumps_typed(value)
|
||||
self._conn.execute(
|
||||
"INSERT OR REPLACE INTO cache (ns, key, expiry, encoding, val) VALUES (?, ?, ?, ?, ?)",
|
||||
(",".join(key[0]), key[1], expiry, encoding, raw),
|
||||
)
|
||||
|
||||
async def aset(self, mapping: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
|
||||
"""Asynchronously set the cached values for the given keys and TTLs."""
|
||||
await asyncio.to_thread(self.set, mapping)
|
||||
|
||||
def clear(self, namespaces: Sequence[Namespace] | None = None) -> None:
|
||||
"""Delete the cached values for the given namespaces.
|
||||
If no namespaces are provided, clear all cached values."""
|
||||
with self._lock, self._conn:
|
||||
if namespaces is None:
|
||||
self._conn.execute("DELETE FROM cache")
|
||||
else:
|
||||
placeholders = ",".join("?" for _ in namespaces)
|
||||
self._conn.execute(
|
||||
f"DELETE FROM cache WHERE (ns) IN ({placeholders})",
|
||||
tuple(",".join(key) for key in namespaces),
|
||||
)
|
||||
|
||||
async def aclear(self, namespaces: Sequence[Namespace] | None = None) -> None:
|
||||
"""Asynchronously delete the cached values for the given namespaces.
|
||||
If no namespaces are provided, clear all cached values."""
|
||||
await asyncio.to_thread(self.clear, namespaces)
|
||||
|
||||
def __del__(self) -> None:
|
||||
try:
|
||||
self._conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
@@ -1,556 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import random
|
||||
import sqlite3
|
||||
import threading
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from contextlib import closing, contextmanager
|
||||
from typing import Any, cast
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
WRITES_IDX_MAP,
|
||||
BaseCheckpointSaver,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
SerializerProtocol,
|
||||
get_checkpoint_id,
|
||||
get_checkpoint_metadata,
|
||||
)
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
from langgraph.checkpoint.sqlite.utils import search_where
|
||||
|
||||
_AIO_ERROR_MSG = (
|
||||
"The SqliteSaver does not support async methods. "
|
||||
"Consider using AsyncSqliteSaver instead.\n"
|
||||
"from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver\n"
|
||||
"Note: AsyncSqliteSaver requires the aiosqlite package to use.\n"
|
||||
"Install with:\n`pip install aiosqlite`\n"
|
||||
"See https://langchain-ai.github.io/langgraph/reference/checkpoints/#langgraph.checkpoint.sqlite.aio.AsyncSqliteSaver"
|
||||
"for more information."
|
||||
)
|
||||
|
||||
|
||||
class SqliteSaver(BaseCheckpointSaver[str]):
|
||||
"""A checkpoint saver that stores checkpoints in a SQLite database.
|
||||
|
||||
Note:
|
||||
This class is meant for lightweight, synchronous use cases
|
||||
(demos and small projects) and does not
|
||||
scale to multiple threads.
|
||||
For a similar sqlite saver with `async` support,
|
||||
consider using [AsyncSqliteSaver][langgraph.checkpoint.sqlite.aio.AsyncSqliteSaver].
|
||||
|
||||
Args:
|
||||
conn (sqlite3.Connection): The SQLite database connection.
|
||||
serde (Optional[SerializerProtocol]): The serializer to use for serializing and deserializing checkpoints. Defaults to JsonPlusSerializerCompat.
|
||||
|
||||
Examples:
|
||||
|
||||
>>> import sqlite3
|
||||
>>> from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
>>> from langgraph.graph import StateGraph
|
||||
>>>
|
||||
>>> builder = StateGraph(int)
|
||||
>>> builder.add_node("add_one", lambda x: x + 1)
|
||||
>>> builder.set_entry_point("add_one")
|
||||
>>> builder.set_finish_point("add_one")
|
||||
>>> # Create a new SqliteSaver instance
|
||||
>>> # Note: check_same_thread=False is OK as the implementation uses a lock
|
||||
>>> # to ensure thread safety.
|
||||
>>> conn = sqlite3.connect("checkpoints.sqlite", check_same_thread=False)
|
||||
>>> memory = SqliteSaver(conn)
|
||||
>>> graph = builder.compile(checkpointer=memory)
|
||||
>>> config = {"configurable": {"thread_id": "1"}}
|
||||
>>> graph.get_state(config)
|
||||
>>> result = graph.invoke(3, config)
|
||||
>>> graph.get_state(config)
|
||||
StateSnapshot(values=4, next=(), config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '0c62ca34-ac19-445d-bbb0-5b4984975b2a'}}, parent_config=None)
|
||||
""" # noqa
|
||||
|
||||
conn: sqlite3.Connection
|
||||
is_setup: bool
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
*,
|
||||
serde: SerializerProtocol | None = None,
|
||||
) -> None:
|
||||
super().__init__(serde=serde)
|
||||
self.jsonplus_serde = JsonPlusSerializer()
|
||||
self.conn = conn
|
||||
self.is_setup = False
|
||||
self.lock = threading.Lock()
|
||||
|
||||
@classmethod
|
||||
@contextmanager
|
||||
def from_conn_string(cls, conn_string: str) -> Iterator[SqliteSaver]:
|
||||
"""Create a new SqliteSaver instance from a connection string.
|
||||
|
||||
Args:
|
||||
conn_string: The SQLite connection string.
|
||||
|
||||
Yields:
|
||||
SqliteSaver: A new SqliteSaver instance.
|
||||
|
||||
Examples:
|
||||
|
||||
In memory:
|
||||
|
||||
with SqliteSaver.from_conn_string(":memory:") as memory:
|
||||
...
|
||||
|
||||
To disk:
|
||||
|
||||
with SqliteSaver.from_conn_string("checkpoints.sqlite") as memory:
|
||||
...
|
||||
"""
|
||||
with closing(
|
||||
sqlite3.connect(
|
||||
conn_string,
|
||||
# https://ricardoanderegg.com/posts/python-sqlite-thread-safety/
|
||||
check_same_thread=False,
|
||||
)
|
||||
) as conn:
|
||||
yield cls(conn)
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Set up the checkpoint database.
|
||||
|
||||
This method creates the necessary tables in the SQLite database if they don't
|
||||
already exist. It is called automatically when needed and should not be called
|
||||
directly by the user.
|
||||
"""
|
||||
if self.is_setup:
|
||||
return
|
||||
|
||||
self.conn.executescript(
|
||||
"""
|
||||
PRAGMA journal_mode=WAL;
|
||||
CREATE TABLE IF NOT EXISTS checkpoints (
|
||||
thread_id TEXT NOT NULL,
|
||||
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
||||
checkpoint_id TEXT NOT NULL,
|
||||
parent_checkpoint_id TEXT,
|
||||
type TEXT,
|
||||
checkpoint BLOB,
|
||||
metadata BLOB,
|
||||
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS writes (
|
||||
thread_id TEXT NOT NULL,
|
||||
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
||||
checkpoint_id TEXT NOT NULL,
|
||||
task_id TEXT NOT NULL,
|
||||
idx INTEGER NOT NULL,
|
||||
channel TEXT NOT NULL,
|
||||
type TEXT,
|
||||
value BLOB,
|
||||
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
self.is_setup = True
|
||||
|
||||
@contextmanager
|
||||
def cursor(self, transaction: bool = True) -> Iterator[sqlite3.Cursor]:
|
||||
"""Get a cursor for the SQLite database.
|
||||
|
||||
This method returns a cursor for the SQLite database. It is used internally
|
||||
by the SqliteSaver and should not be called directly by the user.
|
||||
|
||||
Args:
|
||||
transaction (bool): Whether to commit the transaction when the cursor is closed. Defaults to True.
|
||||
|
||||
Yields:
|
||||
sqlite3.Cursor: A cursor for the SQLite database.
|
||||
"""
|
||||
with self.lock:
|
||||
self.setup()
|
||||
cur = self.conn.cursor()
|
||||
try:
|
||||
yield cur
|
||||
finally:
|
||||
if transaction:
|
||||
self.conn.commit()
|
||||
cur.close()
|
||||
|
||||
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Get a checkpoint tuple from the database.
|
||||
|
||||
This method retrieves a checkpoint tuple from the SQLite database based on the
|
||||
provided config. If the config contains a `checkpoint_id` key, the checkpoint with
|
||||
the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint
|
||||
for the given thread ID is retrieved.
|
||||
|
||||
Args:
|
||||
config: The config to use for retrieving the checkpoint.
|
||||
|
||||
Returns:
|
||||
The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
||||
|
||||
Examples:
|
||||
|
||||
Basic:
|
||||
>>> config = {"configurable": {"thread_id": "1"}}
|
||||
>>> checkpoint_tuple = memory.get_tuple(config)
|
||||
>>> print(checkpoint_tuple)
|
||||
CheckpointTuple(...)
|
||||
|
||||
With checkpoint ID:
|
||||
|
||||
>>> config = {
|
||||
... "configurable": {
|
||||
... "thread_id": "1",
|
||||
... "checkpoint_ns": "",
|
||||
... "checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875",
|
||||
... }
|
||||
... }
|
||||
>>> checkpoint_tuple = memory.get_tuple(config)
|
||||
>>> print(checkpoint_tuple)
|
||||
CheckpointTuple(...)
|
||||
""" # noqa
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
with self.cursor(transaction=False) as cur:
|
||||
# find the latest checkpoint for the thread_id
|
||||
if checkpoint_id := get_checkpoint_id(config):
|
||||
cur.execute(
|
||||
"SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
|
||||
(
|
||||
str(config["configurable"]["thread_id"]),
|
||||
checkpoint_ns,
|
||||
checkpoint_id,
|
||||
),
|
||||
)
|
||||
else:
|
||||
cur.execute(
|
||||
"SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? ORDER BY checkpoint_id DESC LIMIT 1",
|
||||
(str(config["configurable"]["thread_id"]), checkpoint_ns),
|
||||
)
|
||||
# if a checkpoint is found, return it
|
||||
if value := cur.fetchone():
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_id,
|
||||
parent_checkpoint_id,
|
||||
type,
|
||||
checkpoint,
|
||||
metadata,
|
||||
) = value
|
||||
if not get_checkpoint_id(config):
|
||||
config = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint_id,
|
||||
}
|
||||
}
|
||||
# find any pending writes
|
||||
cur.execute(
|
||||
"SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? ORDER BY task_id, idx",
|
||||
(
|
||||
str(config["configurable"]["thread_id"]),
|
||||
checkpoint_ns,
|
||||
str(config["configurable"]["checkpoint_id"]),
|
||||
),
|
||||
)
|
||||
# deserialize the checkpoint and metadata
|
||||
return CheckpointTuple(
|
||||
config,
|
||||
self.serde.loads_typed((type, checkpoint)),
|
||||
cast(
|
||||
CheckpointMetadata,
|
||||
json.loads(metadata) if metadata is not None else {},
|
||||
),
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": parent_checkpoint_id,
|
||||
}
|
||||
}
|
||||
if parent_checkpoint_id
|
||||
else None
|
||||
),
|
||||
[
|
||||
(task_id, channel, self.serde.loads_typed((type, value)))
|
||||
for task_id, channel, type, value in cur
|
||||
],
|
||||
)
|
||||
|
||||
def list(
|
||||
self,
|
||||
config: RunnableConfig | None,
|
||||
*,
|
||||
filter: dict[str, Any] | None = None,
|
||||
before: RunnableConfig | None = None,
|
||||
limit: int | None = None,
|
||||
) -> Iterator[CheckpointTuple]:
|
||||
"""List checkpoints from the database.
|
||||
|
||||
This method retrieves a list of checkpoint tuples from the SQLite database based
|
||||
on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
|
||||
|
||||
Args:
|
||||
config: The config to use for listing the checkpoints.
|
||||
filter: Additional filtering criteria for metadata.
|
||||
before: If provided, only checkpoints before the specified checkpoint ID are returned.
|
||||
limit: The maximum number of checkpoints to return.
|
||||
|
||||
Yields:
|
||||
An iterator of checkpoint tuples.
|
||||
|
||||
Examples:
|
||||
>>> from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
>>> with SqliteSaver.from_conn_string(":memory:") as memory:
|
||||
... # Run a graph, then list the checkpoints
|
||||
>>> config = {"configurable": {"thread_id": "1"}}
|
||||
>>> checkpoints = list(memory.list(config, limit=2))
|
||||
>>> print(checkpoints)
|
||||
[CheckpointTuple(...), CheckpointTuple(...)]
|
||||
|
||||
>>> config = {"configurable": {"thread_id": "1"}}
|
||||
>>> before = {"configurable": {"checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875"}}
|
||||
>>> with SqliteSaver.from_conn_string(":memory:") as memory:
|
||||
... # Run a graph, then list the checkpoints
|
||||
>>> checkpoints = list(memory.list(config, before=before))
|
||||
>>> print(checkpoints)
|
||||
[CheckpointTuple(...), ...]
|
||||
"""
|
||||
where, param_values = search_where(config, filter, before)
|
||||
query = f"""SELECT thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata
|
||||
FROM checkpoints
|
||||
{where}
|
||||
ORDER BY checkpoint_id DESC"""
|
||||
if limit is not None:
|
||||
query += " LIMIT ?"
|
||||
param_values = (*param_values, limit)
|
||||
with self.cursor(transaction=False) as cur, closing(self.conn.cursor()) as wcur:
|
||||
cur.execute(query, param_values)
|
||||
for (
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_id,
|
||||
parent_checkpoint_id,
|
||||
type,
|
||||
checkpoint,
|
||||
metadata,
|
||||
) in cur:
|
||||
wcur.execute(
|
||||
"SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? ORDER BY task_id, idx",
|
||||
(thread_id, checkpoint_ns, checkpoint_id),
|
||||
)
|
||||
yield CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint_id,
|
||||
}
|
||||
},
|
||||
self.serde.loads_typed((type, checkpoint)),
|
||||
cast(
|
||||
CheckpointMetadata,
|
||||
json.loads(metadata) if metadata is not None else {},
|
||||
),
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": parent_checkpoint_id,
|
||||
}
|
||||
}
|
||||
if parent_checkpoint_id
|
||||
else None
|
||||
),
|
||||
[
|
||||
(task_id, channel, self.serde.loads_typed((type, value)))
|
||||
for task_id, channel, type, value in wcur
|
||||
],
|
||||
)
|
||||
|
||||
def put(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
"""Save a checkpoint to the database.
|
||||
|
||||
This method saves a checkpoint to the SQLite database. The checkpoint is associated
|
||||
with the provided config and its parent config (if any).
|
||||
|
||||
Args:
|
||||
config: The config to associate with the checkpoint.
|
||||
checkpoint: The checkpoint to save.
|
||||
metadata: Additional metadata to save with the checkpoint.
|
||||
new_versions: New channel versions as of this write.
|
||||
|
||||
Returns:
|
||||
RunnableConfig: Updated configuration after storing the checkpoint.
|
||||
|
||||
Examples:
|
||||
|
||||
>>> from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
>>> with SqliteSaver.from_conn_string(":memory:") as memory:
|
||||
>>> config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
|
||||
>>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}}
|
||||
>>> saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {})
|
||||
>>> print(saved_config)
|
||||
{'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}}
|
||||
"""
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"]["checkpoint_ns"]
|
||||
type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
|
||||
serialized_metadata = json.dumps(
|
||||
get_checkpoint_metadata(config, metadata), ensure_ascii=False
|
||||
).encode("utf-8", "ignore")
|
||||
with self.cursor() as cur:
|
||||
cur.execute(
|
||||
"INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
(
|
||||
str(config["configurable"]["thread_id"]),
|
||||
checkpoint_ns,
|
||||
checkpoint["id"],
|
||||
config["configurable"].get("checkpoint_id"),
|
||||
type_,
|
||||
serialized_checkpoint,
|
||||
serialized_metadata,
|
||||
),
|
||||
)
|
||||
return {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
}
|
||||
|
||||
def put_writes(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
writes: Sequence[tuple[str, Any]],
|
||||
task_id: str,
|
||||
task_path: str = "",
|
||||
) -> None:
|
||||
"""Store intermediate writes linked to a checkpoint.
|
||||
|
||||
This method saves intermediate writes associated with a checkpoint to the SQLite database.
|
||||
|
||||
Args:
|
||||
config: Configuration of the related checkpoint.
|
||||
writes: List of writes to store, each as (channel, value) pair.
|
||||
task_id: Identifier for the task creating the writes.
|
||||
task_path: Path of the task creating the writes.
|
||||
"""
|
||||
query = (
|
||||
"INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
|
||||
if all(w[0] in WRITES_IDX_MAP for w in writes)
|
||||
else "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
|
||||
)
|
||||
with self.cursor() as cur:
|
||||
cur.executemany(
|
||||
query,
|
||||
[
|
||||
(
|
||||
str(config["configurable"]["thread_id"]),
|
||||
str(config["configurable"]["checkpoint_ns"]),
|
||||
str(config["configurable"]["checkpoint_id"]),
|
||||
task_id,
|
||||
WRITES_IDX_MAP.get(channel, idx),
|
||||
channel,
|
||||
*self.serde.dumps_typed(value),
|
||||
)
|
||||
for idx, (channel, value) in enumerate(writes)
|
||||
],
|
||||
)
|
||||
|
||||
def delete_thread(self, thread_id: str) -> None:
|
||||
"""Delete all checkpoints and writes associated with a thread ID.
|
||||
|
||||
Args:
|
||||
thread_id: The thread ID to delete.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
with self.cursor() as cur:
|
||||
cur.execute(
|
||||
"DELETE FROM checkpoints WHERE thread_id = ?",
|
||||
(str(thread_id),),
|
||||
)
|
||||
cur.execute(
|
||||
"DELETE FROM writes WHERE thread_id = ?",
|
||||
(str(thread_id),),
|
||||
)
|
||||
|
||||
async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Get a checkpoint tuple from the database asynchronously.
|
||||
|
||||
Note:
|
||||
This async method is not supported by the SqliteSaver class.
|
||||
Use get_tuple() instead, or consider using [AsyncSqliteSaver][langgraph.checkpoint.sqlite.aio.AsyncSqliteSaver].
|
||||
"""
|
||||
raise NotImplementedError(_AIO_ERROR_MSG)
|
||||
|
||||
async def alist(
|
||||
self,
|
||||
config: RunnableConfig | None,
|
||||
*,
|
||||
filter: dict[str, Any] | None = None,
|
||||
before: RunnableConfig | None = None,
|
||||
limit: int | None = None,
|
||||
) -> AsyncIterator[CheckpointTuple]:
|
||||
"""List checkpoints from the database asynchronously.
|
||||
|
||||
Note:
|
||||
This async method is not supported by the SqliteSaver class.
|
||||
Use list() instead, or consider using [AsyncSqliteSaver][langgraph.checkpoint.sqlite.aio.AsyncSqliteSaver].
|
||||
"""
|
||||
raise NotImplementedError(_AIO_ERROR_MSG)
|
||||
yield
|
||||
|
||||
async def aput(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
"""Save a checkpoint to the database asynchronously.
|
||||
|
||||
Note:
|
||||
This async method is not supported by the SqliteSaver class.
|
||||
Use put() instead, or consider using [AsyncSqliteSaver][langgraph.checkpoint.sqlite.aio.AsyncSqliteSaver].
|
||||
"""
|
||||
raise NotImplementedError(_AIO_ERROR_MSG)
|
||||
|
||||
def get_next_version(self, current: str | None, channel: None) -> str:
|
||||
"""Generate the next version ID for a channel.
|
||||
|
||||
This method creates a new version identifier for a channel based on its current version.
|
||||
|
||||
Args:
|
||||
current (Optional[str]): The current version identifier of the channel.
|
||||
|
||||
Returns:
|
||||
str: The next version identifier, which is guaranteed to be monotonically increasing.
|
||||
"""
|
||||
if current is None:
|
||||
current_v = 0
|
||||
elif isinstance(current, int):
|
||||
current_v = current
|
||||
else:
|
||||
current_v = int(current.split(".")[0])
|
||||
next_v = current_v + 1
|
||||
next_h = random.random()
|
||||
return f"{next_v:032}.{next_h:016}"
|
||||
@@ -1,631 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import random
|
||||
import threading
|
||||
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, TypeVar, cast
|
||||
|
||||
import aiosqlite
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
WRITES_IDX_MAP,
|
||||
BaseCheckpointSaver,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
SerializerProtocol,
|
||||
get_checkpoint_id,
|
||||
get_checkpoint_metadata,
|
||||
)
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
from langgraph.checkpoint.sqlite.utils import search_where
|
||||
|
||||
T = TypeVar("T", bound=Callable)
|
||||
|
||||
|
||||
class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
||||
"""An asynchronous checkpoint saver that stores checkpoints in a SQLite database.
|
||||
|
||||
This class provides an asynchronous interface for saving and retrieving checkpoints
|
||||
using a SQLite database. It's designed for use in asynchronous environments and
|
||||
offers better performance for I/O-bound operations compared to synchronous alternatives.
|
||||
|
||||
Attributes:
|
||||
conn (aiosqlite.Connection): The asynchronous SQLite database connection.
|
||||
serde (SerializerProtocol): The serializer used for encoding/decoding checkpoints.
|
||||
|
||||
Tip:
|
||||
Requires the [aiosqlite](https://pypi.org/project/aiosqlite/) package.
|
||||
Install it with `pip install aiosqlite`.
|
||||
|
||||
Warning:
|
||||
While this class supports asynchronous checkpointing, it is not recommended
|
||||
for production workloads due to limitations in SQLite's write performance.
|
||||
For production use, consider a more robust database like PostgreSQL.
|
||||
|
||||
Tip:
|
||||
Remember to **close the database connection** after executing your code,
|
||||
otherwise, you may see the graph "hang" after execution (since the program
|
||||
will not exit until the connection is closed).
|
||||
|
||||
The easiest way is to use the `async with` statement as shown in the examples.
|
||||
|
||||
```python
|
||||
async with AsyncSqliteSaver.from_conn_string("checkpoints.sqlite") as saver:
|
||||
# Your code here
|
||||
graph = builder.compile(checkpointer=saver)
|
||||
config = {"configurable": {"thread_id": "thread-1"}}
|
||||
async for event in graph.astream_events(..., config, version="v1"):
|
||||
print(event)
|
||||
```
|
||||
|
||||
Examples:
|
||||
Usage within StateGraph:
|
||||
|
||||
```pycon
|
||||
>>> import asyncio
|
||||
>>>
|
||||
>>> from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||||
>>> from langgraph.graph import StateGraph
|
||||
>>>
|
||||
>>> async def main():
|
||||
>>> builder = StateGraph(int)
|
||||
>>> builder.add_node("add_one", lambda x: x + 1)
|
||||
>>> builder.set_entry_point("add_one")
|
||||
>>> builder.set_finish_point("add_one")
|
||||
>>> async with AsyncSqliteSaver.from_conn_string("checkpoints.db") as memory:
|
||||
>>> graph = builder.compile(checkpointer=memory)
|
||||
>>> coro = graph.ainvoke(1, {"configurable": {"thread_id": "thread-1"}})
|
||||
>>> print(await asyncio.gather(coro))
|
||||
>>>
|
||||
>>> asyncio.run(main())
|
||||
Output: [2]
|
||||
```
|
||||
Raw usage:
|
||||
|
||||
```pycon
|
||||
>>> import asyncio
|
||||
>>> import aiosqlite
|
||||
>>> from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||||
>>>
|
||||
>>> async def main():
|
||||
>>> async with aiosqlite.connect("checkpoints.db") as conn:
|
||||
... saver = AsyncSqliteSaver(conn)
|
||||
... config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
|
||||
... checkpoint = {"ts": "2023-05-03T10:00:00Z", "data": {"key": "value"}, "id": "0c62ca34-ac19-445d-bbb0-5b4984975b2a"}
|
||||
... saved_config = await saver.aput(config, checkpoint, {}, {})
|
||||
... print(saved_config)
|
||||
>>> asyncio.run(main())
|
||||
{'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '0c62ca34-ac19-445d-bbb0-5b4984975b2a'}}
|
||||
```
|
||||
"""
|
||||
|
||||
lock: asyncio.Lock
|
||||
is_setup: bool
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: aiosqlite.Connection,
|
||||
*,
|
||||
serde: SerializerProtocol | None = None,
|
||||
):
|
||||
super().__init__(serde=serde)
|
||||
self.jsonplus_serde = JsonPlusSerializer()
|
||||
self.conn = conn
|
||||
self.lock = asyncio.Lock()
|
||||
self.loop = asyncio.get_running_loop()
|
||||
self.is_setup = False
|
||||
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
async def from_conn_string(
|
||||
cls, conn_string: str
|
||||
) -> AsyncIterator[AsyncSqliteSaver]:
|
||||
"""Create a new AsyncSqliteSaver instance from a connection string.
|
||||
|
||||
Args:
|
||||
conn_string: The SQLite connection string.
|
||||
|
||||
Yields:
|
||||
AsyncSqliteSaver: A new AsyncSqliteSaver instance.
|
||||
"""
|
||||
async with aiosqlite.connect(conn_string) as conn:
|
||||
yield cls(conn)
|
||||
|
||||
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Get a checkpoint tuple from the database.
|
||||
|
||||
This method retrieves a checkpoint tuple from the SQLite database based on the
|
||||
provided config. If the config contains a `checkpoint_id` key, the checkpoint with
|
||||
the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint
|
||||
for the given thread ID is retrieved.
|
||||
|
||||
Args:
|
||||
config: The config to use for retrieving the checkpoint.
|
||||
|
||||
Returns:
|
||||
The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
||||
"""
|
||||
try:
|
||||
# check if we are in the main thread, only bg threads can block
|
||||
# we don't check in other methods to avoid the overhead
|
||||
if asyncio.get_running_loop() is self.loop:
|
||||
raise asyncio.InvalidStateError(
|
||||
"Synchronous calls to AsyncSqliteSaver are only allowed from a "
|
||||
"different thread. From the main thread, use the async interface. "
|
||||
"For example, use `await checkpointer.aget_tuple(...)` or `await "
|
||||
"graph.ainvoke(...)`."
|
||||
)
|
||||
except RuntimeError:
|
||||
pass
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.aget_tuple(config), self.loop
|
||||
).result()
|
||||
|
||||
def list(
|
||||
self,
|
||||
config: RunnableConfig | None,
|
||||
*,
|
||||
filter: dict[str, Any] | None = None,
|
||||
before: RunnableConfig | None = None,
|
||||
limit: int | None = None,
|
||||
) -> Iterator[CheckpointTuple]:
|
||||
"""List checkpoints from the database asynchronously.
|
||||
|
||||
This method retrieves a list of checkpoint tuples from the SQLite database based
|
||||
on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
|
||||
|
||||
Args:
|
||||
config: Base configuration for filtering checkpoints.
|
||||
filter: Additional filtering criteria for metadata.
|
||||
before: If provided, only checkpoints before the specified checkpoint ID are returned.
|
||||
limit: Maximum number of checkpoints to return.
|
||||
|
||||
Yields:
|
||||
An iterator of matching checkpoint tuples.
|
||||
"""
|
||||
try:
|
||||
# check if we are in the main thread, only bg threads can block
|
||||
# we don't check in other methods to avoid the overhead
|
||||
if asyncio.get_running_loop() is self.loop:
|
||||
raise asyncio.InvalidStateError(
|
||||
"Synchronous calls to AsyncSqliteSaver are only allowed from a "
|
||||
"different thread. From the main thread, use the async interface. "
|
||||
"For example, use `checkpointer.alist(...)` or `await "
|
||||
"graph.ainvoke(...)`."
|
||||
)
|
||||
except RuntimeError:
|
||||
pass
|
||||
aiter_ = self.alist(config, filter=filter, before=before, limit=limit)
|
||||
while True:
|
||||
try:
|
||||
yield asyncio.run_coroutine_threadsafe(
|
||||
anext(aiter_), # type: ignore[arg-type] # noqa: F821
|
||||
self.loop,
|
||||
).result()
|
||||
except StopAsyncIteration:
|
||||
break
|
||||
|
||||
def put(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
"""Save a checkpoint to the database.
|
||||
|
||||
This method saves a checkpoint to the SQLite database. The checkpoint is associated
|
||||
with the provided config and its parent config (if any).
|
||||
|
||||
Args:
|
||||
config: The config to associate with the checkpoint.
|
||||
checkpoint: The checkpoint to save.
|
||||
metadata: Additional metadata to save with the checkpoint.
|
||||
new_versions: New channel versions as of this write.
|
||||
|
||||
Returns:
|
||||
RunnableConfig: Updated configuration after storing the checkpoint.
|
||||
"""
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.aput(config, checkpoint, metadata, new_versions), self.loop
|
||||
).result()
|
||||
|
||||
def put_writes(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
writes: Sequence[tuple[str, Any]],
|
||||
task_id: str,
|
||||
task_path: str = "",
|
||||
) -> None:
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.aput_writes(config, writes, task_id, task_path), self.loop
|
||||
).result()
|
||||
|
||||
def delete_thread(self, thread_id: str) -> None:
|
||||
"""Delete all checkpoints and writes associated with a thread ID.
|
||||
|
||||
Args:
|
||||
thread_id: The thread ID to delete.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
try:
|
||||
# check if we are in the main thread, only bg threads can block
|
||||
# we don't check in other methods to avoid the overhead
|
||||
if asyncio.get_running_loop() is self.loop:
|
||||
raise asyncio.InvalidStateError(
|
||||
"Synchronous calls to AsyncSqliteSaver are only allowed from a "
|
||||
"different thread. From the main thread, use the async interface. "
|
||||
"For example, use `checkpointer.alist(...)` or `await "
|
||||
"graph.ainvoke(...)`."
|
||||
)
|
||||
except RuntimeError:
|
||||
pass
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.adelete_thread(thread_id), self.loop
|
||||
).result()
|
||||
|
||||
async def setup(self) -> None:
|
||||
"""Set up the checkpoint database asynchronously.
|
||||
|
||||
This method creates the necessary tables in the SQLite database if they don't
|
||||
already exist. It is called automatically when needed and should not be called
|
||||
directly by the user.
|
||||
"""
|
||||
async with self.lock:
|
||||
if self.is_setup:
|
||||
return
|
||||
await _ensure_connected(self.conn)
|
||||
async with self.conn.executescript(
|
||||
"""
|
||||
PRAGMA journal_mode=WAL;
|
||||
CREATE TABLE IF NOT EXISTS checkpoints (
|
||||
thread_id TEXT NOT NULL,
|
||||
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
||||
checkpoint_id TEXT NOT NULL,
|
||||
parent_checkpoint_id TEXT,
|
||||
type TEXT,
|
||||
checkpoint BLOB,
|
||||
metadata BLOB,
|
||||
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS writes (
|
||||
thread_id TEXT NOT NULL,
|
||||
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
||||
checkpoint_id TEXT NOT NULL,
|
||||
task_id TEXT NOT NULL,
|
||||
idx INTEGER NOT NULL,
|
||||
channel TEXT NOT NULL,
|
||||
type TEXT,
|
||||
value BLOB,
|
||||
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
|
||||
);
|
||||
"""
|
||||
):
|
||||
await self.conn.commit()
|
||||
|
||||
self.is_setup = True
|
||||
|
||||
async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Get a checkpoint tuple from the database asynchronously.
|
||||
|
||||
This method retrieves a checkpoint tuple from the SQLite database based on the
|
||||
provided config. If the config contains a `checkpoint_id` key, the checkpoint with
|
||||
the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint
|
||||
for the given thread ID is retrieved.
|
||||
|
||||
Args:
|
||||
config: The config to use for retrieving the checkpoint.
|
||||
|
||||
Returns:
|
||||
The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
||||
"""
|
||||
await self.setup()
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
async with self.lock, self.conn.cursor() as cur:
|
||||
# find the latest checkpoint for the thread_id
|
||||
if checkpoint_id := get_checkpoint_id(config):
|
||||
await cur.execute(
|
||||
"SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
|
||||
(
|
||||
str(config["configurable"]["thread_id"]),
|
||||
checkpoint_ns,
|
||||
checkpoint_id,
|
||||
),
|
||||
)
|
||||
else:
|
||||
await cur.execute(
|
||||
"SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? ORDER BY checkpoint_id DESC LIMIT 1",
|
||||
(str(config["configurable"]["thread_id"]), checkpoint_ns),
|
||||
)
|
||||
# if a checkpoint is found, return it
|
||||
if value := await cur.fetchone():
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_id,
|
||||
parent_checkpoint_id,
|
||||
type,
|
||||
checkpoint,
|
||||
metadata,
|
||||
) = value
|
||||
if not get_checkpoint_id(config):
|
||||
config = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint_id,
|
||||
}
|
||||
}
|
||||
# find any pending writes
|
||||
await cur.execute(
|
||||
"SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? ORDER BY task_id, idx",
|
||||
(
|
||||
str(config["configurable"]["thread_id"]),
|
||||
checkpoint_ns,
|
||||
str(config["configurable"]["checkpoint_id"]),
|
||||
),
|
||||
)
|
||||
# deserialize the checkpoint and metadata
|
||||
return CheckpointTuple(
|
||||
config,
|
||||
self.serde.loads_typed((type, checkpoint)),
|
||||
cast(
|
||||
CheckpointMetadata,
|
||||
(json.loads(metadata) if metadata is not None else {}),
|
||||
),
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": parent_checkpoint_id,
|
||||
}
|
||||
}
|
||||
if parent_checkpoint_id
|
||||
else None
|
||||
),
|
||||
[
|
||||
(task_id, channel, self.serde.loads_typed((type, value)))
|
||||
async for task_id, channel, type, value in cur
|
||||
],
|
||||
)
|
||||
|
||||
async def alist(
|
||||
self,
|
||||
config: RunnableConfig | None,
|
||||
*,
|
||||
filter: dict[str, Any] | None = None,
|
||||
before: RunnableConfig | None = None,
|
||||
limit: int | None = None,
|
||||
) -> AsyncIterator[CheckpointTuple]:
|
||||
"""List checkpoints from the database asynchronously.
|
||||
|
||||
This method retrieves a list of checkpoint tuples from the SQLite database based
|
||||
on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
|
||||
|
||||
Args:
|
||||
config: Base configuration for filtering checkpoints.
|
||||
filter: Additional filtering criteria for metadata.
|
||||
before: If provided, only checkpoints before the specified checkpoint ID are returned.
|
||||
limit: Maximum number of checkpoints to return.
|
||||
|
||||
Yields:
|
||||
An asynchronous iterator of matching checkpoint tuples.
|
||||
"""
|
||||
await self.setup()
|
||||
where, params = search_where(config, filter, before)
|
||||
query = f"""SELECT thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata
|
||||
FROM checkpoints
|
||||
{where}
|
||||
ORDER BY checkpoint_id DESC"""
|
||||
if limit is not None:
|
||||
query += " LIMIT ?"
|
||||
params = (*params, limit)
|
||||
async with (
|
||||
self.lock,
|
||||
self.conn.execute(query, params) as cur,
|
||||
self.conn.cursor() as wcur,
|
||||
):
|
||||
async for (
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_id,
|
||||
parent_checkpoint_id,
|
||||
type,
|
||||
checkpoint,
|
||||
metadata,
|
||||
) in cur:
|
||||
await wcur.execute(
|
||||
"SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? ORDER BY task_id, idx",
|
||||
(thread_id, checkpoint_ns, checkpoint_id),
|
||||
)
|
||||
yield CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint_id,
|
||||
}
|
||||
},
|
||||
self.serde.loads_typed((type, checkpoint)),
|
||||
cast(
|
||||
CheckpointMetadata,
|
||||
(json.loads(metadata) if metadata is not None else {}),
|
||||
),
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": parent_checkpoint_id,
|
||||
}
|
||||
}
|
||||
if parent_checkpoint_id
|
||||
else None
|
||||
),
|
||||
[
|
||||
(task_id, channel, self.serde.loads_typed((type, value)))
|
||||
async for task_id, channel, type, value in wcur
|
||||
],
|
||||
)
|
||||
|
||||
async def aput(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
"""Save a checkpoint to the database asynchronously.
|
||||
|
||||
This method saves a checkpoint to the SQLite database. The checkpoint is associated
|
||||
with the provided config and its parent config (if any).
|
||||
|
||||
Args:
|
||||
config: The config to associate with the checkpoint.
|
||||
checkpoint: The checkpoint to save.
|
||||
metadata: Additional metadata to save with the checkpoint.
|
||||
new_versions: New channel versions as of this write.
|
||||
|
||||
Returns:
|
||||
RunnableConfig: Updated configuration after storing the checkpoint.
|
||||
"""
|
||||
await self.setup()
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"]["checkpoint_ns"]
|
||||
type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
|
||||
serialized_metadata = json.dumps(
|
||||
get_checkpoint_metadata(config, metadata), ensure_ascii=False
|
||||
).encode("utf-8", "ignore")
|
||||
async with (
|
||||
self.lock,
|
||||
self.conn.execute(
|
||||
"INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
(
|
||||
str(config["configurable"]["thread_id"]),
|
||||
checkpoint_ns,
|
||||
checkpoint["id"],
|
||||
config["configurable"].get("checkpoint_id"),
|
||||
type_,
|
||||
serialized_checkpoint,
|
||||
serialized_metadata,
|
||||
),
|
||||
),
|
||||
):
|
||||
await self.conn.commit()
|
||||
return {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
}
|
||||
|
||||
async def aput_writes(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
writes: Sequence[tuple[str, Any]],
|
||||
task_id: str,
|
||||
task_path: str = "",
|
||||
) -> None:
|
||||
"""Store intermediate writes linked to a checkpoint asynchronously.
|
||||
|
||||
This method saves intermediate writes associated with a checkpoint to the database.
|
||||
|
||||
Args:
|
||||
config: Configuration of the related checkpoint.
|
||||
writes: List of writes to store, each as (channel, value) pair.
|
||||
task_id: Identifier for the task creating the writes.
|
||||
task_path: Path of the task creating the writes.
|
||||
"""
|
||||
query = (
|
||||
"INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
|
||||
if all(w[0] in WRITES_IDX_MAP for w in writes)
|
||||
else "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
|
||||
)
|
||||
await self.setup()
|
||||
async with self.lock, self.conn.cursor() as cur:
|
||||
await cur.executemany(
|
||||
query,
|
||||
[
|
||||
(
|
||||
str(config["configurable"]["thread_id"]),
|
||||
str(config["configurable"]["checkpoint_ns"]),
|
||||
str(config["configurable"]["checkpoint_id"]),
|
||||
task_id,
|
||||
WRITES_IDX_MAP.get(channel, idx),
|
||||
channel,
|
||||
*self.serde.dumps_typed(value),
|
||||
)
|
||||
for idx, (channel, value) in enumerate(writes)
|
||||
],
|
||||
)
|
||||
await self.conn.commit()
|
||||
|
||||
async def adelete_thread(self, thread_id: str) -> None:
|
||||
"""Delete all checkpoints and writes associated with a thread ID.
|
||||
|
||||
Args:
|
||||
thread_id: The thread ID to delete.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
async with self.lock, self.conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"DELETE FROM checkpoints WHERE thread_id = ?",
|
||||
(str(thread_id),),
|
||||
)
|
||||
await cur.execute(
|
||||
"DELETE FROM writes WHERE thread_id = ?",
|
||||
(str(thread_id),),
|
||||
)
|
||||
await self.conn.commit()
|
||||
|
||||
def get_next_version(self, current: str | None, channel: None) -> str:
|
||||
"""Generate the next version ID for a channel.
|
||||
|
||||
This method creates a new version identifier for a channel based on its current version.
|
||||
|
||||
Args:
|
||||
current (Optional[str]): The current version identifier of the channel.
|
||||
|
||||
Returns:
|
||||
str: The next version identifier, which is guaranteed to be monotonically increasing.
|
||||
"""
|
||||
if current is None:
|
||||
current_v = 0
|
||||
elif isinstance(current, int):
|
||||
current_v = current
|
||||
else:
|
||||
current_v = int(current.split(".")[0])
|
||||
next_v = current_v + 1
|
||||
next_h = random.random()
|
||||
return f"{next_v:032}.{next_h:016}"
|
||||
|
||||
|
||||
async def _ensure_connected(conn: aiosqlite.Connection) -> None:
|
||||
if not _CONN_STARTED_CHECK(conn):
|
||||
await conn
|
||||
|
||||
|
||||
def _build_conn_started_check() -> Callable[[aiosqlite.Connection], bool]:
|
||||
is_alive = getattr(aiosqlite.Connection, "is_alive", None)
|
||||
if callable(is_alive):
|
||||
return lambda conn: conn.is_alive() # type: ignore[attr-defined]
|
||||
|
||||
def _started(conn: aiosqlite.Connection) -> bool:
|
||||
thread: threading.Thread | None = getattr(conn, "_thread", None)
|
||||
return False if thread is None else thread.is_alive()
|
||||
|
||||
return _started
|
||||
|
||||
|
||||
_CONN_STARTED_CHECK = _build_conn_started_check()
|
||||
@@ -1,116 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import get_checkpoint_id
|
||||
|
||||
_FILTER_PATTERN = re.compile(r"^[a-zA-Z0-9_.-]+$")
|
||||
|
||||
|
||||
def _validate_filter_key(key: str) -> None:
|
||||
"""Validate that a filter key is safe for use in SQL queries.
|
||||
|
||||
Args:
|
||||
key: The filter key to validate
|
||||
|
||||
Raises:
|
||||
ValueError: If the key contains invalid characters that could enable SQL injection
|
||||
"""
|
||||
# Allow alphanumeric characters, underscores, dots, and hyphens
|
||||
# This covers typical JSON property names while preventing SQL injection
|
||||
if not _FILTER_PATTERN.match(key):
|
||||
raise ValueError(
|
||||
f"Invalid filter key: '{key}'. Filter keys must contain only alphanumeric characters, underscores, dots, and hyphens."
|
||||
)
|
||||
|
||||
|
||||
def _metadata_predicate(
|
||||
metadata_filter: dict[str, Any],
|
||||
) -> tuple[Sequence[str], Sequence[Any]]:
|
||||
"""Return WHERE clause predicates for (a)search() given metadata filter.
|
||||
|
||||
This method returns a tuple of a string and a tuple of values. The string
|
||||
is the parametered WHERE clause predicate (excluding the WHERE keyword):
|
||||
"column1 = ? AND column2 IS ?". The tuple of values contains the values
|
||||
for each of the corresponding parameters.
|
||||
"""
|
||||
|
||||
def _where_value(query_value: Any) -> tuple[str, Any]:
|
||||
"""Return tuple of operator and value for WHERE clause predicate."""
|
||||
if query_value is None:
|
||||
return ("IS ?", None)
|
||||
elif (
|
||||
isinstance(query_value, str)
|
||||
or isinstance(query_value, int)
|
||||
or isinstance(query_value, float)
|
||||
):
|
||||
return ("= ?", query_value)
|
||||
elif isinstance(query_value, bool):
|
||||
return ("= ?", 1 if query_value else 0)
|
||||
elif isinstance(query_value, dict) or isinstance(query_value, list):
|
||||
# query value for JSON object cannot have trailing space after separators (, :)
|
||||
# SQLite json_extract() returns JSON string without whitespace
|
||||
return ("= ?", json.dumps(query_value, separators=(",", ":")))
|
||||
else:
|
||||
return ("= ?", str(query_value))
|
||||
|
||||
predicates = []
|
||||
param_values = []
|
||||
|
||||
# process metadata query
|
||||
for query_key, query_value in metadata_filter.items():
|
||||
_validate_filter_key(query_key)
|
||||
operator, param_value = _where_value(query_value)
|
||||
predicates.append(
|
||||
f"json_extract(CAST(metadata AS TEXT), '$.{query_key}') {operator}"
|
||||
)
|
||||
param_values.append(param_value)
|
||||
|
||||
return (predicates, param_values)
|
||||
|
||||
|
||||
def search_where(
|
||||
config: RunnableConfig | None,
|
||||
filter: dict[str, Any] | None,
|
||||
before: RunnableConfig | None = None,
|
||||
) -> tuple[str, Sequence[Any]]:
|
||||
"""Return WHERE clause predicates for (a)search() given metadata filter
|
||||
and `before` config.
|
||||
|
||||
This method returns a tuple of a string and a tuple of values. The string
|
||||
is the parametered WHERE clause predicate (including the WHERE keyword):
|
||||
"WHERE column1 = ? AND column2 IS ?". The tuple of values contains the
|
||||
values for each of the corresponding parameters.
|
||||
"""
|
||||
wheres = []
|
||||
param_values = []
|
||||
|
||||
# construct predicate for config filter
|
||||
if config is not None:
|
||||
wheres.append("thread_id = ?")
|
||||
param_values.append(config["configurable"]["thread_id"])
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns")
|
||||
if checkpoint_ns is not None:
|
||||
wheres.append("checkpoint_ns = ?")
|
||||
param_values.append(checkpoint_ns)
|
||||
|
||||
if checkpoint_id := get_checkpoint_id(config):
|
||||
wheres.append("checkpoint_id = ?")
|
||||
param_values.append(checkpoint_id)
|
||||
|
||||
# construct predicate for metadata filter
|
||||
if filter:
|
||||
metadata_predicates, metadata_values = _metadata_predicate(filter)
|
||||
wheres.extend(metadata_predicates)
|
||||
param_values.extend(metadata_values)
|
||||
|
||||
# construct predicate for `before`
|
||||
if before is not None:
|
||||
wheres.append("checkpoint_id < ?")
|
||||
param_values.append(get_checkpoint_id(before))
|
||||
|
||||
return ("WHERE " + " AND ".join(wheres) if wheres else "", param_values)
|
||||
@@ -1,4 +0,0 @@
|
||||
from langgraph.store.sqlite.aio import AsyncSqliteStore
|
||||
from langgraph.store.sqlite.base import SqliteStore
|
||||
|
||||
__all__ = ["AsyncSqliteStore", "SqliteStore"]
|
||||
@@ -1,623 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from collections.abc import AsyncIterator, Callable, Iterable, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from types import TracebackType
|
||||
from typing import Any, cast
|
||||
|
||||
import aiosqlite
|
||||
import orjson
|
||||
import sqlite_vec # type: ignore[import-untyped]
|
||||
from langgraph.store.base import (
|
||||
GetOp,
|
||||
ListNamespacesOp,
|
||||
Op,
|
||||
PutOp,
|
||||
Result,
|
||||
SearchOp,
|
||||
TTLConfig,
|
||||
)
|
||||
from langgraph.store.base.batch import AsyncBatchedBaseStore
|
||||
|
||||
from langgraph.store.sqlite.base import (
|
||||
_PLACEHOLDER,
|
||||
BaseSqliteStore,
|
||||
SqliteIndexConfig,
|
||||
_decode_ns_text,
|
||||
_ensure_index_config,
|
||||
_group_ops,
|
||||
_row_to_item,
|
||||
_row_to_search_item,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AsyncSqliteStore(AsyncBatchedBaseStore, BaseSqliteStore):
|
||||
"""Asynchronous SQLite-backed store with optional vector search.
|
||||
|
||||
This class provides an asynchronous interface for storing and retrieving data
|
||||
using a SQLite database with support for vector search capabilities.
|
||||
|
||||
Examples:
|
||||
Basic setup and usage:
|
||||
```python
|
||||
from langgraph.store.sqlite import AsyncSqliteStore
|
||||
|
||||
async with AsyncSqliteStore.from_conn_string(":memory:") as store:
|
||||
await store.setup() # Run migrations
|
||||
|
||||
# Store and retrieve data
|
||||
await store.aput(("users", "123"), "prefs", {"theme": "dark"})
|
||||
item = await store.aget(("users", "123"), "prefs")
|
||||
```
|
||||
|
||||
Vector search using LangChain embeddings:
|
||||
```python
|
||||
from langchain_openai import OpenAIEmbeddings
|
||||
from langgraph.store.sqlite import AsyncSqliteStore
|
||||
|
||||
async with AsyncSqliteStore.from_conn_string(
|
||||
":memory:",
|
||||
index={
|
||||
"dims": 1536,
|
||||
"embed": OpenAIEmbeddings(),
|
||||
"fields": ["text"] # specify which fields to embed
|
||||
}
|
||||
) as store:
|
||||
await store.setup() # Run migrations once
|
||||
|
||||
# Store documents
|
||||
await store.aput(("docs",), "doc1", {"text": "Python tutorial"})
|
||||
await store.aput(("docs",), "doc2", {"text": "TypeScript guide"})
|
||||
await store.aput(("docs",), "doc3", {"text": "Other guide"}, index=False) # don't index
|
||||
|
||||
# Search by similarity
|
||||
results = await store.asearch(("docs",), query="programming guides", limit=2)
|
||||
```
|
||||
|
||||
Warning:
|
||||
Make sure to call `setup()` before first use to create necessary tables and indexes.
|
||||
|
||||
Note:
|
||||
This class requires the aiosqlite package. Install with `pip install aiosqlite`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: aiosqlite.Connection,
|
||||
*,
|
||||
deserializer: Callable[[bytes | str | orjson.Fragment], dict[str, Any]]
|
||||
| None = None,
|
||||
index: SqliteIndexConfig | None = None,
|
||||
ttl: TTLConfig | None = None,
|
||||
):
|
||||
"""Initialize the async SQLite store.
|
||||
|
||||
Args:
|
||||
conn: The SQLite database connection.
|
||||
deserializer: Optional custom deserializer function for values.
|
||||
index: Optional vector search configuration.
|
||||
ttl: Optional time-to-live configuration.
|
||||
"""
|
||||
super().__init__()
|
||||
self._deserializer = deserializer
|
||||
self.conn = conn
|
||||
self.lock = asyncio.Lock()
|
||||
self.loop = asyncio.get_running_loop()
|
||||
self.is_setup = False
|
||||
self.index_config = index
|
||||
if self.index_config:
|
||||
self.embeddings, self.index_config = _ensure_index_config(self.index_config)
|
||||
else:
|
||||
self.embeddings = None
|
||||
self.ttl_config = ttl
|
||||
self._ttl_sweeper_task: asyncio.Task[None] | None = None
|
||||
self._ttl_stop_event = asyncio.Event()
|
||||
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
async def from_conn_string(
|
||||
cls,
|
||||
conn_string: str,
|
||||
*,
|
||||
index: SqliteIndexConfig | None = None,
|
||||
ttl: TTLConfig | None = None,
|
||||
) -> AsyncIterator[AsyncSqliteStore]:
|
||||
"""Create a new AsyncSqliteStore instance from a connection string.
|
||||
|
||||
Args:
|
||||
conn_string: The SQLite connection string.
|
||||
index: Optional vector search configuration.
|
||||
ttl: Optional time-to-live configuration.
|
||||
|
||||
Returns:
|
||||
An AsyncSqliteStore instance wrapped in an async context manager.
|
||||
"""
|
||||
async with aiosqlite.connect(conn_string, isolation_level=None) as conn:
|
||||
yield cls(conn, index=index, ttl=ttl)
|
||||
|
||||
async def setup(self) -> None:
|
||||
"""Set up the store database.
|
||||
|
||||
This method creates the necessary tables in the SQLite database if they don't
|
||||
already exist and runs database migrations. It should be called before first use.
|
||||
"""
|
||||
async with self.lock:
|
||||
if self.is_setup:
|
||||
return
|
||||
|
||||
# Create migrations table if it doesn't exist
|
||||
await self.conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS store_migrations (
|
||||
v INTEGER PRIMARY KEY
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
# Check current migration version
|
||||
async with self.conn.execute(
|
||||
"SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
|
||||
) as cur:
|
||||
row = await cur.fetchone()
|
||||
if row is None:
|
||||
version = -1
|
||||
else:
|
||||
version = row[0]
|
||||
|
||||
# Apply migrations
|
||||
for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1):
|
||||
await self.conn.executescript(sql)
|
||||
await self.conn.execute(
|
||||
"INSERT INTO store_migrations (v) VALUES (?)", (v,)
|
||||
)
|
||||
|
||||
# Apply vector migrations if index config is provided
|
||||
if self.index_config:
|
||||
# Create vector migrations table if it doesn't exist
|
||||
await self.conn.enable_load_extension(True)
|
||||
await self.conn.load_extension(sqlite_vec.loadable_path())
|
||||
await self.conn.enable_load_extension(False)
|
||||
await self.conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS vector_migrations (
|
||||
v INTEGER PRIMARY KEY
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
# Check current vector migration version
|
||||
async with self.conn.execute(
|
||||
"SELECT v FROM vector_migrations ORDER BY v DESC LIMIT 1"
|
||||
) as cur:
|
||||
row = await cur.fetchone()
|
||||
if row is None:
|
||||
version = -1
|
||||
else:
|
||||
version = row[0]
|
||||
|
||||
# Apply vector migrations
|
||||
for v, sql in enumerate(
|
||||
self.VECTOR_MIGRATIONS[version + 1 :], start=version + 1
|
||||
):
|
||||
await self.conn.executescript(sql)
|
||||
await self.conn.execute(
|
||||
"INSERT INTO vector_migrations (v) VALUES (?)", (v,)
|
||||
)
|
||||
|
||||
self.is_setup = True
|
||||
|
||||
@asynccontextmanager
|
||||
async def _cursor(
|
||||
self, *, transaction: bool = True
|
||||
) -> AsyncIterator[aiosqlite.Cursor]:
|
||||
"""Get a cursor for the SQLite database.
|
||||
|
||||
Args:
|
||||
transaction: Whether to use a transaction for database operations.
|
||||
|
||||
Yields:
|
||||
An SQLite cursor object.
|
||||
"""
|
||||
if not self.is_setup:
|
||||
await self.setup()
|
||||
async with self.lock:
|
||||
if transaction:
|
||||
await self.conn.execute("BEGIN")
|
||||
|
||||
async with self.conn.cursor() as cur:
|
||||
try:
|
||||
yield cur
|
||||
finally:
|
||||
if transaction:
|
||||
await self.conn.execute("COMMIT")
|
||||
|
||||
async def sweep_ttl(self) -> int:
|
||||
"""Delete expired store items based on TTL.
|
||||
|
||||
Returns:
|
||||
int: The number of deleted items.
|
||||
"""
|
||||
async with self._cursor() as cur:
|
||||
await cur.execute(
|
||||
"""
|
||||
DELETE FROM store
|
||||
WHERE expires_at IS NOT NULL AND expires_at < CURRENT_TIMESTAMP
|
||||
"""
|
||||
)
|
||||
deleted_count = cur.rowcount
|
||||
return deleted_count
|
||||
|
||||
async def start_ttl_sweeper(
|
||||
self, sweep_interval_minutes: int | None = None
|
||||
) -> asyncio.Task[None]:
|
||||
"""Periodically delete expired store items based on TTL.
|
||||
|
||||
Returns:
|
||||
Task that can be awaited or cancelled.
|
||||
"""
|
||||
if not self.ttl_config:
|
||||
return asyncio.create_task(asyncio.sleep(0))
|
||||
|
||||
if self._ttl_sweeper_task is not None and not self._ttl_sweeper_task.done():
|
||||
return self._ttl_sweeper_task
|
||||
|
||||
self._ttl_stop_event.clear()
|
||||
|
||||
interval = float(
|
||||
sweep_interval_minutes or self.ttl_config.get("sweep_interval_minutes") or 5
|
||||
)
|
||||
logger.info(f"Starting store TTL sweeper with interval {interval} minutes")
|
||||
|
||||
async def _sweep_loop() -> None:
|
||||
while not self._ttl_stop_event.is_set():
|
||||
try:
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._ttl_stop_event.wait(),
|
||||
timeout=interval * 60,
|
||||
)
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
expired_items = await self.sweep_ttl()
|
||||
if expired_items > 0:
|
||||
logger.info(f"Store swept {expired_items} expired items")
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as exc:
|
||||
logger.exception("Store TTL sweep iteration failed", exc_info=exc)
|
||||
|
||||
task = asyncio.create_task(_sweep_loop())
|
||||
task.set_name("ttl_sweeper")
|
||||
self._ttl_sweeper_task = task
|
||||
return task
|
||||
|
||||
async def stop_ttl_sweeper(self, timeout: float | None = None) -> bool:
|
||||
"""Stop the TTL sweeper task if it's running.
|
||||
|
||||
Args:
|
||||
timeout: Maximum time to wait for the task to stop, in seconds.
|
||||
If `None`, wait indefinitely.
|
||||
|
||||
Returns:
|
||||
bool: True if the task was successfully stopped or wasn't running,
|
||||
False if the timeout was reached before the task stopped.
|
||||
"""
|
||||
if self._ttl_sweeper_task is None or self._ttl_sweeper_task.done():
|
||||
return True
|
||||
|
||||
logger.info("Stopping TTL sweeper task")
|
||||
self._ttl_stop_event.set()
|
||||
|
||||
if timeout is not None:
|
||||
try:
|
||||
await asyncio.wait_for(self._ttl_sweeper_task, timeout=timeout)
|
||||
success = True
|
||||
except asyncio.TimeoutError:
|
||||
success = False
|
||||
else:
|
||||
await self._ttl_sweeper_task
|
||||
success = True
|
||||
|
||||
if success:
|
||||
self._ttl_sweeper_task = None
|
||||
logger.info("TTL sweeper task stopped")
|
||||
else:
|
||||
logger.warning("Timed out waiting for TTL sweeper task to stop")
|
||||
|
||||
return success
|
||||
|
||||
async def __aenter__(self) -> AsyncSqliteStore:
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
# Ensure the TTL sweeper task is stopped when exiting the context
|
||||
if hasattr(self, "_ttl_sweeper_task") and self._ttl_sweeper_task is not None:
|
||||
# Set the event to signal the task to stop
|
||||
self._ttl_stop_event.set()
|
||||
# We don't wait for the task to complete here to avoid blocking
|
||||
# The task will clean up itself gracefully
|
||||
|
||||
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
"""Execute a batch of operations asynchronously.
|
||||
|
||||
Args:
|
||||
ops: Iterable of operations to execute.
|
||||
|
||||
Returns:
|
||||
List of operation results.
|
||||
"""
|
||||
grouped_ops, num_ops = _group_ops(ops)
|
||||
results: list[Result] = [None] * num_ops
|
||||
|
||||
async with self._cursor(transaction=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
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
async def _batch_get_ops(
|
||||
self,
|
||||
get_ops: Sequence[tuple[int, GetOp]],
|
||||
results: list[Result],
|
||||
cur: aiosqlite.Cursor,
|
||||
) -> None:
|
||||
"""Process batch GET operations.
|
||||
|
||||
Args:
|
||||
get_ops: Sequence of GET operations.
|
||||
results: List to store results in.
|
||||
cur: Database cursor.
|
||||
"""
|
||||
# Group all queries by namespace to execute all operations for each namespace together
|
||||
namespace_queries = defaultdict(list)
|
||||
for prepared_query in self._get_batch_GET_ops_queries(get_ops):
|
||||
namespace_queries[prepared_query.namespace].append(prepared_query)
|
||||
|
||||
# Process each namespace's operations
|
||||
for namespace, queries in namespace_queries.items():
|
||||
# Execute TTL refresh queries first
|
||||
for query in queries:
|
||||
if query.kind == "refresh":
|
||||
try:
|
||||
await cur.execute(query.query, query.params)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Error executing TTL refresh: \n{query.query}\n{query.params}\n{e}"
|
||||
) from e
|
||||
|
||||
# Then execute GET queries and process results
|
||||
for query in queries:
|
||||
if query.kind == "get":
|
||||
try:
|
||||
await cur.execute(query.query, query.params)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Error executing GET query: \n{query.query}\n{query.params}\n{e}"
|
||||
) from e
|
||||
|
||||
rows = await cur.fetchall()
|
||||
key_to_row = {
|
||||
row[0]: {
|
||||
"key": row[0],
|
||||
"value": row[1],
|
||||
"created_at": row[2],
|
||||
"updated_at": row[3],
|
||||
"expires_at": row[4] if len(row) > 4 else None,
|
||||
"ttl_minutes": row[5] if len(row) > 5 else None,
|
||||
}
|
||||
for row in rows
|
||||
}
|
||||
|
||||
# Process results for this query
|
||||
for idx, key in query.items:
|
||||
row = key_to_row.get(key)
|
||||
if row:
|
||||
results[idx] = _row_to_item(
|
||||
namespace, row, loader=self._deserializer
|
||||
)
|
||||
else:
|
||||
results[idx] = None
|
||||
|
||||
async def _batch_put_ops(
|
||||
self,
|
||||
put_ops: Sequence[tuple[int, PutOp]],
|
||||
cur: aiosqlite.Cursor,
|
||||
) -> None:
|
||||
"""Process batch PUT operations.
|
||||
|
||||
Args:
|
||||
put_ops: Sequence of PUT operations.
|
||||
cur: Database cursor.
|
||||
"""
|
||||
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 = await self.embeddings.aembed_documents(
|
||||
[param[-1] for param in txt_params]
|
||||
)
|
||||
|
||||
# Convert vectors to SQLite-friendly format
|
||||
vector_params = []
|
||||
for (ns, k, pathname, _), vector in zip(txt_params, vectors, strict=False):
|
||||
vector_params.extend(
|
||||
[ns, k, pathname, sqlite_vec.serialize_float32(vector)]
|
||||
)
|
||||
|
||||
queries.append((query, vector_params))
|
||||
|
||||
for query, params in queries:
|
||||
await cur.execute(query, params)
|
||||
|
||||
async def _batch_search_ops(
|
||||
self,
|
||||
search_ops: Sequence[tuple[int, SearchOp]],
|
||||
results: list[Result],
|
||||
cur: aiosqlite.Cursor,
|
||||
) -> None:
|
||||
"""Process batch SEARCH operations.
|
||||
|
||||
Args:
|
||||
search_ops: Sequence of SEARCH operations.
|
||||
results: List to store results in.
|
||||
cur: Database cursor.
|
||||
"""
|
||||
prepared_queries, embedding_requests = self._prepare_batch_search_queries(
|
||||
search_ops
|
||||
)
|
||||
|
||||
# Setup dot_product function if it doesn't exist
|
||||
if embedding_requests and self.embeddings:
|
||||
vectors = await self.embeddings.aembed_documents(
|
||||
[query for _, query in embedding_requests]
|
||||
)
|
||||
|
||||
for (embed_req_idx, _), embedding in zip(
|
||||
embedding_requests, vectors, strict=False
|
||||
):
|
||||
# Find the corresponding query in prepared_queries
|
||||
# The embed_req_idx is the original index in search_ops, which should map to prepared_queries
|
||||
if embed_req_idx < len(prepared_queries):
|
||||
_params_list: list = prepared_queries[embed_req_idx][1]
|
||||
for i, param in enumerate(_params_list):
|
||||
if param is _PLACEHOLDER:
|
||||
_params_list[i] = sqlite_vec.serialize_float32(embedding)
|
||||
else:
|
||||
logger.warning(
|
||||
f"Embedding request index {embed_req_idx} out of bounds for prepared_queries."
|
||||
)
|
||||
|
||||
for (original_op_idx, _), (query, params, needs_refresh) in zip(
|
||||
search_ops, prepared_queries, strict=False
|
||||
):
|
||||
await cur.execute(query, params)
|
||||
rows = await cur.fetchall()
|
||||
|
||||
if needs_refresh and rows and self.ttl_config:
|
||||
keys_to_refresh = []
|
||||
for row_data in rows:
|
||||
# Assuming row_data[0] is prefix (text), row_data[1] is key (text)
|
||||
# These are raw text values directly from the DB.
|
||||
keys_to_refresh.append((row_data[0], row_data[1]))
|
||||
|
||||
if keys_to_refresh:
|
||||
updates_by_prefix = defaultdict(list)
|
||||
for prefix_text, key_text in keys_to_refresh:
|
||||
updates_by_prefix[prefix_text].append(key_text)
|
||||
|
||||
for prefix_text, key_list in updates_by_prefix.items():
|
||||
placeholders = ",".join(["?"] * len(key_list))
|
||||
update_query = f"""
|
||||
UPDATE store
|
||||
SET expires_at = DATETIME(CURRENT_TIMESTAMP, '+' || ttl_minutes || ' minutes')
|
||||
WHERE prefix = ? AND key IN ({placeholders}) AND ttl_minutes IS NOT NULL
|
||||
"""
|
||||
update_params = (prefix_text, *key_list)
|
||||
try:
|
||||
await cur.execute(update_query, update_params)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error during TTL refresh update for search: {e}"
|
||||
)
|
||||
|
||||
# Process rows into items
|
||||
if "score" in query: # Vector search query
|
||||
items = [
|
||||
_row_to_search_item(
|
||||
_decode_ns_text(row[0]), # prefix
|
||||
{
|
||||
"key": row[1], # key
|
||||
"value": row[2], # value
|
||||
"created_at": row[3],
|
||||
"updated_at": row[4],
|
||||
"expires_at": row[5] if len(row) > 5 else None,
|
||||
"ttl_minutes": row[6] if len(row) > 6 else None,
|
||||
"score": row[7] if len(row) > 7 else None,
|
||||
},
|
||||
loader=self._deserializer,
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
else: # Regular search query
|
||||
items = [
|
||||
_row_to_search_item(
|
||||
_decode_ns_text(row[0]), # prefix
|
||||
{
|
||||
"key": row[1], # key
|
||||
"value": row[2], # value
|
||||
"created_at": row[3],
|
||||
"updated_at": row[4],
|
||||
"expires_at": row[5] if len(row) > 5 else None,
|
||||
"ttl_minutes": row[6] if len(row) > 6 else None,
|
||||
},
|
||||
loader=self._deserializer,
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
results[original_op_idx] = items
|
||||
|
||||
async def _batch_list_namespaces_ops(
|
||||
self,
|
||||
list_ops: Sequence[tuple[int, ListNamespacesOp]],
|
||||
results: list[Result],
|
||||
cur: aiosqlite.Cursor,
|
||||
) -> None:
|
||||
"""Process batch LIST NAMESPACES operations.
|
||||
|
||||
Args:
|
||||
list_ops: Sequence of LIST NAMESPACES operations.
|
||||
results: List to store results in.
|
||||
cur: Database cursor.
|
||||
"""
|
||||
queries = self._get_batch_list_namespaces_queries(list_ops)
|
||||
for (query, params), (idx, _) in zip(queries, list_ops, strict=False):
|
||||
await cur.execute(query, params)
|
||||
|
||||
rows = await cur.fetchall()
|
||||
results[idx] = [_decode_ns_text(row[0]) for row in rows]
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user