move to stream mode

This commit is contained in:
Christian Bromann
2026-03-12 15:08:26 -07:00
parent f6286dce38
commit 37a5504433
9 changed files with 62 additions and 74 deletions
+2 -2
View File
@@ -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,
)
)
-26
View File
@@ -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
+9 -2
View File
@@ -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"]
+14 -4
View File
@@ -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(
+29 -8
View File
@@ -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"]
-12
View File
@@ -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,
-12
View File
@@ -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,
+2 -2
View File
@@ -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."""
+6 -6
View File
@@ -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"],
)
)