Compare commits

..
Author SHA1 Message Date
Nuno Campos 2c984dfe4e checkpoint 2.0.10 2025-01-15 11:14:03 -08:00
Nuno CamposandGitHub 09c0d9cec3 Add optional task_path arg for put_writes() (#3049)
- Will be used for sorting pending_sends when available
2025-01-15 11:11:10 -08:00
Nuno Campos 31734fb792 Lint 2025-01-15 11:02:10 -08:00
Nuno Campos eed577ee2a Update sqlite signature 2025-01-15 10:58:46 -08:00
Nuno Campos bba00506ea Lint 2025-01-15 10:57:26 -08:00
Nuno Campos fa12538a4e Lint 2025-01-15 10:50:11 -08:00
Nuno Campos c1a7bb8902 Lint 2025-01-15 10:49:47 -08:00
Nuno Campos 14b500d8a8 Fix 2025-01-15 10:47:09 -08:00
Nuno Campos 053a501db3 Add optional task_path arg for put_writes()
- Will be used for sorting pending_sends when available
2025-01-15 10:42:54 -08:00
Andrew NguonlyandGitHub bbe8eacf8d docs: Update Studio FAQ (#3048)
### Screenshot

![image](https://github.com/user-attachments/assets/6da2faa6-1c62-44e4-96f8-3450f94cb15c)
2025-01-15 10:23:43 -08:00
Nuno CamposandGitHub 50cb387904 Remove duckdb checkpointer and store (#3046)
- duckdb is too buggy to be able to provide reliable checkpointer and
store
2025-01-15 09:23:49 -08:00
Nuno Campos 003b1883c9 Remove duckdb checkpointer and store
- duckdb is too buggy to be able to provide reliable checkpointer and store
2025-01-15 09:15:17 -08:00
Nuno CamposandGitHub d12830e95c Fix tracing hierarchy for imperative api (#3036) 2025-01-15 09:10:45 -08:00
Nuno CamposandGitHub 4426552c26 Fix flaky test output order (#3045) 2025-01-15 09:02:05 -08:00
Vadym BardaandGitHub e2554c9616 langgraph: fix non-empty value check in ensure_config (#3039)
Fixes https://github.com/langchain-ai/langgraph/issues/2890
2025-01-15 16:54:31 +00:00
Nuno CamposandGitHub f1ee650489 Remove CI job to test against core 0.2.x (#3044) 2025-01-15 08:52:58 -08:00
Nuno Campos b7e4656c91 Fix flaky test output order 2025-01-15 08:45:35 -08:00
Nuno Campos 65a41942ef Remove CI job to test against core 0.2.x 2025-01-15 08:39:09 -08:00
Nuno Campos c767d86c9c Remove ci job for removed flag 2025-01-15 08:37:24 -08:00
Nuno Campos 01331ef858 Lint 2025-01-15 08:33:09 -08:00
Nuno CamposandGitHub 24a4c67c52 Merge branch 'main' into nc/14jan/fix-tracing-hierarchy-imperative 2025-01-15 08:31:38 -08:00
Nuno CamposandGitHub c7e43f86d4 Remove send v2 (#3033) 2025-01-15 08:29:05 -08:00
Vadym BardaandGitHub b18d266f2c langgraph: allow model names as strings in create_react_agent (#3031)
```python
agent = create_react_agent("anthropic:claude-3-5-sonnet-latest", [add])
agent.invoke({"messages": [("user", "what's 3 + 5")]})
```
2025-01-15 09:43:37 -05:00
Eugene YurtsevandGitHub 536ee7b37a docs: document inject kwarg (#3026) 2025-01-14 21:41:27 -05:00
Nuno Campos 901273c0e2 Lint 2025-01-14 17:09:30 -08:00
Nuno Campos 7bdbd62611 Fix tracing hierarchy for imperative api 2025-01-14 17:03:45 -08:00
Nuno Campos 0b74e25f72 Remove send v2 2025-01-14 15:49:57 -08:00
Eugene YurtsevandGitHub ff843ab005 docs: internal doc nits (#2946)
Update internal documentation
2025-01-14 23:16:38 +00:00
Vadym BardaandGitHub f7fae7e140 langgraph: fix ismethod check in add_node (#3032)
Fixes #2893 #2965
2025-01-14 17:17:50 -05:00
Eugene YurtsevandGitHub 651ee8bd24 Update local-server.md (#3030) 2025-01-14 16:22:26 -05:00
Eugene YurtsevandGitHub 54bdba2da9 Update introduction.ipynb (#3029) 2025-01-14 16:20:54 -05:00
Eugene YurtsevandGitHub 7e5806cb0b Update local-server.md 2025-01-14 16:13:57 -05:00
Eugene YurtsevandGitHub 66f674fead Update introduction.ipynb 2025-01-14 16:12:20 -05:00
Vadym BardaandGitHub f029d615e6 ci: fix docs build (#3028) 2025-01-14 15:58:03 -05:00
Eugene Yurtsev bd76773b31 x 2025-01-14 15:50:20 -05:00
ccurmeandGitHub 9b6e6d67dc layout (#2972)
Some docs layout improvements to help guide user journey.

Currently we have `Home | Tutorials | How-tos | Concepts | Reference` in
top-level horizontal navigation bar.

Here we make these updates:
- Top-level horizontal navigation bar is just `Home | API Reference`
- Add vertical sidebar to `Home` with sections:
  - Introduction
  - Get started
  - Guides
  - Resources

`Get Started` contains quickstarts for LG and LG Platform / deployment.
These are tutorials in Diataxis terms.

`Guides` contains index pages for how-tos, concepts, tutorials.

Advantage of this organization is that users are directed naturally down
the sidebar from Intro -> Get started -> How-tos, which is roughly how
we expect them to proceed.

This also makes deployment info more accessible as it is highlighted in
the "Getting started" section.

![Screenshot 2025-01-14 at 1 37
01 PM](https://github.com/user-attachments/assets/76dcde25-874b-4def-ad71-5313e399b7d7)
2025-01-14 15:29:59 -05:00
ccurmeandGitHub ecfb3e3850 Merge branch 'main' into eugene/langgraph_nav 2025-01-14 15:21:26 -05:00
Chester Curme ec71cd6fc1 update 2025-01-14 15:09:46 -05:00
Chester Curme 2db6d9d7b5 clean up 2025-01-14 15:08:48 -05:00
Chester Curme 0f5df797af fix link 2025-01-14 15:00:38 -05:00
Chester Curme 5dc3f14e8b populate troubleshootings nav 2025-01-14 14:56:23 -05:00
Chester Curme f25ea9ecaf populate tutorials nav 2025-01-14 14:50:05 -05:00
Chester Curme 9940a9f833 populate how-tos nav 2025-01-14 14:47:42 -05:00
Chester Curme d175d5e4a8 populate concepts nav 2025-01-14 14:42:10 -05:00
Chester Curme fe63440d98 add deployment landing page 2025-01-14 14:40:45 -05:00
vbarda a9073241b2 add css for navbar depth 2025-01-14 13:31:39 -05:00
Chester Curme 0a04262f05 add custom css 2025-01-14 12:46:00 -05:00
Chester Curme 381796c474 hack 2025-01-14 12:39:38 -05:00
Chester Curme 09ae5ce05d update 2025-01-14 11:26:49 -05:00
Chester Curme 59e7751a8c update 2025-01-14 10:48:50 -05:00
Chester Curme b9816e31ff update 2025-01-13 19:12:51 -05:00
Chester Curme 3ea7f9469a update 2025-01-13 18:26:24 -05:00
Chester Curme 24e61fafa9 update 2025-01-13 16:22:58 -05:00
Chester Curme 2b370cf3b7 template applications -> lg plat 2025-01-13 14:21:12 -05:00
Chester Curme 839bea0f33 fixup 2025-01-13 14:16:14 -05:00
Eugene Yurtsev fba5ca7ab9 x 2025-01-09 12:24:30 -05:00
Eugene Yurtsev f141089811 x 2025-01-09 11:11:04 -05:00
Eugene Yurtsev 38487cb158 x 2025-01-09 11:11:04 -05:00
68 changed files with 1355 additions and 6837 deletions
+1 -16
View File
@@ -17,21 +17,11 @@ jobs:
- "3.11"
- "3.12"
- "3.13"
core-version:
- "latest"
ff-send-v2:
- "false"
include:
- python-version: "3.11"
core-version: ">=0.2.42,<0.3.0"
- python-version: "3.11"
core-version: "latest"
ff-send-v2: "true"
defaults:
run:
working-directory: libs/langgraph
name: "test #${{ matrix.python-version }} (langchain-core: ${{ matrix.core-version }}, ff-send-v2: ${{ matrix.ff-send-v2 }})"
name: "test #${{ matrix.python-version }}"
steps:
- uses: actions/checkout@v4
- name: Set up Python ${{ matrix.python-version }} + Poetry ${{ env.POETRY_VERSION }}
@@ -51,14 +41,9 @@ jobs:
shell: bash
run: |
poetry install --with dev
if [ "${{ matrix.core-version }}" != "latest" ]; then
poetry run pip install "langchain-core${{ matrix.core-version }}"
fi
- name: Run tests
shell: bash
env:
LANGGRAPH_FF_SEND_V2: ${{ matrix.ff-send-v2 }}
run: |
make test_parallel
+14 -6
View File
@@ -31,7 +31,6 @@ jobs:
"libs/cli",
"libs/checkpoint",
"libs/checkpoint-sqlite",
"libs/checkpoint-duckdb",
"libs/checkpoint-postgres",
"libs/scheduler-kafka",
]
@@ -44,12 +43,12 @@ jobs:
name: cd ${{ matrix.working-directory }}
strategy:
matrix:
working-directory: [
working-directory:
[
"libs/cli",
"libs/checkpoint",
"libs/checkpoint-sqlite",
"libs/checkpoint-duckdb",
"libs/checkpoint-postgres"
"libs/checkpoint-postgres",
]
uses: ./.github/workflows/_test.yml
with:
@@ -76,7 +75,7 @@ jobs:
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: '3.11'
python-version: "3.11"
- name: Run check_sdk_methods script
run: python .github/scripts/check_sdk_methods.py
@@ -133,7 +132,16 @@ jobs:
ci_success:
name: "CI Success"
needs: [lint, lint-js, test, test-langgraph, test-scheduler-kafka, integration-test, test-js]
needs:
[
lint,
lint-js,
test,
test-langgraph,
test-scheduler-kafka,
integration-test,
test-js,
]
if: |
always()
runs-on: ubuntu-latest
+2 -2
View File
@@ -85,7 +85,7 @@ jobs:
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://(api|web|docs|academy)\.smith\.langchain\.com/.*" \
--check-links-ignore "https://x.com/.*" \
--check-links-ignore "https://github\.com/.*" \
--check-links-ignore "http://localhost:8123/.*" \
@@ -106,7 +106,7 @@ jobs:
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://(api|web|docs|academy)\.smith\.langchain\.com/.*" \
--check-links-ignore "http://localhost:8123/.*" \
--check-links-ignore "http://localhost:2024.*" \
--check-links-ignore "http://127.0.0.1:.*" \
-2
View File
@@ -1,6 +1,4 @@
---
hide:
- navigation
title: Concepts
description: Conceptual Guide for LangGraph
---
+12 -1
View File
@@ -60,7 +60,18 @@ LangGraph Studio (desktop) requires Docker Desktop version 4.24 or higher. Pleas
#### Configuration or environment issues
Another reason your project might fail to start is because your configuration file is defined incorrectly, or you are missing required environment variables.
Another reason your project might fail to start is because your configuration file is defined incorrectly, or you are missing required environment variables.
!!! Important "Note (desktop only)"
LangGraph Studio Desktop automatically populates `LANGCHAIN_*` environment variables for license verification and tracing, regardless of the contents of the `.env` file. All other environment variables defined in `.env` will be read as normal.
#### Incorrect data region (desktop only)
If you receive a license verification error when attempting to start the LangGraph Server, you may be logged into the incorrect LangSmith data region. Ensure that you're logged into the correct LangSmith data region and ensure that the LangSmith account has access to LangGraph platform.
1. In the top right-hand corner, click the user icon and select `Logout`.
1. At the login screen, click the `Data Region` dropdown menu and select the appropriate data region. Then click `Login to LangSmith`.
### How does interrupt work?
-2
View File
@@ -1,6 +1,4 @@
---
hide:
- navigation
title: How-to Guides
description: How to accomplish common tasks in LangGraph
---
-2
View File
@@ -1,7 +1,5 @@
---
hide_comments: true
hide:
- navigation
title: Home
---
+18
View File
@@ -0,0 +1,18 @@
# Deployment
Get started deploying your LangGraph applications locally or on the cloud with
[LangGraph Platform](../concepts/langgraph_platform.md).
## Get Started 🚀 {#quick-start}
- [LangGraph Server Quickstart](../tutorials/langgraph-platform/local-server.md): Launch a LangGraph server locally and interact with it using REST API and LangGraph Studio Web UI.
- [LangGraph Template Quickstart](../concepts/template_applications.md): Start building with LangGraph Platform using a template application.
- [Deploy with LangGraph Cloud Quickstart](../cloud/quick_start.md): Deploy a LangGraph app using LangGraph Cloud.
## Deployment Options
- [Self-Hosted Lite](../concepts/self_hosted.md): A free (up to 1 million nodes executed), limited version of LangGraph Platform that you can run locally or in a self-hosted manner
- [Cloud SaaS](../concepts/langgraph_cloud.md): Hosted as part of LangSmith.
- [Bring Your Own Cloud](../concepts/bring_your_own_cloud.md): We manage the infrastructure, so you don't have to, but the infrastructure all runs within your cloud.
- [Self-Hosted Enterprise](../concepts/self_hosted.md): Completely managed by you.
-2
View File
@@ -1,6 +1,4 @@
---
hide:
- navigation
title: Tutorials
---
+1 -1
View File
@@ -5,7 +5,7 @@
"id": "4a1aae78-88a6-4133-b905-7e46c8e3772f",
"metadata": {},
"source": [
"# 🚀 LangGraph Quick Start\n",
"# 🚀 LangGraph Quickstart\n",
"\n",
"In this tutorial, we will build a support chatbot in LangGraph that can:\n",
"\n",
@@ -1,4 +1,4 @@
# QuickStart: Launch Local LangGraph Server
# Quickstart: Launch Local LangGraph Server
This is a quick start guide to help you get a LangGraph app up and running locally.
+280 -269
View File
@@ -22,7 +22,6 @@ theme:
- header.autohide
- navigation.expand
- navigation.footer
- navigation.indexes
- navigation.instant
- navigation.sections
- navigation.instant.prefetch
@@ -30,7 +29,6 @@ theme:
- navigation.path
- navigation.prune
- navigation.tabs
- navigation.tabs.sticky
- navigation.top
- navigation.tracking
- search.highlight
@@ -89,281 +87,294 @@ plugins:
options:
filters:
- "!^_"
nav:
- Home: index.md
- Tutorials:
- tutorials/index.md
- Quick Start:
- Quick Start: tutorials#quick-start
- tutorials/introduction.ipynb
- tutorials/langgraph-platform/local-server.md
- cloud/quick_start.md
- Chatbots:
- Chatbots: tutorials#chatbots
- tutorials/customer-support/customer-support.ipynb
- tutorials/chatbots/information-gather-prompting.ipynb
- tutorials/code_assistant/langgraph_code_assistant.ipynb
- RAG:
- RAG: tutorials#rag
- tutorials/rag/langgraph_adaptive_rag.ipynb
- tutorials/rag/langgraph_adaptive_rag_local.ipynb
- tutorials/rag/langgraph_agentic_rag.ipynb
- tutorials/rag/langgraph_crag.ipynb
- tutorials/rag/langgraph_crag_local.ipynb
- tutorials/rag/langgraph_self_rag.ipynb
- tutorials/rag/langgraph_self_rag_local.ipynb
- tutorials/sql-agent.ipynb
- Agent Architectures:
- Agent Architectures: tutorials#agent-architectures
- Multi-Agent Systems:
- Multi-Agent Systems: tutorials#multi-agent-systems
- tutorials/multi_agent/multi-agent-collaboration.ipynb
- tutorials/multi_agent/agent_supervisor.ipynb
- tutorials/multi_agent/hierarchical_agent_teams.ipynb
- Planning Agents:
- Planning Agents: tutorials#planning-agents
- tutorials/plan-and-execute/plan-and-execute.ipynb
- tutorials/rewoo/rewoo.ipynb
- tutorials/llm-compiler/LLMCompiler.ipynb
- Reflection & Critique:
- Reflection & Critique: tutorials#reflection-critique
- tutorials/reflection/reflection.ipynb
- tutorials/reflexion/reflexion.ipynb
- tutorials/tot/tot.ipynb
- tutorials/lats/lats.ipynb
- tutorials/self-discover/self-discover.ipynb
- Evaluation & Analysis:
- Evaluation & Analysis: tutorials#evaluation
- tutorials/chatbot-simulation-evaluation/agent-simulation-evaluation.ipynb
- tutorials/chatbot-simulation-evaluation/langsmith-agent-simulation-evaluation.ipynb
- Experimental:
- Experimental: tutorials#experimental
- tutorials/storm/storm.ipynb
- tutorials/tnt-llm/tnt-llm.ipynb
- tutorials/web-navigation/web_voyager.ipynb
- tutorials/usaco/usaco.ipynb
- tutorials/extraction/retries.ipynb
- LangGraph Platform:
- LangGraph Platform: concepts#langgraph-platform
- tutorials/auth/getting_started.md
- tutorials/auth/resource_auth.md
- tutorials/auth/add_auth_server.md
- How-to Guides:
- how-tos/index.md
- LangGraph:
- LangGraph: how-tos#langgraph
- Controllability:
- Controllability: how-tos#controllability
- how-tos/branching.ipynb
- how-tos/map-reduce.ipynb
- how-tos/recursion-limit.ipynb
- how-tos/command.ipynb
- Persistence:
- Persistence: how-tos#persistence
- how-tos/persistence.ipynb
- how-tos/subgraph-persistence.ipynb
- how-tos/cross-thread-persistence.ipynb
- how-tos/persistence_postgres.ipynb
- how-tos/persistence_mongodb.ipynb
- how-tos/persistence_redis.ipynb
- Memory:
- Memory: how-tos#memory
- how-tos/memory/manage-conversation-history.ipynb
- how-tos/memory/delete-messages.ipynb
- how-tos/memory/add-summary-conversation-history.ipynb
- how-tos/memory/semantic-search.ipynb
- Human-in-the-loop:
- Human-in-the-loop: how-tos#human-in-the-loop
- how-tos/human_in_the_loop/breakpoints.ipynb
- how-tos/human_in_the_loop/dynamic_breakpoints.ipynb
- how-tos/human_in_the_loop/edit-graph-state.ipynb
- how-tos/human_in_the_loop/wait-user-input.ipynb
- how-tos/human_in_the_loop/time-travel.ipynb
- how-tos/human_in_the_loop/review-tool-calls.ipynb
- Streaming:
- Streaming: how-tos#streaming
- how-tos/stream-values.ipynb
- how-tos/stream-updates.ipynb
- how-tos/streaming-tokens.ipynb
- how-tos/streaming-tokens-without-langchain.ipynb
- how-tos/streaming-content.ipynb
- how-tos/stream-multiple.ipynb
- how-tos/streaming-events-from-within-tools.ipynb
- how-tos/streaming-events-from-within-tools-without-langchain.ipynb
- how-tos/streaming-from-final-node.ipynb
- how-tos/streaming-subgraphs.ipynb
- how-tos/disable-streaming.ipynb
- Tool calling:
- Tool calling: how-tos#tool-calling
- how-tos/tool-calling.ipynb
- how-tos/tool-calling-errors.ipynb
- how-tos/pass-run-time-values-to-tools.ipynb
- how-tos/update-state-from-tools.ipynb
- how-tos/pass-config-to-tools.ipynb
- how-tos/many-tools.ipynb
- Subgraphs:
- Subgraphs: how-tos#subgraphs
- how-tos/subgraph.ipynb
- how-tos/subgraphs-manage-state.ipynb
- how-tos/subgraph-transform-state.ipynb
- Multi-agent:
- Multi-agent: how-tos#multi-agent
- how-tos/agent-handoffs.ipynb
- how-tos/multi-agent-network.ipynb
- how-tos/multi-agent-multi-turn-convo.ipynb
- State Management:
- State Management: how-tos#state-management
- how-tos/state-model.ipynb
- how-tos/input_output_schema.ipynb
- how-tos/pass_private_state.ipynb
- Other:
- Other: how-tos#other
- how-tos/async.ipynb
- how-tos/visualization.ipynb
- how-tos/configuration.ipynb
- how-tos/node-retries.ipynb
- how-tos/react-agent-structured-output.ipynb
- how-tos/run-id-langsmith.ipynb
- how-tos/return-when-recursion-limit-hits.ipynb
- Prebuilt ReAct Agent:
- Prebuilt ReAct Agent: how-tos#prebuilt-react-agent
- how-tos/create-react-agent.ipynb
- how-tos/create-react-agent-memory.ipynb
- how-tos/create-react-agent-system-prompt.ipynb
- how-tos/create-react-agent-hitl.ipynb
- how-tos/react-agent-from-scratch.ipynb
- LangGraph Platform:
- LangGraph Platform: how-tos#langgraph-platform
- Application Structure:
- Application Structure: how-tos#application-structure
- cloud/deployment/setup.md
- cloud/deployment/setup_pyproject.md
- cloud/deployment/setup_javascript.md
- cloud/deployment/semantic_search.md
- cloud/deployment/custom_docker.md
- cloud/deployment/test_locally.md
- cloud/deployment/graph_rebuild.md
- Deployment:
- Deployment: how-tos#deployment
- cloud/deployment/cloud.md
- how-tos/deploy-self-hosted.md
- how-tos/use-remote-graph.md
- Authentication & Access Control:
- Authentication & Access Control: how-tos#authentication-access-control
- cloud/how-tos/auth/custom_auth_new.md
- cloud/how-tos/auth/openapi_security_new.md
- Assistants:
- Assistants: how-tos#assistants
- cloud/how-tos/configuration_cloud.md
- cloud/how-tos/assistant_versioning.md
- Threads:
- Threads: how-tos#threads
- cloud/how-tos/copy_threads.md
- cloud/how-tos/check_thread_status.md
- Runs:
- Runs: how-tos#runs
- cloud/how-tos/background_run.md
- cloud/how-tos/same-thread.md
- cloud/how-tos/cron_jobs.md
- cloud/how-tos/stateless_runs.md
- Streaming:
- Streaming: how-tos#streaming_1
- cloud/how-tos/stream_values.md
- cloud/how-tos/stream_updates.md
- cloud/how-tos/stream_messages.md
- cloud/how-tos/stream_events.md
- cloud/how-tos/stream_debug.md
- cloud/how-tos/stream_multiple.md
nav:
- Home:
- Introduction: index.md
- Get started:
- Learn the basics: tutorials/introduction.ipynb
- Deployment:
- tutorials/deployment.md
- Local Deploy: tutorials/langgraph-platform/local-server.md
- Template Applications: concepts/template_applications.md # TODO: make tutorial
- Cloud Deploy: cloud/quick_start.md
- Guides:
- How-to Guides:
- how-tos/index.md
- LangGraph:
- LangGraph: how-tos#langgraph
- Controllability:
- Controllability: how-tos#controllability
- how-tos/branching.ipynb
- how-tos/map-reduce.ipynb
- how-tos/recursion-limit.ipynb
- how-tos/command.ipynb
- Persistence:
- Persistence: how-tos#persistence
- how-tos/persistence.ipynb
- how-tos/subgraph-persistence.ipynb
- how-tos/cross-thread-persistence.ipynb
- how-tos/persistence_postgres.ipynb
- how-tos/persistence_mongodb.ipynb
- how-tos/persistence_redis.ipynb
- Memory:
- Memory: how-tos#memory
- how-tos/memory/manage-conversation-history.ipynb
- how-tos/memory/delete-messages.ipynb
- how-tos/memory/add-summary-conversation-history.ipynb
- how-tos/memory/semantic-search.ipynb
- Human-in-the-loop:
- Human-in-the-loop: how-tos#human-in-the-loop_1
- cloud/how-tos/human_in_the_loop_breakpoint.md
- cloud/how-tos/human_in_the_loop_user_input.md
- cloud/how-tos/human_in_the_loop_edit_state.md
- cloud/how-tos/human_in_the_loop_time_travel.md
- cloud/how-tos/human_in_the_loop_review_tool_calls.md
- Double-texting:
- Double-texting: how-tos#double-texting
- cloud/how-tos/interrupt_concurrent.md
- cloud/how-tos/rollback_concurrent.md
- cloud/how-tos/reject_concurrent.md
- cloud/how-tos/enqueue_concurrent.md
- Webhooks:
- cloud/how-tos/webhooks.md
- Cron Jobs:
- cloud/how-tos/cron_jobs.md
- LangGraph Studio:
- LangGraph Studio: how-tos#langgraph-studio
- cloud/how-tos/test_deployment.md
- cloud/how-tos/test_local_deployment.md
- cloud/how-tos/invoke_studio.md
- cloud/how-tos/threads_studio.md
- cloud/how-tos/datasets_studio.md
- Human-in-the-loop: how-tos#human-in-the-loop
- how-tos/human_in_the_loop/breakpoints.ipynb
- how-tos/human_in_the_loop/dynamic_breakpoints.ipynb
- how-tos/human_in_the_loop/edit-graph-state.ipynb
- how-tos/human_in_the_loop/wait-user-input.ipynb
- how-tos/human_in_the_loop/time-travel.ipynb
- how-tos/human_in_the_loop/review-tool-calls.ipynb
- Streaming:
- Streaming: how-tos#streaming
- how-tos/stream-values.ipynb
- how-tos/stream-updates.ipynb
- how-tos/streaming-tokens.ipynb
- how-tos/streaming-tokens-without-langchain.ipynb
- how-tos/streaming-content.ipynb
- how-tos/stream-multiple.ipynb
- how-tos/streaming-events-from-within-tools.ipynb
- how-tos/streaming-events-from-within-tools-without-langchain.ipynb
- how-tos/streaming-from-final-node.ipynb
- how-tos/streaming-subgraphs.ipynb
- how-tos/disable-streaming.ipynb
- Tool calling:
- Tool calling: how-tos#tool-calling
- how-tos/tool-calling.ipynb
- how-tos/tool-calling-errors.ipynb
- how-tos/pass-run-time-values-to-tools.ipynb
- how-tos/update-state-from-tools.ipynb
- how-tos/pass-config-to-tools.ipynb
- how-tos/many-tools.ipynb
- Subgraphs:
- Subgraphs: how-tos#subgraphs
- how-tos/subgraph.ipynb
- how-tos/subgraphs-manage-state.ipynb
- how-tos/subgraph-transform-state.ipynb
- Multi-agent:
- Multi-agent: how-tos#multi-agent
- how-tos/agent-handoffs.ipynb
- how-tos/multi-agent-network.ipynb
- how-tos/multi-agent-multi-turn-convo.ipynb
- State Management:
- State Management: how-tos#state-management
- how-tos/state-model.ipynb
- how-tos/input_output_schema.ipynb
- how-tos/pass_private_state.ipynb
- Other:
- Other: how-tos#other
- how-tos/async.ipynb
- how-tos/visualization.ipynb
- how-tos/configuration.ipynb
- how-tos/node-retries.ipynb
- how-tos/react-agent-structured-output.ipynb
- how-tos/run-id-langsmith.ipynb
- how-tos/return-when-recursion-limit-hits.ipynb
- Prebuilt ReAct Agent:
- Prebuilt ReAct Agent: how-tos#prebuilt-react-agent
- how-tos/create-react-agent.ipynb
- how-tos/create-react-agent-memory.ipynb
- how-tos/create-react-agent-system-prompt.ipynb
- how-tos/create-react-agent-hitl.ipynb
- how-tos/react-agent-from-scratch.ipynb
- LangGraph Platform:
- LangGraph Platform: how-tos#langgraph-platform
- Application Structure:
- Application Structure: how-tos#application-structure
- cloud/deployment/setup.md
- cloud/deployment/setup_pyproject.md
- cloud/deployment/setup_javascript.md
- cloud/deployment/semantic_search.md
- cloud/deployment/custom_docker.md
- cloud/deployment/test_locally.md
- cloud/deployment/graph_rebuild.md
- Deployment:
- Deployment: how-tos#deployment
- cloud/deployment/cloud.md
- how-tos/deploy-self-hosted.md
- how-tos/use-remote-graph.md
- Authentication & Access Control:
- Authentication & Access Control: how-tos#authentication-access-control
- cloud/how-tos/auth/custom_auth_new.md
- cloud/how-tos/auth/openapi_security_new.md
- Assistants:
- Assistants: how-tos#assistants
- cloud/how-tos/configuration_cloud.md
- cloud/how-tos/assistant_versioning.md
- Threads:
- Threads: how-tos#threads
- cloud/how-tos/copy_threads.md
- cloud/how-tos/check_thread_status.md
- Runs:
- Runs: how-tos#runs
- cloud/how-tos/background_run.md
- cloud/how-tos/same-thread.md
- cloud/how-tos/cron_jobs.md
- cloud/how-tos/stateless_runs.md
- Streaming:
- Streaming: how-tos#streaming_1
- cloud/how-tos/stream_values.md
- cloud/how-tos/stream_updates.md
- cloud/how-tos/stream_messages.md
- cloud/how-tos/stream_events.md
- cloud/how-tos/stream_debug.md
- cloud/how-tos/stream_multiple.md
- Human-in-the-loop:
- Human-in-the-loop: how-tos#human-in-the-loop_1
- cloud/how-tos/human_in_the_loop_breakpoint.md
- cloud/how-tos/human_in_the_loop_user_input.md
- cloud/how-tos/human_in_the_loop_edit_state.md
- cloud/how-tos/human_in_the_loop_time_travel.md
- cloud/how-tos/human_in_the_loop_review_tool_calls.md
- Double-texting:
- Double-texting: how-tos#double-texting
- cloud/how-tos/interrupt_concurrent.md
- cloud/how-tos/rollback_concurrent.md
- cloud/how-tos/reject_concurrent.md
- cloud/how-tos/enqueue_concurrent.md
- Webhooks:
- cloud/how-tos/webhooks.md
- Cron Jobs:
- cloud/how-tos/cron_jobs.md
- LangGraph Studio:
- LangGraph Studio: how-tos#langgraph-studio
- cloud/how-tos/test_deployment.md
- cloud/how-tos/test_local_deployment.md
- cloud/how-tos/invoke_studio.md
- cloud/how-tos/threads_studio.md
- cloud/how-tos/datasets_studio.md
- Concepts:
- concepts/index.md
- LangGraph:
- LangGraph: concepts#langgraph
- concepts/high_level.md
- concepts/low_level.md
- concepts/agentic_concepts.md
- concepts/multi_agent.md
- concepts/breakpoints
- concepts/human_in_the_loop.md
- concepts/time-travel.md
- concepts/persistence.md
- concepts/memory.md
- concepts/streaming.md
- LangGraph Platform:
- LangGraph Platform: concepts#langgraph-platform
- High Level:
- High Level: concepts#high-level
- concepts/langgraph_platform.md
- concepts/deployment_options.md
- concepts/plans.md
- concepts/template_applications.md
- Components:
- Components: concepts#components
- concepts/langgraph_server.md
- concepts/langgraph_studio.md
- concepts/langgraph_cli.md
- concepts/sdk.md
- how-tos/use-remote-graph.md
- LangGraph Server:
- LangGraph Server: concepts#langgraph-server
- concepts/application_structure.md
- concepts/assistants.md
- concepts/double_texting.md
- concepts/auth.md
- Deployment Options:
- Deployment Options: concepts#deployment-options
- concepts/self_hosted.md
- concepts/langgraph_cloud.md
- concepts/bring_your_own_cloud.md
- Tutorials:
- tutorials/index.md
- Quick Start:
- Quick Start: tutorials#quick-start
- tutorials/introduction.ipynb
- tutorials/langgraph-platform/local-server.md
- cloud/quick_start.md
- Chatbots:
- Chatbots: tutorials#chatbots
- tutorials/customer-support/customer-support.ipynb
- tutorials/chatbots/information-gather-prompting.ipynb
- tutorials/code_assistant/langgraph_code_assistant.ipynb
- RAG:
- RAG: tutorials#rag
- tutorials/rag/langgraph_adaptive_rag.ipynb
- tutorials/rag/langgraph_adaptive_rag_local.ipynb
- tutorials/rag/langgraph_agentic_rag.ipynb
- tutorials/rag/langgraph_crag.ipynb
- tutorials/rag/langgraph_crag_local.ipynb
- tutorials/rag/langgraph_self_rag.ipynb
- tutorials/rag/langgraph_self_rag_local.ipynb
- tutorials/sql-agent.ipynb
- Agent Architectures:
- Agent Architectures: tutorials#agent-architectures
- Multi-Agent Systems:
- Multi-Agent Systems: tutorials#multi-agent-systems
- tutorials/multi_agent/multi-agent-collaboration.ipynb
- tutorials/multi_agent/agent_supervisor.ipynb
- tutorials/multi_agent/hierarchical_agent_teams.ipynb
- Planning Agents:
- Planning Agents: tutorials#planning-agents
- tutorials/plan-and-execute/plan-and-execute.ipynb
- tutorials/rewoo/rewoo.ipynb
- tutorials/llm-compiler/LLMCompiler.ipynb
- Reflection & Critique:
- Reflection & Critique: tutorials#reflection-critique
- tutorials/reflection/reflection.ipynb
- tutorials/reflexion/reflexion.ipynb
- tutorials/tot/tot.ipynb
- tutorials/lats/lats.ipynb
- tutorials/self-discover/self-discover.ipynb
- Evaluation & Analysis:
- Evaluation & Analysis: tutorials#evaluation
- tutorials/chatbot-simulation-evaluation/agent-simulation-evaluation.ipynb
- tutorials/chatbot-simulation-evaluation/langsmith-agent-simulation-evaluation.ipynb
- Experimental:
- Experimental: tutorials#experimental
- tutorials/storm/storm.ipynb
- tutorials/tnt-llm/tnt-llm.ipynb
- tutorials/web-navigation/web_voyager.ipynb
- tutorials/usaco/usaco.ipynb
- tutorials/extraction/retries.ipynb
- LangGraph Platform:
- LangGraph Platform: concepts#langgraph-platform
- tutorials/auth/getting_started.md
- tutorials/auth/resource_auth.md
- tutorials/auth/add_auth_server.md
- Resources:
- FAQ: concepts/faq.md
- Troubleshooting:
- Troubleshooting: how-tos#troubleshooting
- Troubleshooting: troubleshooting/errors/index.md
- troubleshooting/errors/index.md
- troubleshooting/errors/GRAPH_RECURSION_LIMIT.md
- troubleshooting/errors/INVALID_CONCURRENT_GRAPH_UPDATE.md
- troubleshooting/errors/INVALID_GRAPH_NODE_RETURN_VALUE.md
- troubleshooting/errors/MULTIPLE_SUBGRAPHS.md
- Conceptual Guides:
- concepts/index.md
- LangGraph:
- LangGraph: concepts#langgraph
- concepts/high_level.md
- concepts/low_level.md
- concepts/agentic_concepts.md
- concepts/multi_agent.md
- concepts/human_in_the_loop.md
- concepts/persistence.md
- concepts/memory.md
- concepts/streaming.md
- concepts/faq.md
- LangGraph Platform:
- LangGraph Platform: concepts#langgraph-platform
- High Level:
- High Level: concepts#high-level
- concepts/langgraph_platform.md
- concepts/deployment_options.md
- concepts/plans.md
- concepts/template_applications.md
- Components:
- Components: concepts#components
- concepts/langgraph_server.md
- concepts/langgraph_studio.md
- concepts/langgraph_cli.md
- concepts/sdk.md
- how-tos/use-remote-graph.md
- LangGraph Server:
- LangGraph Server: concepts#langgraph-server
- concepts/application_structure.md
- concepts/assistants.md
- concepts/double_texting.md
- Deployment Options:
- Deployment Options: concepts#deployment-options
- concepts/self_hosted.md
- concepts/langgraph_cloud.md
- concepts/bring_your_own_cloud.md
- Reference:
- "reference/index.md"
- Library:
- Graphs: reference/graphs.md
- Checkpointing: reference/checkpoints.md
- Storage: reference/store.md
- Prebuilt Components: reference/prebuilt.md
- Channels: reference/channels.md
- Errors: reference/errors.md
- Types: reference/types.md
- Constants: reference/constants.md
- LangGraph Platform:
- Server API: "cloud/reference/api/api_ref.md"
- CLI: "cloud/reference/cli.md"
- SDK (Python): "cloud/reference/sdk/python_sdk_ref.md"
- SDK (JS/TS): "cloud/reference/sdk/js_ts_sdk_ref.md"
- RemoteGraph: reference/remote_graph.md
- Environment Variables: "cloud/reference/env_var.md"
- troubleshooting/errors/INVALID_CHAT_HISTORY.md
- LangGraph Academy Course: https://academy.langchain.com/courses/intro-to-langgraph
- API reference:
- Library:
- Graphs: reference/graphs.md
- Checkpointing: reference/checkpoints.md
- Storage: reference/store.md
- Prebuilt components: reference/prebuilt.md
- Channels: reference/channels.md
- Errors: reference/errors.md
- Types: reference/types.md
- Constants: reference/constants.md
- LangGraph Platform:
- Server API: "cloud/reference/api/api_ref.md"
- CLI: "cloud/reference/cli.md"
- SDK (Python): "cloud/reference/sdk/python_sdk_ref.md"
- SDK (JS/TS): "cloud/reference/sdk/js_ts_sdk_ref.md"
- RemoteGraph: reference/remote_graph.md
- Environment variables: "cloud/reference/env_var.md"
markdown_extensions:
- abbr
+10 -11
View File
@@ -34,17 +34,6 @@
color: #1E88E5;
}
.md-sidebar {
display: none;
}
/* Show sidebar on mobile */
@media screen and (max-width: 1220px) {
.md-sidebar--primary {
display: block;
}
}
.md-typeset a:hover {
color: #1565C0;
}
@@ -169,6 +158,16 @@
background-color: #CFC9FA;
color: #000000;
}
/* control the navbar depth */
[data-md-level="2"] .md-nav {
display: none;
}
/* disable the collapse/expand icon in the navar */
.md-nav__icon {
display: none;
}
</style>
{% endblock %}
-35
View File
@@ -1,35 +0,0 @@
.PHONY: test test_watch lint format
######################
# TESTING AND COVERAGE
######################
test:
poetry run pytest tests
test_watch:
poetry run ptw .
######################
# LINTING AND FORMATTING
######################
# Define a variable for Python and notebook files.
PYTHON_FILES=.
MYPY_CACHE=.mypy_cache
lint format: PYTHON_FILES=.
lint_diff format_diff: PYTHON_FILES=$(shell git diff --name-only --relative --diff-filter=d main . | grep -E '\.py$$|\.ipynb$$')
lint_package: PYTHON_FILES=langgraph
lint_tests: PYTHON_FILES=tests
lint_tests: MYPY_CACHE=.mypy_cache_test
lint lint_diff lint_package lint_tests:
poetry run ruff check .
[ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff
[ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES)
[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE)
[ "$(PYTHON_FILES)" = "" ] || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
format format_diff:
poetry run ruff format $(PYTHON_FILES)
poetry run ruff check --select I --fix $(PYTHON_FILES)
-95
View File
@@ -1,95 +0,0 @@
# LangGraph Checkpoint DuckDB
Implementation of LangGraph CheckpointSaver that uses DuckDB.
## Usage
> [!IMPORTANT]
> When using DuckDB checkpointers for the first time, make sure to call `.setup()` method on them to create required tables. See example below.
```python
from langgraph.checkpoint.duckdb import DuckDBSaver
write_config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
read_config = {"configurable": {"thread_id": "1"}}
with DuckDBSaver.from_conn_string(":memory:") as checkpointer:
# call .setup() the first time you're using the checkpointer
checkpointer.setup()
checkpoint = {
"v": 1,
"ts": "2024-07-31T20:14:19.804150+00:00",
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
"channel_values": {
"my_key": "meow",
"node": "node"
},
"channel_versions": {
"__start__": 2,
"my_key": 3,
"start:node": 3,
"node": 3
},
"versions_seen": {
"__input__": {},
"__start__": {
"__start__": 1
},
"node": {
"start:node": 2
}
},
"pending_sends": [],
}
# store checkpoint
checkpointer.put(write_config, checkpoint, {}, {})
# load checkpoint
checkpointer.get(read_config)
# list checkpoints
list(checkpointer.list(read_config))
```
### Async
```python
from langgraph.checkpoint.duckdb.aio import AsyncDuckDBSaver
async with AsyncDuckDBSaver.from_conn_string(":memory:") as checkpointer:
checkpoint = {
"v": 1,
"ts": "2024-07-31T20:14:19.804150+00:00",
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
"channel_values": {
"my_key": "meow",
"node": "node"
},
"channel_versions": {
"__start__": 2,
"my_key": 3,
"start:node": 3,
"node": 3
},
"versions_seen": {
"__input__": {},
"__start__": {
"__start__": 1
},
"node": {
"start:node": 2
}
},
"pending_sends": [],
}
# store checkpoint
await checkpointer.aput(write_config, checkpoint, {}, {})
# load checkpoint
await checkpointer.aget(read_config)
# list checkpoints
[c async for c in checkpointer.alist(read_config)]
```
@@ -1,356 +0,0 @@
import threading
from contextlib import contextmanager
from typing import Any, Iterator, Optional, Sequence
from langchain_core.runnables import RunnableConfig
import duckdb
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
CheckpointTuple,
get_checkpoint_id,
)
from langgraph.checkpoint.duckdb.base import BaseDuckDBSaver
from langgraph.checkpoint.serde.base import SerializerProtocol
class DuckDBSaver(BaseDuckDBSaver):
lock: threading.Lock
def __init__(
self,
conn: duckdb.DuckDBPyConnection,
serde: Optional[SerializerProtocol] = None,
) -> None:
super().__init__(serde=serde)
self.conn = conn
self.lock = threading.Lock()
@classmethod
@contextmanager
def from_conn_string(cls, conn_string: str) -> Iterator["DuckDBSaver"]:
"""Create a new DuckDBSaver instance from a connection string.
Args:
conn_string (str): The DuckDB connection info string.
Returns:
DuckDBSaver: A new DuckDBSaver instance.
"""
with duckdb.connect(conn_string) as conn:
yield cls(conn)
def setup(self) -> None:
"""Set up the checkpoint database asynchronously.
This method creates the necessary tables in the DuckDB database if they don't
already exist and runs database migrations. It MUST be called directly by the user
the first time checkpointer is used.
"""
with self.lock, self.conn.cursor() as cur:
try:
row = cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
).fetchone()
if row is None:
version = -1
else:
version = row[0]
except duckdb.CatalogException:
version = -1
for v, migration in zip(
range(version + 1, len(self.MIGRATIONS)),
self.MIGRATIONS[version + 1 :],
):
cur.execute(migration)
cur.execute("INSERT INTO checkpoint_migrations (v) VALUES (?)", [v])
def list(
self,
config: Optional[RunnableConfig],
*,
filter: Optional[dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
"""List checkpoints from the database.
This method retrieves a list of checkpoint tuples from the DuckDB database based
on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
Args:
config (RunnableConfig): The config to use for listing the checkpoints.
filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. Defaults to None.
before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
Yields:
Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
Examples:
>>> from langgraph.checkpoint.duckdb import DuckDBSaver
>>> with DuckDBSaver.from_conn_string(":memory:") as memory:
... # Run a graph, then list the checkpoints
>>> config = {"configurable": {"thread_id": "1"}}
>>> checkpoints = list(memory.list(config, limit=2))
>>> print(checkpoints)
[CheckpointTuple(...), CheckpointTuple(...)]
>>> config = {"configurable": {"thread_id": "1"}}
>>> before = {"configurable": {"checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875"}}
>>> with DuckDBSaver.from_conn_string(":memory:") as memory:
... # Run a graph, then list the checkpoints
>>> checkpoints = list(memory.list(config, before=before))
>>> print(checkpoints)
[CheckpointTuple(...), ...]
"""
where, args = self._search_where(config, filter, before)
query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
if limit:
query += f" LIMIT {limit}"
# if we change this to use .stream() we need to make sure to close the cursor
with self._cursor() as cur:
cur.execute(query, args)
for value in cur.fetchall():
(
thread_id,
checkpoint,
checkpoint_ns,
checkpoint_id,
parent_checkpoint_id,
metadata,
channel_values,
pending_writes,
pending_sends,
) = value
yield CheckpointTuple(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
self._load_checkpoint(
checkpoint,
channel_values,
pending_sends,
),
self._load_metadata(metadata),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_checkpoint_id
else None
),
self._load_writes(pending_writes),
)
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Get a checkpoint tuple from the database.
This method retrieves a checkpoint tuple from the DuckDB database based on the
provided config. If the config contains a "checkpoint_id" key, the checkpoint with
the matching thread ID and timestamp is retrieved. Otherwise, the latest checkpoint
for the given thread ID is retrieved.
Args:
config (RunnableConfig): The config to use for retrieving the checkpoint.
Returns:
Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
Examples:
Basic:
>>> config = {"configurable": {"thread_id": "1"}}
>>> checkpoint_tuple = memory.get_tuple(config)
>>> print(checkpoint_tuple)
CheckpointTuple(...)
With timestamp:
>>> config = {
... "configurable": {
... "thread_id": "1",
... "checkpoint_ns": "",
... "checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875",
... }
... }
>>> checkpoint_tuple = memory.get_tuple(config)
>>> print(checkpoint_tuple)
CheckpointTuple(...)
""" # noqa
thread_id = config["configurable"]["thread_id"]
checkpoint_id = get_checkpoint_id(config)
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
if checkpoint_id:
args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
where = "WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?"
else:
args = (thread_id, checkpoint_ns)
where = "WHERE thread_id = ? AND checkpoint_ns = ? ORDER BY checkpoint_id DESC LIMIT 1"
with self._cursor() as cur:
cur.execute(
self.SELECT_SQL + where,
args,
)
value = cur.fetchone()
if value:
(
thread_id,
checkpoint,
checkpoint_ns,
checkpoint_id,
parent_checkpoint_id,
metadata,
channel_values,
pending_writes,
pending_sends,
) = value
return CheckpointTuple(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
self._load_checkpoint(
checkpoint,
channel_values,
pending_sends,
),
self._load_metadata(metadata),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_checkpoint_id
else None
),
self._load_writes(pending_writes),
)
def put(
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
new_versions: ChannelVersions,
) -> RunnableConfig:
"""Save a checkpoint to the database.
This method saves a checkpoint to the DuckDB database. The checkpoint is associated
with the provided config and its parent config (if any).
Args:
config (RunnableConfig): The config to associate with the checkpoint.
checkpoint (Checkpoint): The checkpoint to save.
metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
new_versions (ChannelVersions): New channel versions as of this write.
Returns:
RunnableConfig: Updated configuration after storing the checkpoint.
Examples:
>>> from langgraph.checkpoint.duckdb import DuckDBSaver
>>> with DuckDBSaver.from_conn_string(":memory:") as memory:
>>> config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
>>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}}
>>> saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {})
>>> print(saved_config)
{'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}}
"""
configurable = config["configurable"].copy()
thread_id = configurable.pop("thread_id")
checkpoint_ns = configurable.pop("checkpoint_ns")
checkpoint_id = configurable.pop(
"checkpoint_id", configurable.pop("thread_ts", None)
)
copy = checkpoint.copy()
next_config = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint["id"],
}
}
checkpoint_blobs = self._dump_blobs(
thread_id,
checkpoint_ns,
copy.pop("channel_values"), # type: ignore[misc]
new_versions,
)
with self._cursor() as cur:
if checkpoint_blobs:
cur.executemany(self.UPSERT_CHECKPOINT_BLOBS_SQL, checkpoint_blobs)
cur.execute(
self.UPSERT_CHECKPOINTS_SQL,
(
thread_id,
checkpoint_ns,
checkpoint["id"],
checkpoint_id,
self._dump_checkpoint(copy),
self._dump_metadata(metadata),
),
)
return next_config
def put_writes(
self,
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
) -> None:
"""Store intermediate writes linked to a checkpoint.
This method saves intermediate writes associated with a checkpoint to the DuckDB database.
Args:
config (RunnableConfig): Configuration of the related checkpoint.
writes (List[Tuple[str, Any]]): List of writes to store.
task_id (str): Identifier for the task creating the writes.
"""
query = (
self.UPSERT_CHECKPOINT_WRITES_SQL
if all(w[0] in WRITES_IDX_MAP for w in writes)
else self.INSERT_CHECKPOINT_WRITES_SQL
)
with self._cursor() as cur:
cur.executemany(
query,
self._dump_writes(
config["configurable"]["thread_id"],
config["configurable"]["checkpoint_ns"],
config["configurable"]["checkpoint_id"],
task_id,
writes,
),
)
@contextmanager
def _cursor(self) -> Iterator[duckdb.DuckDBPyConnection]:
with self.lock, self.conn.cursor() as cur:
yield cur
__all__ = ["DuckDBSaver", "Conn"]
@@ -1,443 +0,0 @@
import asyncio
from contextlib import asynccontextmanager
from typing import Any, AsyncIterator, Iterator, Optional, Sequence
from langchain_core.runnables import RunnableConfig
import duckdb
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
CheckpointTuple,
get_checkpoint_id,
)
from langgraph.checkpoint.duckdb.base import BaseDuckDBSaver
from langgraph.checkpoint.serde.base import SerializerProtocol
class AsyncDuckDBSaver(BaseDuckDBSaver):
lock: asyncio.Lock
def __init__(
self,
conn: duckdb.DuckDBPyConnection,
serde: Optional[SerializerProtocol] = None,
) -> None:
super().__init__(serde=serde)
self.conn = conn
self.lock = asyncio.Lock()
self.loop = asyncio.get_running_loop()
@classmethod
@asynccontextmanager
async def from_conn_string(
cls,
conn_string: str,
) -> AsyncIterator["AsyncDuckDBSaver"]:
"""Create a new AsyncDuckDBSaver instance from a connection string.
Args:
conn_string (str): The DuckDB connection info string.
Returns:
AsyncDuckDBSaver: A new AsyncDuckDBSaver instance.
"""
with duckdb.connect(conn_string) as conn:
yield cls(conn)
async def setup(self) -> None:
"""Set up the checkpoint database asynchronously.
This method creates the necessary tables in the DuckDB database if they don't
already exist and runs database migrations. It MUST be called directly by the user
the first time checkpointer is used.
"""
async with self.lock:
with self.conn.cursor() as cur:
try:
await asyncio.to_thread(
cur.execute,
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1",
)
row = await asyncio.to_thread(cur.fetchone)
if row is None:
version = -1
else:
version = row[0]
except duckdb.CatalogException:
version = -1
for v, migration in zip(
range(version + 1, len(self.MIGRATIONS)),
self.MIGRATIONS[version + 1 :],
):
await asyncio.to_thread(cur.execute, migration)
await asyncio.to_thread(
cur.execute,
"INSERT INTO checkpoint_migrations (v) VALUES (?)",
[v],
)
async def alist(
self,
config: Optional[RunnableConfig],
*,
filter: Optional[dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
"""List checkpoints from the database asynchronously.
This method retrieves a list of checkpoint tuples from the DuckDB database based
on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
Args:
config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
limit (Optional[int]): Maximum number of checkpoints to return.
Yields:
AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.
"""
where, args = self._search_where(config, filter, before)
query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
if limit:
query += f" LIMIT {limit}"
# if we change this to use .stream() we need to make sure to close the cursor
async with self._cursor() as cur:
await asyncio.to_thread(cur.execute, query, args)
results = await asyncio.to_thread(cur.fetchall)
for value in results:
(
thread_id,
checkpoint,
checkpoint_ns,
checkpoint_id,
parent_checkpoint_id,
metadata,
channel_values,
pending_writes,
pending_sends,
) = value
yield CheckpointTuple(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
await asyncio.to_thread(
self._load_checkpoint,
checkpoint,
channel_values,
pending_sends,
),
self._load_metadata(metadata),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_checkpoint_id
else None
),
await asyncio.to_thread(self._load_writes, pending_writes),
)
async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Get a checkpoint tuple from the database asynchronously.
This method retrieves a checkpoint tuple from the DuckDBdatabase based on the
provided config. If the config contains a "checkpoint_id" key, the checkpoint with
the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
for the given thread ID is retrieved.
Args:
config (RunnableConfig): The config to use for retrieving the checkpoint.
Returns:
Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
"""
thread_id = config["configurable"]["thread_id"]
checkpoint_id = get_checkpoint_id(config)
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
if checkpoint_id:
args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
where = "WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?"
else:
args = (thread_id, checkpoint_ns)
where = "WHERE thread_id = ? AND checkpoint_ns = ? ORDER BY checkpoint_id DESC LIMIT 1"
async with self._cursor() as cur:
await asyncio.to_thread(
cur.execute,
self.SELECT_SQL + where,
args,
)
value = await asyncio.to_thread(cur.fetchone)
if value:
(
thread_id,
checkpoint,
checkpoint_ns,
checkpoint_id,
parent_checkpoint_id,
metadata,
channel_values,
pending_writes,
pending_sends,
) = value
return CheckpointTuple(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
await asyncio.to_thread(
self._load_checkpoint,
checkpoint,
channel_values,
pending_sends,
),
self._load_metadata(metadata),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_checkpoint_id
else None
),
await asyncio.to_thread(self._load_writes, pending_writes),
)
async def aput(
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
new_versions: ChannelVersions,
) -> RunnableConfig:
"""Save a checkpoint to the database asynchronously.
This method saves a checkpoint to the DuckDB database. The checkpoint is associated
with the provided config and its parent config (if any).
Args:
config (RunnableConfig): The config to associate with the checkpoint.
checkpoint (Checkpoint): The checkpoint to save.
metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
new_versions (ChannelVersions): New channel versions as of this write.
Returns:
RunnableConfig: Updated configuration after storing the checkpoint.
"""
configurable = config["configurable"].copy()
thread_id = configurable.pop("thread_id")
checkpoint_ns = configurable.pop("checkpoint_ns")
checkpoint_id = configurable.pop(
"checkpoint_id", configurable.pop("thread_ts", None)
)
copy = checkpoint.copy()
next_config = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint["id"],
}
}
checkpoint_blobs = await asyncio.to_thread(
self._dump_blobs,
thread_id,
checkpoint_ns,
copy.pop("channel_values"), # type: ignore[misc]
new_versions,
)
async with self._cursor() as cur:
if checkpoint_blobs:
await asyncio.to_thread(
cur.executemany, self.UPSERT_CHECKPOINT_BLOBS_SQL, checkpoint_blobs
)
await asyncio.to_thread(
cur.execute,
self.UPSERT_CHECKPOINTS_SQL,
(
thread_id,
checkpoint_ns,
checkpoint["id"],
checkpoint_id,
self._dump_checkpoint(copy),
self._dump_metadata(metadata),
),
)
return next_config
async def aput_writes(
self,
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
) -> None:
"""Store intermediate writes linked to a checkpoint asynchronously.
This method saves intermediate writes associated with a checkpoint to the database.
Args:
config (RunnableConfig): Configuration of the related checkpoint.
writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
"""
query = (
self.UPSERT_CHECKPOINT_WRITES_SQL
if all(w[0] in WRITES_IDX_MAP for w in writes)
else self.INSERT_CHECKPOINT_WRITES_SQL
)
params = await asyncio.to_thread(
self._dump_writes,
config["configurable"]["thread_id"],
config["configurable"]["checkpoint_ns"],
config["configurable"]["checkpoint_id"],
task_id,
writes,
)
async with self._cursor() as cur:
await asyncio.to_thread(cur.executemany, query, params)
@asynccontextmanager
async def _cursor(self) -> AsyncIterator[duckdb.DuckDBPyConnection]:
async with self.lock:
with self.conn.cursor() as cur:
yield cur
def list(
self,
config: Optional[RunnableConfig],
*,
filter: Optional[dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
"""List checkpoints from the database.
This method retrieves a list of checkpoint tuples from the DuckDB database based
on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
Args:
config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
limit (Optional[int]): Maximum number of checkpoints to return.
Yields:
Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples.
"""
try:
# check if we are in the main thread, only bg threads can block
# we don't check in other methods to avoid the overhead
if asyncio.get_running_loop() is self.loop:
raise asyncio.InvalidStateError(
"Synchronous calls to AsyncDuckDBSaver are only allowed from a "
"different thread. From the main thread, use the async interface. "
"For example, use `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_),
self.loop,
).result()
except StopAsyncIteration:
break
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Get a checkpoint tuple from the database.
This method retrieves a checkpoint tuple from the DuckDB database based on the
provided config. If the config contains a "checkpoint_id" key, the checkpoint with
the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
for the given thread ID is retrieved.
Args:
config (RunnableConfig): The config to use for retrieving the checkpoint.
Returns:
Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
"""
try:
# check if we are in the main thread, only bg threads can block
# we don't check in other methods to avoid the overhead
if asyncio.get_running_loop() is self.loop:
raise asyncio.InvalidStateError(
"Synchronous calls to AsyncDuckDBSaver are only allowed from a "
"different thread. From the main thread, use the async interface."
"For example, use `await checkpointer.aget_tuple(...)` or `await "
"graph.ainvoke(...)`."
)
except RuntimeError:
pass
return asyncio.run_coroutine_threadsafe(
self.aget_tuple(config), self.loop
).result()
def put(
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
new_versions: ChannelVersions,
) -> RunnableConfig:
"""Save a checkpoint to the database.
This method saves a checkpoint to the DuckDB database. The checkpoint is associated
with the provided config and its parent config (if any).
Args:
config (RunnableConfig): The config to associate with the checkpoint.
checkpoint (Checkpoint): The checkpoint to save.
metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
new_versions (ChannelVersions): New channel versions as of this write.
Returns:
RunnableConfig: Updated configuration after storing the checkpoint.
"""
return asyncio.run_coroutine_threadsafe(
self.aput(config, checkpoint, metadata, new_versions), self.loop
).result()
def put_writes(
self,
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
) -> None:
"""Store intermediate writes linked to a checkpoint.
This method saves intermediate writes associated with a checkpoint to the database.
Args:
config (RunnableConfig): Configuration of the related checkpoint.
writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
"""
return asyncio.run_coroutine_threadsafe(
self.aput_writes(config, writes, task_id), self.loop
).result()
@@ -1,290 +0,0 @@
import json
import random
from typing import Any, List, Optional, Sequence, Tuple, cast
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
get_checkpoint_id,
)
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.serde.types import TASKS, ChannelProtocol
MetadataInput = Optional[dict[str, Any]]
"""
To add a new migration, add a new string to the MIGRATIONS list.
The position of the migration in the list is the version number.
"""
MIGRATIONS = [
"""CREATE TABLE IF NOT EXISTS checkpoint_migrations (
v INTEGER PRIMARY KEY
);""",
"""CREATE TABLE IF NOT EXISTS checkpoints (
thread_id TEXT NOT NULL,
checkpoint_ns TEXT NOT NULL DEFAULT '',
checkpoint_id TEXT NOT NULL,
parent_checkpoint_id TEXT,
type TEXT,
checkpoint JSON NOT NULL,
metadata JSON NOT NULL DEFAULT '{}',
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
);""",
"""CREATE TABLE IF NOT EXISTS checkpoint_blobs (
thread_id TEXT NOT NULL,
checkpoint_ns TEXT NOT NULL DEFAULT '',
channel TEXT NOT NULL,
version TEXT NOT NULL,
type TEXT NOT NULL,
blob BLOB,
PRIMARY KEY (thread_id, checkpoint_ns, channel, version)
);""",
"""CREATE TABLE IF NOT EXISTS checkpoint_writes (
thread_id TEXT NOT NULL,
checkpoint_ns TEXT NOT NULL DEFAULT '',
checkpoint_id TEXT NOT NULL,
task_id TEXT NOT NULL,
idx INTEGER NOT NULL,
channel TEXT NOT NULL,
type TEXT,
blob BLOB NOT NULL,
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
);""",
]
SELECT_SQL = f"""
select
thread_id,
checkpoint,
checkpoint_ns,
checkpoint_id,
parent_checkpoint_id,
metadata,
(
select array_agg(array[bl.channel::bytea, bl.type::bytea, bl.blob])
from (
SELECT unnest(json_keys(json_extract(checkpoint, '$.channel_versions'))) as key
) cv
inner join checkpoint_blobs bl
on bl.thread_id = checkpoints.thread_id
and bl.checkpoint_ns = checkpoints.checkpoint_ns
and bl.channel = cv.key
and bl.version = json_extract_string(checkpoint, '$.channel_versions.' || cv.key)
) as channel_values,
(
select
array_agg(array[cw.task_id::blob, cw.channel::blob, cw.type::blob, cw.blob])
from checkpoint_writes cw
where cw.thread_id = checkpoints.thread_id
and cw.checkpoint_ns = checkpoints.checkpoint_ns
and cw.checkpoint_id = checkpoints.checkpoint_id
) as pending_writes,
(
select array_agg(array[cw.type::blob, cw.blob])
from checkpoint_writes cw
where cw.thread_id = checkpoints.thread_id
and cw.checkpoint_ns = checkpoints.checkpoint_ns
and cw.checkpoint_id = checkpoints.parent_checkpoint_id
and cw.channel = '{TASKS}'
) as pending_sends
from checkpoints """
UPSERT_CHECKPOINT_BLOBS_SQL = """
INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob)
VALUES (?, ?, ?, ?, ?, ?)
ON CONFLICT (thread_id, checkpoint_ns, channel, version) DO NOTHING
"""
UPSERT_CHECKPOINTS_SQL = """
INSERT INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, checkpoint, metadata)
VALUES (?, ?, ?, ?, ?, ?)
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id)
DO UPDATE SET
checkpoint = EXCLUDED.checkpoint,
metadata = EXCLUDED.metadata;
"""
UPSERT_CHECKPOINT_WRITES_SQL = """
INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO UPDATE SET
channel = EXCLUDED.channel,
type = EXCLUDED.type,
blob = EXCLUDED.blob;
"""
INSERT_CHECKPOINT_WRITES_SQL = """
INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
"""
class BaseDuckDBSaver(BaseCheckpointSaver[str]):
SELECT_SQL = SELECT_SQL
MIGRATIONS = MIGRATIONS
UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL
jsonplus_serde = JsonPlusSerializer()
def _load_checkpoint(
self,
checkpoint_json_str: str,
channel_values: list[tuple[bytes, bytes, bytes]],
pending_sends: list[tuple[bytes, bytes]],
) -> Checkpoint:
checkpoint = json.loads(checkpoint_json_str)
return {
**checkpoint,
"pending_sends": [
self.serde.loads_typed((c.decode(), b)) for c, b in pending_sends or []
],
"channel_values": self._load_blobs(channel_values),
}
def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]:
return {**checkpoint, "pending_sends": []}
def _load_blobs(
self, blob_values: list[tuple[bytes, bytes, bytes]]
) -> dict[str, Any]:
if not blob_values:
return {}
return {
k.decode(): self.serde.loads_typed((t.decode(), v))
for k, t, v in blob_values
if t.decode() != "empty"
}
def _dump_blobs(
self,
thread_id: str,
checkpoint_ns: str,
values: dict[str, Any],
versions: ChannelVersions,
) -> list[tuple[str, str, str, str, str, Optional[bytes]]]:
if not versions:
return []
return [
(
thread_id,
checkpoint_ns,
k,
cast(str, ver),
*(
self.serde.dumps_typed(values[k])
if k in values
else ("empty", None)
),
)
for k, ver in versions.items()
]
def _load_writes(
self, writes: list[tuple[bytes, bytes, bytes, bytes]]
) -> list[tuple[str, str, Any]]:
return (
[
(
tid.decode(),
channel.decode(),
self.serde.loads_typed((t.decode(), v)),
)
for tid, channel, t, v in writes
]
if writes
else []
)
def _dump_writes(
self,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
task_id: str,
writes: Sequence[tuple[str, Any]],
) -> list[tuple[str, str, str, str, int, str, str, bytes]]:
return [
(
thread_id,
checkpoint_ns,
checkpoint_id,
task_id,
WRITES_IDX_MAP.get(channel, idx),
channel,
*self.serde.dumps_typed(value),
)
for idx, (channel, value) in enumerate(writes)
]
def _load_metadata(self, metadata_json_str: str) -> CheckpointMetadata:
return self.jsonplus_serde.loads(metadata_json_str.encode())
def _dump_metadata(self, metadata: CheckpointMetadata) -> str:
serialized_metadata = self.jsonplus_serde.dumps(metadata)
# NOTE: we're using JSON serializer (not msgpack), so we need to remove null characters before writing
return serialized_metadata.decode().replace("\\u0000", "")
def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
if current is None:
current_v = 0
elif isinstance(current, int):
current_v = current
else:
current_v = int(current.split(".")[0])
next_v = current_v + 1
next_h = random.random()
return f"{next_v:032}.{next_h:016}"
def _search_where(
self,
config: Optional[RunnableConfig],
filter: MetadataInput,
before: Optional[RunnableConfig] = None,
) -> Tuple[str, List[Any]]:
"""Return WHERE clause predicates for alist() given config, filter, before.
This method returns a tuple of a string and a tuple of values. The string
is the parametered WHERE clause predicate (including the WHERE keyword):
"WHERE column1 = $1 AND column2 IS $2". The list of values contains the
values for each of the corresponding parameters.
"""
wheres = []
param_values = []
# construct predicate for config filter
if config:
wheres.append("thread_id = ?")
param_values.append(config["configurable"]["thread_id"])
checkpoint_ns = config["configurable"].get("checkpoint_ns")
if checkpoint_ns is not None:
wheres.append("checkpoint_ns = ?")
param_values.append(checkpoint_ns)
if checkpoint_id := get_checkpoint_id(config):
wheres.append("checkpoint_id = ?")
param_values.append(checkpoint_id)
# construct predicate for metadata filter
if filter:
wheres.append("json_contains(metadata, ?)")
param_values.append(json.dumps(filter))
# construct predicate for `before`
if before is not None:
wheres.append("checkpoint_id < ?")
param_values.append(get_checkpoint_id(before))
return (
"WHERE " + " AND ".join(wheres) if wheres else "",
param_values,
)
@@ -1,4 +0,0 @@
from langgraph.store.duckdb.aio import AsyncDuckDBStore
from langgraph.store.duckdb.base import DuckDBStore
__all__ = ["AsyncDuckDBStore", "DuckDBStore"]
@@ -1,195 +0,0 @@
import asyncio
import logging
from contextlib import asynccontextmanager
from typing import (
AsyncIterator,
Iterable,
Sequence,
cast,
)
import duckdb
from langgraph.store.base import GetOp, ListNamespacesOp, Op, PutOp, Result, SearchOp
from langgraph.store.base.batch import AsyncBatchedBaseStore
from langgraph.store.duckdb.base import (
BaseDuckDBStore,
_convert_ns,
_group_ops,
_row_to_item,
)
logger = logging.getLogger(__name__)
class AsyncDuckDBStore(AsyncBatchedBaseStore, BaseDuckDBStore):
def __init__(
self,
conn: duckdb.DuckDBPyConnection,
) -> None:
super().__init__()
self.conn = conn
self.loop = asyncio.get_running_loop()
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
grouped_ops, num_ops = _group_ops(ops)
results: list[Result] = [None] * num_ops
tasks = []
if GetOp in grouped_ops:
tasks.append(
self._batch_get_ops(
cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results
)
)
if PutOp in grouped_ops:
tasks.append(
self._batch_put_ops(
cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp])
)
)
if SearchOp in grouped_ops:
tasks.append(
self._batch_search_ops(
cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
results,
)
)
if ListNamespacesOp in grouped_ops:
tasks.append(
self._batch_list_namespaces_ops(
cast(
Sequence[tuple[int, ListNamespacesOp]],
grouped_ops[ListNamespacesOp],
),
results,
)
)
await asyncio.gather(*tasks)
return results
def batch(self, ops: Iterable[Op]) -> list[Result]:
return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
async def _batch_get_ops(
self,
get_ops: Sequence[tuple[int, GetOp]],
results: list[Result],
) -> None:
cursors = []
for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
cur = self.conn.cursor()
await asyncio.to_thread(cur.execute, query, params)
cursors.append((cur, namespace, items))
for cur, namespace, items in cursors:
rows = await asyncio.to_thread(cur.fetchall)
key_to_row = {row[1]: row for row in rows}
for idx, key in items:
row = key_to_row.get(key)
if row:
results[idx] = _row_to_item(namespace, row)
else:
results[idx] = None
async def _batch_put_ops(
self,
put_ops: Sequence[tuple[int, PutOp]],
) -> None:
queries = self._get_batch_PUT_queries(put_ops)
for query, params in queries:
cur = self.conn.cursor()
await asyncio.to_thread(cur.execute, query, params)
async def _batch_search_ops(
self,
search_ops: Sequence[tuple[int, SearchOp]],
results: list[Result],
) -> None:
queries = self._get_batch_search_queries(search_ops)
cursors: list[tuple[duckdb.DuckDBPyConnection, int]] = []
for (query, params), (idx, _) in zip(queries, search_ops):
cur = self.conn.cursor()
await asyncio.to_thread(cur.execute, query, params)
cursors.append((cur, idx))
for cur, idx in cursors:
rows = await asyncio.to_thread(cur.fetchall)
items = [_row_to_item(_convert_ns(row[0]), row) for row in rows]
results[idx] = items
async def _batch_list_namespaces_ops(
self,
list_ops: Sequence[tuple[int, ListNamespacesOp]],
results: list[Result],
) -> None:
queries = self._get_batch_list_namespaces_queries(list_ops)
cursors: list[tuple[duckdb.DuckDBPyConnection, int]] = []
for (query, params), (idx, _) in zip(queries, list_ops):
cur = self.conn.cursor()
await asyncio.to_thread(cur.execute, query, params)
cursors.append((cur, idx))
for cur, idx in cursors:
rows = cast(list[tuple], await asyncio.to_thread(cur.fetchall))
namespaces = [_convert_ns(row[0]) for row in rows]
results[idx] = namespaces
@classmethod
@asynccontextmanager
async def from_conn_string(
cls,
conn_string: str,
) -> AsyncIterator["AsyncDuckDBStore"]:
"""Create a new AsyncDuckDBStore instance from a connection string.
Args:
conn_string (str): The DuckDB connection info string.
Returns:
AsyncDuckDBStore: A new AsyncDuckDBStore instance.
"""
with duckdb.connect(conn_string) as conn:
yield cls(conn)
async def setup(self) -> None:
"""Set up the store database asynchronously.
This method creates the necessary tables in the DuckDB database if they don't
already exist and runs database migrations. It is called automatically when needed and should not be called
directly by the user.
"""
cur = self.conn.cursor()
try:
await asyncio.to_thread(
cur.execute, "SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
)
row = await asyncio.to_thread(cur.fetchone)
if row is None:
version = -1
else:
version = row[0]
except duckdb.CatalogException:
version = -1
# Create store_migrations table if it doesn't exist
await asyncio.to_thread(
cur.execute,
"""
CREATE TABLE IF NOT EXISTS store_migrations (
v INTEGER PRIMARY KEY
)
""",
)
for v, migration in enumerate(
self.MIGRATIONS[version + 1 :], start=version + 1
):
await asyncio.to_thread(cur.execute, migration)
await asyncio.to_thread(
cur.execute, "INSERT INTO store_migrations (v) VALUES (?)", (v,)
)
@@ -1,408 +0,0 @@
import asyncio
import json
import logging
from collections import defaultdict
from contextlib import contextmanager
from typing import (
Any,
Generic,
Iterable,
Iterator,
Sequence,
TypeVar,
Union,
cast,
)
import duckdb
from langgraph.store.base import (
BaseStore,
GetOp,
Item,
ListNamespacesOp,
Op,
PutOp,
Result,
SearchItem,
SearchOp,
)
logger = logging.getLogger(__name__)
MIGRATIONS = [
"""
CREATE TABLE IF NOT EXISTS store (
prefix TEXT NOT NULL,
key TEXT NOT NULL,
value JSON NOT NULL,
created_at TIMESTAMP DEFAULT now(),
updated_at TIMESTAMP DEFAULT now(),
PRIMARY KEY (prefix, key)
);
""",
"""
CREATE INDEX IF NOT EXISTS store_prefix_idx ON store (prefix);
""",
]
C = TypeVar("C", bound=duckdb.DuckDBPyConnection)
class BaseDuckDBStore(Generic[C]):
MIGRATIONS = MIGRATIONS
conn: C
def _get_batch_GET_ops_queries(
self,
get_ops: Sequence[tuple[int, GetOp]],
) -> list[tuple[str, tuple, tuple[str, ...], list]]:
namespace_groups = defaultdict(list)
for idx, op in get_ops:
namespace_groups[op.namespace].append((idx, op.key))
results = []
for namespace, items in namespace_groups.items():
_, keys = zip(*items)
keys_to_query = ",".join(["?"] * len(keys))
query = f"""
SELECT prefix, key, value, created_at, updated_at
FROM store
WHERE prefix = ? AND key IN ({keys_to_query})
"""
params = (_namespace_to_text(namespace), *keys)
results.append((query, params, namespace, items))
return results
def _get_batch_PUT_queries(
self,
put_ops: Sequence[tuple[int, PutOp]],
) -> list[tuple[str, Sequence]]:
inserts: list[PutOp] = []
deletes: list[PutOp] = []
for _, op in put_ops:
if op.value is None:
deletes.append(op)
else:
inserts.append(op)
queries: list[tuple[str, Sequence]] = []
if deletes:
namespace_groups: dict[tuple[str, ...], list[str]] = defaultdict(list)
for op in deletes:
namespace_groups[op.namespace].append(op.key)
for namespace, keys in namespace_groups.items():
placeholders = ",".join(["?"] * len(keys))
query = (
f"DELETE FROM store WHERE prefix = ? AND key IN ({placeholders})"
)
params = (_namespace_to_text(namespace), *keys)
queries.append((query, params))
if inserts:
values = []
insertion_params = []
for op in inserts:
values.append("(?, ?, ?, now(), now())")
insertion_params.extend(
[
_namespace_to_text(op.namespace),
op.key,
json.dumps(op.value),
]
)
values_str = ",".join(values)
query = f"""
INSERT INTO store (prefix, key, value, created_at, updated_at)
VALUES {values_str}
ON CONFLICT (prefix, key) DO UPDATE
SET value = EXCLUDED.value, updated_at = now()
"""
queries.append((query, insertion_params))
return queries
def _get_batch_search_queries(
self,
search_ops: Sequence[tuple[int, SearchOp]],
) -> list[tuple[str, Sequence]]:
queries: list[tuple[str, Sequence]] = []
for _, op in search_ops:
query = """
SELECT prefix, key, value, created_at, updated_at
FROM store
WHERE prefix LIKE ?
"""
params: list = [f"{_namespace_to_text(op.namespace_prefix)}%"]
if op.filter:
filter_conditions = []
for key, value in op.filter.items():
filter_conditions.append(f"json_extract(value, '$.{key}') = ?")
params.append(json.dumps(value))
query += " AND " + " AND ".join(filter_conditions)
query += " ORDER BY updated_at DESC LIMIT ? OFFSET ?"
params.extend([op.limit, op.offset])
queries.append((query, params))
return queries
def _get_batch_list_namespaces_queries(
self,
list_ops: Sequence[tuple[int, ListNamespacesOp]],
) -> list[tuple[str, Sequence]]:
queries: list[tuple[str, Sequence]] = []
for _, op in list_ops:
query = """
WITH split_prefix AS (
SELECT
prefix,
string_split(prefix, '.') AS parts
FROM store
)
SELECT DISTINCT ON (truncated_prefix)
CASE
WHEN ? IS NOT NULL THEN
array_to_string(array_slice(parts, 1, ?), '.')
ELSE prefix
END AS truncated_prefix,
prefix
FROM split_prefix
"""
params: list[Any] = [op.max_depth, op.max_depth]
conditions = []
if op.match_conditions:
for condition in op.match_conditions:
if condition.match_type == "prefix":
conditions.append("prefix LIKE ?")
params.append(
f"{_namespace_to_text(condition.path, handle_wildcards=True)}%"
)
elif condition.match_type == "suffix":
conditions.append("prefix LIKE ?")
params.append(
f"%{_namespace_to_text(condition.path, handle_wildcards=True)}"
)
else:
logger.warning(
f"Unknown match_type in list_namespaces: {condition.match_type}"
)
if conditions:
query += " WHERE " + " AND ".join(conditions)
query += " ORDER BY prefix LIMIT ? OFFSET ?"
params.extend([op.limit, op.offset])
queries.append((query, params))
return queries
class DuckDBStore(BaseStore, BaseDuckDBStore[duckdb.DuckDBPyConnection]):
def __init__(
self,
conn: duckdb.DuckDBPyConnection,
) -> None:
super().__init__()
self.conn = conn
def batch(self, ops: Iterable[Op]) -> list[Result]:
grouped_ops, num_ops = _group_ops(ops)
results: list[Result] = [None] * num_ops
if GetOp in grouped_ops:
self._batch_get_ops(
cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results
)
if PutOp in grouped_ops:
self._batch_put_ops(cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]))
if SearchOp in grouped_ops:
self._batch_search_ops(
cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
results,
)
if ListNamespacesOp in grouped_ops:
self._batch_list_namespaces_ops(
cast(
Sequence[tuple[int, ListNamespacesOp]],
grouped_ops[ListNamespacesOp],
),
results,
)
return results
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
return await asyncio.get_running_loop().run_in_executor(None, self.batch, ops)
def _batch_get_ops(
self,
get_ops: Sequence[tuple[int, GetOp]],
results: list[Result],
) -> None:
cursors = []
for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
cur = self.conn.cursor()
cur.execute(query, params)
cursors.append((cur, namespace, items))
for cur, namespace, items in cursors:
rows = cur.fetchall()
key_to_row = {row[1]: row for row in rows}
for idx, key in items:
row = key_to_row.get(key)
if row:
results[idx] = _row_to_item(namespace, row)
else:
results[idx] = None
def _batch_put_ops(
self,
put_ops: Sequence[tuple[int, PutOp]],
) -> None:
queries = self._get_batch_PUT_queries(put_ops)
for query, params in queries:
cur = self.conn.cursor()
cur.execute(query, params)
def _batch_search_ops(
self,
search_ops: Sequence[tuple[int, SearchOp]],
results: list[Result],
) -> None:
queries = self._get_batch_search_queries(search_ops)
cursors: list[tuple[duckdb.DuckDBPyConnection, int]] = []
for (query, params), (idx, _) in zip(queries, search_ops):
cur = self.conn.cursor()
cur.execute(query, params)
cursors.append((cur, idx))
for cur, idx in cursors:
rows = cur.fetchall()
items = [_row_to_search_item(_convert_ns(row[0]), row) for row in rows]
results[idx] = items
def _batch_list_namespaces_ops(
self,
list_ops: Sequence[tuple[int, ListNamespacesOp]],
results: list[Result],
) -> None:
queries = self._get_batch_list_namespaces_queries(list_ops)
cursors: list[tuple[duckdb.DuckDBPyConnection, int]] = []
for (query, params), (idx, _) in zip(queries, list_ops):
cur = self.conn.cursor()
cur.execute(query, params)
cursors.append((cur, idx))
for cur, idx in cursors:
rows = cast(list[dict], cur.fetchall())
namespaces = [_convert_ns(row[0]) for row in rows]
results[idx] = namespaces
@classmethod
@contextmanager
def from_conn_string(
cls,
conn_string: str,
) -> Iterator["DuckDBStore"]:
"""Create a new BaseDuckDBStore instance from a connection string.
Args:
conn_string (str): The DuckDB connection info string.
Returns:
DuckDBStore: A new DuckDBStore instance.
"""
with duckdb.connect(conn_string) as conn:
yield cls(conn=conn)
def setup(self) -> None:
"""Set up the store database.
This method creates the necessary tables in the DuckDB database if they don't
already exist and runs database migrations. It is called automatically when needed and should not be called
directly by the user.
"""
with self.conn.cursor() as cur:
try:
cur.execute("SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1")
row = cast(dict, cur.fetchone())
if row is None:
version = -1
else:
version = row["v"]
except duckdb.CatalogException:
version = -1
# Create store_migrations table if it doesn't exist
cur.execute(
"""
CREATE TABLE IF NOT EXISTS store_migrations (
v INTEGER PRIMARY KEY
)
"""
)
for v, migration in enumerate(
self.MIGRATIONS[version + 1 :], start=version + 1
):
cur.execute(migration)
cur.execute("INSERT INTO store_migrations (v) VALUES (?)", (v,))
def _namespace_to_text(
namespace: tuple[str, ...], handle_wildcards: bool = False
) -> str:
"""Convert namespace tuple to text string."""
if handle_wildcards:
namespace = tuple("%" if val == "*" else val for val in namespace)
return ".".join(namespace)
def _row_to_item(
namespace: tuple[str, ...],
row: tuple,
) -> Item:
"""Convert a row from the database into an Item."""
_, key, val, created_at, updated_at = row
return Item(
value=val if isinstance(val, dict) else json.loads(val),
key=key,
namespace=namespace,
created_at=created_at,
updated_at=updated_at,
)
def _row_to_search_item(
namespace: tuple[str, ...],
row: tuple,
) -> SearchItem:
"""Convert a row from the database into an SearchItem."""
# TODO: Add support for search
_, key, val, created_at, updated_at = row
return SearchItem(
value=val if isinstance(val, dict) else json.loads(val),
key=key,
namespace=namespace,
created_at=created_at,
updated_at=updated_at,
)
def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int]:
grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list)
tot = 0
for idx, op in enumerate(ops):
grouped_ops[type(op)].append((idx, op))
tot += 1
return grouped_ops, tot
def _convert_ns(namespace: Union[str, list]) -> tuple[str, ...]:
if isinstance(namespace, list):
return tuple(namespace)
return tuple(namespace.split("."))
-1058
View File
File diff suppressed because it is too large Load Diff
-60
View File
@@ -1,60 +0,0 @@
[tool.poetry]
name = "langgraph-checkpoint-duckdb"
version = "2.0.2"
description = "Library with a DuckDB implementation of LangGraph checkpoint saver."
authors = []
license = "MIT"
readme = "README.md"
repository = "https://www.github.com/langchain-ai/langgraph"
packages = [{ include = "langgraph" }]
[tool.poetry.dependencies]
python = "^3.9.0,<4.0"
langgraph-checkpoint = "^2.0.2"
duckdb = ">=1.1.2"
[tool.poetry.group.dev.dependencies]
ruff = "^0.6.2"
codespell = "^2.2.0"
pytest = "^7.2.1"
anyio = "^4.4.0"
pytest-asyncio = "^0.21.1"
pytest-mock = "^3.11.1"
pytest-watch = "^4.2.0"
mypy = "^1.10.0"
langgraph-checkpoint = {path = "../checkpoint", develop = true}
[tool.pytest.ini_options]
# --strict-markers will raise errors on unknown marks.
# https://docs.pytest.org/en/7.1.x/how-to/mark.html#raising-errors-on-unknown-marks
#
# https://docs.pytest.org/en/7.1.x/reference/reference.html
# --strict-config any warnings encountered while parsing the `pytest`
# section of the configuration file raise errors.
addopts = "--strict-markers --strict-config --durations=5 -vv"
asyncio_mode = "auto"
[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.ruff]
lint.select = [
"E", # pycodestyle
"F", # Pyflakes
"UP", # pyupgrade
"B", # flake8-bugbear
"I", # isort
]
lint.ignore = ["E501", "B008", "UP007", "UP006"]
[tool.mypy]
# https://mypy.readthedocs.io/en/stable/config_file.html
disallow_untyped_defs = "True"
explicit_package_bases = "True"
warn_no_return = "False"
warn_unused_ignores = "True"
warn_redundant_casts = "True"
allow_redefinition = "True"
disable_error_code = "typeddict-item, return-value"
-112
View File
@@ -1,112 +0,0 @@
from typing import Any
import pytest
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
Checkpoint,
CheckpointMetadata,
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.duckdb.aio import AsyncDuckDBSaver
class TestAsyncDuckDBSaver:
@pytest.fixture(autouse=True)
async def setup(self) -> None:
# objects for test setup
self.config_1: RunnableConfig = {
"configurable": {
"thread_id": "thread-1",
# for backwards compatibility testing
"thread_ts": "1",
"checkpoint_ns": "",
}
}
self.config_2: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2",
"checkpoint_ns": "",
}
}
self.config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
}
}
self.chkpnt_1: Checkpoint = empty_checkpoint()
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
self.chkpnt_3: Checkpoint = empty_checkpoint()
self.metadata_1: CheckpointMetadata = {
"source": "input",
"step": 2,
"writes": {},
"score": 1,
}
self.metadata_2: CheckpointMetadata = {
"source": "loop",
"step": 1,
"writes": {"foo": "bar"},
"score": None,
}
self.metadata_3: CheckpointMetadata = {}
async def test_asearch(self) -> None:
async with AsyncDuckDBSaver.from_conn_string(":memory:") as saver:
await saver.setup()
await saver.aput(self.config_1, self.chkpnt_1, self.metadata_1, {})
await saver.aput(self.config_2, self.chkpnt_2, self.metadata_2, {})
await saver.aput(self.config_3, self.chkpnt_3, self.metadata_3, {})
# call method / assertions
query_1 = {"source": "input"} # search by 1 key
query_2 = {
"step": 1,
"writes": {"foo": "bar"},
} # search by multiple keys
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
query_4 = {"source": "update", "step": 1} # no match
search_results_1 = [c async for c in saver.alist(None, filter=query_1)]
assert len(search_results_1) == 1
assert search_results_1[0].metadata == self.metadata_1
search_results_2 = [c async for c in saver.alist(None, filter=query_2)]
assert len(search_results_2) == 1
assert search_results_2[0].metadata == self.metadata_2
search_results_3 = [c async for c in saver.alist(None, filter=query_3)]
assert len(search_results_3) == 3
search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
assert len(search_results_4) == 0
# search by config (defaults to checkpoints across all namespaces)
search_results_5 = [
c
async for c in saver.alist({"configurable": {"thread_id": "thread-2"}})
]
assert len(search_results_5) == 2
assert {
search_results_5[0].config["configurable"]["checkpoint_ns"],
search_results_5[1].config["configurable"]["checkpoint_ns"],
} == {"", "inner"}
# TODO: test before and limit params
async def test_null_chars(self) -> None:
async with AsyncDuckDBSaver.from_conn_string(":memory:") as saver:
await saver.setup()
config = await saver.aput(
self.config_1, self.chkpnt_1, {"my_key": "\x00abc"}, {}
)
assert (await saver.aget_tuple(config)).metadata["my_key"] == "abc" # type: ignore
assert [c async for c in saver.alist(None, filter={"my_key": "abc"})][
0
].metadata["my_key"] == "abc"
@@ -1,517 +0,0 @@
# type: ignore
import uuid
from datetime import datetime
from typing import Any
from unittest.mock import MagicMock
import pytest
from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
from langgraph.store.duckdb import AsyncDuckDBStore
class MockCursor:
def __init__(self, fetch_result: Any) -> None:
self.fetch_result = fetch_result
self.execute = MagicMock()
self.fetchall = MagicMock(return_value=self.fetch_result)
class MockConnection:
def __init__(self) -> None:
self.cursor = MagicMock()
@pytest.fixture
def mock_connection() -> MockConnection:
return MockConnection()
@pytest.fixture
async def store(mock_connection: MockConnection) -> AsyncDuckDBStore:
duck_db_store = AsyncDuckDBStore(mock_connection)
await duck_db_store.setup()
return duck_db_store
async def test_abatch_order(store: AsyncDuckDBStore) -> None:
mock_connection = store.conn
mock_get_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
(
"test.bar",
"key2",
'{"data": "value2"}',
datetime.now(),
datetime.now(),
),
]
)
mock_search_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
]
)
mock_list_namespaces_cursor = MockCursor(
[
("test",),
]
)
failures = []
def cursor_side_effect() -> Any:
cursor = MagicMock()
def execute_side_effect(query: str, *params: Any) -> None:
# My super sophisticated database.
if "WHERE prefix = ? AND key" in query:
cursor.fetchall = mock_get_cursor.fetchall
elif "SELECT prefix, key, value" in query:
cursor.fetchall = mock_search_cursor.fetchall
elif "SELECT DISTINCT ON (truncated_prefix)" in query:
cursor.fetchall = mock_list_namespaces_cursor.fetchall
elif "INSERT INTO " in query:
pass
else:
e = ValueError(f"Unmatched query: {query}")
failures.append(e)
raise e
cursor.execute = MagicMock(side_effect=execute_side_effect)
return cursor
mock_connection.cursor.side_effect = cursor_side_effect # type: ignore
ops = [
GetOp(namespace=("test",), key="key1"),
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
SearchOp(
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
),
ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0),
GetOp(namespace=("test",), key="key3"),
]
results = await store.abatch(ops)
assert not failures
assert len(results) == 5
assert isinstance(results[0], Item)
assert isinstance(results[0].value, dict)
assert results[0].value == {"data": "value1"}
assert results[0].key == "key1"
assert results[1] is None
assert isinstance(results[2], list)
assert len(results[2]) == 1
assert isinstance(results[3], list)
assert results[3] == [("test",)]
assert results[4] is None
ops_reordered = [
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
GetOp(namespace=("test",), key="key2"),
ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0),
PutOp(namespace=("test",), key="key3", value={"data": "value3"}),
GetOp(namespace=("test",), key="key1"),
]
results_reordered = await store.abatch(ops_reordered)
assert not failures
assert len(results_reordered) == 5
assert isinstance(results_reordered[0], list)
assert len(results_reordered[0]) == 1
assert isinstance(results_reordered[1], Item)
assert results_reordered[1].value == {"data": "value2"}
assert results_reordered[1].key == "key2"
assert isinstance(results_reordered[2], list)
assert results_reordered[2] == [("test",)]
assert results_reordered[3] is None
assert isinstance(results_reordered[4], Item)
assert results_reordered[4].value == {"data": "value1"}
assert results_reordered[4].key == "key1"
async def test_batch_get_ops(store: AsyncDuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
(
"test.bar",
"key2",
'{"data": "value2"}',
datetime.now(),
datetime.now(),
),
]
)
mock_connection.cursor.return_value = mock_cursor
ops = [
GetOp(namespace=("test",), key="key1"),
GetOp(namespace=("test",), key="key2"),
GetOp(namespace=("test",), key="key3"),
]
results = await store.abatch(ops)
assert len(results) == 3
assert results[0] is not None
assert results[1] is not None
assert results[2] is None
assert results[0].key == "key1"
assert results[1].key == "key2"
async def test_batch_put_ops(store: AsyncDuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor([])
mock_connection.cursor.return_value = mock_cursor
ops = [
PutOp(namespace=("test",), key="key1", value={"data": "value1"}),
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
PutOp(namespace=("test",), key="key3", value=None),
]
results = await store.abatch(ops)
assert len(results) == 3
assert all(result is None for result in results)
assert mock_cursor.execute.call_count == 2
async def test_batch_search_ops(store: AsyncDuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
(
"test.bar",
"key2",
'{"data": "value2"}',
datetime.now(),
datetime.now(),
),
]
)
mock_connection.cursor.return_value = mock_cursor
ops = [
SearchOp(
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
),
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
]
results = await store.abatch(ops)
assert len(results) == 2
assert len(results[0]) == 2
assert len(results[1]) == 2
async def test_batch_list_namespaces_ops(store: AsyncDuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor([("test.namespace1",), ("test.namespace2",)])
mock_connection.cursor.return_value = mock_cursor
ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)]
results = await store.abatch(ops)
assert len(results) == 1
assert results[0] == [("test", "namespace1"), ("test", "namespace2")]
# The following use the actual DB connection
async def test_basic_store_ops() -> None:
async with AsyncDuckDBStore.from_conn_string(":memory:") as store:
await store.setup()
namespace = ("test", "documents")
item_id = "doc1"
item_value = {"title": "Test Document", "content": "Hello, World!"}
await store.aput(namespace, item_id, item_value)
item = await store.aget(namespace, item_id)
assert item
assert item.namespace == namespace
assert item.key == item_id
assert item.value == item_value
updated_value = {
"title": "Updated Test Document",
"content": "Hello, LangGraph!",
}
await store.aput(namespace, item_id, updated_value)
updated_item = await store.aget(namespace, item_id)
assert updated_item.value == updated_value
assert updated_item.updated_at > item.updated_at
different_namespace = ("test", "other_documents")
item_in_different_namespace = await store.aget(different_namespace, item_id)
assert item_in_different_namespace is None
new_item_id = "doc2"
new_item_value = {"title": "Another Document", "content": "Greetings!"}
await store.aput(namespace, new_item_id, new_item_value)
search_results = await store.asearch(["test"], limit=10)
items = search_results
assert len(items) == 2
assert any(item.key == item_id for item in items)
assert any(item.key == new_item_id for item in items)
namespaces = await store.alist_namespaces(prefix=["test"])
assert ("test", "documents") in namespaces
await store.adelete(namespace, item_id)
await store.adelete(namespace, new_item_id)
deleted_item = await store.aget(namespace, item_id)
assert deleted_item is None
deleted_item = await store.aget(namespace, new_item_id)
assert deleted_item is None
empty_search_results = await store.asearch(["test"], limit=10)
assert len(empty_search_results) == 0
async def test_list_namespaces() -> None:
async with AsyncDuckDBStore.from_conn_string(":memory:") as store:
await store.setup()
test_pref = str(uuid.uuid4())
test_namespaces = [
(test_pref, "test", "documents", "public", test_pref),
(test_pref, "test", "documents", "private", test_pref),
(test_pref, "test", "images", "public", test_pref),
(test_pref, "test", "images", "private", test_pref),
(test_pref, "prod", "documents", "public", test_pref),
(
test_pref,
"prod",
"documents",
"some",
"nesting",
"public",
test_pref,
),
(test_pref, "prod", "documents", "private", test_pref),
]
for namespace in test_namespaces:
await store.aput(namespace, "dummy", {"content": "dummy"})
prefix_result = await store.alist_namespaces(prefix=[test_pref, "test"])
assert len(prefix_result) == 4
assert all([ns[1] == "test" for ns in prefix_result])
specific_prefix_result = await store.alist_namespaces(
prefix=[test_pref, "test", "documents"]
)
assert len(specific_prefix_result) == 2
assert all([ns[1:3] == ("test", "documents") for ns in specific_prefix_result])
suffix_result = await store.alist_namespaces(suffix=["public", test_pref])
assert len(suffix_result) == 4
assert all(ns[-2] == "public" for ns in suffix_result)
prefix_suffix_result = await store.alist_namespaces(
prefix=[test_pref, "test"], suffix=["public", test_pref]
)
assert len(prefix_suffix_result) == 2
assert all(
ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result
)
wildcard_prefix_result = await store.alist_namespaces(
prefix=[test_pref, "*", "documents"]
)
assert len(wildcard_prefix_result) == 5
assert all(ns[2] == "documents" for ns in wildcard_prefix_result)
wildcard_suffix_result = await store.alist_namespaces(
suffix=["*", "public", test_pref]
)
assert len(wildcard_suffix_result) == 4
assert all(ns[-2] == "public" for ns in wildcard_suffix_result)
wildcard_single = await store.alist_namespaces(
suffix=["some", "*", "public", test_pref]
)
assert len(wildcard_single) == 1
assert wildcard_single[0] == (
test_pref,
"prod",
"documents",
"some",
"nesting",
"public",
test_pref,
)
max_depth_result = await store.alist_namespaces(max_depth=3)
assert all([len(ns) <= 3 for ns in max_depth_result])
max_depth_result = await store.alist_namespaces(
max_depth=4, prefix=[test_pref, "*", "documents"]
)
assert (
len(set(tuple(res) for res in max_depth_result))
== len(max_depth_result)
== 5
)
limit_result = await store.alist_namespaces(prefix=[test_pref], limit=3)
assert len(limit_result) == 3
offset_result = await store.alist_namespaces(prefix=[test_pref], offset=3)
assert len(offset_result) == len(test_namespaces) - 3
empty_prefix_result = await store.alist_namespaces(prefix=[test_pref])
assert len(empty_prefix_result) == len(test_namespaces)
assert set(tuple(ns) for ns in empty_prefix_result) == set(
tuple(ns) for ns in test_namespaces
)
for namespace in test_namespaces:
await store.adelete(namespace, "dummy")
async def test_search():
async with AsyncDuckDBStore.from_conn_string(":memory:") as store:
await store.setup()
test_namespaces = [
("test_search", "documents", "user1"),
("test_search", "documents", "user2"),
("test_search", "reports", "department1"),
("test_search", "reports", "department2"),
]
test_items = [
{"title": "Doc 1", "author": "John Doe", "tags": ["important"]},
{"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]},
{"title": "Report A", "author": "John Doe", "tags": ["final"]},
{"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]},
]
empty = await store.asearch(
(
"scoped",
"assistant_id",
"shared",
"6c5356f6-63ab-4158-868d-cd9fd14c736e",
),
limit=10,
offset=0,
)
assert len(empty) == 0
for namespace, item in zip(test_namespaces, test_items):
await store.aput(namespace, f"item_{namespace[-1]}", item)
docs_result = await store.asearch(["test_search", "documents"])
assert len(docs_result) == 2
assert all([item.namespace[1] == "documents" for item in docs_result]), [
item.namespace for item in docs_result
]
reports_result = await store.asearch(["test_search", "reports"])
assert len(reports_result) == 2
assert all(item.namespace[1] == "reports" for item in reports_result)
limited_result = await store.asearch(["test_search"], limit=2)
assert len(limited_result) == 2
offset_result = await store.asearch(["test_search"])
assert len(offset_result) == 4
offset_result = await store.asearch(["test_search"], offset=2)
assert len(offset_result) == 2
assert all(item not in limited_result for item in offset_result)
john_doe_result = await store.asearch(
["test_search"], filter={"author": "John Doe"}
)
assert len(john_doe_result) == 2
assert all(item.value["author"] == "John Doe" for item in john_doe_result)
draft_result = await store.asearch(["test_search"], filter={"tags": ["draft"]})
assert len(draft_result) == 2
assert all("draft" in item.value["tags"] for item in draft_result)
page1 = await store.asearch(["test_search"], limit=2, offset=0)
page2 = await store.asearch(["test_search"], limit=2, offset=2)
all_items = page1 + page2
assert len(all_items) == 4
assert len(set(item.key for item in all_items)) == 4
empty = await store.asearch(
(
"scoped",
"assistant_id",
"shared",
"again",
"maybe",
"some-long",
"6be5cb0e-2eb4-42e6-bb6b-fba3c269db25",
),
limit=10,
offset=0,
)
assert len(empty) == 0
# Test with a namespace beginning with a number (like a UUID)
uuid_namespace = (str(uuid.uuid4()), "documents")
uuid_item_id = "uuid_doc"
uuid_item_value = {
"title": "UUID Document",
"content": "This document has a UUID namespace.",
}
# Insert the item with the UUID namespace
await store.aput(uuid_namespace, uuid_item_id, uuid_item_value)
# Retrieve the item to verify it was stored correctly
retrieved_item = await store.aget(uuid_namespace, uuid_item_id)
assert retrieved_item is not None
assert retrieved_item.namespace == uuid_namespace
assert retrieved_item.key == uuid_item_id
assert retrieved_item.value == uuid_item_value
# Search for the item using the UUID namespace
search_result = await store.asearch([uuid_namespace[0]])
assert len(search_result) == 1
assert search_result[0].key == uuid_item_id
assert search_result[0].value == uuid_item_value
# Clean up: delete the item with the UUID namespace
await store.adelete(uuid_namespace, uuid_item_id)
# Verify the item was deleted
deleted_item = await store.aget(uuid_namespace, uuid_item_id)
assert deleted_item is None
for namespace in test_namespaces:
await store.adelete(namespace, f"item_{namespace[-1]}")
-457
View File
@@ -1,457 +0,0 @@
# type: ignore
import uuid
from datetime import datetime
from typing import Any
from unittest.mock import MagicMock
import pytest
from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
from langgraph.store.duckdb import DuckDBStore
class MockCursor:
def __init__(self, fetch_result: Any) -> None:
self.fetch_result = fetch_result
self.execute = MagicMock()
self.fetchall = MagicMock(return_value=self.fetch_result)
class MockConnection:
def __init__(self) -> None:
self.cursor = MagicMock()
@pytest.fixture
def mock_connection() -> MockConnection:
return MockConnection()
@pytest.fixture
def store(mock_connection: MockConnection) -> DuckDBStore:
duck_db_store = DuckDBStore(mock_connection)
duck_db_store.setup()
return duck_db_store
def test_batch_order(store: DuckDBStore) -> None:
mock_connection = store.conn
mock_get_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
(
"test.bar",
"key2",
'{"data": "value2"}',
datetime.now(),
datetime.now(),
),
]
)
mock_search_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
]
)
mock_list_namespaces_cursor = MockCursor(
[
("test",),
]
)
failures = []
def cursor_side_effect() -> Any:
cursor = MagicMock()
def execute_side_effect(query: str, *params: Any) -> None:
# My super sophisticated database.
if "WHERE prefix = ? AND key" in query:
cursor.fetchall = mock_get_cursor.fetchall
elif "SELECT prefix, key, value" in query:
cursor.fetchall = mock_search_cursor.fetchall
elif "SELECT DISTINCT ON (truncated_prefix)" in query:
cursor.fetchall = mock_list_namespaces_cursor.fetchall
elif "INSERT INTO " in query:
pass
else:
e = ValueError(f"Unmatched query: {query}")
failures.append(e)
raise e
cursor.execute = MagicMock(side_effect=execute_side_effect)
return cursor
mock_connection.cursor.side_effect = cursor_side_effect
ops = [
GetOp(namespace=("test",), key="key1"),
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
SearchOp(
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
),
ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0),
GetOp(namespace=("test",), key="key3"),
]
results = store.batch(ops)
assert not failures
assert len(results) == 5
assert isinstance(results[0], Item)
assert isinstance(results[0].value, dict)
assert results[0].value == {"data": "value1"}
assert results[0].key == "key1"
assert results[1] is None
assert isinstance(results[2], list)
assert len(results[2]) == 1
assert isinstance(results[3], list)
assert results[3] == [("test",)]
assert results[4] is None
ops_reordered = [
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
GetOp(namespace=("test",), key="key2"),
ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0),
PutOp(namespace=("test",), key="key3", value={"data": "value3"}),
GetOp(namespace=("test",), key="key1"),
]
results_reordered = store.batch(ops_reordered)
assert not failures
assert len(results_reordered) == 5
assert isinstance(results_reordered[0], list)
assert len(results_reordered[0]) == 1
assert isinstance(results_reordered[1], Item)
assert results_reordered[1].value == {"data": "value2"}
assert results_reordered[1].key == "key2"
assert isinstance(results_reordered[2], list)
assert results_reordered[2] == [("test",)]
assert results_reordered[3] is None
assert isinstance(results_reordered[4], Item)
assert results_reordered[4].value == {"data": "value1"}
assert results_reordered[4].key == "key1"
def test_batch_get_ops(store: DuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
(
"test.bar",
"key2",
'{"data": "value2"}',
datetime.now(),
datetime.now(),
),
]
)
mock_connection.cursor.return_value = mock_cursor
ops = [
GetOp(namespace=("test",), key="key1"),
GetOp(namespace=("test",), key="key2"),
GetOp(namespace=("test",), key="key3"),
]
results = store.batch(ops)
assert len(results) == 3
assert results[0] is not None
assert results[1] is not None
assert results[2] is None
assert results[0].key == "key1"
assert results[1].key == "key2"
def test_batch_put_ops(store: DuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor([])
mock_connection.cursor.return_value = mock_cursor
ops = [
PutOp(namespace=("test",), key="key1", value={"data": "value1"}),
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
PutOp(namespace=("test",), key="key3", value=None),
]
results = store.batch(ops)
assert len(results) == 3
assert all(result is None for result in results)
assert mock_cursor.execute.call_count == 2
def test_batch_search_ops(store: DuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor(
[
(
"test.foo",
"key1",
'{"data": "value1"}',
datetime.now(),
datetime.now(),
),
(
"test.bar",
"key2",
'{"data": "value2"}',
datetime.now(),
datetime.now(),
),
]
)
mock_connection.cursor.return_value = mock_cursor
ops = [
SearchOp(
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
),
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
]
results = store.batch(ops)
assert len(results) == 2
assert len(results[0]) == 2
assert len(results[1]) == 2
def test_batch_list_namespaces_ops(store: DuckDBStore) -> None:
mock_connection = store.conn
mock_cursor = MockCursor([("test.namespace1",), ("test.namespace2",)])
mock_connection.cursor.return_value = mock_cursor
ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)]
results = store.batch(ops)
assert len(results) == 1
assert results[0] == [("test", "namespace1"), ("test", "namespace2")]
def test_basic_store_ops() -> None:
with DuckDBStore.from_conn_string(":memory:") as store:
store.setup()
namespace = ("test", "documents")
item_id = "doc1"
item_value = {"title": "Test Document", "content": "Hello, World!"}
store.put(namespace, item_id, item_value)
item = store.get(namespace, item_id)
assert item
assert item.namespace == namespace
assert item.key == item_id
assert item.value == item_value
updated_value = {
"title": "Updated Test Document",
"content": "Hello, LangGraph!",
}
store.put(namespace, item_id, updated_value)
updated_item = store.get(namespace, item_id)
assert updated_item.value == updated_value
assert updated_item.updated_at > item.updated_at
different_namespace = ("test", "other_documents")
item_in_different_namespace = store.get(different_namespace, item_id)
assert item_in_different_namespace is None
new_item_id = "doc2"
new_item_value = {"title": "Another Document", "content": "Greetings!"}
store.put(namespace, new_item_id, new_item_value)
search_results = store.search(["test"], limit=10)
items = search_results
assert len(items) == 2
assert any(item.key == item_id for item in items)
assert any(item.key == new_item_id for item in items)
namespaces = store.list_namespaces(prefix=["test"])
assert ("test", "documents") in namespaces
store.delete(namespace, item_id)
store.delete(namespace, new_item_id)
deleted_item = store.get(namespace, item_id)
assert deleted_item is None
deleted_item = store.get(namespace, new_item_id)
assert deleted_item is None
empty_search_results = store.search(["test"], limit=10)
assert len(empty_search_results) == 0
def test_list_namespaces() -> None:
with DuckDBStore.from_conn_string(":memory:") as store:
store.setup()
test_pref = str(uuid.uuid4())
test_namespaces = [
(test_pref, "test", "documents", "public", test_pref),
(test_pref, "test", "documents", "private", test_pref),
(test_pref, "test", "images", "public", test_pref),
(test_pref, "test", "images", "private", test_pref),
(test_pref, "prod", "documents", "public", test_pref),
(
test_pref,
"prod",
"documents",
"some",
"nesting",
"public",
test_pref,
),
(test_pref, "prod", "documents", "private", test_pref),
]
for namespace in test_namespaces:
store.put(namespace, "dummy", {"content": "dummy"})
prefix_result = store.list_namespaces(prefix=[test_pref, "test"])
assert len(prefix_result) == 4
assert all([ns[1] == "test" for ns in prefix_result])
specific_prefix_result = store.list_namespaces(
prefix=[test_pref, "test", "documents"]
)
assert len(specific_prefix_result) == 2
assert all([ns[1:3] == ("test", "documents") for ns in specific_prefix_result])
suffix_result = store.list_namespaces(suffix=["public", test_pref])
assert len(suffix_result) == 4
assert all(ns[-2] == "public" for ns in suffix_result)
prefix_suffix_result = store.list_namespaces(
prefix=[test_pref, "test"], suffix=["public", test_pref]
)
assert len(prefix_suffix_result) == 2
assert all(
ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result
)
wildcard_prefix_result = store.list_namespaces(
prefix=[test_pref, "*", "documents"]
)
assert len(wildcard_prefix_result) == 5
assert all(ns[2] == "documents" for ns in wildcard_prefix_result)
wildcard_suffix_result = store.list_namespaces(
suffix=["*", "public", test_pref]
)
assert len(wildcard_suffix_result) == 4
assert all(ns[-2] == "public" for ns in wildcard_suffix_result)
wildcard_single = store.list_namespaces(
suffix=["some", "*", "public", test_pref]
)
assert len(wildcard_single) == 1
assert wildcard_single[0] == (
test_pref,
"prod",
"documents",
"some",
"nesting",
"public",
test_pref,
)
max_depth_result = store.list_namespaces(max_depth=3)
assert all([len(ns) <= 3 for ns in max_depth_result])
max_depth_result = store.list_namespaces(
max_depth=4, prefix=[test_pref, "*", "documents"]
)
assert (
len(set(tuple(res) for res in max_depth_result))
== len(max_depth_result)
== 5
)
limit_result = store.list_namespaces(prefix=[test_pref], limit=3)
assert len(limit_result) == 3
offset_result = store.list_namespaces(prefix=[test_pref], offset=3)
assert len(offset_result) == len(test_namespaces) - 3
empty_prefix_result = store.list_namespaces(prefix=[test_pref])
assert len(empty_prefix_result) == len(test_namespaces)
assert set(tuple(ns) for ns in empty_prefix_result) == set(
tuple(ns) for ns in test_namespaces
)
for namespace in test_namespaces:
store.delete(namespace, "dummy")
def test_search():
with DuckDBStore.from_conn_string(":memory:") as store:
store.setup()
test_namespaces = [
("test_search", "documents", "user1"),
("test_search", "documents", "user2"),
("test_search", "reports", "department1"),
("test_search", "reports", "department2"),
]
test_items = [
{"title": "Doc 1", "author": "John Doe", "tags": ["important"]},
{"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]},
{"title": "Report A", "author": "John Doe", "tags": ["final"]},
{"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]},
]
for namespace, item in zip(test_namespaces, test_items):
store.put(namespace, f"item_{namespace[-1]}", item)
docs_result = store.search(["test_search", "documents"])
assert len(docs_result) == 2
assert all(
[item.namespace[1] == "documents" for item in docs_result]
), docs_result
reports_result = store.search(["test_search", "reports"])
assert len(reports_result) == 2
assert all(item.namespace[1] == "reports" for item in reports_result)
limited_result = store.search(["test_search"], limit=2)
assert len(limited_result) == 2
offset_result = store.search(["test_search"])
assert len(offset_result) == 4
offset_result = store.search(["test_search"], offset=2)
assert len(offset_result) == 2
assert all(item not in limited_result for item in offset_result)
john_doe_result = store.search(["test_search"], filter={"author": "John Doe"})
assert len(john_doe_result) == 2
assert all(item.value["author"] == "John Doe" for item in john_doe_result)
draft_result = store.search(["test_search"], filter={"tags": ["draft"]})
assert len(draft_result) == 2
assert all("draft" in item.value["tags"] for item in draft_result)
page1 = store.search(["test_search"], limit=2, offset=0)
page2 = store.search(["test_search"], limit=2, offset=2)
all_items = page1 + page2
assert len(all_items) == 4
assert len(set(item.key for item in all_items)) == 4
for namespace in test_namespaces:
store.delete(namespace, f"item_{namespace[-1]}")
-111
View File
@@ -1,111 +0,0 @@
from typing import Any
import pytest
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
Checkpoint,
CheckpointMetadata,
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.duckdb import DuckDBSaver
class TestDuckDBSaver:
@pytest.fixture(autouse=True)
def setup(self) -> None:
# objects for test setup
self.config_1: RunnableConfig = {
"configurable": {
"thread_id": "thread-1",
# for backwards compatibility testing
"thread_ts": "1",
"checkpoint_ns": "",
}
}
self.config_2: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2",
"checkpoint_ns": "",
}
}
self.config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
}
}
self.chkpnt_1: Checkpoint = empty_checkpoint()
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
self.chkpnt_3: Checkpoint = empty_checkpoint()
self.metadata_1: CheckpointMetadata = {
"source": "input",
"step": 2,
"writes": {},
"score": 1,
}
self.metadata_2: CheckpointMetadata = {
"source": "loop",
"step": 1,
"writes": {"foo": "bar"},
"score": None,
}
self.metadata_3: CheckpointMetadata = {}
def test_search(self) -> None:
with DuckDBSaver.from_conn_string(":memory:") as saver:
saver.setup()
# save checkpoints
saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {})
saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
# call method / assertions
query_1 = {"source": "input"} # search by 1 key
query_2 = {
"step": 1,
"writes": {"foo": "bar"},
} # search by multiple keys
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
query_4 = {"source": "update", "step": 1} # no match
search_results_1 = list(saver.list(None, filter=query_1))
assert len(search_results_1) == 1
assert search_results_1[0].metadata == self.metadata_1
search_results_2 = list(saver.list(None, filter=query_2))
assert len(search_results_2) == 1
assert search_results_2[0].metadata == self.metadata_2
search_results_3 = list(saver.list(None, filter=query_3))
assert len(search_results_3) == 3
search_results_4 = list(saver.list(None, filter=query_4))
assert len(search_results_4) == 0
# search by config (defaults to checkpoints across all namespaces)
search_results_5 = list(
saver.list({"configurable": {"thread_id": "thread-2"}})
)
assert len(search_results_5) == 2
assert {
search_results_5[0].config["configurable"]["checkpoint_ns"],
search_results_5[1].config["configurable"]["checkpoint_ns"],
} == {"", "inner"}
# TODO: test before and limit params
def test_null_chars(self) -> None:
with DuckDBSaver.from_conn_string(":memory:") as saver:
saver.setup()
config = saver.put(self.config_1, self.chkpnt_1, {"my_key": "\x00abc"}, {})
assert saver.get_tuple(config).metadata["my_key"] == "abc" # type: ignore
assert (
list(saver.list(None, filter={"my_key": "abc"}))[0].metadata["my_key"] # type: ignore
== "abc"
)
@@ -327,6 +327,7 @@ class PostgresSaver(BasePostgresSaver):
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Store intermediate writes linked to a checkpoint.
@@ -350,6 +351,7 @@ class PostgresSaver(BasePostgresSaver):
config["configurable"]["checkpoint_ns"],
config["configurable"]["checkpoint_id"],
task_id,
task_path,
writes,
),
)
@@ -285,6 +285,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Store intermediate writes linked to a checkpoint asynchronously.
@@ -306,6 +307,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
config["configurable"]["checkpoint_ns"],
config["configurable"]["checkpoint_id"],
task_id,
task_path,
writes,
)
async with self._cursor(pipeline=True) as cur:
@@ -462,6 +464,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Store intermediate writes linked to a checkpoint.
@@ -471,9 +474,10 @@ class AsyncPostgresSaver(BasePostgresSaver):
config (RunnableConfig): Configuration of the related checkpoint.
writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
task_path (str): Path of the task creating the writes.
"""
return asyncio.run_coroutine_threadsafe(
self.aput_writes(config, writes, task_id), self.loop
self.aput_writes(config, writes, task_id, task_path), self.loop
).result()
@@ -66,6 +66,7 @@ MIGRATIONS = [
"""
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id);
""",
"""ALTER TABLE checkpoint_writes ADD COLUMN task_path TEXT NOT NULL DEFAULT '';""",
]
SELECT_SQL = f"""
@@ -94,7 +95,7 @@ select
and cw.checkpoint_id = checkpoints.checkpoint_id
) as pending_writes,
(
select array_agg(array[cw.type::bytea, cw.blob] order by cw.task_id, cw.idx)
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
@@ -119,8 +120,8 @@ UPSERT_CHECKPOINTS_SQL = """
"""
UPSERT_CHECKPOINT_WRITES_SQL = """
INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
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,
@@ -128,8 +129,8 @@ UPSERT_CHECKPOINT_WRITES_SQL = """
"""
INSERT_CHECKPOINT_WRITES_SQL = """
INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
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
"""
@@ -220,14 +221,16 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
checkpoint_ns: str,
checkpoint_id: str,
task_id: str,
task_path: str,
writes: Sequence[tuple[str, Any]],
) -> list[tuple[str, str, str, str, int, str, str, bytes]]:
) -> list[tuple[str, str, str, str, str, int, str, str, bytes]]:
return [
(
thread_id,
checkpoint_ns,
checkpoint_id,
task_id,
task_path,
WRITES_IDX_MAP.get(channel, idx),
channel,
*self.serde.dumps_typed(value),
@@ -74,6 +74,9 @@ MIGRATIONS = [
"""
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id);
""",
"""
ALTER TABLE checkpoint_writes ADD COLUMN task_path TEXT NOT NULL DEFAULT '';
""",
]
SELECT_SQL = f"""
@@ -99,7 +102,7 @@ select
and cw.checkpoint_id = (checkpoint->>'id')
) as pending_writes,
(
select array_agg(array[cw.type::bytea, cw.blob] order by cw.task_id, cw.idx)
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
@@ -125,8 +128,8 @@ UPSERT_CHECKPOINTS_SQL = """
"""
UPSERT_CHECKPOINT_WRITES_SQL = """
INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
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,
@@ -134,8 +137,8 @@ UPSERT_CHECKPOINT_WRITES_SQL = """
"""
INSERT_CHECKPOINT_WRITES_SQL = """
INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
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
"""
@@ -430,6 +433,7 @@ class ShallowPostgresSaver(BasePostgresSaver):
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Store intermediate writes linked to a checkpoint.
@@ -453,6 +457,7 @@ class ShallowPostgresSaver(BasePostgresSaver):
config["configurable"]["checkpoint_ns"],
config["configurable"]["checkpoint_id"],
task_id,
task_path,
writes,
),
)
@@ -747,6 +752,7 @@ class AsyncShallowPostgresSaver(BasePostgresSaver):
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Store intermediate writes linked to a checkpoint asynchronously.
@@ -768,6 +774,7 @@ class AsyncShallowPostgresSaver(BasePostgresSaver):
config["configurable"]["checkpoint_ns"],
config["configurable"]["checkpoint_id"],
task_id,
task_path,
writes,
)
async with self._cursor(pipeline=True) as cur:
@@ -903,6 +910,7 @@ class AsyncShallowPostgresSaver(BasePostgresSaver):
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Store intermediate writes linked to a checkpoint.
@@ -912,7 +920,8 @@ class AsyncShallowPostgresSaver(BasePostgresSaver):
config (RunnableConfig): Configuration of the related checkpoint.
writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
task_path (str): Path of the task creating the writes.
"""
return asyncio.run_coroutine_threadsafe(
self.aput_writes(config, writes, task_id), self.loop
self.aput_writes(config, writes, task_id, task_path), self.loop
).result()
@@ -424,6 +424,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
config: RunnableConfig,
writes: Sequence[Tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Store intermediate writes linked to a checkpoint.
@@ -433,6 +434,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
config (RunnableConfig): Configuration of the related checkpoint.
writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
task_path (str): 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 (?, ?, ?, ?, ?, ?, ?, ?)"
@@ -1,17 +1,8 @@
import asyncio
import random
from collections.abc import AsyncIterator, Iterator, Sequence
from contextlib import asynccontextmanager
from typing import (
Any,
AsyncIterator,
Callable,
Dict,
Iterator,
Optional,
Sequence,
Tuple,
TypeVar,
)
from typing import Any, Callable, Optional, TypeVar
import aiosqlite
from langchain_core.runnables import RunnableConfig
@@ -173,7 +164,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
self,
config: Optional[RunnableConfig],
*,
filter: Optional[Dict[str, Any]] = None,
filter: Optional[dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
@@ -207,7 +198,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
while True:
try:
yield asyncio.run_coroutine_threadsafe(
anext(aiter_),
anext(aiter_), # noqa: F821
self.loop,
).result()
except StopAsyncIteration:
@@ -239,10 +230,14 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
).result()
def put_writes(
self, config: RunnableConfig, writes: Sequence[Tuple[str, Any]], task_id: str
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), self.loop
self.aput_writes(config, writes, task_id, task_path), self.loop
).result()
async def setup(self) -> None:
@@ -372,7 +367,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
self,
config: Optional[RunnableConfig],
*,
filter: Optional[Dict[str, Any]] = None,
filter: Optional[dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
@@ -398,9 +393,11 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
ORDER BY checkpoint_id DESC"""
if limit:
query += f" LIMIT {limit}"
async with self.lock, self.conn.execute(
query, params
) as cur, self.conn.cursor() as wcur:
async with (
self.lock,
self.conn.execute(query, params) as cur,
self.conn.cursor() as wcur,
):
async for (
thread_id,
checkpoint_ns,
@@ -467,16 +464,19 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
checkpoint_ns = config["configurable"]["checkpoint_ns"]
type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
serialized_metadata = self.jsonplus_serde.dumps(metadata)
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,
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()
@@ -491,8 +491,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
async def aput_writes(
self,
config: RunnableConfig,
writes: Sequence[Tuple[str, Any]],
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Store intermediate writes linked to a checkpoint asynchronously.
@@ -502,6 +503,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
config (RunnableConfig): Configuration of the related checkpoint.
writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
task_path (str): 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 (?, ?, ?, ?, ?, ?, ?, ?)"
@@ -1,16 +1,13 @@
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
from datetime import datetime, timezone
from typing import (
from typing import ( # noqa: UP035
Any,
AsyncIterator,
Dict,
Generic,
Iterator,
List,
Literal,
Mapping,
NamedTuple,
Optional,
Sequence,
Tuple,
TypedDict,
TypeVar,
@@ -305,6 +302,7 @@ class BaseCheckpointSaver(Generic[V]):
config: RunnableConfig,
writes: Sequence[Tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Store intermediate writes linked to a checkpoint.
@@ -312,6 +310,7 @@ class BaseCheckpointSaver(Generic[V]):
config (RunnableConfig): Configuration of the related checkpoint.
writes (List[Tuple[str, Any]]): List of writes to store.
task_id (str): Identifier for the task creating the writes.
task_path (str): Path of the task creating the writes.
Raises:
NotImplementedError: Implement this method in your custom checkpoint saver.
@@ -397,6 +396,7 @@ class BaseCheckpointSaver(Generic[V]):
config: RunnableConfig,
writes: Sequence[Tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Asynchronously store intermediate writes linked to a checkpoint.
@@ -404,6 +404,7 @@ class BaseCheckpointSaver(Generic[V]):
config (RunnableConfig): Configuration of the related checkpoint.
writes (List[Tuple[str, Any]]): List of writes to store.
task_id (str): Identifier for the task creating the writes.
task_path (str): Path of the task creating the writes.
Raises:
NotImplementedError: Implement this method in your custom checkpoint saver.
@@ -4,9 +4,10 @@ import pickle
import random
import shutil
from collections import defaultdict
from collections.abc import AsyncIterator, Iterator, Sequence
from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack
from types import TracebackType
from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple, Type
from typing import Any, Optional
from langchain_core.runnables import RunnableConfig
@@ -65,14 +66,15 @@ class MemorySaver(
],
]
writes: defaultdict[
tuple[str, str, str], dict[tuple[str, int], tuple[str, str, tuple[str, bytes]]]
tuple[str, str, str],
dict[tuple[str, int], tuple[str, str, tuple[str, bytes], str]],
]
def __init__(
self,
*,
serde: Optional[SerializerProtocol] = None,
factory: Type[defaultdict] = defaultdict,
factory: type[defaultdict] = defaultdict,
) -> None:
super().__init__(serde=serde)
self.storage = factory(lambda: defaultdict(dict))
@@ -125,24 +127,27 @@ class MemorySaver(
checkpoint, metadata, parent_checkpoint_id = saved
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
if parent_checkpoint_id:
sends = [
w[2]
for w in self.writes[
(thread_id, checkpoint_ns, parent_checkpoint_id)
].values()
if w[1] == TASKS
]
sends = sorted(
(
(*w, k[1])
for k, w in self.writes[
(thread_id, checkpoint_ns, parent_checkpoint_id)
].items()
if w[1] == TASKS
),
key=lambda w: (w[3], w[0], w[4]),
)
else:
sends = []
return CheckpointTuple(
config=config,
checkpoint={
**self.serde.loads_typed(checkpoint),
"pending_sends": [self.serde.loads_typed(s) for s in sends],
"pending_sends": [self.serde.loads_typed(s[2]) for s in sends],
},
metadata=self.serde.loads_typed(metadata),
pending_writes=[
(id, c, self.serde.loads_typed(v)) for id, c, v in writes
(id, c, self.serde.loads_typed(v)) for id, c, v, _ in writes
],
parent_config={
"configurable": {
@@ -160,13 +165,16 @@ class MemorySaver(
checkpoint, metadata, parent_checkpoint_id = checkpoints[checkpoint_id]
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
if parent_checkpoint_id:
sends = [
w[2]
for w in self.writes[
(thread_id, checkpoint_ns, parent_checkpoint_id)
].values()
if w[1] == TASKS
]
sends = sorted(
(
(*w, k[1])
for k, w in self.writes[
(thread_id, checkpoint_ns, parent_checkpoint_id)
].items()
if w[1] == TASKS
),
key=lambda w: (w[3], w[0], w[4]),
)
else:
sends = []
return CheckpointTuple(
@@ -179,11 +187,11 @@ class MemorySaver(
},
checkpoint={
**self.serde.loads_typed(checkpoint),
"pending_sends": [self.serde.loads_typed(s) for s in sends],
"pending_sends": [self.serde.loads_typed(s[2]) for s in sends],
},
metadata=self.serde.loads_typed(metadata),
pending_writes=[
(id, c, self.serde.loads_typed(v)) for id, c, v in writes
(id, c, self.serde.loads_typed(v)) for id, c, v, _ in writes
],
parent_config={
"configurable": {
@@ -200,7 +208,7 @@ class MemorySaver(
self,
config: Optional[RunnableConfig],
*,
filter: Optional[Dict[str, Any]] = None,
filter: Optional[dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
@@ -271,13 +279,16 @@ class MemorySaver(
].values()
if parent_checkpoint_id:
sends = [
w[2]
for w in self.writes[
(thread_id, checkpoint_ns, parent_checkpoint_id)
].values()
if w[1] == TASKS
]
sends = sorted(
(
(*w, k[1])
for k, w in self.writes[
(thread_id, checkpoint_ns, parent_checkpoint_id)
].items()
if w[1] == TASKS
),
key=lambda w: (w[3], w[0], w[4]),
)
else:
sends = []
@@ -291,7 +302,9 @@ class MemorySaver(
},
checkpoint={
**self.serde.loads_typed(checkpoint),
"pending_sends": [self.serde.loads_typed(s) for s in sends],
"pending_sends": [
self.serde.loads_typed(s[2]) for s in sends
],
},
metadata=metadata,
parent_config={
@@ -304,7 +317,7 @@ class MemorySaver(
if parent_checkpoint_id
else None,
pending_writes=[
(id, c, self.serde.loads_typed(v)) for id, c, v in writes
(id, c, self.serde.loads_typed(v)) for id, c, v, _ in writes
],
)
@@ -353,8 +366,9 @@ class MemorySaver(
def put_writes(
self,
config: RunnableConfig,
writes: Sequence[Tuple[str, Any]],
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Save a list of writes to the in-memory storage.
@@ -365,6 +379,7 @@ class MemorySaver(
config (RunnableConfig): The config to associate with the writes.
writes (list[tuple[str, Any]]): The writes to save.
task_id (str): Identifier for the task creating the writes.
task_path (str): Path of the task creating the writes.
Returns:
RunnableConfig: The updated config containing the saved writes' timestamp.
@@ -379,7 +394,12 @@ class MemorySaver(
if inner_key[1] >= 0 and outer_writes_ and inner_key in outer_writes_:
continue
self.writes[outer_key][inner_key] = (task_id, c, self.serde.dumps_typed(v))
self.writes[outer_key][inner_key] = (
task_id,
c,
self.serde.dumps_typed(v),
task_path,
)
async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Asynchronous version of get_tuple.
@@ -399,7 +419,7 @@ class MemorySaver(
self,
config: Optional[RunnableConfig],
*,
filter: Optional[Dict[str, Any]] = None,
filter: Optional[dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
@@ -440,8 +460,9 @@ class MemorySaver(
async def aput_writes(
self,
config: RunnableConfig,
writes: Sequence[Tuple[str, Any]],
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Asynchronous version of put_writes.
@@ -452,9 +473,12 @@ class MemorySaver(
config (RunnableConfig): The config to associate with the writes.
writes (List[Tuple[str, Any]]): The writes to save, each as a (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
return self.put_writes(config, writes, task_id)
task_path (str): Path of the task creating the writes.
Returns:
None
"""
return self.put_writes(config, writes, task_id)
return self.put_writes(config, writes, task_id, task_path)
def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
if current is None:
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-checkpoint"
version = "2.0.9"
version = "2.0.10"
description = "Library with base interfaces for LangGraph checkpoint savers."
authors = []
license = "MIT"
-3
View File
@@ -1,5 +1,4 @@
import sys
from os import getenv
from types import MappingProxyType
from typing import Any, Literal, Mapping, cast
@@ -93,8 +92,6 @@ NS_END = sys.intern(":")
# for checkpoint_ns, for each level, separates the namespace from the task_id
CONF = cast(Literal["configurable"], sys.intern("configurable"))
# key for the configurable dict in RunnableConfig
FF_SEND_V2 = getenv("LANGGRAPH_FF_SEND_V2", "false").lower() == "true"
# temporary flag to enable new Send semantics
NULL_TASK_ID = sys.intern("00000000-0000-0000-0000-000000000000")
# the task_id to use for writes that are not associated with a task
+5 -5
View File
@@ -19,7 +19,7 @@ from typing_extensions import ParamSpec
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import END, START, TAG_HIDDEN
from langgraph.constants import CONF, END, START, TAG_HIDDEN
from langgraph.pregel import Pregel
from langgraph.pregel.call import get_runnable_for_func
from langgraph.pregel.read import PregelNode
@@ -39,11 +39,11 @@ def call(
**kwargs: Any,
) -> concurrent.futures.Future[T]:
from langgraph.constants import CONFIG_KEY_CALL
from langgraph.utils.config import get_configurable
from langgraph.utils.config import get_config
conf = get_configurable()
impl = conf[CONFIG_KEY_CALL]
fut = impl(func, (args, kwargs), retry=retry)
config = get_config()
impl = config[CONF][CONFIG_KEY_CALL]
fut = impl(func, (args, kwargs), retry=retry, callbacks=config["callbacks"])
return fut
+4 -4
View File
@@ -257,7 +257,7 @@ class Graph:
selected by `path`.
Returns:
None
Self: The instance of the graph, allowing for method chaining.
Note: Without typehints on the `path` function's return value (e.g., `-> Literal["foo", "__end__"]:`)
or a path_map, the graph visualization assumes the edge could transition to any node in the graph.
@@ -308,7 +308,7 @@ class Graph:
key (str): The key of the node to set as the entry point.
Returns:
None
Self: The instance of the graph, allowing for method chaining.
"""
return self.add_edge(START, key)
@@ -334,7 +334,7 @@ class Graph:
selected by `path`.
Returns:
None
Self: The instance of the graph, allowing for method chaining.
"""
return self.add_conditional_edges(START, path, path_map, then)
@@ -347,7 +347,7 @@ class Graph:
key (str): The key of the node to set as the finish point.
Returns:
None
Self: The instance of the graph, allowing for method chaining.
"""
return self.add_edge(key, END)
+11 -8
View File
@@ -107,7 +107,6 @@ class StateGraph(Graph):
config_schema (Optional[Type[Any]]): The schema class that defines the configuration.
Use this to expose configurable parameters in your API.
Examples:
>>> from langchain_core.runnables import RunnableConfig
>>> from typing_extensions import Annotated, TypedDict
@@ -243,7 +242,7 @@ class StateGraph(Graph):
ValueError: If the key is already being used as a state key.
Returns:
StateGraph
Self: The instance of the state graph, allowing for method chaining.
"""
...
@@ -267,7 +266,7 @@ class StateGraph(Graph):
ValueError: If the key is already being used as a state key.
Returns:
StateGraph
Self: The instance of the state graph, allowing for method chaining.
"""
...
@@ -290,10 +289,10 @@ class StateGraph(Graph):
metadata (Optional[dict[str, Any]]): The metadata associated with the node. (default: None)
input (Optional[Type[Any]]): The input schema for the node. (default: the graph's input schema)
retry (Optional[RetryPolicy]): The policy for retrying the node. (default: None)
Raises:
ValueError: If the key is already being used as a state key.
Examples:
```pycon
>>> from langgraph.graph import START, StateGraph
@@ -320,7 +319,7 @@ class StateGraph(Graph):
```
Returns:
StateGraph
Self: The instance of the state graph, allowing for method chaining.
"""
if not isinstance(node, str):
action = node
@@ -361,7 +360,11 @@ class StateGraph(Graph):
ends = EMPTY_SEQ
try:
if (isfunction(action) or ismethod(getattr(action, "__call__", None))) and (
if (
isfunction(action)
or ismethod(action)
or ismethod(getattr(action, "__call__", None))
) and (
hints := get_type_hints(getattr(action, "__call__"))
or get_type_hints(action)
):
@@ -412,7 +415,7 @@ class StateGraph(Graph):
ValueError: If the start key is 'END' or if the start key or end key is not present in the graph.
Returns:
StateGraph
Self: The instance of the state graph, allowing for method chaining.
"""
if isinstance(start_key, str):
return super().add_edge(start_key, end_key)
@@ -451,7 +454,7 @@ class StateGraph(Graph):
ValueError: if the sequence contains duplicate node names.
Returns:
StateGraph
Self: The instance of the state graph, allowing for method chaining.
"""
if len(nodes) < 1:
raise ValueError("Sequence requires at least one node.")
@@ -223,7 +223,7 @@ def _validate_chat_history(
@deprecated_parameter("messages_modifier", "0.1.9", "state_modifier", removal="0.3.0")
def create_react_agent(
model: LanguageModelLike,
model: Union[str, LanguageModelLike],
tools: Union[ToolExecutor, Sequence[BaseTool], ToolNode],
*,
state_schema: Optional[StateSchemaType] = None,
@@ -595,6 +595,18 @@ def create_react_agent(
# get the tool functions wrapped in a tool class from the ToolNode
tool_classes = list(tool_node.tools_by_name.values())
if isinstance(model, str):
try:
from langchain.chat_models import ( # type: ignore[import-not-found]
init_chat_model,
)
except ImportError:
raise ImportError(
"Please install langchain (`pip install langchain`) to use '<provider>:<model>' string syntax for `model` parameter."
)
model = cast(BaseChatModel, init_chat_model(model))
tool_calling_enabled = len(tool_classes) > 0
if _should_bind_tools(model, tool_classes) and tool_calling_enabled:
+7 -5
View File
@@ -966,7 +966,7 @@ class Pregel(PregelProtocol):
return patch_checkpoint_map(
next_config, saved.metadata if saved else None
)
# no values, copy checkpoint
# no values, empty checkpoint
if values is None and as_node is None:
next_checkpoint = create_checkpoint(checkpoint, None, step)
# copy checkpoint
@@ -985,6 +985,7 @@ class Pregel(PregelProtocol):
return patch_checkpoint_map(
next_config, saved.metadata if saved else None
)
# no values, copy checkpoint
if values is None and as_node == "__copy__":
next_checkpoint = create_checkpoint(checkpoint, None, step)
# copy checkpoint
@@ -1248,7 +1249,7 @@ class Pregel(PregelProtocol):
return patch_checkpoint_map(
next_config, saved.metadata if saved else None
)
# no values, copy checkpoint
# no values, empty checkpoint
if values is None and as_node is None:
next_checkpoint = create_checkpoint(checkpoint, None, step)
# copy checkpoint
@@ -1267,6 +1268,7 @@ class Pregel(PregelProtocol):
return patch_checkpoint_map(
next_config, saved.metadata if saved else None
)
# no values, copy checkpoint
if values is None and as_node == "__copy__":
next_checkpoint = create_checkpoint(checkpoint, None, step)
# copy checkpoint
@@ -1652,10 +1654,10 @@ class Pregel(PregelProtocol):
else:
get_waiter = None # type: ignore[assignment]
# Similarly to Bulk Synchronous Parallel / Pregel model
# computation proceeds in steps, while there are channel updates
# channel updates from step N are only visible in step N+1
# computation proceeds in steps, while there are channel updates.
# Channel updates from step N are only visible in step N+1
# channels are guaranteed to be immutable for the duration of the step,
# with channel updates applied only at the transition between steps
# with channel updates applied only at the transition between steps.
while loop.tick(input_keys=self.input_channels):
for _ in runner.tick(
loop.tasks.values(),
+22 -108
View File
@@ -19,6 +19,7 @@ from typing import (
)
from uuid import UUID
from langchain_core.callbacks import Callbacks
from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager
from langchain_core.runnables.config import RunnableConfig
@@ -107,18 +108,25 @@ class PregelTaskWrites(NamedTuple):
class Call:
__slots__ = ("func", "input", "retry")
__slots__ = ("func", "input", "retry", "callbacks")
func: Callable
input: Any
retry: Optional[RetryPolicy]
callbacks: Callbacks
def __init__(
self, func: Callable, input: Any, *, retry: Optional[RetryPolicy]
self,
func: Callable,
input: Any,
*,
retry: Optional[RetryPolicy],
callbacks: Callbacks,
) -> None:
self.func = func
self.input = input
self.retry = retry
self.callbacks = callbacks
def should_interrupt(
@@ -228,7 +236,7 @@ def apply_writes(
# sort tasks on path, to ensure deterministic order for update application
# any path parts after the 3rd are ignored for sorting
# (we use them for eg. task ids which aren't good for sorting)
tasks = sorted(tasks, key=lambda t: t.path[:3])
tasks = sorted(tasks, key=lambda t: _tuple_str(t.path[:3]))
# if no task has triggers this is applying writes from the null task only
# so we don't do anything other than update the channels written to
bump_step = any(t.triggers for t in tasks)
@@ -273,7 +281,7 @@ def apply_writes(
for chan, val in task.writes:
if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR):
pass
elif chan == TASKS: # TODO: remove branch in 1.0
elif chan == TASKS:
checkpoint["pending_sends"].append(val)
elif chan in channels:
pending_writes_by_channel[chan].append(val)
@@ -363,8 +371,8 @@ def prepare_next_tasks(
This is the union of all PUSH tasks (Sends) and PULL tasks (nodes triggered
by edges)."""
tasks: list[Union[PregelTask, PregelExecutableTask]] = []
# Consume pending_sends from previous step (legacy version of Send)
for idx, _ in enumerate(checkpoint["pending_sends"]): # TODO: remove branch in 1.0
# Consume pending_sends from previous step
for idx, _ in enumerate(checkpoint["pending_sends"]):
if task := prepare_single_task(
(PUSH, idx),
None,
@@ -400,65 +408,7 @@ def prepare_next_tasks(
manager=manager,
):
tasks.append(task)
# Consume pending Sends from this step (new version of Send)
if any(c == PUSH for _, c, _ in pending_writes):
# group writes by task id
grouped_by_task = defaultdict(list)
for tid, c, _ in pending_writes:
grouped_by_task[tid].append(c)
# prepare send tasks from grouped writes
# 1. start from sends originating from existing tasks
tidx = 0
while tidx < len(tasks):
task = tasks[tidx]
if twrites := grouped_by_task.pop(task.id, None):
for idx, c in enumerate(twrites):
if c != PUSH:
continue
if next_task := prepare_single_task(
(PUSH, task.path, idx, task.id),
None,
checkpoint=checkpoint,
pending_writes=pending_writes,
processes=processes,
channels=channels,
managed=managed,
config=config,
step=step,
for_execution=for_execution,
store=store,
checkpointer=checkpointer,
manager=manager,
):
tasks.append(next_task)
tidx += 1
# key tasks by id
task_map = {t.id: t for t in tasks}
# 2. create new tasks for remaining sends (eg. from update_state)
for tid, writes in grouped_by_task.items():
task = task_map.get(tid)
for idx, c in enumerate(writes):
if c != PUSH:
continue
if next_task := prepare_single_task(
(PUSH, task.path if task else (), idx, tid),
None,
checkpoint=checkpoint,
pending_writes=pending_writes,
processes=processes,
channels=channels,
managed=managed,
config=config,
step=step,
for_execution=for_execution,
store=store,
checkpointer=checkpointer,
manager=manager,
):
task_map[next_task.id] = next_task
else:
task_map = {t.id: t for t in tasks}
return task_map
return {t.id: t for t in tasks}
def prepare_single_task(
@@ -523,9 +473,8 @@ def prepare_single_task(
patch_config(
merge_configs(config, {"metadata": metadata}),
run_name=name,
callbacks=(
manager.get_child(f"graph:step:{step}") if manager else None
),
callbacks=call.callbacks
or (manager.get_child(f"graph:step:{step}") if manager else None),
configurable={
CONFIG_KEY_TASK_ID: task_id,
# deque.extend is thread-safe
@@ -571,8 +520,8 @@ def prepare_single_task(
else:
return PregelTask(task_id, name, task_path[:3])
elif task_path[0] == PUSH:
if len(task_path) == 2: # TODO: remove branch in 1.0
# legacy SEND tasks, executed in superstep n+1
if len(task_path) == 2:
# SEND tasks, executed in superstep n+1
# (PUSH, idx of pending send)
idx = cast(int, task_path[1])
if idx >= len(checkpoint["pending_sends"]):
@@ -601,43 +550,6 @@ def prepare_single_task(
PUSH,
str(idx),
)
elif len(task_path) >= 4:
# new PUSH tasks, executed in superstep n
# (PUSH, parent task path, idx of PUSH write, id of parent task)
task_path_tt = cast(tuple[str, tuple, int, str], task_path)
writes_for_path = [w for w in pending_writes if w[0] == task_path_tt[3]]
if task_path_tt[2] >= len(writes_for_path):
logger.warning(
f"Ignoring invalid write index {task_path[2]} in pending writes"
)
return
packet = writes_for_path[task_path_tt[2]][2]
if packet is None:
return
if not isinstance(packet, Send):
logger.warning(
f"Ignoring invalid packet type {type(packet)} in pending writes"
)
return
if packet.node not in processes:
logger.warning(
f"Ignoring unknown node name {packet.node} in pending writes"
)
return
# create task id
triggers = [PUSH]
checkpoint_ns = (
f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
)
task_id = _uuid5_str(
checkpoint_id,
checkpoint_ns,
str(step),
packet.node,
PUSH,
_tuple_str(task_path[1]),
str(task_path[2]),
)
else:
logger.warning(f"Ignoring invalid PUSH task path {task_path}")
return
@@ -904,7 +816,9 @@ def _uuid5_str(namespace: bytes, *parts: str) -> str:
def _tuple_str(tup: Union[str, int, tuple]) -> str:
"""Generate a string representation of a tuple."""
return (
f"({', '.join(_tuple_str(x) for x in tup)})"
f"~{', '.join(_tuple_str(x) for x in tup)}"
if isinstance(tup, (tuple, list))
else f"{tup:010d}"
if isinstance(tup, int)
else str(tup)
)
+5 -4
View File
@@ -1,3 +1,5 @@
"""Utility to convert a user provided function into a Runnable with a ChannelWrite."""
import sys
import types
from typing import Any, Callable, Optional
@@ -6,10 +8,9 @@ from langgraph.constants import RETURN
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
from langgraph.utils.runnable import RunnableSeq, coerce_to_runnable
"""
Utilities borrowed from cloudpickle.
https://github.com/cloudpipe/cloudpickle/blob/6220b0ce83ffee5e47e06770a1ee38ca9e47c850/cloudpickle/cloudpickle.py#L265
"""
##
# Utilities borrowed from cloudpickle.
# https://github.com/cloudpipe/cloudpickle/blob/6220b0ce83ffee5e47e06770a1ee38ca9e47c850/cloudpickle/cloudpickle.py#L265
def _getattribute(obj: Any, name: str) -> Any:
+11 -4
View File
@@ -51,6 +51,7 @@ class BackgroundExecutor(ContextManager):
def __init__(self, config: RunnableConfig) -> None:
self.stack = ExitStack()
self.executor = self.stack.enter_context(get_executor_for_config(config))
# mapping of Future to (__cancel_on_exit__, __reraise_on_exit__) flags
self.tasks: dict[concurrent.futures.Future, tuple[bool, bool]] = {}
def submit( # type: ignore[valid-type]
@@ -63,15 +64,21 @@ class BackgroundExecutor(ContextManager):
__next_tick__: bool = False,
**kwargs: P.kwargs,
) -> concurrent.futures.Future[T]:
ctx = copy_context()
if __next_tick__:
task = self.executor.submit(next_tick, fn, *args, **kwargs)
task = cast(
concurrent.futures.Future[T],
self.executor.submit(next_tick, ctx.run, fn, *args, **kwargs), # type: ignore[arg-type]
)
else:
task = self.executor.submit(fn, *args, **kwargs)
task = self.executor.submit(ctx.run, fn, *args, **kwargs)
self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__)
# add a callback to remove the task from the tasks dict when it's done
task.add_done_callback(self.done)
return task
def done(self, task: concurrent.futures.Future) -> None:
"""Remove the task from the tasks dict when it's done."""
try:
task.result()
except GraphBubbleUp:
@@ -103,9 +110,9 @@ class BackgroundExecutor(ContextManager):
concurrent.futures.wait(pending)
# shutdown the executor
self.stack.__exit__(exc_type, exc_value, traceback)
# re-raise the first exception that occurred in a task
# if there's already an exception being raised, don't raise another one
if exc_type is None:
# if there's already an exception being raised, don't raise another one
# re-raise the first exception that occurred in a task
for task, (_, reraise) in tasks.items():
if not reraise:
continue
+1 -3
View File
@@ -9,10 +9,8 @@ from langgraph.checkpoint.base import PendingWrite
from langgraph.constants import (
EMPTY_SEQ,
ERROR,
FF_SEND_V2,
INTERRUPT,
NULL_TASK_ID,
PUSH,
RESUME,
RETURN,
SELF,
@@ -83,7 +81,7 @@ def map_command(
sends = [cmd.goto]
for send in sends:
if isinstance(send, Send):
yield (NULL_TASK_ID, PUSH if FF_SEND_V2 else TASKS, send)
yield (NULL_TASK_ID, TASKS, send)
elif isinstance(send, str):
yield (NULL_TASK_ID, f"branch:{START}:{SELF}:{send}", START)
else:
+12 -4
View File
@@ -359,7 +359,13 @@ class PregelLoop(LoopProtocol):
input_keys: Union[str, Sequence[str]],
) -> bool:
"""Execute a single iteration of the Pregel loop.
Returns True if more iterations are needed."""
Args:
input_keys: The key(s) to read input from.
Returns:
True if more iterations are needed.
"""
if self.status != "pending":
raise RuntimeError("Cannot tick when status is no longer 'pending'")
@@ -693,9 +699,6 @@ class PregelLoop(LoopProtocol):
traceback: Optional[TracebackType],
) -> Optional[bool]:
suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested
if suppress or exc_type is None:
# save final output
self.output = read_channels(self.channels, self.output_keys)
if suppress:
# emit one last "values" event, with pending writes applied
if (
@@ -723,8 +726,13 @@ class PregelLoop(LoopProtocol):
"updates",
lambda: iter([{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}]),
)
# save final output
self.output = read_channels(self.channels, self.output_keys)
# suppress interrupt
return True
elif exc_type is None:
# save final output
self.output = read_channels(self.channels, self.output_keys)
def _emit(
self,
+10 -2
View File
@@ -17,6 +17,8 @@ from typing import (
cast,
)
from langchain_core.callbacks import Callbacks
from langgraph.constants import (
CONF,
CONFIG_KEY_CALL,
@@ -148,9 +150,12 @@ class PregelRunner:
input: Any,
*,
retry: Optional[RetryPolicy] = None,
callbacks: Callbacks = None,
) -> concurrent.futures.Future[Any]:
(fut,) = writer(
task, [(PUSH, None)], calls=[Call(func, input, retry=retry)]
task,
[(PUSH, None)],
calls=[Call(func, input, retry=retry, callbacks=callbacks)],
)
assert fut is not None, "writer did not return a future for call"
return fut
@@ -337,9 +342,12 @@ class PregelRunner:
input: Any,
*,
retry: Optional[RetryPolicy] = None,
callbacks: Callbacks = None,
) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]:
(fut,) = writer(
task, [(PUSH, None)], calls=[Call(func, input, retry=retry)]
task,
[(PUSH, None)],
calls=[Call(func, input, retry=retry, callbacks=callbacks)],
)
assert fut is not None, "writer did not return a future for call"
if asyncio.iscoroutinefunction(func):
+3 -3
View File
@@ -14,7 +14,7 @@ from typing import (
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.runnables.utils import ConfigurableFieldSpec
from langgraph.constants import CONF, CONFIG_KEY_SEND, FF_SEND_V2, PUSH, TASKS, Send
from langgraph.constants import CONF, CONFIG_KEY_SEND, TASKS, Send
from langgraph.errors import InvalidUpdateError
from langgraph.utils.runnable import RunnableCallable
@@ -125,7 +125,7 @@ class ChannelWrite(RunnableCallable):
# validate
for w in writes:
if isinstance(w, ChannelWriteEntry):
if w.channel in (TASKS, PUSH):
if w.channel == TASKS:
raise InvalidUpdateError(
"Cannot write to the reserved channel TASKS"
)
@@ -138,7 +138,7 @@ class ChannelWrite(RunnableCallable):
tuples: list[tuple[str, Any]] = []
for w in writes:
if isinstance(w, Send):
tuples.append((PUSH if FF_SEND_V2 else TASKS, w))
tuples.append((TASKS, w))
elif isinstance(w, ChannelWriteTupleEntry):
if ww := w.mapper(w.value):
tuples.extend(ww)
+2 -2
View File
@@ -453,9 +453,9 @@ def interrupt(value: Any) -> Any:
RESUME,
)
from langgraph.errors import GraphInterrupt
from langgraph.utils.config import get_configurable
from langgraph.utils.config import get_config
conf = get_configurable()
conf = get_config()["configurable"]
# track interrupt index
scratchpad: PregelScratchpad = conf[CONFIG_KEY_SCRATCHPAD]
if "interrupt_counter" not in scratchpad:
+13 -6
View File
@@ -132,7 +132,7 @@ def merge_configs(*configs: Optional[RunnableConfig]) -> RunnableConfig:
def patch_config(
config: Optional[RunnableConfig],
*,
callbacks: Optional[Callbacks] = None,
callbacks: Callbacks = None,
recursion_limit: Optional[int] = None,
max_concurrency: Optional[int] = None,
run_name: Optional[str] = None,
@@ -250,6 +250,13 @@ def get_async_callback_manager_for_config(
)
def _is_not_empty(value: Any) -> bool:
if isinstance(value, (list, tuple, dict)):
return len(value) > 0
else:
return value is not None
def ensure_config(*configs: Optional[RunnableConfig]) -> RunnableConfig:
"""Ensure that a config is a dict with all keys present.
@@ -272,20 +279,20 @@ def ensure_config(*configs: Optional[RunnableConfig]) -> RunnableConfig:
{
k: v.copy() if k in COPIABLE_KEYS else v # type: ignore[attr-defined]
for k, v in var_config.items()
if v is not None
if _is_not_empty(v)
},
)
for config in configs:
if config is None:
continue
for k, v in config.items():
if v is not None and k in CONFIG_KEYS:
if _is_not_empty(v) and k in CONFIG_KEYS:
if k == CONF:
empty[k] = cast(dict, v).copy()
else:
empty[k] = v # type: ignore[literal-required]
for k, v in config.items():
if v is not None and k not in CONFIG_KEYS:
if _is_not_empty(v) and k not in CONFIG_KEYS:
empty[CONF][k] = v
for key, value in empty[CONF].items():
if (
@@ -297,7 +304,7 @@ def ensure_config(*configs: Optional[RunnableConfig]) -> RunnableConfig:
return empty
def get_configurable() -> dict[str, Any]:
def get_config() -> RunnableConfig:
if sys.version_info < (3, 11):
try:
if asyncio.current_task():
@@ -307,6 +314,6 @@ def get_configurable() -> dict[str, Any]:
except RuntimeError:
pass
if var_config := var_child_runnable_config.get():
return var_config[CONF]
return var_config
else:
raise RuntimeError("Called get_configurable outside of a runnable context")
+20 -6
View File
@@ -75,7 +75,20 @@ KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = (
),
)
"""List of kwargs that can be passed to functions, and their corresponding
config keys, default values and type annotations."""
config keys, default values and type annotations.
Used to configure keyword arguments that can be injected at runtime
from the config object as kwargs to `invoke`, `ainvoke`, `stream` and `astream`.
For a keyword to be injected from the config object, the function signature
must contain a kwarg with the same name and a matching type annotation.
Each tuple contains:
- the name of the kwarg in the function signature
- the type annotation(s) for the kwarg
- the config key to look for the value in
- the default value for the kwarg
"""
VALID_KINDS = (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY)
@@ -298,21 +311,22 @@ def coerce_to_runnable(
class RunnableSeq(Runnable):
"""A simpler version of RunnableSequence."""
"""Sequence of Runnables, where the output of each is the input of the next.
RunnableSeq is a simpler version of RunnableSequence that is internal to
LangGraph.
"""
def __init__(
self,
*steps: RunnableLike,
name: Optional[str] = None,
) -> None:
"""Create a new RunnableSequence.
"""Create a new RunnableSeq.
Args:
steps: The steps to include in the sequence.
name: The name of the Runnable. Defaults to None.
first: The first Runnable in the sequence. Defaults to None.
middle: The middle Runnables in the sequence. Defaults to None.
last: The last Runnable in the sequence. Defaults to None.
Raises:
ValueError: If the sequence has less than 2 steps.
+162 -84
View File
File diff suppressed because it is too large Load Diff
-1
View File
@@ -27,7 +27,6 @@ jupyter = "^1.0.0"
pytest-xdist = {extras = ["psutil"], version = "^3.6.1"}
pytest-repeat = "^0.9.3"
langgraph-checkpoint = {path = "../checkpoint", develop = true}
langgraph-checkpoint-duckdb = {path = "../checkpoint-duckdb", develop = true}
langgraph-checkpoint-sqlite = {path = "../checkpoint-sqlite", develop = true}
langgraph-checkpoint-postgres = {path = "../checkpoint-postgres", develop = true}
langgraph-sdk = {path = "../sdk-py", develop = true}
File diff suppressed because one or more lines are too long
@@ -329,127 +329,6 @@
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[duckdb_aio]
'''
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[duckdb_aio].1
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'answer': dict({
'anyOf': list([
dict({
'type': 'string',
}),
dict({
'type': 'null',
}),
]),
'default': None,
'title': 'Answer',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
'docs',
]),
'title': 'State',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[duckdb_aio].2
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'answer': dict({
'anyOf': list([
dict({
'type': 'string',
}),
dict({
'type': 'null',
}),
]),
'default': None,
'title': 'Answer',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
'docs',
]),
'title': 'State',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory]
'''
graph TD;
@@ -1513,31 +1392,6 @@
'''
# ---
# name: test_weather_subgraph[duckdb_aio]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
__end__([<p>__end__</p>]):::last
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
subgraph weather_graph
weather_graph_model_node --> weather_graph_weather_node;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[memory]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
-39
View File
@@ -11,8 +11,6 @@ from psycopg_pool import AsyncConnectionPool, ConnectionPool
from pytest_mock import MockerFixture
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.checkpoint.duckdb import DuckDBSaver
from langgraph.checkpoint.duckdb.aio import AsyncDuckDBSaver
from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver
from langgraph.checkpoint.postgres.aio import (
AsyncPostgresSaver,
@@ -21,7 +19,6 @@ from langgraph.checkpoint.postgres.aio import (
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from langgraph.store.base import BaseStore
from langgraph.store.duckdb import AsyncDuckDBStore, DuckDBStore
from langgraph.store.memory import InMemoryStore
from langgraph.store.postgres import AsyncPostgresStore, PostgresStore
@@ -70,20 +67,6 @@ async def _checkpointer_sqlite_aio():
yield checkpointer
@pytest.fixture(scope="function")
def checkpointer_duckdb():
with DuckDBSaver.from_conn_string(":memory:") as checkpointer:
checkpointer.setup()
yield checkpointer
@asynccontextmanager
async def _checkpointer_duckdb_aio():
async with AsyncDuckDBSaver.from_conn_string(":memory:") as checkpointer:
await checkpointer.setup()
yield checkpointer
@pytest.fixture(scope="function")
def checkpointer_postgres():
database = f"test_{uuid4().hex[:16]}"
@@ -281,9 +264,6 @@ async def awith_checkpointer(
elif checkpointer_name == "sqlite_aio":
async with _checkpointer_sqlite_aio() as checkpointer:
yield checkpointer
elif checkpointer_name == "duckdb_aio":
async with _checkpointer_duckdb_aio() as checkpointer:
yield checkpointer
elif checkpointer_name == "postgres_aio":
async with _checkpointer_postgres_aio() as checkpointer:
yield checkpointer
@@ -370,13 +350,6 @@ async def _store_postgres_aio_pool():
await conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _store_duckdb_aio():
async with AsyncDuckDBStore.from_conn_string(":memory:") as store:
await store.setup()
yield store
@pytest.fixture(scope="function")
def store_postgres():
database = f"test_{uuid4().hex[:16]}"
@@ -433,13 +406,6 @@ def store_postgres_pool():
conn.execute(f"DROP DATABASE {database}")
@pytest.fixture(scope="function")
def store_duckdb():
with DuckDBStore.from_conn_string(":memory:") as store:
store.setup()
yield store
@pytest.fixture(scope="function")
def store_in_memory():
yield InMemoryStore()
@@ -460,9 +426,6 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
elif store_name == "postgres_aio_pool":
async with _store_postgres_aio_pool() as store:
yield store
elif store_name == "duckdb_aio":
async with _store_duckdb_aio() as store:
yield store
else:
raise NotImplementedError(f"Unknown store {store_name}")
@@ -500,12 +463,10 @@ ALL_STORES_SYNC = [
"postgres",
"postgres_pipe",
"postgres_pool",
"duckdb",
]
ALL_STORES_ASYNC = [
"in_memory",
"postgres_aio",
"postgres_aio_pipe",
"postgres_aio_pool",
"duckdb_aio",
]
+25 -1
View File
@@ -1,5 +1,6 @@
from langgraph.checkpoint.base import empty_checkpoint
from langgraph.pregel.algo import prepare_next_tasks
from langgraph.constants import PULL, PUSH
from langgraph.pregel.algo import _tuple_str, prepare_next_tasks
from langgraph.pregel.manager import ChannelsManager
@@ -40,3 +41,26 @@ def test_prepare_next_tasks() -> None:
)
# TODO: add more tests
def test_tuple_str() -> None:
push_path_a = (PUSH, 2)
pull_path_a = (PULL, "abc")
push_path_b = (PUSH, push_path_a, 1)
push_path_c = (PUSH, push_path_b, 3)
assert _tuple_str(push_path_a) == f"~{PUSH}, 0000000002"
assert _tuple_str(push_path_b) == f"~{PUSH}, ~{PUSH}, 0000000002, 0000000001"
assert (
_tuple_str(push_path_c)
== f"~{PUSH}, ~{PUSH}, ~{PUSH}, 0000000002, 0000000001, 0000000003"
)
assert _tuple_str(pull_path_a) == f"~{PULL}, abc"
path_list = [push_path_b, push_path_a, pull_path_a, push_path_c]
assert sorted(map(_tuple_str, path_list)) == [
f"~{PULL}, abc",
f"~{PUSH}, 0000000002",
f"~{PUSH}, ~{PUSH}, 0000000002, 0000000001",
f"~{PUSH}, ~{PUSH}, ~{PUSH}, 0000000002, 0000000001, 0000000003",
]
+175 -553
View File
@@ -17,7 +17,7 @@ from langgraph.channels.context import Context
from langgraph.channels.last_value import LastValue
from langgraph.channels.untracked_value import UntrackedValue
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import END, FF_SEND_V2, PULL, PUSH, START
from langgraph.constants import END, PULL, PUSH, START
from langgraph.errors import NodeInterrupt
from langgraph.graph import StateGraph
from langgraph.graph.graph import Graph
@@ -3022,9 +3022,6 @@ def test_state_graph_packets(
{"__interrupt__": ()},
]
if not FF_SEND_V2:
return
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"messages": [
@@ -3042,33 +3039,7 @@ def test_state_graph_packets(
),
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
content="",
additional_kwargs={},
response_metadata={},
id="ai1",
tool_calls=[
{
"name": "search_api",
"args": {"query": "query"},
"id": "tool_call123",
"type": "tool_call",
}
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
),
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config={
"configurable": {
@@ -3081,8 +3052,23 @@ def test_state_graph_packets(
metadata={
"parents": {},
"source": "loop",
"step": 0,
"writes": None,
"step": 1,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "search_api",
"args": {"query": "query"},
"id": "tool_call123",
"type": "tool_call",
}
],
)
}
},
"thread_id": "1",
},
parent_config=(
@@ -3117,7 +3103,7 @@ def test_state_graph_packets(
),
]
},
tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),),
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config={
"configurable": {
@@ -3130,7 +3116,7 @@ def test_state_graph_packets(
metadata={
"parents": {},
"source": "update",
"step": 1,
"step": 2,
"writes": {
"agent": {
"messages": AIMessage(
@@ -3228,16 +3214,26 @@ def test_state_graph_packets(
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
PregelTask(AnyStr(), "tools", (PUSH, 0)),
PregelTask(AnyStr(), "tools", (PUSH, 1)),
),
next=("tools", "tools"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 4,
"writes": {
"agent": {
"messages": AIMessage(
"",
content="",
id="ai2",
tool_calls=[
{
@@ -3255,31 +3251,6 @@ def test_state_graph_packets(
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)),
),
next=("tools", "tools"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 2,
"writes": {
"tools": {
"messages": _AnyIdToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
),
},
},
"thread_id": "1",
},
@@ -3335,7 +3306,7 @@ def test_state_graph_packets(
metadata={
"parents": {},
"source": "update",
"step": 3,
"step": 5,
"writes": {
"agent": {
"messages": AIMessage(content="answer", id="ai2"),
@@ -3401,31 +3372,7 @@ def test_state_graph_packets(
),
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
"",
id="ai1",
tool_calls=[
{
"name": "search_api",
"args": {"query": "query"},
"id": "tool_call123",
"type": "tool_call",
}
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
),
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config={
"configurable": {
@@ -3438,8 +3385,23 @@ def test_state_graph_packets(
metadata={
"parents": {},
"source": "loop",
"step": 0,
"writes": None,
"step": 1,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "search_api",
"args": {"query": "query"},
"id": "tool_call123",
"type": "tool_call",
}
],
)
}
},
"thread_id": "2",
},
parent_config=(
@@ -3474,14 +3436,14 @@ def test_state_graph_packets(
),
]
},
tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),),
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=AnyStr(),
metadata={
"parents": {},
"source": "update",
"step": 1,
"step": 2,
"writes": {
"agent": {
"messages": AIMessage(
@@ -3579,36 +3541,8 @@ def test_state_graph_packets(
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
"",
id="ai2",
tool_calls=[
{
"name": "search_api",
"args": {"query": "another", "idx": 0},
"id": "tool_call234",
"type": "tool_call",
},
{
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
"id": "tool_call567",
"type": "tool_call",
},
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)),
PregelTask(AnyStr(), "tools", (PUSH, 0)),
PregelTask(AnyStr(), "tools", (PUSH, 1)),
),
next=("tools", "tools"),
config={
@@ -3622,14 +3556,25 @@ def test_state_graph_packets(
metadata={
"parents": {},
"source": "loop",
"step": 2,
"step": 4,
"writes": {
"tools": {
"messages": _AnyIdToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
),
"agent": {
"messages": AIMessage(
id="ai2",
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another", "idx": 0},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
},
],
)
},
},
"thread_id": "2",
@@ -3686,7 +3631,7 @@ def test_state_graph_packets(
metadata={
"parents": {},
"source": "update",
"step": 3,
"step": 5,
"writes": {
"agent": {
"messages": AIMessage(content="answer", id="ai2"),
@@ -5860,7 +5805,6 @@ def test_dynamic_interrupt(
)
@pytest.mark.skipif(not FF_SEND_V2, reason="send v2 is not enabled")
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_copy_checkpoint(
request: pytest.FixtureRequest, checkpointer_name: str
@@ -5880,6 +5824,7 @@ def test_copy_checkpoint(
nonlocal tool_two_node_count
tool_two_node_count += 1
if s["market"] == "DE":
time.sleep(0.1)
answer = interrupt("Just because...")
else:
answer = " all good"
@@ -5909,7 +5854,7 @@ def test_copy_checkpoint(
assert run.outputs == {"market": "DE", "my_key": "value one"}
assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
"my_key": "value one all good",
"my_key": "value all good one",
"market": "US",
}
@@ -5940,6 +5885,7 @@ def test_copy_checkpoint(
]
# resume with answer
assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [
{"tool_one": {"my_key": " one"}, "__metadata__": {"cached": True}},
{"tool_two": {"my_key": " my answer"}},
]
@@ -5956,7 +5902,7 @@ def test_copy_checkpoint(
"parents": {},
"source": "loop",
"step": 0,
"writes": {"tool_one": {"my_key": " one"}},
"writes": None,
"thread_id": "1",
},
{
@@ -5972,6 +5918,12 @@ def test_copy_checkpoint(
values={"my_key": "value ⛰️ one", "market": "DE"},
next=("tool_two",),
tasks=(
PregelTask(
id=AnyStr(),
name="tool_one",
path=("__pregel_push", 0),
result={"my_key": " one"},
),
PregelTask(
AnyStr(),
"tool_two",
@@ -5997,7 +5949,7 @@ def test_copy_checkpoint(
"parents": {},
"source": "loop",
"step": 0,
"writes": {"tool_one": {"my_key": " one"}},
"writes": None,
"thread_id": "1",
},
parent_config=(
@@ -6006,13 +5958,25 @@ def test_copy_checkpoint(
else [*tool_two.checkpointer.list(thread1, limit=2)][-1].config
),
)
if "shallow" in checkpointer_name:
return
# clear the interrupt and next tasks
tool_two.update_state(thread1, None)
tool_two.update_state(thread1, None, as_node="__copy__")
# interrupt is cleared, next task is kept
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️ one", "market": "DE"},
next=("tool_two",),
values={"my_key": "value ⛰️", "market": "DE"},
next=(
"tool_one",
"tool_two",
),
tasks=(
PregelTask(
id=AnyStr(),
name="tool_one",
path=("__pregel_push", 0),
),
PregelTask(
AnyStr(),
"tool_two",
@@ -6030,15 +5994,13 @@ def test_copy_checkpoint(
created_at=AnyStr(),
metadata={
"parents": {},
"source": "update",
"source": "fork",
"step": 1,
"writes": {},
"writes": None,
"thread_id": "1",
},
parent_config=(
None
if "shallow" in checkpointer_name
else [*tool_two.checkpointer.list(thread1, limit=2)][-1].config
[*tool_two.checkpointer.list(thread1, limit=2)][-1].parent_config
),
)
@@ -7293,12 +7255,11 @@ def test_branch_then(
)
@pytest.mark.skip("TODO: re-enable in next PR")
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_send_dedupe_on_resume(
request: pytest.FixtureRequest, checkpointer_name: str
) -> None:
if not FF_SEND_V2:
pytest.skip("Send deduplication is only available in Send V2")
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
class InterruptOnce:
@@ -9321,328 +9282,18 @@ def test_send_to_nested_graphs(
# check state
outer_state = graph.get_state(config)
if not FF_SEND_V2:
# update state of dogs joke graph
graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"})
# continue past interrupt
assert sorted(
graph.stream(None, config=config),
key=lambda d: d["generate_joke"]["jokes"][0],
) == [
{"generate_joke": {"jokes": ["Joke about cats - hohoho"]}},
{"generate_joke": {"jokes": ["Joke about turtles - hohoho"]}},
]
return
assert outer_state == StateSnapshot(
values={"subjects": ["cats", "dogs"], "jokes": []},
tasks=(
PregelTask(
id=AnyStr(),
name="__start__",
path=("__pregel_pull", "__start__"),
error=None,
interrupts=(),
state=None,
result={"subjects": ["cats", "dogs"]},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 1),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 2),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
),
next=("generate_joke", "generate_joke"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "input",
"writes": {"__start__": {"subjects": ["cats", "dogs"]}},
"step": -1,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=None,
)
# check state of each of the inner tasks
assert graph.get_state(outer_state.tasks[1].state) == StateSnapshot(
values={"subject": "cats - hohoho", "jokes": []},
next=("generate",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("generate_joke:"): AnyStr(),
}
),
}
},
metadata={
"step": 1,
"source": "loop",
"writes": None,
"parents": {"": AnyStr()},
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"langgraph_checkpoint_ns": AnyStr("generate_joke:"),
"langgraph_node": "generate_joke",
"langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 1],
"langgraph_step": 0,
"langgraph_triggers": [PUSH],
},
created_at=AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("generate_joke:"): AnyStr(),
}
),
}
}
),
tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),),
)
assert graph.get_state(outer_state.tasks[2].state) == StateSnapshot(
values={"subject": "dogs - hohoho", "jokes": []},
next=("generate",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("generate_joke:"): AnyStr(),
}
),
}
},
metadata={
"step": 1,
"source": "loop",
"writes": None,
"parents": {"": AnyStr()},
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"langgraph_checkpoint_ns": AnyStr("generate_joke:"),
"langgraph_node": "generate_joke",
"langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 2],
"langgraph_step": 0,
"langgraph_triggers": [PUSH],
},
created_at=AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("generate_joke:"): AnyStr(),
}
),
}
}
),
tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),),
)
# update state of dogs joke graph
graph.update_state(
outer_state.tasks[2 if FF_SEND_V2 else 1].state, {"subject": "turtles - hohoho"}
)
graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"})
# continue past interrupt
assert sorted(
graph.stream(None, config=config), key=lambda d: d["generate_joke"]["jokes"][0]
graph.stream(None, config=config),
key=lambda d: d["generate_joke"]["jokes"][0],
) == [
{"generate_joke": {"jokes": ["Joke about cats - hohoho"]}},
{"generate_joke": {"jokes": ["Joke about turtles - hohoho"]}},
]
actual_snapshot = graph.get_state(config)
expected_snapshot = StateSnapshot(
values={
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"writes": {
"generate_joke": [
{"jokes": ["Joke about cats - hohoho"]},
{"jokes": ["Joke about turtles - hohoho"]},
]
},
"step": 0,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
)
assert actual_snapshot == expected_snapshot
if "shallow" in checkpointer_name:
return
# test full history
actual_history = list(graph.get_state_history(config))
# get subgraph node state for expected history
expected_history = [
StateSnapshot(
values={
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"writes": {
"generate_joke": [
{"jokes": ["Joke about cats - hohoho"]},
{"jokes": ["Joke about turtles - hohoho"]},
]
},
"step": 0,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
),
StateSnapshot(
values={"jokes": []},
tasks=(
PregelTask(
id=AnyStr(),
name="__start__",
path=("__pregel_pull", "__start__"),
error=None,
interrupts=(),
state=None,
result={"subjects": ["cats", "dogs"]},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 1),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
result={"jokes": ["Joke about cats - hohoho"]},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 2),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
result={"jokes": ["Joke about turtles - hohoho"]},
),
),
next=("__start__", "generate_joke", "generate_joke"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "input",
"writes": {"__start__": {"subjects": ["cats", "dogs"]}},
"step": -1,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=None,
),
]
assert actual_history == expected_history
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_send_react_interrupt(
@@ -9768,9 +9419,6 @@ def test_send_react_interrupt(
}
assert foo_called == 0
if not FF_SEND_V2:
return
# get state should show the pending task
state = graph.get_state(thread1)
assert state == StateSnapshot(
@@ -9799,9 +9447,24 @@ def test_send_react_interrupt(
}
},
metadata={
"step": 0,
"step": 1,
"source": "loop",
"writes": None,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
}
},
"parents": {},
"thread_id": "2",
},
@@ -9818,34 +9481,10 @@ def test_send_react_interrupt(
}
),
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
content="",
additional_kwargs={},
response_metadata={},
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
},
),
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", ("__pregel_pull", "agent"), 2),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -9879,7 +9518,7 @@ def test_send_react_interrupt(
}
},
metadata={
"step": 1,
"step": 2,
"source": "update",
"writes": {
"agent": {
@@ -9966,9 +9605,26 @@ def test_send_react_interrupt(
}
},
metadata={
"step": 0,
"step": 1,
"source": "loop",
"writes": None,
"writes": {
"agent": {
"messages": AIMessage(
content="",
additional_kwargs={},
response_metadata={},
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
}
},
"parents": {},
"thread_id": "3",
},
@@ -9985,32 +9641,10 @@ def test_send_react_interrupt(
}
),
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
"",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
},
),
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", ("__pregel_pull", "agent"), 2),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -10065,7 +9699,7 @@ def test_send_react_interrupt(
}
},
metadata={
"step": 1,
"step": 2,
"source": "update",
"writes": {
"agent": {
@@ -10101,7 +9735,7 @@ def test_send_react_interrupt(
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", (), 0),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -10231,9 +9865,6 @@ def test_send_react_interrupt_control(
}
assert foo_called == 1
if not FF_SEND_V2:
return
# interrupt-update-resume flow
foo_called = 0
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"])
@@ -10284,9 +9915,24 @@ def test_send_react_interrupt_control(
}
},
metadata={
"step": 0,
"step": 1,
"source": "loop",
"writes": None,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
}
},
"parents": {},
"thread_id": "2",
},
@@ -10303,34 +9949,10 @@ def test_send_react_interrupt_control(
}
),
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
content="",
additional_kwargs={},
response_metadata={},
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
},
),
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", ("__pregel_pull", "agent"), 2),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -10364,7 +9986,7 @@ def test_send_react_interrupt_control(
}
},
metadata={
"step": 1,
"step": 2,
"source": "update",
"writes": {
"agent": {
+76 -351
View File
@@ -25,7 +25,7 @@ from typing_extensions import TypedDict
from langgraph.channels.context import Context
from langgraph.channels.last_value import LastValue
from langgraph.channels.untracked_value import UntrackedValue
from langgraph.constants import END, FF_SEND_V2, PULL, PUSH, START
from langgraph.constants import END, PULL, PUSH, START
from langgraph.graph.graph import Graph
from langgraph.graph.message import MessageGraph, add_messages
from langgraph.graph.state import StateGraph
@@ -2740,9 +2740,6 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
{"__interrupt__": ()},
]
if not FF_SEND_V2:
return
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"messages": [
@@ -2760,31 +2757,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
),
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
"",
id="ai1",
tool_calls=[
{
"name": "search_api",
"args": {"query": "query"},
"id": "tool_call123",
"type": "tool_call",
}
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
),
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config={
"configurable": {
@@ -2797,8 +2770,23 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
metadata={
"parents": {},
"source": "loop",
"step": 0,
"writes": None,
"step": 1,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "search_api",
"args": {"query": "query"},
"id": "tool_call123",
"type": "tool_call",
}
],
)
}
},
"thread_id": "1",
},
parent_config=(
@@ -2834,14 +2822,14 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
),
]
},
tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),),
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "update",
"step": 1,
"step": 2,
"writes": {
"agent": {
"messages": AIMessage(
@@ -2941,36 +2929,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
"",
id="ai2",
tool_calls=[
{
"name": "search_api",
"args": {"query": "another", "idx": 0},
"id": "tool_call234",
"type": "tool_call",
},
{
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
"id": "tool_call567",
"type": "tool_call",
},
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)),
PregelTask(AnyStr(), "tools", (PUSH, 0)),
PregelTask(AnyStr(), "tools", (PUSH, 1)),
),
next=("tools", "tools"),
config=tup.config,
@@ -2978,13 +2938,24 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
metadata={
"parents": {},
"source": "loop",
"step": 2,
"step": 4,
"writes": {
"tools": {
"messages": _AnyIdToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
"agent": {
"messages": AIMessage(
id="ai2",
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another", "idx": 0},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
},
],
),
},
},
@@ -3036,7 +3007,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
metadata={
"parents": {},
"source": "update",
"step": 3,
"step": 5,
"writes": {
"agent": {
"messages": AIMessage(content="answer", id="ai2"),
@@ -3103,15 +3074,16 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
),
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "loop",
"step": 1,
"writes": {
"agent": {
"messages": AIMessage(
content="",
additional_kwargs={},
@@ -3126,18 +3098,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
}
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
),
next=("tools",),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "loop",
"step": 0,
"writes": None,
}
},
"thread_id": "2",
},
parent_config=(
@@ -3173,14 +3135,14 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
),
]
},
tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),),
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "update",
"step": 1,
"step": 2,
"writes": {
"agent": {
"messages": AIMessage(
@@ -3280,38 +3242,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
content="",
additional_kwargs={},
response_metadata={},
id="ai2",
tool_calls=[
{
"name": "search_api",
"args": {"query": "another", "idx": 0},
"id": "tool_call234",
"type": "tool_call",
},
{
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
"id": "tool_call567",
"type": "tool_call",
},
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)),
PregelTask(AnyStr(), "tools", (PUSH, 0)),
PregelTask(AnyStr(), "tools", (PUSH, 1)),
),
next=("tools", "tools"),
config=tup.config,
@@ -3319,13 +3251,24 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
metadata={
"parents": {},
"source": "loop",
"step": 2,
"step": 4,
"writes": {
"tools": {
"messages": _AnyIdToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
"agent": {
"messages": AIMessage(
id="ai2",
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another", "idx": 0},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
},
],
),
},
},
@@ -3377,7 +3320,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
metadata={
"parents": {},
"source": "update",
"step": 3,
"step": 5,
"writes": {
"agent": {
"messages": AIMessage(content="answer", id="ai2"),
@@ -6961,83 +6904,9 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
# check state
outer_state = await graph.aget_state(config)
if not FF_SEND_V2:
# update state of dogs joke graph
await graph.aupdate_state(
outer_state.tasks[1].state, {"subject": "turtles - hohoho"}
)
# continue past interrupt
assert await graph.ainvoke(None, config=config) == {
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
}
return
assert outer_state == StateSnapshot(
values={"subjects": ["cats", "dogs"], "jokes": []},
tasks=(
PregelTask(
id=AnyStr(),
name="__start__",
path=("__pregel_pull", "__start__"),
error=None,
interrupts=(),
state=None,
result={"subjects": ["cats", "dogs"]},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 1),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 2),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
),
next=("generate_joke", "generate_joke"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "input",
"writes": {
"__start__": {
"subjects": [
"cats",
"dogs",
],
}
},
"step": -1,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=None,
)
# update state of dogs joke graph
await graph.aupdate_state(
outer_state.tasks[2].state, {"subject": "turtles - hohoho"}
outer_state.tasks[1].state, {"subject": "turtles - hohoho"}
)
# continue past interrupt
@@ -7046,150 +6915,6 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
}
actual_snapshot = await graph.aget_state(config)
expected_snapshot = StateSnapshot(
values={
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"writes": {
"generate_joke": [
{"jokes": ["Joke about cats - hohoho"]},
{"jokes": ["Joke about turtles - hohoho"]},
]
},
"step": 0,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
)
assert actual_snapshot == expected_snapshot
if "shallow" in checkpointer_name:
return
# test full history
actual_history = [c async for c in graph.aget_state_history(config)]
expected_history = [
StateSnapshot(
values={
"subjects": ["cats", "dogs"],
"jokes": [
"Joke about cats - hohoho",
"Joke about turtles - hohoho",
],
},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"writes": {
"generate_joke": [
{"jokes": ["Joke about cats - hohoho"]},
{"jokes": ["Joke about turtles - hohoho"]},
]
},
"step": 0,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
),
StateSnapshot(
values={"jokes": []},
next=("__start__", "generate_joke", "generate_joke"),
tasks=(
PregelTask(
id=AnyStr(),
name="__start__",
path=("__pregel_pull", "__start__"),
error=None,
interrupts=(),
state=None,
result={"subjects": ["cats", "dogs"]},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 1),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
result={"jokes": ["Joke about cats - hohoho"]},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 2),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
result={"jokes": ["Joke about turtles - hohoho"]},
),
),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "input",
"writes": {"__start__": {"subjects": ["cats", "dogs"]}},
"step": -1,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=None,
),
]
assert actual_history == expected_history
@pytest.mark.skipif(
sys.version_info < (3, 11),
+23 -63
View File
@@ -50,13 +50,7 @@ from langgraph.checkpoint.base import (
CheckpointTuple,
)
from langgraph.checkpoint.memory import MemorySaver
from langgraph.constants import (
CONFIG_KEY_NODE_FINISHED,
ERROR,
FF_SEND_V2,
PULL,
START,
)
from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START
from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError
from langgraph.func import entrypoint, task
from langgraph.graph import END, Graph, StateGraph
@@ -1375,31 +1369,17 @@ def test_concurrent_emit_sends() -> None:
builder.add_conditional_edges("1.1", send_for_profit)
builder.add_conditional_edges("2", route_to_three)
graph = builder.compile()
assert graph.invoke(["0"]) == (
[
"0",
"1",
"1.1",
"2|1",
"2|2",
"2|3",
"2|4",
"3",
"3.1",
]
if FF_SEND_V2
else [
"0",
"1",
"1.1",
"3.1",
"2|1",
"2|2",
"2|3",
"2|4",
"3",
]
)
assert graph.invoke(["0"]) == [
"0",
"1",
"1.1",
"3.1",
"2|1",
"2|2",
"2|3",
"2|4",
"3",
]
def test_send_sequences() -> None:
@@ -1438,31 +1418,17 @@ def test_send_sequences() -> None:
builder.add_conditional_edges("1", send_for_fun)
builder.add_conditional_edges("2", route_to_three)
graph = builder.compile()
assert (
graph.invoke(["0"])
== [
"0",
"1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"2|3",
"2|4",
"3",
"3.1",
]
if FF_SEND_V2
else [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"3",
"2|3",
"2|4",
"3",
]
)
assert graph.invoke(["0"]) == [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"3",
"2|3",
"2|4",
"3",
]
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
@@ -4193,12 +4159,6 @@ def test_store_injected(
builder.add_edge("__start__", "node")
N = 500
M = 1
if "duckdb" in store_name:
logger.warning(
"DuckDB store implementation has a known issue that does not"
" support concurrent writes, so we're reducing the test scope"
)
N = M = 1
for i in range(N):
builder.add_node(f"node_{i}", Node(i))
+256 -233
View File
@@ -47,13 +47,7 @@ from langgraph.checkpoint.base import (
CheckpointTuple,
)
from langgraph.checkpoint.memory import MemorySaver
from langgraph.constants import (
CONFIG_KEY_NODE_FINISHED,
ERROR,
FF_SEND_V2,
PULL,
START,
)
from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START
from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
from langgraph.func import entrypoint, task
from langgraph.graph import END, Graph, StateGraph
@@ -878,7 +872,6 @@ async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None:
)
@pytest.mark.skipif(not FF_SEND_V2, reason="send v2 is not enabled")
@pytest.mark.skipif(
sys.version_info < (3, 11),
reason="Python 3.11+ is required for async contextvars support",
@@ -914,7 +907,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
tracer = FakeTracer()
assert await tool_two.ainvoke(
{"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}
{"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}, debug=True
) == {
"my_key": "value one",
"market": "DE",
@@ -927,7 +920,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
assert run.outputs == {"market": "DE", "my_key": "value one"}
assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == {
"my_key": "value one all good",
"my_key": "value all good one",
"market": "US",
}
@@ -964,6 +957,10 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
assert [
c async for c in tool_two.astream(Command(resume=" my answer"), thread2)
] == [
{
"__metadata__": {"cached": True},
"tool_one": {"my_key": " one"},
},
{"tool_two": {"my_key": " my answer"}},
]
@@ -983,7 +980,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
"parents": {},
"source": "loop",
"step": 0,
"writes": {"tool_one": {"my_key": " one"}},
"writes": None,
"thread_id": "1",
},
{
@@ -1000,6 +997,15 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
values={"my_key": "value ⛰️ one", "market": "DE"},
next=("tool_two",),
tasks=(
PregelTask(
AnyStr(),
name="tool_one",
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
result={"my_key": " one"},
),
PregelTask(
AnyStr(),
"tool_two",
@@ -1019,7 +1025,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
"parents": {},
"source": "loop",
"step": 0,
"writes": {"tool_one": {"my_key": " one"}},
"writes": None,
"thread_id": "1",
},
parent_config=(
@@ -1030,14 +1036,25 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
].config
),
)
if "shallow" in checkpointer_name:
# shallow checkpointer doesn't support copy
return
# clear the interrupt and next tasks
await tool_two.aupdate_state(thread1, None)
await tool_two.aupdate_state(thread1, None, as_node="__copy__")
# interrupt is cleared, next task is kept
tup = await tool_two.checkpointer.aget_tuple(thread1)
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️ one", "market": "DE"},
next=("tool_two",),
values={"my_key": "value ⛰️", "market": "DE"},
next=("tool_one", "tool_two"),
tasks=(
PregelTask(
AnyStr(),
"tool_one",
(PUSH, 0),
result=None,
),
PregelTask(
AnyStr(),
"tool_two",
@@ -1049,17 +1066,15 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "update",
"source": "fork",
"step": 1,
"writes": {},
"writes": None,
"thread_id": "1",
},
parent_config=(
None
if "shallow" in checkpointer_name
else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
-1
].config
].parent_config
),
)
@@ -2345,18 +2360,6 @@ async def test_concurrent_emit_sends() -> None:
graph = builder.compile()
assert await graph.ainvoke(["0"]) == (
[
"0",
"1",
"1.1",
"2|1",
"2|2",
"2|3",
"2|4",
"3",
"3.1",
]
if FF_SEND_V2
else [
"0",
"1",
"1.1",
@@ -2370,6 +2373,7 @@ async def test_concurrent_emit_sends() -> None:
)
@pytest.mark.skip("TODO: re-enable in next PR")
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_send_sequences(checkpointer_name: str) -> None:
class Node:
@@ -2407,34 +2411,17 @@ async def test_send_sequences(checkpointer_name: str) -> None:
builder.add_conditional_edges("1", send_for_fun)
builder.add_conditional_edges("2", route_to_three)
graph = builder.compile()
assert (
await graph.ainvoke(["0"])
== [
"0",
"1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"2|3",
"2|4",
"3",
"3.1",
]
if FF_SEND_V2
else [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"3",
"2|3",
"2|4",
"3",
]
)
if not FF_SEND_V2:
return
assert await graph.ainvoke(["0"]) == [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"3",
"2|3",
"2|4",
"3",
]
async with awith_checkpointer(checkpointer_name) as checkpointer:
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["3.1"])
@@ -2442,20 +2429,17 @@ async def test_send_sequences(checkpointer_name: str) -> None:
assert await graph.ainvoke(["0"], thread1) == [
"0",
"1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"2|3",
"2|4",
]
assert await graph.ainvoke(None, thread1) == [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"3",
"2|3",
"2|4",
"3",
"3.1",
]
@@ -2472,6 +2456,7 @@ async def test_imp_task(checkpointer_name: str) -> None:
async def mapper(input: int) -> str:
nonlocal mapper_calls
mapper_calls += 1
await asyncio.sleep(0.1 * input)
return str(input) * 2
@entrypoint(checkpointer=checkpointer)
@@ -2481,7 +2466,8 @@ async def test_imp_task(checkpointer_name: str) -> None:
answer = interrupt("question")
return [m + answer for m in mapped]
thread1 = {"configurable": {"thread_id": "1"}}
tracer = FakeTracer()
thread1 = {"configurable": {"thread_id": "1"}, "callbacks": [tracer]}
assert [c async for c in graph.astream([0, 1], thread1)] == [
{"mapper": "00"},
{"mapper": "11"},
@@ -2497,6 +2483,9 @@ async def test_imp_task(checkpointer_name: str) -> None:
},
]
assert mapper_calls == 2
assert len(tracer.runs) == 1
assert len(tracer.runs[0].child_runs) == 1
assert tracer.runs[0].child_runs[0].name == "graph"
assert await graph.ainvoke(Command(resume="answer"), thread1) == [
"00answer",
@@ -2632,11 +2621,9 @@ async def test_imp_stream_order(checkpointer_name: str) -> None:
]
@pytest.mark.skip("TODO: re-enable in next PR")
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
if not FF_SEND_V2:
pytest.skip("Send deduplication is only available in Send V2")
class InterruptOnce:
ticks: int = 0
@@ -2690,22 +2677,26 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
assert await graph.ainvoke(["0"], thread1, debug=1) == [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"3",
"2|3",
]
assert builder.nodes["2"].runnable.func.ticks == 3
assert builder.nodes["flaky"].runnable.func.ticks == 1
print((await graph.aget_state(thread1)).tasks)
# resume execution
assert await graph.ainvoke(None, thread1, debug=1) == [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
"3.1",
"flaky|4",
"2|3",
"3",
]
# node "2" doesn't get called again, as we recover writes saved before
assert builder.nodes["2"].runnable.func.ticks == 3
@@ -2718,12 +2709,13 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
values=[
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
"3.1",
"flaky|4",
"2|3",
"3",
],
next=(),
config={
@@ -2735,9 +2727,9 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
},
metadata={
"source": "loop",
"writes": {"3": ["3"], "3.1": ["3.1"]},
"writes": {"3": ["3"]},
"thread_id": "1",
"step": 2,
"step": 4,
"parents": {},
},
created_at=AnyStr(),
@@ -2754,12 +2746,14 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
values=[
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"2|3",
"3",
"flaky|4",
"2|3",
],
next=("3", "3.1"),
next=("3",),
config={
"configurable": {
"thread_id": "1",
@@ -2769,17 +2763,9 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
},
metadata={
"source": "loop",
"writes": {
"1": ["1"],
"2": [
["2|Command(goto=Send(node='2', arg=3))"],
["2|Command(goto=Send(node='flaky', arg=4))"],
["2|3"],
],
"flaky": ["flaky|4"],
},
"writes": {"2": ["2|3"], "3": ["3"], "flaky": ["flaky|4"]},
"thread_id": "1",
"step": 1,
"step": 3,
"parents": {},
},
created_at=AnyStr(),
@@ -2800,6 +2786,123 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
state=None,
result=["3"],
),
),
),
StateSnapshot(
values=[
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
],
next=("2", "flaky", "3"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"2": [
["2|Command(goto=Send(node='2', arg=3))"],
["2|Command(goto=Send(node='flaky', arg=4))"],
],
"3.1": ["3.1"],
},
"thread_id": "1",
"step": 2,
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
tasks=(
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
result=["2|3"],
),
PregelTask(
id=AnyStr(),
name="flaky",
path=("__pregel_push", 1),
error=None,
interrupts=(
Interrupt(
value="Bahh", resumable=False, ns=None, when="during"
),
),
state=None,
result=["flaky|4"],
),
PregelTask(
id=AnyStr(),
name="3",
path=("__pregel_pull", "3"),
error=None,
interrupts=(),
state=None,
result=["3"],
),
),
),
StateSnapshot(
values=["0", "1"],
next=("2", "2", "3.1"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"1": ["1"]},
"thread_id": "1",
"step": 1,
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
tasks=(
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
result=["2|Command(goto=Send(node='2', arg=3))"],
),
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", 1),
error=None,
interrupts=(),
state=None,
result=["2|Command(goto=Send(node='flaky', arg=4))"],
),
PregelTask(
id=AnyStr(),
name="3.1",
@@ -2813,7 +2916,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
),
StateSnapshot(
values=["0"],
next=("1", "2", "2", "2", "flaky"),
next=("1",),
config={
"configurable": {
"thread_id": "1",
@@ -2846,50 +2949,6 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
state=None,
result=["1"],
),
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", ("__pregel_pull", "1"), 2),
error=None,
interrupts=(),
state=None,
result=["2|Command(goto=Send(node='2', arg=3))"],
),
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", ("__pregel_pull", "1"), 3),
error=None,
interrupts=(),
state=None,
result=["2|Command(goto=Send(node='flaky', arg=4))"],
),
PregelTask(
id=AnyStr(),
name="2",
path=(
"__pregel_push",
("__pregel_push", ("__pregel_pull", "1"), 2),
2,
),
error=None,
interrupts=(),
state=None,
result=["2|3"],
),
PregelTask(
id=AnyStr(),
name="flaky",
path=(
"__pregel_push",
("__pregel_push", ("__pregel_pull", "1"), 3),
2,
),
error=None,
interrupts=(Interrupt(value="Bahh", when="during"),),
state=None,
result=["flaky|4"],
),
),
),
StateSnapshot(
@@ -3047,9 +3106,6 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
assert foo_called == 0
if not FF_SEND_V2:
return
# get state should show the pending task
state = await graph.aget_state(thread1)
assert state == StateSnapshot(
@@ -3078,9 +3134,24 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
},
metadata={
"step": 0,
"step": 1,
"source": "loop",
"writes": None,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
}
},
"parents": {},
"thread_id": "2",
},
@@ -3097,34 +3168,10 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
),
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
content="",
additional_kwargs={},
response_metadata={},
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
},
),
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", ("__pregel_pull", "agent"), 2),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -3158,7 +3205,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
},
metadata={
"step": 1,
"step": 2,
"source": "update",
"writes": {
"agent": {
@@ -3245,9 +3292,24 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
},
metadata={
"step": 0,
"step": 1,
"source": "loop",
"writes": None,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
}
},
"parents": {},
"thread_id": "3",
},
@@ -3264,32 +3326,10 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
),
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
"",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
},
),
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", ("__pregel_pull", "agent"), 2),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -3344,7 +3384,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
},
metadata={
"step": 1,
"step": 2,
"source": "update",
"writes": {
"agent": {
@@ -3380,7 +3420,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", (), 0),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -3531,9 +3571,6 @@ async def test_send_react_interrupt_control(
}
assert foo_called == 0
if not FF_SEND_V2:
return
# get state should show the pending task
state = await graph.aget_state(thread1)
assert state == StateSnapshot(
@@ -3562,9 +3599,24 @@ async def test_send_react_interrupt_control(
}
},
metadata={
"step": 0,
"step": 1,
"source": "loop",
"writes": None,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
}
},
"parents": {},
"thread_id": "2",
},
@@ -3581,34 +3633,10 @@ async def test_send_react_interrupt_control(
}
),
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
content="",
additional_kwargs={},
response_metadata={},
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
},
),
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", ("__pregel_pull", "agent"), 2),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -3642,7 +3670,7 @@ async def test_send_react_interrupt_control(
}
},
metadata={
"step": 1,
"step": 2,
"source": "update",
"writes": {
"agent": {
@@ -5722,20 +5750,15 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) ->
N = 500
M = 1
if "duckdb" in store_name:
logger.warning(
"DuckDB store implementation has a known issue that does not"
" support concurrent writes, so we're reducing the test scope"
)
N = M = 1
for i in range(N):
builder.add_node(f"node_{i}", Node(i))
builder.add_edge("__start__", f"node_{i}")
async with awith_checkpointer(checkpointer_name) as checkpointer, awith_store(
store_name
) as the_store:
async with (
awith_checkpointer(checkpointer_name) as checkpointer,
awith_store(store_name) as the_store,
):
graph = builder.compile(store=the_store, checkpointer=checkpointer)
# Test batch operations with multiple threads
+11 -1
View File
@@ -79,11 +79,19 @@ def test_state_schema_with_type_hint():
def pre_foo(_) -> FooState:
return {"foo": "bar"}
def pre_bar(_) -> FooState:
return {"foo": "bar"}
class Foo:
def __call__(self, state: FooState) -> OutputState:
assert state.pop("foo") == "bar"
return {"input_state": state}
class Bar:
def my_node(self, state: FooState) -> OutputState:
assert state.pop("foo") == "bar"
return {"input_state": state}
graph = StateGraph(InputState, output=OutputState)
actions = [
complete_hint,
@@ -92,6 +100,8 @@ def test_state_schema_with_type_hint():
miss_all_hint,
pre_foo,
Foo(),
pre_bar,
Bar().my_node,
]
for action in actions:
@@ -112,7 +122,7 @@ def test_state_schema_with_type_hint():
foo_state = FooState(foo="bar")
for i, c in enumerate(graph.stream(input_state, stream_mode="updates")):
node_name = get_name(actions[i])
if node_name == get_name(pre_foo):
if node_name in {"pre_foo", "pre_bar"}:
assert c[node_name] == foo_state
else:
assert c[node_name] == output_state
+12
View File
@@ -20,6 +20,7 @@ from typing_extensions import Annotated, NotRequired, Required, TypedDict
from langgraph.graph import END, StateGraph
from langgraph.graph.graph import CompiledGraph
from langgraph.utils.config import _is_not_empty
from langgraph.utils.fields import (
_is_optional_type,
get_enhanced_type_hints,
@@ -284,3 +285,14 @@ def test_enhanced_type_hints() -> None:
assert hints[0] == ("val_1", str, None, "A description")
assert hints[1] == ("val_2", int, 42, None)
assert hints[2] == ("val_3", str, "default", "Another description")
def test_is_not_empty() -> None:
assert _is_not_empty("foo")
assert _is_not_empty("")
assert _is_not_empty(1)
assert _is_not_empty(0)
assert not _is_not_empty(None)
assert not _is_not_empty([])
assert not _is_not_empty(())
assert not _is_not_empty({})
+2 -4
View File
@@ -9,7 +9,7 @@ import pytest
from aiokafka import AIOKafkaProducer
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import FF_SEND_V2, START
from langgraph.constants import START
from langgraph.errors import NodeInterrupt
from langgraph.graph.state import CompiledStateGraph, StateGraph
from langgraph.scheduler.kafka import serde
@@ -76,10 +76,8 @@ def mk_push_graph(
return builder.compile(checkpointer=checkpointer)
@pytest.mark.skip("TODO: re-enable in next PR")
async def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
if not FF_SEND_V2:
pytest.skip("Test requires FF_SEND_V2")
input = ["0"]
config = {"configurable": {"thread_id": "1"}}
graph = mk_push_graph(acheckpointer)
+2 -4
View File
@@ -8,7 +8,7 @@ from typing import (
import pytest
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import FF_SEND_V2, START
from langgraph.constants import START
from langgraph.errors import NodeInterrupt
from langgraph.graph.state import CompiledStateGraph, StateGraph
from langgraph.scheduler.kafka import serde
@@ -76,10 +76,8 @@ def mk_push_graph(
return builder.compile(checkpointer=checkpointer)
@pytest.mark.skip("TODO: re-enable in next PR")
def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
if not FF_SEND_V2:
pytest.skip("Test requires FF_SEND_V2")
input = ["0"]
config = {"configurable": {"thread_id": "1"}}
graph = mk_push_graph(acheckpointer)