WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-30 19:59:40 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: type safe stream/invoke w/ proper output type coercion (#6961)

Browse Source 
Adding type safe streaming + more robust pydantic + dataclass support on
graph outputs

* type safe streaming via #6931 
* type safe invoke via
https://github.com/langchain-ai/langgraph/pull/6963
* actually thread through types via parametrization in
https://github.com/langchain-ai/langgraph/pull/7009
* deprecation of backwards compatible accessor via
https://github.com/langchain-ai/langgraph/pull/7011

full spec of changes:
https://github.com/langchain-ai/langgraph/issues/7008
proving out that required changes are minimal even when we update the
default to v2: https://github.com/langchain-ai/langchain/pull/35541
This commit is contained in:
Sydney Runkle
2026-03-09 17:52:36 +00:00
committed by GitHub
parent e00a027579
commit 46fed9d161
14 changed files with 2892 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/graph/state.py
+25 -4
View File
@@ -6,6 +6,7 @@ import typing
import warnings
from collections import defaultdict
from collections.abc import Awaitable, Callable, Hashable, Sequence
from dataclasses import is_dataclass
from functools import partial
from inspect import isclass, isfunction, ismethod, signature
from types import FunctionType
@@ -14,6 +15,7 @@ from typing import (
Any,
Generic,
Literal,
TypeVar,
Union,
cast,
get_args,
@@ -1164,6 +1166,20 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
for key, node in self.nodes.items():
compiled.attach_node(key, node)
# Record output/state mappers for v2 stream coercion (pydantic/dataclass only)
compiled._output_mapper = _pick_mapper(
list(output_channels)
if isinstance(output_channels, list)
else [output_channels],
self.output_schema,
)
compiled._state_mapper = _pick_mapper(
list(stream_channels)
if isinstance(stream_channels, list)
else [stream_channels],
self.state_schema,
)
for start, end in self.edges:
compiled.attach_edge(start, end)
@@ -1183,6 +1199,8 @@ class CompiledStateGraph(
):
builder: StateGraph[StateT, ContextT, InputT, OutputT]
schema_to_mapper: dict[type[Any], Callable[[Any], Any] | None]
_output_mapper: Callable[[Any], Any] | None
_state_mapper: Callable[[Any], Any] | None
def __init__(
self,
@@ -1504,12 +1522,15 @@ def _pick_mapper(
) -> Callable[[Any], Any] | None:
if state_keys == ["__root__"]:
return None
if isclass(schema) and issubclass(schema, dict):
return None
return partial(_coerce_state, schema)
if isclass(schema) and (issubclass(schema, BaseModel) or is_dataclass(schema)):
return partial(_coerce_state, schema)
return None
def _coerce_state(schema: type[Any], input: dict[str, Any]) -> dict[str, Any]:
_S = TypeVar("_S")
def _coerce_state(schema: type[_S], input: dict[str, Any]) -> _S:
return schema(**input)
libs/langgraph/langgraph/pregel/debug.py
+8 -37
View File
@@ -7,7 +7,6 @@ from uuid import UUID
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import CheckpointMetadata, PendingWrite
from typing_extensions import TypedDict
from langgraph._internal._config import patch_checkpoint_map
from langgraph._internal._constants import (
@@ -23,42 +22,14 @@ from langgraph._internal._typing import MISSING
from langgraph.channels.base import BaseChannel
from langgraph.constants import TAG_HIDDEN
from langgraph.pregel._io import read_channels
from langgraph.types import PregelExecutableTask, PregelTask, StateSnapshot
__all__ = ("TaskPayload", "TaskResultPayload", "CheckpointTask", "CheckpointPayload")
class TaskPayload(TypedDict):
id: str
name: str
input: Any
triggers: list[str]
class TaskResultPayload(TypedDict):
id: str
name: str
error: str | None
interrupts: list[dict]
result: dict[str, Any]
class CheckpointTask(TypedDict):
id: str
name: str
error: str | None
interrupts: list[dict]
state: StateSnapshot | RunnableConfig | None
class CheckpointPayload(TypedDict):
config: RunnableConfig | None
metadata: CheckpointMetadata
values: dict[str, Any]
next: list[str]
parent_config: RunnableConfig | None
tasks: list[CheckpointTask]
from langgraph.types import (
CheckpointPayload,
PregelExecutableTask,
PregelTask,
StateSnapshot,
TaskPayload,
TaskResultPayload,
)
TASK_NAMESPACE = UUID("6ba7b831-9dad-11d1-80b4-00c04fd430c8")
libs/langgraph/langgraph/pregel/main.py
+386 -64
View File
@@ -22,8 +22,10 @@ from inspect import isclass
from typing import (
Any,
Generic,
Literal,
cast,
get_type_hints,
overload,
)
from uuid import UUID, uuid5
@@ -121,7 +123,10 @@ from langgraph.pregel._checkpoint import (
)
from langgraph.pregel._draw import draw_graph
from langgraph.pregel._io import map_input, read_channels
from langgraph.pregel._loop import AsyncPregelLoop, SyncPregelLoop
from langgraph.pregel._loop import (
AsyncPregelLoop,
SyncPregelLoop,
)
from langgraph.pregel._messages import StreamMessagesHandler
from langgraph.pregel._read import DEFAULT_BOUND, PregelNode
from langgraph.pregel._retry import RetryPolicy
@@ -138,11 +143,13 @@ from langgraph.types import (
Checkpointer,
Command,
Durability,
GraphOutput,
Interrupt,
Send,
StateSnapshot,
StateUpdate,
StreamMode,
StreamPart,
ensure_valid_checkpointer,
)
from langgraph.typing import ContextT, InputT, OutputT, StateT
@@ -993,6 +1000,11 @@ class Pregel(
for name, node in self.get_subgraphs(namespace=namespace, recurse=recurse):
yield name, node
# Mappers for v2 stream coercion (pydantic/dataclass).
# Set by CompiledStateGraph; None for base Pregel.
_output_mapper: Callable[[Any], Any] | None = None
_state_mapper: Callable[[Any], Any] | None = None
def _migrate_checkpoint(self, checkpoint: Checkpoint) -> None:
"""Migrate a saved checkpoint to new channel layout."""
if checkpoint["v"] < 4 and checkpoint.get("pending_sends"):
@@ -2427,6 +2439,7 @@ class Pregel(
durability,
)
@overload
def stream(
self,
input: InputT | Command | None,
@@ -2441,6 +2454,44 @@ class Pregel(
durability: Durability | None = None,
subgraphs: bool = False,
debug: bool | None = None,
version: Literal["v2"],
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[StreamPart[OutputT, StateT]]: ...
@overload
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: StreamMode | Sequence[StreamMode] | None = None,
print_mode: StreamMode | Sequence[StreamMode] = (),
output_keys: str | Sequence[str] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
subgraphs: bool = False,
debug: bool | None = None,
version: Literal["v1"] = ...,
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[dict[str, Any] | Any]: ...
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: StreamMode | Sequence[StreamMode] | None = None,
print_mode: StreamMode | Sequence[StreamMode] = (),
output_keys: str | Sequence[str] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
subgraphs: bool = False,
debug: bool | None = None,
version: Literal["v1", "v2"] = "v1",
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[dict[str, Any] | Any]:
"""Stream graph steps for a single input.
@@ -2602,6 +2653,10 @@ class Pregel(
runtime = parent_runtime.merge(runtime)
config[CONF][CONFIG_KEY_RUNTIME] = runtime
# resolve mappers for v2 stream coercion
_output_mapper = self._output_mapper if version == "v2" else None
_state_mapper = self._state_mapper if version == "v2" else None
with SyncPregelLoop(
input,
stream=StreamProtocol(stream.put, stream_modes),
@@ -2674,7 +2729,14 @@ class Pregel(
):
# emit output
yield from _output(
stream_mode, print_mode, subgraphs, stream.get, queue.Empty
stream_mode,
print_mode,
subgraphs,
stream.get,
queue.Empty,
version,
_output_mapper,
_state_mapper,
)
loop.after_tick()
# wait for checkpoint
@@ -2682,7 +2744,14 @@ class Pregel(
loop._put_checkpoint_fut.result()
# emit output
yield from _output(
stream_mode, print_mode, subgraphs, stream.get, queue.Empty
stream_mode,
print_mode,
subgraphs,
stream.get,
queue.Empty,
version,
_output_mapper,
_state_mapper,
)
# handle exit
if loop.status == "out_of_steps":
@@ -2701,6 +2770,44 @@ class Pregel(
run_manager.on_chain_error(e)
raise
@overload
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: StreamMode | Sequence[StreamMode] | None = None,
print_mode: StreamMode | Sequence[StreamMode] = (),
output_keys: str | Sequence[str] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
subgraphs: bool = False,
debug: bool | None = None,
version: Literal["v2"],
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[StreamPart[OutputT, StateT]]: ...
@overload
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: StreamMode | Sequence[StreamMode] | None = None,
print_mode: StreamMode | Sequence[StreamMode] = (),
output_keys: str | Sequence[str] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
subgraphs: bool = False,
debug: bool | None = None,
version: Literal["v1"] = ...,
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[dict[str, Any] | Any]: ...
async def astream(
self,
input: InputT | Command | None,
@@ -2715,6 +2822,7 @@ class Pregel(
durability: Durability | None = None,
subgraphs: bool = False,
debug: bool | None = None,
version: Literal["v1", "v2"] = "v1",
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[dict[str, Any] | Any]:
"""Asynchronously stream graph steps for a single input.
@@ -2911,6 +3019,10 @@ class Pregel(
runtime = parent_runtime.merge(runtime)
config[CONF][CONFIG_KEY_RUNTIME] = runtime
# resolve mappers for v2 stream coercion
_output_mapper = self._output_mapper if version == "v2" else None
_state_mapper = self._state_mapper if version == "v2" else None
async with AsyncPregelLoop(
input,
stream=StreamProtocol(stream.put_nowait, stream_modes),
@@ -3007,6 +3119,9 @@ class Pregel(
subgraphs,
stream.get_nowait,
asyncio.QueueEmpty,
version,
_output_mapper,
_state_mapper,
):
yield o
loop.after_tick()
@@ -3025,6 +3140,9 @@ class Pregel(
subgraphs,
stream.get_nowait,
asyncio.QueueEmpty,
version,
_output_mapper,
_state_mapper,
):
yield o
# handle exit
@@ -3044,6 +3162,41 @@ class Pregel(
await asyncio.shield(run_manager.on_chain_error(e))
raise
@overload
def invoke(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["values"] = ...,
print_mode: StreamMode | Sequence[StreamMode] = (),
output_keys: str | Sequence[str] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
version: Literal["v2"],
**kwargs: Any,
) -> GraphOutput[OutputT]: ...
@overload
def invoke(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: StreamMode,
print_mode: StreamMode | Sequence[StreamMode] = (),
output_keys: str | Sequence[str] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
version: Literal["v2"],
**kwargs: Any,
) -> list[StreamPart[OutputT, StateT]]: ...
@overload
def invoke(
self,
input: InputT | Command | None,
@@ -3056,6 +3209,23 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
version: Literal["v1"] = ...,
**kwargs: Any,
) -> dict[str, Any] | Any: ...
def invoke(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: StreamMode = "values",
print_mode: StreamMode | Sequence[StreamMode] = (),
output_keys: str | Sequence[str] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
version: Literal["v1", "v2"] = "v1",
**kwargs: Any,
) -> dict[str, Any] | Any:
"""Run the graph with a single input and config.
@@ -3079,6 +3249,9 @@ class Pregel(
- `"sync"`: Changes are persisted synchronously before the next step starts.
- `"async"`: Changes are persisted asynchronously while the next step executes.
- `"exit"`: Changes are persisted only when the graph exits.
version: The streaming format version. `"v1"` (default) returns the
traditional format, `"v2"` returns `StreamPart` typed dicts when
`stream_mode` is not `"values"`.
**kwargs: Additional keyword arguments to pass to the graph run.
Returns:
@@ -3091,39 +3264,64 @@ class Pregel(
chunks: list[dict[str, Any] | Any] = []
interrupts: list[Interrupt] = []
for chunk in self.stream(
input,
config,
context=context,
stream_mode=["updates", "values"]
if stream_mode == "values"
else stream_mode,
print_mode=print_mode,
output_keys=output_keys,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
durability=durability,
**kwargs,
):
if stream_mode == "values":
if len(chunk) == 2:
mode, payload = cast(tuple[StreamMode, Any], chunk)
if version == "v2":
# v2: values stream parts carry interrupts directly
for chunk in self.stream(
input,
config,
context=context,
stream_mode="values" if stream_mode == "values" else stream_mode,
print_mode=print_mode,
output_keys=output_keys,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
durability=durability,
version=version,
**kwargs,
):
if stream_mode == "values":
latest = chunk["data"]
if chunk_ints := chunk.get("interrupts", ()):
interrupts.extend(chunk_ints) # type: ignore[arg-type]
else:
_, mode, payload = cast(
tuple[tuple[str, ...], StreamMode, Any], chunk
)
if (
mode == "updates"
and isinstance(payload, dict)
and (ints := payload.get(INTERRUPT)) is not None
):
interrupts.extend(ints)
elif mode == "values":
latest = payload
else:
chunks.append(chunk)
chunks.append(chunk)
else:
# v1: collect interrupts from updates stream
for chunk in self.stream(
input,
config,
context=context,
stream_mode=(
["updates", "values"] if stream_mode == "values" else stream_mode
),
print_mode=print_mode,
output_keys=output_keys,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
durability=durability,
**kwargs,
):
if stream_mode == "values":
if len(chunk) == 2:
mode, payload = cast(tuple[StreamMode, Any], chunk)
else:
_, mode, payload = cast(
tuple[tuple[str, ...], StreamMode, Any], chunk
)
if (
mode == "updates"
and isinstance(payload, dict)
and (ints := payload.get(INTERRUPT)) is not None
):
interrupts.extend(ints)
elif mode == "values":
latest = payload
else:
chunks.append(chunk)
if stream_mode == "values":
if version == "v2":
return GraphOutput(value=latest, interrupts=tuple(interrupts))
if interrupts:
return (
{**latest, INTERRUPT: interrupts}
@@ -3134,6 +3332,57 @@ class Pregel(
else:
return chunks
@overload
async def ainvoke(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["values"] = ...,
print_mode: StreamMode | Sequence[StreamMode] = (),
output_keys: str | Sequence[str] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
version: Literal["v2"],
**kwargs: Any,
) -> GraphOutput[OutputT]: ...
@overload
async def ainvoke(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: StreamMode,
print_mode: StreamMode | Sequence[StreamMode] = (),
output_keys: str | Sequence[str] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
version: Literal["v2"],
**kwargs: Any,
) -> list[StreamPart[OutputT, StateT]]: ...
@overload
async def ainvoke(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: StreamMode = "values",
print_mode: StreamMode | Sequence[StreamMode] = (),
output_keys: str | Sequence[str] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
version: Literal["v1"] = ...,
**kwargs: Any,
) -> dict[str, Any] | Any: ...
async def ainvoke(
self,
input: InputT | Command | None,
@@ -3146,6 +3395,7 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
version: Literal["v1", "v2"] = "v1",
**kwargs: Any,
) -> dict[str, Any] | Any:
"""Asynchronously run the graph with a single input and config.
@@ -3169,6 +3419,9 @@ class Pregel(
- `"sync"`: Changes are persisted synchronously before the next step starts.
- `"async"`: Changes are persisted asynchronously while the next step executes.
- `"exit"`: Changes are persisted only when the graph exits.
version: The streaming format version. `"v1"` (default) returns the
traditional format, `"v2"` returns `StreamPart` typed dicts when
`stream_mode` is not `"values"`.
**kwargs: Additional keyword arguments to pass to the graph run.
Returns:
@@ -3181,39 +3434,64 @@ class Pregel(
chunks: list[dict[str, Any] | Any] = []
interrupts: list[Interrupt] = []
async for chunk in self.astream(
input,
config,
context=context,
stream_mode=["updates", "values"]
if stream_mode == "values"
else stream_mode,
print_mode=print_mode,
output_keys=output_keys,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
durability=durability,
**kwargs,
):
if stream_mode == "values":
if len(chunk) == 2:
mode, payload = cast(tuple[StreamMode, Any], chunk)
if version == "v2":
# v2: values stream parts carry interrupts directly
async for chunk in self.astream(
input,
config,
context=context,
stream_mode="values" if stream_mode == "values" else stream_mode,
print_mode=print_mode,
output_keys=output_keys,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
durability=durability,
version=version,
**kwargs,
):
if stream_mode == "values":
latest = chunk["data"]
if chunk_ints := chunk.get("interrupts", ()):
interrupts.extend(chunk_ints) # type: ignore[arg-type]
else:
_, mode, payload = cast(
tuple[tuple[str, ...], StreamMode, Any], chunk
)
if (
mode == "updates"
and isinstance(payload, dict)
and (ints := payload.get(INTERRUPT)) is not None
):
interrupts.extend(ints)
elif mode == "values":
latest = payload
else:
chunks.append(chunk)
chunks.append(chunk)
else:
# v1: collect interrupts from updates stream
async for chunk in self.astream(
input,
config,
context=context,
stream_mode=(
["updates", "values"] if stream_mode == "values" else stream_mode
),
print_mode=print_mode,
output_keys=output_keys,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
durability=durability,
**kwargs,
):
if stream_mode == "values":
if len(chunk) == 2:
mode, payload = cast(tuple[StreamMode, Any], chunk)
else:
_, mode, payload = cast(
tuple[tuple[str, ...], StreamMode, Any], chunk
)
if (
mode == "updates"
and isinstance(payload, dict)
and (ints := payload.get(INTERRUPT)) is not None
):
interrupts.extend(ints)
elif mode == "values":
latest = payload
else:
chunks.append(chunk)
if stream_mode == "values":
if version == "v2":
return GraphOutput(value=latest, interrupts=tuple(interrupts))
if interrupts:
return (
{**latest, INTERRUPT: interrupts}
@@ -3278,6 +3556,9 @@ def _output(
stream_subgraphs: bool,
getter: Callable[[], tuple[tuple[str, ...], str, Any]],
empty_exc: type[Exception],
version: Literal["v1", "v2"] = "v1",
output_mapper: Callable[[Any], Any] | None = None,
state_mapper: Callable[[Any], Any] | None = None,
) -> Iterator:
while True:
try:
@@ -3305,7 +3586,23 @@ def _output(
)
)
if mode in stream_mode:
if stream_subgraphs and isinstance(stream_mode, list):
if version == "v2":
if mode == "values":
# pop __interrupt__ into typed field, coerce data
ints: tuple[Interrupt, ...] = ()
if isinstance(payload, dict):
ints = payload.pop(INTERRUPT, ())
if output_mapper:
payload = output_mapper(payload)
yield {"type": mode, "ns": ns, "data": payload, "interrupts": ints}
elif mode in ("checkpoints", "debug"):
# coerce state values in checkpoint/debug payloads
if state_mapper:
_coerce_checkpoint_values(payload, state_mapper)
yield {"type": mode, "ns": ns, "data": payload}
else:
yield {"type": mode, "ns": ns, "data": payload}
elif stream_subgraphs and isinstance(stream_mode, list):
yield (ns, mode, payload)
elif isinstance(stream_mode, list):
yield (mode, payload)
@@ -3315,6 +3612,31 @@ def _output(
yield payload
def _coerce_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> None:
"""Coerce `values` dicts inside checkpoint or debug payloads in-place.
Skips the initial checkpoint (where next contains ``__start__``) because
not all channels are populated yet and coercion would fail.
"""
_START = "__start__"
# debug wrapper: {"type": "checkpoint", "payload": {"values": dict, ...}}
if (
isinstance(payload, dict)
and payload.get("type") == "checkpoint"
and isinstance(payload.get("payload"), dict)
and isinstance(payload["payload"].get("values"), dict)
and _START not in payload["payload"].get("next", ())
):
payload["payload"]["values"] = mapper(payload["payload"]["values"])
# direct checkpoint payload: {"values": dict, ...}
elif (
isinstance(payload, dict)
and isinstance(payload.get("values"), dict)
and _START not in payload.get("next", ())
):
payload["values"] = mapper(payload["values"])
def _coerce_context(
context_schema: type[ContextT] | None, context: Any
) -> ContextT | None:
libs/langgraph/langgraph/pregel/protocol.py
+126 -2
View File
@@ -2,13 +2,21 @@ from __future__ import annotations
from abc import abstractmethod
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
from typing import Any, Generic, cast
from typing import Any, Generic, Literal, cast, overload
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.runnables.graph import Graph as DrawableGraph
from typing_extensions import Self
from langgraph.types import All, Command, StateSnapshot, StateUpdate, StreamMode
from langgraph.types import (
All,
Command,
GraphOutput,
StateSnapshot,
StateUpdate,
StreamMode,
StreamPart,
)
from langgraph.typing import ContextT, InputT, OutputT, StateT
__all__ = ("PregelProtocol", "StreamProtocol")
@@ -96,6 +104,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
as_node: str | None = None,
) -> RunnableConfig: ...
@overload
@abstractmethod
def stream(
self,
@@ -107,8 +116,68 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: bool = False,
version: Literal["v2"],
) -> Iterator[StreamPart[OutputT, StateT]]: ...
@overload
@abstractmethod
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: bool = False,
version: Literal["v1"] = ...,
) -> Iterator[dict[str, Any] | Any]: ...
@abstractmethod
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: bool = False,
version: Literal["v1", "v2"] = "v1",
) -> Iterator[dict[str, Any] | Any]: ...
@overload
@abstractmethod
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: bool = False,
version: Literal["v2"],
) -> AsyncIterator[StreamPart[OutputT, StateT]]: ...
@overload
@abstractmethod
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: bool = False,
version: Literal["v1"] = ...,
) -> AsyncIterator[dict[str, Any] | Any]: ...
@abstractmethod
def astream(
self,
@@ -120,8 +189,35 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: bool = False,
version: Literal["v1", "v2"] = "v1",
) -> AsyncIterator[dict[str, Any] | Any]: ...
@overload
@abstractmethod
def invoke(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
version: Literal["v2"],
) -> GraphOutput[OutputT]: ...
@overload
@abstractmethod
def invoke(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
version: Literal["v1"] = ...,
) -> dict[str, Any] | Any: ...
@abstractmethod
def invoke(
self,
@@ -131,6 +227,33 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
context: ContextT | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
version: Literal["v1", "v2"] = "v1",
) -> dict[str, Any] | Any: ...
@overload
@abstractmethod
async def ainvoke(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
version: Literal["v2"],
) -> GraphOutput[OutputT]: ...
@overload
@abstractmethod
async def ainvoke(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
version: Literal["v1"] = ...,
) -> dict[str, Any] | Any: ...
@abstractmethod
@@ -142,6 +265,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
context: ContextT | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
version: Literal["v1", "v2"] = "v1",
) -> dict[str, Any] | Any: ...
libs/langgraph/langgraph/pregel/remote.py
+165 -6
View File
@@ -7,6 +7,7 @@ from typing import (
Any,
Literal,
cast,
overload,
)
from uuid import UUID
@@ -57,10 +58,12 @@ from langgraph.pregel.protocol import PregelProtocol, StreamProtocol
from langgraph.types import (
All,
Command,
GraphOutput,
Interrupt,
PregelTask,
StateSnapshot,
StreamMode,
StreamPart,
)
logger = logging.getLogger(__name__)
@@ -682,6 +685,7 @@ class RemoteGraph(PregelProtocol):
updated_stream_modes.remove("events")
return (updated_stream_modes, requested_stream_modes, req_single, stream)
@overload
def stream(
self,
input: dict[str, Any] | Any,
@@ -693,6 +697,38 @@ class RemoteGraph(PregelProtocol):
subgraphs: bool = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
version: Literal["v2"],
**kwargs: Any,
) -> Iterator[StreamPart]: ...
@overload
def stream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: bool = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
version: Literal["v1"] = ...,
**kwargs: Any,
) -> Iterator[dict[str, Any] | Any]: ...
def stream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: bool = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
version: Literal["v1", "v2"] = "v1",
**kwargs: Any,
) -> Iterator[dict[str, Any] | Any]:
"""Create a run and stream the results.
@@ -774,10 +810,18 @@ class RemoteGraph(PregelProtocol):
continue
if chunk.event.startswith("messages"):
chunk = chunk._replace(data=tuple(chunk.data)) # type: ignore
chunk = chunk._replace(data=tuple(chunk.data))
# emit chunk
if subgraphs:
if version == "v2":
ints: tuple[Interrupt, ...] = ()
if mode == "values" and isinstance(chunk.data, dict):
ints = tuple(
Interrupt(**i) if isinstance(i, dict) else i
for i in chunk.data.pop(INTERRUPT, ())
)
yield {"type": mode, "ns": ns, "data": chunk.data, "interrupts": ints}
elif subgraphs:
if NS_SEP in chunk.event:
mode, ns_ = chunk.event.split(NS_SEP, 1)
ns = tuple(ns_.split(NS_SEP))
@@ -792,6 +836,38 @@ class RemoteGraph(PregelProtocol):
else:
yield chunk
@overload
def astream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: bool = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
version: Literal["v2"],
**kwargs: Any,
) -> AsyncIterator[StreamPart]: ...
@overload
def astream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: bool = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
version: Literal["v1"] = ...,
**kwargs: Any,
) -> AsyncIterator[dict[str, Any] | Any]: ...
async def astream(
self,
input: dict[str, Any] | Any,
@@ -803,6 +879,7 @@ class RemoteGraph(PregelProtocol):
subgraphs: bool = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
version: Literal["v1", "v2"] = "v1",
**kwargs: Any,
) -> AsyncIterator[dict[str, Any] | Any]:
"""Create a run and stream the results.
@@ -884,10 +961,18 @@ class RemoteGraph(PregelProtocol):
continue
if chunk.event.startswith("messages"):
chunk = chunk._replace(data=tuple(chunk.data)) # type: ignore
chunk = chunk._replace(data=tuple(chunk.data))
# emit chunk
if subgraphs:
if version == "v2":
ints: tuple[Interrupt, ...] = ()
if mode == "values" and isinstance(chunk.data, dict):
ints = tuple(
Interrupt(**i) if isinstance(i, dict) else i
for i in chunk.data.pop(INTERRUPT, ())
)
yield {"type": mode, "ns": ns, "data": chunk.data, "interrupts": ints}
elif subgraphs:
if NS_SEP in chunk.event:
mode, ns_ = chunk.event.split(NS_SEP, 1)
ns = tuple(ns_.split(NS_SEP))
@@ -918,6 +1003,7 @@ class RemoteGraph(PregelProtocol):
) -> AsyncIterator[dict[str, Any]]:
raise NotImplementedError
@overload
def invoke(
self,
input: dict[str, Any] | Any,
@@ -927,6 +1013,34 @@ class RemoteGraph(PregelProtocol):
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
version: Literal["v2"],
**kwargs: Any,
) -> GraphOutput[dict[str, Any]]: ...
@overload
def invoke(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
version: Literal["v1"] = ...,
**kwargs: Any,
) -> dict[str, Any] | Any: ...
def invoke(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
version: Literal["v1", "v2"] = "v1",
**kwargs: Any,
) -> dict[str, Any] | Any:
"""Create a run, wait until it finishes and return the final state.
@@ -937,12 +1051,14 @@ class RemoteGraph(PregelProtocol):
interrupt_before: Interrupt the graph before these nodes.
interrupt_after: Interrupt the graph after these nodes.
headers: Additional headers to pass to the request.
version: The streaming format version. `"v1"` (default) returns the
traditional format, `"v2"` returns `StreamPart` typed dicts.
**kwargs: Additional params to pass to RemoteGraph.stream.
Returns:
The output of the graph.
"""
for chunk in self.stream(
for chunk in self.stream( # type: ignore[misc, call-overload]
input,
config=config,
interrupt_before=interrupt_before,
@@ -950,15 +1066,49 @@ class RemoteGraph(PregelProtocol):
headers=headers,
stream_mode="values",
params=params,
version=version,
**kwargs,
):
pass
try:
if version == "v2":
return GraphOutput(
value=chunk["data"],
interrupts=tuple(chunk.get("interrupts", ())),
)
return chunk
except UnboundLocalError:
logger.warning("No events received from remote graph")
return None
@overload
async def ainvoke(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
version: Literal["v2"],
**kwargs: Any,
) -> GraphOutput[dict[str, Any]]: ...
@overload
async def ainvoke(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
version: Literal["v1"] = ...,
**kwargs: Any,
) -> dict[str, Any] | Any: ...
async def ainvoke(
self,
input: dict[str, Any] | Any,
@@ -968,6 +1118,7 @@ class RemoteGraph(PregelProtocol):
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
version: Literal["v1", "v2"] = "v1",
**kwargs: Any,
) -> dict[str, Any] | Any:
"""Create a run, wait until it finishes and return the final state.
@@ -978,12 +1129,14 @@ class RemoteGraph(PregelProtocol):
interrupt_before: Interrupt the graph before these nodes.
interrupt_after: Interrupt the graph after these nodes.
headers: Additional headers to pass to the request.
version: The streaming format version. `"v1"` (default) returns the
traditional format, `"v2"` returns `StreamPart` typed dicts.
**kwargs: Additional params to pass to RemoteGraph.astream.
Returns:
The output of the graph.
"""
async for chunk in self.astream(
async for chunk in self.astream( # type: ignore[misc, call-overload]
input,
config=config,
interrupt_before=interrupt_before,
@@ -991,10 +1144,16 @@ class RemoteGraph(PregelProtocol):
headers=headers,
stream_mode="values",
params=params,
version=version,
**kwargs,
):
pass
try:
if version == "v2":
return GraphOutput(
value=chunk["data"],
interrupts=tuple(chunk.get("interrupts", ())),
)
return chunk
except UnboundLocalError:
logger.warning("No events received from remote graph")
libs/langgraph/langgraph/types.py
+287 -2
View File
@@ -16,16 +16,25 @@ from typing import (
)
from warnings import warn
from langchain_core.messages import AnyMessage
from langchain_core.runnables import Runnable, RunnableConfig
from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
from typing_extensions import Unpack, deprecated
from typing_extensions import NotRequired, TypeAliasType, TypedDict, Unpack, deprecated
from xxhash import xxh3_128_hexdigest
from langgraph._internal._cache import default_cache_key
from langgraph._internal._constants import INTERRUPT as _INTERRUPT_KEY
from langgraph._internal._fields import get_cached_annotated_keys, get_update_as_tuples
from langgraph._internal._retry import default_retry_on
from langgraph._internal._typing import MISSING, DeprecatedKwargs
from langgraph.warnings import LangGraphDeprecatedSinceV10
from langgraph.warnings import LangGraphDeprecatedSinceV10, LangGraphDeprecatedSinceV11
# Local TypeVars for generic stream TypedDicts.
# We use separate TypeVars here (rather than importing from langgraph.typing)
# because the typing module TypeVars have defaults that cause mypy issues
# when used in standalone type aliases.
StateT = TypeVar("StateT")
OutputT = TypeVar("OutputT")
if TYPE_CHECKING:
from langgraph.pregel.protocol import PregelProtocol
@@ -44,6 +53,19 @@ __all__ = (
"Checkpointer",
"StreamMode",
"StreamWriter",
"StreamPart",
"ValuesStreamPart",
"UpdatesStreamPart",
"MessagesStreamPart",
"CustomStreamPart",
"CheckpointStreamPart",
"TasksStreamPart",
"DebugStreamPart",
"TaskPayload",
"TaskResultPayload",
"CheckpointTask",
"CheckpointPayload",
"DebugPayload",
"RetryPolicy",
"CachePolicy",
"Interrupt",
@@ -56,6 +78,7 @@ __all__ = (
"Durability",
"interrupt",
"Overwrite",
"GraphOutput",
"ensure_valid_checkpointer",
)
@@ -113,6 +136,268 @@ StreamWriter = Callable[[Any], None]
Always injected into nodes if requested as a keyword argument, but it's a no-op
when not using `stream_mode="custom"`."""
class TaskPayload(TypedDict):
"""Payload for a task start event."""
id: str
"""Unique identifier for this task."""
name: str
"""Name of the node being executed."""
input: Any
"""Input data passed to the task."""
triggers: list[str]
"""List of triggers that caused this task to be executed (e.g. channel writes)."""
class TaskResultPayload(TypedDict):
"""Payload for a task result event."""
id: str
"""Unique identifier for this task."""
name: str
"""Name of the node that was executed."""
error: str | None
"""Error message if the task failed, otherwise `None`."""
interrupts: list[dict]
"""List of interrupts that occurred during task execution."""
result: dict[str, Any]
"""Mapping of channel names to the values written by this task."""
class CheckpointTask(TypedDict):
"""A task entry within a `CheckpointPayload`.
The keys present depend on the task's state:
- **Error:** `id`, `name`, `error`, `state`
- **Has result:** `id`, `name`, `result`, `interrupts`, `state`
- **Pending:** `id`, `name`, `interrupts`, `state`
"""
id: str
"""Unique identifier for this task."""
name: str
"""Name of the node being executed."""
error: NotRequired[str]
"""Error message, present only if the task failed."""
result: NotRequired[Any]
"""Result of the task, present only if the task completed successfully."""
interrupts: NotRequired[list[dict]]
"""List of interrupts, present when the task has been interrupted or completed."""
state: StateSnapshot | RunnableConfig | None
"""Snapshot of the subgraph state, or a `RunnableConfig` pointing to it. `None` if not a subgraph."""
class CheckpointPayload(TypedDict, Generic[StateT]):
"""Payload for a checkpoint event."""
config: RunnableConfig | None
"""Configuration for this checkpoint, including the `thread_id` and `checkpoint_id`."""
metadata: CheckpointMetadata
"""Metadata associated with this checkpoint (e.g. step number, source, writes)."""
values: StateT
"""Current state values at the time of this checkpoint."""
next: list[str]
"""Names of the nodes scheduled to execute next."""
parent_config: RunnableConfig | None
"""Configuration of the parent checkpoint, or `None` if this is the first checkpoint."""
tasks: list[CheckpointTask]
"""List of tasks associated with this checkpoint."""
class _DebugCheckpointPayload(TypedDict, Generic[StateT]):
step: int
"""The step number in the graph execution."""
timestamp: str
"""ISO 8601 timestamp of when this event occurred."""
type: Literal["checkpoint"]
"""Event type discriminator, always `"checkpoint"`."""
payload: CheckpointPayload[StateT]
"""The checkpoint payload."""
class _DebugTaskPayload(TypedDict):
step: int
"""The step number in the graph execution."""
timestamp: str
"""ISO 8601 timestamp of when this event occurred."""
type: Literal["task"]
"""Event type discriminator, always `"task"`."""
payload: TaskPayload
"""The task start payload."""
class _DebugTaskResultPayload(TypedDict):
step: int
"""The step number in the graph execution."""
timestamp: str
"""ISO 8601 timestamp of when this event occurred."""
type: Literal["task_result"]
"""Event type discriminator, always `"task_result"`."""
payload: TaskResultPayload
"""The task result payload."""
DebugPayload = TypeAliasType(
"DebugPayload",
_DebugCheckpointPayload[StateT] | _DebugTaskPayload | _DebugTaskResultPayload,
type_params=(StateT,),
)
"""Wrapper payload for debug events. Discriminate on `type`."""
class ValuesStreamPart(TypedDict, Generic[OutputT]):
"""Stream part emitted for `stream_mode="values"`.
`data` contains the full state after each step, as returned by `read_channels()`.
"""
type: Literal["values"]
ns: tuple[str, ...]
data: OutputT
interrupts: tuple[Interrupt, ...]
class UpdatesStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="updates"`.
`data` maps node names to their outputs. May also contain
`__interrupt__` (tuple of `Interrupt` dicts) and `__metadata__` keys.
"""
type: Literal["updates"]
ns: tuple[str, ...]
data: dict[str, Any]
class MessagesStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="messages"`.
`data` is a 2-tuple of `(message, metadata)` where `message` is a
`BaseMessage` (e.g. `AIMessageChunk`) and `metadata` is a dict containing
keys like `langgraph_step`, `langgraph_node`, `langgraph_triggers`, etc.
"""
type: Literal["messages"]
ns: tuple[str, ...]
data: tuple[AnyMessage, dict[str, Any]]
class CustomStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="custom"`.
`data` is whatever value was passed to `StreamWriter` inside a node.
"""
type: Literal["custom"]
ns: tuple[str, ...]
data: Any
class CheckpointStreamPart(TypedDict, Generic[StateT]):
"""Stream part emitted for `stream_mode="checkpoints"`."""
type: Literal["checkpoints"]
ns: tuple[str, ...]
data: CheckpointPayload[StateT]
class TasksStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="tasks"`.
For task start events, `data` is a `TaskPayload` with `id`, `name`,
`input`, and `triggers` keys.
For task result events, `data` is a `TaskResultPayload` with `id`,
`name`, `error`, `interrupts`, and `result` keys.
"""
type: Literal["tasks"]
ns: tuple[str, ...]
data: TaskPayload | TaskResultPayload
class DebugStreamPart(TypedDict, Generic[StateT]):
"""Stream part emitted for `stream_mode="debug"`."""
type: Literal["debug"]
ns: tuple[str, ...]
data: DebugPayload[StateT]
StreamPart = TypeAliasType(
"StreamPart",
ValuesStreamPart[OutputT]
| UpdatesStreamPart
| MessagesStreamPart
| CustomStreamPart
| CheckpointStreamPart[StateT]
| TasksStreamPart
| DebugStreamPart[StateT],
type_params=(OutputT, StateT),
)
"""A discriminated union of all v2 stream part types.
Use `part["type"]` to narrow the type:
```python
async for part in graph.astream(input, version="v2"):
if part["type"] == "values":
part["data"] # OutputT — full state (pydantic/dataclass/dict)
elif part["type"] == "messages":
part["data"] # tuple[BaseMessage, dict] — (message, metadata)
elif part["type"] == "custom":
part["data"] # Any — user-defined
```
"""
@dataclass(frozen=True)
class GraphOutput(Generic[OutputT]):
"""Typed container returned by `invoke()` / `ainvoke()` with `version="v2"`.
Attributes:
value: The final output of the graph (dict, Pydantic model, dataclass, etc.).
interrupts: Any interrupts that occurred during execution.
"""
value: OutputT
interrupts: tuple[Interrupt, ...] = ()
def __getitem__(self, key: str) -> Any:
"""Backward compat: `result['__interrupt__']` and dict-key access."""
warn(
"Accessing GraphOutput via `result[key]` is deprecated. "
"Use `result.value` to access the output value directly, "
"or `result.interrupts` for interrupts.",
LangGraphDeprecatedSinceV11,
stacklevel=2,
)
if key == _INTERRUPT_KEY:
return self.interrupts
if isinstance(self.value, dict):
return self.value[key]
try:
return getattr(self.value, key)
except AttributeError:
raise KeyError(key)
def __contains__(self, key: object) -> bool:
warn(
"Accessing GraphOutput via `key in result` is deprecated. "
"Use `result.value` to access the output value directly, "
"or `result.interrupts` for interrupts.",
LangGraphDeprecatedSinceV11,
stacklevel=2,
)
if key == _INTERRUPT_KEY:
return bool(self.interrupts)
if isinstance(self.value, dict):
return key in self.value
return isinstance(key, str) and hasattr(self.value, key)
_DC_KWARGS = {"kw_only": True, "slots": True, "frozen": True}
libs/langgraph/langgraph/warnings.py
+8
View File
@@ -6,6 +6,7 @@ __all__ = (
"LangGraphDeprecationWarning",
"LangGraphDeprecatedSinceV05",
"LangGraphDeprecatedSinceV10",
"LangGraphDeprecatedSinceV11",
)
@@ -59,3 +60,10 @@ class LangGraphDeprecatedSinceV10(LangGraphDeprecationWarning):
def __init__(self, message: str, *args: object) -> None:
super().__init__(message, *args, since=(1, 0), expected_removal=(2, 0))
class LangGraphDeprecatedSinceV11(LangGraphDeprecationWarning):
"""A specific `LangGraphDeprecationWarning` subclass defining functionality deprecated since LangGraph v1.1.0"""
def __init__(self, message: str, *args: object) -> None:
super().__init__(message, *args, since=(1, 1), expected_removal=(3, 0))
libs/langgraph/tests/test_deprecation.py
+38 -2
View File
@@ -14,8 +14,12 @@ from langgraph.func import entrypoint, task
from langgraph.graph import StateGraph
from langgraph.graph.message import MessageGraph
from langgraph.pregel import NodeBuilder, Pregel
from langgraph.types import Interrupt, RetryPolicy
from langgraph.warnings import LangGraphDeprecatedSinceV05, LangGraphDeprecatedSinceV10
from langgraph.types import GraphOutput, Interrupt, RetryPolicy
from langgraph.warnings import (
LangGraphDeprecatedSinceV05,
LangGraphDeprecatedSinceV10,
LangGraphDeprecatedSinceV11,
)
class PlainState(TypedDict): ...
@@ -197,6 +201,7 @@ def test_deprecated_import() -> None:
@pytest.mark.filterwarnings(
"ignore:`durability` has no effect when no checkpointer is present"
)
@pytest.mark.filterwarnings("ignore:Accessing GraphOutput via")
def test_checkpoint_during_deprecation_state_graph() -> None:
class CheckDurability(TypedDict):
durability: NotRequired[str]
@@ -341,3 +346,34 @@ def test_message_graph_deprecation() -> None:
match="MessageGraph is deprecated in LangGraph v1.0.0, to be removed in v2.0.0. Please use StateGraph with a `messages` key instead.",
):
MessageGraph()
def test_graph_output_getitem_deprecation() -> None:
output = GraphOutput(value={"foo": "bar"})
with pytest.warns(
LangGraphDeprecatedSinceV11,
match=r"Accessing GraphOutput via `result\[key\]` is deprecated",
):
assert output["foo"] == "bar"
def test_graph_output_contains_deprecation() -> None:
output = GraphOutput(value={"foo": "bar"})
with pytest.warns(
LangGraphDeprecatedSinceV11,
match=r"Accessing GraphOutput via `key in result` is deprecated",
):
assert "foo" in output
def test_graph_output_getitem_interrupt_deprecation() -> None:
interrupts = (Interrupt(value="q", id="abc"),)
output = GraphOutput(value={"foo": "bar"}, interrupts=interrupts)
with pytest.warns(
LangGraphDeprecatedSinceV11,
match=r"Accessing GraphOutput via `result\[key\]` is deprecated",
):
assert output["__interrupt__"] == interrupts
libs/langgraph/tests/test_stream_v2.py
+1152
View File
File diff suppressed because it is too large Load Diff
libs/sdk-py/langgraph_sdk/_async/runs.py
+88 -6
View File
@@ -5,12 +5,15 @@ from __future__ import annotations
import builtins
import warnings
from collections.abc import AsyncIterator, Callable, Mapping, Sequence
from typing import Any, overload
from typing import Any, Literal, overload
import httpx
from langgraph_sdk._async.http import HttpClient
from langgraph_sdk._shared.utilities import _get_run_metadata_from_response
from langgraph_sdk._shared.utilities import (
_get_run_metadata_from_response,
_sse_to_v2_dict,
)
from langgraph_sdk.schema import (
All,
BulkCancelRunsStatus,
@@ -33,9 +36,21 @@ from langgraph_sdk.schema import (
RunStatus,
StreamMode,
StreamPart,
StreamPartV2,
StreamVersion,
)
async def _wrap_stream_v2(
raw: AsyncIterator[StreamPart],
) -> AsyncIterator[StreamPartV2]:
"""Wrap a raw SSE stream, converting each event to a v2 dict."""
async for part in raw:
v2 = _sse_to_v2_dict(part.event, part.data)
if v2 is not None:
yield v2
class RunsClient:
"""Client for managing runs in LangGraph.
@@ -81,6 +96,66 @@ class RunsClient:
headers: Mapping[str, str] | None = None,
params: QueryParamTypes | None = None,
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
version: Literal["v1"] = "v1",
) -> AsyncIterator[StreamPart]: ...
@overload
def stream(
self,
thread_id: str,
assistant_id: str,
*,
input: Input | None = None,
command: Command | None = None,
stream_mode: StreamMode | Sequence[StreamMode] = "values",
stream_subgraphs: bool = False,
stream_resumable: bool = False,
metadata: Mapping[str, Any] | None = None,
config: Config | None = None,
context: Context | None = None,
checkpoint: Checkpoint | None = None,
checkpoint_id: str | None = None,
checkpoint_during: bool | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
feedback_keys: Sequence[str] | None = None,
on_disconnect: DisconnectMode | None = None,
webhook: str | None = None,
multitask_strategy: MultitaskStrategy | None = None,
if_not_exists: IfNotExists | None = None,
after_seconds: int | None = None,
headers: Mapping[str, str] | None = None,
params: QueryParamTypes | None = None,
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
version: Literal["v2"],
) -> AsyncIterator[StreamPartV2]: ...
@overload
def stream(
self,
thread_id: None,
assistant_id: str,
*,
input: Input | None = None,
command: Command | None = None,
stream_mode: StreamMode | Sequence[StreamMode] = "values",
stream_subgraphs: bool = False,
stream_resumable: bool = False,
metadata: Mapping[str, Any] | None = None,
config: Config | None = None,
checkpoint_during: bool | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
feedback_keys: Sequence[str] | None = None,
on_disconnect: DisconnectMode | None = None,
on_completion: OnCompletionBehavior | None = None,
if_not_exists: IfNotExists | None = None,
webhook: str | None = None,
after_seconds: int | None = None,
headers: Mapping[str, str] | None = None,
params: QueryParamTypes | None = None,
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
version: Literal["v1"] = "v1",
) -> AsyncIterator[StreamPart]: ...
@overload
@@ -108,7 +183,8 @@ class RunsClient:
headers: Mapping[str, str] | None = None,
params: QueryParamTypes | None = None,
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
) -> AsyncIterator[StreamPart]: ...
version: Literal["v2"],
) -> AsyncIterator[StreamPartV2]: ...
def stream(
self,
@@ -139,7 +215,8 @@ class RunsClient:
params: QueryParamTypes | None = None,
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
durability: Durability | None = None,
) -> AsyncIterator[StreamPart]:
version: StreamVersion = "v1",
) -> AsyncIterator[StreamPart | StreamPartV2]:
"""Create a run and stream the results.
Args:
@@ -180,6 +257,8 @@ class RunsClient:
"async" means checkpoints are persisted async while next graph step executes, replaces checkpoint_during=True
"sync" means checkpoints are persisted sync after graph step executes, replaces checkpoint_during=False
"exit" means checkpoints are only persisted when the run exits, does not save intermediate steps
version: Stream format version. "v1" (default) returns raw SSE StreamPart
NamedTuples. "v2" returns typed dicts with `type`, `ns`, and `data` keys.
Returns:
Asynchronous iterator of stream results.
@@ -222,7 +301,7 @@ class RunsClient:
stacklevel=2,
)
payload = {
payload: dict[str, Any] = {
"input": input,
"command": (
{k: v for k, v in command.items() if v is not None} if command else None
@@ -259,7 +338,7 @@ class RunsClient:
if on_run_created and (metadata := _get_run_metadata_from_response(res)):
on_run_created(metadata)
return self.http.stream(
raw = self.http.stream(
endpoint,
"POST",
json={k: v for k, v in payload.items() if v is not None},
@@ -267,6 +346,9 @@ class RunsClient:
headers=headers,
on_response=on_response if on_run_created else None,
)
if version == "v2":
return _wrap_stream_v2(raw)
return raw
@overload
async def create(
libs/sdk-py/langgraph_sdk/_shared/utilities.py
+18
View File
@@ -107,6 +107,24 @@ def _get_run_metadata_from_response(
return None
def _sse_to_v2_dict(event: str, data: Any) -> dict[str, Any] | None:
"""Convert an SSE event+data pair into a v2 stream part dict.
Returns None for ``end`` events (signals end of stream).
"""
if event == "end":
return None
parts = event.split("|")
event_type = parts[0]
ns = parts[1:] if len(parts) > 1 else []
result: dict[str, Any] = {"type": event_type, "ns": ns, "data": data}
if event_type == "values" and isinstance(data, dict):
result["interrupts"] = data.pop("__interrupt__", [])
else:
result["interrupts"] = []
return result
def _provided_vals(d: Mapping[str, Any]) -> dict[str, Any]:
return {k: v for k, v in d.items() if v is not None}
libs/sdk-py/langgraph_sdk/_sync/runs.py
+88 -7
View File
@@ -5,11 +5,14 @@ from __future__ import annotations
import builtins
import warnings
from collections.abc import Callable, Iterator, Mapping, Sequence
from typing import Any, overload
from typing import Any, Literal, overload
import httpx
from langgraph_sdk._shared.utilities import _get_run_metadata_from_response
from langgraph_sdk._shared.utilities import (
_get_run_metadata_from_response,
_sse_to_v2_dict,
)
from langgraph_sdk._sync.http import SyncHttpClient
from langgraph_sdk.schema import (
All,
@@ -33,9 +36,21 @@ from langgraph_sdk.schema import (
RunStatus,
StreamMode,
StreamPart,
StreamPartV2,
StreamVersion,
)
def _wrap_stream_v2_sync(
raw: Iterator[StreamPart],
) -> Iterator[StreamPartV2]:
"""Wrap a raw SSE stream, converting each event to a v2 dict."""
for part in raw:
v2 = _sse_to_v2_dict(part.event, part.data)
if v2 is not None:
yield v2
class SyncRunsClient:
"""Synchronous client for managing runs in LangGraph.
@@ -80,6 +95,66 @@ class SyncRunsClient:
headers: Mapping[str, str] | None = None,
params: QueryParamTypes | None = None,
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
version: Literal["v1"] = "v1",
) -> Iterator[StreamPart]: ...
@overload
def stream(
self,
thread_id: str,
assistant_id: str,
*,
input: Input | None = None,
command: Command | None = None,
stream_mode: StreamMode | Sequence[StreamMode] = "values",
stream_subgraphs: bool = False,
metadata: Mapping[str, Any] | None = None,
config: Config | None = None,
context: Context | None = None,
checkpoint: Checkpoint | None = None,
checkpoint_id: str | None = None,
checkpoint_during: bool | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
feedback_keys: Sequence[str] | None = None,
on_disconnect: DisconnectMode | None = None,
webhook: str | None = None,
multitask_strategy: MultitaskStrategy | None = None,
if_not_exists: IfNotExists | None = None,
after_seconds: int | None = None,
headers: Mapping[str, str] | None = None,
params: QueryParamTypes | None = None,
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
version: Literal["v2"],
) -> Iterator[StreamPartV2]: ...
@overload
def stream(
self,
thread_id: None,
assistant_id: str,
*,
input: Input | None = None,
command: Command | None = None,
stream_mode: StreamMode | Sequence[StreamMode] = "values",
stream_subgraphs: bool = False,
stream_resumable: bool = False,
metadata: Mapping[str, Any] | None = None,
config: Config | None = None,
context: Context | None = None,
checkpoint_during: bool | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
feedback_keys: Sequence[str] | None = None,
on_disconnect: DisconnectMode | None = None,
on_completion: OnCompletionBehavior | None = None,
if_not_exists: IfNotExists | None = None,
webhook: str | None = None,
after_seconds: int | None = None,
headers: Mapping[str, str] | None = None,
params: QueryParamTypes | None = None,
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
version: Literal["v1"] = "v1",
) -> Iterator[StreamPart]: ...
@overload
@@ -108,7 +183,8 @@ class SyncRunsClient:
headers: Mapping[str, str] | None = None,
params: QueryParamTypes | None = None,
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
) -> Iterator[StreamPart]: ...
version: Literal["v2"],
) -> Iterator[StreamPartV2]: ...
def stream(
self,
@@ -139,7 +215,8 @@ class SyncRunsClient:
params: QueryParamTypes | None = None,
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
durability: Durability | None = None,
) -> Iterator[StreamPart]:
version: StreamVersion = "v1",
) -> Iterator[StreamPart | StreamPartV2]:
"""Create a run and stream the results.
Args:
@@ -179,7 +256,8 @@ class SyncRunsClient:
"async" means checkpoints are persisted async while next graph step executes, replaces checkpoint_during=True
"sync" means checkpoints are persisted sync after graph step executes, replaces checkpoint_during=False
"exit" means checkpoints are only persisted when the run exits, does not save intermediate steps
version: Stream format version. "v1" (default) returns raw SSE StreamPart
NamedTuples. "v2" returns typed dicts with `type`, `ns`, and `data` keys.
Returns:
Iterator of stream results.
@@ -218,7 +296,7 @@ class SyncRunsClient:
DeprecationWarning,
stacklevel=2,
)
payload = {
payload: dict[str, Any] = {
"input": input,
"command": (
{k: v for k, v in command.items() if v is not None} if command else None
@@ -255,7 +333,7 @@ class SyncRunsClient:
if on_run_created and (metadata := _get_run_metadata_from_response(res)):
on_run_created(metadata)
return self.http.stream(
raw = self.http.stream(
endpoint,
"POST",
json={k: v for k, v in payload.items() if v is not None},
@@ -263,6 +341,9 @@ class SyncRunsClient:
headers=headers,
on_response=on_response if on_run_created else None,
)
if version == "v2":
return _wrap_stream_v2_sync(raw)
return raw
@overload
def create(
libs/sdk-py/langgraph_sdk/schema.py
+269
View File
@@ -588,6 +588,275 @@ class StreamPart(NamedTuple):
"""The ID of the event."""
StreamVersion = Literal["v1", "v2"]
"""Stream format version.
- `"v1"`: Traditional format — raw SSE `StreamPart` NamedTuples.
- `"v2"`: Each event is a typed dict with `type`, `ns`, and `data` keys.
"""
# --- Typed payload dicts (JSON-deserialized from the server) ---
class TaskPayload(TypedDict):
"""Payload for a task start event."""
id: str
"""Unique identifier for this task."""
name: str
"""Name of the node being executed."""
input: Any
"""Input data passed to the task."""
triggers: list[str]
"""List of triggers that caused this task to be executed (e.g. channel writes)."""
class TaskResultPayload(TypedDict):
"""Payload for a task result event."""
id: str
"""Unique identifier for this task."""
name: str
"""Name of the node that was executed."""
error: str | None
"""Error message if the task failed, otherwise `None`."""
interrupts: list[dict[str, Any]]
"""List of interrupts that occurred during task execution."""
result: dict[str, Any]
"""Mapping of channel names to the values written by this task."""
class CheckpointTaskPayload(TypedDict):
"""A task entry within a `CheckpointPayload`.
The keys present depend on the task's state:
- **Error:** `id`, `name`, `error`, `state`
- **Has result:** `id`, `name`, `result`, `interrupts`, `state`
- **Pending:** `id`, `name`, `interrupts`, `state`
"""
id: str
"""Unique identifier for this task."""
name: str
"""Name of the node being executed."""
error: NotRequired[str]
"""Error message, present only if the task failed."""
result: NotRequired[Any]
"""Result of the task, present only if the task completed successfully."""
interrupts: NotRequired[list[dict[str, Any]]]
"""List of interrupts, present when the task has been interrupted or completed."""
state: dict[str, Any] | None
"""Snapshot of the subgraph state. `None` if not a subgraph."""
class CheckpointPayload(TypedDict):
"""Payload for a checkpoint event."""
config: dict[str, Any] | None
"""Configuration for this checkpoint, including the `thread_id` and `checkpoint_id`."""
metadata: dict[str, Any]
"""Metadata associated with this checkpoint (e.g. step number, source, writes)."""
values: dict[str, Any]
"""Current state values at the time of this checkpoint."""
next: list[str]
"""Names of the nodes scheduled to execute next."""
parent_config: dict[str, Any] | None
"""Configuration of the parent checkpoint, or `None` if this is the first checkpoint."""
tasks: list[CheckpointTaskPayload]
"""List of tasks associated with this checkpoint."""
class _DebugCheckpointPayload(TypedDict):
step: int
"""The step number in the graph execution."""
timestamp: str
"""ISO 8601 timestamp of when this event occurred."""
type: Literal["checkpoint"]
"""Event type discriminator, always `"checkpoint"`."""
payload: CheckpointPayload
"""The checkpoint payload."""
class _DebugTaskPayload(TypedDict):
step: int
"""The step number in the graph execution."""
timestamp: str
"""ISO 8601 timestamp of when this event occurred."""
type: Literal["task"]
"""Event type discriminator, always `"task"`."""
payload: TaskPayload
"""The task start payload."""
class _DebugTaskResultPayload(TypedDict):
step: int
"""The step number in the graph execution."""
timestamp: str
"""ISO 8601 timestamp of when this event occurred."""
type: Literal["task_result"]
"""Event type discriminator, always `"task_result"`."""
payload: TaskResultPayload
"""The task result payload."""
DebugPayload = _DebugCheckpointPayload | _DebugTaskPayload | _DebugTaskResultPayload
"""Wrapper payload for debug events. Discriminate on `type`."""
class RunMetadataPayload(TypedDict):
"""Payload for the `metadata` control event."""
run_id: str
"""The unique identifier of the run."""
# --- v2 stream part TypedDicts ---
class ValuesStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="values"`."""
type: Literal["values"]
"""Stream part type discriminator."""
ns: list[str]
"""Namespace path of the emitting node (empty for root graph)."""
data: dict[str, Any]
"""Full state values after the step."""
interrupts: list[dict[str, Any]]
"""List of interrupts that occurred during this step."""
class UpdatesStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="updates"`."""
type: Literal["updates"]
"""Stream part type discriminator."""
ns: list[str]
"""Namespace path of the emitting node (empty for root graph)."""
data: dict[str, Any]
"""Mapping of node names to their outputs."""
class MessagesPartialStreamPart(TypedDict):
"""Stream part emitted for partial message chunks (`messages/partial`)."""
type: Literal["messages/partial"]
"""Stream part type discriminator."""
ns: list[str]
"""Namespace path of the emitting node (empty for root graph)."""
data: list[dict[str, Any]]
"""List of partial message chunk dicts."""
class MessagesCompleteStreamPart(TypedDict):
"""Stream part emitted for complete messages (`messages/complete`)."""
type: Literal["messages/complete"]
"""Stream part type discriminator."""
ns: list[str]
"""Namespace path of the emitting node (empty for root graph)."""
data: list[dict[str, Any]]
"""List of complete message dicts."""
class MessagesMetadataStreamPart(TypedDict):
"""Stream part emitted for message metadata (`messages/metadata`)."""
type: Literal["messages/metadata"]
"""Stream part type discriminator."""
ns: list[str]
"""Namespace path of the emitting node (empty for root graph)."""
data: dict[str, Any]
"""Metadata dict for the message (e.g. `langgraph_step`, `langgraph_node`)."""
class MessagesTupleStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="messages"` (raw message+metadata pair)."""
type: Literal["messages"]
"""Stream part type discriminator."""
ns: list[str]
"""Namespace path of the emitting node (empty for root graph)."""
data: list[dict[str, Any]]
"""Two-element list of `[message_dict, metadata_dict]`."""
class CustomStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="custom"`."""
type: Literal["custom"]
"""Stream part type discriminator."""
ns: list[str]
"""Namespace path of the emitting node (empty for root graph)."""
data: Any
"""User-defined data passed to `StreamWriter` inside a node."""
class CheckpointsStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="checkpoints"`."""
type: Literal["checkpoints"]
"""Stream part type discriminator."""
ns: list[str]
"""Namespace path of the emitting node (empty for root graph)."""
data: CheckpointPayload
"""The checkpoint payload."""
class TasksStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="tasks"`."""
type: Literal["tasks"]
"""Stream part type discriminator."""
ns: list[str]
"""Namespace path of the emitting node (empty for root graph)."""
data: TaskPayload | TaskResultPayload
"""Task start or task result payload."""
class DebugStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="debug"`."""
type: Literal["debug"]
"""Stream part type discriminator."""
ns: list[str]
"""Namespace path of the emitting node (empty for root graph)."""
data: DebugPayload
"""The debug event payload."""
class MetadataStreamPart(TypedDict):
"""Control event with `run_id` and other run metadata."""
type: Literal["metadata"]
"""Stream part type discriminator."""
ns: list[str]
"""Namespace path (empty for root graph)."""
data: RunMetadataPayload
"""The run metadata payload."""
StreamPartV2 = (
ValuesStreamPart
| UpdatesStreamPart
| MessagesPartialStreamPart
| MessagesCompleteStreamPart
| MessagesMetadataStreamPart
| MessagesTupleStreamPart
| CustomStreamPart
| CheckpointsStreamPart
| TasksStreamPart
| DebugStreamPart
| MetadataStreamPart
)
"""Discriminated union of all v2 stream part types.
Use `part["type"]` to narrow the type.
"""
class Send(TypedDict):
"""Represents a message to be sent to a specific node in the graph.
libs/sdk-py/tests/test_client_stream.py
+234 -15
View File
@@ -2,18 +2,39 @@ from __future__ import annotations
from collections.abc import Iterator, Sequence
from pathlib import Path
from typing import Any
import httpx
import pytest
from typing_extensions import assert_type
from langgraph_sdk._shared.utilities import _sse_to_v2_dict
from langgraph_sdk.client import HttpClient, SyncHttpClient
from langgraph_sdk.schema import StreamPart
from langgraph_sdk.schema import (
CheckpointPayload,
CheckpointsStreamPart,
CustomStreamPart,
DebugPayload,
DebugStreamPart,
MetadataStreamPart,
RunMetadataPayload,
StreamPart,
StreamPartV2,
TaskPayload,
TaskResultPayload,
TasksStreamPart,
UpdatesStreamPart,
ValuesStreamPart,
)
from langgraph_sdk.sse import BytesLike, BytesLineDecoder, SSEDecoder
with open(Path(__file__).parent / "fixtures" / "response.txt", "rb") as f:
RESPONSE_PAYLOAD = f.read()
# --- test helpers ---
class AsyncListByteStream(httpx.AsyncByteStream):
def __init__(self, chunks: Sequence[bytes], exc: Exception | None = None) -> None:
self._chunks = list(chunks)
@@ -50,6 +71,24 @@ def iter_lines_raw(payload: list[bytes]) -> Iterator[BytesLike]:
yield from decoder.flush()
_V2_REQUIRED_KEYS = {"type", "ns", "data"}
def _assert_v2_shape(part: Any) -> None:
"""Assert a v2 stream part has the required keys and types."""
assert isinstance(part, dict), f"Expected dict, got {type(part)}"
assert part.keys() >= _V2_REQUIRED_KEYS, (
f"Missing keys: {_V2_REQUIRED_KEYS - part.keys()}"
)
assert isinstance(part["type"], str)
assert isinstance(part["ns"], list)
for elem in part["ns"]:
assert isinstance(elem, str)
# --- SSE parsing ---
def test_stream_sse():
for groups in (
[RESPONSE_PAYLOAD],
@@ -69,6 +108,9 @@ def test_stream_sse():
assert len(parts) == 79
# --- HTTP client streaming ---
@pytest.mark.asyncio
async def test_http_client_stream_flushes_trailing_event():
payload = b'event: foo\ndata: {"bar": 1}\n'
@@ -92,6 +134,26 @@ async def test_http_client_stream_flushes_trailing_event():
assert parts == [StreamPart(event="foo", data={"bar": 1})]
def test_sync_http_client_stream_flushes_trailing_event():
payload = b'event: foo\ndata: {"bar": 1}\n'
def handler(request: httpx.Request) -> httpx.Response:
assert request.headers["accept"] == "text/event-stream"
assert request.headers["cache-control"] == "no-store"
return httpx.Response(
200,
headers={"Content-Type": "text/event-stream"},
content=payload,
)
transport = httpx.MockTransport(handler)
with httpx.Client(transport=transport, base_url="https://example.com") as client:
http_client = SyncHttpClient(client)
parts = list(http_client.stream("/stream", "GET"))
assert parts == [StreamPart(event="foo", data={"bar": 1})]
def test_sync_http_client_stream_recovers_after_disconnect():
reconnect_path = "/reconnect"
first_chunks = [
@@ -228,21 +290,178 @@ async def test_http_client_stream_recovers_after_disconnect():
]
def test_sync_http_client_stream_flushes_trailing_event():
payload = b'event: foo\ndata: {"bar": 1}\n'
# --- _sse_to_v2_dict conversion ---
def handler(request: httpx.Request) -> httpx.Response:
assert request.headers["accept"] == "text/event-stream"
assert request.headers["cache-control"] == "no-store"
return httpx.Response(
200,
headers={"Content-Type": "text/event-stream"},
content=payload,
def test_sse_to_v2_dict_basic() -> None:
result = _sse_to_v2_dict("values", {"messages": [{"role": "user"}]})
assert result is not None
_assert_v2_shape(result)
assert result == {
"type": "values",
"ns": [],
"data": {"messages": [{"role": "user"}]},
"interrupts": [],
}
def test_sse_to_v2_dict_with_namespace() -> None:
result = _sse_to_v2_dict("updates|sub:abc", {"key": "val"})
assert result is not None
_assert_v2_shape(result)
assert result == {
"type": "updates",
"ns": ["sub:abc"],
"data": {"key": "val"},
"interrupts": [],
}
def test_sse_to_v2_dict_with_multiple_ns() -> None:
result = _sse_to_v2_dict("custom|parent|child:123", "hello")
assert result is not None
_assert_v2_shape(result)
assert result == {
"type": "custom",
"ns": ["parent", "child:123"],
"data": "hello",
"interrupts": [],
}
def test_sse_to_v2_dict_end_event() -> None:
assert _sse_to_v2_dict("end", None) is None
def test_sse_to_v2_dict_metadata_event() -> None:
result = _sse_to_v2_dict("metadata", {"run_id": "abc-123"})
assert result is not None
_assert_v2_shape(result)
assert result == {
"type": "metadata",
"ns": [],
"data": {"run_id": "abc-123"},
"interrupts": [],
}
def test_sse_to_v2_dict_messages_partial() -> None:
result = _sse_to_v2_dict("messages/partial", [{"type": "ai", "content": "hi"}])
assert result is not None
_assert_v2_shape(result)
assert result == {
"type": "messages/partial",
"ns": [],
"data": [{"type": "ai", "content": "hi"}],
"interrupts": [],
}
def test_sse_to_v2_dict_values_with_interrupts() -> None:
data = {
"messages": [{"role": "user"}],
"__interrupt__": [{"value": "confirm?", "resumable": True}],
}
result = _sse_to_v2_dict("values", data)
assert result is not None
_assert_v2_shape(result)
assert result == {
"type": "values",
"ns": [],
"data": {"messages": [{"role": "user"}]},
"interrupts": [{"value": "confirm?", "resumable": True}],
}
# __interrupt__ should be popped from data
assert "__interrupt__" not in result["data"]
# --- client-side v2 stream wrapping ---
@pytest.mark.asyncio
async def test_async_stream_v2_client_side_conversion() -> None:
from langgraph_sdk._async.runs import _wrap_stream_v2
async def mock_stream() -> Any:
yield StreamPart(event="metadata", data={"run_id": "r1"})
yield StreamPart(
event="values", data={"messages": [{"role": "user", "content": "hi"}]}
)
yield StreamPart(event="updates|sub:abc", data={"node": {"out": 1}})
yield StreamPart(event="end", data=None) # type: ignore[arg-type]
transport = httpx.MockTransport(handler)
with httpx.Client(transport=transport, base_url="https://example.com") as client:
http_client = SyncHttpClient(client)
parts = list(http_client.stream("/stream", "GET"))
parts: list[StreamPartV2] = [part async for part in _wrap_stream_v2(mock_stream())]
assert len(parts) == 3
for part in parts:
_assert_v2_shape(part)
assert parts[0] == {
"type": "metadata",
"ns": [],
"data": {"run_id": "r1"},
"interrupts": [],
}
assert parts[1] == {
"type": "values",
"ns": [],
"data": {"messages": [{"role": "user", "content": "hi"}]},
"interrupts": [],
}
assert parts[2] == {
"type": "updates",
"ns": ["sub:abc"],
"data": {"node": {"out": 1}},
"interrupts": [],
}
assert parts == [StreamPart(event="foo", data={"bar": 1})]
def test_sync_stream_v2_client_side_conversion() -> None:
from langgraph_sdk._sync.runs import _wrap_stream_v2_sync
def mock_stream() -> Any:
yield StreamPart(event="metadata", data={"run_id": "r1"})
yield StreamPart(event="values", data={"state": "full"})
yield StreamPart(event="end", data=None) # type: ignore[arg-type]
parts: list[StreamPartV2] = list(_wrap_stream_v2_sync(mock_stream()))
assert len(parts) == 2
for part in parts:
_assert_v2_shape(part)
assert parts[0] == {
"type": "metadata",
"ns": [],
"data": {"run_id": "r1"},
"interrupts": [],
}
assert parts[1] == {
"type": "values",
"ns": [],
"data": {"state": "full"},
"interrupts": [],
}
# --- type narrowing compile-time checks ---
def _check_v2_type_narrowing(part: StreamPartV2) -> None:
"""Compile-time type narrowing checks — validates mypy narrows the union."""
if part["type"] == "values":
assert_type(part, ValuesStreamPart)
assert_type(part["data"], dict[str, Any])
elif part["type"] == "updates":
assert_type(part, UpdatesStreamPart)
assert_type(part["data"], dict[str, Any])
elif part["type"] == "custom":
assert_type(part, CustomStreamPart)
elif part["type"] == "checkpoints":
assert_type(part, CheckpointsStreamPart)
assert_type(part["data"], CheckpointPayload)
elif part["type"] == "tasks":
assert_type(part, TasksStreamPart)
assert_type(part["data"], TaskPayload | TaskResultPayload)
elif part["type"] == "debug":
assert_type(part, DebugStreamPart)
assert_type(part["data"], DebugPayload)
elif part["type"] == "metadata":
assert_type(part, MetadataStreamPart)
assert_type(part["data"], RunMetadataPayload)