mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-13 21:27:52 +02:00
Bound EventLog / StreamChannel memory with drop-oldest semantics
- EventLog(maxlen=N) caps retention. When the buffer is full, push evicts the oldest item and advances an absolute _first_seq so cursors can detect they've fallen off the back. A lagging cursor raises BufferOverflowError on its next read — mirrors the restored=false signal from the reconnection scenario (§06). - New cursors start at the current head of the buffer, not seq 0. For unbounded logs this is indistinguishable from the old behavior; for bounded logs, new consumers see whatever is still retained. - StreamChannel(name, *, maxlen=N) forwards maxlen to its inner log. - StreamMux(..., max_events=N) sets a default maxlen for every log / channel it binds (main event log plus each transformer projection). Explicit per-log maxlen wins over the mux default. - StreamingHandler.stream() / astream() expose max_events: caller sets the run-wide memory budget; transformer authors can override per-log when they know better. Default unbounded, matching §15 Q3.
This commit is contained in:
@@ -4,7 +4,7 @@ Provides a ``StreamingHandler`` that wraps a compiled graph and exposes
|
||||
ergonomic streaming projections through a transformer pipeline.
|
||||
"""
|
||||
|
||||
from langgraph.stream._event_log import EventLog
|
||||
from langgraph.stream._event_log import BufferOverflowError, EventLog
|
||||
from langgraph.stream._types import ProtocolEvent, StreamTransformer
|
||||
from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream
|
||||
from langgraph.stream.stream_channel import StreamChannel
|
||||
@@ -12,6 +12,7 @@ from langgraph.stream.streaming_handler import StreamingHandler
|
||||
|
||||
__all__ = [
|
||||
"AsyncGraphRunStream",
|
||||
"BufferOverflowError",
|
||||
"EventLog",
|
||||
"GraphRunStream",
|
||||
"ProtocolEvent",
|
||||
|
||||
@@ -1,12 +1,23 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections import deque
|
||||
from collections.abc import AsyncIterator, Callable, Iterator
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class BufferOverflowError(RuntimeError):
|
||||
"""Raised when an EventLog cursor falls off the back of a bounded buffer.
|
||||
|
||||
Mirrors the ``restored: false`` signal from the protocol's reconnection
|
||||
story (§ 06): consumers that fall behind the retention window get an
|
||||
explicit error and can decide to rebuild from a snapshot rather than
|
||||
silently losing events.
|
||||
"""
|
||||
|
||||
|
||||
class EventLog(Generic[T]):
|
||||
"""Append-only buffer that supports multiple independent consumers.
|
||||
|
||||
@@ -28,10 +39,26 @@ class EventLog(Generic[T]):
|
||||
|
||||
Async iteration uses a shared ``asyncio.Event`` — cursors await
|
||||
the event when they catch up, and the producer sets it on each push.
|
||||
|
||||
Bounded mode
|
||||
------------
|
||||
Pass ``maxlen=N`` to cap memory. When the buffer is full, ``push``
|
||||
drops the oldest item to make room. Cursors track an absolute
|
||||
sequence number; a cursor that falls off the back of the retention
|
||||
window raises ``BufferOverflowError`` on its next read.
|
||||
|
||||
New cursors start at the current head of the buffer, not at seq 0
|
||||
— they see whatever is still retained. This matches the protocol's
|
||||
reconnection semantics (§ 06: "missed events can be replayed from
|
||||
a bounded buffer").
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._items: list[T] = []
|
||||
def __init__(self, maxlen: int | None = None) -> None:
|
||||
if maxlen is not None and maxlen <= 0:
|
||||
raise ValueError("EventLog maxlen must be a positive int or None")
|
||||
self._items: deque[T] = deque()
|
||||
self._maxlen: int | None = maxlen
|
||||
self._first_seq = 0 # absolute seq of _items[0]
|
||||
self._closed = False
|
||||
self._error: BaseException | None = None
|
||||
|
||||
@@ -65,9 +92,17 @@ class EventLog(Generic[T]):
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def push(self, item: T) -> None:
|
||||
"""Append *item* and wake all waiting cursors."""
|
||||
"""Append *item* and wake all waiting cursors.
|
||||
|
||||
In bounded mode, evicts the oldest item first if the buffer
|
||||
is full, advancing ``_first_seq`` so cursors can detect that
|
||||
they've fallen off the back of the retention window.
|
||||
"""
|
||||
if self._closed:
|
||||
raise RuntimeError("Cannot push to a closed EventLog")
|
||||
if self._maxlen is not None and len(self._items) >= self._maxlen:
|
||||
self._items.popleft()
|
||||
self._first_seq += 1
|
||||
self._items.append(item)
|
||||
self._notify()
|
||||
|
||||
@@ -112,11 +147,19 @@ class EventLog(Generic[T]):
|
||||
return self._sync_cursor()
|
||||
|
||||
def _sync_cursor(self) -> Iterator[T]:
|
||||
cursor = 0
|
||||
# Start at the current head — if maxlen is None this is 0 (seen everything),
|
||||
# if bounded this is wherever retention currently begins.
|
||||
seq = self._first_seq
|
||||
while True:
|
||||
if cursor < len(self._items):
|
||||
item = self._items[cursor]
|
||||
cursor += 1
|
||||
if seq < self._first_seq:
|
||||
raise BufferOverflowError(
|
||||
f"Cursor fell {self._first_seq - seq} items behind the "
|
||||
f"bounded EventLog's retention window (maxlen={self._maxlen})"
|
||||
)
|
||||
idx = seq - self._first_seq
|
||||
if idx < len(self._items):
|
||||
item = self._items[idx]
|
||||
seq += 1
|
||||
yield item
|
||||
elif self._closed:
|
||||
if self._error is not None:
|
||||
@@ -125,7 +168,7 @@ class EventLog(Generic[T]):
|
||||
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).
|
||||
while cursor >= len(self._items) and not self._closed:
|
||||
while (seq - self._first_seq) >= len(self._items) and not self._closed:
|
||||
if not self._request_more():
|
||||
break
|
||||
else:
|
||||
@@ -149,16 +192,22 @@ class EventLog(Generic[T]):
|
||||
|
||||
async def _async_cursor(self) -> AsyncIterator[T]:
|
||||
assert self._event is not None
|
||||
cursor = 0
|
||||
seq = self._first_seq
|
||||
while True:
|
||||
if cursor < len(self._items):
|
||||
yield self._items[cursor]
|
||||
cursor += 1
|
||||
if seq < self._first_seq:
|
||||
raise BufferOverflowError(
|
||||
f"Cursor fell {self._first_seq - seq} items behind the "
|
||||
f"bounded EventLog's retention window (maxlen={self._maxlen})"
|
||||
)
|
||||
idx = seq - self._first_seq
|
||||
if idx < len(self._items):
|
||||
yield self._items[idx]
|
||||
seq += 1
|
||||
elif self._closed:
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
return
|
||||
else:
|
||||
self._event.clear()
|
||||
if cursor >= len(self._items) and not self._closed:
|
||||
if (seq - self._first_seq) >= len(self._items) and not self._closed:
|
||||
await self._event.wait()
|
||||
|
||||
@@ -33,6 +33,7 @@ class StreamMux:
|
||||
transformers: list[StreamTransformer] | None = None,
|
||||
*,
|
||||
is_async: bool = False,
|
||||
max_events: int | None = None,
|
||||
) -> None:
|
||||
"""Initialize the mux and register *transformers* in order.
|
||||
|
||||
@@ -42,11 +43,18 @@ class StreamMux:
|
||||
keys are recorded in ``self.native_keys``, and any ``EventLog``
|
||||
/ ``StreamChannel`` instances are bound/wired.
|
||||
|
||||
*max_events* sets a default capacity for every ``EventLog`` /
|
||||
``StreamChannel`` the mux binds, including the main event log.
|
||||
Logs that were constructed with an explicit ``maxlen`` keep
|
||||
their own setting — the mux only fills in ``None`` defaults.
|
||||
Unbounded when ``max_events`` is ``None``.
|
||||
|
||||
Raises ``RuntimeError`` if any transformer requires an async run
|
||||
under sync mode, and ``ValueError`` on projection-key conflicts.
|
||||
"""
|
||||
self._is_async = is_async
|
||||
self._events: EventLog[ProtocolEvent] = EventLog()
|
||||
self._default_maxlen = max_events
|
||||
self._events: EventLog[ProtocolEvent] = EventLog(maxlen=max_events)
|
||||
self._events._bind(is_async=is_async)
|
||||
self._transformers: list[StreamTransformer] = []
|
||||
self._channels: list[StreamChannel[Any]] = []
|
||||
@@ -275,6 +283,7 @@ class StreamMux:
|
||||
"""Bind and wire EventLog / StreamChannel instances in *projection*."""
|
||||
for value in projection.values():
|
||||
if isinstance(value, StreamChannel):
|
||||
self._apply_default_maxlen(value._log)
|
||||
value._bind(is_async=self._is_async)
|
||||
self._channels.append(value)
|
||||
channel_name = value.name
|
||||
@@ -287,9 +296,15 @@ class StreamMux:
|
||||
|
||||
value._wire(_make_forward(channel_name))
|
||||
elif isinstance(value, EventLog):
|
||||
self._apply_default_maxlen(value)
|
||||
value._bind(is_async=self._is_async)
|
||||
self._logs.append(value)
|
||||
|
||||
def _apply_default_maxlen(self, log: EventLog[Any]) -> None:
|
||||
"""Fill in the mux's default maxlen if the log hasn't set its own."""
|
||||
if log._maxlen is None and self._default_maxlen is not None:
|
||||
log._maxlen = self._default_maxlen
|
||||
|
||||
def _forward(self, channel_name: str, item: Any) -> None:
|
||||
"""Inject a ProtocolEvent for a StreamChannel push.
|
||||
|
||||
|
||||
@@ -29,9 +29,9 @@ class StreamChannel(Generic[T]):
|
||||
using only StreamChannels don't need ``finalize`` / ``fail`` hooks.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str) -> None:
|
||||
def __init__(self, name: str, *, maxlen: int | None = None) -> None:
|
||||
self.name = name
|
||||
self._log: EventLog[T] = EventLog()
|
||||
self._log: EventLog[T] = EventLog(maxlen=maxlen)
|
||||
self._wire_fn: Callable[[T], None] | None = None
|
||||
|
||||
def _bind(self, *, is_async: bool) -> None:
|
||||
|
||||
@@ -56,6 +56,7 @@ class StreamingHandler:
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
transformers: list[StreamTransformer] | None = None,
|
||||
max_events: int | None = None,
|
||||
) -> GraphRunStream:
|
||||
"""Start a sync streaming run.
|
||||
|
||||
@@ -63,11 +64,19 @@ class StreamingHandler:
|
||||
any projection drives the graph forward — no background thread is
|
||||
used. This matches v1's model where the caller's ``for`` loop is
|
||||
the pump.
|
||||
|
||||
*max_events* caps the retention of every ``EventLog`` /
|
||||
``StreamChannel`` the mux binds (main event log plus each
|
||||
transformer's projection logs) to the given number of items,
|
||||
dropping the oldest when full. Transformers that constructed
|
||||
their own logs with an explicit ``maxlen`` keep their setting.
|
||||
Unbounded when ``None``.
|
||||
"""
|
||||
values_t = ValuesTransformer()
|
||||
mux = StreamMux(
|
||||
[values_t, MessagesTransformer(), *(transformers or ())],
|
||||
is_async=False,
|
||||
max_events=max_events,
|
||||
)
|
||||
|
||||
graph_iter = iter(
|
||||
@@ -92,16 +101,21 @@ class StreamingHandler:
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
transformers: list[StreamTransformer] | None = None,
|
||||
max_events: int | None = None,
|
||||
) -> AsyncGraphRunStream:
|
||||
"""Start an async streaming run.
|
||||
|
||||
Returns an `AsyncGraphRunStream` immediately. A background asyncio
|
||||
task pumps events from the graph into the transformer pipeline.
|
||||
|
||||
*max_events* caps retention of every ``EventLog`` / ``StreamChannel``
|
||||
the mux binds — see ``stream()`` for the full semantics.
|
||||
"""
|
||||
values_t = ValuesTransformer()
|
||||
mux = StreamMux(
|
||||
[values_t, MessagesTransformer(), *(transformers or ())],
|
||||
is_async=True,
|
||||
max_events=max_events,
|
||||
)
|
||||
|
||||
async def pump() -> None:
|
||||
|
||||
@@ -15,6 +15,7 @@ from typing_extensions import TypedDict
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.stream import (
|
||||
BufferOverflowError,
|
||||
EventLog,
|
||||
StreamChannel,
|
||||
StreamingHandler,
|
||||
@@ -1444,3 +1445,145 @@ class TestAsyncTransformerLane:
|
||||
_ = await run.output
|
||||
scores = [x async for x in run.extensions["scores"]]
|
||||
assert scores and all(s == 42 for s in scores)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Bounded EventLog / StreamChannel — memory caps and overflow semantics
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBoundedEventLog:
|
||||
def test_invalid_maxlen_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="positive int or None"):
|
||||
EventLog(maxlen=0)
|
||||
with pytest.raises(ValueError, match="positive int or None"):
|
||||
EventLog(maxlen=-3)
|
||||
|
||||
def test_unbounded_default_preserves_replay(self) -> None:
|
||||
"""Default EventLog (maxlen=None) still replays from seq 0."""
|
||||
log: EventLog[int] = EventLog()
|
||||
log._bind(is_async=False)
|
||||
for i in range(100):
|
||||
log.push(i)
|
||||
log.close()
|
||||
assert list(log) == list(range(100))
|
||||
|
||||
def test_bounded_drops_oldest_on_overflow(self) -> None:
|
||||
"""When bounded, pushing past maxlen evicts the oldest item."""
|
||||
log: EventLog[int] = EventLog(maxlen=3)
|
||||
log._bind(is_async=False)
|
||||
for i in range(5):
|
||||
log.push(i)
|
||||
log.close()
|
||||
# Only the last 3 survive; new cursors start at the current head.
|
||||
assert list(log) == [2, 3, 4]
|
||||
|
||||
def test_new_cursor_starts_at_head_not_zero(self) -> None:
|
||||
"""New cursors see the retained window, not the evicted prefix."""
|
||||
log: EventLog[int] = EventLog(maxlen=2)
|
||||
log._bind(is_async=False)
|
||||
log.push(1)
|
||||
log.push(2)
|
||||
log.push(3) # evicts 1
|
||||
log.close()
|
||||
assert list(log) == [2, 3]
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_cursor_overflow_raises(self) -> None:
|
||||
"""An async cursor that falls behind the retention window raises."""
|
||||
log: EventLog[int] = EventLog(maxlen=2)
|
||||
log._bind(is_async=True)
|
||||
|
||||
log.push(1)
|
||||
cursor = aiter(log)
|
||||
# Advance cursor to seq 1, reading item 1.
|
||||
first = await anext(cursor)
|
||||
assert first == 1
|
||||
# Now push enough to roll the cursor off the back.
|
||||
log.push(2) # buffer: [1, 2] _first_seq=0, cursor at seq=1
|
||||
log.push(3) # buffer: [2, 3] _first_seq=1, cursor at seq=1 still OK
|
||||
log.push(4) # buffer: [3, 4] _first_seq=2, cursor at seq=1 — overflow
|
||||
with pytest.raises(BufferOverflowError, match="fell"):
|
||||
await anext(cursor)
|
||||
|
||||
def test_sync_cursor_sees_all_while_bounded_but_under_cap(self) -> None:
|
||||
"""Bounded mode with pushes under cap behaves identically to unbounded."""
|
||||
log: EventLog[int] = EventLog(maxlen=100)
|
||||
log._bind(is_async=False)
|
||||
log.push(1)
|
||||
log.push(2)
|
||||
log.close()
|
||||
assert list(log) == [1, 2]
|
||||
|
||||
|
||||
class TestStreamChannelMaxlen:
|
||||
def test_maxlen_passes_through_to_inner_log(self) -> None:
|
||||
ch: StreamChannel[int] = StreamChannel("ch", maxlen=2)
|
||||
ch._bind(is_async=False)
|
||||
ch.push(1)
|
||||
ch.push(2)
|
||||
ch.push(3)
|
||||
ch._close()
|
||||
assert list(ch) == [2, 3]
|
||||
|
||||
|
||||
class TestMuxMaxEventsDefault:
|
||||
def test_mux_fills_in_default_when_log_has_none(self) -> None:
|
||||
class Simple(StreamTransformer):
|
||||
def __init__(self) -> None:
|
||||
self.log: EventLog[int] = EventLog()
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"out": self.log}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
return True
|
||||
|
||||
t = Simple()
|
||||
StreamMux([t], max_events=10)
|
||||
assert t.log._maxlen == 10
|
||||
|
||||
def test_explicit_log_maxlen_wins_over_mux_default(self) -> None:
|
||||
class Explicit(StreamTransformer):
|
||||
def __init__(self) -> None:
|
||||
self.log: EventLog[int] = EventLog(maxlen=3)
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"out": self.log}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
return True
|
||||
|
||||
t = Explicit()
|
||||
StreamMux([t], max_events=1000)
|
||||
assert t.log._maxlen == 3 # transformer author's setting stands
|
||||
|
||||
def test_main_event_log_respects_max_events(self) -> None:
|
||||
mux = StreamMux([], max_events=5)
|
||||
assert mux._events._maxlen == 5
|
||||
|
||||
def test_max_events_default_cascades_to_channels(self) -> None:
|
||||
class WithChannel(StreamTransformer):
|
||||
def __init__(self) -> None:
|
||||
self.ch: StreamChannel[int] = StreamChannel("out")
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"out": self.ch}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
return True
|
||||
|
||||
t = WithChannel()
|
||||
StreamMux([t], max_events=7)
|
||||
assert t.ch._log._maxlen == 7
|
||||
|
||||
def test_handler_propagates_max_events(self) -> None:
|
||||
graph = _build_simple_graph()
|
||||
handler = StreamingHandler(graph)
|
||||
run = handler.stream({"value": "x", "items": []}, max_events=50)
|
||||
# Main log inherits the default.
|
||||
assert run._mux._events._maxlen == 50
|
||||
# Native projections (values log, messages log) inherit too.
|
||||
for log in run._mux._logs:
|
||||
assert log._maxlen == 50
|
||||
_ = run.output
|
||||
|
||||
Reference in New Issue
Block a user