Compare commits

...
11 Commits
Author SHA1 Message Date
Elior Nataf Lackritz 8fd97561d6 perf(checkpoint-postgres): skip the async saver's lock for connection pools (#9638)
Fixes #7259

`AsyncPostgresSaver` takes its instance lock around every database call.
The lock is needed when the saver shares one connection, which serves
one call at a time, but the saver also took it with an
`AsyncConnectionPool`, which gives each call its own connection. So
every saver call in the process waited for the one before it: concurrent
runs, and within one run the background checkpoint saves and each task's
`put_writes`.

With a pool, `_cursor` no longer takes the lock. A single connection and
a shared `pipe` still do.

Python 3.14 on an Apple M5 Pro, Postgres 16,
`AsyncConnectionPool(max_size=10)`, median of 5 rounds that alternate
main and this branch (3 rounds with 1 ms round trips):

| workload | main | this PR | speedup |
|---|---|---|---|
| 8 concurrent one-turn agents on one saver | 87.1 ms | 49.8 ms | 1.75x
|
| 500-way `Send` fan-out | 246 ms | 144 ms | 1.71x |
| 200 sequential nodes | 265 ms | 190 ms | 1.39x |
| one agent turn (7 supersteps) | 12.3 ms | 10.4 ms | 1.18x |
| 8 concurrent agents, 1 ms round trips | 234 ms | 60 ms | 3.86x |
| 500-way fan-out, 1 ms round trips | 773 ms | 167 ms | 4.63x |
| 200 sequential nodes, 1 ms round trips | 696 ms | 346 ms | 2.01x |

Process CPU stays about the same. A burst of calls larger than the pool
now waits in `psycopg_pool` (30 s timeout by default) instead of on the
lock. Calls that must be ordered already are: the loop saves a
checkpoint after the previous one and after its DeltaChannel writes.

The sync `PostgresSaver` takes the same lock with a `ConnectionPool`.
Removing it there let up to 19 executor threads contend for the GIL and
the pool: 1.3x to 1.5x faster, but 1.2x to 2.4x more CPU, so this PR
leaves it. `AsyncShallowPostgresSaver` is deprecated and also keeps its
lock.

@glaziermag reported this in #7259 with a load test, and @jitokim
proposed the same fix in #7269.

