mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-29 11:19:54 +02:00
Compare commits
58
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2c984dfe4e | ||
|
|
09c0d9cec3 | ||
|
|
31734fb792 | ||
|
|
eed577ee2a | ||
|
|
bba00506ea | ||
|
|
fa12538a4e | ||
|
|
c1a7bb8902 | ||
|
|
14b500d8a8 | ||
|
|
053a501db3 | ||
|
|
bbe8eacf8d | ||
|
|
50cb387904 | ||
|
|
003b1883c9 | ||
|
|
d12830e95c | ||
|
|
4426552c26 | ||
|
|
e2554c9616 | ||
|
|
f1ee650489 | ||
|
|
b7e4656c91 | ||
|
|
65a41942ef | ||
|
|
c767d86c9c | ||
|
|
01331ef858 | ||
|
|
24a4c67c52 | ||
|
|
c7e43f86d4 | ||
|
|
b18d266f2c | ||
|
|
536ee7b37a | ||
|
|
901273c0e2 | ||
|
|
7bdbd62611 | ||
|
|
0b74e25f72 | ||
|
|
ff843ab005 | ||
|
|
f7fae7e140 | ||
|
|
651ee8bd24 | ||
|
|
54bdba2da9 | ||
|
|
7e5806cb0b | ||
|
|
66f674fead | ||
|
|
f029d615e6 | ||
|
|
bd76773b31 | ||
|
|
9b6e6d67dc | ||
|
|
ecfb3e3850 | ||
|
|
ec71cd6fc1 | ||
|
|
2db6d9d7b5 | ||
|
|
0f5df797af | ||
|
|
5dc3f14e8b | ||
|
|
f25ea9ecaf | ||
|
|
9940a9f833 | ||
|
|
d175d5e4a8 | ||
|
|
fe63440d98 | ||
|
|
a9073241b2 | ||
|
|
0a04262f05 | ||
|
|
381796c474 | ||
|
|
09ae5ce05d | ||
|
|
59e7751a8c | ||
|
|
b9816e31ff | ||
|
|
3ea7f9469a | ||
|
|
24e61fafa9 | ||
|
|
2b370cf3b7 | ||
|
|
839bea0f33 | ||
|
|
fba5ca7ab9 | ||
|
|
f141089811 | ||
|
|
38487cb158 |
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:.*" \
|
||||
|
||||
@@ -1,6 +1,4 @@
|
||||
---
|
||||
hide:
|
||||
- navigation
|
||||
title: Concepts
|
||||
description: Conceptual Guide for LangGraph
|
||||
---
|
||||
|
||||
@@ -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?
|
||||
|
||||
|
||||
@@ -1,6 +1,4 @@
|
||||
---
|
||||
hide:
|
||||
- navigation
|
||||
title: How-to Guides
|
||||
description: How to accomplish common tasks in LangGraph
|
||||
---
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
---
|
||||
hide_comments: true
|
||||
hide:
|
||||
- navigation
|
||||
title: Home
|
||||
---
|
||||
|
||||
|
||||
@@ -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.
|
||||
@@ -1,6 +1,4 @@
|
||||
---
|
||||
hide:
|
||||
- navigation
|
||||
title: Tutorials
|
||||
---
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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 %}
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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("."))
|
||||
Generated
-1058
File diff suppressed because it is too large
Load Diff
@@ -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"
|
||||
@@ -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]}")
|
||||
@@ -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]}")
|
||||
@@ -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,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"
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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.
|
||||
|
||||
Generated
+162
-84
File diff suppressed because it is too large
Load Diff
@@ -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'}}}%%
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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": {
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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({})
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user