mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-30 11:49:38 +02:00
Tighten streaming run stream API and close review footguns
- AsyncGraphRunStream.output/interrupted/interrupts are now methods (await run.output()), not properties returning coroutines. Forgetting `await` now fails at type-check time and at runtime on the common operations (bool/len/iter), instead of silently yielding a live coroutine that's truthy, lenless, and never awaited. - interrupted/interrupts re-raise the run's error on both lanes so a failed run doesn't silently return the last-known interrupt state. - Narrow the async pump catch from BaseException to Exception so CancelledError / KeyboardInterrupt propagate. - Wrap run.extensions with types.MappingProxyType so users can't add or remove projection keys behind the mux's back. - Add ValuesTransformer.error accessor; run stream stops reaching into _log._error. - Tighten StreamingHandler graph type from Any to Pregel and widen convert_to_protocol_event to accept StreamPart. - Projection-conflict ValueError now names the transformer that owns each colliding key, not just the incoming transformer. - Replace deprecated asyncio.get_event_loop() in the async iteration test with asyncio.create_task. - Document wall-clock semantics of ProtocolEvent.params.timestamp, the subgraph-namespace drop in MessagesTransformer, and the transformer-pipeline bypass for StreamChannel auto-forwarded events. - Add tests for the new error-raising behavior on interrupted / interrupts and for the read-only extensions contract.
This commit is contained in:
Nick Hollon
@@ -1,30 +1,32 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Any
|
||||
from typing import Any, cast
|
||||
|
||||
from langgraph.stream._types import ProtocolEvent, _ProtocolEventParams
|
||||
from langgraph.types import StreamPart
|
||||
|
||||
|
||||
def convert_to_protocol_event(part: dict[str, Any]) -> ProtocolEvent:
|
||||
"""Convert a v2 StreamPart dict to a ProtocolEvent.
|
||||
def convert_to_protocol_event(part: StreamPart) -> ProtocolEvent:
|
||||
"""Convert a v2 StreamPart to a ProtocolEvent.
|
||||
|
||||
Args:
|
||||
part: A stream part dict with keys `type`, `ns`, `data`, and
|
||||
part: A stream part with keys `type`, `ns`, `data`, and
|
||||
optionally `interrupts` (present on values events).
|
||||
|
||||
Returns:
|
||||
The equivalent ProtocolEvent.
|
||||
"""
|
||||
part_dict = cast(dict[str, Any], part)
|
||||
params: _ProtocolEventParams = {
|
||||
"namespace": list(part["ns"]),
|
||||
"namespace": list(part_dict["ns"]),
|
||||
"timestamp": int(time.time() * 1000),
|
||||
"data": part["data"],
|
||||
"data": part_dict["data"],
|
||||
}
|
||||
if "interrupts" in part:
|
||||
params["interrupts"] = part["interrupts"]
|
||||
if "interrupts" in part_dict:
|
||||
params["interrupts"] = part_dict["interrupts"]
|
||||
return {
|
||||
"type": "event",
|
||||
"method": part["type"],
|
||||
"method": part_dict["type"],
|
||||
"params": params,
|
||||
}
|
||||
|
||||
@@ -198,7 +198,10 @@ class EventLog(Generic[T]):
|
||||
return
|
||||
elif self._request_more is not None:
|
||||
# Pull from the producer until this log gets a new item
|
||||
# or the graph is exhausted (which closes the log).
|
||||
# or the graph is exhausted (which closes the log). A push
|
||||
# may evict the item this cursor was about to read; in that
|
||||
# case the inner loop breaks and the outer `seq < _first_seq`
|
||||
# check catches the overflow on the next iteration.
|
||||
while (seq - self._first_seq) >= len(self._items) and not self._closed:
|
||||
if not self._request_more():
|
||||
break
|
||||
|
||||
@@ -78,6 +78,7 @@ class StreamMux:
|
||||
|
||||
self.extensions: dict[str, Any] = {}
|
||||
self.native_keys: set[str] = set()
|
||||
self._projection_owners: dict[str, str] = {}
|
||||
|
||||
for transformer in transformers or ():
|
||||
self._register(transformer)
|
||||
@@ -103,14 +104,21 @@ class StreamMux:
|
||||
)
|
||||
conflicts = set(projection) & set(self.extensions)
|
||||
if conflicts:
|
||||
attributions = ", ".join(
|
||||
f"{key!r} (owned by {self._projection_owners[key]})"
|
||||
for key in sorted(conflicts)
|
||||
)
|
||||
raise ValueError(
|
||||
f"Transformer {type(transformer).__name__} returned "
|
||||
f"projection keys that conflict with already-registered "
|
||||
f"keys: {conflicts}"
|
||||
f"keys: {attributions}"
|
||||
)
|
||||
self._transformers.append(transformer)
|
||||
self._bind_and_wire(projection)
|
||||
self.extensions.update(projection)
|
||||
owner_name = type(transformer).__name__
|
||||
for key in projection:
|
||||
self._projection_owners[key] = owner_name
|
||||
if getattr(transformer, "_native", False):
|
||||
self.native_keys.update(projection.keys())
|
||||
|
||||
|
||||
@@ -12,7 +12,12 @@ _logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class _ProtocolEventParams(TypedDict):
|
||||
"""Parameters for a protocol event."""
|
||||
"""Parameters for a protocol event.
|
||||
|
||||
`timestamp` is wall-clock milliseconds since the epoch and can go
|
||||
backwards across NTP adjustments — use `ProtocolEvent.seq` for
|
||||
ordering.
|
||||
"""
|
||||
|
||||
namespace: list[str]
|
||||
timestamp: int
|
||||
@@ -25,6 +30,8 @@ class ProtocolEvent(TypedDict):
|
||||
|
||||
Wraps a raw stream part (values, messages, custom, etc.) in a uniform
|
||||
envelope with a monotonic sequence number assigned by the StreamMux.
|
||||
Consumers that need a total order across events should use `seq`, not
|
||||
`params.timestamp` (which is wall-clock and not monotonic).
|
||||
"""
|
||||
|
||||
type: Literal["event"]
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from collections.abc import AsyncIterator, Iterator, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Any
|
||||
|
||||
from langgraph.stream._convert import convert_to_protocol_event
|
||||
@@ -44,7 +45,7 @@ class GraphRunStream:
|
||||
"""
|
||||
self._graph_iter = graph_iter
|
||||
self._mux = mux
|
||||
self.extensions = mux.extensions
|
||||
self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions)
|
||||
self._values_transformer = values_transformer
|
||||
self._exhausted = False
|
||||
# Native-transformer projections also show up as direct attributes.
|
||||
@@ -98,20 +99,35 @@ class GraphRunStream:
|
||||
def output(self) -> dict[str, Any] | None:
|
||||
"""Block until the run completes and return the final state."""
|
||||
self._pump_all()
|
||||
if self._values_transformer._log._error is not None:
|
||||
raise self._values_transformer._log._error
|
||||
err = self._values_transformer.error
|
||||
if err is not None:
|
||||
raise err
|
||||
return self._values_transformer._latest
|
||||
|
||||
@property
|
||||
def interrupted(self) -> bool:
|
||||
"""Block until the run completes, then return whether it was interrupted."""
|
||||
"""Block until the run completes, then return whether it was interrupted.
|
||||
|
||||
Raises:
|
||||
BaseException: If the run ended with an error.
|
||||
"""
|
||||
self._pump_all()
|
||||
err = self._values_transformer.error
|
||||
if err is not None:
|
||||
raise err
|
||||
return self._values_transformer._interrupted
|
||||
|
||||
@property
|
||||
def interrupts(self) -> list[Any]:
|
||||
"""Block until the run completes, then return interrupt payloads."""
|
||||
"""Block until the run completes, then return interrupt payloads.
|
||||
|
||||
Raises:
|
||||
BaseException: If the run ended with an error.
|
||||
"""
|
||||
self._pump_all()
|
||||
err = self._values_transformer.error
|
||||
if err is not None:
|
||||
raise err
|
||||
return self._values_transformer._interrupts
|
||||
|
||||
def __iter__(self) -> Iterator[ProtocolEvent]:
|
||||
@@ -150,67 +166,76 @@ class AsyncGraphRunStream:
|
||||
pump_task: Background task pumping graph events into the mux.
|
||||
"""
|
||||
self._mux = mux
|
||||
self.extensions = mux.extensions
|
||||
self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions)
|
||||
self._values_transformer = values_transformer
|
||||
self._pump_task = pump_task
|
||||
# Native-transformer projections also show up as direct attributes.
|
||||
for key in mux.native_keys:
|
||||
setattr(self, key, mux.extensions[key])
|
||||
|
||||
@property
|
||||
def output(self) -> Any:
|
||||
"""Return an awaitable that resolves to the final state.
|
||||
async def output(self) -> dict[str, Any] | None:
|
||||
"""Wait for the run to complete and return the final state.
|
||||
|
||||
Methods (not properties) on the async lane so `run.output` without
|
||||
`await` raises at type-check time instead of silently yielding a
|
||||
coroutine object that's truthy, lenless, and never awaited.
|
||||
|
||||
The pump routes any Exception into `mux.afail`, which surfaces on
|
||||
`ValuesTransformer.error`. CancelledError / KeyboardInterrupt
|
||||
propagate so cancellation isn't silently dropped.
|
||||
|
||||
Example:
|
||||
```python
|
||||
output = await run.output
|
||||
output = await run.output()
|
||||
```
|
||||
"""
|
||||
return self._get_output()
|
||||
|
||||
async def _get_output(self) -> dict[str, Any] | None:
|
||||
Raises:
|
||||
BaseException: If the run ended with an error.
|
||||
"""
|
||||
try:
|
||||
await self._pump_task
|
||||
except BaseException:
|
||||
except Exception:
|
||||
pass
|
||||
if self._values_transformer._log._error is not None:
|
||||
raise self._values_transformer._log._error
|
||||
if (err := self._values_transformer.error) is not None:
|
||||
raise err
|
||||
return self._values_transformer._latest
|
||||
|
||||
@property
|
||||
def interrupted(self) -> Any:
|
||||
"""Return an awaitable that resolves to whether the run was interrupted.
|
||||
async def interrupted(self) -> bool:
|
||||
"""Wait for the run to complete and return whether it was interrupted.
|
||||
|
||||
Example:
|
||||
```python
|
||||
interrupted = await run.interrupted
|
||||
interrupted = await run.interrupted()
|
||||
```
|
||||
"""
|
||||
return self._get_interrupted()
|
||||
|
||||
async def _get_interrupted(self) -> bool:
|
||||
Raises:
|
||||
BaseException: If the run ended with an error.
|
||||
"""
|
||||
try:
|
||||
await self._pump_task
|
||||
except BaseException:
|
||||
except Exception:
|
||||
pass
|
||||
if (err := self._values_transformer.error) is not None:
|
||||
raise err
|
||||
return self._values_transformer._interrupted
|
||||
|
||||
@property
|
||||
def interrupts(self) -> Any:
|
||||
"""Return an awaitable that resolves to interrupt payloads.
|
||||
async def interrupts(self) -> list[Any]:
|
||||
"""Wait for the run to complete and return interrupt payloads.
|
||||
|
||||
Example:
|
||||
```python
|
||||
interrupts = await run.interrupts
|
||||
interrupts = await run.interrupts()
|
||||
```
|
||||
"""
|
||||
return self._get_interrupts()
|
||||
|
||||
async def _get_interrupts(self) -> list[Any]:
|
||||
Raises:
|
||||
BaseException: If the run ended with an error.
|
||||
"""
|
||||
try:
|
||||
await self._pump_task
|
||||
except BaseException:
|
||||
except Exception:
|
||||
pass
|
||||
if (err := self._values_transformer.error) is not None:
|
||||
raise err
|
||||
return self._values_transformer._interrupts
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[ProtocolEvent]:
|
||||
|
||||
@@ -17,6 +17,14 @@ class StreamChannel(Generic[T]):
|
||||
`ProtocolEvent` into the main event stream using the channel's
|
||||
name as the method.
|
||||
|
||||
Auto-forwarded events bypass the transformer pipeline — other
|
||||
transformers' `process()` / `aprocess()` methods do not see
|
||||
`custom:<name>` events produced by a channel push. This prevents a
|
||||
transformer that pushes to its own channel during `process()` from
|
||||
re-triggering itself, but it also means filter- or tap-style
|
||||
transformers cannot observe channel output from peer transformers.
|
||||
Consumers that need that should iterate the main event stream.
|
||||
|
||||
In-process consumers iterate the channel directly (`for item in ch`
|
||||
or `async for item in ch`). Remote SDK clients subscribe via
|
||||
`session.subscribe("custom:<channelName>")`.
|
||||
|
||||
@@ -6,6 +6,7 @@ from typing import Any
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.stream._convert import convert_to_protocol_event
|
||||
from langgraph.stream._mux import StreamMux
|
||||
from langgraph.stream._types import StreamTransformer
|
||||
@@ -38,15 +39,16 @@ class StreamingHandler:
|
||||
print(state)
|
||||
output = run.output
|
||||
|
||||
# Async
|
||||
# Async — terminal accessors are methods so a missing `await`
|
||||
# fails loudly instead of silently yielding a coroutine.
|
||||
run = await handler.astream(input_data)
|
||||
async for state in run.values:
|
||||
print(state)
|
||||
output = await run.output
|
||||
output = await run.output()
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, graph: Any) -> None:
|
||||
def __init__(self, graph: Pregel) -> None:
|
||||
"""Initialize the handler.
|
||||
|
||||
Args:
|
||||
|
||||
@@ -11,6 +11,9 @@ class ValuesTransformer(StreamTransformer):
|
||||
|
||||
Native transformer — projection keys are exposed as direct
|
||||
attributes on the run stream (e.g. `run.values`).
|
||||
|
||||
Only root-namespace values events are captured; subgraph state
|
||||
snapshots are ignored.
|
||||
"""
|
||||
|
||||
_native = True
|
||||
@@ -24,11 +27,18 @@ class ValuesTransformer(StreamTransformer):
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"values": self._log}
|
||||
|
||||
@property
|
||||
def error(self) -> BaseException | None:
|
||||
"""The error that ended the run, or `None` if it succeeded.
|
||||
|
||||
Set by the mux when it auto-fails the projection log.
|
||||
"""
|
||||
return self._log._error
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
if event["method"] != "values":
|
||||
return True
|
||||
params = event["params"]
|
||||
# Only capture root namespace events
|
||||
if params["namespace"]:
|
||||
return True
|
||||
self._latest = params["data"]
|
||||
@@ -47,6 +57,11 @@ class MessagesTransformer(StreamTransformer):
|
||||
A follow-on PR will replace this with a richer transformer that
|
||||
produces ChatModelStream objects using the protocol handler.
|
||||
|
||||
Only root-namespace messages events are captured; tokens emitted
|
||||
from subgraphs are dropped from the `messages` projection. Consumers
|
||||
that need subgraph tokens should iterate the raw event stream or
|
||||
register a custom transformer.
|
||||
|
||||
Native transformer — projection keys are exposed as direct
|
||||
attributes on the run stream (e.g. `run.messages`).
|
||||
"""
|
||||
@@ -63,7 +78,6 @@ class MessagesTransformer(StreamTransformer):
|
||||
if event["method"] != "messages":
|
||||
return True
|
||||
params = event["params"]
|
||||
# Only capture root namespace events
|
||||
if params["namespace"]:
|
||||
return True
|
||||
self._log.push(params["data"])
|
||||
|
||||
@@ -176,8 +176,9 @@ class TestEventLog:
|
||||
log.push(i)
|
||||
log.close()
|
||||
|
||||
asyncio.get_event_loop().call_soon(lambda: asyncio.ensure_future(producer()))
|
||||
producer_task = asyncio.create_task(producer())
|
||||
items = [item async for item in log]
|
||||
await producer_task
|
||||
assert items == [0, 1, 2]
|
||||
|
||||
@pytest.mark.anyio
|
||||
@@ -420,6 +421,16 @@ class TestStreamingHandlerSync:
|
||||
assert run.values is run.extensions["values"]
|
||||
assert run.messages is run.extensions["messages"]
|
||||
|
||||
def test_extensions_is_read_only(self) -> None:
|
||||
"""`run.extensions` must reject mutations so users can't corrupt mux state."""
|
||||
graph = _build_simple_graph()
|
||||
handler = StreamingHandler(graph)
|
||||
run = handler.stream({"value": "x", "items": []})
|
||||
with pytest.raises(TypeError):
|
||||
run.extensions["new_key"] = object() # type: ignore[index]
|
||||
with pytest.raises(TypeError):
|
||||
del run.extensions["values"] # type: ignore[attr-defined]
|
||||
|
||||
def test_custom_stream_events(self) -> None:
|
||||
graph = _build_custom_stream_graph()
|
||||
handler = StreamingHandler(graph)
|
||||
@@ -452,6 +463,22 @@ class TestStreamingHandlerSyncErrors:
|
||||
with pytest.raises(ValueError, match="boom"):
|
||||
list(run)
|
||||
|
||||
def test_error_propagation_interrupted(self) -> None:
|
||||
"""`run.interrupted` should raise on a failed run, not silently return False."""
|
||||
graph = _build_error_graph()
|
||||
handler = StreamingHandler(graph)
|
||||
run = handler.stream({"value": "x", "items": []})
|
||||
with pytest.raises(ValueError, match="boom"):
|
||||
_ = run.interrupted
|
||||
|
||||
def test_error_propagation_interrupts(self) -> None:
|
||||
"""`run.interrupts` should raise on a failed run."""
|
||||
graph = _build_error_graph()
|
||||
handler = StreamingHandler(graph)
|
||||
run = handler.stream({"value": "x", "items": []})
|
||||
with pytest.raises(ValueError, match="boom"):
|
||||
_ = run.interrupts
|
||||
|
||||
|
||||
class TestStreamingHandlerSyncInterrupt:
|
||||
def test_interrupted(self) -> None:
|
||||
@@ -490,7 +517,7 @@ class TestStreamingHandlerAsync:
|
||||
graph = _build_simple_graph()
|
||||
handler = StreamingHandler(graph)
|
||||
run = await handler.astream({"value": "x", "items": []})
|
||||
output = await run.output
|
||||
output = await run.output()
|
||||
assert output is not None
|
||||
assert output["value"] == "xAB"
|
||||
assert output["items"] == ["a", "b"]
|
||||
@@ -512,7 +539,7 @@ class TestStreamingHandlerAsync:
|
||||
graph = _build_simple_graph()
|
||||
handler = StreamingHandler(graph)
|
||||
run = await handler.astream({"value": "x", "items": []})
|
||||
_ = await run.output
|
||||
_ = await run.output()
|
||||
assert "values" in run.extensions
|
||||
assert "messages" in run.extensions
|
||||
assert run.values is run.extensions["values"]
|
||||
@@ -527,7 +554,7 @@ class TestStreamingHandlerAsyncErrors:
|
||||
handler = StreamingHandler(graph)
|
||||
run = await handler.astream({"value": "x", "items": []})
|
||||
with pytest.raises(ValueError, match="boom"):
|
||||
await run.output
|
||||
await run.output()
|
||||
|
||||
@pytest.mark.anyio
|
||||
@NEEDS_CONTEXTVARS
|
||||
@@ -549,6 +576,26 @@ class TestStreamingHandlerAsyncErrors:
|
||||
async for _ in run:
|
||||
pass
|
||||
|
||||
@pytest.mark.anyio
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_error_propagation_interrupted(self) -> None:
|
||||
"""`await run.interrupted()` should raise on a failed async run."""
|
||||
graph = _build_error_graph()
|
||||
handler = StreamingHandler(graph)
|
||||
run = await handler.astream({"value": "x", "items": []})
|
||||
with pytest.raises(ValueError, match="boom"):
|
||||
await run.interrupted()
|
||||
|
||||
@pytest.mark.anyio
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_error_propagation_interrupts(self) -> None:
|
||||
"""`await run.interrupts()` should raise on a failed async run."""
|
||||
graph = _build_error_graph()
|
||||
handler = StreamingHandler(graph)
|
||||
run = await handler.astream({"value": "x", "items": []})
|
||||
with pytest.raises(ValueError, match="boom"):
|
||||
await run.interrupts()
|
||||
|
||||
|
||||
class TestStreamingHandlerAsyncInterrupt:
|
||||
@pytest.mark.anyio
|
||||
@@ -560,9 +607,9 @@ class TestStreamingHandlerAsyncInterrupt:
|
||||
{"value": "x", "items": []},
|
||||
{"configurable": {"thread_id": "t2"}},
|
||||
)
|
||||
_ = await run.output
|
||||
assert await run.interrupted is True
|
||||
assert len(await run.interrupts) > 0
|
||||
_ = await run.output()
|
||||
assert await run.interrupted() is True
|
||||
assert len(await run.interrupts()) > 0
|
||||
|
||||
|
||||
class TestStreamingHandlerAsyncCustom:
|
||||
@@ -1042,7 +1089,10 @@ class TestCustomTransformer:
|
||||
|
||||
graph = _build_simple_graph()
|
||||
handler = StreamingHandler(graph)
|
||||
with pytest.raises(ValueError, match="conflict.*{'values'}"):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"conflict.*'values'.*ValuesTransformer",
|
||||
):
|
||||
handler.stream(
|
||||
{"value": "x", "items": []},
|
||||
transformers=[ConflictTransformer()],
|
||||
@@ -1442,7 +1492,7 @@ class TestAsyncTransformerLane:
|
||||
{"value": "x", "items": []},
|
||||
transformers=[Scorer()],
|
||||
)
|
||||
_ = await run.output
|
||||
_ = await run.output()
|
||||
scores = [x async for x in run.extensions["scores"]]
|
||||
assert scores and all(s == 42 for s in scores)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user