initial pass - surfacing interrupts for stream_mode='values'

This commit is contained in:
Sydney Runkle
2025-04-22 10:15:19 -07:00
parent 38d806733d
commit 9a7b1fa12a
3 changed files with 39 additions and 22 deletions
+11 -4
View File
@@ -2732,10 +2732,11 @@ class Pregel(PregelProtocol):
If stream_mode is not "values", it returns a list of output chunks.
"""
output_keys = output_keys if output_keys is not None else self.output_channels
if stream_mode == "values":
latest: Union[dict[str, Any], Any] = None
else:
chunks = []
latest: Union[dict[str, Any], Any] = None
chunks: list[Union[dict[str, Any], Any]] = []
interrupts: list[Interrupt] = []
for chunk in self.stream(
input,
config,
@@ -2748,10 +2749,16 @@ class Pregel(PregelProtocol):
**kwargs,
):
if stream_mode == "values":
if isinstance(chunk, dict) and (ints := chunk.get(INTERRUPT)) is not None:
interrupts.extend(ints)
latest = chunk
else:
chunks.append(chunk)
if stream_mode == "values":
if len(interrupts) > 0:
return {
INTERRUPT: interrupts
}
return latest
else:
return chunks
+17 -17
View File
@@ -910,23 +910,23 @@ class PregelLoop(LoopProtocol):
):
return
if writes[0][0] == INTERRUPT:
self._emit(
"updates",
lambda: iter(
[
{
INTERRUPT: tuple(
v
for w in writes
if w[0] == INTERRUPT
for v in (
w[1] if isinstance(w[1], Sequence) else (w[1],)
)
)
}
]
),
)
interrupts = [
{
INTERRUPT: tuple(
v
for w in writes
if w[0] == INTERRUPT
for v in (
w[1] if isinstance(w[1], Sequence) else (w[1],)
)
)
}
]
stream_modes = self.stream.modes if self.stream else []
if "updates" in stream_modes:
self._emit("updates", lambda: iter(interrupts))
elif "values" in stream_modes:
self._emit("values", lambda: iter(interrupts))
elif writes[0][0] != ERROR:
self._emit(
"updates",
+11 -1
View File
@@ -24,6 +24,8 @@ from typing_extensions import Self
from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
from langgraph.utils.fields import get_update_as_tuples
import hashlib
import uuid
if TYPE_CHECKING:
from langgraph.pregel.protocol import PregelProtocol
@@ -144,6 +146,13 @@ class Interrupt:
when: Literal["during"] = dataclasses.field(default="during", repr=False)
@property
def interrupt_id(self) -> str:
"""Generate a unique ID for the interrupt based on its namespace."""
identifier = uuid.uuid4().bytes if self.ns is None else ''.join(self.ns).encode()
return hashlib.sha256(identifier).hexdigest()
class StateUpdate(NamedTuple):
values: Optional[dict[str, Any]]
as_node: Optional[str] = None
@@ -483,11 +492,12 @@ def interrupt(value: Any) -> Any:
CONFIG_KEY_SEND,
NS_SEP,
RESUME,
CONF,
)
from langgraph.errors import GraphInterrupt
from langgraph.utils.config import get_config
conf = get_config()["configurable"]
conf = get_config()[CONF]
# track interrupt index
scratchpad: PregelScratchpad = conf[CONFIG_KEY_SCRATCHPAD]
idx = scratchpad.interrupt_counter()