From 8d2bf018715bc87b22b8094367c071b6ce48a200 Mon Sep 17 00:00:00 2001 From: StreetLamb Date: Sun, 11 Aug 2024 11:34:53 +0800 Subject: [PATCH 1/4] Upgrade langgraph version to v0.2.3. Add langgraph-checkpoint-postgres --- backend/poetry.lock | 183 +++++++++++++++++++++++------------------ backend/pyproject.toml | 3 +- 2 files changed, 103 insertions(+), 83 deletions(-) diff --git a/backend/poetry.lock b/backend/poetry.lock index f25c11b..1f2f49a 100644 --- a/backend/poetry.lock +++ b/backend/poetry.lock @@ -2290,17 +2290,48 @@ extended-testing = ["beautifulsoup4 (>=4.12.3,<5.0.0)", "lxml (>=4.9.3,<6.0)"] [[package]] name = "langgraph" -version = "0.1.9" +version = "0.2.3" description = "Building stateful, multi-actor applications with LLMs" optional = false python-versions = "<4.0,>=3.9.0" files = [ - {file = "langgraph-0.1.9-py3-none-any.whl", hash = "sha256:b3b5698686ae71fbf0cb2439f34d8a840f061c0e5ddc76d618674c3611ed787a"}, - {file = "langgraph-0.1.9.tar.gz", hash = "sha256:9ab6150d4b46089f8ea484fc68b1b28e0dd3adb7e383f0b8520ec04b7f6d5938"}, + {file = "langgraph-0.2.3-py3-none-any.whl", hash = "sha256:b725ec8627d79d7b44ac57e416a8128712cea609b6b67996576744a84ad26f9e"}, + {file = "langgraph-0.2.3.tar.gz", hash = "sha256:acc4ff58c572a689e83d62b37c1944a56336878ddb91268bf9daeeb3493fcb1c"}, ] [package.dependencies] -langchain-core = ">=0.2.19,<0.3" +langchain-core = ">=0.2.27,<0.3" +langgraph-checkpoint = ">=1.0.2,<2.0.0" + +[[package]] +name = "langgraph-checkpoint" +version = "1.0.2" +description = "Library with base interfaces for LangGraph checkpoint savers." +optional = false +python-versions = "<4.0.0,>=3.9.0" +files = [ + {file = "langgraph_checkpoint-1.0.2-py3-none-any.whl", hash = "sha256:c16cc3ee8b52f47799b4e9ad9981793a9cf582f77f6d03aed9d72a78dc618590"}, + {file = "langgraph_checkpoint-1.0.2.tar.gz", hash = "sha256:7f46b033888923ae4521ca2c9ccfcba6aacc7121888d77a6da1fd41ac2768d52"}, +] + +[package.dependencies] +langchain-core = ">=0.2.22,<0.3" + +[[package]] +name = "langgraph-checkpoint-postgres" +version = "1.0.3" +description = "Library with a Postgres implementation of LangGraph checkpoint saver." +optional = false +python-versions = "<4.0.0,>=3.9.0" +files = [ + {file = "langgraph_checkpoint_postgres-1.0.3-py3-none-any.whl", hash = "sha256:dd5b80c87099f9a051b1e14d1e93c5c0863561c67ec15adc58d25d803a407f8b"}, + {file = "langgraph_checkpoint_postgres-1.0.3.tar.gz", hash = "sha256:00803a3ce16fe3b7b33140b36e39ce88ad226a22f6faea013d5cabdc48957b7e"}, +] + +[package.dependencies] +langgraph-checkpoint = ">=1.0.1,<2.0.0" +orjson = ">=3.10.1" +psycopg = {version = ">=3.1.19", extras = ["binary"]} [[package]] name = "langsmith" @@ -3472,100 +3503,88 @@ files = [ [[package]] name = "psycopg" -version = "3.1.18" +version = "3.2.1" description = "PostgreSQL database adapter for Python" optional = false -python-versions = ">=3.7" +python-versions = ">=3.8" files = [ - {file = "psycopg-3.1.18-py3-none-any.whl", hash = "sha256:4d5a0a5a8590906daa58ebd5f3cfc34091377354a1acced269dd10faf55da60e"}, - {file = "psycopg-3.1.18.tar.gz", hash = "sha256:31144d3fb4c17d78094d9e579826f047d4af1da6a10427d91dfcfb6ecdf6f12b"}, + {file = "psycopg-3.2.1-py3-none-any.whl", hash = "sha256:ece385fb413a37db332f97c49208b36cf030ff02b199d7635ed2fbd378724175"}, + {file = "psycopg-3.2.1.tar.gz", hash = "sha256:dc8da6dc8729dacacda3cc2f17d2c9397a70a66cf0d2b69c91065d60d5f00cb7"}, ] [package.dependencies] -psycopg-binary = {version = "3.1.18", optional = true, markers = "implementation_name != \"pypy\" and extra == \"binary\""} -typing-extensions = ">=4.1" +psycopg-binary = {version = "3.2.1", optional = true, markers = "implementation_name != \"pypy\" and extra == \"binary\""} +typing-extensions = ">=4.4" tzdata = {version = "*", markers = "sys_platform == \"win32\""} [package.extras] -binary = ["psycopg-binary (==3.1.18)"] -c = ["psycopg-c (==3.1.18)"] -dev = ["black (>=24.1.0)", "codespell (>=2.2)", "dnspython (>=2.1)", "flake8 (>=4.0)", "mypy (>=1.4.1)", "types-setuptools (>=57.4)", "wheel (>=0.37)"] +binary = ["psycopg-binary (==3.2.1)"] +c = ["psycopg-c (==3.2.1)"] +dev = ["ast-comments (>=1.1.2)", "black (>=24.1.0)", "codespell (>=2.2)", "dnspython (>=2.1)", "flake8 (>=4.0)", "mypy (>=1.6)", "types-setuptools (>=57.4)", "wheel (>=0.37)"] docs = ["Sphinx (>=5.0)", "furo (==2022.6.21)", "sphinx-autobuild (>=2021.3.14)", "sphinx-autodoc-typehints (>=1.12)"] pool = ["psycopg-pool"] -test = ["anyio (>=3.6.2,<4.0)", "mypy (>=1.4.1)", "pproxy (>=2.7)", "pytest (>=6.2.5)", "pytest-cov (>=3.0)", "pytest-randomly (>=3.5)"] +test = ["anyio (>=4.0)", "mypy (>=1.6)", "pproxy (>=2.7)", "pytest (>=6.2.5)", "pytest-cov (>=3.0)", "pytest-randomly (>=3.5)"] [[package]] name = "psycopg-binary" -version = "3.1.18" +version = "3.2.1" description = "PostgreSQL database adapter for Python -- C optimisation distribution" optional = false -python-versions = ">=3.7" +python-versions = ">=3.8" files = [ - {file = "psycopg_binary-3.1.18-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:5c323103dfa663b88204cf5f028e83c77d7a715f9b6f51d2bbc8184b99ddd90a"}, - {file = "psycopg_binary-3.1.18-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:887f8d856c91510148be942c7acd702ccf761a05f59f8abc123c22ab77b5a16c"}, - {file = "psycopg_binary-3.1.18-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d322ba72cde4ca2eefc2196dad9ad7e52451acd2f04e3688d590290625d0c970"}, - {file = "psycopg_binary-3.1.18-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:489aa4fe5a0b653b68341e9e44af247dedbbc655326854aa34c163ef1bcb3143"}, - {file = "psycopg_binary-3.1.18-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:55ff0948457bfa8c0d35c46e3a75193906d1c275538877ba65907fd67aa059ad"}, - {file = "psycopg_binary-3.1.18-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b15e3653c82384b043d820fc637199b5c6a36b37fa4a4943e0652785bb2bad5d"}, - {file = "psycopg_binary-3.1.18-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:f8ff3bc08b43f36fdc24fedb86d42749298a458c4724fb588c4d76823ac39f54"}, - {file = "psycopg_binary-3.1.18-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:1729d0e3dfe2546d823841eb7a3d003144189d6f5e138ee63e5227f8b75276a5"}, - {file = "psycopg_binary-3.1.18-cp310-cp310-musllinux_1_1_ppc64le.whl", hash = "sha256:13bcd3742112446037d15e360b27a03af4b5afcf767f5ee374ef8f5dd7571b31"}, - {file = "psycopg_binary-3.1.18-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:320047e3d3554b857e16c2b6b615a85e0db6a02426f4d203a4594a2f125dfe57"}, - {file = "psycopg_binary-3.1.18-cp310-cp310-win_amd64.whl", hash = "sha256:888a72c2aca4316ca6d4a619291b805677bae99bba2f6e31a3c18424a48c7e4d"}, - {file = "psycopg_binary-3.1.18-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:4e4de16a637ec190cbee82e0c2dc4860fed17a23a35f7a1e6dc479a5c6876722"}, - {file = "psycopg_binary-3.1.18-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:6432047b8b24ef97e3fbee1d1593a0faaa9544c7a41a2c67d1f10e7621374c83"}, - {file = "psycopg_binary-3.1.18-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9d684227ef8212e27da5f2aff9d4d303cc30b27ac1702d4f6881935549486dd5"}, - {file = "psycopg_binary-3.1.18-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:67284e2e450dc7a9e4d76e78c0bd357dc946334a3d410defaeb2635607f632cd"}, - {file = "psycopg_binary-3.1.18-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1c9b6bd7fb5c6638cb32469674707649b526acfe786ba6d5a78ca4293d87bae4"}, - {file = "psycopg_binary-3.1.18-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7121acc783c4e86d2d320a7fb803460fab158a7f0a04c5e8c5d49065118c1e73"}, - {file = "psycopg_binary-3.1.18-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:e28ff8f3de7b56588c2a398dc135fd9f157d12c612bd3daa7e6ba9872337f6f5"}, - {file = "psycopg_binary-3.1.18-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:c84a0174109f329eeda169004c7b7ca2e884a6305acab4a39600be67f915ed38"}, - {file = "psycopg_binary-3.1.18-cp311-cp311-musllinux_1_1_ppc64le.whl", hash = "sha256:531381f6647fc267383dca88dbe8a70d0feff433a8e3d0c4939201fea7ae1b82"}, - {file = "psycopg_binary-3.1.18-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:b293e01057e63c3ac0002aa132a1071ce0fdb13b9ee2b6b45d3abdb3525c597d"}, - {file = "psycopg_binary-3.1.18-cp311-cp311-win_amd64.whl", hash = "sha256:780a90bcb69bf27a8b08bc35b958e974cb6ea7a04cdec69e737f66378a344d68"}, - {file = "psycopg_binary-3.1.18-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:87dd9154b757a5fbf6d590f6f6ea75f4ad7b764a813ae04b1d91a70713f414a1"}, - {file = "psycopg_binary-3.1.18-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f876ebbf92db70125f6375f91ab4bc6b27648aa68f90d661b1fc5affb4c9731c"}, - {file = "psycopg_binary-3.1.18-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:258d2f0cb45e4574f8b2fe7c6d0a0e2eb58903a4fd1fbaf60954fba82d595ab7"}, - {file = "psycopg_binary-3.1.18-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:bd27f713f2e5ef3fd6796e66c1a5203a27a30ecb847be27a78e1df8a9a5ae68c"}, - {file = "psycopg_binary-3.1.18-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:c38a4796abf7380f83b1653c2711cb2449dd0b2e5aca1caa75447d6fa5179c69"}, - {file = "psycopg_binary-3.1.18-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b2f7f95746efd1be2dc240248cc157f4315db3fd09fef2adfcc2a76e24aa5741"}, - {file = "psycopg_binary-3.1.18-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:4085f56a8d4fc8b455e8f44380705c7795be5317419aa5f8214f315e4205d804"}, - {file = "psycopg_binary-3.1.18-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:2e2484ae835dedc80cdc7f1b1a939377dc967fed862262cfd097aa9f50cade46"}, - {file = "psycopg_binary-3.1.18-cp312-cp312-musllinux_1_1_ppc64le.whl", hash = "sha256:3c2b039ae0c45eee4cd85300ef802c0f97d0afc78350946a5d0ec77dd2d7e834"}, - {file = "psycopg_binary-3.1.18-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8f54978c4b646dec77fefd8485fa82ec1a87807f334004372af1aaa6de9539a5"}, - {file = "psycopg_binary-3.1.18-cp312-cp312-win_amd64.whl", hash = "sha256:9ffcbbd389e486d3fd83d30107bbf8b27845a295051ccabde240f235d04ed921"}, - {file = "psycopg_binary-3.1.18-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:c76659ae29a84f2c14f56aad305dd00eb685bd88f8c0a3281a9a4bc6bd7d2aa7"}, - {file = "psycopg_binary-3.1.18-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3c7afcd6f1d55992f26d9ff7b0bd4ee6b475eb43aa3f054d67d32e09f18b0065"}, - {file = "psycopg_binary-3.1.18-cp37-cp37m-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:639dd78ac09b144b0119076783cb64e1128cc8612243e9701d1503c816750b2e"}, - {file = "psycopg_binary-3.1.18-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e1cf59e0bb12e031a48bb628aae32df3d0c98fd6c759cb89f464b1047f0ca9c8"}, - {file = "psycopg_binary-3.1.18-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e262398e5d51563093edf30612cd1e20fedd932ad0994697d7781ca4880cdc3d"}, - {file = "psycopg_binary-3.1.18-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:59701118c7d8842e451f1e562d08e8708b3f5d14974eefbce9374badd723c4ae"}, - {file = "psycopg_binary-3.1.18-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:dea4a59da7850192fdead9da888e6b96166e90608cf39e17b503f45826b16f84"}, - {file = "psycopg_binary-3.1.18-cp37-cp37m-musllinux_1_1_ppc64le.whl", hash = "sha256:4575da95fc441244a0e2ebaf33a2b2f74164603341d2046b5cde0a9aa86aa7e2"}, - {file = "psycopg_binary-3.1.18-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:812726266ab96de681f2c7dbd6b734d327f493a78357fcc16b2ac86ff4f4e080"}, - {file = "psycopg_binary-3.1.18-cp37-cp37m-win_amd64.whl", hash = "sha256:3e7ce4d988112ca6c75765c7f24c83bdc476a6a5ce00878df6c140ca32c3e16d"}, - {file = "psycopg_binary-3.1.18-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:02bd4da45d5ee9941432e2e9bf36fa71a3ac21c6536fe7366d1bd3dd70d6b1e7"}, - {file = "psycopg_binary-3.1.18-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:39242546383f6b97032de7af30edb483d237a0616f6050512eee7b218a2aa8ee"}, - {file = "psycopg_binary-3.1.18-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d46ae44d66bf6058a812467f6ae84e4e157dee281bfb1cfaeca07dee07452e85"}, - {file = "psycopg_binary-3.1.18-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ad35ac7fd989184bf4d38a87decfb5a262b419e8ba8dcaeec97848817412c64a"}, - {file = "psycopg_binary-3.1.18-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:247474af262bdd5559ee6e669926c4f23e9cf53dae2d34c4d991723c72196404"}, - {file = "psycopg_binary-3.1.18-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6ebecbf2406cd6875bdd2453e31067d1bd8efe96705a9489ef37e93b50dc6f09"}, - {file = "psycopg_binary-3.1.18-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:1859aeb2133f5ecdd9cbcee155f5e38699afc06a365f903b1512c765fd8d457e"}, - {file = "psycopg_binary-3.1.18-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:da917f6df8c6b2002043193cb0d74cc173b3af7eb5800ad69c4e1fbac2a71c30"}, - {file = "psycopg_binary-3.1.18-cp38-cp38-musllinux_1_1_ppc64le.whl", hash = "sha256:9e24e7b6a68a51cc3b162d0339ae4e1263b253e887987d5c759652f5692b5efe"}, - {file = "psycopg_binary-3.1.18-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:e252d66276c992319ed6cd69a3ffa17538943954075051e992143ccbf6dc3d3e"}, - {file = "psycopg_binary-3.1.18-cp38-cp38-win_amd64.whl", hash = "sha256:5d6e860edf877d4413e4a807e837d55e3a7c7df701e9d6943c06e460fa6c058f"}, - {file = "psycopg_binary-3.1.18-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:eea5f14933177ffe5c40b200f04f814258cc14b14a71024ad109f308e8bad414"}, - {file = "psycopg_binary-3.1.18-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:824a1bfd0db96cc6bef2d1e52d9e0963f5bf653dd5bc3ab519a38f5e6f21c299"}, - {file = "psycopg_binary-3.1.18-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a87e9eeb80ce8ec8c2783f29bce9a50bbcd2e2342a340f159c3326bf4697afa1"}, - {file = "psycopg_binary-3.1.18-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:91074f78a9f890af5f2c786691575b6b93a4967ad6b8c5a90101f7b8c1a91d9c"}, - {file = "psycopg_binary-3.1.18-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e05f6825f8db4428782135e6986fec79b139210398f3710ed4aa6ef41473c008"}, - {file = "psycopg_binary-3.1.18-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0f68ac2364a50d4cf9bb803b4341e83678668f1881a253e1224574921c69868c"}, - {file = "psycopg_binary-3.1.18-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:7ac1785d67241d5074f8086705fa68e046becea27964267ab3abd392481d7773"}, - {file = "psycopg_binary-3.1.18-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:cd2a9f7f0d4dacc5b9ce7f0e767ae6cc64153264151f50698898c42cabffec0c"}, - {file = "psycopg_binary-3.1.18-cp39-cp39-musllinux_1_1_ppc64le.whl", hash = "sha256:3e4b0bb91da6f2238dbd4fbb4afc40dfb4f045bb611b92fce4d381b26413c686"}, - {file = "psycopg_binary-3.1.18-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:74e498586b72fb819ca8ea82107747d0cb6e00ae685ea6d1ab3f929318a8ce2d"}, - {file = "psycopg_binary-3.1.18-cp39-cp39-win_amd64.whl", hash = "sha256:d4422af5232699f14b7266a754da49dc9bcd45eba244cf3812307934cd5d6679"}, + {file = "psycopg_binary-3.2.1-cp310-cp310-macosx_12_0_x86_64.whl", hash = "sha256:cad2de17804c4cfee8640ae2b279d616bb9e4734ac3c17c13db5e40982bd710d"}, + {file = "psycopg_binary-3.2.1-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:592b27d6c46a40f9eeaaeea7c1fef6f3c60b02c634365eb649b2d880669f149f"}, + {file = "psycopg_binary-3.2.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9a997efbaadb5e1a294fb5760e2f5643d7b8e4e3fe6cb6f09e6d605fd28e0291"}, + {file = "psycopg_binary-3.2.1-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c1d2b6438fb83376f43ebb798bf0ad5e57bc56c03c9c29c85bc15405c8c0ac5a"}, + {file = "psycopg_binary-3.2.1-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b1f087bd84bdcac78bf9f024ebdbfacd07fc0a23ec8191448a50679e2ac4a19e"}, + {file = "psycopg_binary-3.2.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:415c3b72ea32119163255c6504085f374e47ae7345f14bc3f0ef1f6e0976a879"}, + {file = "psycopg_binary-3.2.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:f092114f10f81fb6bae544a0ec027eb720e2d9c74a4fcdaa9dd3899873136935"}, + {file = "psycopg_binary-3.2.1-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:06a7aae34edfe179ddc04da005e083ff6c6b0020000399a2cbf0a7121a8a22ea"}, + {file = "psycopg_binary-3.2.1-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:0b018631e5c80ce9bc210b71ea885932f9cca6db131e4df505653d7e3873a938"}, + {file = "psycopg_binary-3.2.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:f8a509aeaac364fa965454e80cd110fe6d48ba2c80f56c9b8563423f0b5c3cfd"}, + {file = "psycopg_binary-3.2.1-cp310-cp310-win_amd64.whl", hash = "sha256:413977d18412ff83486eeb5875eb00b185a9391c57febac45b8993bf9c0ff489"}, + {file = "psycopg_binary-3.2.1-cp311-cp311-macosx_12_0_x86_64.whl", hash = "sha256:62b1b7b07e00ee490afb39c0a47d8282a9c2822c7cfed9553a04b0058adf7e7f"}, + {file = "psycopg_binary-3.2.1-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:f8afb07114ea9b924a4a0305ceb15354ccf0ef3c0e14d54b8dbeb03e50182dd7"}, + {file = "psycopg_binary-3.2.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:40bb515d042f6a345714ec0403df68ccf13f73b05e567837d80c886c7c9d3805"}, + {file = "psycopg_binary-3.2.1-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6418712ba63cebb0c88c050b3997185b0ef54173b36568522d5634ac06153040"}, + {file = "psycopg_binary-3.2.1-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:101472468d59c74bb8565fab603e032803fd533d16be4b2d13da1bab8deb32a3"}, + {file = "psycopg_binary-3.2.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:aa3931f308ab4a479d0ee22dc04bea867a6365cac0172e5ddcba359da043854b"}, + {file = "psycopg_binary-3.2.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:dc314a47d44fe1a8069b075a64abffad347a3a1d8652fed1bab5d3baea37acb2"}, + {file = "psycopg_binary-3.2.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:cc304a46be1e291031148d9d95c12451ffe783ff0cc72f18e2cc7ec43cdb8c68"}, + {file = "psycopg_binary-3.2.1-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:6f9e13600647087df5928875559f0eb8f496f53e6278b7da9511b4b3d0aff960"}, + {file = "psycopg_binary-3.2.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:b140182830c76c74d17eba27df3755a46442ce8d4fb299e7f1cf2f74a87c877b"}, + {file = "psycopg_binary-3.2.1-cp311-cp311-win_amd64.whl", hash = "sha256:3c838806eeb99af39f934b7999e35f947a8e577997cc892c12b5053a97a9057f"}, + {file = "psycopg_binary-3.2.1-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:7066d3dca196ed0dc6172f9777b2d62e4f138705886be656cccff2d555234d60"}, + {file = "psycopg_binary-3.2.1-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:28ada5f610468c57d8a4a055a8ea915d0085a43d794266c4f3b9d02f4288f4db"}, + {file = "psycopg_binary-3.2.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2e8213bf50af073b1aa8dc3cff123bfeedac86332a16c1b7274910bc88a847c7"}, + {file = "psycopg_binary-3.2.1-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:74d623261655a169bc84a9669890975c229f2fa6e19a7f2d10a77675dcf1a707"}, + {file = "psycopg_binary-3.2.1-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:42781ba94e8842ee98bca5a7d0c44cc9d067500fedca2d6a90fa3609b6d16b42"}, + {file = "psycopg_binary-3.2.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:33e6669091d09f8ba36e10ce678a6d9916e110446236a9b92346464a3565635e"}, + {file = "psycopg_binary-3.2.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:b09e8a576a2ac69d695032ee76f31e03b30781828b5dd6d18c6a009e5a3d1c35"}, + {file = "psycopg_binary-3.2.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:8f28ff0cb9f1defdc4a6f8c958bf6787274247e7dfeca811f6e2f56602695fb1"}, + {file = "psycopg_binary-3.2.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:4c84fcac8a3a3479ac14673095cc4e1fdba2935499f72c436785ac679bec0d1a"}, + {file = "psycopg_binary-3.2.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:950fd666ec9e9fe6a8eeb2b5a8f17301790e518953730ad44d715b59ffdbc67f"}, + {file = "psycopg_binary-3.2.1-cp312-cp312-win_amd64.whl", hash = "sha256:334046a937bb086c36e2c6889fe327f9f29bfc085d678f70fac0b0618949f674"}, + {file = "psycopg_binary-3.2.1-cp38-cp38-macosx_12_0_x86_64.whl", hash = "sha256:1d6833f607f3fc7b22226a9e121235d3b84c0eda1d3caab174673ef698f63788"}, + {file = "psycopg_binary-3.2.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1d353e028b8f848b9784450fc2abf149d53a738d451eab3ee4c85703438128b9"}, + {file = "psycopg_binary-3.2.1-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f34e369891f77d0738e5d25727c307d06d5344948771e5379ea29c76c6d84555"}, + {file = "psycopg_binary-3.2.1-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0ab58213cc976a1666f66bc1cb2e602315cd753b7981a8e17237ac2a185bd4a1"}, + {file = "psycopg_binary-3.2.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b0104a72a17aa84b3b7dcab6c84826c595355bf54bb6ea6d284dcb06d99c6801"}, + {file = "psycopg_binary-3.2.1-cp38-cp38-musllinux_1_2_aarch64.whl", hash = "sha256:059cbd4e6da2337e17707178fe49464ed01de867dc86c677b30751755ec1dc51"}, + {file = "psycopg_binary-3.2.1-cp38-cp38-musllinux_1_2_i686.whl", hash = "sha256:73f9c9b984be9c322b5ec1515b12df1ee5896029f5e72d46160eb6517438659c"}, + {file = "psycopg_binary-3.2.1-cp38-cp38-musllinux_1_2_ppc64le.whl", hash = "sha256:af0469c00f24c4bec18c3d2ede124bf62688d88d1b8a5f3c3edc2f61046fe0d7"}, + {file = "psycopg_binary-3.2.1-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:463d55345f73ff391df8177a185ad57b552915ad33f5cc2b31b930500c068b22"}, + {file = "psycopg_binary-3.2.1-cp38-cp38-win_amd64.whl", hash = "sha256:302b86f92c0d76e99fe1b5c22c492ae519ce8b98b88d37ef74fda4c9e24c6b46"}, + {file = "psycopg_binary-3.2.1-cp39-cp39-macosx_12_0_x86_64.whl", hash = "sha256:0879b5d76b7d48678d31278242aaf951bc2d69ca4e4d7cef117e4bbf7bfefda9"}, + {file = "psycopg_binary-3.2.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f99e59f8a5f4dcd9cbdec445f3d8ac950a492fc0e211032384d6992ed3c17eb7"}, + {file = "psycopg_binary-3.2.1-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:84837e99353d16c6980603b362d0f03302d4b06c71672a6651f38df8a482923d"}, + {file = "psycopg_binary-3.2.1-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7ce965caf618061817f66c0906f0452aef966c293ae0933d4fa5a16ea6eaf5bb"}, + {file = "psycopg_binary-3.2.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:78c2007caf3c90f08685c5378e3ceb142bafd5636be7495f7d86ec8a977eaeef"}, + {file = "psycopg_binary-3.2.1-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:7a84b5eb194a258116154b2a4ff2962ea60ea52de089508db23a51d3d6b1c7d1"}, + {file = "psycopg_binary-3.2.1-cp39-cp39-musllinux_1_2_i686.whl", hash = "sha256:4a42b8f9ab39affcd5249b45cac763ac3cf12df962b67e23fd15a2ee2932afe5"}, + {file = "psycopg_binary-3.2.1-cp39-cp39-musllinux_1_2_ppc64le.whl", hash = "sha256:788ffc43d7517c13e624c83e0e553b7b8823c9655e18296566d36a829bfb373f"}, + {file = "psycopg_binary-3.2.1-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:21927f41c4d722ae8eb30d62a6ce732c398eac230509af5ba1749a337f8a63e2"}, + {file = "psycopg_binary-3.2.1-cp39-cp39-win_amd64.whl", hash = "sha256:921f0c7f39590763d64a619de84d1b142587acc70fd11cbb5ba8fa39786f3073"}, ] [[package]] @@ -5314,4 +5333,4 @@ repair = ["scipy (>=1.6.3)"] [metadata] lock-version = "2.0" python-versions = ">=3.10,<3.13" -content-hash = "7335e20fadc21f74bfb3d9fa6462eba11fcbcc33fe73f2d176c58de85ea5f8ab" +content-hash = "292c7427979bea5552753d8ac04c8f2f25e13ce8a802b06acd51d8f9277cb209" diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 788a1e5..2d99499 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -25,7 +25,7 @@ sqlmodel = "0.0.21" bcrypt = "4.0.1" pydantic-settings = "^2.2.1" sentry-sdk = {extras = ["fastapi"], version = "^2.8.0"} -langgraph = "0.1.9" +langgraph = "0.2.3" langchain-openai = "0.1.17" grandalf = "^0.8" langchain = "0.2.12" @@ -53,6 +53,7 @@ celery-stubs = "^0.1.3" pymupdf = "^1.24.7" psycopg-pool = "^3.2.2" langchain-ollama = "0.1.1" +langgraph-checkpoint-postgres = "^1.0.3" [tool.poetry.group.dev.dependencies] pytest = "^7.4.3" From 51976b7b0564c8522d299aa9d09f3ac8d5458923 Mon Sep 17 00:00:00 2001 From: StreetLamb Date: Tue, 13 Aug 2024 23:07:17 +0800 Subject: [PATCH 2/4] Update models for new checkpointer --- backend/app/models.py | 60 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 50 insertions(+), 10 deletions(-) diff --git a/backend/app/models.py b/backend/app/models.py index cc0a5e6..14529c1 100644 --- a/backend/app/models.py +++ b/backend/app/models.py @@ -10,12 +10,14 @@ Column, DateTime, PrimaryKeyConstraint, + String, UniqueConstraint, func, ) from sqlalchemy import ( Enum as SQLEnum, ) +from sqlalchemy.dialects.postgresql import JSONB from sqlmodel import Field, Relationship, SQLModel from app.core.graph.messages import ChatResponse @@ -361,12 +363,23 @@ class ToolDefinitionValidate(SQLModel): class Checkpoint(SQLModel, table=True): __tablename__ = "checkpoints" - __table_args__ = (PrimaryKeyConstraint("thread_id", "thread_ts"),) + __table_args__ = ( + PrimaryKeyConstraint("thread_id", "checkpoint_id", "checkpoint_ns"), + ) thread_id: UUID = Field(foreign_key="thread.id", primary_key=True) - thread_ts: UUID = Field(primary_key=True) - parent_ts: UUID | None - checkpoint: bytes - metadata_: bytes = Field(sa_column_kwargs={"name": "metadata"}) + checkpoint_ns: str = Field( + sa_column=Column( + "checkpoint_ns", String, nullable=False, server_default="", primary_key=True + ), + ) + checkpoint_id: UUID = Field(primary_key=True) + parent_checkpoint_id: UUID | None + type: str | None + checkpoint: dict[Any, Any] = Field(default_factory=dict, sa_column=Column(JSONB)) + metadata_: dict[Any, Any] = Field( + default_factory=dict, + sa_column=Column("metadata", JSONB, nullable=False, server_default="{}"), + ) thread: Thread = Relationship(back_populates="checkpoints") created_at: datetime | None = Field( sa_column=Column( @@ -378,22 +391,49 @@ class Checkpoint(SQLModel, table=True): ) +class CheckpointBlobs(SQLModel, table=True): + __tablename__ = "checkpoint_blobs" + __table_args__ = ( + PrimaryKeyConstraint("thread_id", "checkpoint_ns", "channel", "version"), + ) + thread_id: UUID = Field(foreign_key="thread.id", primary_key=True) + checkpoint_ns: str = Field( + sa_column=Column( + "checkpoint_ns", String, nullable=False, server_default="", primary_key=True + ), + ) + channel: str = Field(primary_key=True) + version: str = Field(primary_key=True) + type: str + blob: bytes | None + + class CheckpointOut(SQLModel): thread_id: UUID - thread_ts: UUID + checkpoint_id: UUID checkpoint: bytes created_at: datetime class Write(SQLModel, table=True): - __tablename__ = "writes" - __table_args__ = (PrimaryKeyConstraint("thread_id", "thread_ts", "task_id", "idx"),) + __tablename__ = "checkpoint_writes" + __table_args__ = ( + PrimaryKeyConstraint( + "thread_id", "checkpoint_ns", "checkpoint_id", "task_id", "idx" + ), + ) thread_id: UUID = Field(foreign_key="thread.id", primary_key=True) - thread_ts: UUID = Field(primary_key=True) + checkpoint_ns: str = Field( + sa_column=Column( + "checkpoint_ns", String, nullable=False, server_default="", primary_key=True + ), + ) + checkpoint_id: UUID = Field(primary_key=True) task_id: UUID = Field(primary_key=True) idx: int = Field(primary_key=True) channel: str - value: bytes + type: str | None + blob: bytes thread: Thread = Relationship(back_populates="writes") From 96500b4b129645f1771fc6122017ff7d5edccfc1 Mon Sep 17 00:00:00 2001 From: StreetLamb Date: Tue, 13 Aug 2024 23:08:02 +0800 Subject: [PATCH 3/4] Create migration file for new checkpointer --- .../20f584dc80d2_upgrade_checkpointer.py | 105 ++++++++++++++++++ 1 file changed, 105 insertions(+) create mode 100644 backend/app/alembic/versions/20f584dc80d2_upgrade_checkpointer.py diff --git a/backend/app/alembic/versions/20f584dc80d2_upgrade_checkpointer.py b/backend/app/alembic/versions/20f584dc80d2_upgrade_checkpointer.py new file mode 100644 index 0000000..a2afcb9 --- /dev/null +++ b/backend/app/alembic/versions/20f584dc80d2_upgrade_checkpointer.py @@ -0,0 +1,105 @@ +from alembic import op +import sqlalchemy as sa +import sqlmodel.sql.sqltypes +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = '20f584dc80d2' +down_revision = '38a9c73bfce2' +branch_labels = None +depends_on = None + +def upgrade(): + # Create new tables + op.create_table('checkpoint_blobs', + sa.Column('thread_id', sa.Uuid(), nullable=False), + sa.Column('checkpoint_ns', sa.String(), nullable=False, server_default=''), + sa.Column('channel', sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column('version', sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column('type', sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column('blob', sa.LargeBinary(), nullable=True), + sa.ForeignKeyConstraint(['thread_id'], ['thread.id']), + sa.PrimaryKeyConstraint('thread_id', 'checkpoint_ns', 'channel', 'version') + ) + op.create_table('checkpoint_writes', + sa.Column('thread_id', sa.Uuid(), nullable=False), + sa.Column('checkpoint_ns', sa.String(), nullable=False, server_default=''), + sa.Column('checkpoint_id', sa.Uuid(), nullable=False), + sa.Column('task_id', sa.Uuid(), nullable=False), + sa.Column('idx', sa.Integer(), nullable=False), + sa.Column('channel', sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column('type', sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.Column('blob', sa.LargeBinary(), nullable=False), + sa.ForeignKeyConstraint(['thread_id'], ['thread.id']), + sa.PrimaryKeyConstraint('thread_id', 'checkpoint_ns', 'checkpoint_id', 'task_id', 'idx') + ) + + # Drop the old table + op.drop_table('writes') + + # Rename and recreate the checkpoints table + op.rename_table('checkpoints', 'checkpoints_old') + + op.create_table( + 'checkpoints', + sa.Column('thread_id', sa.Uuid(), nullable=False), + sa.Column('checkpoint_ns', sa.String(), nullable=False, server_default=''), + sa.Column('checkpoint_id', sa.Uuid(), nullable=False), + sa.Column('parent_checkpoint_id', sa.Uuid(), nullable=True), + sa.Column('type', sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.Column('checkpoint', postgresql.JSONB(astext_type=sa.Text()), nullable=True), + sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default='{}'), + sa.Column('created_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()), + sa.ForeignKeyConstraint(['thread_id'], ['thread.id']), + sa.PrimaryKeyConstraint('thread_id', 'checkpoint_ns', 'checkpoint_id') + ) + + # Drop the old checkpoints table + op.drop_table('checkpoints_old') + + # Clear the threads table + op.execute('DELETE FROM thread') + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + + # Recreate the old checkpoints table + op.create_table( + 'checkpoints_old', + sa.Column('thread_id', sa.Uuid(), nullable=False), + sa.Column('thread_ts', sa.Uuid(), nullable=False), + sa.Column('parent_ts', sa.Uuid(), nullable=True), + sa.Column('checkpoint', sa.LargeBinary(), nullable=True), + sa.Column('metadata', sa.LargeBinary(), nullable=False, server_default=sa.text("'\\x'::bytea")), + sa.Column('created_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()), + sa.ForeignKeyConstraint(['thread_id'], ['thread.id']), + sa.PrimaryKeyConstraint('thread_id', 'thread_ts') + ) + + # Drop the new checkpoints table + op.drop_table('checkpoints') + + # Rename the old table back to 'checkpoints' + op.rename_table('checkpoints_old', 'checkpoints') + + # Recreate the old 'writes' table + op.create_table('writes', + sa.Column('thread_id', sa.UUID(), autoincrement=False, nullable=False), + sa.Column('thread_ts', sa.UUID(), autoincrement=False, nullable=False), + sa.Column('task_id', sa.UUID(), autoincrement=False, nullable=False), + sa.Column('idx', sa.INTEGER(), autoincrement=False, nullable=False), + sa.Column('channel', sa.VARCHAR(), autoincrement=False, nullable=False), + sa.Column('value', postgresql.BYTEA(), autoincrement=False, nullable=False), + sa.ForeignKeyConstraint(['thread_id'], ['thread.id'], name='writes_thread_id_fkey'), + sa.PrimaryKeyConstraint('thread_id', 'thread_ts', 'task_id', 'idx', name='writes_pkey') + ) + + # Drop the new tables + op.drop_table('checkpoint_writes') + op.drop_table('checkpoint_blobs') + + # Clear the threads table + op.execute('DELETE FROM thread') + # ### end Alembic commands ### \ No newline at end of file From 3a6a42977f4a945aed6b5f028fca74b324a60130 Mon Sep 17 00:00:00 2001 From: StreetLamb Date: Tue, 13 Aug 2024 23:09:25 +0800 Subject: [PATCH 4/4] Refactor code to use new AsyncPostgresSaver. Delete old PostgresSaver class. --- backend/app/core/config.py | 8 + backend/app/core/graph/build.py | 11 +- backend/app/core/graph/checkpoint/postgres.py | 578 ------------------ backend/app/core/graph/checkpoint/utils.py | 8 +- 4 files changed, 20 insertions(+), 585 deletions(-) delete mode 100644 backend/app/core/graph/checkpoint/postgres.py diff --git a/backend/app/core/config.py b/backend/app/core/config.py index 65583c5..7377590 100644 --- a/backend/app/core/config.py +++ b/backend/app/core/config.py @@ -2,6 +2,7 @@ import warnings from typing import Annotated, Any, Literal +from psycopg.rows import dict_row from pydantic import ( AnyUrl, BeforeValidator, @@ -66,6 +67,13 @@ def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn: path=self.POSTGRES_DB, ) + # For checkpointer + SQLALCHEMY_CONNECTION_KWARGS: dict[str, Any] = { + "autocommit": True, + "prepare_threshold": 0, + "row_factory": dict_row, + } + @computed_field # type: ignore[misc] @property def PG_DATABASE_URI(self) -> str: diff --git a/backend/app/core/graph/build.py b/backend/app/core/graph/build.py index e8e18a9..024aead 100644 --- a/backend/app/core/graph/build.py +++ b/backend/app/core/graph/build.py @@ -13,7 +13,8 @@ ) from langchain_core.runnables import RunnableLambda from langchain_core.runnables.config import RunnableConfig -from langgraph.checkpoint import BaseCheckpointSaver +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver from langgraph.graph import END, StateGraph from langgraph.graph.graph import CompiledGraph from langgraph.prebuilt import ( @@ -22,7 +23,6 @@ from psycopg import AsyncConnection from app.core.config import settings -from app.core.graph.checkpoint.postgres import PostgresSaver from app.core.graph.members import ( GraphLeader, GraphMember, @@ -471,8 +471,11 @@ async def generator( ] try: - async with await AsyncConnection.connect(settings.PG_DATABASE_URI) as conn: - checkpointer = PostgresSaver(async_connection=conn) + async with await AsyncConnection.connect( + settings.PG_DATABASE_URI, + **settings.SQLALCHEMY_CONNECTION_KWARGS, + ) as conn: + checkpointer = AsyncPostgresSaver(conn=conn) if team.workflow == "hierarchical": teams = convert_hierarchical_team_to_dict(team, members) team_leader = list(teams.keys())[0] diff --git a/backend/app/core/graph/checkpoint/postgres.py b/backend/app/core/graph/checkpoint/postgres.py deleted file mode 100644 index c8c1ff3..0000000 --- a/backend/app/core/graph/checkpoint/postgres.py +++ /dev/null @@ -1,578 +0,0 @@ -"""Implementation of a langgraph checkpoint saver using Postgres.""" -from collections.abc import AsyncGenerator, AsyncIterator, Generator, Sequence -from contextlib import asynccontextmanager, contextmanager -from typing import Any, List # noqa: UP035 - -import psycopg -from langchain_core.runnables import RunnableConfig -from langgraph.checkpoint import BaseCheckpointSaver -from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, CheckpointTuple -from langgraph.serde.jsonplus import JsonPlusSerializer -from psycopg_pool import AsyncConnectionPool, ConnectionPool - - -class JsonAndBinarySerializer(JsonPlusSerializer): - def _default(self, obj: Any) -> Any: - if isinstance(obj, bytes | bytearray): - return self._encode_constructor_args( - obj.__class__, method="fromhex", args=[obj.hex()] - ) - return super()._default(obj) # type: ignore[no-untyped-call] - - def dumps(self, obj: Any) -> tuple[str, bytes]: # type: ignore[override] - if isinstance(obj, bytes): - return "bytes", obj - elif isinstance(obj, bytearray): - return "bytearray", obj - - return "json", super().dumps(obj) - - def loads(self, s: tuple[str, bytes]) -> Any: # type: ignore[override] - if s[0] == "bytes": - return s[1] - elif s[0] == "bytearray": - return bytearray(s[1]) - elif s[0] == "json": - return super().loads(s[1]) - else: - raise NotImplementedError(f"Unknown serialization type: {s[0]}") - - -@contextmanager -def _get_sync_connection( - connection: psycopg.Connection | ConnectionPool | None, # type: ignore[type-arg] -) -> Generator[psycopg.Connection, None, None]: # type: ignore[type-arg] - """Get the connection to the Postgres database.""" - if isinstance(connection, psycopg.Connection): - yield connection - elif isinstance(connection, ConnectionPool): - with connection.connection() as conn: - yield conn - else: - raise ValueError( - "Invalid sync connection object. Please initialize the check pointer " - f"with an appropriate sync connection object. " - f"Got {type(connection)}." - ) - - -@asynccontextmanager -async def _get_async_connection( - connection: psycopg.AsyncConnection | AsyncConnectionPool | None, # type: ignore[type-arg] -) -> AsyncGenerator[psycopg.AsyncConnection, None]: # type: ignore[type-arg] - """Get the connection to the Postgres database.""" - if isinstance(connection, psycopg.AsyncConnection): - yield connection - elif isinstance(connection, AsyncConnectionPool): - async with connection.connection() as conn: - yield conn - else: - raise ValueError( - "Invalid async connection object. Please initialize the check pointer " - f"with an appropriate async connection object. " - f"Got {type(connection)}." - ) - - -class PostgresSaver(BaseCheckpointSaver): - sync_connection: psycopg.Connection | ConnectionPool | None = None # type: ignore[type-arg] - """The synchronous connection or pool to the Postgres database. - - If providing a connection object, please ensure that the connection is open - and remember to close the connection when done. - """ - async_connection: psycopg.AsyncConnection | AsyncConnectionPool | None = None # type: ignore[type-arg] - """The asynchronous connection or pool to the Postgres database. - - If providing a connection object, please ensure that the connection is open - and remember to close the connection when done. - """ - - def __init__( - self, - sync_connection: psycopg.Connection | ConnectionPool | None = None, # type: ignore[type-arg] - async_connection: psycopg.AsyncConnection | AsyncConnectionPool | None = None, # type: ignore[type-arg] - ): - super().__init__(serde=JsonPlusSerializer()) - self.sync_connection = sync_connection - self.async_connection = async_connection - - @contextmanager - def _get_sync_connection(self) -> Generator[psycopg.Connection, None, None]: # type: ignore[type-arg] - """Get the connection to the Postgres database.""" - with _get_sync_connection(self.sync_connection) as connection: - yield connection - - @asynccontextmanager - async def _get_async_connection( - self, - ) -> AsyncGenerator[psycopg.AsyncConnection, None]: # type: ignore[type-arg] - """Get the connection to the Postgres database.""" - async with _get_async_connection(self.async_connection) as connection: - yield connection - - CREATE_TABLES_QUERY = """ - CREATE TABLE IF NOT EXISTS checkpoints ( - thread_id TEXT NOT NULL, - thread_ts TEXT NOT NULL, - parent_ts TEXT, - checkpoint BYTEA NOT NULL, - metadata BYTEA NOT NULL, - PRIMARY KEY (thread_id, thread_ts) - ); - CREATE TABLE IF NOT EXISTS writes ( - thread_id TEXT NOT NULL, - thread_ts TEXT NOT NULL, - task_id TEXT NOT NULL, - idx INTEGER NOT NULL, - channel TEXT NOT NULL, - value BYTEA, - PRIMARY KEY (thread_id, thread_ts, task_id, idx) - ); - """ - - @staticmethod - def create_tables(connection: psycopg.Connection | ConnectionPool, /) -> None: # type: ignore[type-arg] - """Create the schema for the checkpoint saver.""" - with _get_sync_connection(connection) as conn: - with conn.cursor() as cur: - cur.execute(PostgresSaver.CREATE_TABLES_QUERY) - - @staticmethod - async def acreate_tables( - connection: psycopg.AsyncConnection | AsyncConnectionPool, # type: ignore[type-arg] - /, - ) -> None: - """Create the schema for the checkpoint saver.""" - async with _get_async_connection(connection) as conn: - async with conn.cursor() as cur: - await cur.execute(PostgresSaver.CREATE_TABLES_QUERY) - - @staticmethod - def drop_tables(connection: psycopg.Connection, /) -> None: # type: ignore[type-arg] - """Drop the table for the checkpoint saver.""" - with connection.cursor() as cur: - cur.execute("DROP TABLE IF EXISTS checkpoints, writes;") - - @staticmethod - async def adrop_tables(connection: psycopg.AsyncConnection, /) -> None: # type: ignore[type-arg] - """Drop the table for the checkpoint saver.""" - async with connection.cursor() as cur: - await cur.execute("DROP TABLE IF EXISTS checkpoints, writes;") - - UPSERT_CHECKPOINT_QUERY = """ - INSERT INTO checkpoints - (thread_id, thread_ts, parent_ts, checkpoint, metadata) - VALUES - (%s, %s, %s, %s, %s) - ON CONFLICT (thread_id, thread_ts) - DO UPDATE SET checkpoint = EXCLUDED.checkpoint, - metadata = EXCLUDED.metadata; - """ - - def put( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - ) -> RunnableConfig: - """Put the checkpoint for the given configuration. - Args: - config: The configuration for the checkpoint. - A dict with a `configurable` key which is a dict with - a `thread_id` key and an optional `thread_ts` key. - For example, { 'configurable': { 'thread_id': 'test_thread' } } - checkpoint: The checkpoint to persist. - Returns: - The RunnableConfig that describes the checkpoint that was just created. - It'll contain the `thread_id` and `thread_ts` of the checkpoint. - """ - thread_id = config["configurable"]["thread_id"] - parent_ts = config["configurable"].get("thread_ts") - with self._get_sync_connection() as conn: - with conn.cursor() as cur: - cur.execute( - self.UPSERT_CHECKPOINT_QUERY, - ( - thread_id, - checkpoint["id"], - parent_ts if parent_ts else None, - self.serde.dumps(checkpoint), - self.serde.dumps(metadata), - ), - ) - - return { - "configurable": { - "thread_id": thread_id, - "thread_ts": checkpoint["id"], - }, - } - - async def aput( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - ) -> RunnableConfig: - """Put the checkpoint for the given configuration. - Args: - config: The configuration for the checkpoint. - A dict with a `configurable` key which is a dict with - a `thread_id` key and an optional `thread_ts` key. - For example, { 'configurable': { 'thread_id': 'test_thread' } } - checkpoint: The checkpoint to persist. - Returns: - The RunnableConfig that describes the checkpoint that was just created. - It'll contain the `thread_id` and `thread_ts` of the checkpoint. - """ - thread_id = config["configurable"]["thread_id"] - parent_ts = config["configurable"].get("thread_ts") - async with self._get_async_connection() as conn: - async with conn.cursor() as cur: - await cur.execute( - self.UPSERT_CHECKPOINT_QUERY, - ( - thread_id, - checkpoint["id"], - parent_ts if parent_ts else None, - self.serde.dumps(checkpoint), - self.serde.dumps(metadata), - ), - ) - - return { - "configurable": { - "thread_id": thread_id, - "thread_ts": checkpoint["id"], - }, - } - - UPSERT_WRITES_QUERY = """ - INSERT INTO writes - (thread_id, thread_ts, task_id, idx, channel, value) - VALUES - (%s, %s, %s, %s, %s, %s) - ON CONFLICT (thread_id, thread_ts, task_id, idx) - DO UPDATE SET value = EXCLUDED.value; - """ - - def put_writes( - self, - config: RunnableConfig, - writes: Sequence[tuple[str, Any]], - task_id: str, - ) -> None: - with self._get_sync_connection() as conn: - with conn.cursor() as cur: - cur.executemany( - self.UPSERT_WRITES_QUERY, - [ - ( - str(config["configurable"]["thread_id"]), - str(config["configurable"]["thread_ts"]), - task_id, - idx, - channel, - self.serde.dumps(value), - ) - for idx, (channel, value) in enumerate(writes) - ], - ) - conn.commit() - - async def aput_writes( - self, - config: RunnableConfig, - writes: Sequence[tuple[str, Any]], - task_id: str, - ) -> None: - async with self._get_async_connection() as conn: - async with conn.cursor() as cur: - await cur.executemany( - self.UPSERT_WRITES_QUERY, - [ - ( - str(config["configurable"]["thread_id"]), - str(config["configurable"]["thread_ts"]), - task_id, - idx, - channel, - self.serde.dumps(value), - ) - for idx, (channel, value) in enumerate(writes) - ], - ) - await conn.commit() - - LIST_CHECKPOINTS_QUERY_STR = """ - SELECT checkpoint, metadata, thread_ts, parent_ts - FROM checkpoints - {where} - ORDER BY thread_ts DESC - """ - - def list( - self, - config: RunnableConfig | None, - *, - filter: dict[str, Any] | None = None, - before: RunnableConfig | None = None, - limit: int | None = None, - ) -> Generator[CheckpointTuple, None, None]: - """Get all the checkpoints for the given configuration.""" - where, args = self._search_where(config, filter, before) - query = self.LIST_CHECKPOINTS_QUERY_STR.format(where=where) - if limit: - query += f" LIMIT {limit}" - with self._get_sync_connection() as conn: - with conn.cursor() as cur: - thread_id = config["configurable"]["thread_id"] # type: ignore[index] - cur.execute(query, tuple(args)) - for value in cur: - checkpoint, metadata, thread_ts, parent_ts = value - yield CheckpointTuple( - config={ - "configurable": { - "thread_id": thread_id, - "thread_ts": thread_ts, - } - }, - checkpoint=self.serde.loads(checkpoint), - metadata=self.serde.loads(metadata), - parent_config={ - "configurable": { - "thread_id": thread_id, - "thread_ts": thread_ts, - } - } - if parent_ts - else None, - ) - - async def alist( - self, - config: RunnableConfig | None, - *, - filter: dict[str, Any] | None = None, - before: RunnableConfig | None = None, - limit: int | None = None, - ) -> AsyncIterator[CheckpointTuple]: - """Get all the checkpoints for the given configuration.""" - where, args = self._search_where(config, filter, before) - query = self.LIST_CHECKPOINTS_QUERY_STR.format(where=where) - if limit: - query += f" LIMIT {limit}" - async with self._get_async_connection() as conn: - async with conn.cursor() as cur: - thread_id = config["configurable"]["thread_id"] # type: ignore[index] - await cur.execute(query, tuple(args)) - async for value in cur: - checkpoint, metadata, thread_ts, parent_ts = value - yield CheckpointTuple( - config={ - "configurable": { - "thread_id": thread_id, - "thread_ts": thread_ts, - } - }, - checkpoint=self.serde.loads(checkpoint), - metadata=self.serde.loads(metadata), - parent_config={ - "configurable": { - "thread_id": thread_id, - "thread_ts": thread_ts, - } - } - if parent_ts - else None, - ) - - GET_CHECKPOINT_BY_TS_QUERY = """ - SELECT checkpoint, metadata, thread_ts, parent_ts - FROM checkpoints - WHERE thread_id = %(thread_id)s AND thread_ts = %(thread_ts)s - """ - - GET_CHECKPOINT_QUERY = """ - SELECT checkpoint, metadata, thread_ts, parent_ts - FROM checkpoints - WHERE thread_id = %(thread_id)s - ORDER BY thread_ts DESC LIMIT 1 - """ - - def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: - """Get the checkpoint tuple for the given configuration. - Args: - config: The configuration for the checkpoint. - A dict with a `configurable` key which is a dict with - a `thread_id` key and an optional `thread_ts` key. - For example, { 'configurable': { 'thread_id': 'test_thread' } } - Returns: - The checkpoint tuple for the given configuration if it exists, - otherwise None. - If thread_ts is None, the latest checkpoint is returned if it exists. - """ - thread_id = config["configurable"]["thread_id"] - thread_ts = config["configurable"].get("thread_ts") - with self._get_sync_connection() as conn: - with conn.cursor() as cur: - # find the latest checkpoint for the thread_id - if thread_ts: - cur.execute( - self.GET_CHECKPOINT_BY_TS_QUERY, - { - "thread_id": thread_id, - "thread_ts": thread_ts, - }, - ) - else: - cur.execute( - self.GET_CHECKPOINT_QUERY, - { - "thread_id": thread_id, - }, - ) - - # if a checkpoint is found, return it - if value := cur.fetchone(): - checkpoint, metadata, thread_ts, parent_ts = value - if not config["configurable"].get("thread_ts"): - config = { - "configurable": { - "thread_id": thread_id, - "thread_ts": thread_ts, - } - } - - # find any pending writes - cur.execute( - "SELECT task_id, channel, value FROM writes WHERE thread_id = %(thread_id)s AND thread_ts = %(thread_ts)s", - { - "thread_id": thread_id, - "thread_ts": thread_ts, - }, - ) - # deserialize the checkpoint and metadata - return CheckpointTuple( - config=config, - checkpoint=self.serde.loads(checkpoint), - metadata=self.serde.loads(metadata), - parent_config={ - "configurable": { - "thread_id": thread_id, - "thread_ts": parent_ts, - } - } - if parent_ts - else None, - pending_writes=[ - (task_id, channel, self.serde.loads(value)) - for task_id, channel, value in cur - ], - ) - return None - - async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: - """Get the checkpoint tuple for the given configuration. - Args: - config: The configuration for the checkpoint. - A dict with a `configurable` key which is a dict with - a `thread_id` key and an optional `thread_ts` key. - For example, { 'configurable': { 'thread_id': 'test_thread' } } - Returns: - The checkpoint tuple for the given configuration if it exists, - otherwise None. - If thread_ts is None, the latest checkpoint is returned if it exists. - """ - thread_id = config["configurable"]["thread_id"] - thread_ts = config["configurable"].get("thread_ts") - async with self._get_async_connection() as conn: - async with conn.cursor() as cur: - # find the latest checkpoint for the thread_id - if thread_ts: - await cur.execute( - self.GET_CHECKPOINT_BY_TS_QUERY, - { - "thread_id": thread_id, - "thread_ts": thread_ts, - }, - ) - else: - await cur.execute( - self.GET_CHECKPOINT_QUERY, - { - "thread_id": thread_id, - }, - ) - # if a checkpoint is found, return it - if value := await cur.fetchone(): - checkpoint, metadata, thread_ts, parent_ts = value - if not config["configurable"].get("thread_ts"): - config = { - "configurable": { - "thread_id": thread_id, - "thread_ts": thread_ts, - } - } - - # find any pending writes - await cur.execute( - "SELECT task_id, channel, value FROM writes WHERE thread_id = %(thread_id)s AND thread_ts = %(thread_ts)s", - { - "thread_id": thread_id, - "thread_ts": thread_ts, - }, - ) - # deserialize the checkpoint and metadata - return CheckpointTuple( - config=config, - checkpoint=self.serde.loads(checkpoint), - metadata=self.serde.loads(metadata), - parent_config={ - "configurable": { - "thread_id": thread_id, - "thread_ts": parent_ts, - } - } - if parent_ts - else None, - pending_writes=[ - (task_id, channel, self.serde.loads(value)) - async for task_id, channel, value in cur - ], - ) - return None - - def _search_where( - self, - config: RunnableConfig | None, - filter: dict[str, Any] | None = None, - before: RunnableConfig | None = None, - ) -> tuple[str, List[Any]]: # noqa: UP006 - """Return WHERE clause predicates for given config, filter, and before parameters. - Args: - config (Optional[RunnableConfig]): The config to use for filtering. - filter (Optional[Dict[str, Any]]): Additional filtering criteria. - before (Optional[RunnableConfig]): A config to limit results before a certain timestamp. - Returns: - Tuple[str, Sequence[Any]]: A tuple containing the WHERE clause and parameter values. - """ - wheres = [] - param_values = [] - - # Add predicate for config - if config is not None: - wheres.append("thread_id = %s ") - param_values.append(config["configurable"]["thread_id"]) - - if filter: - raise NotImplementedError() - - # Add predicate for limiting results before a certain timestamp - if before is not None: - wheres.append("thread_ts < %s") - param_values.append(before["configurable"]["thread_ts"]) - - where_clause = "WHERE " + " AND ".join(wheres) if wheres else "" - return where_clause, param_values diff --git a/backend/app/core/graph/checkpoint/utils.py b/backend/app/core/graph/checkpoint/utils.py index 1629c1c..a0173fb 100644 --- a/backend/app/core/graph/checkpoint/utils.py +++ b/backend/app/core/graph/checkpoint/utils.py @@ -5,10 +5,10 @@ from langchain_core.documents import Document from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage from langgraph.checkpoint.base import CheckpointTuple +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver from psycopg import AsyncConnection from app.core.config import settings -from app.core.graph.checkpoint.postgres import PostgresSaver from app.core.graph.messages import ChatResponse @@ -106,8 +106,10 @@ async def get_checkpoint_tuples(thread_id: str) -> CheckpointTuple | None: Returns: CheckpointTuple: The latest checkpoint tuple. """ - async with await AsyncConnection.connect(settings.PG_DATABASE_URI) as conn: - checkpointer = PostgresSaver(async_connection=conn) + async with await AsyncConnection.connect( + settings.PG_DATABASE_URI, **settings.SQLALCHEMY_CONNECTION_KWARGS + ) as conn: + checkpointer = AsyncPostgresSaver(conn=conn) checkpoint_tuple = await checkpointer.aget_tuple( {"configurable": {"thread_id": thread_id}} )