From 37a55044338beab18f952ea5102dfb23d4fecef1 Mon Sep 17 00:00:00 2001 From: Christian Bromann Date: Thu, 12 Mar 2026 15:06:14 -0700 Subject: [PATCH] move to stream mode --- libs/langgraph/langgraph/pregel/main.py | 4 +-- libs/langgraph/langgraph/pregel/remote.py | 26 ---------------- libs/langgraph/langgraph/types.py | 11 +++++-- libs/langgraph/tests/test_remote_graph.py | 18 ++++++++--- libs/langgraph/tests/test_stream_v2.py | 37 ++++++++++++++++++----- libs/sdk-py/langgraph_sdk/_async/runs.py | 12 -------- libs/sdk-py/langgraph_sdk/_sync/runs.py | 12 -------- libs/sdk-py/langgraph_sdk/schema.py | 4 +-- libs/sdk-py/tests/test_client_stream.py | 12 ++++---- 9 files changed, 62 insertions(+), 74 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 5bf5b4337..37646a891 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -2614,7 +2614,7 @@ class Pregel( stream.put, subgraphs, parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None, - dedupe_metadata=version == "v2", + dedupe_metadata="compact" in stream_modes, ) ) @@ -2966,7 +2966,7 @@ class Pregel( stream_put, subgraphs, parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None, - dedupe_metadata=version == "v2", + dedupe_metadata="compact" in stream_modes, ) ) diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 6c1e18a08..9c6a367b4 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -71,8 +71,6 @@ logger = logging.getLogger(__name__) __all__ = ("RemoteGraph", "RemoteException") -_STREAM_PROTOCOL_CONFIG_KEY = "__stream_protocol_version__" - _CONF_DROPLIST = frozenset( ( CONFIG_KEY_CHECKPOINT_MAP, @@ -801,17 +799,6 @@ class RemoteGraph(PregelProtocol): stream_modes, requested, req_single, stream = self._get_stream_modes( stream_mode, config ) - stream_protocol_version = cast( - Literal["v1", "v2"] | None, - kwargs.pop( - "stream_protocol_version", - sanitized_config.get("configurable", {}).get( - _STREAM_PROTOCOL_CONFIG_KEY - ), - ), - ) - if version == "v2" and stream_protocol_version is None: - stream_protocol_version = "v2" if isinstance(input, Command): command: CommandSDK | None = cast(CommandSDK, asdict(input)) input = None @@ -837,7 +824,6 @@ class RemoteGraph(PregelProtocol): _merge_tracing_headers(headers) if self.distributed_tracing else headers ), params=params, - stream_protocol_version=stream_protocol_version, **kwargs, ): # split mode and ns @@ -972,17 +958,6 @@ class RemoteGraph(PregelProtocol): stream_modes, requested, req_single, stream = self._get_stream_modes( stream_mode, config ) - stream_protocol_version = cast( - Literal["v1", "v2"] | None, - kwargs.pop( - "stream_protocol_version", - sanitized_config.get("configurable", {}).get( - _STREAM_PROTOCOL_CONFIG_KEY - ), - ), - ) - if version == "v2" and stream_protocol_version is None: - stream_protocol_version = "v2" if isinstance(input, Command): command: CommandSDK | None = cast(CommandSDK, asdict(input)) input = None @@ -1008,7 +983,6 @@ class RemoteGraph(PregelProtocol): _merge_tracing_headers(headers) if self.distributed_tracing else headers ), params=params, - stream_protocol_version=stream_protocol_version, **kwargs, ): # split mode and ns diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 980bc07df..4aaec1b14 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -116,7 +116,14 @@ def ensure_valid_checkpointer(checkpointer: Checkpointer) -> Checkpointer: StreamMode = Literal[ - "values", "updates", "checkpoints", "tasks", "debug", "messages", "custom" + "values", + "updates", + "checkpoints", + "tasks", + "debug", + "messages", + "custom", + "compact", ] """How the stream method should emit outputs. @@ -277,7 +284,7 @@ class MessagesStreamPart(TypedDict): `data` is a 2-tuple of `(message, metadata)` where `message` is a `BaseMessage` (e.g. `AIMessageChunk`) and `metadata` is either a dict containing keys like `langgraph_step`, `langgraph_node`, `langgraph_triggers`, etc. or - `None` for deduplicated follow-up chunks in `version="v2"` streams. + `None` for deduplicated follow-up chunks when `stream_mode` includes `"compact"`. """ type: Literal["messages"] diff --git a/libs/langgraph/tests/test_remote_graph.py b/libs/langgraph/tests/test_remote_graph.py index 6d930dbfe..49dd56341 100644 --- a/libs/langgraph/tests/test_remote_graph.py +++ b/libs/langgraph/tests/test_remote_graph.py @@ -917,7 +917,7 @@ def test_stream_restores_messages_and_merges_values_patch(): remote_pregel.stream( {"input": "data"}, config={"configurable": {"thread_id": "thread_1"}}, - stream_mode=["messages", "values"], + stream_mode=["messages", "values", "compact"], subgraphs=True, version="v2", ) @@ -948,7 +948,12 @@ def test_stream_restores_messages_and_merges_values_patch(): }, ] _, kwargs = mock_sync_client.runs.stream.call_args - assert kwargs["stream_protocol_version"] == "v2" + assert set(kwargs["stream_mode"]) == { + "messages-tuple", + "values", + "compact", + "updates", + } @pytest.mark.anyio @@ -1199,7 +1204,7 @@ async def test_astream_restores_messages_and_merges_values_patch(): async for part in remote_pregel.astream( {"input": "data"}, config={"configurable": {"thread_id": "thread_1"}}, - stream_mode=["messages", "values"], + stream_mode=["messages", "values", "compact"], subgraphs=True, version="v2", ): @@ -1230,7 +1235,12 @@ async def test_astream_restores_messages_and_merges_values_patch(): }, ] _, kwargs = mock_async_client.runs.stream.call_args - assert kwargs["stream_protocol_version"] == "v2" + assert set(kwargs["stream_mode"]) == { + "messages-tuple", + "values", + "compact", + "updates", + } @pytest.mark.skip( diff --git a/libs/langgraph/tests/test_stream_v2.py b/libs/langgraph/tests/test_stream_v2.py index f2fdd9b15..4e692976d 100644 --- a/libs/langgraph/tests/test_stream_v2.py +++ b/libs/langgraph/tests/test_stream_v2.py @@ -190,6 +190,19 @@ class TestV1BackwardsCompat: assert len(metadata) >= 3 assert all(isinstance(meta, dict) for meta in metadata) + def test_stream_v1_messages_compact_dedupes_metadata(self) -> None: + graph = _make_streaming_messages_graph().compile() + chunks = list(graph.stream(_MSG_INPUT, stream_mode=["messages", "compact"])) + metadata = [ + meta + for mode, payload in chunks + if mode == "messages" + for _message, meta in [payload] + ] + assert len(metadata) >= 3 + assert isinstance(metadata[0], dict) + assert all(meta is None for meta in metadata[1:]) + # --- v2 sync stream --- @@ -228,12 +241,18 @@ class TestV2Stream: assert isinstance(data, tuple) and len(data) == 2 message, metadata = data assert isinstance(message, BaseMessage) - assert metadata is None or isinstance(metadata, dict) - assert any(isinstance(c["data"][1], dict) for c in msg_chunks) + assert isinstance(metadata, dict) + assert "langgraph_node" in metadata - def test_messages_streaming_dedupes_metadata(self) -> None: + def test_messages_streaming_compact_dedupes_metadata(self) -> None: graph = _make_streaming_messages_graph().compile() - chunks = list(graph.stream(_MSG_INPUT, stream_mode="messages", version="v2")) + chunks = list( + graph.stream( + _MSG_INPUT, + stream_mode=["messages", "compact"], + version="v2", + ) + ) msg_chunks = [c for c in chunks if c["type"] == "messages"] assert len(msg_chunks) >= 3 first_message, first_metadata = msg_chunks[0]["data"] @@ -581,16 +600,18 @@ class TestV2StreamAsync: assert isinstance(data, tuple) and len(data) == 2 message, metadata = data assert isinstance(message, BaseMessage) - assert metadata is None or isinstance(metadata, dict) - assert any(isinstance(c["data"][1], dict) for c in msg_chunks) + assert isinstance(metadata, dict) + assert "langgraph_node" in metadata @pytest.mark.anyio - async def test_messages_streaming_dedupes_metadata(self) -> None: + async def test_messages_streaming_compact_dedupes_metadata(self) -> None: graph = _make_streaming_messages_graph().compile() chunks = [ c async for c in graph.astream( - _MSG_INPUT, stream_mode="messages", version="v2" + _MSG_INPUT, + stream_mode=["messages", "compact"], + version="v2", ) ] msg_chunks = [c for c in chunks if c["type"] == "messages"] diff --git a/libs/sdk-py/langgraph_sdk/_async/runs.py b/libs/sdk-py/langgraph_sdk/_async/runs.py index 05d65cf0e..c6d1fa162 100644 --- a/libs/sdk-py/langgraph_sdk/_async/runs.py +++ b/libs/sdk-py/langgraph_sdk/_async/runs.py @@ -85,7 +85,6 @@ class RunsClient: checkpoint: Checkpoint | None = None, checkpoint_id: str | None = None, checkpoint_during: bool | None = None, - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, feedback_keys: Sequence[str] | None = None, @@ -125,7 +124,6 @@ class RunsClient: multitask_strategy: MultitaskStrategy | None = None, if_not_exists: IfNotExists | None = None, after_seconds: int | None = None, - stream_protocol_version: StreamVersion | None = None, headers: Mapping[str, str] | None = None, params: QueryParamTypes | None = None, on_run_created: Callable[[RunCreateMetadata], None] | None = None, @@ -146,7 +144,6 @@ class RunsClient: metadata: Mapping[str, Any] | None = None, config: Config | None = None, checkpoint_during: bool | None = None, - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, feedback_keys: Sequence[str] | None = None, @@ -183,7 +180,6 @@ class RunsClient: if_not_exists: IfNotExists | None = None, webhook: str | None = None, after_seconds: int | None = None, - stream_protocol_version: StreamVersion | None = None, headers: Mapping[str, str] | None = None, params: QueryParamTypes | None = None, on_run_created: Callable[[RunCreateMetadata], None] | None = None, @@ -215,7 +211,6 @@ class RunsClient: multitask_strategy: MultitaskStrategy | None = None, if_not_exists: IfNotExists | None = None, after_seconds: int | None = None, - stream_protocol_version: StreamVersion | None = None, headers: Mapping[str, str] | None = None, params: QueryParamTypes | None = None, on_run_created: Callable[[RunCreateMetadata], None] | None = None, @@ -317,7 +312,6 @@ class RunsClient: "stream_mode": stream_mode, "stream_subgraphs": stream_subgraphs, "stream_resumable": stream_resumable, - "stream_protocol_version": stream_protocol_version, "assistant_id": assistant_id, "interrupt_before": interrupt_before, "interrupt_after": interrupt_after, @@ -399,7 +393,6 @@ class RunsClient: checkpoint: Checkpoint | None = None, checkpoint_id: str | None = None, checkpoint_during: bool | None = None, - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, webhook: str | None = None, @@ -427,7 +420,6 @@ class RunsClient: checkpoint: Checkpoint | None = None, checkpoint_id: str | None = None, checkpoint_during: bool | None = None, # deprecated - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, webhook: str | None = None, @@ -565,7 +557,6 @@ class RunsClient: "stream_mode": stream_mode, "stream_subgraphs": stream_subgraphs, "stream_resumable": stream_resumable, - "stream_protocol_version": stream_protocol_version, "config": config, "context": context, "metadata": metadata, @@ -653,7 +644,6 @@ class RunsClient: config: Config | None = None, context: Context | None = None, checkpoint_during: bool | None = None, - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, webhook: str | None = None, @@ -680,7 +670,6 @@ class RunsClient: checkpoint: Checkpoint | None = None, checkpoint_id: str | None = None, checkpoint_during: bool | None = None, # deprecated - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, webhook: str | None = None, @@ -796,7 +785,6 @@ class RunsClient: "config": config, "context": context, "metadata": metadata, - "stream_protocol_version": stream_protocol_version, "assistant_id": assistant_id, "interrupt_before": interrupt_before, "interrupt_after": interrupt_after, diff --git a/libs/sdk-py/langgraph_sdk/_sync/runs.py b/libs/sdk-py/langgraph_sdk/_sync/runs.py index 54f3976e4..1b52b8113 100644 --- a/libs/sdk-py/langgraph_sdk/_sync/runs.py +++ b/libs/sdk-py/langgraph_sdk/_sync/runs.py @@ -84,7 +84,6 @@ class SyncRunsClient: checkpoint: Checkpoint | None = None, checkpoint_id: str | None = None, checkpoint_during: bool | None = None, - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, feedback_keys: Sequence[str] | None = None, @@ -115,7 +114,6 @@ class SyncRunsClient: checkpoint: Checkpoint | None = None, checkpoint_id: str | None = None, checkpoint_during: bool | None = None, - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, feedback_keys: Sequence[str] | None = None, @@ -145,7 +143,6 @@ class SyncRunsClient: config: Config | None = None, context: Context | None = None, checkpoint_during: bool | None = None, - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, feedback_keys: Sequence[str] | None = None, @@ -175,7 +172,6 @@ class SyncRunsClient: config: Config | None = None, context: Context | None = None, checkpoint_during: bool | None = None, - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, feedback_keys: Sequence[str] | None = None, @@ -206,7 +202,6 @@ class SyncRunsClient: checkpoint: Checkpoint | None = None, checkpoint_id: str | None = None, checkpoint_during: bool | None = None, # deprecated - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, feedback_keys: Sequence[str] | None = None, @@ -312,7 +307,6 @@ class SyncRunsClient: "stream_mode": stream_mode, "stream_subgraphs": stream_subgraphs, "stream_resumable": stream_resumable, - "stream_protocol_version": stream_protocol_version, "assistant_id": assistant_id, "interrupt_before": interrupt_before, "interrupt_after": interrupt_after, @@ -366,7 +360,6 @@ class SyncRunsClient: config: Config | None = None, context: Context | None = None, checkpoint_during: bool | None = None, - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, webhook: str | None = None, @@ -395,7 +388,6 @@ class SyncRunsClient: checkpoint: Checkpoint | None = None, checkpoint_id: str | None = None, checkpoint_during: bool | None = None, - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, webhook: str | None = None, @@ -423,7 +415,6 @@ class SyncRunsClient: checkpoint: Checkpoint | None = None, checkpoint_id: str | None = None, checkpoint_during: bool | None = None, # deprecated - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, webhook: str | None = None, @@ -561,7 +552,6 @@ class SyncRunsClient: "stream_mode": stream_mode, "stream_subgraphs": stream_subgraphs, "stream_resumable": stream_resumable, - "stream_protocol_version": stream_protocol_version, "config": config, "context": context, "metadata": metadata, @@ -675,7 +665,6 @@ class SyncRunsClient: checkpoint_during: bool | None = None, # deprecated checkpoint: Checkpoint | None = None, checkpoint_id: str | None = None, - stream_protocol_version: StreamVersion | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, webhook: str | None = None, @@ -793,7 +782,6 @@ class SyncRunsClient: "config": config, "context": context, "metadata": metadata, - "stream_protocol_version": stream_protocol_version, "assistant_id": assistant_id, "interrupt_before": interrupt_before, "interrupt_after": interrupt_after, diff --git a/libs/sdk-py/langgraph_sdk/schema.py b/libs/sdk-py/langgraph_sdk/schema.py index ada89ac24..60271a107 100644 --- a/libs/sdk-py/langgraph_sdk/schema.py +++ b/libs/sdk-py/langgraph_sdk/schema.py @@ -58,6 +58,7 @@ StreamMode = Literal[ "debug", "custom", "messages-tuple", + "compact", ] """ Defines the mode of streaming: @@ -69,6 +70,7 @@ Defines the mode of streaming: - "tasks": Stream task start and finish events. - "debug": Stream detailed debug information. - "custom": Stream custom events. +- "compact": Enable compact streaming payloads for other selected modes. """ DisconnectMode = Literal["cancel", "continue"] @@ -532,8 +534,6 @@ class RunCreate(TypedDict): """List of node names to interrupt execution after.""" webhook: str | None """URL to send webhook notifications about the run's progress.""" - stream_protocol_version: StreamVersion | None - """Opt into an alternate stream wire protocol version.""" multitask_strategy: MultitaskStrategy | None """Strategy for handling concurrent runs on the same thread.""" diff --git a/libs/sdk-py/tests/test_client_stream.py b/libs/sdk-py/tests/test_client_stream.py index ac4908daa..98d5f094b 100644 --- a/libs/sdk-py/tests/test_client_stream.py +++ b/libs/sdk-py/tests/test_client_stream.py @@ -393,12 +393,12 @@ def test_sse_to_v2_dict_values_patch() -> None: @pytest.mark.asyncio -async def test_async_runs_stream_includes_stream_protocol_version(): +async def test_async_runs_stream_includes_compact_mode(): async def handler(request: httpx.Request) -> httpx.Response: assert request.method == "POST" assert request.url.path == "/runs/stream" body = json.loads(request.content) - assert body["stream_protocol_version"] == "v2" + assert body["stream_mode"] == ["values", "compact"] return httpx.Response( 200, headers={"Content-Type": "text/event-stream"}, @@ -416,7 +416,7 @@ async def test_async_runs_stream_includes_stream_protocol_version(): thread_id=None, assistant_id="agent", input={"messages": []}, - stream_protocol_version="v2", + stream_mode=["values", "compact"], ) ] @@ -425,12 +425,12 @@ async def test_async_runs_stream_includes_stream_protocol_version(): assert parts[0].data is None -def test_sync_runs_stream_includes_stream_protocol_version(): +def test_sync_runs_stream_includes_compact_mode(): def handler(request: httpx.Request) -> httpx.Response: assert request.method == "POST" assert request.url.path == "/runs/stream" body = json.loads(request.content) - assert body["stream_protocol_version"] == "v2" + assert body["stream_mode"] == ["values", "compact"] return httpx.Response( 200, headers={"Content-Type": "text/event-stream"}, @@ -445,7 +445,7 @@ def test_sync_runs_stream_includes_stream_protocol_version(): thread_id=None, assistant_id="agent", input={"messages": []}, - stream_protocol_version="v2", + stream_mode=["values", "compact"], ) )