WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-06 09:47:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(langgraph): backwards compat type safe streaming (#6931)

Browse Source 
## Type-safe stream parts for v2 streaming

### Review recommendations

* Don't fear the diff! It's not so bad, I swear! The PR description
below gives a nice overview of changes
* Check out my explicit comments below, those should help to orient you
to the important changes :)
* On a first pass, ignore the test files! Just check out the new types,
overloads, and minor logical changes (diff behavior based on the flag)

## Summary

Adds a `stream_version="v2"` option that emits typed `{"type", "ns",
"data"}` dicts instead of raw tuples/SSE events. Each stream mode gets
its own `TypedDict` with a `Literal` type field, enabling full type
narrowing on `part["type"]`.

This is **opt in**, so it's **non-breaking**!!

```python
async for part in graph.astream(
    inputs, stream_mode=["messages", "custom"], stream_version="v2"
):
    if part["type"] == "messages":
        msg, metadata = part["data"]  # tuple[AnyMessage, dict] ✅
        print(msg.content)
    elif part["type"] == "custom":
        part["data"]  # Any ✅
```

Before v2 you'd get `tuple[str, Any]` with no way to narrow `data` based
on mode.

### What changed

**`langgraph` (core):** New `StreamPart` discriminated union + per-mode
TypedDicts in `types.py`. Stream-emit code in `pregel/` refactored to
use the new types. `RemoteGraph` gains a `stream_version` param.

**`sdk-py`:** Client-side v2 wrapper that converts SSE events into typed
dicts. No server API changes — v2 is purely a client-side rewrite of the
stream format.

### Stream part types

#### `langgraph` (core)

| `type` | `data` |
|---|---|
| `"values"` | `dict[str, Any]` — full state after each step |
| `"updates"` | `dict[str, Any]` — node name → output |
| `"messages"` | `tuple[AnyMessage, dict]` — message + metadata |
| `"custom"` | `Any` — whatever was passed to `StreamWriter` |
| `"tasks"` | `TaskPayload \| TaskResultPayload` |
| `"checkpoints"` | `CheckpointPayload` |
| `"debug"` | `DebugPayload` |

#### `sdk-py` (additional types from SSE events)

| `type` | `data` | Description |
|---|---|---|
| `"messages/partial"` | `list[dict]` | Partial message chunks |
| `"messages/complete"` | `list[dict]` | Complete messages |
| `"messages/metadata"` | `dict` | Message metadata |
| `"metadata"` | `RunMetadataPayload` | Run-level metadata (`run_id`,
etc.) |

All parts share the shape `{"type": Literal[...], "ns": list[str],
"data": ...}`.

## Release plan

* Release as a part of langgraph 1.1
* Before release, I'd like to do more experimentation with support for
pydantic + dataclasses and/or input/output runtime validation, as that
would help resolve the lack of typing for `values` mode.

### Notes

