mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-09 19:27:54 +02:00
move to stream mode
This commit is contained in:
@@ -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,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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"],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user