Compare commits

...
Author SHA1 Message Date
Christian Bromann 741c6f8d50 cr 2026-03-18 15:46:35 -07:00
Christian Bromann 37a5504433 move to stream mode 2026-03-12 15:08:26 -07:00
Christian Bromann f6286dce38 feat(langgraph): add protocol improvements for better streaming 2026-03-12 15:08:25 -07:00
8 changed files with 421 additions and 18 deletions
+10 -1
View File
@@ -25,7 +25,7 @@ except ImportError:
_StreamingCallbackHandler = object # type: ignore _StreamingCallbackHandler = object # type: ignore
T = TypeVar("T") T = TypeVar("T")
Meta = tuple[tuple[str, ...], dict[str, Any]] Meta = tuple[tuple[str, ...], dict[str, Any] | None]
def _state_values(obj: Any) -> Sequence[Any]: def _state_values(obj: Any) -> Sequence[Any]:
@@ -56,6 +56,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
subgraphs: bool, subgraphs: bool,
*, *,
parent_ns: tuple[str, ...] | None = None, parent_ns: tuple[str, ...] | None = None,
dedupe_metadata: bool = False,
) -> None: ) -> None:
"""Configure the handler to stream messages from LLMs and nodes. """Configure the handler to stream messages from LLMs and nodes.
@@ -84,8 +85,10 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
self.stream = stream self.stream = stream
self.subgraphs = subgraphs self.subgraphs = subgraphs
self.metadata: dict[UUID, Meta] = {} self.metadata: dict[UUID, Meta] = {}
self.emitted_metadata: set[UUID] = set()
self.seen: set[int | str] = set() self.seen: set[int | str] = set()
self.parent_ns = parent_ns self.parent_ns = parent_ns
self.dedupe_metadata = dedupe_metadata
def _emit(self, meta: Meta, message: BaseMessage, *, dedupe: bool = False) -> None: def _emit(self, meta: Meta, message: BaseMessage, *, dedupe: bool = False) -> None:
if dedupe and message.id in self.seen: if dedupe and message.id in self.seen:
@@ -155,6 +158,10 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
if not isinstance(chunk, ChatGenerationChunk): if not isinstance(chunk, ChatGenerationChunk):
return return
if meta := self.metadata.get(run_id): if meta := self.metadata.get(run_id):
if self.dedupe_metadata and run_id in self.emitted_metadata:
meta = (meta[0], None)
else:
self.emitted_metadata.add(run_id)
self._emit(meta, chunk.message) self._emit(meta, chunk.message)
def on_llm_end( def on_llm_end(
@@ -170,6 +177,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
gen = response.generations[0][0] gen = response.generations[0][0]
if isinstance(gen, ChatGeneration): if isinstance(gen, ChatGeneration):
self._emit(meta, gen.message, dedupe=True) self._emit(meta, gen.message, dedupe=True)
self.emitted_metadata.discard(run_id)
self.metadata.pop(run_id, None) self.metadata.pop(run_id, None)
def on_llm_error( def on_llm_error(
@@ -180,6 +188,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
parent_run_id: UUID | None = None, parent_run_id: UUID | None = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
self.emitted_metadata.discard(run_id)
self.metadata.pop(run_id, None) self.metadata.pop(run_id, None)
def on_chain_start( def on_chain_start(
+2
View File
@@ -2614,6 +2614,7 @@ class Pregel(
stream.put, stream.put,
subgraphs, subgraphs,
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None, parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
dedupe_metadata="compact" in stream_modes,
) )
) )
@@ -2965,6 +2966,7 @@ class Pregel(
stream_put, stream_put,
subgraphs, subgraphs,
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None, parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
dedupe_metadata="compact" in stream_modes,
) )
) )
+59 -12
View File
@@ -109,6 +109,45 @@ class RemoteException(Exception):
pass pass
def _restore_message_metadata(
data: Any, metadata_by_message_id: dict[str, dict[str, Any]]
) -> Any:
"""Restore deduplicated message metadata using the message id as cache key."""
if not (isinstance(data, list) and len(data) == 2):
return data
message, metadata = data
if not isinstance(message, dict):
return data
message_id = message.get("id")
if isinstance(message_id, str):
if isinstance(metadata, dict):
metadata_by_message_id[message_id] = metadata
else:
metadata = metadata_by_message_id.get(message_id)
return (message, metadata)
def _merge_values_patch(
ns: tuple[str, ...],
mode: str,
data: Any,
values_by_ns: dict[tuple[str, ...], dict[str, Any]],
) -> tuple[str, Any]:
"""Merge `values-patch` events back into full values snapshots."""
if mode != "values-patch" or not isinstance(data, dict):
return mode, data
values = data.get("values")
if not isinstance(values, dict):
return "values", values if values is not None else {}
merged = dict(values_by_ns.get(ns, {}))
merged.update(values)
for key in data.get("deleted_keys", ()):
if isinstance(key, str):
merged.pop(key, None)
values_by_ns[ns] = merged
return "values", merged
class RemoteGraph(PregelProtocol): class RemoteGraph(PregelProtocol):
"""The `RemoteGraph` class is a client implementation for calling remote """The `RemoteGraph` class is a client implementation for calling remote
APIs that implement the LangGraph Server API specification. APIs that implement the LangGraph Server API specification.
@@ -766,6 +805,8 @@ class RemoteGraph(PregelProtocol):
else: else:
command = None command = None
thread_id = sanitized_config.get("configurable", {}).pop("thread_id", None) thread_id = sanitized_config.get("configurable", {}).pop("thread_id", None)
message_metadata_by_id: dict[str, dict[str, Any]] = {}
values_by_ns: dict[tuple[str, ...], dict[str, Any]] = {}
for chunk in sync_client.runs.stream( for chunk in sync_client.runs.stream(
thread_id=thread_id, thread_id=thread_id,
@@ -798,6 +839,11 @@ class RemoteGraph(PregelProtocol):
if caller_ns := (config or {}).get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS): if caller_ns := (config or {}).get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS):
caller_ns = tuple(caller_ns.split(NS_SEP)) caller_ns = tuple(caller_ns.split(NS_SEP))
ns = caller_ns + ns ns = caller_ns + ns
mode, data = _merge_values_patch(ns, mode, chunk.data, values_by_ns)
if mode != chunk.event:
chunk = chunk._replace(data=data)
elif data is not chunk.data:
chunk = chunk._replace(data=data)
# stream to parent stream # stream to parent stream
if stream is not None and mode in stream.modes: if stream is not None and mode in stream.modes:
stream((ns, mode, chunk.data)) stream((ns, mode, chunk.data))
@@ -815,7 +861,9 @@ class RemoteGraph(PregelProtocol):
continue continue
if chunk.event.startswith("messages"): if chunk.event.startswith("messages"):
chunk = chunk._replace(data=tuple(chunk.data)) chunk = chunk._replace(
data=_restore_message_metadata(chunk.data, message_metadata_by_id)
)
# emit chunk # emit chunk
if version == "v2": if version == "v2":
@@ -827,11 +875,6 @@ class RemoteGraph(PregelProtocol):
) )
yield {"type": mode, "ns": ns, "data": chunk.data, "interrupts": ints} yield {"type": mode, "ns": ns, "data": chunk.data, "interrupts": ints}
elif subgraphs: elif subgraphs:
if NS_SEP in chunk.event:
mode, ns_ = chunk.event.split(NS_SEP, 1)
ns = tuple(ns_.split(NS_SEP))
else:
mode, ns = chunk.event, ()
if req_single: if req_single:
yield ns, chunk.data yield ns, chunk.data
else: else:
@@ -921,6 +964,8 @@ class RemoteGraph(PregelProtocol):
else: else:
command = None command = None
thread_id = sanitized_config.get("configurable", {}).pop("thread_id", None) thread_id = sanitized_config.get("configurable", {}).pop("thread_id", None)
message_metadata_by_id: dict[str, dict[str, Any]] = {}
values_by_ns: dict[tuple[str, ...], dict[str, Any]] = {}
async for chunk in client.runs.stream( async for chunk in client.runs.stream(
thread_id=thread_id, thread_id=thread_id,
@@ -953,6 +998,11 @@ class RemoteGraph(PregelProtocol):
if caller_ns := (config or {}).get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS): if caller_ns := (config or {}).get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS):
caller_ns = tuple(caller_ns.split(NS_SEP)) caller_ns = tuple(caller_ns.split(NS_SEP))
ns = caller_ns + ns ns = caller_ns + ns
mode, data = _merge_values_patch(ns, mode, chunk.data, values_by_ns)
if mode != chunk.event:
chunk = chunk._replace(data=data)
elif data is not chunk.data:
chunk = chunk._replace(data=data)
# stream to parent stream # stream to parent stream
if stream is not None and mode in stream.modes: if stream is not None and mode in stream.modes:
stream((ns, mode, chunk.data)) stream((ns, mode, chunk.data))
@@ -970,7 +1020,9 @@ class RemoteGraph(PregelProtocol):
continue continue
if chunk.event.startswith("messages"): if chunk.event.startswith("messages"):
chunk = chunk._replace(data=tuple(chunk.data)) chunk = chunk._replace(
data=_restore_message_metadata(chunk.data, message_metadata_by_id)
)
# emit chunk # emit chunk
if version == "v2": if version == "v2":
@@ -982,11 +1034,6 @@ class RemoteGraph(PregelProtocol):
) )
yield {"type": mode, "ns": ns, "data": chunk.data, "interrupts": ints} yield {"type": mode, "ns": ns, "data": chunk.data, "interrupts": ints}
elif subgraphs: elif subgraphs:
if NS_SEP in chunk.event:
mode, ns_ = chunk.event.split(NS_SEP, 1)
ns = tuple(ns_.split(NS_SEP))
else:
mode, ns = chunk.event, ()
if req_single: if req_single:
yield ns, chunk.data yield ns, chunk.data
else: else:
+13 -5
View File
@@ -116,7 +116,14 @@ def ensure_valid_checkpointer(checkpointer: Checkpointer) -> Checkpointer:
StreamMode = Literal[ StreamMode = Literal[
"values", "updates", "checkpoints", "tasks", "debug", "messages", "custom" "values",
"updates",
"checkpoints",
"tasks",
"debug",
"messages",
"custom",
"compact",
] ]
"""How the stream method should emit outputs. """How the stream method should emit outputs.
@@ -275,13 +282,14 @@ class MessagesStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="messages"`. """Stream part emitted for `stream_mode="messages"`.
`data` is a 2-tuple of `(message, metadata)` where `message` is a `data` is a 2-tuple of `(message, metadata)` where `message` is a
`BaseMessage` (e.g. `AIMessageChunk`) and `metadata` is a dict containing `BaseMessage` (e.g. `AIMessageChunk`) and `metadata` is either a dict containing
keys like `langgraph_step`, `langgraph_node`, `langgraph_triggers`, etc. keys like `langgraph_step`, `langgraph_node`, `langgraph_triggers`, etc. or
`None` for deduplicated follow-up chunks when `stream_mode` includes `"compact"`.
""" """
type: Literal["messages"] type: Literal["messages"]
ns: tuple[str, ...] ns: tuple[str, ...]
data: tuple[AnyMessage, dict[str, Any]] data: tuple[AnyMessage, dict[str, Any] | None]
class CustomStreamPart(TypedDict): class CustomStreamPart(TypedDict):
@@ -346,7 +354,7 @@ async for part in graph.astream(input, version="v2"):
if part["type"] == "values": if part["type"] == "values":
part["data"] # OutputT — full state (pydantic/dataclass/dict) part["data"] # OutputT — full state (pydantic/dataclass/dict)
elif part["type"] == "messages": elif part["type"] == "messages":
part["data"] # tuple[BaseMessage, dict] — (message, metadata) part["data"] # tuple[BaseMessage, dict | None] — (message, metadata)
elif part["type"] == "custom": elif part["type"] == "custom":
part["data"] # Any — user-defined part["data"] # Any — user-defined
``` ```
+151
View File
@@ -882,6 +882,80 @@ def test_stream_sanitizes_thread_id():
assert not passed_config["configurable"] assert not passed_config["configurable"]
def test_stream_restores_messages_and_merges_values_patch():
mock_sync_client = MagicMock()
mock_sync_client.runs.stream.return_value = [
StreamPart(
event="messages|tools:call_1",
data=[
{"id": "msg-1", "type": "AIMessageChunk", "content": "hel"},
{
"langgraph_checkpoint_ns": "tools:call_1",
"langgraph_node": "agent",
},
],
),
StreamPart(
event="messages|tools:call_1",
data=[
{"id": "msg-1", "type": "AIMessageChunk", "content": "lo"},
None,
],
),
StreamPart(
event="values|tools:call_1",
data={"messages": [{"type": "human", "content": "hi"}], "count": 1},
),
StreamPart(
event="values-patch|tools:call_1",
data={"values": {"count": 2}, "deleted_keys": ["messages"]},
),
]
remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client)
parts = list(
remote_pregel.stream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode=["messages", "values", "compact"],
subgraphs=True,
version="v2",
)
)
message_parts = [part for part in parts if part["type"] == "messages"]
assert message_parts[0]["data"][1] == {
"langgraph_checkpoint_ns": "tools:call_1",
"langgraph_node": "agent",
}
assert message_parts[1]["data"][1] == {
"langgraph_checkpoint_ns": "tools:call_1",
"langgraph_node": "agent",
}
value_parts = [part for part in parts if part["type"] == "values"]
assert value_parts == [
{
"type": "values",
"ns": ("tools:call_1",),
"data": {"messages": [{"type": "human", "content": "hi"}], "count": 1},
"interrupts": (),
},
{
"type": "values",
"ns": ("tools:call_1",),
"data": {"count": 2},
"interrupts": (),
},
]
_, kwargs = mock_sync_client.runs.stream.call_args
assert set(kwargs["stream_mode"]) == {
"messages-tuple",
"values",
"compact",
"updates",
}
@pytest.mark.anyio @pytest.mark.anyio
async def test_ainvoke(): async def test_ainvoke():
# set up test # set up test
@@ -1092,6 +1166,83 @@ def test_stream_context_base_model():
assert kwargs["context"] == ctx assert kwargs["context"] == ctx
@pytest.mark.anyio
async def test_astream_restores_messages_and_merges_values_patch():
mock_async_client = MagicMock()
async_iter = MagicMock()
async_iter.__aiter__.return_value = [
StreamPart(
event="messages|tools:call_1",
data=[
{"id": "msg-1", "type": "AIMessageChunk", "content": "hel"},
{
"langgraph_checkpoint_ns": "tools:call_1",
"langgraph_node": "agent",
},
],
),
StreamPart(
event="messages|tools:call_1",
data=[
{"id": "msg-1", "type": "AIMessageChunk", "content": "lo"},
None,
],
),
StreamPart(
event="values|tools:call_1",
data={"messages": [{"type": "human", "content": "hi"}], "count": 1},
),
StreamPart(
event="values-patch|tools:call_1",
data={"values": {"count": 2}, "deleted_keys": ["messages"]},
),
]
mock_async_client.runs.stream.return_value = async_iter
remote_pregel = RemoteGraph("test_graph_id", client=mock_async_client)
parts = []
async for part in remote_pregel.astream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode=["messages", "values", "compact"],
subgraphs=True,
version="v2",
):
parts.append(part)
message_parts = [part for part in parts if part["type"] == "messages"]
assert message_parts[0]["data"][1] == {
"langgraph_checkpoint_ns": "tools:call_1",
"langgraph_node": "agent",
}
assert message_parts[1]["data"][1] == {
"langgraph_checkpoint_ns": "tools:call_1",
"langgraph_node": "agent",
}
value_parts = [part for part in parts if part["type"] == "values"]
assert value_parts == [
{
"type": "values",
"ns": ("tools:call_1",),
"data": {"messages": [{"type": "human", "content": "hi"}], "count": 1},
"interrupts": (),
},
{
"type": "values",
"ns": ("tools:call_1",),
"data": {"count": 2},
"interrupts": (),
},
]
_, kwargs = mock_async_client.runs.stream.call_args
assert set(kwargs["stream_mode"]) == {
"messages-tuple",
"values",
"compact",
"updates",
}
@pytest.mark.skip( @pytest.mark.skip(
"Unskip this test to manually test the LangSmith Deployment integration" "Unskip this test to manually test the LangSmith Deployment integration"
) )
+82
View File
@@ -90,6 +90,25 @@ def _make_messages_graph() -> StateGraph[
return builder return builder
def _make_streaming_messages_graph() -> StateGraph[
MessagesState, None, MessagesState, MessagesState
]:
model = FakeChatModel(messages=[AIMessage(content="hello world", id="ai-1")])
def call_model(state: MessagesState) -> dict[str, Any]:
streamed = model.stream(state["messages"])
message = next(streamed)
for chunk in streamed:
message += chunk
return {"messages": message}
builder = StateGraph(MessagesState, input_schema=MessagesState)
builder.add_node("call_model", call_model)
builder.add_edge(START, "call_model")
builder.add_edge("call_model", END)
return builder
def _make_custom_graph() -> Any: def _make_custom_graph() -> Any:
@entrypoint() @entrypoint()
def graph(inputs: Any, *, writer: StreamWriter) -> Any: def graph(inputs: Any, *, writer: StreamWriter) -> Any:
@@ -164,6 +183,26 @@ class TestV1BackwardsCompat:
ns, _data = chunk ns, _data = chunk
assert isinstance(ns, tuple) assert isinstance(ns, tuple)
def test_stream_v1_messages_keep_metadata_on_every_chunk(self) -> None:
graph = _make_streaming_messages_graph().compile()
chunks = list(graph.stream(_MSG_INPUT, stream_mode="messages"))
metadata = [meta for _message, meta in chunks]
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 --- # --- v2 sync stream ---
@@ -205,6 +244,26 @@ class TestV2Stream:
assert isinstance(metadata, dict) assert isinstance(metadata, dict)
assert "langgraph_node" in metadata assert "langgraph_node" in metadata
def test_messages_streaming_compact_dedupes_metadata(self) -> None:
graph = _make_streaming_messages_graph().compile()
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"]
assert isinstance(first_message, BaseMessage)
assert isinstance(first_metadata, dict)
assert "langgraph_node" in first_metadata
for chunk in msg_chunks[1:]:
message, metadata = chunk["data"]
assert isinstance(message, BaseMessage)
assert metadata is None
def test_custom(self) -> None: def test_custom(self) -> None:
graph = _make_custom_graph() graph = _make_custom_graph()
chunks = list(graph.stream({"key": "val"}, stream_mode="custom", version="v2")) chunks = list(graph.stream({"key": "val"}, stream_mode="custom", version="v2"))
@@ -544,6 +603,29 @@ class TestV2StreamAsync:
assert isinstance(metadata, dict) assert isinstance(metadata, dict)
assert "langgraph_node" in metadata assert "langgraph_node" in metadata
@NEEDS_CONTEXTVARS
@pytest.mark.anyio
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", "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"]
assert isinstance(first_message, BaseMessage)
assert isinstance(first_metadata, dict)
assert "langgraph_node" in first_metadata
for chunk in msg_chunks[1:]:
message, metadata = chunk["data"]
assert isinstance(message, BaseMessage)
assert metadata is None
@NEEDS_CONTEXTVARS @NEEDS_CONTEXTVARS
@pytest.mark.anyio @pytest.mark.anyio
async def test_custom(self) -> None: async def test_custom(self) -> None:
+23
View File
@@ -58,6 +58,7 @@ StreamMode = Literal[
"debug", "debug",
"custom", "custom",
"messages-tuple", "messages-tuple",
"compact",
] ]
""" """
Defines the mode of streaming: Defines the mode of streaming:
@@ -69,6 +70,7 @@ Defines the mode of streaming:
- "tasks": Stream task start and finish events. - "tasks": Stream task start and finish events.
- "debug": Stream detailed debug information. - "debug": Stream detailed debug information.
- "custom": Stream custom events. - "custom": Stream custom events.
- "compact": Enable compact streaming payloads for other selected modes.
""" """
DisconnectMode = Literal["cancel", "continue"] DisconnectMode = Literal["cancel", "continue"]
@@ -733,6 +735,26 @@ class ValuesStreamPart(TypedDict):
"""List of interrupts that occurred during this step.""" """List of interrupts that occurred during this step."""
class ValuesPatchPayload(TypedDict):
"""Incremental patch payload for subgraph `values` updates."""
values: dict[str, Any]
"""Only the changed fields since the previous `values` event for this namespace."""
deleted_keys: NotRequired[list[str]]
"""Optional list of keys that were removed from the previous values snapshot."""
class ValuesPatchStreamPart(TypedDict):
"""Stream part emitted for incremental subgraph value patches (`values-patch`)."""
type: Literal["values-patch"]
"""Stream part type discriminator."""
ns: list[str]
"""Namespace path of the emitting node (empty for root graph)."""
data: ValuesPatchPayload
"""Incremental state patch for the namespace."""
class UpdatesStreamPart(TypedDict): class UpdatesStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="updates"`.""" """Stream part emitted for `stream_mode="updates"`."""
@@ -845,6 +867,7 @@ class MetadataStreamPart(TypedDict):
StreamPartV2 = ( StreamPartV2 = (
ValuesStreamPart ValuesStreamPart
| ValuesPatchStreamPart
| UpdatesStreamPart | UpdatesStreamPart
| MessagesPartialStreamPart | MessagesPartialStreamPart
| MessagesCompleteStreamPart | MessagesCompleteStreamPart
+81
View File
@@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import json
from collections.abc import Iterator, Sequence from collections.abc import Iterator, Sequence
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
@@ -8,7 +9,9 @@ import httpx
import pytest import pytest
from typing_extensions import assert_type from typing_extensions import assert_type
from langgraph_sdk._async.runs import RunsClient
from langgraph_sdk._shared.utilities import _sse_to_v2_dict from langgraph_sdk._shared.utilities import _sse_to_v2_dict
from langgraph_sdk._sync.runs import SyncRunsClient
from langgraph_sdk.client import HttpClient, SyncHttpClient from langgraph_sdk.client import HttpClient, SyncHttpClient
from langgraph_sdk.schema import ( from langgraph_sdk.schema import (
CheckpointPayload, CheckpointPayload,
@@ -24,6 +27,7 @@ from langgraph_sdk.schema import (
TaskResultPayload, TaskResultPayload,
TasksStreamPart, TasksStreamPart,
UpdatesStreamPart, UpdatesStreamPart,
ValuesPatchStreamPart,
ValuesStreamPart, ValuesStreamPart,
) )
from langgraph_sdk.sse import BytesLike, BytesLineDecoder, SSEDecoder from langgraph_sdk.sse import BytesLike, BytesLineDecoder, SSEDecoder
@@ -375,6 +379,81 @@ def test_sse_to_v2_dict_values_with_interrupts() -> None:
assert "__interrupt__" not in result["data"] assert "__interrupt__" not in result["data"]
def test_sse_to_v2_dict_values_patch() -> None:
payload = {"values": {"count": 2}, "deleted_keys": ["stale"]}
result = _sse_to_v2_dict("values-patch|tools:call_1", payload)
assert result is not None
_assert_v2_shape(result)
assert result == {
"type": "values-patch",
"ns": ["tools:call_1"],
"data": {"values": {"count": 2}, "deleted_keys": ["stale"]},
"interrupts": [],
}
@pytest.mark.asyncio
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_mode"] == ["values", "compact"]
return httpx.Response(
200,
headers={"Content-Type": "text/event-stream"},
content=b"event: end\ndata: null\n\n",
)
transport = httpx.MockTransport(handler)
async with httpx.AsyncClient(
transport=transport, base_url="https://example.com"
) as client:
runs_client = RunsClient(HttpClient(client))
parts = [
part
async for part in runs_client.stream(
thread_id=None,
assistant_id="agent",
input={"messages": []},
stream_mode=["values", "compact"],
)
]
assert len(parts) == 1
assert parts[0].event == "end"
assert parts[0].data is None
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_mode"] == ["values", "compact"]
return httpx.Response(
200,
headers={"Content-Type": "text/event-stream"},
content=b"event: end\ndata: null\n\n",
)
transport = httpx.MockTransport(handler)
with httpx.Client(transport=transport, base_url="https://example.com") as client:
runs_client = SyncRunsClient(SyncHttpClient(client))
parts = list(
runs_client.stream(
thread_id=None,
assistant_id="agent",
input={"messages": []},
stream_mode=["values", "compact"],
)
)
assert len(parts) == 1
assert parts[0].event == "end"
assert parts[0].data is None
# --- client-side v2 stream wrapping --- # --- client-side v2 stream wrapping ---
@@ -448,6 +527,8 @@ def _check_v2_type_narrowing(part: StreamPartV2) -> None:
if part["type"] == "values": if part["type"] == "values":
assert_type(part, ValuesStreamPart) assert_type(part, ValuesStreamPart)
assert_type(part["data"], dict[str, Any]) assert_type(part["data"], dict[str, Any])
elif part["type"] == "values-patch":
assert_type(part, ValuesPatchStreamPart)
elif part["type"] == "updates": elif part["type"] == "updates":
assert_type(part, UpdatesStreamPart) assert_type(part, UpdatesStreamPart)
assert_type(part["data"], dict[str, Any]) assert_type(part["data"], dict[str, Any])