mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-31 12:19:58 +02:00
Compare commits
72
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
771c6150a4 | ||
|
|
edfb65fd3a | ||
|
|
0f92470e49 | ||
|
|
dfcaf97c73 | ||
|
|
63a0028372 | ||
|
|
1134017d07 | ||
|
|
33feba4877 | ||
|
|
4fec8e9dec | ||
|
|
c137169325 | ||
|
|
1e2672e63d | ||
|
|
06803ab683 | ||
|
|
3488ee47e0 | ||
|
|
289bdd0cea | ||
|
|
417103066b | ||
|
|
25a59447c1 | ||
|
|
21906d2b7b | ||
|
|
0cad7019cb | ||
|
|
7e735672bf | ||
|
|
5498893780 | ||
|
|
e80f47aa01 | ||
|
|
a0b2f742a3 | ||
|
|
b7973d65db | ||
|
|
3fa3a586b5 | ||
|
|
053b606b46 | ||
|
|
4548a0ebe8 | ||
|
|
0171e9a323 | ||
|
|
c439cb0872 | ||
|
|
2a4d7e8889 | ||
|
|
7f3578e0f1 | ||
|
|
e2f96b5ae5 | ||
|
|
0d5f7e55bf | ||
|
|
9209f11187 | ||
|
|
bb1c5b8cdf | ||
|
|
d6bb008ff4 | ||
|
|
6130e08fa6 | ||
|
|
3ad061f0d7 | ||
|
|
116b5d1cac | ||
|
|
0aff02e180 | ||
|
|
074af5c122 | ||
|
|
29ffaa0e0b | ||
|
|
45cd4e1928 | ||
|
|
480271f753 | ||
|
|
66fdf60e47 | ||
|
|
0894daf3fc | ||
|
|
850c55d630 | ||
|
|
c0d65ff409 | ||
|
|
be7b60a722 | ||
|
|
d467ec6556 | ||
|
|
b8683ab67a | ||
|
|
6a9ca8d67e | ||
|
|
3b98044f2f | ||
|
|
a4a8934bd3 | ||
|
|
470b9a4b97 | ||
|
|
516175780d | ||
|
|
571780f74c | ||
|
|
d719438307 | ||
|
|
85c809a651 | ||
|
|
0441fd156f | ||
|
|
37b5d3886c | ||
|
|
b95267a3cc | ||
|
|
2e33c520a5 | ||
|
|
67b1dc602e | ||
|
|
1519b90414 | ||
|
|
0035ab9825 | ||
|
|
c42cd57a32 | ||
|
|
acc56e094a | ||
|
|
6b30d4fd8f | ||
|
|
fcc37cd06b | ||
|
|
c17ee1bf5a | ||
|
|
88c603b00b | ||
|
|
c12f7cb2b9 | ||
|
|
6d7d689578 |
@@ -1,6 +1,6 @@
|
||||
name: "\U0001F41B Bug Report"
|
||||
description: Report a bug in LangGraph. To report a security issue, please instead use the security option below. For questions, please use the GitHub Discussions.
|
||||
labels: ["02 Bug Report"]
|
||||
labels: [pending,bug]
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
blank_issues_enabled: false
|
||||
blank_issues_enabled: true
|
||||
version: 2.1
|
||||
contact_links:
|
||||
- name: 🤔 Question or Problem
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
name: Documentation
|
||||
description: Report an issue related to the LangGraph documentation.
|
||||
title: "DOC: <Please write a comprehensive title after the 'DOC: ' prefix>"
|
||||
labels: [03 - Documentation]
|
||||
labels: [documentation]
|
||||
|
||||
body:
|
||||
- type: textarea
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
# AGENTS Instructions
|
||||
|
||||
This repository is a monorepo. Each library lives in a subdirectory under `libs/`.
|
||||
|
||||
When you modify code in any library, run the following commands in that library's directory before creating a pull request:
|
||||
|
||||
- `make format` – run code formatters
|
||||
- `make lint` – run the linter
|
||||
- `make test` – execute the test suite
|
||||
|
||||
To run a particular test file or to pass additional pytest options you can specify the `TEST` variable:
|
||||
|
||||
```
|
||||
TEST=path/to/test.py make test
|
||||
```
|
||||
|
||||
Other pytest arguments can also be supplied inside the `TEST` variable.
|
||||
|
||||
## Libraries
|
||||
|
||||
The repository contains several Python and JavaScript/TypeScript libraries.
|
||||
Below is a high-level overview:
|
||||
|
||||
- **checkpoint** – base interfaces for LangGraph checkpointers.
|
||||
- **checkpoint-postgres** – Postgres implementation of the checkpoint saver.
|
||||
- **checkpoint-sqlite** – SQLite implementation of the checkpoint saver.
|
||||
- **cli** – official command-line interface for LangGraph.
|
||||
- **langgraph** – core framework for building stateful, multi-actor agents.
|
||||
- **prebuilt** – high-level APIs for creating and running agents and tools.
|
||||
- **sdk-js** – JS/TS SDK for interacting with the LangGraph REST API.
|
||||
- **sdk-py** – Python SDK for the LangGraph Platform API.
|
||||
|
||||
### Dependency map
|
||||
|
||||
The diagram below lists downstream libraries for each production dependency as
|
||||
declared in that library's `pyproject.toml` (or `package.json`).
|
||||
|
||||
```text
|
||||
checkpoint
|
||||
├── checkpoint-postgres
|
||||
├── checkpoint-sqlite
|
||||
├── prebuilt
|
||||
└── langgraph
|
||||
|
||||
prebuilt
|
||||
└── langgraph
|
||||
|
||||
sdk-py
|
||||
├── langgraph
|
||||
└── cli
|
||||
|
||||
sdk-js (standalone)
|
||||
```
|
||||
|
||||
Changes to a library may impact all of its dependents shown above.
|
||||
@@ -12,7 +12,6 @@
|
||||
[](https://pepy.tech/project/langgraph)
|
||||
[](https://github.com/langchain-ai/langgraph/issues)
|
||||
[](https://langchain-ai.github.io/langgraph/)
|
||||
[](https://gitmcp.io/langchain-ai/langgraph)
|
||||
|
||||
Trusted by companies shaping the future of agents – including Klarna, Replit, Elastic, and more – LangGraph is a low-level orchestration framework for building, managing, and deploying long-running, stateful agents.
|
||||
|
||||
|
||||
@@ -62,6 +62,15 @@ Starting from the `LangGraph Platform` view...
|
||||
1. In the panel, select the `Server` tab to view server logs for the revision. Server logs are only available after a revision has been deployed.
|
||||
1. Within the `Server` tab, adjust the date/time range picker as needed. By default, the date/time range picker is set to the `Last 7 days`.
|
||||
|
||||
## View Deployment Metrics
|
||||
|
||||
Starting from the <a href="https://smith.langchain.com/" target="_blank">LangSmith UI</a>...
|
||||
|
||||
1. In the left-hand navigation panel, select `LangGraph Platform`. The `LangGraph Platform` view contains a list of existing LangGraph Platform deployments.
|
||||
1. Select an existing deployment to monitor.
|
||||
1. Select the `Monitoring` tab to view the deployment metrics. See a list of [all available metrics](../../concepts/langgraph_control_plane.md#monitoring).
|
||||
1. Within the `Monitoring` tab, use the date/time range picker as needed. By default, the date/time range picker is set to the `Last 15 minutes`.
|
||||
|
||||
## Interrupt Revision
|
||||
|
||||
Interrupting a revision will stop deployment of the revision.
|
||||
|
||||
@@ -20,7 +20,7 @@ my-app/
|
||||
|-- openai_agent.py # code for your graph
|
||||
```
|
||||
|
||||
where the graph is defined in `openai_agent.py`.
|
||||
where the graph is defined in `openai_agent.py`.
|
||||
|
||||
### No rebuild
|
||||
|
||||
@@ -28,11 +28,11 @@ In the standard LangGraph API configuration, the server uses the compiled graph
|
||||
|
||||
```python
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langgraph.graph import END, START, StateGraph, MessagesState
|
||||
from langgraph.graph import END, START, MessageGraph
|
||||
|
||||
model = ChatOpenAI(temperature=0)
|
||||
|
||||
graph_workflow = StateGraph(MessagesState)
|
||||
graph_workflow = MessageGraph()
|
||||
|
||||
graph_workflow.add_node("agent", model)
|
||||
graph_workflow.add_edge("agent", END)
|
||||
@@ -61,7 +61,7 @@ To make your graph rebuild on each new run with custom configuration, you need t
|
||||
from typing import Annotated
|
||||
from typing_extensions import TypedDict
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langgraph.graph import END, START
|
||||
from langgraph.graph import END, START, MessageGraph
|
||||
from langgraph.graph.state import StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.prebuilt import ToolNode
|
||||
@@ -144,4 +144,4 @@ Finally, you need to specify the path to your graph-making function (`make_graph
|
||||
}
|
||||
```
|
||||
|
||||
See more info on LangGraph API configuration file [here](../reference/cli.md#configuration-file)
|
||||
See more info on LangGraph API configuration file [here](../reference/cli.md#configuration-file)
|
||||
@@ -30,18 +30,16 @@ Before deploying, review the [conceptual guide for the Self-Hosted Control Plane
|
||||
1. `LangGraphPlatform CRD`: A CRD for LangGraph Platform deployments. This contains the spec for managing an instance of a LangGraph platform deployment.
|
||||
1. `operator`: This operator handles changes to your LangGraph Platform CRDs.
|
||||
1. `host-backend`: This is the [control plane](../../concepts/langgraph_control_plane.md).
|
||||
1. Two additional images will be used by the chart.
|
||||
1. Two additional images will be used by the chart. Use the images that are specified in the latest release.
|
||||
|
||||
hostBackendImage:
|
||||
repository: "docker.io/langchain/hosted-langserve-backend"
|
||||
pullPolicy: IfNotPresent
|
||||
tag: "0.9.80"
|
||||
operatorImage:
|
||||
repository: "docker.io/langchain/langgraph-operator"
|
||||
pullPolicy: IfNotPresent
|
||||
tag: "aa9dff4"
|
||||
|
||||
1. In your `langsmith_config.yaml` file, enable the `langgraphPlatform` option. Note that you must also have a valid ingress setup:
|
||||
1. In your config file for langsmith (usually `langsmith_config.yaml`, enable the `langgraphPlatform` option. Note that you must also have a valid ingress setup:
|
||||
|
||||
config:
|
||||
langgraphPlatform:
|
||||
|
||||
@@ -3818,6 +3818,14 @@
|
||||
"title": "Filter",
|
||||
"description": "Optional dictionary of key-value pairs to filter results."
|
||||
},
|
||||
"query": {
|
||||
"type": [
|
||||
"string",
|
||||
"null"
|
||||
],
|
||||
"title": "Query",
|
||||
"description": "Query string for semantic/vector search."
|
||||
},
|
||||
"limit": {
|
||||
"type": "integer",
|
||||
"default": 10,
|
||||
|
||||
@@ -123,3 +123,12 @@ Defaults to `''`.
|
||||
Set `REDIS_CLUSTER` to `True` to enable Redis Cluster mode. When enabled, the system will connect to Redis using cluster mode. This is useful when connecting to a Redis Cluster deployment.
|
||||
|
||||
Defaults to `False`.
|
||||
|
||||
## `MOUNT_PREFIX`
|
||||
|
||||
!!! info "Only Allowed in Self-Hosted Deployments"
|
||||
The `MOUNT_PREFIX` environment variable is only allowed in Self-Hosted Deployment models, LangGraph Platform SaaS will not allow this environment variable.
|
||||
|
||||
Set `MOUNT_PREFIX` to serve the LangGraph Server under a specific path prefix. This is useful for deployments where the server is behind a reverse proxy or load balancer that requires a specific path prefix.
|
||||
|
||||
For example, if the server is to be served under `https://example.com/langgraph`, set `MOUNT_PREFIX` to `/langgraph`.
|
||||
|
||||
@@ -19,6 +19,7 @@ From the control plane UI, you can:
|
||||
- Update a deployment.
|
||||
- Update environment variables for a deployment.
|
||||
- View build and server logs of a deployment.
|
||||
- View deployment metrics like CPU and memory usage.
|
||||
- Delete a deployment.
|
||||
|
||||
The Control Plane UI is embedded in [LangSmith](https://docs.smith.langchain.com/langgraph_cloud).
|
||||
@@ -88,6 +89,15 @@ Infrastructure for deployments and revisions are provisioned and deployed asynch
|
||||
|
||||
The control plane and [LangGraph Data Plane](./langgraph_data_plane.md) "listener" application coordinate to achieve asynchronous deployments.
|
||||
|
||||
### Monitoring
|
||||
|
||||
After a deployment is ready, the control plane monitors the deployment and records various metrics, such as:
|
||||
|
||||
- CPU and memory usage of the deployment.
|
||||
- Number of container restarts.
|
||||
|
||||
These metrics are displayed as charts in the Control Plane UI.
|
||||
|
||||
### LangSmith Integration
|
||||
|
||||
A [LangSmith](https://docs.smith.langchain.com/) tracing project is automatically created for each deployment. The tracing project has the same name as the deployment. When creating a deployment, the `LANGCHAIN_TRACING` and `LANGSMITH_API_KEY`/`LANGCHAIN_API_KEY` environment variables do not need to be specified; they are set automatically by the control plane.
|
||||
|
||||
@@ -470,9 +470,51 @@ If the checkpointer is used with asynchronous graph execution (i.e. executing th
|
||||
|
||||
### Serializer
|
||||
|
||||
When checkpointers save the graph state, they need to serialize the channel values in the state. This is done using serializer objects.
|
||||
When checkpointers save the graph state, they need to serialize the channel values in the state. This is done using serializer objects.
|
||||
`langgraph_checkpoint` defines [protocol][langgraph.checkpoint.serde.base.SerializerProtocol] for implementing serializers provides a default implementation ([JsonPlusSerializer][langgraph.checkpoint.serde.jsonplus.JsonPlusSerializer]) that handles a wide variety of types, including LangChain and LangGraph primitives, datetimes, enums and more.
|
||||
|
||||
#### Serialization with `pickle`
|
||||
|
||||
The default serializer, [`JsonPlusSerializer`][langgraph.checkpoint.serde.jsonplus.JsonPlusSerializer], uses ormsgpack and JSON under the hood, which is not suitable for all types of objects.
|
||||
|
||||
If you want to fallback to pickle for objects not currently supported by our msgpack encoder (such as Pandas dataframes),
|
||||
you can use the `pickle_fallback` argument of the `JsonPlusSerializer`:
|
||||
|
||||
```python
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
# ... Define the graph ...
|
||||
graph.compile(
|
||||
checkpointer=MemorySaver(serde=JsonPlusSerializer(pickle_fallback=True))
|
||||
)
|
||||
```
|
||||
|
||||
#### Encryption
|
||||
|
||||
Checkpointers can optionally encrypt all persisted state. To enable this, pass an instance of [`EncryptedSerializer`][langgraph.checkpoint.serde.encrypted.EncryptedSerializer] to the `serde` argument of any `BaseCheckpointSaver` implementation. The easiest way to create an encrypted serializer is via [`from_pycryptodome_aes`][langgraph.checkpoint.serde.encrypted.EncryptedSerializer.from_pycryptodome_aes], which reads the AES key from the `LANGGRAPH_AES_KEY` environment variable (or accepts a `key` argument):
|
||||
|
||||
```python
|
||||
import sqlite3
|
||||
|
||||
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
|
||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
|
||||
serde = EncryptedSerializer.from_pycryptodome_aes() # reads LANGGRAPH_AES_KEY
|
||||
checkpointer = SqliteSaver(sqlite3.connect("checkpoint.db"), serde=serde)
|
||||
```
|
||||
|
||||
```python
|
||||
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
|
||||
serde = EncryptedSerializer.from_pycryptodome_aes()
|
||||
checkpointer = PostgresSaver.from_conn_string("postgresql://...", serde=serde)
|
||||
checkpointer.setup()
|
||||
```
|
||||
|
||||
When running on LangGraph Platform, encryption is automatically enabled whenever `LANGGRAPH_AES_KEY` is present, so you only need to provide the environment variable. Other encryption schemes can be used by implementing [`CipherProtocol`][langgraph.checkpoint.serde.base.CipherProtocol] and supplying it to `EncryptedSerializer`.
|
||||
|
||||
## Capabilities
|
||||
|
||||
### Human-in-the-loop
|
||||
|
||||
@@ -59,8 +59,9 @@ The main question when adding subgraphs is how the parent graph and subgraph com
|
||||
response = model.invoke(state["subgraph_messages"])
|
||||
return {"subgraph_messages": response}
|
||||
|
||||
subgraph_builder = StateGraph(State)
|
||||
subgraph_builder.add_node(call_model)
|
||||
subgraph_builder = StateGraph(SubgraphMessagesState)
|
||||
subgraph_builder.add_node("call_model_from_subgraph", call_model)
|
||||
subgraph_builder.add_edge(START, "call_model_from_subgraph")
|
||||
...
|
||||
# highlight-next-line
|
||||
subgraph = subgraph_builder.compile()
|
||||
|
||||
@@ -1107,10 +1107,10 @@
|
||||
"source": [
|
||||
"### Use in production\n",
|
||||
"\n",
|
||||
"In production, you would want to use a checkpointer backed by a database:\n",
|
||||
"In production, you would want to use a store backed by a database:\n",
|
||||
"\n",
|
||||
"```python\n",
|
||||
"from langgraph.checkpoint.postgres import PostgresSaver\n",
|
||||
"from langgraph.store.postgres import PostgresStore\n",
|
||||
"\n",
|
||||
"DB_URI = \"postgresql://postgres:postgres@localhost:5442/postgres?sslmode=disable\"\n",
|
||||
"# highlight-next-line\n",
|
||||
|
||||
@@ -12,12 +12,18 @@
|
||||
options:
|
||||
members:
|
||||
- SerializerProtocol
|
||||
- CipherProtocol
|
||||
|
||||
::: langgraph.checkpoint.serde.jsonplus
|
||||
options:
|
||||
members:
|
||||
- JsonPlusSerializer
|
||||
|
||||
::: langgraph.checkpoint.serde.encrypted
|
||||
options:
|
||||
members:
|
||||
- EncryptedSerializer
|
||||
|
||||
::: langgraph.checkpoint.memory
|
||||
|
||||
::: langgraph.checkpoint.sqlite
|
||||
@@ -32,4 +38,4 @@
|
||||
::: langgraph.checkpoint.postgres.aio
|
||||
options:
|
||||
members:
|
||||
- AsyncPostgresSaver
|
||||
- AsyncPostgresSaver
|
||||
|
||||
+1
-3
@@ -580,9 +580,7 @@
|
||||
" ]\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"evaluator = prompt | ChatOpenAI(model=\"gpt-4-turbo-preview\").with_structured_output(\n",
|
||||
" RedTeamingResult, method=\"function_calling\"\n",
|
||||
")\n",
|
||||
"evaluator = prompt | ChatOpenAI(model=\"gpt-4o\").with_structured_output(RedTeamingResult)\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def did_resist(run, example):\n",
|
||||
|
||||
@@ -89,7 +89,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"execution_count": 2,
|
||||
"id": "baf669a0-04ee-492d-80d8-8fcb658ed128",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
@@ -313,8 +313,8 @@
|
||||
"\n",
|
||||
" builder.add_edge(\"finalizer\", END)\n",
|
||||
"\n",
|
||||
" # These functions let the step be used in a\n",
|
||||
" # StateGraph with 'messages' as the key.\n",
|
||||
" # These functions let the step be used in a MessageGraph\n",
|
||||
" # or a StateGraph with 'messages' as the key.\n",
|
||||
" def encode(x: Union[Sequence[AnyMessage], PromptValue]) -> dict:\n",
|
||||
" \"\"\"Ensure the input is the correct format.\"\"\"\n",
|
||||
" if isinstance(x, PromptValue):\n",
|
||||
|
||||
@@ -164,7 +164,7 @@ llm = init_chat_model("anthropic:claude-3-5-sonnet-latest")
|
||||
```
|
||||
-->
|
||||
|
||||
```python
|
||||
```python hl_lines="36 37"
|
||||
from typing import Annotated
|
||||
|
||||
from langchain.chat_models import init_chat_model
|
||||
@@ -206,4 +206,4 @@ graph = graph_builder.compile(checkpointer=memory)
|
||||
|
||||
## Next steps
|
||||
|
||||
In the next tutorial, you will [add human-in-the-loop to the chatbot](./4-human-in-the-loop.md) to handle situations where it may need guidance or verification before proceeding.
|
||||
In the next tutorial, you will [add human-in-the-loop to the chatbot](./4-human-in-the-loop.md) to handle situations where it may need guidance or verification before proceeding.
|
||||
|
||||
@@ -471,7 +471,7 @@
|
||||
"\n",
|
||||
"_get_pass(\"TAVILY_API_KEY\")\n",
|
||||
"\n",
|
||||
"calculate = get_math_tool(ChatOpenAI(model=\"gpt-4-turbo-preview\"))\n",
|
||||
"calculate = get_math_tool(ChatOpenAI(model=\"gpt-4o\"))\n",
|
||||
"search = TavilySearchResults(\n",
|
||||
" max_results=1,\n",
|
||||
" description='tavily_search_results_json(query=\"the search query\") - a search engine.',\n",
|
||||
@@ -540,11 +540,11 @@
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"================================\u001b[1m System Message \u001b[0m================================\n",
|
||||
"================================\u001B[1m System Message \u001B[0m================================\n",
|
||||
"\n",
|
||||
"Given a user query, create a plan to solve it with the utmost parallelizability. Each plan should comprise an action from the following \u001b[33;1m\u001b[1;3m{num_tools}\u001b[0m types:\n",
|
||||
"\u001b[33;1m\u001b[1;3m{tool_descriptions}\u001b[0m\n",
|
||||
"\u001b[33;1m\u001b[1;3m{num_tools}\u001b[0m. join(): Collects and combines results from prior actions.\n",
|
||||
"Given a user query, create a plan to solve it with the utmost parallelizability. Each plan should comprise an action from the following \u001B[33;1m\u001B[1;3m{num_tools}\u001B[0m types:\n",
|
||||
"\u001B[33;1m\u001B[1;3m{tool_descriptions}\u001B[0m\n",
|
||||
"\u001B[33;1m\u001B[1;3m{num_tools}\u001B[0m. join(): Collects and combines results from prior actions.\n",
|
||||
"\n",
|
||||
" - An LLM agent is called upon invoking join() to either finalize the user query or wait until the plans are executed.\n",
|
||||
" - join should always be the last action in the plan, and will be called in two scenarios:\n",
|
||||
@@ -561,11 +561,11 @@
|
||||
" - Only use the provided action types. If a query cannot be addressed using these, invoke the join action for the next steps.\n",
|
||||
" - Never introduce new actions other than the ones provided.\n",
|
||||
"\n",
|
||||
"=============================\u001b[1m Messages Placeholder \u001b[0m=============================\n",
|
||||
"=============================\u001B[1m Messages Placeholder \u001B[0m=============================\n",
|
||||
"\n",
|
||||
"\u001b[33;1m\u001b[1;3m{messages}\u001b[0m\n",
|
||||
"\u001B[33;1m\u001B[1;3m{messages}\u001B[0m\n",
|
||||
"\n",
|
||||
"================================\u001b[1m System Message \u001b[0m================================\n",
|
||||
"================================\u001B[1m System Message \u001B[0m================================\n",
|
||||
"\n",
|
||||
"Remember, ONLY respond with the task list in the correct format! E.g.:\n",
|
||||
"idx. tool(arg_name=args)\n",
|
||||
@@ -1030,7 +1030,7 @@
|
||||
"joiner_prompt = hub.pull(\"wfh/llm-compiler-joiner\").partial(\n",
|
||||
" examples=\"\"\n",
|
||||
") # You can optionally add examples\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-4-turbo-preview\")\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-4o\")\n",
|
||||
"\n",
|
||||
"runnable = joiner_prompt | llm.with_structured_output(\n",
|
||||
" JoinOutputs, method=\"function_calling\"\n",
|
||||
|
||||
@@ -135,7 +135,6 @@
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from langchain import hub\n",
|
||||
"from langchain_openai import ChatOpenAI\n",
|
||||
"\n",
|
||||
"from langgraph.prebuilt import create_react_agent\n",
|
||||
|
||||
@@ -90,7 +90,11 @@
|
||||
"id": "9ac1c2cd-81fb-40eb-8ba1-e9197800cba6",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Create Index"
|
||||
"## Create Index\n",
|
||||
"\n",
|
||||
"Set up a vector database using OpenAI Embeddings and the Chroma vector database. \n",
|
||||
"Input URLs of blog posts related to agents, prompt engineering, and large language models (LLMs). \n",
|
||||
"Generate vector indices for use in Retrieval-Augmented Generation (RAG)."
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -159,6 +163,21 @@
|
||||
"</div>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "6cdd5ac0-fa18-4ee9-8051-062a0c56268f",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Router for Query Analysis\n",
|
||||
"\n",
|
||||
"Let’s start with Routing. First, assign the query analysis to the LLM.\n",
|
||||
"\n",
|
||||
"Create a RouteQuery data model and specify it in a structured format for the LLM. The decision for routing should be embedded in the prompt. You need to clearly define which parts of the document should be directed to RAG based on the topic.\n",
|
||||
"\n",
|
||||
"While you could automate this process by having the LLM summarize the RAG documents again, it’s more cost-effective to manually manage this when dealing with large documents, as automation could become expensive.\n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
@@ -219,6 +238,18 @@
|
||||
"print(question_router.invoke({\"question\": \"What are the types of agent memory?\"}))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "cb248c94-0b0c-4d86-8565-32aa8d7424e4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Retrieval Grader\n",
|
||||
"\n",
|
||||
"After performing retrieval, evaluate the results. Although you initially decided to use RAG based on the query, the retrieved documents might not be satisfactory. Assess whether the retrieved documents are sufficiently relevant to the query.\n",
|
||||
"\n",
|
||||
"For this, rely on the LLM to evaluate the relevance, providing a binary ‘yes’ or ‘no’ decision."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
@@ -309,6 +340,17 @@
|
||||
"print(generation)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "cb0ab54a-4a4f-45fa-b1c5-cea1bf4c59d5",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Hallucination Grader\n",
|
||||
"\n",
|
||||
"Verify if the LLM produced any hallucinations by comparing its output to the retrieved facts. \n",
|
||||
"Provide the LLM’s evaluation in a binary ‘yes’ or ‘no’ format.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
@@ -357,6 +399,16 @@
|
||||
"hallucination_grader.invoke({\"documents\": docs, \"generation\": generation})"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "4f58502a-c25f-4d80-a402-5583b0cd3e41",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Answer Grader\n",
|
||||
"\n",
|
||||
"Evaluate the answer finally."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
@@ -405,6 +457,18 @@
|
||||
"answer_grader.invoke({\"question\": question, \"generation\": generation})"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "af77946c-2646-4039-86b0-e2fde1ab7459",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Question Rewriting\n",
|
||||
"\n",
|
||||
"The original question from user was directly used in RAG. \n",
|
||||
"However, the user’s question might not be in a form suitable for RAG. \n",
|
||||
"To improve retrieval, rephrase the question to ensure it aligns better with vector similarity search."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
@@ -450,7 +514,9 @@
|
||||
"id": "d07c0b31-b919-4498-869f-9673125c2473",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Web Search Tool"
|
||||
"## Web Search Tool\n",
|
||||
"\n",
|
||||
"Use Tavily Search tool to get information from the web."
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
@@ -185,7 +185,7 @@
|
||||
"\n",
|
||||
"\n",
|
||||
"# LLM with function call\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
|
||||
"structured_llm_grader = llm.with_structured_output(GradeDocuments)\n",
|
||||
"\n",
|
||||
"# Prompt\n",
|
||||
|
||||
@@ -364,16 +364,6 @@ markdown_extensions:
|
||||
hooks:
|
||||
- _scripts/notebook_hooks.py
|
||||
extra:
|
||||
consent:
|
||||
title: Cookie consent
|
||||
actions:
|
||||
- accept
|
||||
- reject
|
||||
description: >-
|
||||
We use cookies to recognize your repeated visits and preferences, as well
|
||||
as to measure the effectiveness of our documentation and whether users
|
||||
find what they're searching for. <strong>Clicking "Accept" makes our
|
||||
documentation better. Thank you!</strong> ❤️
|
||||
social:
|
||||
- icon: fontawesome/brands/js
|
||||
link: https://langchain-ai.github.io/langgraphjs/
|
||||
@@ -381,25 +371,6 @@ extra:
|
||||
link: https://github.com/langchain-ai/langgraph
|
||||
- icon: fontawesome/brands/twitter
|
||||
link: https://twitter.com/LangChainAI
|
||||
analytics:
|
||||
provider: google
|
||||
property: G-G8X6ELZYE0
|
||||
feedback:
|
||||
title: Was this page helpful?
|
||||
ratings:
|
||||
- icon: material/emoticon-happy-outline
|
||||
name: This page was helpful
|
||||
data: 1
|
||||
note: >-
|
||||
Thanks for your feedback!
|
||||
- icon: material/emoticon-sad-outline
|
||||
name: This page could be improved
|
||||
data: 0
|
||||
note: >-
|
||||
Thanks for your feedback! Please help us improve this page by adding to the discussion below.
|
||||
shared_analytics:
|
||||
provider: google
|
||||
property: G-47WX3HKKY2
|
||||
validation:
|
||||
# https://www.mkdocs.org/user-guide/configuration/
|
||||
# We are still raising for omitted files because they determine the breadcrumbs for pages.
|
||||
|
||||
Generated
+3060
-3059
File diff suppressed because it is too large
Load Diff
@@ -184,7 +184,7 @@
|
||||
"\n",
|
||||
"\n",
|
||||
"# LLM with function call\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
|
||||
"structured_llm_router = llm.with_structured_output(RouteQuery)\n",
|
||||
"\n",
|
||||
"# Prompt\n",
|
||||
@@ -235,7 +235,7 @@
|
||||
"\n",
|
||||
"\n",
|
||||
"# LLM with function call\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
|
||||
"structured_llm_grader = llm.with_structured_output(GradeDocuments)\n",
|
||||
"\n",
|
||||
"# Prompt\n",
|
||||
@@ -328,7 +328,7 @@
|
||||
"\n",
|
||||
"\n",
|
||||
"# LLM with function call\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
|
||||
"structured_llm_grader = llm.with_structured_output(GradeHallucinations)\n",
|
||||
"\n",
|
||||
"# Prompt\n",
|
||||
@@ -376,7 +376,7 @@
|
||||
"\n",
|
||||
"\n",
|
||||
"# LLM with function call\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
|
||||
"structured_llm_grader = llm.with_structured_output(GradeAnswer)\n",
|
||||
"\n",
|
||||
"# Prompt\n",
|
||||
|
||||
@@ -200,11 +200,11 @@
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"********************Prompt[rlm/rag-prompt]********************\n",
|
||||
"================================\u001b[1m Human Message \u001b[0m=================================\n",
|
||||
"================================\u001B[1m Human Message \u001B[0m=================================\n",
|
||||
"\n",
|
||||
"You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.\n",
|
||||
"Question: \u001b[33;1m\u001b[1;3m{question}\u001b[0m \n",
|
||||
"Context: \u001b[33;1m\u001b[1;3m{context}\u001b[0m \n",
|
||||
"Question: \u001B[33;1m\u001B[1;3m{question}\u001B[0m \n",
|
||||
"Context: \u001B[33;1m\u001B[1;3m{context}\u001B[0m \n",
|
||||
"Answer:\n"
|
||||
]
|
||||
}
|
||||
@@ -244,7 +244,7 @@
|
||||
" binary_score: str = Field(description=\"Relevance score 'yes' or 'no'\")\n",
|
||||
"\n",
|
||||
" # LLM\n",
|
||||
" model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n",
|
||||
" model = ChatOpenAI(temperature=0, model=\"gpt-4o\", streaming=True)\n",
|
||||
"\n",
|
||||
" # LLM with tool and validation\n",
|
||||
" llm_with_tool = model.with_structured_output(grade)\n",
|
||||
|
||||
@@ -171,7 +171,7 @@
|
||||
"\n",
|
||||
"\n",
|
||||
"# LLM with function call\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
|
||||
"structured_llm_grader = llm.with_structured_output(GradeDocuments)\n",
|
||||
"\n",
|
||||
"# Prompt\n",
|
||||
|
||||
@@ -191,7 +191,7 @@
|
||||
"\n",
|
||||
"\n",
|
||||
"# LLM with function call\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
|
||||
"structured_llm_grader = llm.with_structured_output(GradeDocuments)\n",
|
||||
"\n",
|
||||
"# Prompt\n",
|
||||
@@ -284,7 +284,7 @@
|
||||
"\n",
|
||||
"\n",
|
||||
"# LLM with function call\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
|
||||
"structured_llm_grader = llm.with_structured_output(GradeHallucinations)\n",
|
||||
"\n",
|
||||
"# Prompt\n",
|
||||
@@ -332,7 +332,7 @@
|
||||
"\n",
|
||||
"\n",
|
||||
"# LLM with function call\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
|
||||
"structured_llm_grader = llm.with_structured_output(GradeAnswer)\n",
|
||||
"\n",
|
||||
"# Prompt\n",
|
||||
|
||||
@@ -33,7 +33,9 @@
|
||||
"id": "a384cc48-0425-4e8f-aafc-cfb8e56025c9",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": ["%pip install -qU langchain-pinecone langchain-openai langchainhub langgraph"]
|
||||
"source": [
|
||||
"%pip install -qU langchain-pinecone langchain-openai langchainhub langgraph"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
@@ -51,7 +53,9 @@
|
||||
"id": "ccc3dae5-1df6-48ca-af8a-50f0e6128876",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": ["import os\n\nos.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\nos.environ[\"LANGCHAIN_ENDPOINT\"] = \"https://api.smith.langchain.com\"\nos.environ[\"LANGCHAIN_API_KEY\"] = \"<your-api-key>\""]
|
||||
"source": [
|
||||
"import os\n\nos.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\nos.environ[\"LANGCHAIN_ENDPOINT\"] = \"https://api.smith.langchain.com\"\nos.environ[\"LANGCHAIN_API_KEY\"] = \"<your-api-key>\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
@@ -59,7 +63,9 @@
|
||||
"id": "88637820",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": ["import os\n\nos.environ[\"LANGCHAIN_PROJECT\"] = \"pinecone-devconnect\""]
|
||||
"source": [
|
||||
"import os\n\nos.environ[\"LANGCHAIN_PROJECT\"] = \"pinecone-devconnect\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
@@ -77,7 +83,9 @@
|
||||
"id": "565a6d44-2c9f-4fff-b1ec-eea05df9350d",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": ["from langchain_openai import OpenAIEmbeddings\nfrom langchain_pinecone import PineconeVectorStore\n\n# use pinecone movies database\n\n# Add to vectorDB\nvectorstore = PineconeVectorStore(\n embedding=OpenAIEmbeddings(),\n index_name=\"sample-movies\",\n text_key=\"summary\",\n)\nretriever = vectorstore.as_retriever()"]
|
||||
"source": [
|
||||
"from langchain_openai import OpenAIEmbeddings\nfrom langchain_pinecone import PineconeVectorStore\n\n# use pinecone movies database\n\n# Add to vectorDB\nvectorstore = PineconeVectorStore(\n embedding=OpenAIEmbeddings(),\n index_name=\"sample-movies\",\n text_key=\"summary\",\n)\nretriever = vectorstore.as_retriever()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
@@ -104,7 +112,9 @@
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": ["docs = retriever.invoke(\"James Cameron\")\nfor doc in docs:\n print(\"# \" + doc.metadata[\"title\"])\n print(doc.page_content)\n print()"]
|
||||
"source": [
|
||||
"docs = retriever.invoke(\"James Cameron\")\nfor doc in docs:\n print(\"# \" + doc.metadata[\"title\"])\n print(doc.page_content)\n print()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
@@ -120,7 +130,32 @@
|
||||
"id": "1fafad21-60cc-483e-92a3-6a7edb1838e3",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": ["### Retrieval Grader\n\nfrom langchain import hub\nfrom langchain_core.pydantic_v1 import BaseModel, Field\nfrom langchain_openai import ChatOpenAI\n\n\n# Data model\nclass GradeDocuments(BaseModel):\n \"\"\"Binary score for relevance check on retrieved documents.\"\"\"\n\n binary_score: str = Field(\n description=\"Documents are relevant to the question, 'yes' or 'no'\"\n )\n\n\n# https://smith.langchain.com/hub/efriis/self-rag-retrieval-grader\ngrade_prompt = hub.pull(\"efriis/self-rag-retrieval-grader\")\n\n# LLM with function call\nllm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\nstructured_llm_grader = llm.with_structured_output(GradeDocuments)\n\nretrieval_grader = grade_prompt | structured_llm_grader"]
|
||||
"source": [
|
||||
"### Retrieval Grader\n",
|
||||
"\n",
|
||||
"from langchain import hub\n",
|
||||
"from langchain_core.pydantic_v1 import BaseModel, Field\n",
|
||||
"from langchain_openai import ChatOpenAI\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# Data model\n",
|
||||
"class GradeDocuments(BaseModel):\n",
|
||||
" \"\"\"Binary score for relevance check on retrieved documents.\"\"\"\n",
|
||||
"\n",
|
||||
" binary_score: str = Field(\n",
|
||||
" description=\"Documents are relevant to the question, 'yes' or 'no'\"\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# https://smith.langchain.com/hub/efriis/self-rag-retrieval-grader\n",
|
||||
"grade_prompt = hub.pull(\"efriis/self-rag-retrieval-grader\")\n",
|
||||
"\n",
|
||||
"# LLM with function call\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
|
||||
"structured_llm_grader = llm.with_structured_output(GradeDocuments)\n",
|
||||
"\n",
|
||||
"retrieval_grader = grade_prompt | structured_llm_grader"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
@@ -137,7 +172,9 @@
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": ["# Test the retrieval grader\nquestion = \"movies starring jason momoa\"\ndocs = retriever.invoke(question)\ndoc_txt = docs[0].page_content\nprint(doc_txt)\nprint(retrieval_grader.invoke({\"question\": question, \"document\": doc_txt}))"]
|
||||
"source": [
|
||||
"# Test the retrieval grader\nquestion = \"movies starring jason momoa\"\ndocs = retriever.invoke(question)\ndoc_txt = docs[0].page_content\nprint(doc_txt)\nprint(retrieval_grader.invoke({\"question\": question, \"document\": doc_txt}))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
@@ -163,7 +200,9 @@
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": ["### Generate\n\nfrom langchain import hub\nfrom langchain_core.output_parsers import StrOutputParser\n\n# Prompt\nprompt = hub.pull(\"rlm/rag-prompt\")\n\n# LLM\nllm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0)\n\n# Chain\nrag_chain = prompt | llm | StrOutputParser()\n\n# Run\ngeneration = rag_chain.invoke({\"context\": docs, \"question\": question})\nprint(generation)"]
|
||||
"source": [
|
||||
"### Generate\n\nfrom langchain import hub\nfrom langchain_core.output_parsers import StrOutputParser\n\n# Prompt\nprompt = hub.pull(\"rlm/rag-prompt\")\n\n# LLM\nllm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0)\n\n# Chain\nrag_chain = prompt | llm | StrOutputParser()\n\n# Run\ngeneration = rag_chain.invoke({\"context\": docs, \"question\": question})\nprint(generation)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
@@ -189,7 +228,30 @@
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": ["### Hallucination Grader\n\n\n# Data model\nclass GradeHallucinations(BaseModel):\n \"\"\"Binary score for hallucination present in generation answer.\"\"\"\n\n binary_score: str = Field(\n description=\"Answer is grounded in the facts, 'yes' or 'no'\"\n )\n\n\n# LLM with function call\nllm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\nstructured_llm_grader = llm.with_structured_output(GradeHallucinations)\n\n# https://smith.langchain.com/hub/efriis/self-rag-hallucination-grader\nhallucination_prompt = hub.pull(\"efriis/self-rag-hallucination-grader\")\n\nhallucination_grader = hallucination_prompt | structured_llm_grader\nprint(generation)\nhallucination_grader.invoke({\"documents\": docs, \"generation\": generation})"]
|
||||
"source": [
|
||||
"### Hallucination Grader\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# Data model\n",
|
||||
"class GradeHallucinations(BaseModel):\n",
|
||||
" \"\"\"Binary score for hallucination present in generation answer.\"\"\"\n",
|
||||
"\n",
|
||||
" binary_score: str = Field(\n",
|
||||
" description=\"Answer is grounded in the facts, 'yes' or 'no'\"\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# LLM with function call\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
|
||||
"structured_llm_grader = llm.with_structured_output(GradeHallucinations)\n",
|
||||
"\n",
|
||||
"# https://smith.langchain.com/hub/efriis/self-rag-hallucination-grader\n",
|
||||
"hallucination_prompt = hub.pull(\"efriis/self-rag-hallucination-grader\")\n",
|
||||
"\n",
|
||||
"hallucination_grader = hallucination_prompt | structured_llm_grader\n",
|
||||
"print(generation)\n",
|
||||
"hallucination_grader.invoke({\"documents\": docs, \"generation\": generation})"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
@@ -216,7 +278,31 @@
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": ["### Answer Grader\n\n\n# Data model\nclass GradeAnswer(BaseModel):\n \"\"\"Binary score to assess answer addresses question.\"\"\"\n\n binary_score: str = Field(\n description=\"Answer addresses the question, 'yes' or 'no'\"\n )\n\n\n# LLM with function call\nllm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\nstructured_llm_grader = llm.with_structured_output(GradeAnswer)\n\n# Prompt\nanswer_prompt = hub.pull(\"efriis/self-rag-answer-grader\")\n\nanswer_grader = answer_prompt | structured_llm_grader\nprint(question)\nprint(generation)\nanswer_grader.invoke({\"question\": question, \"generation\": generation})"]
|
||||
"source": [
|
||||
"### Answer Grader\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# Data model\n",
|
||||
"class GradeAnswer(BaseModel):\n",
|
||||
" \"\"\"Binary score to assess answer addresses question.\"\"\"\n",
|
||||
"\n",
|
||||
" binary_score: str = Field(\n",
|
||||
" description=\"Answer addresses the question, 'yes' or 'no'\"\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# LLM with function call\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
|
||||
"structured_llm_grader = llm.with_structured_output(GradeAnswer)\n",
|
||||
"\n",
|
||||
"# Prompt\n",
|
||||
"answer_prompt = hub.pull(\"efriis/self-rag-answer-grader\")\n",
|
||||
"\n",
|
||||
"answer_grader = answer_prompt | structured_llm_grader\n",
|
||||
"print(question)\n",
|
||||
"print(generation)\n",
|
||||
"answer_grader.invoke({\"question\": question, \"generation\": generation})"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
@@ -242,7 +328,9 @@
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": ["### Question Re-writer\n\n# LLM\nllm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n\n# Prompt\nre_write_prompt = hub.pull(\"efriis/self-rag-question-rewriter\")\n\nquestion_rewriter = re_write_prompt | llm | StrOutputParser()\nprint(question)\nquestion_rewriter.invoke({\"question\": question})"]
|
||||
"source": [
|
||||
"### Question Re-writer\n\n# LLM\nllm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n\n# Prompt\nre_write_prompt = hub.pull(\"efriis/self-rag-question-rewriter\")\n\nquestion_rewriter = re_write_prompt | llm | StrOutputParser()\nprint(question)\nquestion_rewriter.invoke({\"question\": question})"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
@@ -262,7 +350,9 @@
|
||||
"id": "f1617e9e-66a8-4c1a-a1fe-cc936284c085",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": ["from typing import List\n\nfrom typing_extensions import TypedDict\n\n\nclass GraphState(TypedDict):\n \"\"\"\n Represents the state of our graph.\n\n Attributes:\n question: question\n generation: LLM generation\n documents: list of documents\n \"\"\"\n\n question: str\n generation: str\n documents: List[str]"]
|
||||
"source": [
|
||||
"from typing import List\n\nfrom typing_extensions import TypedDict\n\n\nclass GraphState(TypedDict):\n \"\"\"\n Represents the state of our graph.\n\n Attributes:\n question: question\n generation: LLM generation\n documents: list of documents\n \"\"\"\n\n question: str\n generation: str\n documents: List[str]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
@@ -270,7 +360,9 @@
|
||||
"id": "add509d8-6682-4127-8d95-13dd37d79702",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": ["### Nodes\n\n\ndef retrieve(state):\n \"\"\"\n Retrieve documents\n\n Args:\n state (dict): The current graph state\n\n Returns:\n state (dict): New key added to state, documents, that contains retrieved documents\n \"\"\"\n print(\"---RETRIEVE---\")\n question = state[\"question\"]\n\n # Retrieval\n documents = retriever.invoke(question)\n return {\"documents\": documents, \"question\": question}\n\n\ndef generate(state):\n \"\"\"\n Generate answer\n\n Args:\n state (dict): The current graph state\n\n Returns:\n state (dict): New key added to state, generation, that contains LLM generation\n \"\"\"\n print(\"---GENERATE---\")\n question = state[\"question\"]\n documents = state[\"documents\"]\n\n # RAG generation\n generation = rag_chain.invoke({\"context\": documents, \"question\": question})\n return {\"documents\": documents, \"question\": question, \"generation\": generation}\n\n\ndef grade_documents(state):\n \"\"\"\n Determines whether the retrieved documents are relevant to the question.\n\n Args:\n state (dict): The current graph state\n\n Returns:\n state (dict): Updates documents key with only filtered relevant documents\n \"\"\"\n\n print(\"---CHECK DOCUMENT RELEVANCE TO QUESTION---\")\n question = state[\"question\"]\n documents = state[\"documents\"]\n\n # Score each doc\n filtered_docs = []\n for d in documents:\n score = retrieval_grader.invoke(\n {\"question\": question, \"document\": d.page_content}\n )\n grade = score.binary_score\n if grade == \"yes\":\n print(\"---GRADE: DOCUMENT RELEVANT---\")\n filtered_docs.append(d)\n else:\n print(\"---GRADE: DOCUMENT NOT RELEVANT---\")\n continue\n return {\"documents\": filtered_docs, \"question\": question}\n\n\ndef transform_query(state):\n \"\"\"\n Transform the query to produce a better question.\n\n Args:\n state (dict): The current graph state\n\n Returns:\n state (dict): Updates question key with a re-phrased question\n \"\"\"\n\n print(\"---TRANSFORM QUERY---\")\n question = state[\"question\"]\n documents = state[\"documents\"]\n\n # Re-write question\n better_question = question_rewriter.invoke({\"question\": question})\n return {\"documents\": documents, \"question\": better_question}"]
|
||||
"source": [
|
||||
"### Nodes\n\n\ndef retrieve(state):\n \"\"\"\n Retrieve documents\n\n Args:\n state (dict): The current graph state\n\n Returns:\n state (dict): New key added to state, documents, that contains retrieved documents\n \"\"\"\n print(\"---RETRIEVE---\")\n question = state[\"question\"]\n\n # Retrieval\n documents = retriever.invoke(question)\n return {\"documents\": documents, \"question\": question}\n\n\ndef generate(state):\n \"\"\"\n Generate answer\n\n Args:\n state (dict): The current graph state\n\n Returns:\n state (dict): New key added to state, generation, that contains LLM generation\n \"\"\"\n print(\"---GENERATE---\")\n question = state[\"question\"]\n documents = state[\"documents\"]\n\n # RAG generation\n generation = rag_chain.invoke({\"context\": documents, \"question\": question})\n return {\"documents\": documents, \"question\": question, \"generation\": generation}\n\n\ndef grade_documents(state):\n \"\"\"\n Determines whether the retrieved documents are relevant to the question.\n\n Args:\n state (dict): The current graph state\n\n Returns:\n state (dict): Updates documents key with only filtered relevant documents\n \"\"\"\n\n print(\"---CHECK DOCUMENT RELEVANCE TO QUESTION---\")\n question = state[\"question\"]\n documents = state[\"documents\"]\n\n # Score each doc\n filtered_docs = []\n for d in documents:\n score = retrieval_grader.invoke(\n {\"question\": question, \"document\": d.page_content}\n )\n grade = score.binary_score\n if grade == \"yes\":\n print(\"---GRADE: DOCUMENT RELEVANT---\")\n filtered_docs.append(d)\n else:\n print(\"---GRADE: DOCUMENT NOT RELEVANT---\")\n continue\n return {\"documents\": filtered_docs, \"question\": question}\n\n\ndef transform_query(state):\n \"\"\"\n Transform the query to produce a better question.\n\n Args:\n state (dict): The current graph state\n\n Returns:\n state (dict): Updates question key with a re-phrased question\n \"\"\"\n\n print(\"---TRANSFORM QUERY---\")\n question = state[\"question\"]\n documents = state[\"documents\"]\n\n # Re-write question\n better_question = question_rewriter.invoke({\"question\": question})\n return {\"documents\": documents, \"question\": better_question}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
@@ -278,7 +370,9 @@
|
||||
"id": "09fc91b4",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": ["### Edges\n\n\ndef decide_to_generate(state):\n \"\"\"\n Determines whether to generate an answer, or re-generate a question.\n\n Args:\n state (dict): The current graph state\n\n Returns:\n str: Binary decision for next node to call\n \"\"\"\n\n print(\"---ASSESS GRADED DOCUMENTS---\")\n state[\"question\"]\n filtered_documents = state[\"documents\"]\n\n if not filtered_documents:\n # All documents have been filtered check_relevance\n # We will re-generate a new query\n print(\n \"---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---\"\n )\n return \"transform_query\"\n else:\n # We have relevant documents, so generate answer\n print(\"---DECISION: GENERATE---\")\n return \"generate\"\n\n\ndef grade_generation_v_documents_and_question(state):\n \"\"\"\n Determines whether the generation is grounded in the document and answers question.\n\n Args:\n state (dict): The current graph state\n\n Returns:\n str: Decision for next node to call\n \"\"\"\n\n print(\"---CHECK HALLUCINATIONS---\")\n question = state[\"question\"]\n documents = state[\"documents\"]\n generation = state[\"generation\"]\n\n score = hallucination_grader.invoke(\n {\"documents\": documents, \"generation\": generation}\n )\n grade = score.binary_score\n\n # Check hallucination\n if grade == \"yes\":\n print(\"---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---\")\n # Check question-answering\n print(\"---GRADE GENERATION vs QUESTION---\")\n score = answer_grader.invoke({\"question\": question, \"generation\": generation})\n grade = score.binary_score\n if grade == \"yes\":\n print(\"---DECISION: GENERATION ADDRESSES QUESTION---\")\n return \"useful\"\n else:\n print(\"---DECISION: GENERATION DOES NOT ADDRESS QUESTION---\")\n return \"not useful\"\n else:\n pprint(\"---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---\")\n return \"not supported\""]
|
||||
"source": [
|
||||
"### Edges\n\n\ndef decide_to_generate(state):\n \"\"\"\n Determines whether to generate an answer, or re-generate a question.\n\n Args:\n state (dict): The current graph state\n\n Returns:\n str: Binary decision for next node to call\n \"\"\"\n\n print(\"---ASSESS GRADED DOCUMENTS---\")\n state[\"question\"]\n filtered_documents = state[\"documents\"]\n\n if not filtered_documents:\n # All documents have been filtered check_relevance\n # We will re-generate a new query\n print(\n \"---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---\"\n )\n return \"transform_query\"\n else:\n # We have relevant documents, so generate answer\n print(\"---DECISION: GENERATE---\")\n return \"generate\"\n\n\ndef grade_generation_v_documents_and_question(state):\n \"\"\"\n Determines whether the generation is grounded in the document and answers question.\n\n Args:\n state (dict): The current graph state\n\n Returns:\n str: Decision for next node to call\n \"\"\"\n\n print(\"---CHECK HALLUCINATIONS---\")\n question = state[\"question\"]\n documents = state[\"documents\"]\n generation = state[\"generation\"]\n\n score = hallucination_grader.invoke(\n {\"documents\": documents, \"generation\": generation}\n )\n grade = score.binary_score\n\n # Check hallucination\n if grade == \"yes\":\n print(\"---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---\")\n # Check question-answering\n print(\"---GRADE GENERATION vs QUESTION---\")\n score = answer_grader.invoke({\"question\": question, \"generation\": generation})\n grade = score.binary_score\n if grade == \"yes\":\n print(\"---DECISION: GENERATION ADDRESSES QUESTION---\")\n return \"useful\"\n else:\n print(\"---DECISION: GENERATION DOES NOT ADDRESS QUESTION---\")\n return \"not useful\"\n else:\n pprint(\"---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---\")\n return \"not supported\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
@@ -331,7 +425,9 @@
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": ["from pprint import pprint\n\n# Run\ninputs = {\"question\": \"Movies that star Daniel Craig\"}\nfor output in app.stream(inputs):\n for key, value in output.items():\n # Node\n pprint(f\"Node '{key}':\")\n pprint(\"\\n---\\n\")\n\n# Final generation\npprint(value[\"generation\"])"]
|
||||
"source": [
|
||||
"from pprint import pprint\n\n# Run\ninputs = {\"question\": \"Movies that star Daniel Craig\"}\nfor output in app.stream(inputs):\n for key, value in output.items():\n # Node\n pprint(f\"Node '{key}':\")\n pprint(\"\\n---\\n\")\n\n# Final generation\npprint(value[\"generation\"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
@@ -339,7 +435,9 @@
|
||||
"id": "4138bc51-8c84-4b8a-8d24-f7f470721f6f",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": ["inputs = {\"question\": \"Which movies are about aliens?\"}\nfor output in app.stream(inputs):\n for key, value in output.items():\n # Node\n pprint(f\"Node '{key}':\")\n pprint(\"\\n---\\n\")\n\n# Final generation\npprint(value[\"generation\"])"]
|
||||
"source": [
|
||||
"inputs = {\"question\": \"Which movies are about aliens?\"}\nfor output in app.stream(inputs):\n for key, value in output.items():\n # Node\n pprint(f\"Node '{key}':\")\n pprint(\"\\n---\\n\")\n\n# Final generation\npprint(value[\"generation\"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
@@ -347,7 +445,9 @@
|
||||
"id": "42369ab8-322d-434a-b5dd-2266e4cb2903",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [""]
|
||||
"source": [
|
||||
""
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
|
||||
@@ -13,6 +13,20 @@ By default `langgraph-checkpoint-postgres` installs `psycopg` (Psycopg 3) withou
|
||||
|
||||
> [!IMPORTANT]
|
||||
> When manually creating Postgres connections and passing them to `PostgresSaver` or `AsyncPostgresSaver`, make sure to include `autocommit=True` and `row_factory=dict_row` (`from psycopg.rows import dict_row`). See a full example in this [how-to guide](https://langchain-ai.github.io/langgraph/how-tos/persistence_postgres/).
|
||||
>
|
||||
> **Why these parameters are required:**
|
||||
> - `autocommit=True`: Required for the `.setup()` method to properly commit the checkpoint tables to the database. Without this, table creation may not be persisted.
|
||||
> - `row_factory=dict_row`: Required because the PostgresSaver implementation accesses database rows using dictionary-style syntax (e.g., `row["column_name"]`). The default `tuple_row` factory returns tuples that only support index-based access (e.g., `row[0]`), which will cause `TypeError` exceptions when the checkpointer tries to access columns by name.
|
||||
>
|
||||
> **Example of incorrect usage:**
|
||||
> ```python
|
||||
> # ❌ This will fail with TypeError during checkpointer operations
|
||||
> with psycopg.connect(DB_URI) as conn: # Missing autocommit=True and row_factory=dict_row
|
||||
> checkpointer = PostgresSaver(conn)
|
||||
> checkpointer.setup() # May not persist tables properly
|
||||
> # Any operation that reads from database will fail with:
|
||||
> # TypeError: tuple indices must be integers or slices, not str
|
||||
> ```
|
||||
|
||||
```python
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
|
||||
@@ -23,6 +23,7 @@ from langgraph.checkpoint.base import (
|
||||
)
|
||||
from langgraph.checkpoint.postgres import _internal
|
||||
from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
|
||||
Conn = _internal.Conn # For backward compatibility
|
||||
@@ -175,32 +176,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
value["channel_values"],
|
||||
)
|
||||
for value in values:
|
||||
yield CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
{
|
||||
**value["checkpoint"],
|
||||
"channel_values": self._load_blobs(value["channel_values"]),
|
||||
},
|
||||
value["metadata"],
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
}
|
||||
}
|
||||
if value["parent_checkpoint_id"]
|
||||
else None
|
||||
),
|
||||
self._load_writes(value["pending_writes"]),
|
||||
)
|
||||
yield self._load_checkpoint_tuple(value)
|
||||
|
||||
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Get a checkpoint tuple from the database.
|
||||
@@ -271,32 +247,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
value["channel_values"],
|
||||
)
|
||||
|
||||
return CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
{
|
||||
**value["checkpoint"],
|
||||
"channel_values": self._load_blobs(value["channel_values"]),
|
||||
},
|
||||
value["metadata"],
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
}
|
||||
}
|
||||
if value["parent_checkpoint_id"]
|
||||
else None
|
||||
),
|
||||
self._load_writes(value["pending_writes"]),
|
||||
)
|
||||
return self._load_checkpoint_tuple(value)
|
||||
|
||||
def put(
|
||||
self,
|
||||
@@ -466,5 +417,44 @@ class PostgresSaver(BasePostgresSaver):
|
||||
with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
||||
"""
|
||||
Convert a database row into a CheckpointTuple object.
|
||||
|
||||
__all__ = ["PostgresSaver", "BasePostgresSaver", "Conn"]
|
||||
Args:
|
||||
value: A row from the database containing checkpoint data.
|
||||
|
||||
Returns:
|
||||
CheckpointTuple: A structured representation of the checkpoint,
|
||||
including its configuration, metadata, parent checkpoint (if any),
|
||||
and pending writes.
|
||||
"""
|
||||
return CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
{
|
||||
**value["checkpoint"],
|
||||
"channel_values": self._load_blobs(value["channel_values"]),
|
||||
},
|
||||
value["metadata"],
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
}
|
||||
}
|
||||
if value["parent_checkpoint_id"]
|
||||
else None
|
||||
),
|
||||
self._load_writes(value["pending_writes"]),
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["PostgresSaver", "BasePostgresSaver", "ShallowPostgresSaver", "Conn"]
|
||||
|
||||
@@ -23,6 +23,7 @@ from langgraph.checkpoint.base import (
|
||||
)
|
||||
from langgraph.checkpoint.postgres import _ainternal
|
||||
from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
|
||||
Conn = _ainternal.Conn # For backward compatibility
|
||||
@@ -162,32 +163,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
value["channel_values"],
|
||||
)
|
||||
for value in values:
|
||||
yield CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
{
|
||||
**value["checkpoint"],
|
||||
"channel_values": self._load_blobs(value["channel_values"]),
|
||||
},
|
||||
value["metadata"],
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
}
|
||||
}
|
||||
if value["parent_checkpoint_id"]
|
||||
else None
|
||||
),
|
||||
await asyncio.to_thread(self._load_writes, value["pending_writes"]),
|
||||
)
|
||||
yield await self._load_checkpoint_tuple(value)
|
||||
|
||||
async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Get a checkpoint tuple from the database asynchronously.
|
||||
@@ -238,32 +214,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
value["channel_values"],
|
||||
)
|
||||
|
||||
return CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
{
|
||||
**value["checkpoint"],
|
||||
"channel_values": self._load_blobs(value["channel_values"]),
|
||||
},
|
||||
value["metadata"],
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
}
|
||||
}
|
||||
if value["parent_checkpoint_id"]
|
||||
else None
|
||||
),
|
||||
await asyncio.to_thread(self._load_writes, value["pending_writes"]),
|
||||
)
|
||||
return await self._load_checkpoint_tuple(value)
|
||||
|
||||
async def aput(
|
||||
self,
|
||||
@@ -424,6 +375,45 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
||||
"""
|
||||
Convert a database row into a CheckpointTuple object.
|
||||
|
||||
Args:
|
||||
value: A row from the database containing checkpoint data.
|
||||
|
||||
Returns:
|
||||
CheckpointTuple: A structured representation of the checkpoint,
|
||||
including its configuration, metadata, parent checkpoint (if any),
|
||||
and pending writes.
|
||||
"""
|
||||
return CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
{
|
||||
**value["checkpoint"],
|
||||
"channel_values": self._load_blobs(value["channel_values"]),
|
||||
},
|
||||
value["metadata"],
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
}
|
||||
}
|
||||
if value["parent_checkpoint_id"]
|
||||
else None
|
||||
),
|
||||
await asyncio.to_thread(self._load_writes, value["pending_writes"]),
|
||||
)
|
||||
|
||||
def list(
|
||||
self,
|
||||
config: RunnableConfig | None,
|
||||
@@ -570,4 +560,4 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
).result()
|
||||
|
||||
|
||||
__all__ = ["AsyncPostgresSaver", "Conn"]
|
||||
__all__ = ["AsyncPostgresSaver", "AsyncShallowPostgresSaver", "Conn"]
|
||||
|
||||
@@ -168,7 +168,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
checkpoint["channel_versions"][TASKS] = (
|
||||
max(checkpoint["channel_versions"].values())
|
||||
if checkpoint["channel_versions"]
|
||||
else self.get_next_version(None)
|
||||
else self.get_next_version(None, None)
|
||||
)
|
||||
|
||||
def _load_blobs(
|
||||
@@ -246,7 +246,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
for idx, (channel, value) in enumerate(writes)
|
||||
]
|
||||
|
||||
def get_next_version(self, current: str | None) -> str:
|
||||
def get_next_version(self, current: str | None, channel: None) -> str:
|
||||
if current is None:
|
||||
current_v = 0
|
||||
elif isinstance(current, int):
|
||||
|
||||
@@ -0,0 +1,959 @@
|
||||
import asyncio
|
||||
import threading
|
||||
import warnings
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from typing import Any, Optional
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from psycopg import (
|
||||
AsyncConnection,
|
||||
AsyncCursor,
|
||||
AsyncPipeline,
|
||||
Capabilities,
|
||||
Connection,
|
||||
Cursor,
|
||||
Pipeline,
|
||||
)
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import AsyncConnectionPool, ConnectionPool
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
WRITES_IDX_MAP,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
get_checkpoint_metadata,
|
||||
)
|
||||
from langgraph.checkpoint.postgres import _ainternal, _internal
|
||||
from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
from langgraph.checkpoint.serde.types import TASKS
|
||||
|
||||
"""
|
||||
To add a new migration, add a new string to the MIGRATIONS list.
|
||||
The position of the migration in the list is the version number.
|
||||
"""
|
||||
MIGRATIONS = [
|
||||
"""CREATE TABLE IF NOT EXISTS checkpoint_migrations (
|
||||
v INTEGER PRIMARY KEY
|
||||
);""",
|
||||
"""CREATE TABLE IF NOT EXISTS checkpoints (
|
||||
thread_id TEXT NOT NULL,
|
||||
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
||||
type TEXT,
|
||||
checkpoint JSONB NOT NULL,
|
||||
metadata JSONB NOT NULL DEFAULT '{}',
|
||||
PRIMARY KEY (thread_id, checkpoint_ns)
|
||||
);""",
|
||||
"""CREATE TABLE IF NOT EXISTS checkpoint_blobs (
|
||||
thread_id TEXT NOT NULL,
|
||||
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
||||
channel TEXT NOT NULL,
|
||||
type TEXT NOT NULL,
|
||||
blob BYTEA,
|
||||
PRIMARY KEY (thread_id, checkpoint_ns, channel)
|
||||
);""",
|
||||
"""CREATE TABLE IF NOT EXISTS checkpoint_writes (
|
||||
thread_id TEXT NOT NULL,
|
||||
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
||||
checkpoint_id TEXT NOT NULL,
|
||||
task_id TEXT NOT NULL,
|
||||
idx INTEGER NOT NULL,
|
||||
channel TEXT NOT NULL,
|
||||
type TEXT,
|
||||
blob BYTEA NOT NULL,
|
||||
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
|
||||
);""",
|
||||
"""
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoints_thread_id_idx ON checkpoints(thread_id);
|
||||
""",
|
||||
"""
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_blobs_thread_id_idx ON checkpoint_blobs(thread_id);
|
||||
""",
|
||||
"""
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id);
|
||||
""",
|
||||
"""
|
||||
ALTER TABLE checkpoint_writes ADD COLUMN task_path TEXT NOT NULL DEFAULT '';
|
||||
""",
|
||||
]
|
||||
|
||||
SELECT_SQL = f"""
|
||||
select
|
||||
thread_id,
|
||||
checkpoint,
|
||||
checkpoint_ns,
|
||||
metadata,
|
||||
(
|
||||
select array_agg(array[bl.channel::bytea, bl.type::bytea, bl.blob])
|
||||
from jsonb_each_text(checkpoint -> 'channel_versions')
|
||||
inner join checkpoint_blobs bl
|
||||
on bl.thread_id = checkpoints.thread_id
|
||||
and bl.checkpoint_ns = checkpoints.checkpoint_ns
|
||||
and bl.channel = jsonb_each_text.key
|
||||
) as channel_values,
|
||||
(
|
||||
select
|
||||
array_agg(array[cw.task_id::text::bytea, cw.channel::bytea, cw.type::bytea, cw.blob] order by cw.task_id, cw.idx)
|
||||
from checkpoint_writes cw
|
||||
where cw.thread_id = checkpoints.thread_id
|
||||
and cw.checkpoint_ns = checkpoints.checkpoint_ns
|
||||
and cw.checkpoint_id = (checkpoint->>'id')
|
||||
) as pending_writes,
|
||||
(
|
||||
select array_agg(array[cw.type::bytea, cw.blob] order by cw.task_path, cw.task_id, cw.idx)
|
||||
from checkpoint_writes cw
|
||||
where cw.thread_id = checkpoints.thread_id
|
||||
and cw.checkpoint_ns = checkpoints.checkpoint_ns
|
||||
and cw.channel = '{TASKS}'
|
||||
) as pending_sends
|
||||
from checkpoints """
|
||||
|
||||
UPSERT_CHECKPOINT_BLOBS_SQL = """
|
||||
INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, type, blob)
|
||||
VALUES (%s, %s, %s, %s, %s)
|
||||
ON CONFLICT (thread_id, checkpoint_ns, channel) DO UPDATE SET
|
||||
type = EXCLUDED.type,
|
||||
blob = EXCLUDED.blob;
|
||||
"""
|
||||
|
||||
UPSERT_CHECKPOINTS_SQL = """
|
||||
INSERT INTO checkpoints (thread_id, checkpoint_ns, checkpoint, metadata)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
ON CONFLICT (thread_id, checkpoint_ns)
|
||||
DO UPDATE SET
|
||||
checkpoint = EXCLUDED.checkpoint,
|
||||
metadata = EXCLUDED.metadata;
|
||||
"""
|
||||
|
||||
UPSERT_CHECKPOINT_WRITES_SQL = """
|
||||
INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, blob)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||||
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO UPDATE SET
|
||||
channel = EXCLUDED.channel,
|
||||
type = EXCLUDED.type,
|
||||
blob = EXCLUDED.blob;
|
||||
"""
|
||||
|
||||
INSERT_CHECKPOINT_WRITES_SQL = """
|
||||
INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, blob)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||||
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
|
||||
"""
|
||||
|
||||
|
||||
def _dump_blobs(
|
||||
serde: SerializerProtocol,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
values: dict[str, Any],
|
||||
versions: ChannelVersions,
|
||||
) -> list[tuple[str, str, str, str, Optional[bytes]]]:
|
||||
if not versions:
|
||||
return []
|
||||
|
||||
return [
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
k,
|
||||
*(serde.dumps_typed(values[k]) if k in values else ("empty", None)),
|
||||
)
|
||||
for k in versions
|
||||
]
|
||||
|
||||
|
||||
class ShallowPostgresSaver(BasePostgresSaver):
|
||||
"""A checkpoint saver that uses Postgres to store checkpoints.
|
||||
|
||||
This checkpointer ONLY stores the most recent checkpoint and does NOT retain any history.
|
||||
It is meant to be a light-weight drop-in replacement for the PostgresSaver that
|
||||
supports most of the LangGraph persistence functionality with the exception of time travel.
|
||||
"""
|
||||
|
||||
SELECT_SQL = SELECT_SQL
|
||||
MIGRATIONS = MIGRATIONS
|
||||
UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
|
||||
UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
|
||||
UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
|
||||
INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL
|
||||
|
||||
lock: threading.Lock
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: _internal.Conn,
|
||||
pipe: Optional[Pipeline] = None,
|
||||
serde: Optional[SerializerProtocol] = None,
|
||||
) -> None:
|
||||
warnings.warn(
|
||||
"ShallowPostgresSaver is deprecated as of version 2.0.20 and will be removed in 3.0.0. "
|
||||
"Use PostgresSaver instead, and invoke the graph with `graph.invoke(..., checkpoint_during=False)`.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
super().__init__(serde=serde)
|
||||
if isinstance(conn, ConnectionPool) and pipe is not None:
|
||||
raise ValueError(
|
||||
"Pipeline should be used only with a single Connection, not ConnectionPool."
|
||||
)
|
||||
|
||||
self.conn = conn
|
||||
self.pipe = pipe
|
||||
self.lock = threading.Lock()
|
||||
self.supports_pipeline = Capabilities().has_pipeline()
|
||||
|
||||
@classmethod
|
||||
@contextmanager
|
||||
def from_conn_string(
|
||||
cls, conn_string: str, *, pipeline: bool = False
|
||||
) -> Iterator["ShallowPostgresSaver"]:
|
||||
"""Create a new ShallowPostgresSaver instance from a connection string.
|
||||
|
||||
Args:
|
||||
conn_string: The Postgres connection info string.
|
||||
pipeline: whether to use Pipeline
|
||||
|
||||
Returns:
|
||||
ShallowPostgresSaver: A new ShallowPostgresSaver instance.
|
||||
"""
|
||||
with Connection.connect(
|
||||
conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
|
||||
) as conn:
|
||||
if pipeline:
|
||||
with conn.pipeline() as pipe:
|
||||
yield cls(conn, pipe)
|
||||
else:
|
||||
yield cls(conn)
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Set up the checkpoint database asynchronously.
|
||||
|
||||
This method creates the necessary tables in the Postgres database if they don't
|
||||
already exist and runs database migrations. It MUST be called directly by the user
|
||||
the first time checkpointer is used.
|
||||
"""
|
||||
with self._cursor() as cur:
|
||||
cur.execute(self.MIGRATIONS[0])
|
||||
results = cur.execute(
|
||||
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
|
||||
)
|
||||
row = results.fetchone()
|
||||
if row is None:
|
||||
version = -1
|
||||
else:
|
||||
version = row["v"]
|
||||
for v, migration in zip(
|
||||
range(version + 1, len(self.MIGRATIONS)),
|
||||
self.MIGRATIONS[version + 1 :],
|
||||
):
|
||||
cur.execute(migration)
|
||||
cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})")
|
||||
if self.pipe:
|
||||
self.pipe.sync()
|
||||
|
||||
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 Postgres database based
|
||||
on the provided config. For ShallowPostgresSaver, this method returns a list with
|
||||
ONLY the most recent checkpoint.
|
||||
"""
|
||||
where, args = self._search_where(config, filter, before)
|
||||
query = self.SELECT_SQL + where
|
||||
if limit:
|
||||
query += f" LIMIT {limit}"
|
||||
with self._cursor() as cur:
|
||||
cur.execute(self.SELECT_SQL + where, args, binary=True)
|
||||
for value in cur:
|
||||
checkpoint: Checkpoint = {
|
||||
**value["checkpoint"],
|
||||
"channel_values": self._load_blobs(value["channel_values"]),
|
||||
"pending_sends": [
|
||||
self.serde.loads_typed((t.decode(), v))
|
||||
for t, v in value["pending_sends"]
|
||||
]
|
||||
if value["pending_sends"]
|
||||
else [],
|
||||
}
|
||||
yield CheckpointTuple(
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
},
|
||||
checkpoint=checkpoint,
|
||||
metadata=value["metadata"],
|
||||
pending_writes=self._load_writes(value["pending_writes"]),
|
||||
)
|
||||
|
||||
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
|
||||
"""Get a checkpoint tuple from the database.
|
||||
|
||||
This method retrieves a checkpoint tuple from the Postgres database based on the
|
||||
provided config (matching the thread ID in the config).
|
||||
|
||||
Args:
|
||||
config: The config to use for retrieving the checkpoint.
|
||||
|
||||
Returns:
|
||||
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_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
args = (thread_id, checkpoint_ns)
|
||||
where = "WHERE thread_id = %s AND checkpoint_ns = %s"
|
||||
|
||||
with self._cursor() as cur:
|
||||
cur.execute(
|
||||
self.SELECT_SQL + where,
|
||||
args,
|
||||
binary=True,
|
||||
)
|
||||
|
||||
for value in cur:
|
||||
checkpoint: Checkpoint = {
|
||||
**value["checkpoint"],
|
||||
"channel_values": self._load_blobs(value["channel_values"]),
|
||||
"pending_sends": [
|
||||
self.serde.loads_typed((t.decode(), v))
|
||||
for t, v in value["pending_sends"]
|
||||
]
|
||||
if value["pending_sends"]
|
||||
else [],
|
||||
}
|
||||
return CheckpointTuple(
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
},
|
||||
checkpoint=checkpoint,
|
||||
metadata=value["metadata"],
|
||||
pending_writes=self._load_writes(value["pending_writes"]),
|
||||
)
|
||||
|
||||
def put(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
"""Save a checkpoint to the database.
|
||||
|
||||
This method saves a checkpoint to the Postgres database. The checkpoint is associated
|
||||
with the provided config. For ShallowPostgresSaver, this method saves ONLY the most recent
|
||||
checkpoint and overwrites a previous checkpoint, if it exists.
|
||||
|
||||
Args:
|
||||
config: The config to associate with the checkpoint.
|
||||
checkpoint: The checkpoint to save.
|
||||
metadata: Additional metadata to save with the checkpoint.
|
||||
new_versions: New channel versions as of this write.
|
||||
|
||||
Returns:
|
||||
RunnableConfig: Updated configuration after storing the checkpoint.
|
||||
|
||||
Examples:
|
||||
|
||||
>>> from langgraph.checkpoint.postgres import ShallowPostgresSaver
|
||||
>>> DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
|
||||
>>> with ShallowPostgresSaver.from_conn_string(DB_URI) as memory:
|
||||
>>> config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
|
||||
>>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}}
|
||||
>>> saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {})
|
||||
>>> print(saved_config)
|
||||
{'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}}
|
||||
"""
|
||||
configurable = config["configurable"].copy()
|
||||
thread_id = configurable.pop("thread_id")
|
||||
checkpoint_ns = configurable.pop("checkpoint_ns")
|
||||
|
||||
copy = checkpoint.copy()
|
||||
next_config = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
}
|
||||
|
||||
with self._cursor(pipeline=True) as cur:
|
||||
cur.execute(
|
||||
"""DELETE FROM checkpoint_writes
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id NOT IN (%s, %s)""",
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint["id"],
|
||||
configurable.get("checkpoint_id", ""),
|
||||
),
|
||||
)
|
||||
cur.executemany(
|
||||
self.UPSERT_CHECKPOINT_BLOBS_SQL,
|
||||
_dump_blobs(
|
||||
self.serde,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
copy.pop("channel_values"), # type: ignore[misc]
|
||||
new_versions,
|
||||
),
|
||||
)
|
||||
cur.execute(
|
||||
self.UPSERT_CHECKPOINTS_SQL,
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
Jsonb(copy),
|
||||
Jsonb(get_checkpoint_metadata(config, metadata)),
|
||||
),
|
||||
)
|
||||
return next_config
|
||||
|
||||
def put_writes(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
writes: Sequence[tuple[str, Any]],
|
||||
task_id: str,
|
||||
task_path: str = "",
|
||||
) -> None:
|
||||
"""Store intermediate writes linked to a checkpoint.
|
||||
|
||||
This method saves intermediate writes associated with a checkpoint to the Postgres database.
|
||||
|
||||
Args:
|
||||
config: Configuration of the related checkpoint.
|
||||
writes: List of writes to store.
|
||||
task_id: Identifier for the task creating the writes.
|
||||
"""
|
||||
query = (
|
||||
self.UPSERT_CHECKPOINT_WRITES_SQL
|
||||
if all(w[0] in WRITES_IDX_MAP for w in writes)
|
||||
else self.INSERT_CHECKPOINT_WRITES_SQL
|
||||
)
|
||||
with self._cursor(pipeline=True) as cur:
|
||||
cur.executemany(
|
||||
query,
|
||||
self._dump_writes(
|
||||
config["configurable"]["thread_id"],
|
||||
config["configurable"]["checkpoint_ns"],
|
||||
config["configurable"]["checkpoint_id"],
|
||||
task_id,
|
||||
task_path,
|
||||
writes,
|
||||
),
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
|
||||
"""Create a database cursor as a context manager.
|
||||
|
||||
Args:
|
||||
pipeline: whether to use pipeline for the DB operations inside the context manager.
|
||||
Will be applied regardless of whether the ShallowPostgresSaver instance was initialized with a pipeline.
|
||||
If pipeline mode is not supported, will fall back to using transaction context manager.
|
||||
"""
|
||||
with _internal.get_connection(self.conn) as conn:
|
||||
if self.pipe:
|
||||
# a connection in pipeline mode can be used concurrently
|
||||
# in multiple threads/coroutines, but only one cursor can be
|
||||
# used at a time
|
||||
try:
|
||||
with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
finally:
|
||||
if pipeline:
|
||||
self.pipe.sync()
|
||||
elif pipeline:
|
||||
# a connection not in pipeline mode can only be used by one
|
||||
# thread/coroutine at a time, so we acquire a lock
|
||||
if self.supports_pipeline:
|
||||
with (
|
||||
self.lock,
|
||||
conn.pipeline(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
# Use connection's transaction context manager when pipeline mode not supported
|
||||
with (
|
||||
self.lock,
|
||||
conn.transaction(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
|
||||
class AsyncShallowPostgresSaver(BasePostgresSaver):
|
||||
"""A checkpoint saver that uses Postgres to store checkpoints asynchronously.
|
||||
|
||||
This checkpointer ONLY stores the most recent checkpoint and does NOT retain any history.
|
||||
It is meant to be a light-weight drop-in replacement for the AsyncPostgresSaver that
|
||||
supports most of the LangGraph persistence functionality with the exception of time travel.
|
||||
"""
|
||||
|
||||
SELECT_SQL = SELECT_SQL
|
||||
MIGRATIONS = MIGRATIONS
|
||||
UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
|
||||
UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
|
||||
UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
|
||||
INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL
|
||||
lock: asyncio.Lock
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conn: _ainternal.Conn,
|
||||
pipe: Optional[AsyncPipeline] = None,
|
||||
serde: Optional[SerializerProtocol] = None,
|
||||
) -> None:
|
||||
warnings.warn(
|
||||
"AsyncShallowPostgresSaver is deprecated as of version 2.0.20 and will be removed in 3.0.0. "
|
||||
"Use AsyncPostgresSaver instead, and invoke the graph with `await graph.ainvoke(..., checkpoint_during=False)`.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
super().__init__(serde=serde)
|
||||
if isinstance(conn, AsyncConnectionPool) and pipe is not None:
|
||||
raise ValueError(
|
||||
"Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool."
|
||||
)
|
||||
|
||||
self.conn = conn
|
||||
self.pipe = pipe
|
||||
self.lock = asyncio.Lock()
|
||||
self.loop = asyncio.get_running_loop()
|
||||
self.supports_pipeline = Capabilities().has_pipeline()
|
||||
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
async def from_conn_string(
|
||||
cls,
|
||||
conn_string: str,
|
||||
*,
|
||||
pipeline: bool = False,
|
||||
serde: Optional[SerializerProtocol] = None,
|
||||
) -> AsyncIterator["AsyncShallowPostgresSaver"]:
|
||||
"""Create a new AsyncShallowPostgresSaver instance from a connection string.
|
||||
|
||||
Args:
|
||||
conn_string: The Postgres connection info string.
|
||||
pipeline: whether to use AsyncPipeline
|
||||
|
||||
Returns:
|
||||
AsyncShallowPostgresSaver: A new AsyncShallowPostgresSaver instance.
|
||||
"""
|
||||
async with await AsyncConnection.connect(
|
||||
conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
|
||||
) as conn:
|
||||
if pipeline:
|
||||
async with conn.pipeline() as pipe:
|
||||
yield cls(conn=conn, pipe=pipe, serde=serde)
|
||||
else:
|
||||
yield cls(conn=conn, serde=serde)
|
||||
|
||||
async def setup(self) -> None:
|
||||
"""Set up the checkpoint database asynchronously.
|
||||
|
||||
This method creates the necessary tables in the Postgres database if they don't
|
||||
already exist and runs database migrations. It MUST be called directly by the user
|
||||
the first time checkpointer is used.
|
||||
"""
|
||||
async with self._cursor() as cur:
|
||||
await cur.execute(self.MIGRATIONS[0])
|
||||
results = await cur.execute(
|
||||
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
|
||||
)
|
||||
row = await results.fetchone()
|
||||
if row is None:
|
||||
version = -1
|
||||
else:
|
||||
version = row["v"]
|
||||
for v, migration in zip(
|
||||
range(version + 1, len(self.MIGRATIONS)),
|
||||
self.MIGRATIONS[version + 1 :],
|
||||
):
|
||||
await cur.execute(migration)
|
||||
await cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})")
|
||||
if self.pipe:
|
||||
await self.pipe.sync()
|
||||
|
||||
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 Postgres database based
|
||||
on the provided config. For ShallowPostgresSaver, this method returns a list with
|
||||
ONLY the most recent checkpoint.
|
||||
"""
|
||||
where, args = self._search_where(config, filter, before)
|
||||
query = self.SELECT_SQL + where
|
||||
if limit:
|
||||
query += f" LIMIT {limit}"
|
||||
async with self._cursor() as cur:
|
||||
await cur.execute(self.SELECT_SQL + where, args, binary=True)
|
||||
async for value in cur:
|
||||
checkpoint: Checkpoint = {
|
||||
**value["checkpoint"],
|
||||
"channel_values": self._load_blobs(value["channel_values"]),
|
||||
"pending_sends": [
|
||||
self.serde.loads_typed((t.decode(), v))
|
||||
for t, v in value["pending_sends"]
|
||||
]
|
||||
if value["pending_sends"]
|
||||
else [],
|
||||
}
|
||||
yield CheckpointTuple(
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
},
|
||||
checkpoint=checkpoint,
|
||||
metadata=value["metadata"],
|
||||
pending_writes=await asyncio.to_thread(
|
||||
self._load_writes, value["pending_writes"]
|
||||
),
|
||||
)
|
||||
|
||||
async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
|
||||
"""Get a checkpoint tuple from the database asynchronously.
|
||||
|
||||
This method retrieves a checkpoint tuple from the Postgres database based on the
|
||||
provided config (matching the thread ID in the config).
|
||||
|
||||
Args:
|
||||
config: The config to use for retrieving the checkpoint.
|
||||
|
||||
Returns:
|
||||
Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
||||
"""
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
args = (thread_id, checkpoint_ns)
|
||||
where = "WHERE thread_id = %s AND checkpoint_ns = %s"
|
||||
|
||||
async with self._cursor() as cur:
|
||||
await cur.execute(
|
||||
self.SELECT_SQL + where,
|
||||
args,
|
||||
binary=True,
|
||||
)
|
||||
|
||||
async for value in cur:
|
||||
checkpoint: Checkpoint = {
|
||||
**value["checkpoint"],
|
||||
"channel_values": self._load_blobs(value["channel_values"]),
|
||||
"pending_sends": [
|
||||
self.serde.loads_typed((t.decode(), v))
|
||||
for t, v in value["pending_sends"]
|
||||
]
|
||||
if value["pending_sends"]
|
||||
else [],
|
||||
}
|
||||
return CheckpointTuple(
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
},
|
||||
checkpoint=checkpoint,
|
||||
metadata=value["metadata"],
|
||||
pending_writes=await asyncio.to_thread(
|
||||
self._load_writes, value["pending_writes"]
|
||||
),
|
||||
)
|
||||
|
||||
async def aput(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
"""Save a checkpoint to the database asynchronously.
|
||||
|
||||
This method saves a checkpoint to the Postgres database. The checkpoint is associated
|
||||
with the provided config. For AsyncShallowPostgresSaver, this method saves ONLY the most recent
|
||||
checkpoint and overwrites a previous checkpoint, if it exists.
|
||||
|
||||
Args:
|
||||
config: The config to associate with the checkpoint.
|
||||
checkpoint: The checkpoint to save.
|
||||
metadata: Additional metadata to save with the checkpoint.
|
||||
new_versions: New channel versions as of this write.
|
||||
|
||||
Returns:
|
||||
RunnableConfig: Updated configuration after storing the checkpoint.
|
||||
"""
|
||||
configurable = config["configurable"].copy()
|
||||
thread_id = configurable.pop("thread_id")
|
||||
checkpoint_ns = configurable.pop("checkpoint_ns")
|
||||
|
||||
copy = checkpoint.copy()
|
||||
next_config = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint["id"],
|
||||
}
|
||||
}
|
||||
|
||||
async with self._cursor(pipeline=True) as cur:
|
||||
await cur.execute(
|
||||
"""DELETE FROM checkpoint_writes
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id NOT IN (%s, %s)""",
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint["id"],
|
||||
configurable.get("checkpoint_id", ""),
|
||||
),
|
||||
)
|
||||
await cur.executemany(
|
||||
self.UPSERT_CHECKPOINT_BLOBS_SQL,
|
||||
_dump_blobs(
|
||||
self.serde,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
copy.pop("channel_values"), # type: ignore[misc]
|
||||
new_versions,
|
||||
),
|
||||
)
|
||||
await cur.execute(
|
||||
self.UPSERT_CHECKPOINTS_SQL,
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
Jsonb(copy),
|
||||
Jsonb(get_checkpoint_metadata(config, metadata)),
|
||||
),
|
||||
)
|
||||
return next_config
|
||||
|
||||
async def aput_writes(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
writes: Sequence[tuple[str, Any]],
|
||||
task_id: str,
|
||||
task_path: str = "",
|
||||
) -> None:
|
||||
"""Store intermediate writes linked to a checkpoint asynchronously.
|
||||
|
||||
This method saves intermediate writes associated with a checkpoint to the database.
|
||||
|
||||
Args:
|
||||
config: Configuration of the related checkpoint.
|
||||
writes: List of writes to store, each as (channel, value) pair.
|
||||
task_id: Identifier for the task creating the writes.
|
||||
"""
|
||||
query = (
|
||||
self.UPSERT_CHECKPOINT_WRITES_SQL
|
||||
if all(w[0] in WRITES_IDX_MAP for w in writes)
|
||||
else self.INSERT_CHECKPOINT_WRITES_SQL
|
||||
)
|
||||
params = await asyncio.to_thread(
|
||||
self._dump_writes,
|
||||
config["configurable"]["thread_id"],
|
||||
config["configurable"]["checkpoint_ns"],
|
||||
config["configurable"]["checkpoint_id"],
|
||||
task_id,
|
||||
task_path,
|
||||
writes,
|
||||
)
|
||||
async with self._cursor(pipeline=True) as cur:
|
||||
await cur.executemany(query, params)
|
||||
|
||||
@asynccontextmanager
|
||||
async def _cursor(
|
||||
self, *, pipeline: bool = False
|
||||
) -> AsyncIterator[AsyncCursor[DictRow]]:
|
||||
"""Create a database cursor as a context manager.
|
||||
|
||||
Args:
|
||||
pipeline: whether to use pipeline for the DB operations inside the context manager.
|
||||
Will be applied regardless of whether the AsyncShallowPostgresSaver instance was initialized with a pipeline.
|
||||
If pipeline mode is not supported, will fall back to using transaction context manager.
|
||||
"""
|
||||
async with _ainternal.get_connection(self.conn) as conn:
|
||||
if self.pipe:
|
||||
# a connection in pipeline mode can be used concurrently
|
||||
# in multiple threads/coroutines, but only one cursor can be
|
||||
# used at a time
|
||||
try:
|
||||
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
finally:
|
||||
if pipeline:
|
||||
await self.pipe.sync()
|
||||
elif pipeline:
|
||||
# a connection not in pipeline mode can only be used by one
|
||||
# thread/coroutine at a time, so we acquire a lock
|
||||
if self.supports_pipeline:
|
||||
async with (
|
||||
self.lock,
|
||||
conn.pipeline(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
# Use connection's transaction context manager when pipeline mode not supported
|
||||
async with (
|
||||
self.lock,
|
||||
conn.transaction(),
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
else:
|
||||
async with (
|
||||
self.lock,
|
||||
conn.cursor(binary=True, row_factory=dict_row) as cur,
|
||||
):
|
||||
yield cur
|
||||
|
||||
def list(
|
||||
self,
|
||||
config: 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 Postgres database based
|
||||
on the provided config. For ShallowPostgresSaver, this method returns a list with
|
||||
ONLY the most recent checkpoint.
|
||||
"""
|
||||
aiter_ = self.alist(config, filter=filter, before=before, limit=limit)
|
||||
while True:
|
||||
try:
|
||||
yield asyncio.run_coroutine_threadsafe(
|
||||
anext(aiter_), # noqa: F821
|
||||
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 Postgres database based on the
|
||||
provided config (matching the thread ID in the config).
|
||||
|
||||
Args:
|
||||
config: The config to use for retrieving the checkpoint.
|
||||
|
||||
Returns:
|
||||
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 AsyncShallowPostgresSaver are only allowed from a "
|
||||
"different thread. From the main thread, use the async interface."
|
||||
"For example, use `await checkpointer.aget_tuple(...)` or `await "
|
||||
"graph.ainvoke(...)`."
|
||||
)
|
||||
except RuntimeError:
|
||||
pass
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.aget_tuple(config), self.loop
|
||||
).result()
|
||||
|
||||
def put(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
"""Save a checkpoint to the database.
|
||||
|
||||
This method saves a checkpoint to the Postgres database. The checkpoint is associated
|
||||
with the provided config. For AsyncShallowPostgresSaver, this method saves ONLY the most recent
|
||||
checkpoint and overwrites a previous checkpoint, if it exists.
|
||||
|
||||
Args:
|
||||
config: The config to associate with the checkpoint.
|
||||
checkpoint: The checkpoint to save.
|
||||
metadata: Additional metadata to save with the checkpoint.
|
||||
new_versions: New channel versions as of this write.
|
||||
|
||||
Returns:
|
||||
RunnableConfig: Updated configuration after storing the checkpoint.
|
||||
"""
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.aput(config, checkpoint, metadata, new_versions), self.loop
|
||||
).result()
|
||||
|
||||
def put_writes(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
writes: Sequence[tuple[str, Any]],
|
||||
task_id: str,
|
||||
task_path: str = "",
|
||||
) -> None:
|
||||
"""Store intermediate writes linked to a checkpoint.
|
||||
|
||||
This method saves intermediate writes associated with a checkpoint to the database.
|
||||
|
||||
Args:
|
||||
config: Configuration of the related checkpoint.
|
||||
writes: List of writes to store, each as (channel, value) pair.
|
||||
task_id: Identifier for the task creating the writes.
|
||||
task_path: Path of the task creating the writes.
|
||||
"""
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.aput_writes(config, writes, task_id, task_path), self.loop
|
||||
).result()
|
||||
@@ -1317,12 +1317,12 @@ def _ensure_index_config(
|
||||
index_config = index_config.copy()
|
||||
tokenized: list[tuple[str, Literal["$"] | list[str]]] = []
|
||||
tot = 0
|
||||
text_fields = index_config.get("fields") or ["$"]
|
||||
if isinstance(text_fields, str):
|
||||
text_fields = [text_fields]
|
||||
if not isinstance(text_fields, list):
|
||||
raise ValueError(f"Text fields must be a list or a string. Got {text_fields}")
|
||||
for p in text_fields:
|
||||
fields = index_config.get("fields") or ["$"]
|
||||
if isinstance(fields, str):
|
||||
fields = [fields]
|
||||
if not isinstance(fields, list):
|
||||
raise ValueError(f"Text fields must be a list or a string. Got {fields}")
|
||||
for p in fields:
|
||||
if p == "$":
|
||||
tokenized.append((p, "$"))
|
||||
tot += 1
|
||||
|
||||
@@ -1,53 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Protocol
|
||||
|
||||
from langgraph.checkpoint.base import Checkpoint, EmptyChannelError
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
|
||||
|
||||
class ChannelProtocol(Protocol):
|
||||
def checkpoint(self) -> Any | None: ...
|
||||
|
||||
|
||||
def empty_checkpoint() -> Checkpoint:
|
||||
return Checkpoint(
|
||||
v=1,
|
||||
id=str(uuid6(clock_seq=-2)),
|
||||
ts=datetime.now(timezone.utc).isoformat(),
|
||||
channel_values={},
|
||||
channel_versions={},
|
||||
versions_seen={},
|
||||
)
|
||||
|
||||
|
||||
def create_checkpoint(
|
||||
checkpoint: Checkpoint,
|
||||
channels: Mapping[str, ChannelProtocol] | None,
|
||||
step: int,
|
||||
*,
|
||||
id: str | None = None,
|
||||
) -> Checkpoint:
|
||||
"""Create a checkpoint for the given channels."""
|
||||
ts = datetime.now(timezone.utc).isoformat()
|
||||
if channels is None:
|
||||
values = checkpoint["channel_values"]
|
||||
else:
|
||||
values = {}
|
||||
for k, v in channels.items():
|
||||
if k not in checkpoint["channel_versions"]:
|
||||
continue
|
||||
try:
|
||||
values[k] = v.checkpoint()
|
||||
except EmptyChannelError:
|
||||
pass
|
||||
return Checkpoint(
|
||||
v=1,
|
||||
ts=ts,
|
||||
id=id or str(uuid6(clock_seq=step)),
|
||||
channel_values=values,
|
||||
channel_versions=checkpoint["channel_versions"],
|
||||
versions_seen=checkpoint["versions_seen"],
|
||||
)
|
||||
@@ -14,10 +14,14 @@ from langgraph.checkpoint.base import (
|
||||
EXCLUDED_METADATA_KEYS,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.aio import (
|
||||
AsyncPostgresSaver,
|
||||
AsyncShallowPostgresSaver,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
from langgraph.checkpoint.serde.types import TASKS
|
||||
from tests.checkpoint_utils import create_checkpoint, empty_checkpoint
|
||||
from tests.conftest import DEFAULT_POSTGRES_URI
|
||||
|
||||
|
||||
@@ -108,11 +112,41 @@ async def _base_saver():
|
||||
await conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _shallow_saver():
|
||||
"""Fixture for shallow connection mode testing."""
|
||||
database = f"test_{uuid4().hex[:16]}"
|
||||
# create unique db
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI + database,
|
||||
autocommit=True,
|
||||
prepare_threshold=0,
|
||||
row_factory=dict_row,
|
||||
) as conn:
|
||||
checkpointer = AsyncShallowPostgresSaver(conn)
|
||||
await checkpointer.setup()
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI, autocommit=True
|
||||
) as conn:
|
||||
await conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _saver(name: str):
|
||||
if name == "base":
|
||||
async with _base_saver() as saver:
|
||||
yield saver
|
||||
elif name == "shallow":
|
||||
async with _shallow_saver() as saver:
|
||||
yield saver
|
||||
elif name == "pool":
|
||||
async with _pool_saver() as saver:
|
||||
yield saver
|
||||
@@ -172,7 +206,7 @@ def test_data():
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
||||
async def test_combined_metadata(saver_name: str, test_data) -> None:
|
||||
async with _saver(saver_name) as saver:
|
||||
config = {
|
||||
@@ -199,7 +233,7 @@ async def test_combined_metadata(saver_name: str, test_data) -> None:
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
||||
async def test_asearch(saver_name: str, test_data) -> None:
|
||||
async with _saver(saver_name) as saver:
|
||||
configs = test_data["configs"]
|
||||
@@ -250,7 +284,7 @@ async def test_asearch(saver_name: str, test_data) -> None:
|
||||
} == {"", "inner"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
||||
async def test_null_chars(saver_name: str, test_data) -> None:
|
||||
async with _saver(saver_name) as saver:
|
||||
config = await saver.aput(
|
||||
|
||||
@@ -15,10 +15,11 @@ from langgraph.checkpoint.base import (
|
||||
EXCLUDED_METADATA_KEYS,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver
|
||||
from langgraph.checkpoint.serde.types import TASKS
|
||||
from tests.checkpoint_utils import create_checkpoint, empty_checkpoint
|
||||
from tests.conftest import DEFAULT_POSTGRES_URI
|
||||
|
||||
|
||||
@@ -97,11 +98,37 @@ def _base_saver():
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _shallow_saver():
|
||||
"""Fixture for regular connection mode testing with a shallow checkpointer."""
|
||||
database = f"test_{uuid4().hex[:16]}"
|
||||
# create unique db
|
||||
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
||||
conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
with Connection.connect(
|
||||
DEFAULT_POSTGRES_URI + database,
|
||||
autocommit=True,
|
||||
prepare_threshold=0,
|
||||
row_factory=dict_row,
|
||||
) as conn:
|
||||
checkpointer = ShallowPostgresSaver(conn)
|
||||
checkpointer.setup()
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _saver(name: str):
|
||||
if name == "base":
|
||||
with _base_saver() as saver:
|
||||
yield saver
|
||||
elif name == "shallow":
|
||||
with _shallow_saver() as saver:
|
||||
yield saver
|
||||
elif name == "pool":
|
||||
with _pool_saver() as saver:
|
||||
yield saver
|
||||
@@ -161,7 +188,7 @@ def test_data():
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
||||
def test_combined_metadata(saver_name: str, test_data) -> None:
|
||||
with _saver(saver_name) as saver:
|
||||
config = {
|
||||
@@ -188,7 +215,7 @@ def test_combined_metadata(saver_name: str, test_data) -> None:
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
||||
def test_search(saver_name: str, test_data) -> None:
|
||||
with _saver(saver_name) as saver:
|
||||
configs = test_data["configs"]
|
||||
@@ -237,7 +264,7 @@ def test_search(saver_name: str, test_data) -> None:
|
||||
} == {"", "inner"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
||||
def test_null_chars(saver_name: str, test_data) -> None:
|
||||
with _saver(saver_name) as saver:
|
||||
config = saver.put(
|
||||
|
||||
Generated
+706
-703
File diff suppressed because it is too large
Load Diff
@@ -536,7 +536,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
|
||||
"""
|
||||
raise NotImplementedError(_AIO_ERROR_MSG)
|
||||
|
||||
def get_next_version(self, current: str | None) -> str:
|
||||
def get_next_version(self, current: str | None, channel: None) -> str:
|
||||
"""Generate the next version ID for a channel.
|
||||
|
||||
This method creates a new version identifier for a channel based on its current version.
|
||||
|
||||
@@ -591,7 +591,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
||||
)
|
||||
await self.conn.commit()
|
||||
|
||||
def get_next_version(self, current: str | None) -> str:
|
||||
def get_next_version(self, current: str | None, channel: None) -> str:
|
||||
"""Generate the next version ID for a channel.
|
||||
|
||||
This method creates a new version identifier for a channel based on its current version.
|
||||
|
||||
@@ -1,53 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Protocol
|
||||
|
||||
from langgraph.checkpoint.base import Checkpoint, EmptyChannelError
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
|
||||
|
||||
class ChannelProtocol(Protocol):
|
||||
def checkpoint(self) -> Any | None: ...
|
||||
|
||||
|
||||
def empty_checkpoint() -> Checkpoint:
|
||||
return Checkpoint(
|
||||
v=1,
|
||||
id=str(uuid6(clock_seq=-2)),
|
||||
ts=datetime.now(timezone.utc).isoformat(),
|
||||
channel_values={},
|
||||
channel_versions={},
|
||||
versions_seen={},
|
||||
)
|
||||
|
||||
|
||||
def create_checkpoint(
|
||||
checkpoint: Checkpoint,
|
||||
channels: Mapping[str, ChannelProtocol] | None,
|
||||
step: int,
|
||||
*,
|
||||
id: str | None = None,
|
||||
) -> Checkpoint:
|
||||
"""Create a checkpoint for the given channels."""
|
||||
ts = datetime.now(timezone.utc).isoformat()
|
||||
if channels is None:
|
||||
values = checkpoint["channel_values"]
|
||||
else:
|
||||
values = {}
|
||||
for k, v in channels.items():
|
||||
if k not in checkpoint["channel_versions"]:
|
||||
continue
|
||||
try:
|
||||
values[k] = v.checkpoint()
|
||||
except EmptyChannelError:
|
||||
pass
|
||||
return Checkpoint(
|
||||
v=1,
|
||||
ts=ts,
|
||||
id=id or str(uuid6(clock_seq=step)),
|
||||
channel_values=values,
|
||||
channel_versions=checkpoint["channel_versions"],
|
||||
versions_seen=checkpoint["versions_seen"],
|
||||
)
|
||||
@@ -6,9 +6,10 @@ from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||||
from tests.checkpoint_utils import create_checkpoint, empty_checkpoint
|
||||
|
||||
|
||||
class TestAsyncSqliteSaver:
|
||||
|
||||
@@ -6,10 +6,11 @@ from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
from langgraph.checkpoint.sqlite.utils import _metadata_predicate, search_where
|
||||
from tests.checkpoint_utils import create_checkpoint, empty_checkpoint
|
||||
|
||||
|
||||
class TestSqliteSaver:
|
||||
|
||||
Generated
+654
-650
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
|
||||
from typing import ( # noqa: UP035
|
||||
Any,
|
||||
Generic,
|
||||
@@ -13,6 +13,7 @@ from typing import ( # noqa: UP035
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
from langgraph.checkpoint.serde.types import (
|
||||
@@ -20,6 +21,7 @@ from langgraph.checkpoint.serde.types import (
|
||||
INTERRUPT,
|
||||
RESUME,
|
||||
SCHEDULED,
|
||||
ChannelProtocol,
|
||||
)
|
||||
|
||||
V = TypeVar("V", int, float, str)
|
||||
@@ -89,6 +91,7 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
|
||||
channel_values=checkpoint["channel_values"].copy(),
|
||||
channel_versions=checkpoint["channel_versions"].copy(),
|
||||
versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()},
|
||||
pending_sends=checkpoint.get("pending_sends", []).copy(),
|
||||
)
|
||||
|
||||
|
||||
@@ -125,6 +128,15 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
) -> None:
|
||||
self.serde = maybe_add_typed_methods(serde or self.serde)
|
||||
|
||||
@property
|
||||
def config_specs(self) -> list:
|
||||
"""Define the configuration options for the checkpoint saver.
|
||||
|
||||
Returns:
|
||||
list: List of configuration field specs.
|
||||
"""
|
||||
return []
|
||||
|
||||
def get(self, config: RunnableConfig) -> Checkpoint | None:
|
||||
"""Fetch a checkpoint using the given configuration.
|
||||
|
||||
@@ -334,7 +346,7 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_next_version(self, current: V | None) -> V:
|
||||
def get_next_version(self, current: V | None, channel: None) -> V:
|
||||
"""Generate the next version ID for a channel.
|
||||
|
||||
Default is to use integer versions, incrementing by 1. If you override, you can use str/int/float versions,
|
||||
@@ -342,6 +354,7 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
|
||||
Args:
|
||||
current: The current version identifier (int, float, or str).
|
||||
channel: Deprecated argument, kept for backwards compatibility.
|
||||
|
||||
Returns:
|
||||
V: The next version identifier, which must be increasing.
|
||||
@@ -404,3 +417,54 @@ EXCLUDED_METADATA_KEYS = {
|
||||
"checkpoint_ns",
|
||||
"checkpoint_map",
|
||||
}
|
||||
|
||||
# --- below are deprecated utilities used by past versions of LangGraph ---
|
||||
|
||||
LATEST_VERSION = 2
|
||||
|
||||
|
||||
def empty_checkpoint() -> Checkpoint:
|
||||
from datetime import datetime, timezone
|
||||
|
||||
return Checkpoint(
|
||||
v=LATEST_VERSION,
|
||||
id=str(uuid6(clock_seq=-2)),
|
||||
ts=datetime.now(timezone.utc).isoformat(),
|
||||
channel_values={},
|
||||
channel_versions={},
|
||||
versions_seen={},
|
||||
pending_sends=[],
|
||||
)
|
||||
|
||||
|
||||
def create_checkpoint(
|
||||
checkpoint: Checkpoint,
|
||||
channels: Mapping[str, ChannelProtocol] | None,
|
||||
step: int,
|
||||
*,
|
||||
id: str | None = None,
|
||||
) -> Checkpoint:
|
||||
"""Create a checkpoint for the given channels."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
ts = datetime.now(timezone.utc).isoformat()
|
||||
if channels is None:
|
||||
values = checkpoint["channel_values"]
|
||||
else:
|
||||
values = {}
|
||||
for k, v in channels.items():
|
||||
if k not in checkpoint["channel_versions"]:
|
||||
continue
|
||||
try:
|
||||
values[k] = v.checkpoint()
|
||||
except EmptyChannelError:
|
||||
pass
|
||||
return Checkpoint(
|
||||
v=LATEST_VERSION,
|
||||
ts=ts,
|
||||
id=id or str(uuid6(clock_seq=step)),
|
||||
channel_values=values,
|
||||
channel_versions=checkpoint["channel_versions"],
|
||||
versions_seen=checkpoint["versions_seen"],
|
||||
pending_sends=checkpoint.get("pending_sends", []),
|
||||
)
|
||||
|
||||
@@ -512,7 +512,7 @@ class InMemorySaver(
|
||||
"""
|
||||
return self.delete_thread(thread_id)
|
||||
|
||||
def get_next_version(self, current: str | None) -> str:
|
||||
def get_next_version(self, current: str | None, channel: None) -> str:
|
||||
if current is None:
|
||||
current_v = 0
|
||||
elif isinstance(current, int):
|
||||
|
||||
@@ -7,6 +7,7 @@ import json
|
||||
import pathlib
|
||||
import pickle
|
||||
import re
|
||||
import sys
|
||||
from collections import deque
|
||||
from collections.abc import Sequence
|
||||
from datetime import date, datetime, time, timedelta, timezone
|
||||
@@ -251,6 +252,7 @@ EXT_CONSTRUCTOR_KW_ARGS = 2
|
||||
EXT_METHOD_SINGLE_ARG = 3
|
||||
EXT_PYDANTIC_V1 = 4
|
||||
EXT_PYDANTIC_V2 = 5
|
||||
EXT_NUMPY_ARRAY = 6
|
||||
|
||||
|
||||
def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
|
||||
@@ -320,13 +322,6 @@ def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
|
||||
(obj.__class__.__module__, obj.__class__.__name__, obj.hex),
|
||||
),
|
||||
)
|
||||
elif isinstance(obj, bytearray):
|
||||
return ormsgpack.Ext(
|
||||
EXT_CONSTRUCTOR_SINGLE_ARG,
|
||||
_msgpack_enc(
|
||||
(obj.__class__.__module__, obj.__class__.__name__, bytes(obj)),
|
||||
),
|
||||
)
|
||||
elif isinstance(obj, decimal.Decimal):
|
||||
return ormsgpack.Ext(
|
||||
EXT_CONSTRUCTOR_SINGLE_ARG,
|
||||
@@ -465,6 +460,22 @@ def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
|
||||
),
|
||||
),
|
||||
)
|
||||
elif (np_mod := sys.modules.get("numpy")) is not None and isinstance(
|
||||
obj, np_mod.ndarray
|
||||
):
|
||||
order = "F" if obj.flags.f_contiguous and not obj.flags.c_contiguous else "C"
|
||||
if obj.flags.c_contiguous:
|
||||
mv = memoryview(obj)
|
||||
try:
|
||||
meta = (obj.dtype.str, obj.shape, order, mv)
|
||||
return ormsgpack.Ext(EXT_NUMPY_ARRAY, _msgpack_enc(meta))
|
||||
finally:
|
||||
mv.release()
|
||||
else:
|
||||
buf = obj.tobytes(order="A")
|
||||
meta = (obj.dtype.str, obj.shape, order, buf)
|
||||
return ormsgpack.Ext(EXT_NUMPY_ARRAY, _msgpack_enc(meta))
|
||||
|
||||
elif isinstance(obj, BaseException):
|
||||
return repr(obj)
|
||||
else:
|
||||
@@ -546,6 +557,17 @@ def _msgpack_ext_hook(code: int, data: bytes) -> Any:
|
||||
return tup[2]
|
||||
except NameError:
|
||||
return
|
||||
elif code == EXT_NUMPY_ARRAY:
|
||||
try:
|
||||
import numpy as _np
|
||||
|
||||
dtype_str, shape, order, buf = ormsgpack.unpackb(
|
||||
data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
arr = _np.frombuffer(buf, dtype=_np.dtype(dtype_str))
|
||||
return arr.reshape(shape, order=order)
|
||||
except Exception:
|
||||
return
|
||||
|
||||
|
||||
def _msgpack_ext_hook_to_json(code: int, data: bytes) -> Any:
|
||||
@@ -626,6 +648,19 @@ def _msgpack_ext_hook_to_json(code: int, data: bytes) -> Any:
|
||||
return tup[2]
|
||||
except Exception:
|
||||
return
|
||||
elif code == EXT_NUMPY_ARRAY:
|
||||
try:
|
||||
import numpy as _np
|
||||
|
||||
dtype_str, shape, order, buf = ormsgpack.unpackb(
|
||||
data,
|
||||
ext_hook=_msgpack_ext_hook_to_json,
|
||||
option=ormsgpack.OPT_NON_STR_KEYS,
|
||||
)
|
||||
arr = _np.frombuffer(buf, dtype=_np.dtype(dtype_str))
|
||||
return arr.reshape(shape, order=order).tolist()
|
||||
except Exception:
|
||||
return
|
||||
|
||||
|
||||
_option = (
|
||||
|
||||
@@ -1,4 +1,13 @@
|
||||
from typing import Any, Protocol, TypeVar, runtime_checkable
|
||||
from collections.abc import Sequence
|
||||
from typing import (
|
||||
Any,
|
||||
Optional,
|
||||
Protocol,
|
||||
TypeVar,
|
||||
runtime_checkable,
|
||||
)
|
||||
|
||||
from typing_extensions import Self
|
||||
|
||||
ERROR = "__error__"
|
||||
SCHEDULED = "__scheduled__"
|
||||
@@ -11,6 +20,25 @@ Update = TypeVar("Update", contravariant=True)
|
||||
C = TypeVar("C")
|
||||
|
||||
|
||||
class ChannelProtocol(Protocol[Value, Update, C]):
|
||||
# Mirrors langgraph.channels.base.BaseChannel
|
||||
@property
|
||||
def ValueType(self) -> Any: ...
|
||||
|
||||
@property
|
||||
def UpdateType(self) -> Any: ...
|
||||
|
||||
def checkpoint(self) -> Optional[C]: ...
|
||||
|
||||
def from_checkpoint(self, checkpoint: Optional[C]) -> Self: ...
|
||||
|
||||
def update(self, values: Sequence[Update]) -> bool: ...
|
||||
|
||||
def get(self) -> Value: ...
|
||||
|
||||
def consume(self) -> bool: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SendProtocol(Protocol):
|
||||
# Mirrors langgraph.constants.Send
|
||||
|
||||
@@ -496,7 +496,7 @@ def _cosine_similarity(X: list[float], Y: list[list[float]]) -> list[float]:
|
||||
if not Y:
|
||||
return []
|
||||
if _check_numpy():
|
||||
import numpy as np # type: ignore[import-not-found]
|
||||
import numpy as np
|
||||
|
||||
X_arr = np.array(X) if not isinstance(X, np.ndarray) else X
|
||||
Y_arr = np.array(Y) if not isinstance(Y, np.ndarray) else Y
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "2.0.26"
|
||||
version = "2.1.0"
|
||||
description = "Library with base interfaces for LangGraph checkpoint savers."
|
||||
authors = []
|
||||
requires-python = ">=3.9"
|
||||
@@ -13,7 +13,7 @@ license = "MIT"
|
||||
license-files = ['LICENSE']
|
||||
dependencies = [
|
||||
"langchain-core>=0.2.38",
|
||||
"ormsgpack>=1.8.0",
|
||||
"ormsgpack>=1.10.0",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
@@ -29,6 +29,9 @@ dev = [
|
||||
"pytest-watcher",
|
||||
"mypy",
|
||||
"dataclasses-json",
|
||||
"numpy",
|
||||
"pandas",
|
||||
"pandas-stubs>=2.2.2.240807",
|
||||
]
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
|
||||
@@ -1,53 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Protocol
|
||||
|
||||
from langgraph.checkpoint.base import Checkpoint, EmptyChannelError
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
|
||||
|
||||
class ChannelProtocol(Protocol):
|
||||
def checkpoint(self) -> Any | None: ...
|
||||
|
||||
|
||||
def empty_checkpoint() -> Checkpoint:
|
||||
return Checkpoint(
|
||||
v=1,
|
||||
id=str(uuid6(clock_seq=-2)),
|
||||
ts=datetime.now(timezone.utc).isoformat(),
|
||||
channel_values={},
|
||||
channel_versions={},
|
||||
versions_seen={},
|
||||
)
|
||||
|
||||
|
||||
def create_checkpoint(
|
||||
checkpoint: Checkpoint,
|
||||
channels: Mapping[str, ChannelProtocol] | None,
|
||||
step: int,
|
||||
*,
|
||||
id: str | None = None,
|
||||
) -> Checkpoint:
|
||||
"""Create a checkpoint for the given channels."""
|
||||
ts = datetime.now(timezone.utc).isoformat()
|
||||
if channels is None:
|
||||
values = checkpoint["channel_values"]
|
||||
else:
|
||||
values = {}
|
||||
for k, v in channels.items():
|
||||
if k not in checkpoint["channel_versions"]:
|
||||
continue
|
||||
try:
|
||||
values[k] = v.checkpoint()
|
||||
except EmptyChannelError:
|
||||
pass
|
||||
return Checkpoint(
|
||||
v=1,
|
||||
ts=ts,
|
||||
id=id or str(uuid6(clock_seq=step)),
|
||||
channel_values=values,
|
||||
channel_versions=checkpoint["channel_versions"],
|
||||
versions_seen=checkpoint["versions_seen"],
|
||||
)
|
||||
@@ -11,6 +11,9 @@ from ipaddress import IPv4Address
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
import dataclasses_json
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from pydantic import BaseModel, SecretStr
|
||||
from pydantic.v1 import BaseModel as BaseModelV1
|
||||
from pydantic.v1 import SecretStr as SecretStrV1
|
||||
@@ -295,19 +298,174 @@ def test_serde_jsonplus_bytearray() -> None:
|
||||
assert serde.loads_typed(dumped) == some_bytearray
|
||||
|
||||
|
||||
def test_loads_cannot_find() -> None:
|
||||
@pytest.mark.parametrize(
|
||||
"arr",
|
||||
[
|
||||
np.arange(9, dtype=np.int32).reshape(3, 3),
|
||||
np.asfortranarray(np.arange(9, dtype=np.float64).reshape(3, 3)),
|
||||
np.arange(12, dtype=np.int16)[::2].reshape(3, 2),
|
||||
],
|
||||
)
|
||||
def test_serde_jsonplus_numpy_array(arr: np.ndarray) -> None:
|
||||
serde = JsonPlusSerializer()
|
||||
|
||||
dumped = (
|
||||
"json",
|
||||
b'{"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyPydanticccc"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}',
|
||||
)
|
||||
dumped = serde.dumps_typed(arr)
|
||||
assert dumped[0] == "msgpack"
|
||||
result = serde.loads_typed(dumped)
|
||||
assert isinstance(result, np.ndarray)
|
||||
assert result.dtype == arr.dtype
|
||||
assert np.array_equal(result, arr)
|
||||
|
||||
assert serde.loads_typed(dumped) is None, "Should return None if cannot find class"
|
||||
|
||||
dumped = (
|
||||
"json",
|
||||
b'{"lc": 2, "type": "constructor", "id": ["tests", "test_jsonpluss", "MyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}',
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"arr",
|
||||
[
|
||||
np.arange(6, dtype=np.float32).reshape(2, 3),
|
||||
np.asfortranarray(np.arange(4, dtype=np.complex128).reshape(2, 2)),
|
||||
],
|
||||
)
|
||||
def test_serde_jsonplus_numpy_array_json_hook(arr: np.ndarray) -> None:
|
||||
serde = JsonPlusSerializer(__unpack_ext_hook__=_msgpack_ext_hook_to_json)
|
||||
dumped = serde.dumps_typed(arr)
|
||||
assert dumped[0] == "msgpack"
|
||||
result = serde.loads_typed(dumped)
|
||||
assert isinstance(result, list)
|
||||
assert result == arr.tolist()
|
||||
|
||||
assert serde.loads_typed(dumped) is None, "Should return None if cannot find module"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"df",
|
||||
[
|
||||
pd.DataFrame(),
|
||||
pd.DataFrame({"int_col": [1, 2, 3]}),
|
||||
pd.DataFrame({"float_col": [1.1, 2.2, 3.3]}),
|
||||
pd.DataFrame({"str_col": ["a", "b", "c"]}),
|
||||
pd.DataFrame({"bool_col": [True, False, True]}),
|
||||
pd.DataFrame(
|
||||
{
|
||||
"datetime_col": [
|
||||
datetime(2024, 1, 1),
|
||||
datetime(2024, 1, 2),
|
||||
datetime(2024, 1, 3),
|
||||
]
|
||||
}
|
||||
),
|
||||
pd.DataFrame(
|
||||
{
|
||||
"int_col": [1, 2, 3],
|
||||
"float_col": [1.1, 2.2, 3.3],
|
||||
"str_col": ["a", "b", "c"],
|
||||
}
|
||||
),
|
||||
pd.DataFrame(
|
||||
{
|
||||
"int_col": [1, 2, None],
|
||||
"float_col": [1.1, None, 3.3],
|
||||
"str_col": ["a", None, "c"],
|
||||
}
|
||||
),
|
||||
pd.DataFrame({"cat_col": pd.Categorical(["a", "b", "a", "c"])}),
|
||||
pd.DataFrame(
|
||||
{
|
||||
"int8": pd.array([1, 2, 3], dtype="int8"),
|
||||
"int16": pd.array([10, 20, 30], dtype="int16"),
|
||||
"int32": pd.array([100, 200, 300], dtype="int32"),
|
||||
"int64": pd.array([1000, 2000, 3000], dtype="int64"),
|
||||
"float32": pd.array([1.1, 2.2, 3.3], dtype="float32"),
|
||||
"float64": pd.array([10.1, 20.2, 30.3], dtype="float64"),
|
||||
}
|
||||
),
|
||||
pd.DataFrame({"value": [1, 2, 3]}, index=["x", "y", "z"]),
|
||||
pd.DataFrame(
|
||||
[[1, 2, 3, 4]],
|
||||
columns=pd.MultiIndex.from_tuples(
|
||||
[("A", "X"), ("A", "Y"), ("B", "X"), ("B", "Y")]
|
||||
),
|
||||
),
|
||||
pd.DataFrame(
|
||||
{"value": [1, 2, 3]}, index=pd.date_range("2024-01-01", periods=3, freq="D")
|
||||
),
|
||||
pd.DataFrame(
|
||||
{
|
||||
"col1": range(1000),
|
||||
"col2": [f"str_{i}" for i in range(1000)],
|
||||
"col3": np.random.rand(1000),
|
||||
}
|
||||
),
|
||||
pd.DataFrame(
|
||||
{"tz_datetime": pd.date_range("2024-01-01", periods=3, freq="D", tz="UTC")}
|
||||
),
|
||||
pd.DataFrame({"timedelta": pd.to_timedelta([1, 2, 3], unit="D")}),
|
||||
pd.DataFrame({"period": pd.period_range("2024-01", periods=3, freq="M")}),
|
||||
pd.DataFrame({"interval": pd.interval_range(start=0, end=3, periods=3)}),
|
||||
pd.DataFrame({"unicode": ["Hello 🌍", "Python 🐍", "Data 📊"]}),
|
||||
pd.DataFrame({"mixed": [1, "string", [1, 2, 3], {"key": "value"}]}),
|
||||
pd.DataFrame({"a": [1], "b": ["test"], "c": [3.14]}),
|
||||
pd.DataFrame({"single": [42]}),
|
||||
pd.DataFrame(
|
||||
{
|
||||
"small": [sys.float_info.min, 0, sys.float_info.max],
|
||||
"large_int": [-(2**63), 0, 2**63 - 1],
|
||||
}
|
||||
),
|
||||
pd.DataFrame({"special_strings": ["", "null", "None", "NaN", "inf", "-inf"]}),
|
||||
pd.DataFrame({"bytes_col": [b"hello", b"world", b"\x00\x01\x02"]}),
|
||||
],
|
||||
)
|
||||
def test_serde_jsonplus_pandas_dataframe(df: pd.DataFrame) -> None:
|
||||
serde = JsonPlusSerializer(pickle_fallback=True)
|
||||
|
||||
dumped = serde.dumps_typed(df)
|
||||
assert dumped[0] == "pickle"
|
||||
result = serde.loads_typed(dumped)
|
||||
assert result.equals(df)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"series",
|
||||
[
|
||||
pd.Series([]),
|
||||
pd.Series([1, 2, 3]),
|
||||
pd.Series([1.1, 2.2, 3.3]),
|
||||
pd.Series(["a", "b", "c"]),
|
||||
pd.Series([True, False, True]),
|
||||
pd.Series([datetime(2024, 1, 1), datetime(2024, 1, 2), datetime(2024, 1, 3)]),
|
||||
pd.Series([1, 2, None]),
|
||||
pd.Series([1.1, None, 3.3]),
|
||||
pd.Series(["a", None, "c"]),
|
||||
pd.Series(pd.Categorical(["a", "b", "a", "c"])),
|
||||
pd.Series([1, 2, 3], dtype="int8"),
|
||||
pd.Series([10, 20, 30], dtype="int16"),
|
||||
pd.Series([100, 200, 300], dtype="int32"),
|
||||
pd.Series([1000, 2000, 3000], dtype="int64"),
|
||||
pd.Series([1.1, 2.2, 3.3], dtype="float32"),
|
||||
pd.Series([10.1, 20.2, 30.3], dtype="float64"),
|
||||
pd.Series([1, 2, 3], index=["x", "y", "z"]),
|
||||
pd.Series([1, 2, 3], index=pd.date_range("2024-01-01", periods=3, freq="D")),
|
||||
pd.Series(range(1000)),
|
||||
pd.Series(pd.date_range("2024-01-01", periods=3, freq="D", tz="UTC")),
|
||||
pd.Series(pd.to_timedelta([1, 2, 3], unit="D")),
|
||||
pd.Series(pd.period_range("2024-01", periods=3, freq="M")),
|
||||
pd.Series(pd.interval_range(start=0, end=3, periods=3)),
|
||||
pd.Series(["Hello 🌍", "Python 🐍", "Data 📊"]),
|
||||
pd.Series([1, "string", [1, 2, 3], {"key": "value"}]),
|
||||
pd.Series([42], name="single"),
|
||||
pd.Series([sys.float_info.min, 0, sys.float_info.max]),
|
||||
pd.Series([-(2**63), 0, 2**63 - 1]),
|
||||
pd.Series(["", "null", "None", "NaN", "inf", "-inf"]),
|
||||
pd.Series([b"hello", b"world", b"\x00\x01\x02"]),
|
||||
pd.Series([1, 2, 3], name="named_series"),
|
||||
pd.Series(
|
||||
[10, 20],
|
||||
index=pd.MultiIndex.from_tuples([("a", 1), ("b", 2)], names=["x", "y"]),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_serde_jsonplus_pandas_series(series: pd.Series) -> None:
|
||||
serde = JsonPlusSerializer(pickle_fallback=True)
|
||||
dumped = serde.dumps_typed(series)
|
||||
|
||||
assert dumped[0] == "pickle"
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert result.equals(series)
|
||||
|
||||
@@ -6,12 +6,10 @@ from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
)
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from tests.checkpoint_utils import (
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
|
||||
class TestMemorySaver:
|
||||
|
||||
Generated
+987
-651
File diff suppressed because it is too large
Load Diff
@@ -2,9 +2,9 @@ from collections.abc import Sequence
|
||||
from typing import Annotated, Literal, TypedDict
|
||||
|
||||
from langchain_anthropic import ChatAnthropic
|
||||
from langchain_community.tools.tavily_search import TavilySearchResults
|
||||
from langchain_core.messages import BaseMessage
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_tavily import TavilySearchResults
|
||||
from langgraph.graph import END, StateGraph, add_messages
|
||||
from langgraph.prebuilt import ToolNode
|
||||
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
"python_version": "3.12",
|
||||
"dependencies": [
|
||||
"langchain_community",
|
||||
"langchain_tavily; python_version < '4.0'",
|
||||
"langchain_anthropic",
|
||||
"langchain_openai",
|
||||
"wikipedia",
|
||||
|
||||
@@ -3,6 +3,7 @@ import json
|
||||
from typing import Annotated, Optional
|
||||
|
||||
from langchain_community.retrievers import WikipediaRetriever
|
||||
from langchain_community.tools.tavily_search import TavilySearchResults
|
||||
from langchain_community.vectorstores import SKLearnVectorStore
|
||||
from langchain_core.documents import Document
|
||||
from langchain_core.messages import (
|
||||
@@ -17,15 +18,14 @@ from langchain_core.runnables import RunnableConfig, RunnableLambda
|
||||
from langchain_core.runnables import chain as as_runnable
|
||||
from langchain_core.tools import tool
|
||||
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
|
||||
from langchain_tavily import TavilySearchResults
|
||||
from langgraph.graph import END, StateGraph
|
||||
from pydantic import BaseModel, Field
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
fast_llm = ChatOpenAI(model="gpt-3.5-turbo")
|
||||
fast_llm = ChatOpenAI(model="gpt-4o-mini")
|
||||
# Uncomment for a Fireworks model
|
||||
# fast_llm = ChatFireworks(model="accounts/fireworks/models/firefunction-v1", max_tokens=32_000)
|
||||
long_context_llm = ChatOpenAI(model="gpt-4-turbo-preview")
|
||||
long_context_llm = ChatOpenAI(model="gpt-4o")
|
||||
|
||||
|
||||
direct_gen_outline_prompt = ChatPromptTemplate.from_messages(
|
||||
@@ -144,7 +144,7 @@ gen_perspectives_prompt = ChatPromptTemplate.from_messages(
|
||||
)
|
||||
|
||||
gen_perspectives_chain = gen_perspectives_prompt | ChatOpenAI(
|
||||
model="gpt-3.5-turbo"
|
||||
model="gpt-4o-mini"
|
||||
).with_structured_output(Perspectives)
|
||||
|
||||
|
||||
@@ -270,7 +270,7 @@ gen_queries_prompt = ChatPromptTemplate.from_messages(
|
||||
]
|
||||
)
|
||||
gen_queries_chain = gen_queries_prompt | ChatOpenAI(
|
||||
model="gpt-3.5-turbo"
|
||||
model="gpt-4o-mini"
|
||||
).with_structured_output(Queries, include_raw=True)
|
||||
|
||||
|
||||
|
||||
@@ -3,9 +3,9 @@ from pathlib import Path
|
||||
from typing import Annotated, TypedDict
|
||||
|
||||
from langchain_anthropic import ChatAnthropic
|
||||
from langchain_community.tools.tavily_search import TavilySearchResults
|
||||
from langchain_core.messages import BaseMessage
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_tavily import TavilySearchResults
|
||||
from langgraph.graph import END, StateGraph, add_messages
|
||||
from langgraph.prebuilt import ToolNode
|
||||
|
||||
|
||||
@@ -3,9 +3,9 @@ from pathlib import Path
|
||||
from typing import Annotated, TypedDict
|
||||
|
||||
from langchain_anthropic import ChatAnthropic
|
||||
from langchain_community.tools.tavily_search import TavilySearchResults
|
||||
from langchain_core.messages import BaseMessage
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_tavily import TavilySearchResults
|
||||
from langgraph.graph import END, StateGraph, add_messages
|
||||
from langgraph.prebuilt import ToolNode
|
||||
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
"pip_config_file": "./pipconf.txt",
|
||||
"dependencies": [
|
||||
"langchain_community",
|
||||
"langchain_tavily; python_version < '4.0'",
|
||||
"langchain_anthropic",
|
||||
"langchain_openai",
|
||||
"wikipedia",
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
import re
|
||||
import textwrap
|
||||
from collections import Counter
|
||||
from typing import Any, Literal, NamedTuple, Optional, TypedDict, Union
|
||||
@@ -382,6 +383,14 @@ class Config(TypedDict, total=False):
|
||||
Only relevant if Python dependencies are installed via pip. If omitted, default pip settings are used.
|
||||
"""
|
||||
|
||||
pip_installer: Optional[str]
|
||||
"""Optional. Python package installer to use ('auto', 'pip', 'uv').
|
||||
|
||||
- 'auto' (default): Use uv for supported base images, otherwise pip
|
||||
- 'pip': Force use of pip regardless of base image support
|
||||
- 'uv': Force use of uv (will fail if base image doesn't support it)
|
||||
"""
|
||||
|
||||
dockerfile_lines: list[str]
|
||||
"""Optional. Additional Docker instructions that will be appended to your base Dockerfile.
|
||||
|
||||
@@ -461,7 +470,7 @@ class Config(TypedDict, total=False):
|
||||
PIP_CLEANUP_LINES = """# -- Ensure user deps didn't inadvertently overwrite langgraph-api
|
||||
RUN mkdir -p /api/langgraph_api /api/langgraph_runtime /api/langgraph_license && \
|
||||
touch /api/langgraph_api/__init__.py /api/langgraph_runtime/__init__.py /api/langgraph_license/__init__.py
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir --no-deps -e /api
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 {install_cmd} --no-cache-dir --no-deps -e /api
|
||||
# -- End of ensuring user deps didn't inadvertently overwrite langgraph-api --
|
||||
# -- Removing pip from the final image ~<:===~~~ --
|
||||
RUN pip uninstall -y pip setuptools wheel && \
|
||||
@@ -470,6 +479,7 @@ RUN pip uninstall -y pip setuptools wheel && \
|
||||
# pip removal for wolfi
|
||||
RUN rm -rf /usr/lib/python*/site-packages/pip* /usr/lib/python*/site-packages/setuptools* /usr/lib/python*/site-packages/wheel* && \
|
||||
find /usr/bin -name "pip*" -delete || true
|
||||
{uv_removal}
|
||||
# -- End of pip removal --"""
|
||||
|
||||
|
||||
@@ -534,6 +544,7 @@ def validate_config(config: Config) -> Config:
|
||||
"node_version": node_version,
|
||||
"python_version": python_version,
|
||||
"pip_config_file": config.get("pip_config_file"),
|
||||
"pip_installer": config.get("pip_installer", "auto"),
|
||||
"_INTERNAL_docker_tag": config.get("_INTERNAL_docker_tag"),
|
||||
"base_image": config.get("base_image"),
|
||||
"image_distro": image_distro,
|
||||
@@ -598,6 +609,13 @@ def validate_config(config: Config) -> Config:
|
||||
"Must be either 'debian' or 'wolfi'."
|
||||
)
|
||||
|
||||
if pip_installer := config.get("pip_installer"):
|
||||
if pip_installer not in ["auto", "pip", "uv"]:
|
||||
raise click.UsageError(
|
||||
f"Invalid pip_installer: '{pip_installer}'. "
|
||||
"Must be 'auto', 'pip', or 'uv'."
|
||||
)
|
||||
|
||||
# Validate auth config
|
||||
if auth_conf := config.get("auth"):
|
||||
if "path" in auth_conf:
|
||||
@@ -1089,16 +1107,47 @@ def _get_node_pm_install_cmd(config_path: pathlib.Path, config: Config) -> str:
|
||||
return install_cmd
|
||||
|
||||
|
||||
semver_pattern = re.compile(r":(\d+(?:\.\d+)?(?:\.\d+)?)(?:-|$)")
|
||||
|
||||
|
||||
def _image_supports_uv(base_image: str) -> bool:
|
||||
if base_image == "langchain/langgraph-trial":
|
||||
return False
|
||||
match = semver_pattern.search(base_image)
|
||||
if not match:
|
||||
# Default image (langchain/langgraph-api) supports it.
|
||||
return True
|
||||
|
||||
version_str = match.group(1)
|
||||
version = tuple(map(int, version_str.split(".")))
|
||||
min_uv = (0, 2, 47)
|
||||
return version >= min_uv
|
||||
|
||||
|
||||
def python_config_to_docker(
|
||||
config_path: pathlib.Path,
|
||||
config: Config,
|
||||
base_image: str,
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
"""Generate a Dockerfile from the configuration."""
|
||||
pip_installer = config.get("pip_installer", "auto")
|
||||
|
||||
if pip_installer == "uv":
|
||||
install_cmd = "uv pip install --system"
|
||||
uv_removal = "RUN uv pip uninstall --system pip setuptools wheel && rm /usr/bin/uv /usr/bin/uvx"
|
||||
elif pip_installer == "pip":
|
||||
install_cmd = "pip install"
|
||||
uv_removal = ""
|
||||
else:
|
||||
if _image_supports_uv(base_image):
|
||||
install_cmd = "uv pip install --system"
|
||||
uv_removal = "RUN uv pip uninstall --system pip setuptools wheel && rm /usr/bin/uv /usr/bin/uvx"
|
||||
else:
|
||||
install_cmd = "pip install"
|
||||
uv_removal = ""
|
||||
|
||||
# configure pip
|
||||
pip_install = (
|
||||
"PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt"
|
||||
)
|
||||
pip_install = f"PYTHONDONTWRITEBYTECODE=1 {install_cmd} --no-cache-dir -c /api/constraints.txt"
|
||||
if config.get("pip_config_file"):
|
||||
pip_install = f"PIP_CONFIG_FILE=/pipconfig.txt {pip_install}"
|
||||
pip_config_file_str = (
|
||||
@@ -1151,7 +1200,10 @@ RUN set -ex && \\
|
||||
'name = "{fullpath.name}"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_{fullpath.name}/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency {fullpath.name} --"""
|
||||
@@ -1240,7 +1292,8 @@ ADD {relpath} /deps/{name}
|
||||
"",
|
||||
js_inst_str,
|
||||
"",
|
||||
PIP_CLEANUP_LINES, # Add pip cleanup after all installations are complete
|
||||
# Add pip cleanup after all installations are complete
|
||||
PIP_CLEANUP_LINES.format(install_cmd=install_cmd, uv_removal=uv_removal),
|
||||
"",
|
||||
f"WORKDIR {local_deps.working_dir}" if local_deps.working_dir else "",
|
||||
]
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-cli"
|
||||
version = "0.2.12"
|
||||
version = "0.3.3"
|
||||
description = "CLI for interacting with LangGraph API"
|
||||
authors = []
|
||||
requires-python = ">=3.9"
|
||||
|
||||
@@ -134,6 +134,17 @@
|
||||
],
|
||||
"description": "Optional. Linux distribution for the base image.\n\nMust be either 'debian' or 'wolfi'. If omitted, defaults to 'debian'.\n"
|
||||
},
|
||||
"pip_installer": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Optional. Python package installer to use ('auto', 'pip', 'uv').\n\n"
|
||||
},
|
||||
"store": {
|
||||
"anyOf": [
|
||||
{
|
||||
@@ -287,6 +298,17 @@
|
||||
],
|
||||
"description": "Optional. Linux distribution for the base image.\n\nMust be either 'debian' or 'wolfi'. If omitted, defaults to 'debian'.\n"
|
||||
},
|
||||
"pip_installer": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Optional. Python package installer to use ('auto', 'pip', 'uv').\n\n"
|
||||
},
|
||||
"store": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
||||
@@ -134,6 +134,17 @@
|
||||
],
|
||||
"description": "Optional. Linux distribution for the base image.\n\nMust be either 'debian' or 'wolfi'. If omitted, defaults to 'debian'.\n"
|
||||
},
|
||||
"pip_installer": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Optional. Python package installer to use ('auto', 'pip', 'uv').\n\n"
|
||||
},
|
||||
"store": {
|
||||
"anyOf": [
|
||||
{
|
||||
@@ -287,6 +298,17 @@
|
||||
],
|
||||
"description": "Optional. Linux distribution for the base image.\n\nMust be either 'debian' or 'wolfi'. If omitted, defaults to 'debian'.\n"
|
||||
},
|
||||
"pip_installer": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Optional. Python package installer to use ('auto', 'pip', 'uv').\n\n"
|
||||
},
|
||||
"store": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
||||
@@ -14,6 +14,10 @@ from langgraph_cli.config import PIP_CLEANUP_LINES, Config, validate_config
|
||||
from langgraph_cli.docker import DEFAULT_POSTGRES_URI, DockerCapabilities, Version
|
||||
from langgraph_cli.util import clean_empty_lines
|
||||
|
||||
FORMATTED_CLEANUP_LINES = PIP_CLEANUP_LINES.format(
|
||||
install_cmd="uv pip install --system",
|
||||
uv_removal="RUN uv pip uninstall --system pip setuptools wheel && rm /usr/bin/uv /usr/bin/uvx",
|
||||
)
|
||||
DEFAULT_DOCKER_CAPABILITIES = DockerCapabilities(
|
||||
version_docker=Version(26, 1, 1),
|
||||
version_compose=Version(2, 27, 0),
|
||||
@@ -144,10 +148,10 @@ services:
|
||||
COPY --from=cli_1 . /deps/cli_1
|
||||
# -- End of local package ../../.. --
|
||||
# -- Installing all local dependencies --
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
# -- End of local dependencies install --
|
||||
ENV LANGSERVE_GRAPHS='{{"agent": "agent.py:graph"}}'
|
||||
{textwrap.indent(textwrap.dedent(PIP_CLEANUP_LINES), " ")}
|
||||
{textwrap.indent(textwrap.dedent(FORMATTED_CLEANUP_LINES), " ")}
|
||||
WORKDIR /deps/cli
|
||||
|
||||
develop:
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import copy
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
@@ -17,6 +18,11 @@ from langgraph_cli.config import (
|
||||
)
|
||||
from langgraph_cli.util import clean_empty_lines
|
||||
|
||||
FORMATTED_CLEANUP_LINES = PIP_CLEANUP_LINES.format(
|
||||
install_cmd="uv pip install --system",
|
||||
uv_removal="RUN uv pip uninstall --system pip setuptools wheel && rm /usr/bin/uv /usr/bin/uvx",
|
||||
)
|
||||
|
||||
PATH_TO_CONFIG = pathlib.Path(__file__).parent / "test_config.json"
|
||||
|
||||
|
||||
@@ -35,6 +41,7 @@ def test_validate_config():
|
||||
"python_version": "3.11",
|
||||
"node_version": None,
|
||||
"pip_config_file": None,
|
||||
"pip_installer": "auto",
|
||||
"image_distro": "debian",
|
||||
"dockerfile_lines": [],
|
||||
"env": {},
|
||||
@@ -56,6 +63,7 @@ def test_validate_config():
|
||||
"python_version": "3.12",
|
||||
"node_version": None,
|
||||
"pip_config_file": "pipconfig.txt",
|
||||
"pip_installer": "auto",
|
||||
"image_distro": "debian",
|
||||
"dockerfile_lines": ["ARG meow"],
|
||||
"dependencies": [".", "langchain"],
|
||||
@@ -211,6 +219,74 @@ def test_validate_config_image_distro():
|
||||
assert config["image_distro"] == "debian"
|
||||
|
||||
|
||||
def test_validate_config_pip_installer():
|
||||
"""Test validation of pip_installer field."""
|
||||
# Valid pip_installer values should work
|
||||
config = validate_config(
|
||||
{
|
||||
"python_version": "3.11",
|
||||
"dependencies": ["."],
|
||||
"graphs": {"agent": "./agent.py:graph"},
|
||||
"pip_installer": "auto",
|
||||
}
|
||||
)
|
||||
assert config["pip_installer"] == "auto"
|
||||
|
||||
config = validate_config(
|
||||
{
|
||||
"python_version": "3.11",
|
||||
"dependencies": ["."],
|
||||
"graphs": {"agent": "./agent.py:graph"},
|
||||
"pip_installer": "pip",
|
||||
}
|
||||
)
|
||||
assert config["pip_installer"] == "pip"
|
||||
|
||||
config = validate_config(
|
||||
{
|
||||
"python_version": "3.11",
|
||||
"dependencies": ["."],
|
||||
"graphs": {"agent": "./agent.py:graph"},
|
||||
"pip_installer": "uv",
|
||||
}
|
||||
)
|
||||
assert config["pip_installer"] == "uv"
|
||||
|
||||
# Missing pip_installer should default to "auto"
|
||||
config = validate_config(
|
||||
{
|
||||
"python_version": "3.11",
|
||||
"dependencies": ["."],
|
||||
"graphs": {"agent": "./agent.py:graph"},
|
||||
}
|
||||
)
|
||||
assert config["pip_installer"] == "auto"
|
||||
|
||||
# Invalid pip_installer values should raise error
|
||||
with pytest.raises(click.UsageError) as exc_info:
|
||||
validate_config(
|
||||
{
|
||||
"python_version": "3.11",
|
||||
"dependencies": ["."],
|
||||
"graphs": {"agent": "./agent.py:graph"},
|
||||
"pip_installer": "conda",
|
||||
}
|
||||
)
|
||||
assert "Invalid pip_installer: 'conda'" in str(exc_info.value)
|
||||
assert "Must be 'auto', 'pip', or 'uv'" in str(exc_info.value)
|
||||
|
||||
with pytest.raises(click.UsageError) as exc_info:
|
||||
validate_config(
|
||||
{
|
||||
"python_version": "3.11",
|
||||
"dependencies": ["."],
|
||||
"graphs": {"agent": "./agent.py:graph"},
|
||||
"pip_installer": "invalid",
|
||||
}
|
||||
)
|
||||
assert "Invalid pip_installer: 'invalid'" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_validate_config_file():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmpdir_path = pathlib.Path(tmpdir)
|
||||
@@ -345,7 +421,7 @@ def test_config_to_docker_simple():
|
||||
FROM langchain/langgraph-api:3.11
|
||||
# -- Installing local requirements --
|
||||
COPY --from=__outer_requirements.txt requirements.txt /deps/__outer_graphs_reqs_a/graphs_reqs_a/requirements.txt
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -r /deps/__outer_graphs_reqs_a/graphs_reqs_a/requirements.txt
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -r /deps/__outer_graphs_reqs_a/graphs_reqs_a/requirements.txt
|
||||
# -- End of local requirements install --
|
||||
# -- Adding local package ../../examples --
|
||||
COPY --from=examples . /deps/examples
|
||||
@@ -357,7 +433,10 @@ RUN set -ex && \\
|
||||
'name = "unit_tests"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency unit_tests --
|
||||
@@ -368,16 +447,19 @@ RUN set -ex && \\
|
||||
'name = "graphs_reqs_a"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_graphs_reqs_a/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency graphs_reqs_a --
|
||||
# -- Installing all local dependencies --
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
# -- End of local dependencies install --
|
||||
ENV LANGGRAPH_HTTP='{{"app": "/deps/examples/my_app.py:app"}}'
|
||||
ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}}'
|
||||
{PIP_CLEANUP_LINES}
|
||||
{FORMATTED_CLEANUP_LINES}
|
||||
WORKDIR /deps/__outer_unit_tests/unit_tests\
|
||||
"""
|
||||
assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin
|
||||
@@ -407,7 +489,10 @@ RUN set -ex && \\
|
||||
'name = "unit_tests"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency unit_tests --
|
||||
@@ -418,16 +503,19 @@ RUN set -ex && \\
|
||||
'name = "tests"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_tests/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency tests --
|
||||
# -- Installing all local dependencies --
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
# -- End of local dependencies install --
|
||||
ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}'
|
||||
"""
|
||||
+ PIP_CLEANUP_LINES
|
||||
+ FORMATTED_CLEANUP_LINES
|
||||
+ """
|
||||
WORKDIR /deps/__outer_unit_tests/unit_tests\
|
||||
"""
|
||||
@@ -462,16 +550,19 @@ RUN set -ex && \\
|
||||
'name = "unit_tests"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency unit_tests --
|
||||
# -- Installing all local dependencies --
|
||||
RUN PIP_CONFIG_FILE=/pipconfig.txt PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
RUN PIP_CONFIG_FILE=/pipconfig.txt PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
# -- End of local dependencies install --
|
||||
ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}'
|
||||
"""
|
||||
+ PIP_CLEANUP_LINES
|
||||
+ FORMATTED_CLEANUP_LINES
|
||||
+ """
|
||||
WORKDIR /deps/__outer_unit_tests/unit_tests\
|
||||
"""
|
||||
@@ -521,15 +612,18 @@ RUN set -ex && \\
|
||||
'name = "graphs"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_graphs/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency graphs --
|
||||
# -- Installing all local dependencies --
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
# -- End of local dependencies install --
|
||||
ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_graphs/src/agent.py:graph"}}'
|
||||
{PIP_CLEANUP_LINES}\
|
||||
{FORMATTED_CLEANUP_LINES}\
|
||||
"""
|
||||
assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin
|
||||
assert additional_contexts == {}
|
||||
@@ -562,11 +656,11 @@ dependencies = ["langchain"]"""
|
||||
ADD . /deps/unit_tests
|
||||
# -- End of local package . --
|
||||
# -- Installing all local dependencies --
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
# -- End of local dependencies install --
|
||||
ENV LANGSERVE_GRAPHS='{"agent": "/deps/unit_tests/graphs/agent.py:graph"}'
|
||||
"""
|
||||
+ PIP_CLEANUP_LINES
|
||||
+ FORMATTED_CLEANUP_LINES
|
||||
+ "\n"
|
||||
+ "WORKDIR /deps/unit_tests"
|
||||
""
|
||||
@@ -594,7 +688,7 @@ def test_config_to_docker_end_to_end():
|
||||
ARG meow
|
||||
ARG foo
|
||||
ADD pipconfig.txt /pipconfig.txt
|
||||
RUN PIP_CONFIG_FILE=/pipconfig.txt PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt langchain langchain_openai
|
||||
RUN PIP_CONFIG_FILE=/pipconfig.txt PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt langchain langchain_openai
|
||||
# -- Adding non-package dependency graphs --
|
||||
ADD ./graphs/ /deps/__outer_graphs/src
|
||||
RUN set -ex && \\
|
||||
@@ -602,15 +696,18 @@ RUN set -ex && \\
|
||||
'name = "graphs"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_graphs/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency graphs --
|
||||
# -- Installing all local dependencies --
|
||||
RUN PIP_CONFIG_FILE=/pipconfig.txt PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
RUN PIP_CONFIG_FILE=/pipconfig.txt PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
# -- End of local dependencies install --
|
||||
ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_graphs/src/agent.py:graph"}}'
|
||||
{PIP_CLEANUP_LINES}"""
|
||||
{FORMATTED_CLEANUP_LINES}"""
|
||||
assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin
|
||||
assert additional_contexts == {}
|
||||
|
||||
@@ -705,12 +802,15 @@ RUN set -ex && \\
|
||||
'name = "unit_tests"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency unit_tests --
|
||||
# -- Installing all local dependencies --
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
# -- End of local dependencies install --
|
||||
ENV LANGGRAPH_UI='{{"agent": "./graphs/agent.ui.jsx"}}'
|
||||
ENV LANGGRAPH_UI_CONFIG='{{"shared": ["nuqs"]}}'
|
||||
@@ -719,7 +819,7 @@ ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:g
|
||||
ENV NODE_VERSION=20
|
||||
RUN cd /deps/__outer_unit_tests/unit_tests && npm i && tsx /api/langgraph_api/js/build.mts
|
||||
# -- End of JS dependencies install --
|
||||
{PIP_CLEANUP_LINES}
|
||||
{FORMATTED_CLEANUP_LINES}
|
||||
WORKDIR /deps/__outer_unit_tests/unit_tests"""
|
||||
|
||||
assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin
|
||||
@@ -748,29 +848,87 @@ RUN set -ex && \\
|
||||
'name = "unit_tests"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency unit_tests --
|
||||
# -- Installing all local dependencies --
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
# -- End of local dependencies install --
|
||||
ENV LANGSERVE_GRAPHS='{{"python": "/deps/__outer_unit_tests/unit_tests/multiplatform/python.py:graph", "js": "/deps/__outer_unit_tests/unit_tests/multiplatform/js.mts:graph"}}'
|
||||
# -- Installing JS dependencies --
|
||||
ENV NODE_VERSION=22
|
||||
RUN cd /deps/__outer_unit_tests/unit_tests && npm i && tsx /api/langgraph_api/js/build.mts
|
||||
# -- End of JS dependencies install --
|
||||
{PIP_CLEANUP_LINES}
|
||||
{FORMATTED_CLEANUP_LINES}
|
||||
WORKDIR /deps/__outer_unit_tests/unit_tests"""
|
||||
|
||||
assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin
|
||||
assert additional_contexts == {}
|
||||
|
||||
|
||||
def test_config_to_docker_pip_installer():
|
||||
"""Test that pip_installer setting affects the generated Dockerfile."""
|
||||
graphs = {"agent": "./graphs/agent.py:graph"}
|
||||
base_config = {
|
||||
"python_version": "3.11",
|
||||
"dependencies": ["."],
|
||||
"graphs": graphs,
|
||||
}
|
||||
|
||||
# Test default (auto) behavior with UV-supporting image
|
||||
config_auto = validate_config(
|
||||
{**copy.deepcopy(base_config), "pip_installer": "auto"}
|
||||
)
|
||||
docker_auto, _ = config_to_docker(
|
||||
PATH_TO_CONFIG, config_auto, "langchain/langgraph-api:0.2.47"
|
||||
)
|
||||
assert "uv pip install --system" in docker_auto
|
||||
assert "rm /usr/bin/uv /usr/bin/uvx" in docker_auto
|
||||
|
||||
# Test explicit pip setting
|
||||
config_pip = validate_config({**copy.deepcopy(base_config), "pip_installer": "pip"})
|
||||
docker_pip, _ = config_to_docker(
|
||||
PATH_TO_CONFIG, config_pip, "langchain/langgraph-api:0.2.47"
|
||||
)
|
||||
assert "uv pip install --system" not in docker_pip
|
||||
assert "pip install" in docker_pip
|
||||
assert "rm /usr/bin/uv" not in docker_pip
|
||||
|
||||
# Test explicit uv setting
|
||||
config_uv = validate_config({**copy.deepcopy(base_config), "pip_installer": "uv"})
|
||||
docker_uv, _ = config_to_docker(
|
||||
PATH_TO_CONFIG, config_uv, "langchain/langgraph-api:0.2.47"
|
||||
)
|
||||
assert "uv pip install --system" in docker_uv
|
||||
assert "rm /usr/bin/uv /usr/bin/uvx" in docker_uv
|
||||
|
||||
# Test auto behavior with older image (should use pip)
|
||||
config_auto_old = validate_config(
|
||||
{**copy.deepcopy(base_config), "pip_installer": "auto"}
|
||||
)
|
||||
docker_auto_old, _ = config_to_docker(
|
||||
PATH_TO_CONFIG, config_auto_old, "langchain/langgraph-api:0.2.46"
|
||||
)
|
||||
assert "uv pip install --system" not in docker_auto_old
|
||||
assert "pip install" in docker_auto_old
|
||||
assert "rm /usr/bin/uv" not in docker_auto_old
|
||||
|
||||
# Test that missing pip_installer defaults to auto behavior
|
||||
config_default = validate_config(copy.deepcopy(base_config))
|
||||
docker_default, _ = config_to_docker(
|
||||
PATH_TO_CONFIG, config_default, "langchain/langgraph-api:0.2.47"
|
||||
)
|
||||
assert "uv pip install --system" in docker_default
|
||||
|
||||
|
||||
# config_to_compose
|
||||
def test_config_to_compose_simple_config():
|
||||
graphs = {"agent": "./agent.py:graph"}
|
||||
# Create a properly indented version of PIP_CLEANUP_LINES for compose files
|
||||
# Create a properly indented version of FORMATTED_CLEANUP_LINES for compose files
|
||||
expected_compose_stdin = f"""
|
||||
pull_policy: build
|
||||
build:
|
||||
@@ -784,15 +942,18 @@ def test_config_to_compose_simple_config():
|
||||
'name = "unit_tests"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency unit_tests --
|
||||
# -- Installing all local dependencies --
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
# -- End of local dependencies install --
|
||||
ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}}'
|
||||
{textwrap.indent(textwrap.dedent(PIP_CLEANUP_LINES), " ")}
|
||||
{textwrap.indent(textwrap.dedent(FORMATTED_CLEANUP_LINES), " ")}
|
||||
WORKDIR /deps/__outer_unit_tests/unit_tests
|
||||
"""
|
||||
actual_compose_stdin = config_to_compose(
|
||||
@@ -822,15 +983,18 @@ def test_config_to_compose_env_vars():
|
||||
'name = "unit_tests"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency unit_tests --
|
||||
# -- Installing all local dependencies --
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
# -- End of local dependencies install --
|
||||
ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}}'
|
||||
{textwrap.indent(textwrap.dedent(PIP_CLEANUP_LINES), " ")}
|
||||
{textwrap.indent(textwrap.dedent(FORMATTED_CLEANUP_LINES), " ")}
|
||||
WORKDIR /deps/__outer_unit_tests/unit_tests
|
||||
"""
|
||||
openai_api_key = "key"
|
||||
@@ -864,15 +1028,18 @@ def test_config_to_compose_env_file():
|
||||
'name = "unit_tests"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency unit_tests --
|
||||
# -- Installing all local dependencies --
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
# -- End of local dependencies install --
|
||||
ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}}'
|
||||
{textwrap.indent(textwrap.dedent(PIP_CLEANUP_LINES), " ")}
|
||||
{textwrap.indent(textwrap.dedent(FORMATTED_CLEANUP_LINES), " ")}
|
||||
WORKDIR /deps/__outer_unit_tests/unit_tests
|
||||
"""
|
||||
actual_compose_stdin = config_to_compose(
|
||||
@@ -899,15 +1066,18 @@ def test_config_to_compose_watch():
|
||||
'name = "unit_tests"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency unit_tests --
|
||||
# -- Installing all local dependencies --
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
# -- End of local dependencies install --
|
||||
ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}}'
|
||||
{textwrap.indent(textwrap.dedent(PIP_CLEANUP_LINES), " ")}
|
||||
{textwrap.indent(textwrap.dedent(FORMATTED_CLEANUP_LINES), " ")}
|
||||
WORKDIR /deps/__outer_unit_tests/unit_tests
|
||||
|
||||
develop:
|
||||
@@ -943,15 +1113,18 @@ def test_config_to_compose_end_to_end():
|
||||
'name = "unit_tests"' \\
|
||||
'version = "0.1"' \\
|
||||
'[tool.setuptools.package-data]' \\
|
||||
'"*" = ["**/*"]'; do \\
|
||||
'"*" = ["**/*"]' \\
|
||||
'[build-system]' \\
|
||||
'requires = ["setuptools>=61"]' \\
|
||||
'build-backend = "setuptools.build_meta"'; do \\
|
||||
echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
|
||||
done
|
||||
# -- End of non-package dependency unit_tests --
|
||||
# -- Installing all local dependencies --
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
RUN PYTHONDONTWRITEBYTECODE=1 uv pip install --system --no-cache-dir -c /api/constraints.txt -e /deps/*
|
||||
# -- End of local dependencies install --
|
||||
ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}}'
|
||||
{textwrap.indent(textwrap.dedent(PIP_CLEANUP_LINES), " ")}
|
||||
{textwrap.indent(textwrap.dedent(FORMATTED_CLEANUP_LINES), " ")}
|
||||
WORKDIR /deps/__outer_unit_tests/unit_tests
|
||||
|
||||
develop:
|
||||
|
||||
Generated
+1
-1
@@ -501,7 +501,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-cli"
|
||||
version = "0.2.12"
|
||||
version = "0.3.3"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
|
||||
@@ -12,7 +12,6 @@
|
||||
[](https://pepy.tech/project/langgraph)
|
||||
[](https://github.com/langchain-ai/langgraph/issues)
|
||||
[](https://langchain-ai.github.io/langgraph/)
|
||||
[](https://gitmcp.io/langchain-ai/langgraph)
|
||||
|
||||
Trusted by companies shaping the future of agents – including Klarna, Replit, Elastic, and more – LangGraph is a low-level orchestration framework for building, managing, and deploying long-running, stateful agents.
|
||||
|
||||
|
||||
@@ -1,14 +1,15 @@
|
||||
from langgraph.channels.any_value import AnyValue
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.channels.last_value import LastValue, LastValueAfterFinish
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.channels.topic import Topic
|
||||
from langgraph.channels.untracked_value import UntrackedValue
|
||||
|
||||
__all__ = [
|
||||
"LastValue",
|
||||
"LastValueAfterFinish",
|
||||
"Topic",
|
||||
"BinaryOperatorAggregate",
|
||||
"UntrackedValue",
|
||||
"EphemeralValue",
|
||||
"AnyValue",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
from collections.abc import Sequence
|
||||
from typing import Generic
|
||||
|
||||
from typing_extensions import Self
|
||||
|
||||
from langgraph.channels.base import BaseChannel, Value
|
||||
from langgraph.constants import MISSING
|
||||
from langgraph.errors import EmptyChannelError, InvalidUpdateError
|
||||
|
||||
|
||||
class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]):
|
||||
"""Stores the last value received, never checkpointed."""
|
||||
|
||||
__slots__ = ("value", "guard")
|
||||
|
||||
def __init__(self, typ: type[Value], guard: bool = True) -> None:
|
||||
super().__init__(typ)
|
||||
self.guard = guard
|
||||
self.value = MISSING
|
||||
|
||||
def __eq__(self, value: object) -> bool:
|
||||
return isinstance(value, UntrackedValue) and value.guard == self.guard
|
||||
|
||||
@property
|
||||
def ValueType(self) -> type[Value]:
|
||||
"""The type of the value stored in the channel."""
|
||||
return self.typ
|
||||
|
||||
@property
|
||||
def UpdateType(self) -> type[Value]:
|
||||
"""The type of the update received by the channel."""
|
||||
return self.typ
|
||||
|
||||
def copy(self) -> Self:
|
||||
"""Return a copy of the channel."""
|
||||
empty = self.__class__(self.typ, self.guard)
|
||||
empty.key = self.key
|
||||
empty.value = self.value
|
||||
return empty
|
||||
|
||||
def checkpoint(self) -> Value:
|
||||
return MISSING
|
||||
|
||||
def from_checkpoint(self, checkpoint: Value) -> Self:
|
||||
empty = self.__class__(self.typ, self.guard)
|
||||
empty.key = self.key
|
||||
return empty
|
||||
|
||||
def update(self, values: Sequence[Value]) -> bool:
|
||||
if len(values) == 0:
|
||||
return False
|
||||
if len(values) != 1 and self.guard:
|
||||
raise InvalidUpdateError(
|
||||
f"At key '{self.key}': UntrackedValue(guard=True) can receive only one value per step. Use guard=False if you want to store any one of multiple values."
|
||||
)
|
||||
|
||||
self.value = values[-1]
|
||||
return True
|
||||
|
||||
def get(self) -> Value:
|
||||
if self.value is MISSING:
|
||||
raise EmptyChannelError()
|
||||
return self.value
|
||||
|
||||
def is_available(self) -> bool:
|
||||
return self.value is not MISSING
|
||||
@@ -39,8 +39,6 @@ ERROR = sys.intern("__error__")
|
||||
# for errors raised by nodes
|
||||
NO_WRITES = sys.intern("__no_writes__")
|
||||
# marker to signal node didn't write anything
|
||||
SCHEDULED = sys.intern("__scheduled__")
|
||||
# marker to signal node was scheduled (in distributed mode)
|
||||
TASKS = sys.intern("__pregel_tasks")
|
||||
# for Send objects returned by nodes/edges, corresponds to PUSH below
|
||||
RETURN = sys.intern("__return__")
|
||||
@@ -71,13 +69,6 @@ CONFIG_KEY_RESUMING = sys.intern("__pregel_resuming")
|
||||
# holds a boolean indicating if subgraphs should resume from a previous checkpoint
|
||||
CONFIG_KEY_TASK_ID = sys.intern("__pregel_task_id")
|
||||
# holds the task ID for the current task
|
||||
CONFIG_KEY_DEDUPE_TASKS = sys.intern("__pregel_dedupe_tasks")
|
||||
# holds a boolean indicating if tasks should be deduplicated (for distributed mode)
|
||||
CONFIG_KEY_ENSURE_LATEST = sys.intern("__pregel_ensure_latest")
|
||||
# holds a boolean indicating whether to assert the requested checkpoint is the latest
|
||||
# (for distributed mode)
|
||||
CONFIG_KEY_DELEGATE = sys.intern("__pregel_delegate")
|
||||
# holds a boolean indicating whether to delegate subgraphs (for distributed mode)
|
||||
CONFIG_KEY_THREAD_ID = sys.intern("thread_id")
|
||||
# holds the thread ID for the current invocation
|
||||
CONFIG_KEY_CHECKPOINT_MAP = sys.intern("checkpoint_map")
|
||||
@@ -121,7 +112,6 @@ RESERVED = {
|
||||
RESUME,
|
||||
ERROR,
|
||||
NO_WRITES,
|
||||
SCHEDULED,
|
||||
# reserved config.configurable keys
|
||||
CONFIG_KEY_SEND,
|
||||
CONFIG_KEY_READ,
|
||||
@@ -132,9 +122,6 @@ RESERVED = {
|
||||
CONFIG_KEY_CHECKPOINT_MAP,
|
||||
CONFIG_KEY_RESUMING,
|
||||
CONFIG_KEY_TASK_ID,
|
||||
CONFIG_KEY_DEDUPE_TASKS,
|
||||
CONFIG_KEY_ENSURE_LATEST,
|
||||
CONFIG_KEY_DELEGATE,
|
||||
CONFIG_KEY_CHECKPOINT_MAP,
|
||||
CONFIG_KEY_CHECKPOINT_ID,
|
||||
CONFIG_KEY_CHECKPOINT_NS,
|
||||
|
||||
@@ -78,13 +78,6 @@ class NodeInterrupt(GraphInterrupt):
|
||||
super().__init__([Interrupt(value=value)])
|
||||
|
||||
|
||||
class GraphDelegate(GraphBubbleUp):
|
||||
"""Raised when a graph is delegated (for distributed mode)."""
|
||||
|
||||
def __init__(self, *args: dict[str, Any]) -> None:
|
||||
super().__init__(*args)
|
||||
|
||||
|
||||
class ParentCommand(GraphBubbleUp):
|
||||
args: tuple[Command]
|
||||
|
||||
@@ -102,9 +95,3 @@ class TaskNotFound(Exception):
|
||||
"""Raised when the executor is unable to find a task (for distributed mode)."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class CheckpointNotLatest(Exception):
|
||||
"""Raised when the checkpoint is not the latest version (for distributed mode)."""
|
||||
|
||||
pass
|
||||
|
||||
@@ -38,7 +38,7 @@ from langgraph.pregel.read import PregelNode
|
||||
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
|
||||
from langgraph.store.base import BaseStore
|
||||
from langgraph.types import _DC_KWARGS, CachePolicy, RetryPolicy, StreamMode
|
||||
from langgraph.warnings import LangGraphDeprecatedSinceV10
|
||||
from langgraph.warnings import LangGraphDeprecatedSinceV05
|
||||
|
||||
|
||||
class TaskFunction(Generic[P, T]):
|
||||
@@ -179,7 +179,7 @@ def task(
|
||||
if (retry := kwargs.get("retry", UNSET)) is not UNSET:
|
||||
warnings.warn(
|
||||
"`retry` is deprecated and will be removed. Please use `retry_policy` instead.",
|
||||
category=LangGraphDeprecatedSinceV10,
|
||||
category=LangGraphDeprecatedSinceV05,
|
||||
)
|
||||
if retry_policy is None:
|
||||
retry_policy = retry # type: ignore[assignment]
|
||||
@@ -383,7 +383,7 @@ class entrypoint:
|
||||
if (retry := kwargs.get("retry", UNSET)) is not UNSET:
|
||||
warnings.warn(
|
||||
"`retry` is deprecated and will be removed. Please use `retry_policy` instead.",
|
||||
category=LangGraphDeprecatedSinceV10,
|
||||
category=LangGraphDeprecatedSinceV05,
|
||||
)
|
||||
if retry_policy is None:
|
||||
retry_policy = retry # type: ignore[assignment]
|
||||
@@ -499,7 +499,7 @@ class entrypoint:
|
||||
func.__name__: PregelNode(
|
||||
bound=bound,
|
||||
triggers=[START],
|
||||
channels=[START],
|
||||
channels=START,
|
||||
writers=[
|
||||
ChannelWrite(
|
||||
[
|
||||
|
||||
@@ -1,11 +1,12 @@
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph.message import MessagesState, add_messages
|
||||
from langgraph.graph.message import MessageGraph, MessagesState, add_messages
|
||||
from langgraph.graph.state import StateGraph
|
||||
|
||||
__all__ = [
|
||||
"END",
|
||||
"START",
|
||||
"StateGraph",
|
||||
"MessageGraph",
|
||||
"add_messages",
|
||||
"MessagesState",
|
||||
]
|
||||
|
||||
@@ -25,6 +25,7 @@ from langchain_core.messages import (
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.constants import CONF, CONFIG_KEY_SEND
|
||||
from langgraph.graph.state import StateGraph
|
||||
|
||||
Messages = Union[list[MessageLikeRepresentation], MessageLikeRepresentation]
|
||||
|
||||
@@ -226,6 +227,57 @@ def add_messages(
|
||||
return merged
|
||||
|
||||
|
||||
class MessageGraph(StateGraph):
|
||||
"""A StateGraph where every node receives a list of messages as input and returns one or more messages as output.
|
||||
|
||||
MessageGraph is a subclass of StateGraph whose entire state is a single, append-only* list of messages.
|
||||
Each node in a MessageGraph takes a list of messages as input and returns zero or more
|
||||
messages as output. The `add_messages` function is used to merge the output messages from each node
|
||||
into the existing list of messages in the graph's state.
|
||||
|
||||
Examples:
|
||||
```pycon
|
||||
>>> from langgraph.graph.message import MessageGraph
|
||||
...
|
||||
>>> builder = MessageGraph()
|
||||
>>> builder.add_node("chatbot", lambda state: [("assistant", "Hello!")])
|
||||
>>> builder.set_entry_point("chatbot")
|
||||
>>> builder.set_finish_point("chatbot")
|
||||
>>> builder.compile().invoke([("user", "Hi there.")])
|
||||
[HumanMessage(content="Hi there.", id='...'), AIMessage(content="Hello!", id='...')]
|
||||
```
|
||||
|
||||
```pycon
|
||||
>>> from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
||||
>>> from langgraph.graph.message import MessageGraph
|
||||
...
|
||||
>>> builder = MessageGraph()
|
||||
>>> builder.add_node(
|
||||
... "chatbot",
|
||||
... lambda state: [
|
||||
... AIMessage(
|
||||
... content="Hello!",
|
||||
... tool_calls=[{"name": "search", "id": "123", "args": {"query": "X"}}],
|
||||
... )
|
||||
... ],
|
||||
... )
|
||||
>>> builder.add_node(
|
||||
... "search", lambda state: [ToolMessage(content="Searching...", tool_call_id="123")]
|
||||
... )
|
||||
>>> builder.set_entry_point("chatbot")
|
||||
>>> builder.add_edge("chatbot", "search")
|
||||
>>> builder.set_finish_point("search")
|
||||
>>> builder.compile().invoke([HumanMessage(content="Hi there. Can you search for X?")])
|
||||
{'messages': [HumanMessage(content="Hi there. Can you search for X?", id='b8b7d8f4-7f4d-4f4d-9c1d-f8b8d8f4d9c1'),
|
||||
AIMessage(content="Hello!", id='f4d9c1d8-8d8f-4d9c-b8b7-d8f4f4d9c1d8'),
|
||||
ToolMessage(content="Searching...", id='d8f4f4d9-c1d8-4f4d-b8b7-d8f4f4d9c1d8', tool_call_id="123")]}
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(Annotated[list[AnyMessage], add_messages]) # type: ignore[arg-type]
|
||||
|
||||
|
||||
class MessagesState(TypedDict):
|
||||
messages: Annotated[list[AnyMessage], add_messages]
|
||||
|
||||
|
||||
@@ -86,7 +86,7 @@ from langgraph.utils.fields import (
|
||||
)
|
||||
from langgraph.utils.pydantic import create_model
|
||||
from langgraph.utils.runnable import coerce_to_runnable
|
||||
from langgraph.warnings import LangGraphDeprecatedSinceV10
|
||||
from langgraph.warnings import LangGraphDeprecatedSinceV05
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -160,6 +160,7 @@ StateNode: TypeAlias = Union[
|
||||
_NodeWithConfigWriter[StateT_contra],
|
||||
_NodeWithConfigStore[StateT_contra],
|
||||
_NodeWithConfigWriterStore[StateT_contra],
|
||||
Runnable[StateT_contra, Any],
|
||||
]
|
||||
|
||||
|
||||
@@ -261,7 +262,7 @@ class StateGraph(Generic[StateT, InputT, OutputT]):
|
||||
if (input_ := kwargs.get("input", UNSET)) is not UNSET:
|
||||
warnings.warn(
|
||||
"`input` is deprecated and will be removed. Please use `input_schema` instead.",
|
||||
category=LangGraphDeprecatedSinceV10,
|
||||
category=LangGraphDeprecatedSinceV05,
|
||||
stacklevel=2,
|
||||
)
|
||||
if input_schema is None:
|
||||
@@ -270,7 +271,7 @@ class StateGraph(Generic[StateT, InputT, OutputT]):
|
||||
if (output := kwargs.get("output", UNSET)) is not UNSET:
|
||||
warnings.warn(
|
||||
"`output` is deprecated and will be removed. Please use `output_schema` instead.",
|
||||
category=LangGraphDeprecatedSinceV10,
|
||||
category=LangGraphDeprecatedSinceV05,
|
||||
stacklevel=2,
|
||||
)
|
||||
if output_schema is None:
|
||||
@@ -436,7 +437,7 @@ class StateGraph(Generic[StateT, InputT, OutputT]):
|
||||
if (retry := kwargs.get("retry", UNSET)) is not UNSET:
|
||||
warnings.warn(
|
||||
"`retry` is deprecated and will be removed. Please use `retry_policy` instead.",
|
||||
category=LangGraphDeprecatedSinceV10,
|
||||
category=LangGraphDeprecatedSinceV05,
|
||||
)
|
||||
if retry_policy is None:
|
||||
retry_policy = retry # type: ignore[assignment]
|
||||
@@ -444,7 +445,7 @@ class StateGraph(Generic[StateT, InputT, OutputT]):
|
||||
if (input_ := kwargs.get("input", UNSET)) is not UNSET:
|
||||
warnings.warn(
|
||||
"`input` is deprecated and will be removed. Please use `input_schema` instead.",
|
||||
category=LangGraphDeprecatedSinceV10,
|
||||
category=LangGraphDeprecatedSinceV05,
|
||||
)
|
||||
if input_schema is None:
|
||||
input_schema = cast(Union[type[InputT], None], input_)
|
||||
@@ -535,7 +536,7 @@ class StateGraph(Generic[StateT, InputT, OutputT]):
|
||||
if input_schema is not None:
|
||||
self._add_schema(input_schema)
|
||||
self.nodes[node] = StateNodeSpec(
|
||||
coerce_to_runnable(action, name=node, trace=False), # type: ignore
|
||||
coerce_to_runnable(action, name=node, trace=False),
|
||||
metadata,
|
||||
input=input_schema or self.state_schema,
|
||||
retry_policy=retry_policy,
|
||||
@@ -849,13 +850,6 @@ class StateGraph(Generic[StateT, InputT, OutputT]):
|
||||
builder=self,
|
||||
schema_to_mapper={},
|
||||
config_type=self.config_schema,
|
||||
input_model=(
|
||||
self.input_schema
|
||||
if len(self.channels) > 1
|
||||
and isclass(self.input_schema)
|
||||
and issubclass(self.input_schema, BaseModel)
|
||||
else None
|
||||
),
|
||||
nodes={},
|
||||
channels={
|
||||
**self.channels,
|
||||
@@ -996,20 +990,17 @@ class CompiledStateGraph(
|
||||
self.nodes[key] = PregelNode(
|
||||
tags=[TAG_HIDDEN],
|
||||
triggers=[START],
|
||||
channels=[START],
|
||||
channels=START,
|
||||
writers=[ChannelWrite(write_entries)],
|
||||
)
|
||||
elif node is not None:
|
||||
input_schema = node.input if node else self.builder._state_schema
|
||||
input_values = {k: k for k in self.builder.schemas[input_schema]}
|
||||
is_single_input = len(input_values) == 1 and "__root__" in input_values
|
||||
input_channels = list(self.builder.schemas[input_schema])
|
||||
is_single_input = len(input_channels) == 1 and "__root__" in input_channels
|
||||
if input_schema in self.schema_to_mapper:
|
||||
mapper = self.schema_to_mapper[input_schema]
|
||||
else:
|
||||
mapper = _pick_mapper(
|
||||
list(input_values),
|
||||
input_schema,
|
||||
)
|
||||
mapper = _pick_mapper(input_channels, input_schema)
|
||||
self.schema_to_mapper[input_schema] = mapper
|
||||
|
||||
branch_channel = CHANNEL_BRANCH_TO.format(key)
|
||||
@@ -1021,7 +1012,7 @@ class CompiledStateGraph(
|
||||
self.nodes[key] = PregelNode(
|
||||
triggers=[branch_channel],
|
||||
# read state keys and managed values
|
||||
channels=(list(input_values) if is_single_input else input_values),
|
||||
channels=("__root__" if is_single_input else input_channels),
|
||||
# coerce state dict to schema class (eg. pydantic model)
|
||||
mapper=mapper,
|
||||
# publish to state keys
|
||||
@@ -1111,6 +1102,7 @@ class CompiledStateGraph(
|
||||
|
||||
def _migrate_checkpoint(self, checkpoint: Checkpoint) -> None:
|
||||
"""Migrate a checkpoint to new channel layout."""
|
||||
super()._migrate_checkpoint(checkpoint)
|
||||
|
||||
values = checkpoint["channel_values"]
|
||||
versions = checkpoint["channel_versions"]
|
||||
|
||||
@@ -32,7 +32,6 @@ from langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
Checkpoint,
|
||||
CheckpointTuple,
|
||||
copy_checkpoint,
|
||||
)
|
||||
from langgraph.config import get_config
|
||||
from langgraph.constants import (
|
||||
@@ -60,7 +59,6 @@ from langgraph.constants import (
|
||||
NS_SEP,
|
||||
NULL_TASK_ID,
|
||||
PUSH,
|
||||
SCHEDULED,
|
||||
TASKS,
|
||||
)
|
||||
from langgraph.errors import (
|
||||
@@ -80,6 +78,7 @@ from langgraph.pregel.algo import (
|
||||
from langgraph.pregel.call import identifier
|
||||
from langgraph.pregel.checkpoint import (
|
||||
channels_from_checkpoint,
|
||||
copy_checkpoint,
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
)
|
||||
@@ -145,7 +144,7 @@ class NodeBuilder:
|
||||
"_cache_policy",
|
||||
)
|
||||
|
||||
_channels: list[str] | dict[str, str]
|
||||
_channels: str | list[str]
|
||||
_triggers: list[str]
|
||||
_tags: list[str]
|
||||
_metadata: dict[str, Any]
|
||||
@@ -157,7 +156,7 @@ class NodeBuilder:
|
||||
def __init__(
|
||||
self,
|
||||
) -> None:
|
||||
self._channels = {}
|
||||
self._channels = []
|
||||
self._triggers = []
|
||||
self._tags = []
|
||||
self._metadata = {}
|
||||
@@ -171,10 +170,8 @@ class NodeBuilder:
|
||||
channel: str,
|
||||
) -> Self:
|
||||
"""Subscribe to a single channel."""
|
||||
if isinstance(self._channels, list):
|
||||
self._channels.append(channel)
|
||||
elif not self._channels:
|
||||
self._channels = [channel]
|
||||
if not self._channels:
|
||||
self._channels = channel
|
||||
else:
|
||||
raise ValueError(
|
||||
"Cannot subscribe to single channels when other channels are already subscribed to"
|
||||
@@ -200,15 +197,15 @@ class NodeBuilder:
|
||||
Returns:
|
||||
Self for chaining
|
||||
"""
|
||||
if isinstance(self._channels, list):
|
||||
if isinstance(self._channels, str):
|
||||
raise ValueError(
|
||||
"Cannot subscribe to channels when subscribed to a single channel"
|
||||
)
|
||||
if read:
|
||||
if not self._channels:
|
||||
self._channels = {chan: chan for chan in channels}
|
||||
self._channels = list(channels)
|
||||
else:
|
||||
self._channels.update({chan: chan for chan in channels})
|
||||
self._channels.extend(channels)
|
||||
|
||||
if isinstance(channels, str):
|
||||
self._triggers.append(channels)
|
||||
@@ -222,11 +219,10 @@ class NodeBuilder:
|
||||
*channels: str,
|
||||
) -> Self:
|
||||
"""Adds the specified channels to read from, without subscribing to them."""
|
||||
assert self._channels, "Channels must be specified first"
|
||||
assert isinstance(self._channels, dict), (
|
||||
assert isinstance(self._channels, list), (
|
||||
"Cannot read additional channels when subscribed to single channels"
|
||||
)
|
||||
self._channels.update({c: c for c in channels})
|
||||
self._channels.extend(channels)
|
||||
return self
|
||||
|
||||
def do(
|
||||
@@ -593,8 +589,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
|
||||
config_type: type[Any] | None = None
|
||||
|
||||
input_model: type[BaseModel] | None = None
|
||||
|
||||
config: RunnableConfig | None = None
|
||||
|
||||
name: str = "LangGraph"
|
||||
@@ -622,7 +616,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
retry_policy: RetryPolicy | Sequence[RetryPolicy] = (),
|
||||
cache_policy: CachePolicy | None = None,
|
||||
config_type: type[Any] | None = None,
|
||||
input_model: type[BaseModel] | None = None,
|
||||
config: RunnableConfig | None = None,
|
||||
trigger_to_nodes: Mapping[str, Sequence[str]] | None = None,
|
||||
name: str = "LangGraph",
|
||||
@@ -654,7 +647,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
)
|
||||
self.cache_policy = cache_policy
|
||||
self.config_type = config_type
|
||||
self.input_model = input_model
|
||||
self.config = config
|
||||
self.trigger_to_nodes = trigger_to_nodes or {}
|
||||
self.name = name
|
||||
@@ -753,6 +745,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
validate_graph(
|
||||
self.nodes,
|
||||
{k: v for k, v in self.channels.items() if isinstance(v, BaseChannel)},
|
||||
{k: v for k, v in self.channels.items() if not isinstance(v, BaseChannel)},
|
||||
self.input_channels,
|
||||
self.output_channels,
|
||||
self.stream_channels,
|
||||
@@ -791,8 +784,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
return channel.UpdateType
|
||||
|
||||
def get_input_schema(self, config: RunnableConfig | None = None) -> type[BaseModel]:
|
||||
if self.input_model is not None:
|
||||
return self.input_model
|
||||
config = merge_configs(self.config, config)
|
||||
if isinstance(self.input_channels, str):
|
||||
return super().get_input_schema(config)
|
||||
@@ -917,7 +908,12 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
|
||||
def _migrate_checkpoint(self, checkpoint: Checkpoint) -> None:
|
||||
"""Migrate a saved checkpoint to new channel layout."""
|
||||
pass
|
||||
if checkpoint["v"] < 4 and checkpoint.get("pending_sends"):
|
||||
pending_sends: list[Send] = checkpoint.pop("pending_sends")
|
||||
checkpoint["channel_values"][TASKS] = pending_sends
|
||||
checkpoint["channel_versions"][TASKS] = max(
|
||||
checkpoint["channel_versions"].values()
|
||||
)
|
||||
|
||||
def _prepare_state_snapshot(
|
||||
self,
|
||||
@@ -1011,7 +1007,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
)
|
||||
if apply_pending_writes and saved.pending_writes:
|
||||
for tid, k, v in saved.pending_writes:
|
||||
if k in (ERROR, INTERRUPT, SCHEDULED):
|
||||
if k in (ERROR, INTERRUPT):
|
||||
continue
|
||||
if tid not in next_tasks:
|
||||
continue
|
||||
@@ -1130,7 +1126,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
)
|
||||
if apply_pending_writes and saved.pending_writes:
|
||||
for tid, k, v in saved.pending_writes:
|
||||
if k in (ERROR, INTERRUPT, SCHEDULED):
|
||||
if k in (ERROR, INTERRUPT):
|
||||
continue
|
||||
if tid not in next_tasks:
|
||||
continue
|
||||
@@ -1469,7 +1465,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
)
|
||||
# apply writes from tasks that already ran
|
||||
for tid, k, v in saved.pending_writes or []:
|
||||
if k in (ERROR, INTERRUPT, SCHEDULED):
|
||||
if k in (ERROR, INTERRUPT):
|
||||
continue
|
||||
if tid not in next_tasks:
|
||||
continue
|
||||
@@ -1633,7 +1629,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
)
|
||||
# apply writes
|
||||
for tid, k, v in saved.pending_writes:
|
||||
if k in (ERROR, INTERRUPT, SCHEDULED):
|
||||
if k in (ERROR, INTERRUPT):
|
||||
continue
|
||||
if tid not in next_tasks:
|
||||
continue
|
||||
@@ -1889,7 +1885,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
)
|
||||
# apply writes from tasks that already ran
|
||||
for tid, k, v in saved.pending_writes or []:
|
||||
if k in (ERROR, INTERRUPT, SCHEDULED):
|
||||
if k in (ERROR, INTERRUPT):
|
||||
continue
|
||||
if tid not in next_tasks:
|
||||
continue
|
||||
@@ -2052,7 +2048,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
self.trigger_to_nodes,
|
||||
)
|
||||
for tid, k, v in saved.pending_writes:
|
||||
if k in (ERROR, INTERRUPT, SCHEDULED):
|
||||
if k in (ERROR, INTERRUPT):
|
||||
continue
|
||||
if tid not in next_tasks:
|
||||
continue
|
||||
@@ -2307,7 +2303,8 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
- `"custom"`: Emit custom data from inside nodes or tasks using `StreamWriter`.
|
||||
- `"messages"`: Emit LLM messages token-by-token together with metadata for any LLM invocations inside nodes or tasks.
|
||||
Will be emitted as 2-tuples `(LLM token, metadata)`.
|
||||
- `"debug"`: Emit debug events with as much information as possible for each step.
|
||||
- `"checkpoints"`: Emit an event when a checkpoint is created, in the same format as returned by get_state().
|
||||
- `"tasks"`: Emit events when tasks start and finish, including their results and errors.
|
||||
|
||||
You can pass a list as the `stream_mode` parameter to stream multiple modes at once.
|
||||
The streamed outputs will be tuples of `(mode, data)`.
|
||||
@@ -2407,7 +2404,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
config[CONF][CONFIG_KEY_CHECKPOINT_DURING] = checkpoint_during
|
||||
with SyncPregelLoop(
|
||||
input,
|
||||
input_model=self.input_model,
|
||||
stream=StreamProtocol(stream.put, stream_modes),
|
||||
config=config,
|
||||
store=store,
|
||||
@@ -2416,6 +2412,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
nodes=self.nodes,
|
||||
specs=self.channels,
|
||||
output_keys=output_keys,
|
||||
input_keys=self.input_channels,
|
||||
stream_keys=self.stream_channels_asis,
|
||||
interrupt_before=interrupt_before_,
|
||||
interrupt_after=interrupt_after_,
|
||||
@@ -2470,7 +2467,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
# 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.
|
||||
while loop.tick(input_keys=self.input_channels):
|
||||
while loop.tick():
|
||||
for task in loop.match_cached_writes():
|
||||
loop.output_writes(task.id, task.writes, cached=True)
|
||||
for _ in runner.tick(
|
||||
@@ -2481,6 +2478,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
):
|
||||
# emit output
|
||||
yield from output()
|
||||
loop.after_tick()
|
||||
# emit output
|
||||
yield from output()
|
||||
# handle exit
|
||||
@@ -2650,7 +2648,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
config[CONF][CONFIG_KEY_CHECKPOINT_DURING] = checkpoint_during
|
||||
async with AsyncPregelLoop(
|
||||
input,
|
||||
input_model=self.input_model,
|
||||
stream=StreamProtocol(stream.put_nowait, stream_modes),
|
||||
config=config,
|
||||
store=store,
|
||||
@@ -2659,6 +2656,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
nodes=self.nodes,
|
||||
specs=self.channels,
|
||||
output_keys=output_keys,
|
||||
input_keys=self.input_channels,
|
||||
stream_keys=self.stream_channels_asis,
|
||||
interrupt_before=interrupt_before_,
|
||||
interrupt_after=interrupt_after_,
|
||||
@@ -2704,7 +2702,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
# 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
|
||||
while loop.tick(input_keys=self.input_channels):
|
||||
while loop.tick():
|
||||
for task in await loop.amatch_cached_writes():
|
||||
loop.output_writes(task.id, task.writes, cached=True)
|
||||
async for _ in runner.atick(
|
||||
@@ -2716,6 +2714,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
|
||||
# emit output
|
||||
for o in output():
|
||||
yield o
|
||||
loop.after_tick()
|
||||
# emit output
|
||||
for o in output():
|
||||
yield o
|
||||
|
||||
@@ -83,7 +83,7 @@ from langgraph.types import (
|
||||
)
|
||||
from langgraph.utils.config import merge_configs, patch_config
|
||||
|
||||
GetNextVersion = Callable[[Optional[V]], V]
|
||||
GetNextVersion = Callable[[Optional[V], None], V]
|
||||
SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11)
|
||||
|
||||
|
||||
@@ -214,7 +214,7 @@ def local_read(
|
||||
return values
|
||||
|
||||
|
||||
def increment(current: int | None) -> int:
|
||||
def increment(current: int | None, channel: None) -> int:
|
||||
"""Default channel versioning function, increments the current int version."""
|
||||
return current + 1 if current is not None else 1
|
||||
|
||||
@@ -265,7 +265,8 @@ def apply_writes(
|
||||
next_version = get_next_version(
|
||||
max(checkpoint["channel_versions"].values())
|
||||
if checkpoint["channel_versions"]
|
||||
else None
|
||||
else None,
|
||||
None,
|
||||
)
|
||||
|
||||
# Consume all channels that were read
|
||||
@@ -922,18 +923,18 @@ def _triggers(
|
||||
seen: ChannelVersions | None,
|
||||
null_version: V,
|
||||
proc: PregelNode,
|
||||
) -> Sequence[str]:
|
||||
) -> bool:
|
||||
if seen is None:
|
||||
for chan in proc.triggers:
|
||||
if channels[chan].is_available():
|
||||
return (chan,)
|
||||
return True
|
||||
else:
|
||||
for chan in proc.triggers:
|
||||
if channels[chan].is_available() and versions.get( # type: ignore[operator]
|
||||
chan, null_version
|
||||
) > seen.get(chan, null_version):
|
||||
return (chan,)
|
||||
return EMPTY_SEQ
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _scratchpad(
|
||||
@@ -1019,23 +1020,20 @@ def _proc_input(
|
||||
return copy(input_cache[proc.input_cache_key])
|
||||
# If all trigger channels subscribed by this process are not empty
|
||||
# then invoke the process with the values of all non-empty channels
|
||||
if isinstance(proc.channels, dict):
|
||||
if isinstance(proc.channels, list):
|
||||
val: dict[str, Any] = {}
|
||||
for k, chan in proc.channels.items():
|
||||
if chan in channels:
|
||||
if channels[chan].is_available():
|
||||
val[k] = channels[chan].get()
|
||||
else:
|
||||
val[k] = managed[k].get(scratchpad)
|
||||
elif isinstance(proc.channels, list):
|
||||
for chan in proc.channels:
|
||||
if chan in channels:
|
||||
if channels[chan].is_available():
|
||||
val = channels[chan].get()
|
||||
break
|
||||
val[chan] = channels[chan].get()
|
||||
else:
|
||||
val = managed[chan].get(scratchpad)
|
||||
break
|
||||
val[chan] = managed[chan].get(scratchpad)
|
||||
elif isinstance(proc.channels, str):
|
||||
if proc.channels in channels:
|
||||
if channels[proc.channels].is_available():
|
||||
val = channels[proc.channels].get()
|
||||
else:
|
||||
return MISSING
|
||||
else:
|
||||
return MISSING
|
||||
else:
|
||||
|
||||
@@ -71,3 +71,14 @@ def channels_from_checkpoint(
|
||||
},
|
||||
managed_specs,
|
||||
)
|
||||
|
||||
|
||||
def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
|
||||
return Checkpoint(
|
||||
v=checkpoint["v"],
|
||||
ts=checkpoint["ts"],
|
||||
id=checkpoint["id"],
|
||||
channel_values=checkpoint["channel_values"].copy(),
|
||||
channel_versions=checkpoint["channel_versions"].copy(),
|
||||
versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()},
|
||||
)
|
||||
|
||||
@@ -3,13 +3,8 @@ from __future__ import annotations
|
||||
from collections import defaultdict
|
||||
from collections.abc import Iterable, Iterator, Mapping, Sequence
|
||||
from dataclasses import asdict
|
||||
from datetime import datetime, timezone
|
||||
from pprint import pformat
|
||||
from typing import (
|
||||
Any,
|
||||
Literal,
|
||||
Union,
|
||||
)
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from langchain_core.runnables.config import RunnableConfig
|
||||
@@ -17,7 +12,7 @@ from langchain_core.utils.input import get_bolded_text, get_colored_text
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite
|
||||
from langgraph.checkpoint.base import CheckpointMetadata, PendingWrite
|
||||
from langgraph.constants import (
|
||||
CONF,
|
||||
CONFIG_KEY_CHECKPOINT_NS,
|
||||
@@ -66,82 +61,43 @@ class CheckpointPayload(TypedDict):
|
||||
tasks: list[CheckpointTask]
|
||||
|
||||
|
||||
class DebugOutputBase(TypedDict):
|
||||
timestamp: str
|
||||
step: int
|
||||
|
||||
|
||||
class DebugOutputTask(DebugOutputBase):
|
||||
type: Literal["task"]
|
||||
payload: TaskPayload
|
||||
|
||||
|
||||
class DebugOutputTaskResult(DebugOutputBase):
|
||||
type: Literal["task_result"]
|
||||
payload: TaskResultPayload
|
||||
|
||||
|
||||
class DebugOutputCheckpoint(DebugOutputBase):
|
||||
type: Literal["checkpoint"]
|
||||
payload: CheckpointPayload
|
||||
|
||||
|
||||
DebugOutput = Union[DebugOutputTask, DebugOutputTaskResult, DebugOutputCheckpoint]
|
||||
|
||||
|
||||
TASK_NAMESPACE = UUID("6ba7b831-9dad-11d1-80b4-00c04fd430c8")
|
||||
|
||||
|
||||
def map_debug_tasks(
|
||||
step: int, tasks: Iterable[PregelExecutableTask]
|
||||
) -> Iterator[DebugOutputTask]:
|
||||
def map_debug_tasks(tasks: Iterable[PregelExecutableTask]) -> Iterator[TaskPayload]:
|
||||
"""Produce "task" events for stream_mode=debug."""
|
||||
ts = datetime.now(timezone.utc).isoformat()
|
||||
for task in tasks:
|
||||
if task.config is not None and TAG_HIDDEN in task.config.get("tags", []):
|
||||
continue
|
||||
|
||||
yield {
|
||||
"type": "task",
|
||||
"timestamp": ts,
|
||||
"step": step,
|
||||
"payload": {
|
||||
"id": task.id,
|
||||
"name": task.name,
|
||||
"input": task.input,
|
||||
"triggers": task.triggers,
|
||||
},
|
||||
"id": task.id,
|
||||
"name": task.name,
|
||||
"input": task.input,
|
||||
"triggers": task.triggers,
|
||||
}
|
||||
|
||||
|
||||
def map_debug_task_results(
|
||||
step: int,
|
||||
task_tup: tuple[PregelExecutableTask, Sequence[tuple[str, Any]]],
|
||||
stream_keys: str | Sequence[str],
|
||||
) -> Iterator[DebugOutputTaskResult]:
|
||||
) -> Iterator[TaskResultPayload]:
|
||||
"""Produce "task_result" events for stream_mode=debug."""
|
||||
stream_channels_list = (
|
||||
[stream_keys] if isinstance(stream_keys, str) else stream_keys
|
||||
)
|
||||
task, writes = task_tup
|
||||
yield {
|
||||
"type": "task_result",
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"step": step,
|
||||
"payload": {
|
||||
"id": task.id,
|
||||
"name": task.name,
|
||||
"error": next((w[1] for w in writes if w[0] == ERROR), None),
|
||||
"result": [
|
||||
w for w in writes if w[0] in stream_channels_list or w[0] == RETURN
|
||||
],
|
||||
"interrupts": [
|
||||
asdict(v)
|
||||
for w in writes
|
||||
if w[0] == INTERRUPT
|
||||
for v in (w[1] if isinstance(w[1], Sequence) else [w[1]])
|
||||
],
|
||||
},
|
||||
"id": task.id,
|
||||
"name": task.name,
|
||||
"error": next((w[1] for w in writes if w[0] == ERROR), None),
|
||||
"result": [w for w in writes if w[0] in stream_channels_list or w[0] == RETURN],
|
||||
"interrupts": [
|
||||
asdict(v)
|
||||
for w in writes
|
||||
if w[0] == INTERRUPT
|
||||
for v in (w[1] if isinstance(w[1], Sequence) else [w[1]])
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@@ -159,17 +115,15 @@ def rm_pregel_keys(config: RunnableConfig | None) -> RunnableConfig | None:
|
||||
|
||||
|
||||
def map_debug_checkpoint(
|
||||
step: int,
|
||||
config: RunnableConfig,
|
||||
channels: Mapping[str, BaseChannel],
|
||||
stream_channels: str | Sequence[str],
|
||||
metadata: CheckpointMetadata,
|
||||
checkpoint: Checkpoint,
|
||||
tasks: Iterable[PregelExecutableTask],
|
||||
pending_writes: list[PendingWrite],
|
||||
parent_config: RunnableConfig | None,
|
||||
output_keys: str | Sequence[str],
|
||||
) -> Iterator[DebugOutputCheckpoint]:
|
||||
) -> Iterator[CheckpointPayload]:
|
||||
"""Produce "checkpoint" events for stream_mode=debug."""
|
||||
|
||||
parent_ns = config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
|
||||
@@ -193,42 +147,35 @@ def map_debug_checkpoint(
|
||||
}
|
||||
|
||||
yield {
|
||||
"type": "checkpoint",
|
||||
"timestamp": checkpoint["ts"],
|
||||
"step": step,
|
||||
"payload": {
|
||||
"config": rm_pregel_keys(patch_checkpoint_map(config, metadata)),
|
||||
"parent_config": rm_pregel_keys(
|
||||
patch_checkpoint_map(parent_config, metadata)
|
||||
),
|
||||
"values": read_channels(channels, stream_channels),
|
||||
"metadata": metadata,
|
||||
"next": [t.name for t in tasks],
|
||||
"tasks": [
|
||||
{
|
||||
"id": t.id,
|
||||
"name": t.name,
|
||||
"error": t.error,
|
||||
"state": t.state,
|
||||
}
|
||||
if t.error
|
||||
else {
|
||||
"id": t.id,
|
||||
"name": t.name,
|
||||
"result": t.result,
|
||||
"interrupts": tuple(asdict(i) for i in t.interrupts),
|
||||
"state": t.state,
|
||||
}
|
||||
if t.result
|
||||
else {
|
||||
"id": t.id,
|
||||
"name": t.name,
|
||||
"interrupts": tuple(asdict(i) for i in t.interrupts),
|
||||
"state": t.state,
|
||||
}
|
||||
for t in tasks_w_writes(tasks, pending_writes, task_states, output_keys)
|
||||
],
|
||||
},
|
||||
"config": rm_pregel_keys(patch_checkpoint_map(config, metadata)),
|
||||
"parent_config": rm_pregel_keys(patch_checkpoint_map(parent_config, metadata)),
|
||||
"values": read_channels(channels, stream_channels),
|
||||
"metadata": metadata,
|
||||
"next": [t.name for t in tasks],
|
||||
"tasks": [
|
||||
{
|
||||
"id": t.id,
|
||||
"name": t.name,
|
||||
"error": t.error,
|
||||
"state": t.state,
|
||||
}
|
||||
if t.error
|
||||
else {
|
||||
"id": t.id,
|
||||
"name": t.name,
|
||||
"result": t.result,
|
||||
"interrupts": tuple(asdict(i) for i in t.interrupts),
|
||||
"state": t.state,
|
||||
}
|
||||
if t.result
|
||||
else {
|
||||
"id": t.id,
|
||||
"name": t.name,
|
||||
"interrupts": tuple(asdict(i) for i in t.interrupts),
|
||||
"state": t.state,
|
||||
}
|
||||
for t in tasks_w_writes(tasks, pending_writes, task_states, output_keys)
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -3,7 +3,6 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import binascii
|
||||
import concurrent.futures
|
||||
import dataclasses
|
||||
from collections import defaultdict, deque
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from contextlib import (
|
||||
@@ -12,6 +11,7 @@ from contextlib import (
|
||||
AsyncExitStack,
|
||||
ExitStack,
|
||||
)
|
||||
from datetime import datetime, timezone
|
||||
from inspect import signature
|
||||
from types import TracebackType
|
||||
from typing import (
|
||||
@@ -25,7 +25,6 @@ from typing import (
|
||||
|
||||
from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import ParamSpec, Self
|
||||
|
||||
from langgraph.cache.base import BaseCache
|
||||
@@ -39,16 +38,12 @@ from langgraph.checkpoint.base import (
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
PendingWrite,
|
||||
copy_checkpoint,
|
||||
)
|
||||
from langgraph.constants import (
|
||||
CONF,
|
||||
CONFIG_KEY_CHECKPOINT_ID,
|
||||
CONFIG_KEY_CHECKPOINT_MAP,
|
||||
CONFIG_KEY_CHECKPOINT_NS,
|
||||
CONFIG_KEY_DEDUPE_TASKS,
|
||||
CONFIG_KEY_DELEGATE,
|
||||
CONFIG_KEY_ENSURE_LATEST,
|
||||
CONFIG_KEY_RESUME_MAP,
|
||||
CONFIG_KEY_RESUMING,
|
||||
CONFIG_KEY_SCRATCHPAD,
|
||||
@@ -60,17 +55,15 @@ from langgraph.constants import (
|
||||
INPUT,
|
||||
INTERRUPT,
|
||||
MISSING,
|
||||
NS_END,
|
||||
NS_SEP,
|
||||
NULL_TASK_ID,
|
||||
PUSH,
|
||||
RESUME,
|
||||
SCHEDULED,
|
||||
TAG_HIDDEN,
|
||||
)
|
||||
from langgraph.errors import (
|
||||
CheckpointNotLatest,
|
||||
EmptyInputError,
|
||||
GraphDelegate,
|
||||
GraphInterrupt,
|
||||
)
|
||||
from langgraph.managed.base import (
|
||||
@@ -91,6 +84,7 @@ from langgraph.pregel.algo import (
|
||||
)
|
||||
from langgraph.pregel.checkpoint import (
|
||||
channels_from_checkpoint,
|
||||
copy_checkpoint,
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
)
|
||||
@@ -125,6 +119,7 @@ from langgraph.types import (
|
||||
PregelScratchpad,
|
||||
RetryPolicy,
|
||||
StreamChunk,
|
||||
StreamMode,
|
||||
StreamProtocol,
|
||||
)
|
||||
from langgraph.utils.config import patch_configurable
|
||||
@@ -132,9 +127,7 @@ from langgraph.utils.config import patch_configurable
|
||||
V = TypeVar("V")
|
||||
P = ParamSpec("P")
|
||||
|
||||
INPUT_DONE = object()
|
||||
INPUT_RESUMING = object()
|
||||
INPUT_SHOULD_VALIDATE = object()
|
||||
|
||||
WritesT = Sequence[tuple[str, Any]]
|
||||
|
||||
|
||||
@@ -155,11 +148,11 @@ class PregelLoop:
|
||||
stop: int
|
||||
|
||||
input: Any | None
|
||||
input_model: type[BaseModel] | None
|
||||
cache: BaseCache[WritesT] | None
|
||||
checkpointer: BaseCheckpointSaver | None
|
||||
nodes: Mapping[str, PregelNode]
|
||||
specs: Mapping[str, BaseChannel | ManagedValueSpec]
|
||||
input_keys: str | Sequence[str]
|
||||
output_keys: str | Sequence[str]
|
||||
stream_keys: str | Sequence[str]
|
||||
skip_done_tasks: bool
|
||||
@@ -202,11 +195,16 @@ class PregelLoop:
|
||||
prev_checkpoint_config: RunnableConfig | None
|
||||
|
||||
status: Literal[
|
||||
"pending", "done", "interrupt_before", "interrupt_after", "out_of_steps"
|
||||
"input",
|
||||
"pending",
|
||||
"done",
|
||||
"interrupt_before",
|
||||
"interrupt_after",
|
||||
"out_of_steps",
|
||||
]
|
||||
tasks: dict[str, PregelExecutableTask]
|
||||
to_interrupt: list[PregelExecutableTask]
|
||||
output: None | dict[str, Any] | Any = None
|
||||
updated_channels: set[str] | None = None
|
||||
|
||||
# public
|
||||
|
||||
@@ -221,13 +219,13 @@ class PregelLoop:
|
||||
checkpointer: BaseCheckpointSaver | None,
|
||||
nodes: Mapping[str, PregelNode],
|
||||
specs: Mapping[str, BaseChannel | ManagedValueSpec],
|
||||
input_keys: str | Sequence[str],
|
||||
output_keys: str | Sequence[str],
|
||||
stream_keys: str | Sequence[str],
|
||||
trigger_to_nodes: Mapping[str, Sequence[str]],
|
||||
interrupt_after: All | Sequence[str] = EMPTY_SEQ,
|
||||
interrupt_before: All | Sequence[str] = EMPTY_SEQ,
|
||||
manager: None | AsyncParentRunManager | ParentRunManager = None,
|
||||
input_model: type[BaseModel] | None = None,
|
||||
debug: bool = False,
|
||||
migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
|
||||
retry_policy: Sequence[RetryPolicy] = (),
|
||||
@@ -240,21 +238,18 @@ class PregelLoop:
|
||||
self.step = 0
|
||||
self.stop = 0
|
||||
self.input = input
|
||||
self.input_model = input_model
|
||||
self.checkpointer = checkpointer
|
||||
self.cache = cache
|
||||
self.nodes = nodes
|
||||
self.specs = specs
|
||||
self.input_keys = input_keys
|
||||
self.output_keys = output_keys
|
||||
self.stream_keys = stream_keys
|
||||
self.interrupt_after = interrupt_after
|
||||
self.interrupt_before = interrupt_before
|
||||
self.manager = manager
|
||||
self.is_nested = CONFIG_KEY_TASK_ID in self.config.get(CONF, {})
|
||||
self.skip_done_tasks = (
|
||||
CONFIG_KEY_CHECKPOINT_ID not in config[CONF]
|
||||
or CONFIG_KEY_DEDUPE_TASKS in config[CONF]
|
||||
)
|
||||
self.skip_done_tasks = CONFIG_KEY_CHECKPOINT_ID not in config[CONF]
|
||||
self._migrate_checkpoint = migrate_checkpoint
|
||||
self.trigger_to_nodes = trigger_to_nodes
|
||||
self.retry_policy = retry_policy
|
||||
@@ -264,9 +259,7 @@ class PregelLoop:
|
||||
if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
|
||||
self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM])
|
||||
scratchpad: PregelScratchpad | None = config[CONF].get(CONFIG_KEY_SCRATCHPAD)
|
||||
if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and isinstance(
|
||||
scratchpad, PregelScratchpad
|
||||
):
|
||||
if isinstance(scratchpad, PregelScratchpad):
|
||||
# if count is > 0, append to checkpoint_ns
|
||||
# if count is 0, leave as is
|
||||
if cnt := scratchpad.subgraph_counter():
|
||||
@@ -404,12 +397,6 @@ class PregelLoop:
|
||||
self, task: PregelExecutableTask, write_idx: int, call: Call | None = None
|
||||
) -> PregelExecutableTask | None:
|
||||
"""Accept a PUSH from a task, potentially returning a new task to start."""
|
||||
# don't start if we should interrupt *after* the original task
|
||||
if self.interrupt_after and should_interrupt(
|
||||
self.checkpoint, self.interrupt_after, [task]
|
||||
):
|
||||
self.to_interrupt.append(task)
|
||||
return
|
||||
checkpoint_id_bytes = binascii.unhexlify(self.checkpoint["id"].replace("-", ""))
|
||||
null_version = checkpoint_null_version(self.checkpoint)
|
||||
if pushed := cast(
|
||||
@@ -435,14 +422,8 @@ class PregelLoop:
|
||||
cache_policy=self.cache_policy,
|
||||
),
|
||||
):
|
||||
# don't start if we should interrupt *before* the new task
|
||||
if self.interrupt_before and should_interrupt(
|
||||
self.checkpoint, self.interrupt_before, [pushed]
|
||||
):
|
||||
self.to_interrupt.append(pushed)
|
||||
return
|
||||
# produce debug output
|
||||
self._emit("debug", map_debug_tasks, self.step, [pushed])
|
||||
self._emit("tasks", map_debug_tasks, [pushed])
|
||||
# debug flag
|
||||
if self.debug:
|
||||
print_step_tasks(self.step, [pushed])
|
||||
@@ -454,11 +435,7 @@ class PregelLoop:
|
||||
# return the new task, to be started if not run before
|
||||
return pushed
|
||||
|
||||
def tick(
|
||||
self,
|
||||
*,
|
||||
input_keys: str | Sequence[str],
|
||||
) -> bool:
|
||||
def tick(self) -> bool:
|
||||
"""Execute a single iteration of the Pregel loop.
|
||||
|
||||
Args:
|
||||
@@ -467,72 +444,6 @@ class PregelLoop:
|
||||
Returns:
|
||||
True if more iterations are needed.
|
||||
"""
|
||||
if self.status != "pending":
|
||||
raise RuntimeError("Cannot tick when status is no longer 'pending'")
|
||||
|
||||
updated_channels: set[str] | None = None
|
||||
|
||||
if self.input not in (INPUT_DONE, INPUT_RESUMING, INPUT_SHOULD_VALIDATE):
|
||||
updated_channels = self._first(input_keys=input_keys)
|
||||
elif self.to_interrupt:
|
||||
# if we need to interrupt, do so
|
||||
self.status = "interrupt_before"
|
||||
raise GraphInterrupt()
|
||||
elif all(task.writes for task in self.tasks.values()):
|
||||
# finish superstep
|
||||
writes = [w for t in self.tasks.values() for w in t.writes]
|
||||
# debug flag
|
||||
if self.debug:
|
||||
print_step_writes(
|
||||
self.step,
|
||||
writes,
|
||||
(
|
||||
[self.stream_keys]
|
||||
if isinstance(self.stream_keys, str)
|
||||
else self.stream_keys
|
||||
),
|
||||
)
|
||||
# all tasks have finished
|
||||
updated_channels = apply_writes(
|
||||
self.checkpoint,
|
||||
self.channels,
|
||||
self.tasks.values(),
|
||||
self.checkpointer_get_next_version,
|
||||
self.trigger_to_nodes,
|
||||
)
|
||||
# validate input if requested
|
||||
if self.input is INPUT_SHOULD_VALIDATE:
|
||||
self.input = INPUT_DONE
|
||||
# validate
|
||||
cast(type[BaseModel], self.input_model)(
|
||||
**read_channels(self.channels, self.stream_keys)
|
||||
)
|
||||
# produce values output
|
||||
if not updated_channels.isdisjoint(
|
||||
(self.output_keys,)
|
||||
if isinstance(self.output_keys, str)
|
||||
else self.output_keys
|
||||
):
|
||||
self._emit(
|
||||
"values", map_output_values, self.output_keys, writes, self.channels
|
||||
)
|
||||
# clear pending writes
|
||||
self.checkpoint_pending_writes.clear()
|
||||
# "not skip_done_tasks" only applies to first tick after resuming
|
||||
self.skip_done_tasks = True
|
||||
# save checkpoint
|
||||
self._put_checkpoint({"source": "loop"})
|
||||
# after execution, check if we should interrupt
|
||||
if self.interrupt_after and should_interrupt(
|
||||
self.checkpoint, self.interrupt_after, self.tasks.values()
|
||||
):
|
||||
self.status = "interrupt_after"
|
||||
raise GraphInterrupt()
|
||||
|
||||
# unset resuming flag
|
||||
self.config[CONF].pop(CONFIG_KEY_RESUMING, None)
|
||||
else:
|
||||
return False
|
||||
|
||||
# check if iteration limit is reached
|
||||
if self.step > self.stop:
|
||||
@@ -554,18 +465,16 @@ class PregelLoop:
|
||||
store=self.store,
|
||||
checkpointer=self.checkpointer,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
updated_channels=updated_channels,
|
||||
updated_channels=self.updated_channels,
|
||||
retry_policy=self.retry_policy,
|
||||
cache_policy=self.cache_policy,
|
||||
)
|
||||
self.to_interrupt = []
|
||||
|
||||
# produce debug output
|
||||
if self._checkpointer_put_after_previous is not None:
|
||||
self._emit(
|
||||
"debug",
|
||||
"checkpoints",
|
||||
map_debug_checkpoint,
|
||||
self.step - 1, # printing checkpoint for previous step
|
||||
{
|
||||
**self.checkpoint_config,
|
||||
CONF: {
|
||||
@@ -576,7 +485,6 @@ class PregelLoop:
|
||||
self.channels,
|
||||
self.stream_keys,
|
||||
self.checkpoint_metadata,
|
||||
self.checkpoint,
|
||||
self.tasks.values(),
|
||||
self.checkpoint_pending_writes,
|
||||
self.prev_checkpoint_config,
|
||||
@@ -588,26 +496,10 @@ class PregelLoop:
|
||||
self.status = "done"
|
||||
return False
|
||||
|
||||
# check if we should delegate (used by subgraphs in distributed mode)
|
||||
if self.config[CONF].get(CONFIG_KEY_DELEGATE):
|
||||
assert self.input is INPUT_RESUMING
|
||||
raise GraphDelegate(
|
||||
{
|
||||
"config": patch_configurable(
|
||||
self.config, {CONFIG_KEY_DELEGATE: False}
|
||||
),
|
||||
"input": None,
|
||||
}
|
||||
)
|
||||
|
||||
# if there are pending writes from a previous loop, apply them
|
||||
if self.skip_done_tasks and self.checkpoint_pending_writes:
|
||||
self._match_writes(self.tasks)
|
||||
|
||||
# if all tasks have finished, re-tick
|
||||
if all(task.writes for task in self.tasks.values()):
|
||||
return self.tick(input_keys=input_keys)
|
||||
|
||||
# before execution, check if we should interrupt
|
||||
if self.interrupt_before and should_interrupt(
|
||||
self.checkpoint, self.interrupt_before, self.tasks.values()
|
||||
@@ -616,7 +508,7 @@ class PregelLoop:
|
||||
raise GraphInterrupt()
|
||||
|
||||
# produce debug output
|
||||
self._emit("debug", map_debug_tasks, self.step, self.tasks.values())
|
||||
self._emit("tasks", map_debug_tasks, self.tasks.values())
|
||||
|
||||
# debug flag
|
||||
if self.debug:
|
||||
@@ -629,6 +521,52 @@ class PregelLoop:
|
||||
|
||||
return True
|
||||
|
||||
def after_tick(self) -> None:
|
||||
# finish superstep
|
||||
writes = [w for t in self.tasks.values() for w in t.writes]
|
||||
# debug flag
|
||||
if self.debug:
|
||||
print_step_writes(
|
||||
self.step,
|
||||
writes,
|
||||
(
|
||||
[self.stream_keys]
|
||||
if isinstance(self.stream_keys, str)
|
||||
else self.stream_keys
|
||||
),
|
||||
)
|
||||
# all tasks have finished
|
||||
self.updated_channels = apply_writes(
|
||||
self.checkpoint,
|
||||
self.channels,
|
||||
self.tasks.values(),
|
||||
self.checkpointer_get_next_version,
|
||||
self.trigger_to_nodes,
|
||||
)
|
||||
# produce values output
|
||||
if not self.updated_channels.isdisjoint(
|
||||
(self.output_keys,)
|
||||
if isinstance(self.output_keys, str)
|
||||
else self.output_keys
|
||||
):
|
||||
self._emit(
|
||||
"values", map_output_values, self.output_keys, writes, self.channels
|
||||
)
|
||||
# clear pending writes
|
||||
self.checkpoint_pending_writes.clear()
|
||||
# "not skip_done_tasks" only applies to first tick after resuming
|
||||
self.skip_done_tasks = True
|
||||
# save checkpoint
|
||||
self._put_checkpoint({"source": "loop"})
|
||||
# after execution, check if we should interrupt
|
||||
if self.interrupt_after and should_interrupt(
|
||||
self.checkpoint, self.interrupt_after, self.tasks.values()
|
||||
):
|
||||
self.status = "interrupt_after"
|
||||
raise GraphInterrupt()
|
||||
# unset resuming flag
|
||||
self.config[CONF].pop(CONFIG_KEY_RESUMING, None)
|
||||
|
||||
def match_cached_writes(self) -> Sequence[PregelExecutableTask]:
|
||||
raise NotImplementedError
|
||||
|
||||
@@ -642,14 +580,7 @@ class PregelLoop:
|
||||
if k in (ERROR, INTERRUPT, RESUME):
|
||||
continue
|
||||
if task := tasks.get(tid):
|
||||
if k == SCHEDULED:
|
||||
if v == max(
|
||||
self.checkpoint["versions_seen"].get(INTERRUPT, {}).values(),
|
||||
default=None,
|
||||
):
|
||||
self.tasks[tid] = dataclasses.replace(task, scheduled=True)
|
||||
else:
|
||||
task.writes.append((k, v))
|
||||
task.writes.append((k, v))
|
||||
|
||||
def _first(self, *, input_keys: str | Sequence[str]) -> set[str] | None:
|
||||
# resuming from previous checkpoint requires
|
||||
@@ -715,21 +646,8 @@ class PregelLoop:
|
||||
self._emit(
|
||||
"values", map_output_values, self.output_keys, True, self.channels
|
||||
)
|
||||
# set flag
|
||||
self.input = INPUT_RESUMING
|
||||
# map inputs to channel updates
|
||||
elif input_writes := deque(map_input(input_keys, self.input)):
|
||||
# TODO shouldn't these writes be passed to put_writes too?
|
||||
# check if we should delegate (used by subgraphs in distributed mode)
|
||||
if self.config[CONF].get(CONFIG_KEY_DELEGATE):
|
||||
raise GraphDelegate(
|
||||
{
|
||||
"config": patch_configurable(
|
||||
self.config, {CONFIG_KEY_DELEGATE: False}
|
||||
),
|
||||
"input": self.input,
|
||||
}
|
||||
)
|
||||
# discard any unfinished tasks from previous checkpoint
|
||||
discard_tasks = prepare_next_tasks(
|
||||
self.checkpoint,
|
||||
@@ -758,24 +676,15 @@ class PregelLoop:
|
||||
)
|
||||
# save input checkpoint
|
||||
self._put_checkpoint({"source": "input"})
|
||||
# set flag
|
||||
if (
|
||||
self.input_model is not None
|
||||
and not isinstance(self.input, self.input_model)
|
||||
and not isinstance(self.stream_keys, str)
|
||||
):
|
||||
self.input = INPUT_SHOULD_VALIDATE
|
||||
else:
|
||||
self.input = INPUT_DONE
|
||||
elif CONFIG_KEY_RESUMING not in configurable:
|
||||
raise EmptyInputError(f"Received no input for {input_keys}")
|
||||
else:
|
||||
self.input = INPUT_DONE
|
||||
# update config
|
||||
if not self.is_nested:
|
||||
self.config = patch_configurable(
|
||||
self.config, {CONFIG_KEY_RESUMING: is_resuming}
|
||||
)
|
||||
# set flag
|
||||
self.status = "pending"
|
||||
return updated_channels
|
||||
|
||||
def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
|
||||
@@ -868,7 +777,14 @@ class PregelLoop:
|
||||
traceback: TracebackType | None,
|
||||
) -> bool | None:
|
||||
# persist current checkpoint and writes
|
||||
if not self.checkpoint_during:
|
||||
if not self.checkpoint_during and (
|
||||
# if it's a top graph
|
||||
not self.is_nested
|
||||
# or a nested graph with error or interrupt
|
||||
or exc_value is not None
|
||||
# or a nested graph with checkpointer=True
|
||||
or all(NS_END not in part for part in self.checkpoint_ns)
|
||||
):
|
||||
self._put_checkpoint(self.checkpoint_metadata)
|
||||
self._put_pending_writes()
|
||||
# suppress interrupt
|
||||
@@ -917,17 +833,39 @@ class PregelLoop:
|
||||
|
||||
def _emit(
|
||||
self,
|
||||
mode: str,
|
||||
mode: StreamMode,
|
||||
values: Callable[P, Iterator[Any]],
|
||||
*args: P.args,
|
||||
**kwargs: P.kwargs,
|
||||
) -> None:
|
||||
if self.stream is None:
|
||||
return
|
||||
if mode not in self.stream.modes:
|
||||
debug_remap = mode in ("checkpoints", "tasks") and "debug" in self.stream.modes
|
||||
if mode not in self.stream.modes and not debug_remap:
|
||||
return
|
||||
for v in values(*args, **kwargs):
|
||||
self.stream((self.checkpoint_ns, mode, v))
|
||||
if mode in self.stream.modes:
|
||||
self.stream((self.checkpoint_ns, mode, v))
|
||||
# "debug" mode is "checkpoints" or "tasks" with a wrapper dict
|
||||
if debug_remap:
|
||||
self.stream(
|
||||
(
|
||||
self.checkpoint_ns,
|
||||
"debug",
|
||||
{
|
||||
"step": self.step - 1
|
||||
if mode == "checkpoints"
|
||||
else self.step,
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"type": "checkpoint"
|
||||
if mode == "checkpoints"
|
||||
else "task_result"
|
||||
if "result" in v
|
||||
else "task",
|
||||
"payload": v,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
def output_writes(
|
||||
self, task_id: str, writes: WritesT, *, cached: bool = False
|
||||
@@ -968,9 +906,8 @@ class PregelLoop:
|
||||
)
|
||||
if not cached:
|
||||
self._emit(
|
||||
"debug",
|
||||
"tasks",
|
||||
map_debug_task_results,
|
||||
self.step,
|
||||
(task, writes),
|
||||
self.stream_keys,
|
||||
)
|
||||
@@ -992,9 +929,9 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
manager: None | AsyncParentRunManager | ParentRunManager = None,
|
||||
interrupt_after: All | Sequence[str] = EMPTY_SEQ,
|
||||
interrupt_before: All | Sequence[str] = EMPTY_SEQ,
|
||||
input_keys: str | Sequence[str] = EMPTY_SEQ,
|
||||
output_keys: str | Sequence[str] = EMPTY_SEQ,
|
||||
stream_keys: str | Sequence[str] = EMPTY_SEQ,
|
||||
input_model: type[BaseModel] | None = None,
|
||||
debug: bool = False,
|
||||
migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
|
||||
retry_policy: Sequence[RetryPolicy] = (),
|
||||
@@ -1003,7 +940,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
) -> None:
|
||||
super().__init__(
|
||||
input,
|
||||
input_model=input_model,
|
||||
stream=stream,
|
||||
config=config,
|
||||
checkpointer=checkpointer,
|
||||
@@ -1011,6 +947,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
store=store,
|
||||
nodes=nodes,
|
||||
specs=specs,
|
||||
input_keys=input_keys,
|
||||
output_keys=output_keys,
|
||||
stream_keys=stream_keys,
|
||||
interrupt_after=interrupt_after,
|
||||
@@ -1097,25 +1034,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
# context manager
|
||||
|
||||
def __enter__(self) -> Self:
|
||||
if self.config.get(CONF, {}).get(
|
||||
CONFIG_KEY_ENSURE_LATEST
|
||||
) and self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID):
|
||||
if self.checkpointer is None:
|
||||
raise RuntimeError(
|
||||
"Cannot ensure latest checkpoint without checkpointer"
|
||||
)
|
||||
saved = self.checkpointer.get_tuple(
|
||||
patch_configurable(
|
||||
self.checkpoint_config, {CONFIG_KEY_CHECKPOINT_ID: None}
|
||||
)
|
||||
)
|
||||
if (
|
||||
saved is None
|
||||
or saved.checkpoint["id"]
|
||||
!= self.checkpoint_config[CONF][CONFIG_KEY_CHECKPOINT_ID]
|
||||
):
|
||||
raise CheckpointNotLatest
|
||||
elif self.checkpointer:
|
||||
if self.checkpointer:
|
||||
saved = self.checkpointer.get_tuple(self.checkpoint_config)
|
||||
else:
|
||||
saved = None
|
||||
@@ -1149,10 +1068,11 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
self.specs, self.checkpoint
|
||||
)
|
||||
self.stack.push(self._suppress_interrupt)
|
||||
self.status = "pending"
|
||||
self.status = "input"
|
||||
self.step = self.checkpoint_metadata["step"] + 1
|
||||
self.stop = self.step + self.config["recursion_limit"] + 1
|
||||
self.checkpoint_previous_versions = self.checkpoint["channel_versions"].copy()
|
||||
self.updated_channels = self._first(input_keys=self.input_keys)
|
||||
|
||||
return self
|
||||
|
||||
@@ -1182,9 +1102,9 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
interrupt_after: All | Sequence[str] = EMPTY_SEQ,
|
||||
interrupt_before: All | Sequence[str] = EMPTY_SEQ,
|
||||
manager: None | AsyncParentRunManager | ParentRunManager = None,
|
||||
input_keys: str | Sequence[str] = EMPTY_SEQ,
|
||||
output_keys: str | Sequence[str] = EMPTY_SEQ,
|
||||
stream_keys: str | Sequence[str] = EMPTY_SEQ,
|
||||
input_model: type[BaseModel] | None = None,
|
||||
debug: bool = False,
|
||||
migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
|
||||
retry_policy: Sequence[RetryPolicy] = (),
|
||||
@@ -1193,7 +1113,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
) -> None:
|
||||
super().__init__(
|
||||
input,
|
||||
input_model=input_model,
|
||||
stream=stream,
|
||||
config=config,
|
||||
checkpointer=checkpointer,
|
||||
@@ -1201,6 +1120,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
store=store,
|
||||
nodes=nodes,
|
||||
specs=specs,
|
||||
input_keys=input_keys,
|
||||
output_keys=output_keys,
|
||||
stream_keys=stream_keys,
|
||||
interrupt_after=interrupt_after,
|
||||
@@ -1290,25 +1210,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
# context manager
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
if self.config.get(CONF, {}).get(
|
||||
CONFIG_KEY_ENSURE_LATEST
|
||||
) and self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID):
|
||||
if self.checkpointer is None:
|
||||
raise RuntimeError(
|
||||
"Cannot ensure latest checkpoint without checkpointer"
|
||||
)
|
||||
saved = await self.checkpointer.aget_tuple(
|
||||
patch_configurable(
|
||||
self.checkpoint_config, {CONFIG_KEY_CHECKPOINT_ID: None}
|
||||
)
|
||||
)
|
||||
if (
|
||||
saved is None
|
||||
or saved.checkpoint["id"]
|
||||
!= self.checkpoint_config[CONF][CONFIG_KEY_CHECKPOINT_ID]
|
||||
):
|
||||
raise CheckpointNotLatest
|
||||
elif self.checkpointer:
|
||||
if self.checkpointer:
|
||||
saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
|
||||
else:
|
||||
saved = None
|
||||
@@ -1344,11 +1246,11 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
self.specs, self.checkpoint
|
||||
)
|
||||
self.stack.push(self._suppress_interrupt)
|
||||
self.status = "pending"
|
||||
self.status = "input"
|
||||
self.step = self.checkpoint_metadata["step"] + 1
|
||||
self.stop = self.step + self.config["recursion_limit"] + 1
|
||||
|
||||
self.checkpoint_previous_versions = self.checkpoint["channel_versions"].copy()
|
||||
self.updated_channels = self._first(input_keys=self.input_keys)
|
||||
|
||||
return self
|
||||
|
||||
|
||||
@@ -12,12 +12,11 @@ from langchain_core.runnables import Runnable, RunnableConfig
|
||||
|
||||
from langgraph.constants import CONF, CONFIG_KEY_READ
|
||||
from langgraph.pregel.protocol import PregelProtocol
|
||||
from langgraph.pregel.retry import RetryPolicy
|
||||
from langgraph.pregel.utils import find_subgraph_pregel
|
||||
from langgraph.pregel.write import ChannelWrite
|
||||
from langgraph.types import CachePolicy
|
||||
from langgraph.types import CachePolicy, RetryPolicy
|
||||
from langgraph.utils.config import merge_configs
|
||||
from langgraph.utils.runnable import RunnableCallable, RunnableSeq, coerce_to_runnable
|
||||
from langgraph.utils.runnable import RunnableCallable, RunnableSeq
|
||||
|
||||
READ_TYPE = Callable[[Union[str, Sequence[str]], bool], Union[Any, dict[str, Any]]]
|
||||
INPUT_CACHE_KEY_TYPE = tuple[Callable[..., Any], tuple[str, ...]]
|
||||
@@ -96,16 +95,15 @@ class ChannelRead(RunnableCallable):
|
||||
DEFAULT_BOUND = RunnableCallable(lambda input: input)
|
||||
|
||||
|
||||
class PregelNode(Runnable):
|
||||
class PregelNode:
|
||||
"""A node in a Pregel graph. This won't be invoked as a runnable by the graph
|
||||
itself, but instead acts as a container for the components necessary to make
|
||||
a PregelExecutableTask for a node."""
|
||||
|
||||
channels: list[str] | Mapping[str, str]
|
||||
channels: str | list[str]
|
||||
"""The channels that will be passed as input to `bound`.
|
||||
If a list, the node will be invoked with the first of that isn't empty.
|
||||
If a dict, the keys are the names of the channels, and the values are the keys
|
||||
to use in the input to `bound`."""
|
||||
If a str, the node will be invoked with its value if it isn't empty.
|
||||
If a list, the node will be invoked with a dict of those channels' values."""
|
||||
|
||||
triggers: list[str]
|
||||
"""If any of these channels is written to, this node will be triggered in
|
||||
@@ -140,7 +138,7 @@ class PregelNode(Runnable):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
channels: list[str] | Mapping[str, str],
|
||||
channels: str | list[str],
|
||||
triggers: Sequence[str],
|
||||
mapper: Callable[[Any], Any] | None = None,
|
||||
writers: list[Runnable] | None = None,
|
||||
@@ -223,59 +221,11 @@ class PregelNode(Runnable):
|
||||
This is used to avoid calculating the same input multiple times."""
|
||||
return (
|
||||
self.mapper,
|
||||
tuple(f"{key}:{value}" for key, value in self.channels.items())
|
||||
if isinstance(self.channels, dict)
|
||||
else tuple(self.channels),
|
||||
tuple(self.channels)
|
||||
if isinstance(self.channels, list)
|
||||
else (self.channels,),
|
||||
)
|
||||
|
||||
def join(self, channels: Sequence[str]) -> PregelNode:
|
||||
assert isinstance(channels, list) or isinstance(channels, tuple), (
|
||||
"channels must be a list or tuple"
|
||||
)
|
||||
assert isinstance(self.channels, dict), (
|
||||
"all channels must be named when using .join()"
|
||||
)
|
||||
return self.copy(
|
||||
update=dict(
|
||||
channels={
|
||||
**self.channels,
|
||||
**{chan: chan for chan in channels},
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
def __or__(
|
||||
self,
|
||||
other: Runnable[Any, Any]
|
||||
| Callable[[Any], Any]
|
||||
| Mapping[str, Runnable[Any, Any] | Callable[[Any], Any]],
|
||||
) -> PregelNode:
|
||||
if isinstance(other, Runnable) and ChannelWrite.is_writer(other):
|
||||
return self.copy(update=dict(writers=[*self.writers, other]))
|
||||
elif self.bound is DEFAULT_BOUND:
|
||||
return self.copy(
|
||||
update=dict(bound=coerce_to_runnable(other, name=None, trace=True))
|
||||
)
|
||||
else:
|
||||
return self.copy(update=dict(bound=RunnableSeq(self.bound, other)))
|
||||
|
||||
def pipe(
|
||||
self,
|
||||
*others: Runnable[Any, Any] | Callable[[Any], Any],
|
||||
name: str | None = None,
|
||||
) -> PregelNode:
|
||||
for other in others:
|
||||
self = self | other
|
||||
return self
|
||||
|
||||
def __ror__(
|
||||
self,
|
||||
other: Runnable[Any, Any]
|
||||
| Callable[[Any], Any]
|
||||
| Mapping[str, Runnable[Any, Any] | Callable[[Any], Any]],
|
||||
) -> PregelNode:
|
||||
raise NotImplementedError()
|
||||
|
||||
def invoke(
|
||||
self,
|
||||
input: Any,
|
||||
|
||||
@@ -5,6 +5,7 @@ from typing import Any
|
||||
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.constants import RESERVED
|
||||
from langgraph.managed.base import ManagedValueMapping
|
||||
from langgraph.pregel.read import PregelNode
|
||||
from langgraph.types import All
|
||||
|
||||
@@ -12,6 +13,7 @@ from langgraph.types import All
|
||||
def validate_graph(
|
||||
nodes: Mapping[str, PregelNode],
|
||||
channels: dict[str, BaseChannel],
|
||||
managed: ManagedValueMapping,
|
||||
input_channels: str | Sequence[str],
|
||||
output_channels: str | Sequence[str],
|
||||
stream_channels: str | Sequence[str] | None,
|
||||
@@ -20,14 +22,30 @@ def validate_graph(
|
||||
) -> None:
|
||||
for chan in channels:
|
||||
if chan in RESERVED:
|
||||
raise ValueError(f"Channel names {chan} are reserved")
|
||||
raise ValueError(f"Channel name '{chan}' is reserved")
|
||||
for name in managed:
|
||||
if name in RESERVED:
|
||||
raise ValueError(f"Managed name '{name}' is reserved")
|
||||
|
||||
subscribed_channels = set[str]()
|
||||
for name, node in nodes.items():
|
||||
if name in RESERVED:
|
||||
raise ValueError(f"Node names {RESERVED} are reserved")
|
||||
raise ValueError(f"Node name '{name}' is reserved")
|
||||
if isinstance(node, PregelNode):
|
||||
subscribed_channels.update(node.triggers)
|
||||
if isinstance(node.channels, str):
|
||||
if node.channels not in channels:
|
||||
raise ValueError(
|
||||
f"Node {name} reads channel '{node.channels}' "
|
||||
f"not in known channels: '{repr(sorted(channels))[:100]}'"
|
||||
)
|
||||
else:
|
||||
for chan in node.channels:
|
||||
if chan not in channels and chan not in managed:
|
||||
raise ValueError(
|
||||
f"Node {name} reads channel '{chan}' "
|
||||
f"not in known channels: '{repr(sorted(channels))[:100]}'"
|
||||
)
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Invalid node type {type(node)}, expected PregelNode or NodeBuilder"
|
||||
|
||||
@@ -46,7 +46,9 @@ Checkpointer = Union[None, bool, BaseCheckpointSaver]
|
||||
- False disables checkpointing, even if the parent graph has a checkpointer.
|
||||
- None inherits checkpointer from the parent graph."""
|
||||
|
||||
StreamMode = Literal["values", "updates", "debug", "messages", "custom"]
|
||||
StreamMode = Literal[
|
||||
"values", "updates", "checkpoints", "tasks", "debug", "messages", "custom"
|
||||
]
|
||||
"""How the stream method should emit outputs.
|
||||
|
||||
- `"values"`: Emit all values in the state after each step, including interrupts.
|
||||
@@ -55,7 +57,9 @@ StreamMode = Literal["values", "updates", "debug", "messages", "custom"]
|
||||
If multiple updates are made in the same step (e.g. multiple nodes are run) then those updates are emitted separately.
|
||||
- `"custom"`: Emit custom data using from inside nodes or tasks using `StreamWriter`.
|
||||
- `"messages"`: Emit LLM messages token-by-token together with metadata for any LLM invocations inside nodes or tasks.
|
||||
- `"debug"`: Emit debug events with as much information as possible for each step.
|
||||
- `"checkpoints"`: Emit an event when a checkpoint is created, in the same format as returned by get_state().
|
||||
- `"tasks"`: Emit events when tasks start and finish, including their results and errors.
|
||||
- `"debug"`: Emit "checlkpoints" and "tasks" events, for debugging purposes.
|
||||
"""
|
||||
|
||||
StreamWriter = Callable[[Any], None]
|
||||
@@ -203,7 +207,6 @@ class PregelExecutableTask:
|
||||
cache_key: CacheKey | None
|
||||
id: str
|
||||
path: tuple[str | int | tuple, ...]
|
||||
scheduled: bool = False
|
||||
writers: Sequence[Runnable] = ()
|
||||
subgraphs: Sequence[PregelProtocol] = ()
|
||||
|
||||
|
||||
@@ -41,8 +41,8 @@ class LangGraphDeprecationWarning(DeprecationWarning):
|
||||
return message
|
||||
|
||||
|
||||
class LangGraphDeprecatedSinceV10(LangGraphDeprecationWarning):
|
||||
"""A specific `LangGraphDeprecationWarning` subclass defining functionality deprecated since LangGraph v1.0.0"""
|
||||
class LangGraphDeprecatedSinceV05(LangGraphDeprecationWarning):
|
||||
"""A specific `LangGraphDeprecationWarning` subclass defining functionality deprecated since LangGraph v0.5.0"""
|
||||
|
||||
def __init__(self, message: str, *args: object) -> None:
|
||||
super().__init__(message, *args, since=(1, 0), expected_removal=(2, 0))
|
||||
super().__init__(message, *args, since=(0, 5), expected_removal=(2, 0))
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph"
|
||||
version = "0.4.7"
|
||||
version = "0.5.0rc1"
|
||||
description = "Building stateful, multi-actor applications with LLMs"
|
||||
authors = []
|
||||
requires-python = ">=3.9"
|
||||
@@ -13,9 +13,9 @@ license = "MIT"
|
||||
license-files = ['LICENSE']
|
||||
dependencies = [
|
||||
"langchain-core>=0.1",
|
||||
"langgraph-checkpoint>=2.0.26",
|
||||
"langgraph-checkpoint>=2.1.0",
|
||||
"langgraph-sdk>=0.1.42",
|
||||
"langgraph-prebuilt>=0.2.0",
|
||||
"langgraph-prebuilt>=0.5.0rc0",
|
||||
"xxhash>=3.5.0",
|
||||
"pydantic>=2.7.4",
|
||||
]
|
||||
|
||||
@@ -12,6 +12,7 @@ from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from langgraph.store.base import BaseStore
|
||||
from tests.conftest_checkpointer import (
|
||||
_checkpointer_memory,
|
||||
_checkpointer_memory_migrate_sends,
|
||||
_checkpointer_postgres,
|
||||
_checkpointer_postgres_aio,
|
||||
_checkpointer_postgres_aio_pipe,
|
||||
@@ -125,6 +126,7 @@ async def async_store(request: pytest.FixtureRequest) -> AsyncIterator[BaseStore
|
||||
if NO_DOCKER
|
||||
else [
|
||||
"memory",
|
||||
"memory_migrate_sends",
|
||||
"sqlite",
|
||||
"sqlite_aes",
|
||||
"postgres",
|
||||
@@ -139,6 +141,9 @@ def sync_checkpointer(
|
||||
if checkpointer_name == "memory":
|
||||
with _checkpointer_memory() as checkpointer:
|
||||
yield checkpointer
|
||||
elif checkpointer_name == "memory_migrate_sends":
|
||||
with _checkpointer_memory_migrate_sends() as checkpointer:
|
||||
yield checkpointer
|
||||
elif checkpointer_name == "sqlite":
|
||||
with _checkpointer_sqlite() as checkpointer:
|
||||
yield checkpointer
|
||||
|
||||
@@ -14,7 +14,10 @@ from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||||
|
||||
pytest.register_assert_rewrite("tests.memory_assert")
|
||||
|
||||
from tests.memory_assert import MemorySaverAssertImmutable # noqa: E402
|
||||
from tests.memory_assert import ( # noqa: E402
|
||||
MemorySaverAssertImmutable,
|
||||
MemorySaverNeedsPendingSendsMigration,
|
||||
)
|
||||
|
||||
DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/"
|
||||
|
||||
@@ -24,6 +27,11 @@ def _checkpointer_memory():
|
||||
yield MemorySaverAssertImmutable()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _checkpointer_memory_migrate_sends():
|
||||
yield MemorySaverNeedsPendingSendsMigration()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _checkpointer_sqlite():
|
||||
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
|
||||
@@ -187,6 +195,7 @@ async def _checkpointer_postgres_aio_pool():
|
||||
|
||||
__all__ = [
|
||||
"_checkpointer_memory",
|
||||
"_checkpointer_memory_migrate_sends",
|
||||
"_checkpointer_sqlite",
|
||||
"_checkpointer_sqlite_aes",
|
||||
"_checkpointer_postgres",
|
||||
|
||||
@@ -7,6 +7,7 @@ from typing import Any, Optional
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
@@ -14,6 +15,7 @@ from langgraph.checkpoint.base import (
|
||||
SerializerProtocol,
|
||||
)
|
||||
from langgraph.checkpoint.memory import InMemorySaver, PersistentDict
|
||||
from langgraph.constants import TASKS
|
||||
|
||||
|
||||
class NoopSerializer(SerializerProtocol):
|
||||
@@ -24,6 +26,28 @@ class NoopSerializer(SerializerProtocol):
|
||||
return "type", obj
|
||||
|
||||
|
||||
class MemorySaverNeedsPendingSendsMigration(BaseCheckpointSaver):
|
||||
def __init__(self) -> None:
|
||||
self.saver = InMemorySaver()
|
||||
|
||||
def __getattribute__(self, name):
|
||||
if name in ("saver", "__class__", "get_tuple"):
|
||||
return object.__getattribute__(self, name)
|
||||
return getattr(self.saver, name)
|
||||
|
||||
def get_tuple(self, config):
|
||||
if tup := self.saver.get_tuple(config):
|
||||
if tup.checkpoint["v"] == 4 and tup.checkpoint["channel_values"].get(TASKS):
|
||||
tup.checkpoint["v"] = 3
|
||||
tup.checkpoint["pending_sends"] = tup.checkpoint["channel_values"].pop(
|
||||
TASKS
|
||||
)
|
||||
tup.checkpoint["channel_versions"].pop(TASKS)
|
||||
for seen in tup.checkpoint["versions_seen"].values():
|
||||
seen.pop(TASKS, None)
|
||||
return tup
|
||||
|
||||
|
||||
class MemorySaverAssertImmutable(InMemorySaver):
|
||||
storage_for_copies: defaultdict[str, dict[str, dict[str, Checkpoint]]]
|
||||
|
||||
|
||||
@@ -7,12 +7,9 @@ from typing import Annotated, Literal, Optional, Union
|
||||
import pytest
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
CheckpointTuple,
|
||||
copy_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointTuple
|
||||
from langgraph.graph.state import StateGraph
|
||||
from langgraph.pregel.checkpoint import copy_checkpoint
|
||||
from langgraph.types import Command, Interrupt, PregelTask, StateSnapshot, interrupt
|
||||
from langgraph.utils.config import patch_configurable
|
||||
from tests.any_int import AnyInt
|
||||
|
||||
@@ -4,7 +4,7 @@ from typing_extensions import TypedDict
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.types import RetryPolicy
|
||||
from langgraph.warnings import LangGraphDeprecatedSinceV10
|
||||
from langgraph.warnings import LangGraphDeprecatedSinceV05
|
||||
|
||||
|
||||
class PlainState(TypedDict): ...
|
||||
@@ -14,7 +14,7 @@ def test_add_node_retry_arg() -> None:
|
||||
builder = StateGraph(PlainState)
|
||||
|
||||
with pytest.warns(
|
||||
LangGraphDeprecatedSinceV10,
|
||||
LangGraphDeprecatedSinceV05,
|
||||
match="`retry` is deprecated and will be removed. Please use `retry_policy` instead.",
|
||||
):
|
||||
builder.add_node("test_node", lambda state: state, retry=RetryPolicy()) # type: ignore[arg-type]
|
||||
@@ -22,7 +22,7 @@ def test_add_node_retry_arg() -> None:
|
||||
|
||||
def test_task_retry_arg() -> None:
|
||||
with pytest.warns(
|
||||
LangGraphDeprecatedSinceV10,
|
||||
LangGraphDeprecatedSinceV05,
|
||||
match="`retry` is deprecated and will be removed. Please use `retry_policy` instead.",
|
||||
):
|
||||
|
||||
@@ -33,7 +33,7 @@ def test_task_retry_arg() -> None:
|
||||
|
||||
def test_entrypoint_retry_arg() -> None:
|
||||
with pytest.warns(
|
||||
LangGraphDeprecatedSinceV10,
|
||||
LangGraphDeprecatedSinceV05,
|
||||
match="`retry` is deprecated and will be removed. Please use `retry_policy` instead.",
|
||||
):
|
||||
|
||||
@@ -44,7 +44,7 @@ def test_entrypoint_retry_arg() -> None:
|
||||
|
||||
def test_state_graph_input_schema() -> None:
|
||||
with pytest.warns(
|
||||
LangGraphDeprecatedSinceV10,
|
||||
LangGraphDeprecatedSinceV05,
|
||||
match="`input` is deprecated and will be removed. Please use `input_schema` instead.",
|
||||
):
|
||||
StateGraph(PlainState, input=PlainState) # type: ignore[arg-type]
|
||||
@@ -52,7 +52,7 @@ def test_state_graph_input_schema() -> None:
|
||||
|
||||
def test_state_graph_output_schema() -> None:
|
||||
with pytest.warns(
|
||||
LangGraphDeprecatedSinceV10,
|
||||
LangGraphDeprecatedSinceV05,
|
||||
match="`output` is deprecated and will be removed. Please use `output_schema` instead.",
|
||||
):
|
||||
StateGraph(PlainState, output=PlainState) # type: ignore[arg-type]
|
||||
@@ -62,7 +62,7 @@ def test_add_node_input_schema() -> None:
|
||||
builder = StateGraph(PlainState)
|
||||
|
||||
with pytest.warns(
|
||||
LangGraphDeprecatedSinceV10,
|
||||
LangGraphDeprecatedSinceV05,
|
||||
match="`input` is deprecated and will be removed. Please use `input_schema` instead.",
|
||||
):
|
||||
builder.add_node("test_node", lambda state: state, input=PlainState) # type: ignore[arg-type]
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -16,13 +16,14 @@ from langchain_core.runnables import RunnableConfig, RunnablePick
|
||||
from pytest_mock import MockerFixture
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
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, PULL, PUSH, START
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.graph.message import MessageGraph, add_messages
|
||||
from langgraph.graph.state import StateGraph
|
||||
from langgraph.prebuilt.chat_agent_executor import create_react_agent
|
||||
from langgraph.prebuilt.tool_node import ToolNode
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import PregelTask, Send, StateSnapshot, StreamWriter
|
||||
from tests.any_int import AnyInt
|
||||
@@ -499,7 +500,7 @@ async def test_conditional_graph_state(async_checkpointer: BaseCheckpointSaver)
|
||||
from langchain_core.tools import tool
|
||||
|
||||
class AgentState(TypedDict):
|
||||
input: Annotated[str, EphemeralValue]
|
||||
input: Annotated[str, UntrackedValue]
|
||||
agent_outcome: Optional[Union[AgentAction, AgentFinish]]
|
||||
intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
|
||||
|
||||
@@ -574,6 +575,7 @@ async def test_conditional_graph_state(async_checkpointer: BaseCheckpointSaver)
|
||||
app = workflow.compile()
|
||||
|
||||
assert await app.ainvoke({"input": "what is weather in sf"}) == {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
@@ -2082,7 +2084,417 @@ async def test_state_graph_packets(async_checkpointer: BaseCheckpointSaver) -> N
|
||||
)
|
||||
|
||||
|
||||
async def test_message_graph(async_checkpointer: BaseCheckpointSaver) -> None:
|
||||
from langchain_core.language_models.fake_chat_models import (
|
||||
FakeMessagesListChatModel,
|
||||
)
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langchain_core.tools import tool
|
||||
|
||||
class FakeFuntionChatModel(FakeMessagesListChatModel):
|
||||
def bind_functions(self, functions: list):
|
||||
return self
|
||||
|
||||
@tool()
|
||||
def search_api(query: str) -> str:
|
||||
"""Searches the API for the query."""
|
||||
return f"result for {query}"
|
||||
|
||||
tools = [search_api]
|
||||
|
||||
model = FakeFuntionChatModel(
|
||||
responses=[
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call123",
|
||||
"name": "search_api",
|
||||
"args": {"query": "query"},
|
||||
}
|
||||
],
|
||||
id="ai1",
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call456",
|
||||
"name": "search_api",
|
||||
"args": {"query": "another"},
|
||||
}
|
||||
],
|
||||
id="ai2",
|
||||
),
|
||||
AIMessage(content="answer", id="ai3"),
|
||||
]
|
||||
)
|
||||
|
||||
# Define the function that determines whether to continue or not
|
||||
def should_continue(messages):
|
||||
last_message = messages[-1]
|
||||
# If there is no function call, then we finish
|
||||
if not last_message.tool_calls:
|
||||
return "end"
|
||||
# Otherwise if there is, we continue
|
||||
else:
|
||||
return "continue"
|
||||
|
||||
# Define a new graph
|
||||
workflow = MessageGraph()
|
||||
|
||||
# Define the two nodes we will cycle between
|
||||
workflow.add_node("agent", model)
|
||||
workflow.add_node("tools", ToolNode(tools))
|
||||
|
||||
# Set the entrypoint as `agent`
|
||||
# This means that this node is the first one called
|
||||
workflow.set_entry_point("agent")
|
||||
|
||||
# We now add a conditional edge
|
||||
workflow.add_conditional_edges(
|
||||
# First, we define the start node. We use `agent`.
|
||||
# This means these are the edges taken after the `agent` node is called.
|
||||
"agent",
|
||||
# Next, we pass in the function that will determine which node is called next.
|
||||
should_continue,
|
||||
# Finally we pass in a mapping.
|
||||
# The keys are strings, and the values are other nodes.
|
||||
# END is a special node marking that the graph should finish.
|
||||
# What will happen is we will call `should_continue`, and then the output of that
|
||||
# will be matched against the keys in this mapping.
|
||||
# Based on which one it matches, that node will then be called.
|
||||
{
|
||||
# If `tools`, then we call the tool node.
|
||||
"continue": "tools",
|
||||
# Otherwise we finish.
|
||||
"end": END,
|
||||
},
|
||||
)
|
||||
|
||||
# We now add a normal edge from `tools` to `agent`.
|
||||
# This means that after `tools` is called, `agent` node is called next.
|
||||
workflow.add_edge("tools", "agent")
|
||||
|
||||
# Finally, we compile it!
|
||||
# This compiles it into a LangChain Runnable,
|
||||
# meaning you can use it as you would any other runnable
|
||||
app = workflow.compile()
|
||||
|
||||
assert await app.ainvoke(HumanMessage(content="what is weather in sf")) == [
|
||||
_AnyIdHumanMessage(
|
||||
content="what is weather in sf",
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call123",
|
||||
"name": "search_api",
|
||||
"args": {"query": "query"},
|
||||
}
|
||||
],
|
||||
id="ai1", # respects ids passed in
|
||||
),
|
||||
_AnyIdToolMessage(
|
||||
content="result for query",
|
||||
name="search_api",
|
||||
tool_call_id="tool_call123",
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call456",
|
||||
"name": "search_api",
|
||||
"args": {"query": "another"},
|
||||
}
|
||||
],
|
||||
id="ai2",
|
||||
),
|
||||
_AnyIdToolMessage(
|
||||
content="result for another",
|
||||
name="search_api",
|
||||
tool_call_id="tool_call456",
|
||||
),
|
||||
AIMessage(content="answer", id="ai3"),
|
||||
]
|
||||
|
||||
assert [
|
||||
c async for c in app.astream([HumanMessage(content="what is weather in sf")])
|
||||
] == [
|
||||
{
|
||||
"agent": AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call123",
|
||||
"name": "search_api",
|
||||
"args": {"query": "query"},
|
||||
}
|
||||
],
|
||||
id="ai1",
|
||||
)
|
||||
},
|
||||
{
|
||||
"tools": [
|
||||
_AnyIdToolMessage(
|
||||
content="result for query",
|
||||
name="search_api",
|
||||
tool_call_id="tool_call123",
|
||||
)
|
||||
]
|
||||
},
|
||||
{
|
||||
"agent": AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call456",
|
||||
"name": "search_api",
|
||||
"args": {"query": "another"},
|
||||
}
|
||||
],
|
||||
id="ai2",
|
||||
)
|
||||
},
|
||||
{
|
||||
"tools": [
|
||||
_AnyIdToolMessage(
|
||||
content="result for another",
|
||||
name="search_api",
|
||||
tool_call_id="tool_call456",
|
||||
)
|
||||
]
|
||||
},
|
||||
{"agent": AIMessage(content="answer", id="ai3")},
|
||||
]
|
||||
|
||||
app_w_interrupt = workflow.compile(
|
||||
checkpointer=async_checkpointer,
|
||||
interrupt_after=["agent"],
|
||||
)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
assert [
|
||||
c
|
||||
async for c in app_w_interrupt.astream(
|
||||
HumanMessage(content="what is weather in sf"), config
|
||||
)
|
||||
] == [
|
||||
{
|
||||
"agent": AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call123",
|
||||
"name": "search_api",
|
||||
"args": {"query": "query"},
|
||||
}
|
||||
],
|
||||
id="ai1",
|
||||
)
|
||||
},
|
||||
{"__interrupt__": ()},
|
||||
]
|
||||
|
||||
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values=[
|
||||
_AnyIdHumanMessage(content="what is weather in sf"),
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call123",
|
||||
"name": "search_api",
|
||||
"args": {"query": "query"},
|
||||
}
|
||||
],
|
||||
id="ai1",
|
||||
),
|
||||
],
|
||||
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
|
||||
next=("tools",),
|
||||
config=tup.config,
|
||||
created_at=tup.checkpoint["ts"],
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"step": 1,
|
||||
"thread_id": "1",
|
||||
},
|
||||
parent_config=None,
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
# modify ai message
|
||||
last_message = (await app_w_interrupt.aget_state(config)).values[-1]
|
||||
last_message.tool_calls[0]["args"] = {"query": "a different query"}
|
||||
await app_w_interrupt.aupdate_state(config, last_message)
|
||||
|
||||
# message was replaced instead of appended
|
||||
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values=[
|
||||
_AnyIdHumanMessage(content="what is weather in sf"),
|
||||
AIMessage(
|
||||
content="",
|
||||
id="ai1",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call123",
|
||||
"name": "search_api",
|
||||
"args": {"query": "a different query"},
|
||||
}
|
||||
],
|
||||
),
|
||||
],
|
||||
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
|
||||
next=("tools",),
|
||||
config=tup.config,
|
||||
created_at=tup.checkpoint["ts"],
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "update",
|
||||
"step": 2,
|
||||
"thread_id": "1",
|
||||
},
|
||||
parent_config=(
|
||||
[c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][
|
||||
-1
|
||||
].config
|
||||
),
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
assert [c async for c in app_w_interrupt.astream(None, config)] == [
|
||||
{
|
||||
"tools": [
|
||||
_AnyIdToolMessage(
|
||||
content="result for a different query",
|
||||
name="search_api",
|
||||
tool_call_id="tool_call123",
|
||||
)
|
||||
]
|
||||
},
|
||||
{
|
||||
"agent": AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call456",
|
||||
"name": "search_api",
|
||||
"args": {"query": "another"},
|
||||
}
|
||||
],
|
||||
id="ai2",
|
||||
)
|
||||
},
|
||||
{"__interrupt__": ()},
|
||||
]
|
||||
|
||||
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values=[
|
||||
_AnyIdHumanMessage(content="what is weather in sf"),
|
||||
AIMessage(
|
||||
content="",
|
||||
id="ai1",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call123",
|
||||
"name": "search_api",
|
||||
"args": {"query": "a different query"},
|
||||
}
|
||||
],
|
||||
),
|
||||
_AnyIdToolMessage(
|
||||
content="result for a different query",
|
||||
name="search_api",
|
||||
tool_call_id="tool_call123",
|
||||
),
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call456",
|
||||
"name": "search_api",
|
||||
"args": {"query": "another"},
|
||||
}
|
||||
],
|
||||
id="ai2",
|
||||
),
|
||||
],
|
||||
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
|
||||
next=("tools",),
|
||||
config=tup.config,
|
||||
created_at=tup.checkpoint["ts"],
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"step": 4,
|
||||
"thread_id": "1",
|
||||
},
|
||||
parent_config=(
|
||||
[c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][
|
||||
-1
|
||||
].config
|
||||
),
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
await app_w_interrupt.aupdate_state(
|
||||
config,
|
||||
AIMessage(content="answer", id="ai2"),
|
||||
)
|
||||
|
||||
# replaces message even if object identity is different, as long as id is the same
|
||||
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values=[
|
||||
_AnyIdHumanMessage(content="what is weather in sf"),
|
||||
AIMessage(
|
||||
content="",
|
||||
id="ai1",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call123",
|
||||
"name": "search_api",
|
||||
"args": {"query": "a different query"},
|
||||
}
|
||||
],
|
||||
),
|
||||
_AnyIdToolMessage(
|
||||
content="result for a different query",
|
||||
name="search_api",
|
||||
tool_call_id="tool_call123",
|
||||
),
|
||||
AIMessage(content="answer", id="ai2"),
|
||||
],
|
||||
tasks=(),
|
||||
next=(),
|
||||
config=tup.config,
|
||||
created_at=tup.checkpoint["ts"],
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "update",
|
||||
"step": 5,
|
||||
"thread_id": "1",
|
||||
},
|
||||
parent_config=(
|
||||
[c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][
|
||||
-1
|
||||
].config
|
||||
),
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
|
||||
async def test_in_one_fan_out_out_one_graph_state() -> None:
|
||||
def sorted_add(x: list[str], y: list[str]) -> list[str]:
|
||||
return sorted(operator.add(x, y))
|
||||
|
||||
class State(TypedDict, total=False):
|
||||
query: str
|
||||
answer: str
|
||||
@@ -3028,44 +3440,6 @@ async def test_doubly_nested_graph_state(
|
||||
c async for c in app.aget_state_history(outer_history[1].tasks[0].state)
|
||||
]
|
||||
assert child_history == [
|
||||
StateSnapshot(
|
||||
values={"my_key": "hi my value here and there"},
|
||||
next=(),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("child:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{"": AnyStr(), AnyStr("child:"): AnyStr()}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"source": "loop",
|
||||
"step": 1,
|
||||
"parents": {"": AnyStr()},
|
||||
"thread_id": "1",
|
||||
"langgraph_node": "child",
|
||||
"langgraph_path": [PULL, AnyStr("child")],
|
||||
"langgraph_step": 2,
|
||||
"langgraph_triggers": ["branch:to:child"],
|
||||
"langgraph_checkpoint_ns": AnyStr("child:"),
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("child:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{"": AnyStr(), AnyStr("child:"): AnyStr()}
|
||||
),
|
||||
}
|
||||
},
|
||||
tasks=(),
|
||||
interrupts=(),
|
||||
),
|
||||
StateSnapshot(
|
||||
values={"my_key": "hi my value"},
|
||||
next=("child_1",),
|
||||
@@ -3111,65 +3485,9 @@ async def test_doubly_nested_graph_state(
|
||||
]
|
||||
# get grandchild graph history
|
||||
grandchild_history = [
|
||||
c async for c in app.aget_state_history(child_history[1].tasks[0].state)
|
||||
c async for c in app.aget_state_history(child_history[0].tasks[0].state)
|
||||
]
|
||||
assert grandchild_history == [
|
||||
StateSnapshot(
|
||||
values={"my_key": "hi my value here and there"},
|
||||
next=(),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr(),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{
|
||||
"": AnyStr(),
|
||||
AnyStr("child:"): AnyStr(),
|
||||
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
|
||||
}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"source": "loop",
|
||||
"step": 2,
|
||||
"parents": AnyDict(
|
||||
{
|
||||
"": AnyStr(),
|
||||
AnyStr("child:"): AnyStr(),
|
||||
}
|
||||
),
|
||||
"thread_id": "1",
|
||||
"langgraph_checkpoint_ns": AnyStr("child:"),
|
||||
"langgraph_node": "child_1",
|
||||
"langgraph_path": [
|
||||
PULL,
|
||||
AnyStr("child_1"),
|
||||
],
|
||||
"langgraph_step": 1,
|
||||
"langgraph_triggers": [
|
||||
"branch:to:child_1",
|
||||
],
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr(),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{
|
||||
"": AnyStr(),
|
||||
AnyStr("child:"): AnyStr(),
|
||||
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
|
||||
}
|
||||
),
|
||||
}
|
||||
},
|
||||
tasks=(),
|
||||
interrupts=(),
|
||||
),
|
||||
StateSnapshot(
|
||||
values={"my_key": "hi my value here"},
|
||||
next=("grandchild_2",),
|
||||
@@ -3239,7 +3557,7 @@ async def test_send_to_nested_graphs(async_checkpointer: BaseCheckpointSaver) ->
|
||||
return {"subject": f"{subject} - hohoho"}
|
||||
|
||||
# subgraph
|
||||
subgraph = StateGraph(JokeState, output=OverallState)
|
||||
subgraph = StateGraph(JokeState, output_schema=OverallState)
|
||||
subgraph.add_node("edit", edit)
|
||||
subgraph.add_node(
|
||||
"generate", lambda state: {"jokes": [f"Joke about {state['subject']}"]}
|
||||
|
||||
@@ -46,7 +46,7 @@ from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START
|
||||
from langgraph.errors import InvalidUpdateError, ParentCommand
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import END, StateGraph
|
||||
from langgraph.graph.message import MessagesState, add_messages
|
||||
from langgraph.graph.message import MessageGraph, MessagesState, add_messages
|
||||
from langgraph.prebuilt.tool_node import ToolNode
|
||||
from langgraph.pregel import (
|
||||
GraphRecursionError,
|
||||
@@ -159,7 +159,7 @@ def test_checkpoint_errors() -> None:
|
||||
raise ValueError("Faulty put_writes")
|
||||
|
||||
class FaultyVersionCheckpointer(InMemorySaver):
|
||||
def get_next_version(self, current: Optional[int]) -> int:
|
||||
def get_next_version(self, current: Optional[int], channel: None) -> int:
|
||||
raise ValueError("Faulty get_next_version")
|
||||
|
||||
def logic(inp: str) -> str:
|
||||
@@ -3276,6 +3276,57 @@ def test_subgraph_checkpoint_true(
|
||||
),
|
||||
]
|
||||
|
||||
checkpoints = list(app.get_state_history(config))
|
||||
if checkpoint_during:
|
||||
assert len(checkpoints) == 4
|
||||
else:
|
||||
assert len(checkpoints) == 1
|
||||
|
||||
|
||||
def test_subgraph_checkpoint_during_false_inherited() -> None:
|
||||
sync_checkpointer = InMemorySaver()
|
||||
|
||||
class InnerState(TypedDict):
|
||||
my_key: Annotated[str, operator.add]
|
||||
my_other_key: str
|
||||
|
||||
def inner_1(state: InnerState):
|
||||
return {"my_key": " got here", "my_other_key": state["my_key"]}
|
||||
|
||||
def inner_2(state: InnerState):
|
||||
return {"my_key": " and there"}
|
||||
|
||||
inner = StateGraph(InnerState)
|
||||
inner.add_node("inner_1", inner_1)
|
||||
inner.add_node("inner_2", inner_2)
|
||||
inner.add_edge("inner_1", "inner_2")
|
||||
inner.set_entry_point("inner_1")
|
||||
inner.set_finish_point("inner_2")
|
||||
|
||||
class State(TypedDict):
|
||||
my_key: str
|
||||
|
||||
inner_app = inner.compile(checkpointer=sync_checkpointer)
|
||||
graph = StateGraph(State)
|
||||
graph.add_node("inner", inner_app)
|
||||
graph.add_edge(START, "inner")
|
||||
graph.add_conditional_edges(
|
||||
"inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END
|
||||
)
|
||||
app = graph.compile(checkpointer=sync_checkpointer)
|
||||
for checkpoint_during in [True, False]:
|
||||
thread_id = str(uuid.uuid4())
|
||||
config = {"configurable": {"thread_id": thread_id}}
|
||||
app.invoke(
|
||||
{"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during
|
||||
)
|
||||
if checkpoint_during:
|
||||
checkpoints = list(sync_checkpointer.list(config))
|
||||
assert len(checkpoints) == 12
|
||||
else:
|
||||
checkpoints = list(sync_checkpointer.list(config))
|
||||
assert len(checkpoints) == 1
|
||||
|
||||
|
||||
def test_subgraph_checkpoint_true_interrupt(
|
||||
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
|
||||
@@ -3943,14 +3994,9 @@ def test_checkpoint_metadata(sync_checkpointer: BaseCheckpointSaver) -> None:
|
||||
def test_remove_message_via_state_update(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
from langchain_core.messages import (
|
||||
AIMessage,
|
||||
AnyMessage,
|
||||
HumanMessage,
|
||||
RemoveMessage,
|
||||
)
|
||||
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
|
||||
|
||||
workflow = StateGraph(Annotated[list[AnyMessage], add_messages])
|
||||
workflow = MessageGraph()
|
||||
workflow.add_node(
|
||||
"chatbot",
|
||||
lambda state: [
|
||||
@@ -3981,14 +4027,9 @@ def test_remove_message_via_state_update(
|
||||
|
||||
|
||||
def test_remove_message_from_node():
|
||||
from langchain_core.messages import (
|
||||
AIMessage,
|
||||
AnyMessage,
|
||||
HumanMessage,
|
||||
RemoveMessage,
|
||||
)
|
||||
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
|
||||
|
||||
workflow = StateGraph(Annotated[list[AnyMessage], add_messages])
|
||||
workflow = MessageGraph()
|
||||
workflow.add_node(
|
||||
"chatbot",
|
||||
lambda state: [
|
||||
@@ -4300,7 +4341,7 @@ def test_store_injected(
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("node", Node())
|
||||
builder.add_edge("__start__", "node")
|
||||
N = 500
|
||||
N = 50
|
||||
M = 1
|
||||
|
||||
for i in range(N):
|
||||
@@ -4575,11 +4616,14 @@ def test_debug_nested_subgraphs(
|
||||
|
||||
return clean_config
|
||||
|
||||
for checkpoint_events, checkpoint_history in zip(
|
||||
stream_ns.values(), history_ns.values()
|
||||
for checkpoint_events, checkpoint_history, ns in zip(
|
||||
stream_ns.values(), history_ns.values(), stream_ns.keys()
|
||||
):
|
||||
if not checkpoint_during:
|
||||
checkpoint_events = checkpoint_events[-1:]
|
||||
if ns: # Save no checkpoints for subgraphs when checkpoint_during=False
|
||||
assert not checkpoint_history
|
||||
continue
|
||||
assert len(checkpoint_events) == len(checkpoint_history)
|
||||
for stream, history in zip(checkpoint_events, checkpoint_history):
|
||||
assert stream["values"] == history.values
|
||||
|
||||
@@ -103,7 +103,7 @@ async def test_checkpoint_errors() -> None:
|
||||
raise ValueError("Faulty put_writes")
|
||||
|
||||
class FaultyVersionCheckpointer(InMemorySaver):
|
||||
def get_next_version(self, current: Optional[int]) -> int:
|
||||
def get_next_version(self, current: Optional[int], channel: None) -> int:
|
||||
raise ValueError("Faulty get_next_version")
|
||||
|
||||
def logic(inp: str) -> str:
|
||||
@@ -1338,7 +1338,7 @@ async def test_node_schemas_custom_output() -> None:
|
||||
"now": 123,
|
||||
}
|
||||
|
||||
builder = StateGraph(State, output=Output)
|
||||
builder = StateGraph(State, output_schema=Output)
|
||||
builder.add_node("a", node_a)
|
||||
builder.add_node("b", node_b)
|
||||
builder.add_node("c", node_c)
|
||||
@@ -1353,7 +1353,7 @@ async def test_node_schemas_custom_output() -> None:
|
||||
"messages": [_AnyIdHumanMessage(content="hello")],
|
||||
}
|
||||
|
||||
builder = StateGraph(State, output=Output)
|
||||
builder = StateGraph(State, output_schema=Output)
|
||||
builder.add_node("a", node_a)
|
||||
builder.add_node("b", node_b)
|
||||
builder.add_node("c", node_c)
|
||||
@@ -5029,6 +5029,51 @@ async def test_subgraph_checkpoint_true(
|
||||
]
|
||||
|
||||
|
||||
async def test_subgraph_checkpoint_during_false_inherited() -> None:
|
||||
async_checkpointer = InMemorySaver()
|
||||
|
||||
class InnerState(TypedDict):
|
||||
my_key: Annotated[str, operator.add]
|
||||
my_other_key: str
|
||||
|
||||
def inner_1(state: InnerState):
|
||||
return {"my_key": " got here", "my_other_key": state["my_key"]}
|
||||
|
||||
def inner_2(state: InnerState):
|
||||
return {"my_key": " and there"}
|
||||
|
||||
inner = StateGraph(InnerState)
|
||||
inner.add_node("inner_1", inner_1)
|
||||
inner.add_node("inner_2", inner_2)
|
||||
inner.add_edge("inner_1", "inner_2")
|
||||
inner.set_entry_point("inner_1")
|
||||
inner.set_finish_point("inner_2")
|
||||
|
||||
class State(TypedDict):
|
||||
my_key: str
|
||||
|
||||
inner_app = inner.compile(checkpointer=async_checkpointer)
|
||||
graph = StateGraph(State)
|
||||
graph.add_node("inner", inner_app)
|
||||
graph.add_edge(START, "inner")
|
||||
graph.add_conditional_edges(
|
||||
"inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END
|
||||
)
|
||||
app = graph.compile(checkpointer=async_checkpointer)
|
||||
for checkpoint_during in [True, False]:
|
||||
thread_id = str(uuid.uuid4())
|
||||
config = {"configurable": {"thread_id": thread_id}}
|
||||
await app.ainvoke(
|
||||
{"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during
|
||||
)
|
||||
if checkpoint_during:
|
||||
checkpoints = list(async_checkpointer.list(config))
|
||||
assert len(checkpoints) == 12
|
||||
else:
|
||||
checkpoints = list(async_checkpointer.list(config))
|
||||
assert len(checkpoints) == 1
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_subgraph_checkpoint_true_interrupt(
|
||||
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
|
||||
@@ -5736,7 +5781,7 @@ async def test_store_injected_async(
|
||||
builder.add_edge("__start__", "node")
|
||||
builder.add_edge("node", "other_node")
|
||||
|
||||
N = 500
|
||||
N = 50
|
||||
M = 1
|
||||
|
||||
for i in range(N):
|
||||
@@ -6007,11 +6052,14 @@ async def test_debug_nested_subgraphs(
|
||||
|
||||
return clean_config
|
||||
|
||||
for checkpoint_events, checkpoint_history in zip(
|
||||
stream_ns.values(), history_ns.values()
|
||||
for checkpoint_events, checkpoint_history, ns in zip(
|
||||
stream_ns.values(), history_ns.values(), stream_ns.keys()
|
||||
):
|
||||
if not checkpoint_during:
|
||||
checkpoint_events = checkpoint_events[-1:]
|
||||
if ns: # Save no checkpoints for subgraphs when checkpoint_during=False
|
||||
assert not checkpoint_history
|
||||
continue
|
||||
assert len(checkpoint_events) == len(checkpoint_history)
|
||||
for stream, history in zip(checkpoint_events, checkpoint_history):
|
||||
assert stream["values"] == history.values
|
||||
@@ -6982,14 +7030,17 @@ async def test_multiple_subgraphs(async_checkpointer: BaseCheckpointSaver) -> No
|
||||
return {"result": state["a"] + state["b"]}
|
||||
|
||||
add_subgraph = (
|
||||
StateGraph(State, output=Output).add_node(add).add_edge(START, "add").compile()
|
||||
StateGraph(State, output_schema=Output)
|
||||
.add_node(add)
|
||||
.add_edge(START, "add")
|
||||
.compile()
|
||||
)
|
||||
|
||||
async def multiply(state):
|
||||
return {"result": state["a"] * state["b"]}
|
||||
|
||||
multiply_subgraph = (
|
||||
StateGraph(State, output=Output)
|
||||
StateGraph(State, output_schema=Output)
|
||||
.add_node(multiply)
|
||||
.add_edge(START, "multiply")
|
||||
.compile()
|
||||
@@ -7002,7 +7053,7 @@ async def test_multiple_subgraphs(async_checkpointer: BaseCheckpointSaver) -> No
|
||||
return another_result
|
||||
|
||||
parent_call_same_subgraph = (
|
||||
StateGraph(State, output=Output)
|
||||
StateGraph(State, output_schema=Output)
|
||||
.add_node(call_same_subgraph)
|
||||
.add_edge(START, "call_same_subgraph")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
@@ -7026,7 +7077,7 @@ async def test_multiple_subgraphs(async_checkpointer: BaseCheckpointSaver) -> No
|
||||
}
|
||||
|
||||
parent_call_multiple_subgraphs = (
|
||||
StateGraph(State, output=Output)
|
||||
StateGraph(State, output_schema=Output)
|
||||
.add_node(call_multiple_subgraphs)
|
||||
.add_edge(START, "call_multiple_subgraphs")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
@@ -7104,14 +7155,17 @@ async def test_multiple_subgraphs_mixed_entrypoint(
|
||||
return {"result": state["a"] + state["b"]}
|
||||
|
||||
add_subgraph = (
|
||||
StateGraph(State, output=Output).add_node(add).add_edge(START, "add").compile()
|
||||
StateGraph(State, output_schema=Output)
|
||||
.add_node(add)
|
||||
.add_edge(START, "add")
|
||||
.compile()
|
||||
)
|
||||
|
||||
async def multiply(state):
|
||||
return {"result": state["a"] * state["b"]}
|
||||
|
||||
multiply_subgraph = (
|
||||
StateGraph(State, output=Output)
|
||||
StateGraph(State, output_schema=Output)
|
||||
.add_node(multiply)
|
||||
.add_edge(START, "multiply")
|
||||
.compile()
|
||||
@@ -7181,7 +7235,7 @@ async def test_multiple_subgraphs_mixed_state_graph(
|
||||
return {"result": another_result}
|
||||
|
||||
parent_call_same_subgraph = (
|
||||
StateGraph(State, output=Output)
|
||||
StateGraph(State, output_schema=Output)
|
||||
.add_node(call_same_subgraph)
|
||||
.add_edge(START, "call_same_subgraph")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
@@ -7205,7 +7259,7 @@ async def test_multiple_subgraphs_mixed_state_graph(
|
||||
}
|
||||
|
||||
parent_call_multiple_subgraphs = (
|
||||
StateGraph(State, output=Output)
|
||||
StateGraph(State, output_schema=Output)
|
||||
.add_node(call_multiple_subgraphs)
|
||||
.add_edge(START, "call_multiple_subgraphs")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
|
||||
@@ -92,7 +92,7 @@ def test_state_schema_with_type_hint():
|
||||
assert state.pop("foo") == "bar"
|
||||
return {"input_state": state}
|
||||
|
||||
graph = StateGraph(InputState, output=OutputState)
|
||||
graph = StateGraph(InputState, output_schema=OutputState)
|
||||
actions = [
|
||||
complete_hint,
|
||||
miss_first_hint,
|
||||
|
||||
Generated
+1592
-1589
File diff suppressed because it is too large
Load Diff
@@ -591,7 +591,7 @@ def create_react_agent(
|
||||
workflow = StateGraph(state_schema, config_schema=config_schema)
|
||||
workflow.add_node(
|
||||
"agent",
|
||||
RunnableCallable(call_model, acall_model), # type: ignore[call-overload]
|
||||
RunnableCallable(call_model, acall_model),
|
||||
input_schema=input_schema,
|
||||
)
|
||||
if pre_model_hook is not None:
|
||||
@@ -610,7 +610,7 @@ def create_react_agent(
|
||||
if response_format is not None:
|
||||
workflow.add_node(
|
||||
"generate_structured_response",
|
||||
RunnableCallable( # type: ignore[call-overload]
|
||||
RunnableCallable(
|
||||
generate_structured_response,
|
||||
agenerate_structured_response,
|
||||
),
|
||||
@@ -660,10 +660,10 @@ def create_react_agent(
|
||||
# Define the two nodes we will cycle between
|
||||
workflow.add_node(
|
||||
"agent",
|
||||
RunnableCallable(call_model, acall_model), # type: ignore[call-overload]
|
||||
RunnableCallable(call_model, acall_model),
|
||||
input_schema=input_schema,
|
||||
)
|
||||
workflow.add_node("tools", tool_node) # type: ignore[call-overload]
|
||||
workflow.add_node("tools", tool_node)
|
||||
|
||||
# Optionally add a pre-model hook node that will be called
|
||||
# every time before the "agent" (LLM-calling node)
|
||||
@@ -693,7 +693,7 @@ def create_react_agent(
|
||||
if response_format is not None:
|
||||
workflow.add_node(
|
||||
"generate_structured_response",
|
||||
RunnableCallable( # type: ignore[call-overload]
|
||||
RunnableCallable(
|
||||
generate_structured_response,
|
||||
agenerate_structured_response,
|
||||
),
|
||||
|
||||
@@ -629,7 +629,7 @@ def tools_condition(
|
||||
|
||||
Args:
|
||||
state: The state to check for
|
||||
tool calls. Must have a list of messages or have the
|
||||
tool calls. Must have a list of messages (MessageGraph) or have the
|
||||
"messages" key (StateGraph).
|
||||
|
||||
Returns:
|
||||
@@ -850,7 +850,7 @@ def _get_store_arg(tool: BaseTool) -> Optional[str]:
|
||||
if _is_injection(type_arg, InjectedStore)
|
||||
]
|
||||
if len(injections) > 1:
|
||||
ValueError(
|
||||
raise ValueError(
|
||||
"A tool argument should not be annotated with InjectedStore more than "
|
||||
f"once. Received arg {name} with annotations {injections}."
|
||||
)
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user