mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-07 18:27:52 +02:00
fix(message): full add_messages parity for _messages_delta_reducer
Add REMOVE_ALL_MESSAGES sentinel handling and BaseMessageChunk coercion, completing parity with add_messages. REMOVE_ALL_MESSAGES resets state and discards all preceding writes; chunks are coerced to full messages via message_chunk_to_message. Both properties are batching-invariant.
This commit is contained in:
@@ -249,18 +249,15 @@ def _messages_delta_reducer(
|
||||
) -> list[AnyMessage]:
|
||||
"""**Experimental.** Batch reducer for use with `DeltaChannel`.
|
||||
|
||||
Processes all writes in one pass — dedup by ID, `RemoveMessage`
|
||||
tombstoning — without calling `add_messages`.
|
||||
Provides full `add_messages` parity: dedup by ID, `RemoveMessage`
|
||||
tombstoning, `REMOVE_ALL_MESSAGES` reset, `BaseMessageChunk` coercion,
|
||||
and UUID assignment for ID-less messages — all in a single batched pass.
|
||||
|
||||
This reducer is batching-invariant, as required by `DeltaChannel`:
|
||||
`reducer(reducer(state, xs), ys) == reducer(state, xs + ys)`.
|
||||
|
||||
Raw dict / string / tuple inputs are coerced to typed `BaseMessage`
|
||||
objects so that HTTP-driven graphs work without a separate coercion
|
||||
step. Messages without IDs are assigned UUIDs (matching `add_messages`
|
||||
behavior) so that message eviction and `RemoveMessage` tombstoning work
|
||||
correctly. `REMOVE_ALL_MESSAGES` and `BaseMessageChunk` conversion are
|
||||
not handled here.
|
||||
objects so that HTTP-driven graphs work without a separate coercion step.
|
||||
|
||||
Example::
|
||||
|
||||
@@ -281,14 +278,31 @@ def _messages_delta_reducer(
|
||||
flat.extend(w)
|
||||
else:
|
||||
flat.append(w)
|
||||
# Steady state: the reducer's own output is already typed, so skip
|
||||
# `convert_to_messages` on state when the first element is a BaseMessage.
|
||||
# Steady state: the reducer's own output is already typed BaseMessages
|
||||
# (never chunks), so skip convert_to_messages on the fast path.
|
||||
# Only raw input (initial dicts, deserialized blobs) hits the slow path.
|
||||
if state and isinstance(state[0], BaseMessage):
|
||||
state_msgs = state
|
||||
state_msgs = cast("list[AnyMessage]", state)
|
||||
else:
|
||||
state_msgs = cast("list[AnyMessage]", convert_to_messages(state))
|
||||
msgs = cast("list[AnyMessage]", convert_to_messages(flat))
|
||||
state_msgs = cast(
|
||||
"list[AnyMessage]",
|
||||
[message_chunk_to_message(cast(BaseMessageChunk, m)) for m in convert_to_messages(state)],
|
||||
)
|
||||
# Coerce chunks to full messages — streaming nodes can emit BaseMessageChunk.
|
||||
msgs = cast(
|
||||
"list[AnyMessage]",
|
||||
[message_chunk_to_message(cast(BaseMessageChunk, m)) for m in convert_to_messages(flat)],
|
||||
)
|
||||
|
||||
# REMOVE_ALL_MESSAGES resets everything; find the last sentinel and
|
||||
# discard all state plus all writes before it.
|
||||
remove_all_idx = None
|
||||
for idx, m in enumerate(msgs):
|
||||
if isinstance(m, RemoveMessage) and m.id == REMOVE_ALL_MESSAGES:
|
||||
remove_all_idx = idx
|
||||
if remove_all_idx is not None:
|
||||
state_msgs = []
|
||||
msgs = msgs[remove_all_idx + 1 :]
|
||||
|
||||
# Build index and assign missing IDs in one pass (parity with add_messages
|
||||
# so that eviction and RemoveMessage tombstoning work on ID-less messages).
|
||||
|
||||
@@ -3,7 +3,7 @@ from collections.abc import Sequence
|
||||
from typing import Annotated
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
|
||||
from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, RemoveMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from typing_extensions import NotRequired, TypedDict
|
||||
@@ -16,7 +16,7 @@ from langgraph.channels.topic import Topic
|
||||
from langgraph.channels.untracked_value import UntrackedValue
|
||||
from langgraph.errors import EmptyChannelError, InvalidUpdateError
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
from langgraph.graph.message import REMOVE_ALL_MESSAGES, _messages_delta_reducer
|
||||
from langgraph.graph.state import _get_channel
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
@@ -306,6 +306,31 @@ def test_messages_delta_reducer_assigns_uuid_to_id_less_messages() -> None:
|
||||
assert result2[0].content == "hi"
|
||||
|
||||
|
||||
def test_messages_delta_reducer_remove_all_messages() -> None:
|
||||
"""REMOVE_ALL_MESSAGES sentinel clears all state and preceding writes."""
|
||||
state = [HumanMessage(content="old", id="h1"), AIMessage(content="prior", id="a1")]
|
||||
|
||||
# Sentinel mid-batch: everything before it (including state) is discarded.
|
||||
result = _messages_delta_reducer(state, [[RemoveMessage(id=REMOVE_ALL_MESSAGES), HumanMessage(content="fresh", id="h2")]])
|
||||
assert len(result) == 1
|
||||
assert result[0].content == "fresh"
|
||||
|
||||
# Batching-invariant: split across two calls must equal one combined call.
|
||||
step1 = _messages_delta_reducer(state, [[RemoveMessage(id=REMOVE_ALL_MESSAGES)]])
|
||||
step2 = _messages_delta_reducer(step1, [[HumanMessage(content="fresh", id="h2")]])
|
||||
assert step2 == result
|
||||
|
||||
|
||||
def test_messages_delta_reducer_coerces_message_chunks() -> None:
|
||||
"""BaseMessageChunk writes are coerced to full messages."""
|
||||
chunk = AIMessageChunk(content="hello", id="a1")
|
||||
result = _messages_delta_reducer([], [[chunk]])
|
||||
assert len(result) == 1
|
||||
assert not isinstance(result[0], AIMessageChunk)
|
||||
assert result[0].content == "hello"
|
||||
assert result[0].id == "a1"
|
||||
|
||||
|
||||
def test_delta_channel_checkpoint_returns_missing() -> None:
|
||||
"""checkpoint() always returns MISSING regardless of state.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user