## Test plan
- [x] New `test_pool_saver_reads_while_a_list_is_open`: with a pool,
`aget_tuple` completes while an `alist` iteration is open. On main it
waits for the lock and times out after 5 s.
- [x] `checkpoint-postgres` suite (427 passed, 3 skipped), the
checkpoint conformance suite on a pool, and langgraph's Postgres tests
against this saver (2,832 passed).
2026-11-13 00:11:29 +00:00
Elior Nataf Lackritz ba0acae562 perf(langgraph): skip message conversion in add_messages for plain messages (#9631)
`add_messages` ran `convert_to_messages` and `message_chunk_to_message`
over every message already in the state on every update, plus a final
filter pass. Both conversions return a plain message unchanged, so on a
long history this is overhead: 0.18 ms per call at 1,000 messages. An
agent turn calls it once per update and once more for each conditional
edge that reads `messages`.

The list now passes through as is when every item is already a plain
message, checked on the few distinct types in the list rather than per
item, and the filter runs only when a message was removed. Anything else
(dicts, tuples, strings, chunks, or a mix) goes through the same
conversion as before.

Median of alternating rounds on main and on this branch, Python 3.14, an
agent turn being 3 tool calls on a thread with 1,000 messages:

| workload | main | this PR | speedup |
|---|---|---|---|
| `add_messages` on 1,000 messages plus 1 new one | 184 us | 69 us |
2.7x |
| agent turn, no checkpointer | 4.80 ms | 3.38 ms | 1.42x |
| agent turn, `InMemorySaver`, `durability="exit"` | 11.0 ms | 8.7 ms |
1.25x |
| agent turn, `InMemorySaver`, default durability | 22.4 ms | 21.1 ms |
1.06x |

With a checkpoint every step, saving the whole history at each step
dominates the turn, so the gain is small there.

@jkennedyvz trimmed the same per-message conversion loop in #6969.

## Test plan
- [x] Missing ids are still assigned and the result is still a new list.
Two new tests in `test_messages_state.py` cover message-likes and chunks
on both sides, and that the result is not the input list. They pass on
main too.
2026-11-13 00:10:04 +00:00
Elior Nataf Lackritz cd95887ec5 perf(checkpoint): stop InMemoryStore search at the requested page (#9633)
`InMemoryStore.search` built an `(item, vectors)` pair for every item in
every matching namespace, then kept `candidates[offset:offset + limit]`.
A `limit=10` search on a 10,000-item namespace took 1.9 ms, nearly all
of it building candidates the slice then dropped. Paging a namespace 100
items at a time, as deepagents' `StoreBackend` does for `ls`, `glob` and
`grep`, rebuilt the whole listing for every page.

Without a filter, and without a query the store can embed, the result is
exactly that slice, so `_filter_items` now stops once it has `offset +
limit` matches. The namespace prefix test also becomes one tuple
compare: a namespace shorter than the prefix slices to itself, so it can
never equal the prefix.

Median of four alternating rounds on main and on this branch, Python
3.14:

| workload | main | this PR | speedup |
|---|---|---|---|
| `search(ns, limit=10)` on a 1,000-item namespace | 70 µs | 4 µs | 17x
|
| same, 10,000 items | 1.92 ms | 4 µs | about 490x |
| graph node that searches 1,000 memories (`limit=5`) and adds one, 20
steps | 3.41 ms | 1.86 ms | 1.83x |
| same, 10,000 memories | 39.4 ms | 1.83 ms | 21.5x |
| paging 5,000 items 100 at a time | 47.1 ms | 30.1 ms | 1.57x |
| `search` over 10,000 one-item namespaces | 23.4 ms | 15.5 ms | 1.51x |
| `search` with a filter, `get`, `put` | | | about 1x |

A filter keeps the full scan, so a filter that raises on some item (a
range operator on a missing key) still raises the same way. Negative
`offset` or `limit` keep it too, since a slice counts those from the
end.

## Test plan
- [x] `test_search_pages_match_full_listing` pages plain and filtered
searches across nested and look-alike namespaces and compares each page
with the full listing. It passes on main too, since the results are
unchanged.
2026-11-13 00:08:22 +00:00
Elior Nataf Lackritz 96686e5700 perf(checkpoint-postgres): load no pending writes without a thread hop (#9639)
`AsyncPostgresSaver` deserializes a checkpoint's pending writes in
`asyncio.to_thread`. The hand-off costs about 50 us per checkpoint, and
it ran even when there was nothing to load: the usual case for the
latest checkpoint of a finished run, which `aget_state` and the start of
every run read.

It now returns an empty list without the hop when the row has no pending
writes.

Python 3.14 on an Apple M5 Pro, Postgres 16, median of 5 rounds that
alternate main and this branch:

| workload | main | this PR | speedup |
|---|---|---|---|
| `aget_tuple` of a checkpoint with no pending writes | 264 us | 214 us
| 1.24x |
| `aget_state`, 200-message thread | 2.46 ms | 2.13 ms | 1.16x |
| an agent turn on that thread | 23.8 ms | 23.7 ms | about 1x |

With 1 ms round trips the round trip dominates each call, so wall time
stays about 1x, while process CPU per `aget_tuple` is still 1.36x lower.

`_load_writes(None)` already returned an empty list, so the result is
the same.

## Test plan
- [x] `checkpoint-postgres` suite (426 passed, 3 skipped) and the
checkpoint conformance suite (`get_tuple`, `list`).
2026-11-12 18:55:18 -05:00
Elior Nataf Lackritz acff17b041 perf(checkpoint): dedupe batched store reads with a dict (#9636)
`AsyncBatchedBaseStore` collects the store calls made in the same event
loop tick into one batch and dedupes them first. `_dedupe_ops` found an
earlier equal read with `dedupped.index(op)`, a linear scan per read, so
a batch of n reads cost O(n²) comparisons: 1.1 ms for 500 concurrent
`aget` calls, 16 ms for 2,000. A fan-out where every task reads the
store fills exactly such a batch.

Reads are now looked up in a dict keyed by the op. An op that can't be
hashed (a `SearchOp` with a filter dict) falls back to the old scan, so
it dedupes exactly as before.

Median of four alternating rounds on main and on this branch, Python
3.14:

| workload | main | this PR | speedup |
|---|---|---|---|
| `_dedupe_ops` on 500 gets | 1.13 ms | 31 µs | 36x |
| 2,000 concurrent `aget`, `AsyncSqliteStore` | 17.6 ms | 8.0 ms | 2.2x
|
| 2,000 concurrent `aget`, `AsyncPostgresStore` | 20.2 ms | 11.3 ms |
1.78x |
| 500 concurrent `aget`, either store | | | 1.25x |
| 500-way `Send` fan-out where each task reads and writes the store | |
| about 1x |

## Test plan
- [x] `test_dedupe_ops_maps_repeated_reads_and_puts_to_one_op` covers
repeated gets, filtered searches, list ops and overwritten puts. It
passes on main too, since the result is unchanged.
2026-11-12 18:55:11 -05:00
Elior Nataf Lackritz 46b2ef0e1e perf(langgraph): filter pending writes to the given tasks before reading their statuses (#9634)
Each functional API task call (`accept_push`) restores any output the
pushed task saved in an earlier run. It did that by running
`read_task_statuses` over every pending write, which builds the status
of every task that has written, and then looked up the one task. On
resume every earlier task's writes are pending, so an entrypoint that
calls n tasks paid for n full status reads, O(n²): 196 µs per task call
at 1,000 pending writes. A sync entrypoint running tasks in parallel
hits it too, since writes land while tasks are still being pushed.

`_reapply_writes_to_succeeded_nodes` now hands `read_task_statuses` only
the writes of the tasks it was given. Statuses are per task, so those
tasks get exactly the statuses they got before.

Median of four alternating rounds on main and on this branch, Python
3.14, `InMemorySaver`:

| workload | main | this PR | speedup |
|---|---|---|---|
| matching saved writes for one task call, 1,000 pending writes | 196 µs
| 9.8 µs | 20x |
| resume an entrypoint that ran 500 tasks and then called `interrupt()`
| 153 ms | 35.4 ms | 4.3x |
| same, 100 tasks | 9.8 ms | 6.1 ms | 1.6x |
| entrypoint running 500 tasks in parallel, sync | 134 ms | 90.4 ms |
1.48x |
| 100 tasks in parallel or one after another, sync or async | | | about
1x |

Async entrypoints push all their tasks before any write lands, so they
barely move. In the 500-task parallel run, the biggest remaining cost in
the profile is the runner's stop check, which #9630 addresses.

## Test plan
- [x] Existing coverage: `test_imp_task` resumes an entrypoint and
checks its tasks don't run again, and the full suite passes.
2026-11-12 18:55:04 -05:00
Elior Nataf Lackritz 496885e569 perf(langgraph): match update keys against a set on wide states (#9632)
`_get_updates` (run on every node's return value) and
`map_output_updates` (run for every task in `stream_mode="updates"`)
checked each written key against a list of all state keys, and
`map_output_updates` also looped over every key to look for duplicate
writes. On a state with hundreds of keys that is a list scan for every
written key.

Both now check against a frozenset: `_get_updates` builds it once when
the node is attached, `map_output_updates` once per call. The updates
keep the order of the node's writes.

Median of four alternating rounds on main and on this branch, Python
3.14:

| workload | main | this PR | speedup |
|---|---|---|---|
| `map_output_updates`, 20 writes on a 300-key state | 23.3 us | 4.0 us
| 5.8x |
| 30 nodes writing 20 keys each on a 300-key state, sync | 10.7 ms | 9.6
ms | 1.12x |
| same, async | 12.9 ms | 11.7 ms | 1.10x |
| `bench` `wide_state_25x300` (17 keys), sync | 3.84 ms | 3.74 ms |
about 1x |

States with a handful of keys see no change. Compiling builds one extra
frozenset per node, about 0.5 us each (`pydantic_state_25x300` compiles
in 117 us instead of 114 us).

## Test plan
- [x] The same keys pass the check, so no new test. Existing tests cover
the update paths, including a node that writes the same key twice in
`updates` mode (`test_multiple_writes_same_channel_from_same_node`).
2026-11-12 18:54:57 -05:00
Elior Nataf Lackritz 23a77c57fd perf(langgraph): check only failed futures for the runner stop condition (#9630)
When a task in a superstep finished, `FuturesDict.on_done` ran the stop
check over every future that had finished so far. The check calls
`cancelled()` and `exception()` on each one, and both take the future's
lock for a `concurrent.futures.Future`. A step with n parallel tasks
paid about n²/2 of these calls, about 125,000 for a 500-task `Send`
fan-out.

Only a future that raised can meet the stop condition, and a finished
future never changes, so `FuturesDict` now keeps the failed futures in
their own set and checks only those. The single-task path, which adds a
failed future straight to `done`, adds it to the new set too, so the
check sees the same failures as before, including ones an error handler
takes over.

Median of four alternating rounds on main and on this branch, Python
3.14:

| workload | main | this PR | speedup |
|---|---|---|---|
| 500 futures finishing in one step, `FuturesDict` alone | 28.5 ms |
0.72 ms | 39x |
| 500-way `Send` fan-out, sync, `InMemorySaver` | 77.3 ms | 33.9 ms |
2.28x |
| 500-way `Send` fan-out, sync, no checkpointer | 65.2 ms | 28.1 ms |
2.32x |
| 500-way `Send` fan-out, async | 36.4 ms | 35.8 ms | about 1x |
| `bench` `fanout_to_subgraph_100x`, sync | 116 ms | 113 ms | about 1x |

Async barely moves because asyncio futures answer these calls without a
lock, and the 100-way benchmark spends its time inside the subgraphs.

@AaronG18 (#8860), @nankingjing (#8347) and @014-code (#8270) proposed
checking fewer futures here, and @HumphreySun98 reported the problem in
#8240.

## Test plan
- [x] `test_runner_futures.py` pins the stop decision: a failed task
wakes the waiter before the other tasks finish, and a failure an error
handler took over does not. Both tests pass on main too, since the
behavior is unchanged.
2026-11-12 18:54:50 -05:00
Elior Nataf Lackritz 40f5cc7f1f perf(langgraph): find untracked and delta channels once per run (#9629)
The Pregel loop looked up channel types in two hot places. `put_writes`
scanned every channel for an `UntrackedValue` on each task's writes, and
`_put_checkpoint` walked every channel twice per superstep to find the
`DeltaChannel`s. A graph has a channel per node plus one per state key,
and channels are ABCs, so each check goes through
`ABCMeta.__instancecheck__`: a 200-node graph paid about 630 of them per
step.

Channel types can't change during a run, so the loop now works out both
answers once from `specs` and reuses them.

Median of four alternating rounds on main and on this branch, Python
3.14:

| workload | main | this PR | speedup |
|---|---|---|---|
| channel-type checks per step, 211 channels | 38 us | under 0.1 us | |
| `bench` `sequential_1000`, sync | 219 ms | 93 ms | 2.36x |
| `bench` `sequential_1000`, async | 370 ms | 250 ms | 1.48x |
| 200 sequential nodes, `InMemorySaver` | 38.7 ms | 29.4 ms | 1.31x |
| 200 sequential nodes, no checkpointer | 19.4 ms | 12.9 ms | 1.50x |
| agent turn on a model node + `ToolNode` graph | 2.51 ms | 2.48 ms |
about 1x |

The gain grows with the number of nodes and state keys. Small graphs see
no change.

`self.channels` is built from the same `specs`, with the same keys in
the same order, so the same writes are dropped, the same `Send` packets
are sanitized, and the same delta counters and snapshots come out.

@AaronG18 proposed the `UntrackedValue` half of this in #8858, and
@jkennedyvz in #7121.

## Test plan
- [x] No behavior change, so no new test. `test_untracked_value`,
`test_send_with_untracked_value` and the `test_delta_channel_*` files
cover these paths.
2026-11-12 18:54:43 -05:00
Elior Nataf Lackritz ef6490a521 perf(prebuilt): only format the call_model error message when raising (#9628)
Fixes #9328

`_get_model_input_state` in `create_react_agent` built its
missing-messages error string before checking whether it needed it. The
f-string includes `{state}`, so every model call ran `repr()` over the
whole message history and threw the result away. The cost grows with the
history, so a long tool loop pays it on every step.

The message is now built only when the error is raised, with the same
text.

Median of four alternating rounds on main and on this branch, Python
3.14:

| workload | main | this PR | speedup |
|---|---|---|---|
| building the model input, 1,000 messages | 7.6 ms | 0.08 ms | 95x |
| one agent turn (3 tool calls) on a 1,000-message thread,
`InMemorySaver` | 60.7 ms | 29.6 ms | 2.05x |
| same turn on a 200-message thread | 14.1 ms | 8.1 ms | 1.75x |
| `bench` `react_agent_100x`, sync | 319 ms | 53 ms | 6.0x |
| `bench` `react_agent_100x`, async | 390 ms | 126 ms | 3.1x |

The saving is the cost of that `repr()`, so it grows with the size of
the history.

Thanks @deadstrobe5, who reported this in #9328 and proposed the same
change in #9329.

## Test plan
- [x] Behavior is unchanged: the same `ValueError` with the same text
when the input has no messages. The `create_react_agent` tests in
`libs/prebuilt` cover the call path.
2026-11-12 18:53:53 -05:00
Richard Scarrott a054ed08dc release(langgraph): 1.2.21 (#9683)
Bumps langgraph from 1.2.20 to 1.2.21.

Includes:

- Read a checkpoint namespace the graph doesn't declare, such as a
deepagents subagent's `tools:<task_id>`, with the parent graph instead
of raising `Subgraph tools not found`. The Python Agent Server no longer
returns 400 on `/history?checkpoint_ns=...` for these. Updates to such a
namespace still raise (#9614)
- Match whole namespace segments in `get_subgraphs` / `aget_subgraphs`.
With `namespace=` and `recurse=True` they now yield only the subgraph at
that namespace, and so does the Agent Server's `GET
/assistants/{id}/subgraphs/{namespace}?recurse=true` (#9615)
- Say each `interrupt()` call has its own id in the
`UnknownInterruptError` docstring (#9600)
2026-11-12 17:47:18 +00:00
17 changed files with 285 additions and 57 deletions

No files matched your search

@@ -3,7 +3,7 @@ from __future__ import annotations
import asyncio
from collections import defaultdict
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
from contextlib import asynccontextmanager
from contextlib import asynccontextmanager, nullcontext
from typing import Any, cast
from langchain_core.runnables import RunnableConfig
@@ -425,7 +425,12 @@ class AsyncPostgresSaver(BasePostgresSaver):
Will be applied regardless of whether the AsyncPostgresSaver instance was initialized with a pipeline.
If pipeline mode is not supported, will fall back to using transaction context manager.
"""
async with self.lock, _ainternal.get_connection(self.conn) as conn:
# A pool hands each call its own connection, so only a shared one needs
# the lock.
lock = (
nullcontext() if isinstance(self.conn, AsyncConnectionPool) else self.lock
)
async with lock, _ainternal.get_connection(self.conn) as conn:
if self.pipe:
# a connection in pipeline mode can be used concurrently
# in multiple threads/coroutines, but only one cursor can be
@@ -594,7 +599,11 @@ class AsyncPostgresSaver(BasePostgresSaver):
if value["parent_checkpoint_id"]
else None
),
await asyncio.to_thread(self._load_writes, value["pending_writes"]),
(
await asyncio.to_thread(self._load_writes, value["pending_writes"])
if value["pending_writes"]
else []
),
)
def list(
@@ -1,5 +1,6 @@
# type: ignore
import asyncio
from contextlib import asynccontextmanager
from typing import Any
from uuid import uuid4
@@ -417,3 +418,33 @@ async def test_delta_channel_chain_reconstruction(saver_name: str) -> None:
assert msgs[1].content == "reply-1"
assert msgs[2].content == "there"
assert msgs[3].content == "reply-3"
async def test_pool_saver_reads_while_a_list_is_open(test_data) -> None:
async with _pool_saver() as saver:
config = await saver.aput(
test_data["configs"][0],
test_data["checkpoints"][0],
test_data["metadata"][0],
{},
)
listing = saver.alist(None)
try:
assert await anext(listing) is not None
assert await asyncio.wait_for(saver.aget_tuple(config), timeout=5)
finally:
await listing.aclose()
async def test_base_saver_serializes_concurrent_calls(test_data) -> None:
async with _base_saver() as saver:
config = await saver.aput(
test_data["configs"][0],
test_data["checkpoints"][0],
test_data["metadata"][0],
{},
)
await asyncio.gather(
*(saver.aput_writes(config, [("x", i)], str(uuid4())) for i in range(20)),
*(saver.aget_tuple(config) for _ in range(20)),
)
@@ -302,14 +302,20 @@ def _dedupe_ops(values: list[Op]) -> tuple[list[int] | None, list[Op]]:
dedupped: list[Op] = []
listen: list[int] = []
puts: dict[tuple[tuple[str, ...], str], int] = {}
reads: dict[Op, int] = {}
for op in values:
if isinstance(op, (GetOp, SearchOp, ListNamespacesOp)):
try:
listen.append(dedupped.index(op))
except ValueError:
listen.append(len(dedupped))
ix = reads.setdefault(op, len(dedupped))
except TypeError: # an op holding a dict or a list is unhashable
try:
ix = dedupped.index(op)
except ValueError:
ix = len(dedupped)
if ix == len(dedupped):
dedupped.append(op)
listen.append(ix)
elif isinstance(op, PutOp):
putkey = (op.namespace, op.key)
if putkey in puts:
@@ -248,13 +248,20 @@ class InMemoryStore(BaseStore):
for key, filter_value in op.filter.items()
)
# Without a filter or a query embedding, the caller keeps only
# filtered[offset:offset + limit], so stop collecting there. Negative
# bounds keep the full scan, since a slice counts them from the end.
stop = (
op.offset + op.limit
if not op.filter
and not (op.query and self.embeddings)
and min(op.offset, op.limit) >= 0
else None
)
n = len(namespace_prefix)
filtered = []
for namespace in self._data:
if not (
namespace[: len(namespace_prefix)] == namespace_prefix
if len(namespace) >= len(namespace_prefix)
else False
):
if namespace[:n] != namespace_prefix:
continue
for key, item in self._data[namespace].items():
@@ -263,6 +270,8 @@ class InMemoryStore(BaseStore):
filtered.append((item, list(embeddings.values())))
else:
filtered.append((item, []))
if len(filtered) == stop:
return filtered
return filtered
def _embed_search_queries(
+61 -1
View File
@@ -22,7 +22,7 @@ from langgraph.store.base import (
get_text_at_path,
validate_op_namespace,
)
from langgraph.store.base.batch import AsyncBatchedBaseStore
from langgraph.store.base.batch import AsyncBatchedBaseStore, _dedupe_ops
from langgraph.store.memory import InMemoryStore
from tests.embed_test_utils import CharacterEmbeddings
@@ -75,6 +75,29 @@ async def test_async_batch_store_resilience() -> None:
assert result2.value == doc2
def test_dedupe_ops_maps_repeated_reads_and_puts_to_one_op() -> None:
ops = [
GetOp(("a",), "k"),
PutOp(("a",), "k", {"v": 1}),
GetOp(("a",), "k"),
SearchOp(("a",), filter={"x": 1}),
SearchOp(("a",), filter={"x": 1}),
PutOp(("a",), "k", {"v": 2}),
GetOp(("b",), "k"),
ListNamespacesOp(match_conditions=(MatchCondition("prefix", ("a",)),)),
ListNamespacesOp(match_conditions=(MatchCondition("prefix", ("a",)),)),
]
listen, dedupped = _dedupe_ops(ops)
assert dedupped == [
GetOp(("a",), "k"),
PutOp(("a",), "k", {"v": 2}),
SearchOp(("a",), filter={"x": 1}),
GetOp(("b",), "k"),
ListNamespacesOp(match_conditions=(MatchCondition("prefix", ("a",)),)),
]
assert listen == [0, 1, 0, 2, 2, 1, 3, 4, 4]
def test_get_text_at_path() -> None:
nested_data = {
"name": "test",
@@ -437,6 +460,43 @@ def test_list_namespaces_no_conditions() -> None:
assert sorted(result) == sorted(expected)
def test_search_pages_match_full_listing() -> None:
store = InMemoryStore()
items_per_namespace = 7
for ns in [("a",), ("a", "b"), ("ab",), ("a", "c", "d")]:
for i in range(items_per_namespace):
store.put(ns, f"k{i}", {"i": i, "even": i % 2 == 0})
def keys(items: list) -> list[tuple[tuple[str, ...], str]]:
return [(item.namespace, item.key) for item in items]
full = keys(store.search(("a",), limit=100))
assert {ns for ns, _ in full} == {("a",), ("a", "b"), ("a", "c", "d")}
assert len(full) == 3 * items_per_namespace
paged = [
key
for offset in range(0, len(full), 4)
for key in keys(store.search(("a",), limit=4, offset=offset))
]
assert paged == full
even = keys(store.search(("a",), filter={"even": True}, limit=100))
assert (
keys(store.search(("a",), filter={"even": True}, limit=3, offset=2))
== even[2:5]
)
def test_vector_search_scores_every_item(fake_embeddings: CharacterEmbeddings) -> None:
store = InMemoryStore(
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
)
for i in range(4):
store.put(("test",), f"other{i}", {"text": f"zzz qqq {i}"})
store.put(("test",), "match", {"text": "apple pie"})
results = store.search(("test",), query="apple pie", limit=1)
assert [r.key for r in results] == ["match"]
def test_list_namespaces_empty_store() -> None:
store = InMemoryStore()
+18 -9
View File
@@ -194,14 +194,8 @@ def add_messages(
if not isinstance(right, list):
right = [right] # type: ignore[assignment]
# coerce to message
left = [
message_chunk_to_message(cast(BaseMessageChunk, m))
for m in convert_to_messages(left)
]
right = [
message_chunk_to_message(cast(BaseMessageChunk, m))
for m in convert_to_messages(right)
]
left = _coerce_to_messages(left)
right = _coerce_to_messages(right)
# assign missing ids
for m in left:
if m.id is None:
@@ -234,7 +228,8 @@ def add_messages(
merged_by_id[m.id] = len(merged)
merged.append(m)
merged = [m for m in merged if m.id not in ids_to_remove]
if ids_to_remove:
merged = [m for m in merged if m.id not in ids_to_remove]
if format == "langchain-openai":
merged = _format_messages(merged)
@@ -247,6 +242,20 @@ def add_messages(
return merged
def _coerce_to_messages(messages: list) -> list[Any]:
# Both conversions return a plain message unchanged, and a long history is
# all plain messages, so check the few distinct types instead of each item.
if all(
issubclass(t, BaseMessage) and not issubclass(t, BaseMessageChunk)
for t in set(map(type, messages))
):
return messages
return [
message_chunk_to_message(cast(BaseMessageChunk, m))
for m in convert_to_messages(messages)
]
@deprecated(
"MessageGraph is deprecated in langgraph 1.0.0, to be removed in 2.0.0. Please use StateGraph with a `messages` key instead.",
category=None,
+7 -3
View File
@@ -1452,6 +1452,8 @@ class CompiledStateGraph(
output_keys = list(self.builder.channels) + [
k for k, v in self.builder.managed.items()
]
# for membership checks, which run on every key of every node update
output_key_set = frozenset(output_keys)
def _get_updates(
input: None | dict | Any,
@@ -1459,12 +1461,12 @@ class CompiledStateGraph(
if input is None:
return None
elif isinstance(input, dict):
return [(k, v) for k, v in input.items() if k in output_keys]
return [(k, v) for k, v in input.items() if k in output_key_set]
elif isinstance(input, Command):
if input.graph == Command.PARENT:
return None
return [
(k, v) for k, v in input._update_as_tuples() if k in output_keys
(k, v) for k, v in input._update_as_tuples() if k in output_key_set
]
elif (
isinstance(input, (list, tuple))
@@ -1477,7 +1479,9 @@ class CompiledStateGraph(
if i.graph == Command.PARENT:
continue
updates.extend(
(k, v) for k, v in i._update_as_tuples() if k in output_keys
(k, v)
for k, v in i._update_as_tuples()
if k in output_key_set
)
else:
updates.extend(_get_updates(i) or ())
+9 -9
View File
@@ -130,35 +130,35 @@ def map_output_updates(
]
if not output_tasks:
return
# a set: a wide state would otherwise scan every key for each write
output_keys = (
frozenset(output_channels) if not isinstance(output_channels, str) else None
)
updated: list[tuple[str, Any]] = []
for task, writes in output_tasks:
rtn = next((value for chan, value in writes if chan == RETURN), MISSING)
if rtn is not MISSING:
updated.append((task.name, rtn))
elif isinstance(output_channels, str):
elif output_keys is None:
updated.extend(
(task.name, value) for chan, value in writes if chan == output_channels
)
elif any(chan in output_channels for chan, _ in writes):
elif any(chan in output_keys for chan, _ in writes):
counts = Counter(chan for chan, _ in writes)
if any(counts[chan] > 1 for chan in output_channels):
if any(n > 1 for chan, n in counts.items() if chan in output_keys):
updated.extend(
(
task.name,
{chan: value},
)
for chan, value in writes
if chan in output_channels
if chan in output_keys
)
else:
updated.append(
(
task.name,
{
chan: value
for chan, value in writes
if chan in output_channels
},
{chan: value for chan, value in writes if chan in output_keys},
)
)
grouped: dict[str, Any] = {t.name: [] for t, _ in output_tasks}
+20 -14
View File
@@ -358,6 +358,14 @@ class PregelLoop:
self.retry_policy = retry_policy
self.cache_policy = cache_policy
self.durability = durability
# Channel types are fixed for the run, so per-step code reads these
# instead of scanning every channel.
self._has_untracked_channels = any(
isinstance(spec, UntrackedValue) for spec in specs.values()
)
self._delta_channel_names = [
k for k, spec in specs.items() if isinstance(spec, DeltaChannel)
]
self._has_graph_lifecycle_callbacks = has_graph_lifecycle_callbacks
self._graph_lifecycle_events = deque()
if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
@@ -496,9 +504,7 @@ class PregelLoop:
writes_to_save = writes
# check if any writes are to an UntrackedValue channel
if any(
isinstance(channel, UntrackedValue) for channel in self.channels.values()
):
if self._has_untracked_channels:
# we do not persist untracked values in checkpoints
writes_to_save = [
# sanitize UntrackedValues that are nested within Send packets
@@ -840,10 +846,12 @@ class PregelLoop:
Unfinished (failed or interrupted) tasks keep empty writes, so the
runner re-executes them or routes them to error handlers.
"""
for tid, status in read_task_statuses(self.checkpoint_pending_writes).items():
if task := tasks.get(tid):
task.writes.extend(status.output)
self._reapplied_task_ids.add(tid)
# Filter first: this runs once per task call, and reading the status of
# every task that wrote would make runs with many task calls quadratic.
writes = (w for w in self.checkpoint_pending_writes if w[0] in tasks)
for tid, status in read_task_statuses(writes).items():
tasks[tid].writes.extend(status.output)
self._reapplied_task_ids.add(tid)
def _resume_error_handlers_if_applicable(self) -> None:
"""On resume, schedule error handlers for tasks that failed in a prior run.
@@ -1300,9 +1308,7 @@ class PregelLoop:
)
new_counters: dict[str, tuple[int, int]] = {}
updated = self.updated_channels or set()
for ch_name, ch in self.channels.items():
if not isinstance(ch, DeltaChannel):
continue
for ch_name in self._delta_channel_names:
u, s = prev_counters.get(ch_name, (0, 0))
s += 1
if ch_name in updated:
@@ -1322,7 +1328,9 @@ class PregelLoop:
# create new checkpoint
channels_to_snapshot = (
delta_channels_to_snapshot(
self.channels, new_counters, self.checkpoint["channel_versions"]
{k: self.channels[k] for k in self._delta_channel_names},
new_counters,
self.checkpoint["channel_versions"],
)
| self._delta_channels_forced_snapshot
if do_checkpoint
@@ -1351,9 +1359,7 @@ class PregelLoop:
elif "counters_since_delta_snapshot" in self.checkpoint_metadata:
del self.checkpoint_metadata["counters_since_delta_snapshot"]
# sanitize TASK channel in the checkpoint before saving (durability=="exit")
if TASKS in self.checkpoint["channel_values"] and any(
isinstance(channel, UntrackedValue) for channel in self.channels.values()
):
if TASKS in self.checkpoint["channel_values"] and self._has_untracked_channels:
sanitized_tasks = [
sanitize_untracked_values_in_send(value, self.channels)
if isinstance(value, Send)
+9 -1
View File
@@ -78,6 +78,7 @@ class FuturesDict(Generic[F, E], dict[F, PregelExecutableTask | None]):
should_stop: Callable[[set[F]], bool]
counter: int
done: set[F]
failed: set[F]
lock: threading.Lock
def __init__(
@@ -97,6 +98,7 @@ class FuturesDict(Generic[F, E], dict[F, PregelExecutableTask | None]):
self.should_stop = should_stop
self.counter = 0
self.done: set[F] = set()
self.failed: set[F] = set()
def __setitem__(
self,
@@ -123,9 +125,13 @@ class FuturesDict(Generic[F, E], dict[F, PregelExecutableTask | None]):
with self.lock:
self.done.add(fut)
self.counter -= 1
# Only a future that raised can meet the stop condition, so check
# those instead of every done future (quadratic in wide steps).
if not fut.cancelled() and fut.exception() is not None:
self.failed.add(fut)
# Wake waiter when all tracked futures are done, or when runner-level
# stop condition is met (for example, a non-handled fatal exception).
if self.counter == 0 or self.should_stop(self.done):
if self.counter == 0 or self.should_stop(self.failed):
self.event.set()
@@ -234,6 +240,7 @@ class PregelRunner:
fut: concurrent.futures.Future = concurrent.futures.Future()
fut.set_exception(exc)
futures.done.add(fut)
futures.failed.add(fut)
elif reraise and id(exc) not in self._handled_exception_ids:
# let the caller stream this task's result before it raises
yield
@@ -427,6 +434,7 @@ class PregelRunner:
fut: asyncio.Future = loop.create_future()
fut.set_exception(exc)
futures.done.add(fut)
futures.failed.add(fut)
elif reraise and id(exc) not in self._handled_exception_ids:
# let the caller stream this task's result before it raises
yield
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "langgraph"
version = "1.2.20"
version = "1.2.21"
description = "Building stateful, multi-actor applications with LLMs"
authors = []
requires-python = ">=3.10"
@@ -5,6 +5,7 @@ import langchain_core
import pytest
from langchain_core.messages import (
AIMessage,
AIMessageChunk,
AnyMessage,
HumanMessage,
RemoveMessage,
@@ -168,6 +169,41 @@ def test_non_list_inputs():
assert result == expected_result
def test_message_likes_and_chunks_are_converted_on_both_sides():
left = [
HumanMessage(content="Hello", id="1"),
{"role": "assistant", "content": "Hi", "id": "2"},
]
right = [AIMessageChunk(content="chunk", id="3"), ("user", "Thanks")]
result = add_messages(left, right)
assert [type(m) for m in result] == [
HumanMessage,
AIMessage,
AIMessage,
HumanMessage,
]
assert result[:3] == [
HumanMessage(content="Hello", id="1"),
AIMessage(content="Hi", id="2"),
AIMessage(content="chunk", id="3"),
]
assert result[3].content == "Thanks"
def test_chunk_only_lists_are_converted():
left = [AIMessageChunk(content="a", id="1")]
right = [AIMessageChunk(content="b", id="2")]
result = add_messages(left, right)
assert [type(m) for m in result] == [AIMessage, AIMessage]
def test_result_is_a_new_list():
left = [HumanMessage(content="Hello", id="1")]
result = add_messages(left, [])
assert result == left
assert result is not left
def test_delete_all():
left = [
HumanMessage(content="Hello", id="1"),
@@ -0,0 +1,50 @@
import concurrent.futures
import threading
import weakref
from functools import partial
from langgraph.pregel._runner import FuturesDict, _should_stop_others
def _on_done(task: object, exc: BaseException | None) -> None:
pass
def _futures(handled_exception_ids: set[int]) -> FuturesDict:
return FuturesDict(
event=threading.Event(),
callback=weakref.ref(_on_done),
should_stop=partial(
_should_stop_others, handled_exception_ids=handled_exception_ids
),
future_type=concurrent.futures.Future,
)
def test_failed_task_wakes_waiter_before_others_finish() -> None:
futures = _futures(set())
ok, bad, pending = (concurrent.futures.Future() for _ in range(3))
for fut in (ok, bad, pending):
futures[fut] = object()
ok.set_result(None)
assert not futures.event.is_set()
bad.set_exception(ValueError("boom"))
assert futures.event.is_set()
def test_handled_failure_waits_for_all_tasks() -> None:
handled: set[int] = set()
futures = _futures(handled)
bad, pending = (concurrent.futures.Future() for _ in range(2))
for fut in (bad, pending):
futures[fut] = object()
exc = ValueError("handled by an error handler")
handled.add(id(exc))
bad.set_exception(exc)
assert not futures.event.is_set()
pending.set_result(None)
assert futures.event.is_set()
+1 -1
View File
@@ -1440,7 +1440,7 @@ wheels = [
[[package]]
name = "langgraph"
version = "1.2.20"
version = "1.2.21"
source = { editable = "." }
dependencies = [
{ name = "langchain-core" },
@@ -638,15 +638,15 @@ def create_react_agent(
messages = (
_get_state_value(state, "llm_input_messages")
) or _get_state_value(state, "messages")
error_msg = f"Expected input to call_model to have 'llm_input_messages' or 'messages' key, but got {state}"
expected_keys = "'llm_input_messages' or 'messages' key"
else:
messages = _get_state_value(state, "messages")
error_msg = (
f"Expected input to call_model to have 'messages' key, but got {state}"
)
expected_keys = "'messages' key"
if messages is None:
raise ValueError(error_msg)
raise ValueError(
f"Expected input to call_model to have {expected_keys}, but got {state}"
)
_validate_chat_history(messages)
# we're passing messages under `messages` key, as this is expected by the prompt
+1 -1
View File
@@ -381,7 +381,7 @@ wheels = [
[[package]]
name = "langgraph"
version = "1.2.20"
version = "1.2.21"
source = { editable = "../langgraph" }
dependencies = [
{ name = "langchain-core" },
+1 -1
View File
@@ -394,7 +394,7 @@ wheels = [
[[package]]
name = "langgraph"
version = "1.2.20"
version = "1.2.21"
source = { editable = "../langgraph" }
dependencies = [
{ name = "langchain-core" },