- Docs need a mass update to cover v2 streaming usage and the new types
- Changing the default `stream_version` to `"v2"` in a future release
would be breaking but backwards compatible (users can pin `"v1"` to keep
current behavior)
- I called out specifically relevant parts of the code in comments on
the PR :)
This commit is contained in:
Sydney Runkle
2026-02-27 09:37:09 -05:00
committed by GitHub
parent 353b0d8fb4
commit 9023ad9931
11 changed files with 1753 additions and 894 deletions
Download Patch File Download Diff File Expand all files Collapse all files
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
+171 -325
View File
@@ -20,7 +20,6 @@ from dataclasses import is_dataclass
from functools import partial
from inspect import isclass
from typing import (
TYPE_CHECKING,
Any,
Generic,
Literal,
@@ -51,9 +50,6 @@ from langgraph.store.base import BaseStore
from pydantic import BaseModel, TypeAdapter
from typing_extensions import Self, Unpack, deprecated, is_typeddict
if TYPE_CHECKING:
from langchain_core.messages import AnyMessage
from langgraph._internal._config import (
ensure_config,
merge_configs,
@@ -148,6 +144,7 @@ from langgraph.types import (
StateSnapshot,
StateUpdate,
StreamMode,
StreamPart,
ensure_valid_checkpointer,
)
from langgraph.typing import ContextT, InputT, OutputT, StateT
@@ -2417,16 +2414,17 @@ class Pregel(
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["values"],
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: Literal[False] = False,
subgraphs: bool = False,
debug: bool | None = None,
stream_version: Literal["v2"],
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[OutputT]: ...
) -> Iterator[StreamPart]: ...
@overload
def stream(
@@ -2435,160 +2433,17 @@ class Pregel(
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["updates"],
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: Literal[False] = False,
subgraphs: bool = False,
debug: bool | None = None,
stream_version: Literal["v1"] = ...,
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[dict[str, Any]]: ...
@overload
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["messages"],
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: Literal[False] = False,
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[tuple[AnyMessage, dict[str, Any]]]: ...
@overload
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["custom"],
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: Literal[False] = False,
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[Any]: ...
@overload
def stream(
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,
subgraphs: Literal[True],
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[tuple[tuple[str, ...], OutputT]]: ...
@overload
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["updates"],
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: Literal[True],
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[tuple[tuple[str, ...], dict[str, Any]]]: ...
@overload
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["messages"],
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: Literal[True],
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[tuple[tuple[str, ...], tuple[AnyMessage, dict[str, Any]]]]: ...
@overload
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["custom"],
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: Literal[True],
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[tuple[tuple[str, ...], Any]]: ...
@overload
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: list[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,
subgraphs: Literal[False] = False,
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[tuple[str, Any]]: ...
@overload
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: list[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,
subgraphs: Literal[True],
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[tuple[tuple[str, ...], str, Any]]: ...
) -> Iterator[dict[str, Any] | Any]: ...
def stream(
self,
@@ -2604,6 +2459,7 @@ class Pregel(
durability: Durability | None = None,
subgraphs: bool = False,
debug: bool | None = None,
stream_version: Literal["v1", "v2"] = "v1",
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[dict[str, Any] | Any]:
"""Stream graph steps for a single input.
@@ -2837,7 +2693,12 @@ 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,
stream_version,
)
loop.after_tick()
# wait for checkpoint
@@ -2845,7 +2706,12 @@ 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,
stream_version,
)
# handle exit
if loop.status == "out_of_steps":
@@ -2871,16 +2737,17 @@ class Pregel(
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["values"],
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: Literal[False] = False,
subgraphs: bool = False,
debug: bool | None = None,
stream_version: Literal["v2"],
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[OutputT]: ...
) -> AsyncIterator[StreamPart]: ...
@overload
def astream(
@@ -2889,160 +2756,17 @@ class Pregel(
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["updates"],
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: Literal[False] = False,
subgraphs: bool = False,
debug: bool | None = None,
stream_version: Literal["v1"] = ...,
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[dict[str, Any]]: ...
@overload
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["messages"],
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: Literal[False] = False,
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[tuple[AnyMessage, dict[str, Any]]]: ...
@overload
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["custom"],
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: Literal[False] = False,
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[Any]: ...
@overload
def astream(
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,
subgraphs: Literal[True],
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[tuple[tuple[str, ...], OutputT]]: ...
@overload
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["updates"],
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: Literal[True],
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[tuple[tuple[str, ...], dict[str, Any]]]: ...
@overload
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["messages"],
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: Literal[True],
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[tuple[tuple[str, ...], tuple[AnyMessage, dict[str, Any]]]]: ...
@overload
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["custom"],
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: Literal[True],
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[tuple[tuple[str, ...], Any]]: ...
@overload
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: list[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,
subgraphs: Literal[False] = False,
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[tuple[str, Any]]: ...
@overload
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: list[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,
subgraphs: Literal[True],
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[tuple[tuple[str, ...], str, Any]]: ...
) -> AsyncIterator[dict[str, Any] | Any]: ...
async def astream(
self,
@@ -3058,6 +2782,7 @@ class Pregel(
durability: Durability | None = None,
subgraphs: bool = False,
debug: bool | None = None,
stream_version: Literal["v1", "v2"] = "v1",
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[dict[str, Any] | Any]:
"""Asynchronously stream graph steps for a single input.
@@ -3350,6 +3075,7 @@ class Pregel(
subgraphs,
stream.get_nowait,
asyncio.QueueEmpty,
stream_version,
):
yield o
loop.after_tick()
@@ -3368,6 +3094,7 @@ class Pregel(
subgraphs,
stream.get_nowait,
asyncio.QueueEmpty,
stream_version,
):
yield o
# handle exit
@@ -3387,6 +3114,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,
stream_version: Literal["v2"],
**kwargs: Any,
) -> dict[str, Any]: ...
@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,
stream_version: Literal["v2"],
**kwargs: Any,
) -> list[StreamPart]: ...
@overload
def invoke(
self,
input: InputT | Command | None,
@@ -3399,6 +3161,23 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
stream_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,
stream_version: Literal["v1", "v2"] = "v1",
**kwargs: Any,
) -> dict[str, Any] | Any:
"""Run the graph with a single input and config.
@@ -3422,6 +3201,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.
stream_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:
@@ -3439,30 +3221,33 @@ class Pregel(
config,
context=context,
stream_mode=(
["updates", "values"] # type: ignore[arg-type]
if stream_mode == "values"
else 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,
stream_version=stream_version, # type: ignore[arg-type]
**kwargs,
):
if stream_mode == "values":
if len(chunk) == 2:
mode, payload = cast(tuple[StreamMode, Any], chunk)
if stream_version == "v2":
mode = chunk["type"]
payload = chunk["data"]
else:
_, mode, payload = cast(
tuple[tuple[str, ...], StreamMode, Any], chunk
)
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)
interrupts.extend(ints) # type: ignore[arg-type]
elif mode == "values":
latest = payload
else:
@@ -3479,6 +3264,41 @@ 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,
stream_version: Literal["v2"],
**kwargs: Any,
) -> dict[str, Any]: ...
@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,
stream_version: Literal["v2"],
**kwargs: Any,
) -> list[StreamPart]: ...
@overload
async def ainvoke(
self,
input: InputT | Command | None,
@@ -3491,6 +3311,23 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
stream_version: Literal["v1"] = ...,
**kwargs: Any,
) -> dict[str, Any] | Any: ...
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,
stream_version: Literal["v1", "v2"] = "v1",
**kwargs: Any,
) -> dict[str, Any] | Any:
"""Asynchronously run the graph with a single input and config.
@@ -3514,6 +3351,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.
stream_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:
@@ -3531,30 +3371,33 @@ class Pregel(
config,
context=context,
stream_mode=(
["updates", "values"] # type: ignore[arg-type]
if stream_mode == "values"
else 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,
stream_version=stream_version, # type: ignore[arg-type]
**kwargs,
):
if stream_mode == "values":
if len(chunk) == 2:
mode, payload = cast(tuple[StreamMode, Any], chunk)
if stream_version == "v2":
mode = chunk["type"]
payload = chunk["data"]
else:
_, mode, payload = cast(
tuple[tuple[str, ...], StreamMode, Any], chunk
)
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)
interrupts.extend(ints) # type: ignore[arg-type]
elif mode == "values":
latest = payload
else:
@@ -3625,6 +3468,7 @@ def _output(
stream_subgraphs: bool,
getter: Callable[[], tuple[tuple[str, ...], str, Any]],
empty_exc: type[Exception],
stream_version: Literal["v1", "v2"] = "v1",
) -> Iterator:
while True:
try:
@@ -3652,7 +3496,9 @@ def _output(
)
)
if mode in stream_mode:
if stream_subgraphs and isinstance(stream_mode, list):
if stream_version == "v2":
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)
libs/langgraph/langgraph/pregel/protocol.py
+81 -241
View File
@@ -2,16 +2,20 @@ from __future__ import annotations
from abc import abstractmethod
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
from typing import TYPE_CHECKING, Any, Generic, Literal, cast, overload
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
if TYPE_CHECKING:
from langchain_core.messages import AnyMessage
from langgraph.types import (
All,
Command,
StateSnapshot,
StateUpdate,
StreamMode,
StreamPart,
)
from langgraph.typing import ContextT, InputT, OutputT, StateT
__all__ = ("PregelProtocol", "StreamProtocol")
@@ -107,11 +111,12 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["values"],
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
) -> Iterator[OutputT]: ...
subgraphs: bool = False,
stream_version: Literal["v2"],
) -> Iterator[StreamPart]: ...
@overload
@abstractmethod
@@ -121,123 +126,12 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["updates"],
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
) -> Iterator[dict[str, Any]]: ...
@overload
@abstractmethod
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["messages"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
) -> Iterator[tuple[AnyMessage, dict[str, Any]]]: ...
@overload
@abstractmethod
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["custom"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
) -> Iterator[Any]: ...
@overload
@abstractmethod
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["values"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
) -> Iterator[tuple[tuple[str, ...], OutputT]]: ...
@overload
@abstractmethod
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["updates"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
) -> Iterator[tuple[tuple[str, ...], dict[str, Any]]]: ...
@overload
@abstractmethod
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["messages"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
) -> Iterator[tuple[tuple[str, ...], tuple[AnyMessage, dict[str, Any]]]]: ...
@overload
@abstractmethod
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["custom"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
) -> Iterator[tuple[tuple[str, ...], Any]]: ...
@overload
@abstractmethod
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: list[StreamMode],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
) -> Iterator[tuple[str, Any]]: ...
@overload
@abstractmethod
def stream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: list[StreamMode],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
) -> Iterator[tuple[tuple[str, ...], str, Any]]: ...
subgraphs: bool = False,
stream_version: Literal["v1"] = ...,
) -> Iterator[dict[str, Any] | Any]: ...
@abstractmethod
def stream(
@@ -250,6 +144,7 @@ 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,
stream_version: Literal["v1", "v2"] = "v1",
) -> Iterator[dict[str, Any] | Any]: ...
@overload
@@ -260,11 +155,12 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["values"],
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
) -> AsyncIterator[OutputT]: ...
subgraphs: bool = False,
stream_version: Literal["v2"],
) -> AsyncIterator[StreamPart]: ...
@overload
@abstractmethod
@@ -274,123 +170,12 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["updates"],
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
) -> AsyncIterator[dict[str, Any]]: ...
@overload
@abstractmethod
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["messages"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
) -> AsyncIterator[tuple[AnyMessage, dict[str, Any]]]: ...
@overload
@abstractmethod
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["custom"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
) -> AsyncIterator[Any]: ...
@overload
@abstractmethod
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["values"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
) -> AsyncIterator[tuple[tuple[str, ...], OutputT]]: ...
@overload
@abstractmethod
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["updates"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
) -> AsyncIterator[tuple[tuple[str, ...], dict[str, Any]]]: ...
@overload
@abstractmethod
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["messages"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
) -> AsyncIterator[tuple[tuple[str, ...], tuple[AnyMessage, dict[str, Any]]]]: ...
@overload
@abstractmethod
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: Literal["custom"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
) -> AsyncIterator[tuple[tuple[str, ...], Any]]: ...
@overload
@abstractmethod
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: list[StreamMode],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
) -> AsyncIterator[tuple[str, Any]]: ...
@overload
@abstractmethod
def astream(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
context: ContextT | None = None,
stream_mode: list[StreamMode],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
) -> AsyncIterator[tuple[tuple[str, ...], str, Any]]: ...
subgraphs: bool = False,
stream_version: Literal["v1"] = ...,
) -> AsyncIterator[dict[str, Any] | Any]: ...
@abstractmethod
def astream(
@@ -403,8 +188,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,
stream_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,
stream_version: Literal["v2"],
) -> dict[str, 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,
stream_version: Literal["v1"] = ...,
) -> dict[str, Any] | Any: ...
@abstractmethod
def invoke(
self,
@@ -414,6 +226,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,
stream_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,
stream_version: Literal["v2"],
) -> dict[str, 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,
stream_version: Literal["v1"] = ...,
) -> dict[str, Any] | Any: ...
@abstractmethod
@@ -425,6 +264,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,
stream_version: Literal["v1", "v2"] = "v1",
) -> dict[str, Any] | Any: ...
libs/langgraph/langgraph/pregel/remote.py
+97 -262
View File
@@ -4,7 +4,6 @@ import logging
from collections.abc import AsyncIterator, Iterator, Sequence
from dataclasses import asdict
from typing import (
TYPE_CHECKING,
Any,
Literal,
cast,
@@ -43,9 +42,6 @@ from langgraph_sdk.schema import (
)
from typing_extensions import Self
if TYPE_CHECKING:
from langchain_core.messages import AnyMessage
from langgraph._internal._config import merge_configs
from langgraph._internal._constants import (
CONF,
@@ -66,6 +62,7 @@ from langgraph.types import (
PregelTask,
StateSnapshot,
StreamMode,
StreamPart,
)
logger = logging.getLogger(__name__)
@@ -693,14 +690,15 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["values"],
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
subgraphs: bool = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
stream_version: Literal["v2"],
**kwargs: Any,
) -> Iterator[dict[str, Any]]: ...
) -> Iterator[StreamPart]: ...
@overload
def stream(
@@ -708,134 +706,15 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["updates"],
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
subgraphs: bool = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
stream_version: Literal["v1"] = ...,
**kwargs: Any,
) -> Iterator[dict[str, Any]]: ...
@overload
def stream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["messages"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> Iterator[tuple[AnyMessage, dict[str, Any]]]: ...
@overload
def stream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["custom"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> Iterator[Any]: ...
@overload
def stream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["values"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> Iterator[tuple[tuple[str, ...], dict[str, Any]]]: ...
@overload
def stream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["updates"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> Iterator[tuple[tuple[str, ...], dict[str, Any]]]: ...
@overload
def stream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["messages"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> Iterator[tuple[tuple[str, ...], tuple[AnyMessage, dict[str, Any]]]]: ...
@overload
def stream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["custom"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> Iterator[tuple[tuple[str, ...], Any]]: ...
@overload
def stream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: list[StreamMode],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> Iterator[tuple[str, Any]]: ...
@overload
def stream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: list[StreamMode],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> Iterator[tuple[tuple[str, ...], str, Any]]: ...
) -> Iterator[dict[str, Any] | Any]: ...
def stream(
self,
@@ -848,6 +727,7 @@ class RemoteGraph(PregelProtocol):
subgraphs: bool = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
stream_version: Literal["v1", "v2"] = "v1",
**kwargs: Any,
) -> Iterator[dict[str, Any] | Any]:
"""Create a run and stream the results.
@@ -929,10 +809,12 @@ 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 stream_version == "v2":
yield {"type": mode, "ns": ns, "data": chunk.data}
elif subgraphs:
if NS_SEP in chunk.event:
mode, ns_ = chunk.event.split(NS_SEP, 1)
ns = tuple(ns_.split(NS_SEP))
@@ -953,14 +835,15 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["values"],
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
subgraphs: bool = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
stream_version: Literal["v2"],
**kwargs: Any,
) -> AsyncIterator[dict[str, Any]]: ...
) -> AsyncIterator[StreamPart]: ...
@overload
def astream(
@@ -968,134 +851,15 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["updates"],
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
subgraphs: bool = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
stream_version: Literal["v1"] = ...,
**kwargs: Any,
) -> AsyncIterator[dict[str, Any]]: ...
@overload
def astream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["messages"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> AsyncIterator[tuple[AnyMessage, dict[str, Any]]]: ...
@overload
def astream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["custom"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> AsyncIterator[Any]: ...
@overload
def astream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["values"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> AsyncIterator[tuple[tuple[str, ...], dict[str, Any]]]: ...
@overload
def astream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["updates"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> AsyncIterator[tuple[tuple[str, ...], dict[str, Any]]]: ...
@overload
def astream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["messages"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> AsyncIterator[tuple[tuple[str, ...], tuple[AnyMessage, dict[str, Any]]]]: ...
@overload
def astream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: Literal["custom"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> AsyncIterator[tuple[tuple[str, ...], Any]]: ...
@overload
def astream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: list[StreamMode],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[False] = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> AsyncIterator[tuple[str, Any]]: ...
@overload
def astream(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
stream_mode: list[StreamMode],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: Literal[True],
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
**kwargs: Any,
) -> AsyncIterator[tuple[tuple[str, ...], str, Any]]: ...
) -> AsyncIterator[dict[str, Any] | Any]: ...
async def astream(
self,
@@ -1108,6 +872,7 @@ class RemoteGraph(PregelProtocol):
subgraphs: bool = False,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
stream_version: Literal["v1", "v2"] = "v1",
**kwargs: Any,
) -> AsyncIterator[dict[str, Any] | Any]:
"""Create a run and stream the results.
@@ -1189,10 +954,12 @@ 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 stream_version == "v2":
yield {"type": mode, "ns": ns, "data": chunk.data}
elif subgraphs:
if NS_SEP in chunk.event:
mode, ns_ = chunk.event.split(NS_SEP, 1)
ns = tuple(ns_.split(NS_SEP))
@@ -1223,6 +990,7 @@ class RemoteGraph(PregelProtocol):
) -> AsyncIterator[dict[str, Any]]:
raise NotImplementedError
@overload
def invoke(
self,
input: dict[str, Any] | Any,
@@ -1232,6 +1000,34 @@ class RemoteGraph(PregelProtocol):
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
stream_version: Literal["v2"],
**kwargs: Any,
) -> 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,
stream_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,
stream_version: Literal["v1", "v2"] = "v1",
**kwargs: Any,
) -> dict[str, Any] | Any:
"""Create a run, wait until it finishes and return the final state.
@@ -1242,12 +1038,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.
stream_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,
@@ -1255,15 +1053,46 @@ class RemoteGraph(PregelProtocol):
headers=headers,
stream_mode="values",
params=params,
stream_version=stream_version,
**kwargs,
):
pass
try:
if stream_version == "v2":
return chunk["data"]
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,
stream_version: Literal["v2"],
**kwargs: Any,
) -> 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,
stream_version: Literal["v1"] = ...,
**kwargs: Any,
) -> dict[str, Any] | Any: ...
async def ainvoke(
self,
input: dict[str, Any] | Any,
@@ -1273,6 +1102,7 @@ class RemoteGraph(PregelProtocol):
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
params: QueryParamTypes | None = None,
stream_version: Literal["v1", "v2"] = "v1",
**kwargs: Any,
) -> dict[str, Any] | Any:
"""Create a run, wait until it finishes and return the final state.
@@ -1283,12 +1113,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.
stream_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,
@@ -1296,10 +1128,13 @@ class RemoteGraph(PregelProtocol):
headers=headers,
stream_mode="values",
params=params,
stream_version=stream_version,
**kwargs,
):
pass
try:
if stream_version == "v2":
return chunk["data"]
return chunk
except UnboundLocalError:
logger.warning("No events received from remote graph")
libs/langgraph/langgraph/types.py
+191 -1
View File
@@ -16,9 +16,10 @@ 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, TypedDict, Unpack, deprecated
from xxhash import xxh3_128_hexdigest
from langgraph._internal._cache import default_cache_key
@@ -44,6 +45,19 @@ __all__ = (
"Checkpointer",
"StreamMode",
"StreamWriter",
"StreamPart",
"ValuesStreamPart",
"UpdatesStreamPart",
"MessagesStreamPart",
"CustomStreamPart",
"CheckpointStreamPart",
"TasksStreamPart",
"DebugStreamPart",
"TaskPayload",
"TaskResultPayload",
"CheckpointTask",
"CheckpointPayload",
"DebugPayload",
"RetryPolicy",
"CachePolicy",
"Interrupt",
@@ -113,6 +127,182 @@ 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
name: str
input: Any
triggers: list[str]
class TaskResultPayload(TypedDict):
"""Payload for a task result event."""
id: str
name: str
error: str | None
interrupts: list[dict]
result: dict[str, Any]
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
name: str
error: NotRequired[str]
result: NotRequired[Any]
interrupts: NotRequired[list[dict]]
state: StateSnapshot | RunnableConfig | None
class CheckpointPayload(TypedDict):
"""Payload for a checkpoint event."""
config: RunnableConfig | None
metadata: CheckpointMetadata
values: dict[str, Any]
next: list[str]
parent_config: RunnableConfig | None
tasks: list[CheckpointTask]
class _DebugCheckpointPayload(TypedDict):
step: int
timestamp: str
type: Literal["checkpoint"]
payload: CheckpointPayload
class _DebugTaskPayload(TypedDict):
step: int
timestamp: str
type: Literal["task"]
payload: TaskPayload
class _DebugTaskResultPayload(TypedDict):
step: int
timestamp: str
type: Literal["task_result"]
payload: TaskResultPayload
DebugPayload = _DebugCheckpointPayload | _DebugTaskPayload | _DebugTaskResultPayload
"""Wrapper payload for debug events. Discriminate on `type`."""
class ValuesStreamPart(TypedDict):
"""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: dict[str, Any]
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):
"""Stream part emitted for `stream_mode="checkpoints"`."""
type: Literal["checkpoints"]
ns: tuple[str, ...]
data: CheckpointPayload
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):
"""Stream part emitted for `stream_mode="debug"`."""
type: Literal["debug"]
ns: tuple[str, ...]
data: DebugPayload
StreamPart = (
ValuesStreamPart
| UpdatesStreamPart
| MessagesStreamPart
| CustomStreamPart
| CheckpointStreamPart
| TasksStreamPart
| DebugStreamPart
)
"""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"] # dict[str, Any] — full state
elif part["type"] == "messages":
part["data"] # tuple[BaseMessage, dict] — (message, metadata)
elif part["type"] == "custom":
part["data"] # Any — user-defined
```
"""
_DC_KWARGS = {"kw_only": True, "slots": True, "frozen": True}
libs/langgraph/tests/test_stream_v2.py
+622
View File
@@ -0,0 +1,622 @@
"""Tests for v2 streaming format (StreamPart TypedDicts).
This file is checked by mypy directly — no subprocess workarounds.
Type-narrowing is validated via `assert_type` calls in `_check_type_narrowing`.
"""
from __future__ import annotations
import operator
import sys
from typing import Annotated, Any
import pytest
from langchain_core.messages import AIMessage, BaseMessage
from langgraph.checkpoint.memory import InMemorySaver
from typing_extensions import TypedDict, assert_type
from langgraph.constants import END, START
from langgraph.func import entrypoint
from langgraph.graph import StateGraph
from langgraph.graph.message import MessagesState
from langgraph.types import (
CheckpointPayload,
CheckpointStreamPart,
CustomStreamPart,
DebugPayload,
DebugStreamPart,
MessagesStreamPart,
StreamPart,
StreamWriter,
TaskPayload,
TaskResultPayload,
TasksStreamPart,
UpdatesStreamPart,
ValuesStreamPart,
)
from tests.fake_chat import FakeChatModel
NEEDS_CONTEXTVARS = pytest.mark.skipif(
sys.version_info < (3, 11),
reason="Python 3.11+ is required for async contextvars support",
)
# --- state and graph builders ---
class SimpleState(TypedDict):
value: str
items: Annotated[list[str], operator.add]
_SIMPLE_INPUT: SimpleState = {"value": "x", "items": []}
_MSG_INPUT: MessagesState = {"messages": "hi"}
def _make_simple_graph() -> StateGraph[SimpleState, None, SimpleState, SimpleState]:
def node_a(state: SimpleState) -> dict[str, Any]:
return {"value": state["value"] + "_a", "items": ["a"]}
def node_b(state: SimpleState) -> dict[str, Any]:
return {"value": state["value"] + "_b", "items": ["b"]}
builder = StateGraph(SimpleState, input_schema=SimpleState)
builder.add_node("node_a", node_a)
builder.add_node("node_b", node_b)
builder.add_edge(START, "node_a")
builder.add_edge("node_a", "node_b")
builder.add_edge("node_b", END)
return builder
def _make_messages_graph() -> StateGraph[
MessagesState, None, MessagesState, MessagesState
]:
model = FakeChatModel(messages=[AIMessage(content="hello world")])
def call_model(state: MessagesState) -> dict[str, Any]:
return {"messages": model.invoke(state["messages"])}
builder = StateGraph(MessagesState, input_schema=MessagesState)
builder.add_node("call_model", call_model)
builder.add_edge(START, "call_model")
builder.add_edge("call_model", END)
return builder
def _make_custom_graph() -> Any:
@entrypoint()
def graph(inputs: Any, *, writer: StreamWriter) -> Any:
writer("hello")
writer(42)
return inputs
return graph
def _make_subgraph() -> Any:
inner = _make_simple_graph().compile()
outer_builder = StateGraph(SimpleState, input_schema=SimpleState)
outer_builder.add_node("inner", inner)
outer_builder.add_edge(START, "inner")
outer_builder.add_edge("inner", END)
return outer_builder.compile()
# --- shared assertion helpers ---
_STREAM_PART_KEYS = {"type", "ns", "data"}
def _assert_stream_part_shape(part: StreamPart) -> None:
"""Assert a v2 stream part has the required keys and correct types."""
assert isinstance(part, dict), f"Expected dict, got {type(part)}"
assert _STREAM_PART_KEYS <= part.keys(), (
f"Missing keys: {_STREAM_PART_KEYS - part.keys()}"
)
assert isinstance(part["type"], str)
assert isinstance(part["ns"], tuple)
for elem in part["ns"]:
assert isinstance(elem, str)
# --- v1 backwards compatibility ---
class TestV1BackwardsCompat:
def test_stream_default_is_v1(self) -> None:
graph = _make_simple_graph().compile()
chunks = list(graph.stream(_SIMPLE_INPUT))
for chunk in chunks:
assert isinstance(chunk, dict)
def test_stream_v1_updates_mode(self) -> None:
graph = _make_simple_graph().compile()
chunks = list(graph.stream(_SIMPLE_INPUT, stream_mode="updates"))
assert len(chunks) == 2
assert "node_a" in chunks[0]
assert "node_b" in chunks[1]
def test_stream_v1_list_mode(self) -> None:
graph = _make_simple_graph().compile()
chunks = list(graph.stream(_SIMPLE_INPUT, stream_mode=["values", "updates"]))
for chunk in chunks:
assert isinstance(chunk, tuple) and len(chunk) == 2
mode, _data = chunk
assert mode in ("values", "updates")
def test_stream_v1_subgraphs(self) -> None:
graph = _make_simple_graph().compile()
chunks = list(
graph.stream(_SIMPLE_INPUT, stream_mode="updates", subgraphs=True)
)
for chunk in chunks:
assert isinstance(chunk, tuple) and len(chunk) == 2
ns, _data = chunk
assert isinstance(ns, tuple)
# --- v2 sync stream ---
class TestV2Stream:
def test_values(self) -> None:
graph = _make_simple_graph().compile()
chunks = list(
graph.stream(_SIMPLE_INPUT, stream_mode="values", stream_version="v2")
)
assert len(chunks) >= 1
for c in chunks:
_assert_stream_part_shape(c)
assert c["type"] == "values"
assert c["ns"] == ()
assert isinstance(c["data"], dict)
def test_updates(self) -> None:
graph = _make_simple_graph().compile()
chunks = list(
graph.stream(_SIMPLE_INPUT, stream_mode="updates", stream_version="v2")
)
assert len(chunks) == 2
for c in chunks:
_assert_stream_part_shape(c)
assert c["type"] == "updates"
assert c["ns"] == ()
assert "node_a" in chunks[0]["data"]
assert "node_b" in chunks[1]["data"]
def test_messages(self) -> None:
graph = _make_messages_graph().compile()
chunks = list(
graph.stream(_MSG_INPUT, stream_mode="messages", stream_version="v2")
)
msg_chunks = [c for c in chunks if c["type"] == "messages"]
assert len(msg_chunks) >= 1
for c in msg_chunks:
_assert_stream_part_shape(c)
assert c["ns"] == ()
data = c["data"]
assert isinstance(data, tuple) and len(data) == 2
message, metadata = data
assert isinstance(message, BaseMessage)
assert isinstance(metadata, dict)
assert "langgraph_node" in metadata
def test_custom(self) -> None:
graph = _make_custom_graph()
chunks = list(
graph.stream({"key": "val"}, stream_mode="custom", stream_version="v2")
)
custom = [c for c in chunks if c["type"] == "custom"]
assert len(custom) == 2
for c in custom:
_assert_stream_part_shape(c)
assert custom[0]["data"] == "hello"
assert custom[1]["data"] == 42
def test_multiple_modes(self) -> None:
graph = _make_simple_graph().compile()
chunks = list(
graph.stream(
_SIMPLE_INPUT,
stream_mode=["values", "updates"],
stream_version="v2",
)
)
types_seen = {c["type"] for c in chunks}
assert {"values", "updates"} <= types_seen
for c in chunks:
_assert_stream_part_shape(c)
def test_subgraphs_ns(self) -> None:
outer = _make_subgraph()
chunks = list(
outer.stream(
_SIMPLE_INPUT,
stream_mode="updates",
subgraphs=True,
stream_version="v2",
)
)
for c in chunks:
_assert_stream_part_shape(c)
root = [c for c in chunks if c["ns"] == ()]
sub = [c for c in chunks if c["ns"] != ()]
assert len(root) >= 1
assert len(sub) >= 1
def test_checkpoints(self) -> None:
graph = _make_simple_graph().compile(checkpointer=InMemorySaver())
config: Any = {"configurable": {"thread_id": "test-v2-ckpt"}}
chunks = list(
graph.stream(
_SIMPLE_INPUT,
config,
stream_mode="checkpoints",
stream_version="v2",
)
)
ckpt = [c for c in chunks if c["type"] == "checkpoints"]
assert len(ckpt) >= 1
for c in ckpt:
_assert_stream_part_shape(c)
assert c["ns"] == ()
payload = c["data"]
assert {"config", "metadata", "values", "next", "tasks"} <= payload.keys()
def test_tasks(self) -> None:
graph = _make_simple_graph().compile(checkpointer=InMemorySaver())
config: Any = {"configurable": {"thread_id": "test-v2-tasks"}}
chunks = list(
graph.stream(
_SIMPLE_INPUT,
config,
stream_mode="tasks",
stream_version="v2",
)
)
tasks = [c for c in chunks if c["type"] == "tasks"]
assert len(tasks) >= 2
for c in tasks:
_assert_stream_part_shape(c)
assert c["ns"] == ()
assert "id" in c["data"] and "name" in c["data"]
starts = [c for c in tasks if "triggers" in c["data"]]
results = [c for c in tasks if "result" in c["data"]]
assert len(starts) >= 2
assert len(results) >= 2
def test_debug(self) -> None:
graph = _make_simple_graph().compile(checkpointer=InMemorySaver())
config: Any = {"configurable": {"thread_id": "test-v2-debug"}}
chunks = list(
graph.stream(
_SIMPLE_INPUT,
config,
stream_mode="debug",
stream_version="v2",
)
)
debug = [c for c in chunks if c["type"] == "debug"]
assert len(debug) >= 1
for c in debug:
_assert_stream_part_shape(c)
assert c["ns"] == ()
envelope = c["data"]
assert {"step", "timestamp", "type", "payload"} <= envelope.keys()
assert envelope["type"] in ("checkpoint", "task", "task_result")
def test_subgraphs_param_does_not_change_format(self) -> None:
"""In v2, subgraphs=True/False should not change the output format."""
graph = _make_simple_graph().compile()
chunks_no_sub = list(
graph.stream(
_SIMPLE_INPUT,
stream_mode="updates",
subgraphs=False,
stream_version="v2",
)
)
chunks_with_sub = list(
graph.stream(
_SIMPLE_INPUT,
stream_mode="updates",
subgraphs=True,
stream_version="v2",
)
)
for c in chunks_no_sub + chunks_with_sub:
_assert_stream_part_shape(c)
# --- v2 sync invoke ---
class TestV2Invoke:
def test_values_default(self) -> None:
graph = _make_simple_graph().compile()
result = graph.invoke(_SIMPLE_INPUT, stream_version="v2")
assert isinstance(result, dict)
assert result["value"] == "x_a_b"
assert result["items"] == ["a", "b"]
def test_updates_mode(self) -> None:
graph = _make_simple_graph().compile()
result = graph.invoke(_SIMPLE_INPUT, stream_mode="updates", stream_version="v2")
assert isinstance(result, list) and len(result) == 2
for chunk in result:
_assert_stream_part_shape(chunk)
assert chunk["type"] == "updates"
assert chunk["ns"] == ()
assert "node_a" in result[0]["data"]
assert "node_b" in result[1]["data"]
def test_multiple_modes(self) -> None:
graph = _make_simple_graph().compile()
modes: Any = ["values", "updates"]
result = graph.invoke(_SIMPLE_INPUT, stream_mode=modes, stream_version="v2")
assert isinstance(result, list)
types_seen = {c["type"] for c in result}
assert {"values", "updates"} <= types_seen
for c in result:
_assert_stream_part_shape(c)
def test_v1_default_unchanged(self) -> None:
graph = _make_simple_graph().compile()
result = graph.invoke(_SIMPLE_INPUT)
assert isinstance(result, dict)
assert result["value"] == "x_a_b"
assert result["items"] == ["a", "b"]
def test_v1_updates_unchanged(self) -> None:
graph = _make_simple_graph().compile()
result = graph.invoke(_SIMPLE_INPUT, stream_mode="updates")
assert isinstance(result, list)
for chunk in result:
assert "node_a" in chunk or "node_b" in chunk
# --- v2 async stream ---
class TestV2StreamAsync:
@pytest.mark.anyio
async def test_values(self) -> None:
graph = _make_simple_graph().compile()
chunks = [
c
async for c in graph.astream(
_SIMPLE_INPUT, stream_mode="values", stream_version="v2"
)
]
assert len(chunks) >= 1
for c in chunks:
_assert_stream_part_shape(c)
assert c["type"] == "values"
assert c["ns"] == ()
assert isinstance(c["data"], dict)
@pytest.mark.anyio
async def test_updates(self) -> None:
graph = _make_simple_graph().compile()
chunks = [
c
async for c in graph.astream(
_SIMPLE_INPUT, stream_mode="updates", stream_version="v2"
)
]
assert len(chunks) == 2
for c in chunks:
_assert_stream_part_shape(c)
assert c["type"] == "updates"
assert c["ns"] == ()
assert "node_a" in chunks[0]["data"]
assert "node_b" in chunks[1]["data"]
@pytest.mark.anyio
async def test_messages(self) -> None:
graph = _make_messages_graph().compile()
chunks = [
c
async for c in graph.astream(
_MSG_INPUT, stream_mode="messages", stream_version="v2"
)
]
msg_chunks = [c for c in chunks if c["type"] == "messages"]
assert len(msg_chunks) >= 1
for c in msg_chunks:
_assert_stream_part_shape(c)
assert c["ns"] == ()
data = c["data"]
assert isinstance(data, tuple) and len(data) == 2
message, metadata = data
assert isinstance(message, BaseMessage)
assert isinstance(metadata, dict)
assert "langgraph_node" in metadata
@NEEDS_CONTEXTVARS
@pytest.mark.anyio
async def test_custom(self) -> None:
graph = _make_custom_graph()
chunks = [
c
async for c in graph.astream(
{"key": "val"}, stream_mode="custom", stream_version="v2"
)
]
custom = [c for c in chunks if c["type"] == "custom"]
assert len(custom) == 2
for c in custom:
_assert_stream_part_shape(c)
assert custom[0]["data"] == "hello"
assert custom[1]["data"] == 42
@pytest.mark.anyio
async def test_multiple_modes(self) -> None:
graph = _make_simple_graph().compile()
chunks = [
c
async for c in graph.astream(
_SIMPLE_INPUT,
stream_mode=["values", "updates"],
stream_version="v2",
)
]
types_seen = {c["type"] for c in chunks}
assert {"values", "updates"} <= types_seen
for c in chunks:
_assert_stream_part_shape(c)
@pytest.mark.anyio
async def test_subgraphs_ns(self) -> None:
outer = _make_subgraph()
chunks = [
c
async for c in outer.astream(
_SIMPLE_INPUT,
stream_mode="updates",
subgraphs=True,
stream_version="v2",
)
]
for c in chunks:
_assert_stream_part_shape(c)
root = [c for c in chunks if c["ns"] == ()]
sub = [c for c in chunks if c["ns"] != ()]
assert len(root) >= 1
assert len(sub) >= 1
@pytest.mark.anyio
async def test_checkpoints(self) -> None:
graph = _make_simple_graph().compile(checkpointer=InMemorySaver())
config: Any = {"configurable": {"thread_id": "test-v2-ckpt-async"}}
chunks = [
c
async for c in graph.astream(
_SIMPLE_INPUT,
config,
stream_mode="checkpoints",
stream_version="v2",
)
]
ckpt = [c for c in chunks if c["type"] == "checkpoints"]
assert len(ckpt) >= 1
for c in ckpt:
_assert_stream_part_shape(c)
assert c["ns"] == ()
payload = c["data"]
assert {"config", "metadata", "values", "next", "tasks"} <= payload.keys()
@pytest.mark.anyio
async def test_tasks(self) -> None:
graph = _make_simple_graph().compile(checkpointer=InMemorySaver())
config: Any = {"configurable": {"thread_id": "test-v2-tasks-async"}}
chunks = [
c
async for c in graph.astream(
_SIMPLE_INPUT,
config,
stream_mode="tasks",
stream_version="v2",
)
]
tasks = [c for c in chunks if c["type"] == "tasks"]
assert len(tasks) >= 2
for c in tasks:
_assert_stream_part_shape(c)
assert c["ns"] == ()
assert "id" in c["data"] and "name" in c["data"]
starts = [c for c in tasks if "triggers" in c["data"]]
results = [c for c in tasks if "result" in c["data"]]
assert len(starts) >= 2
assert len(results) >= 2
@pytest.mark.anyio
async def test_debug(self) -> None:
graph = _make_simple_graph().compile(checkpointer=InMemorySaver())
config: Any = {"configurable": {"thread_id": "test-v2-debug-async"}}
chunks = [
c
async for c in graph.astream(
_SIMPLE_INPUT,
config,
stream_mode="debug",
stream_version="v2",
)
]
debug = [c for c in chunks if c["type"] == "debug"]
assert len(debug) >= 1
for c in debug:
_assert_stream_part_shape(c)
assert c["ns"] == ()
envelope = c["data"]
assert {"step", "timestamp", "type", "payload"} <= envelope.keys()
assert envelope["type"] in ("checkpoint", "task", "task_result")
# --- v2 async invoke ---
class TestV2InvokeAsync:
@pytest.mark.anyio
async def test_values_default(self) -> None:
graph = _make_simple_graph().compile()
result = await graph.ainvoke(_SIMPLE_INPUT, stream_version="v2")
assert isinstance(result, dict)
assert result["value"] == "x_a_b"
assert result["items"] == ["a", "b"]
@pytest.mark.anyio
async def test_updates_mode(self) -> None:
graph = _make_simple_graph().compile()
result = await graph.ainvoke(
_SIMPLE_INPUT, stream_mode="updates", stream_version="v2"
)
assert isinstance(result, list) and len(result) == 2
for chunk in result:
_assert_stream_part_shape(chunk)
assert chunk["type"] == "updates"
assert chunk["ns"] == ()
assert "node_a" in result[0]["data"]
assert "node_b" in result[1]["data"]
@pytest.mark.anyio
async def test_multiple_modes(self) -> None:
graph = _make_simple_graph().compile()
modes: Any = ["values", "updates"]
result = await graph.ainvoke(
_SIMPLE_INPUT, stream_mode=modes, stream_version="v2"
)
assert isinstance(result, list)
types_seen = {c["type"] for c in result}
assert {"values", "updates"} <= types_seen
for c in result:
_assert_stream_part_shape(c)
# --- type narrowing compile-time checks ---
# These assert_type calls verify that mypy narrows the union correctly.
def _check_type_narrowing(part: StreamPart) -> None:
"""Compile-time type narrowing checks — never called at runtime."""
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"] == "messages":
assert_type(part, MessagesStreamPart)
elif part["type"] == "custom":
assert_type(part, CustomStreamPart)
elif part["type"] == "checkpoints":
assert_type(part, CheckpointStreamPart)
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)
assert_type(part["ns"], tuple[str, ...])
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,
stream_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,
stream_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,
stream_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]: ...
stream_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]:
stream_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
stream_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 stream_version == "v2":
return _wrap_stream_v2(raw)
return raw
@overload
async def create(
libs/sdk-py/langgraph_sdk/_shared/utilities.py
+13
View File
@@ -107,6 +107,19 @@ 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 []
return {"type": event_type, "ns": ns, "data": data}
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,
stream_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,
stream_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,
stream_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]: ...
stream_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]:
stream_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
stream_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 stream_version == "v2":
return _wrap_stream_v2_sync(raw)
return raw
@overload
def create(
libs/sdk-py/langgraph_sdk/schema.py
+200
View File
@@ -588,6 +588,206 @@ 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
name: str
input: Any
triggers: list[str]
class TaskResultPayload(TypedDict):
"""Payload for a task result event."""
id: str
name: str
error: str | None
interrupts: list[dict[str, Any]]
result: dict[str, Any]
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
name: str
error: NotRequired[str]
result: NotRequired[Any]
interrupts: NotRequired[list[dict[str, Any]]]
state: dict[str, Any] | None
class CheckpointPayload(TypedDict):
"""Payload for a checkpoint event."""
config: dict[str, Any] | None
metadata: dict[str, Any]
values: dict[str, Any]
next: list[str]
parent_config: dict[str, Any] | None
tasks: list[CheckpointTaskPayload]
class _DebugCheckpointPayload(TypedDict):
step: int
timestamp: str
type: Literal["checkpoint"]
payload: CheckpointPayload
class _DebugTaskPayload(TypedDict):
step: int
timestamp: str
type: Literal["task"]
payload: TaskPayload
class _DebugTaskResultPayload(TypedDict):
step: int
timestamp: str
type: Literal["task_result"]
payload: TaskResultPayload
DebugPayload = _DebugCheckpointPayload | _DebugTaskPayload | _DebugTaskResultPayload
"""Wrapper payload for debug events. Discriminate on ``type``."""
class RunMetadataPayload(TypedDict):
"""Payload for the ``metadata`` control event."""
run_id: str
# --- v2 stream part TypedDicts ---
class ValuesStreamPart(TypedDict):
"""Stream part emitted for ``stream_mode="values"``."""
type: Literal["values"]
ns: list[str]
data: dict[str, Any]
class UpdatesStreamPart(TypedDict):
"""Stream part emitted for ``stream_mode="updates"``."""
type: Literal["updates"]
ns: list[str]
data: dict[str, Any]
class MessagesPartialStreamPart(TypedDict):
"""Stream part emitted for partial message chunks (``messages/partial``)."""
type: Literal["messages/partial"]
ns: list[str]
data: list[dict[str, Any]]
class MessagesCompleteStreamPart(TypedDict):
"""Stream part emitted for complete messages (``messages/complete``)."""
type: Literal["messages/complete"]
ns: list[str]
data: list[dict[str, Any]]
class MessagesMetadataStreamPart(TypedDict):
"""Stream part emitted for message metadata (``messages/metadata``)."""
type: Literal["messages/metadata"]
ns: list[str]
data: dict[str, Any]
class MessagesTupleStreamPart(TypedDict):
"""Stream part emitted for ``stream_mode="messages"`` (raw message+metadata pair)."""
type: Literal["messages"]
ns: list[str]
data: list[dict[str, Any]]
class CustomStreamPart(TypedDict):
"""Stream part emitted for ``stream_mode="custom"``."""
type: Literal["custom"]
ns: list[str]
data: Any
class CheckpointsStreamPart(TypedDict):
"""Stream part emitted for ``stream_mode="checkpoints"``."""
type: Literal["checkpoints"]
ns: list[str]
data: CheckpointPayload
class TasksStreamPart(TypedDict):
"""Stream part emitted for ``stream_mode="tasks"``."""
type: Literal["tasks"]
ns: list[str]
data: TaskPayload | TaskResultPayload
class DebugStreamPart(TypedDict):
"""Stream part emitted for ``stream_mode="debug"``."""
type: Literal["debug"]
ns: list[str]
data: DebugPayload
class MetadataStreamPart(TypedDict):
"""Control event with ``run_id`` and other run metadata."""
type: Literal["metadata"]
ns: list[str]
data: RunMetadataPayload
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
+194 -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,138 @@ 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"}]},
}
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"},
}
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",
}
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"},
}
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"}],
}
# --- 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"}}
assert parts[1] == {
"type": "values",
"ns": [],
"data": {"messages": [{"role": "user", "content": "hi"}]},
}
assert parts[2] == {
"type": "updates",
"ns": ["sub:abc"],
"data": {"node": {"out": 1}},
}
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"}}
assert parts[1] == {"type": "values", "ns": [], "data": {"state": "full"}}
# --- 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)