mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
975a87c85e | ||
|
|
a7351aa134 | ||
|
|
7a4ce06a37 | ||
|
|
803a268b39 |
@@ -0,0 +1,807 @@
|
||||
"""Protocol-native content-block message handler for StreamingHandler.
|
||||
|
||||
Emits structured content-block lifecycle events (message-start,
|
||||
content-block-start/delta/finish, message-finish) instead of raw
|
||||
``(AIMessageChunk, metadata)`` tuples. The existing
|
||||
:class:`~langgraph.pregel._messages.StreamMessagesHandler` is NOT
|
||||
modified — this handler is only activated when
|
||||
``__protocol_messages_stream`` is ``True`` in the run's configurable.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, TypeVar, cast
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
from langchain_core.callbacks import BaseCallbackHandler
|
||||
from langchain_core.messages import AIMessageChunk, BaseMessage
|
||||
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult
|
||||
|
||||
from langgraph._internal._constants import NS_SEP
|
||||
from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
|
||||
from langgraph.pregel.protocol import StreamChunk
|
||||
from langgraph.stream._types import (
|
||||
ContentBlockDeltaData,
|
||||
ContentBlockFinishData,
|
||||
ContentBlockStartData,
|
||||
FinishReason,
|
||||
InvalidToolCallBlock,
|
||||
MessageErrorData,
|
||||
MessageStartData,
|
||||
ReasoningBlock,
|
||||
TextBlock,
|
||||
ToolCallBlock,
|
||||
UsageInfo,
|
||||
)
|
||||
|
||||
try:
|
||||
from langchain_core.tracers._streaming import _StreamingCallbackHandler
|
||||
except ImportError:
|
||||
_StreamingCallbackHandler = object # type: ignore
|
||||
|
||||
T = TypeVar("T")
|
||||
Meta = tuple[tuple[str, ...], dict[str, Any]]
|
||||
|
||||
PROTOCOL_MESSAGES_STREAM_KEY = "__protocol_messages_stream"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Content-block accumulation helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# A "compatible content block" is a dict matching one of the protocol block
|
||||
# TypedDicts (TextBlock, ReasoningBlock, ToolCallChunkBlock, etc.).
|
||||
CompatBlock = dict[str, Any]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ProtocolRunState:
|
||||
"""Per-run state for tracking the active message lifecycle."""
|
||||
|
||||
message_id: str | None = None
|
||||
started: bool = False
|
||||
blocks: dict[int, CompatBlock] = field(default_factory=dict)
|
||||
usage: dict[str, Any] | None = None
|
||||
|
||||
|
||||
def _accumulate_block(accumulated: CompatBlock, delta: CompatBlock) -> CompatBlock:
|
||||
"""Merge *delta* into *accumulated*, returning the updated block."""
|
||||
btype = accumulated.get("type", "text")
|
||||
if btype == "text" and delta.get("type", "text") == "text":
|
||||
accumulated["text"] = accumulated.get("text", "") + delta.get("text", "")
|
||||
elif btype == "reasoning" and delta.get("type") == "reasoning":
|
||||
accumulated["reasoning"] = accumulated.get("reasoning", "") + delta.get(
|
||||
"reasoning", ""
|
||||
)
|
||||
elif btype == "tool_call_chunk" and delta.get("type") == "tool_call_chunk":
|
||||
accumulated["args"] = accumulated.get("args", "") + delta.get("args", "")
|
||||
if delta.get("id") is not None:
|
||||
accumulated["id"] = delta["id"]
|
||||
if delta.get("name") is not None:
|
||||
accumulated["name"] = delta["name"]
|
||||
return accumulated
|
||||
|
||||
|
||||
def _delta_block(previous: CompatBlock, current: CompatBlock) -> CompatBlock | None:
|
||||
"""Compute the delta between *previous* and *current*.
|
||||
|
||||
Returns ``None`` if there is nothing new to emit.
|
||||
"""
|
||||
btype = current.get("type", "text")
|
||||
if btype == "text":
|
||||
prev_text = previous.get("text", "")
|
||||
cur_text = current.get("text", "")
|
||||
delta_text = cur_text[len(prev_text) :]
|
||||
if not delta_text:
|
||||
return None
|
||||
return TextBlock(type="text", text=delta_text)
|
||||
elif btype == "reasoning":
|
||||
prev_r = previous.get("reasoning", "")
|
||||
cur_r = current.get("reasoning", "")
|
||||
delta_r = cur_r[len(prev_r) :]
|
||||
if not delta_r:
|
||||
return None
|
||||
return ReasoningBlock(type="reasoning", reasoning=delta_r)
|
||||
elif btype == "tool_call_chunk":
|
||||
prev_args = previous.get("args", "")
|
||||
cur_args = current.get("args", "")
|
||||
delta_args = cur_args[len(prev_args) :]
|
||||
has_meta = current.get("id") is not None or current.get("name") is not None
|
||||
if not delta_args and not has_meta:
|
||||
return None
|
||||
result: CompatBlock = {"type": "tool_call_chunk", "args": delta_args}
|
||||
if current.get("id") is not None and previous.get("id") is None:
|
||||
result["id"] = current["id"]
|
||||
if current.get("name") is not None and previous.get("name") is None:
|
||||
result["name"] = current["name"]
|
||||
return result
|
||||
# Unrecognized block type — pass through unchanged
|
||||
return current
|
||||
|
||||
|
||||
def _finalize_block(block: CompatBlock) -> CompatBlock:
|
||||
"""Convert a ``tool_call_chunk`` block to a finalized ``tool_call`` or
|
||||
``invalid_tool_call`` block. Other block types pass through unchanged.
|
||||
"""
|
||||
if block.get("type") != "tool_call_chunk":
|
||||
return block
|
||||
raw_args = block.get("args", "{}")
|
||||
try:
|
||||
parsed_args = json.loads(raw_args) if raw_args else {}
|
||||
return ToolCallBlock(
|
||||
type="tool_call",
|
||||
id=block.get("id", ""),
|
||||
name=block.get("name", ""),
|
||||
args=parsed_args,
|
||||
)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return InvalidToolCallBlock(
|
||||
type="invalid_tool_call",
|
||||
id=block.get("id"),
|
||||
name=block.get("name"),
|
||||
args=raw_args,
|
||||
error="Failed to parse tool call arguments as JSON",
|
||||
)
|
||||
|
||||
|
||||
def _normalize_finish_reason(value: Any) -> FinishReason:
|
||||
"""Map provider-specific stop reasons to protocol finish reasons."""
|
||||
if value == "length":
|
||||
return "length"
|
||||
if value == "content_filter":
|
||||
return "content_filter"
|
||||
if value in ("tool_use", "tool_calls"):
|
||||
return "tool_use"
|
||||
# "end_turn", "stop", None, and anything else → "stop"
|
||||
return "stop"
|
||||
|
||||
|
||||
def _accumulate_usage(
|
||||
current: dict[str, Any] | None, delta: Any
|
||||
) -> dict[str, Any] | None:
|
||||
"""Accumulate usage metadata from streamed chunks."""
|
||||
if not isinstance(delta, dict):
|
||||
return current
|
||||
if current is None:
|
||||
return dict(delta)
|
||||
for key in ("input_tokens", "output_tokens", "total_tokens", "cached_tokens"):
|
||||
if key in delta:
|
||||
current[key] = current.get(key, 0) + delta[key]
|
||||
# Merge detail dicts
|
||||
for detail_key in ("input_token_details", "output_token_details"):
|
||||
if detail_key in delta and isinstance(delta[detail_key], dict):
|
||||
if detail_key not in current:
|
||||
current[detail_key] = {}
|
||||
current[detail_key].update(delta[detail_key])
|
||||
return current
|
||||
|
||||
|
||||
def _to_protocol_usage(usage: dict[str, Any] | None) -> UsageInfo | None:
|
||||
"""Convert LangChain usage metadata to protocol ``UsageInfo``."""
|
||||
if usage is None:
|
||||
return None
|
||||
result: dict[str, Any] = {}
|
||||
if "input_tokens" in usage:
|
||||
result["input_tokens"] = usage["input_tokens"]
|
||||
if "output_tokens" in usage:
|
||||
result["output_tokens"] = usage["output_tokens"]
|
||||
if "total_tokens" in usage:
|
||||
result["total_tokens"] = usage["total_tokens"]
|
||||
if "cached_tokens" in usage:
|
||||
result["cached_tokens"] = usage["cached_tokens"]
|
||||
return UsageInfo(**result) if result else None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Extracting content blocks from LangChain messages
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _extract_blocks_from_chunk(msg: AIMessageChunk) -> list[tuple[int, CompatBlock]]:
|
||||
"""Extract ``(index, block)`` pairs from an ``AIMessageChunk``.
|
||||
|
||||
LangChain stores content in several places:
|
||||
- ``content: str`` — a single text block at index 0
|
||||
- ``content: list[dict]`` — explicit content blocks with their own types
|
||||
- ``tool_call_chunks`` — separate list for streamed tool call deltas
|
||||
"""
|
||||
blocks: list[tuple[int, CompatBlock]] = []
|
||||
content = msg.content
|
||||
if isinstance(content, str) and content:
|
||||
blocks.append((0, dict(TextBlock(type="text", text=content))))
|
||||
elif isinstance(content, list):
|
||||
for i, item in enumerate(content):
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
ctype = item.get("type", "")
|
||||
if ctype == "text" and item.get("text"):
|
||||
blocks.append(
|
||||
(
|
||||
item.get("index", i),
|
||||
dict(TextBlock(type="text", text=item["text"])),
|
||||
)
|
||||
)
|
||||
elif ctype in ("reasoning_content", "reasoning", "thinking"):
|
||||
reasoning_text = (
|
||||
item.get("reasoning_content")
|
||||
or item.get("reasoning")
|
||||
or item.get("thinking", "")
|
||||
)
|
||||
if reasoning_text:
|
||||
blocks.append(
|
||||
(
|
||||
item.get("index", i),
|
||||
dict(
|
||||
ReasoningBlock(
|
||||
type="reasoning", reasoning=reasoning_text
|
||||
)
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# Tool call chunks live in a separate field
|
||||
for tc in msg.tool_call_chunks or []:
|
||||
idx = tc.get("index")
|
||||
if idx is None:
|
||||
# Assign indices after text content blocks
|
||||
idx = len(blocks)
|
||||
block: CompatBlock = {"type": "tool_call_chunk", "args": tc.get("args", "")}
|
||||
if tc.get("id") is not None:
|
||||
block["id"] = tc["id"]
|
||||
if tc.get("name") is not None:
|
||||
block["name"] = tc["name"]
|
||||
blocks.append((idx, block))
|
||||
|
||||
return blocks
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# The handler
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class StreamProtocolMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
"""Callback handler that emits content-block protocol events.
|
||||
|
||||
Activated when ``__protocol_messages_stream`` is ``True`` in the run's
|
||||
configurable metadata. Emits ``StreamChunk`` tuples of the form
|
||||
``(namespace, "messages", data)`` where *data* is one of the
|
||||
``MessagesData`` event types (``message-start``, ``content-block-start``,
|
||||
etc.).
|
||||
"""
|
||||
|
||||
run_inline = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stream: Callable[[StreamChunk], None],
|
||||
subgraphs: bool,
|
||||
*,
|
||||
parent_ns: tuple[str, ...] | None = None,
|
||||
) -> None:
|
||||
self.stream = stream
|
||||
self.subgraphs = subgraphs
|
||||
self.parent_ns = parent_ns
|
||||
# Per-run metadata: run_id → (namespace, metadata_dict)
|
||||
self.metadata: dict[UUID, Meta] = {}
|
||||
# Per-run protocol state for streamed messages
|
||||
self.protocol_runs: dict[UUID, _ProtocolRunState] = {}
|
||||
# Stable message ID mapping: run_id → message_id
|
||||
self.stable_message_ids: dict[UUID, str] = {}
|
||||
# Seen message IDs for deduplication of chain-emitted messages
|
||||
self.seen: set[str | int] = set()
|
||||
|
||||
def _emit(self, meta: Meta, data: Any) -> None:
|
||||
"""Emit a protocol event as a StreamChunk.
|
||||
|
||||
The node name from *meta* is embedded at ``"__node__"`` so the
|
||||
stream pump can lift it into ``params.node`` without changing the
|
||||
``StreamChunk`` tuple shape.
|
||||
"""
|
||||
node = meta[1].get("langgraph_node")
|
||||
if node and isinstance(data, dict):
|
||||
data = {**data, "__node__": node}
|
||||
self.stream((meta[0], "messages", data))
|
||||
|
||||
# -- Chat model callbacks -----------------------------------------------
|
||||
|
||||
def on_chat_model_start(
|
||||
self,
|
||||
serialized: dict[str, Any],
|
||||
messages: list[list[BaseMessage]],
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: UUID | None = None,
|
||||
tags: list[str] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
if metadata and (not tags or (TAG_NOSTREAM not in tags)):
|
||||
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
|
||||
:-1
|
||||
]
|
||||
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
|
||||
return
|
||||
if tags:
|
||||
if filtered := [t for t in tags if not t.startswith("seq:step")]:
|
||||
metadata["tags"] = filtered
|
||||
self.metadata[run_id] = (ns, metadata)
|
||||
self.protocol_runs[run_id] = _ProtocolRunState()
|
||||
|
||||
def on_llm_new_token(
|
||||
self,
|
||||
token: str,
|
||||
*,
|
||||
chunk: ChatGenerationChunk | None = None,
|
||||
run_id: UUID,
|
||||
parent_run_id: UUID | None = None,
|
||||
tags: list[str] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
if not isinstance(chunk, ChatGenerationChunk):
|
||||
return
|
||||
meta = self.metadata.get(run_id)
|
||||
if meta is None:
|
||||
return
|
||||
state = self.protocol_runs.get(run_id)
|
||||
if state is None:
|
||||
return
|
||||
|
||||
msg = chunk.message
|
||||
if not isinstance(msg, AIMessageChunk):
|
||||
return
|
||||
|
||||
# Emit message-start on first token
|
||||
if not state.started:
|
||||
message_id = self._normalize_message_id(msg, run_id)
|
||||
state.message_id = message_id
|
||||
state.started = True
|
||||
start_data = dict(
|
||||
MessageStartData(
|
||||
event="message-start",
|
||||
role="ai",
|
||||
)
|
||||
)
|
||||
if message_id:
|
||||
start_data["message_id"] = message_id
|
||||
self._emit(meta, start_data)
|
||||
|
||||
# Extract content blocks from this chunk
|
||||
extracted = _extract_blocks_from_chunk(msg)
|
||||
for idx, delta_block in extracted:
|
||||
if idx not in state.blocks:
|
||||
# New block — emit content-block-start
|
||||
state.blocks[idx] = dict(delta_block)
|
||||
# Start block has empty content placeholder
|
||||
start_block = _make_start_block(delta_block)
|
||||
self._emit(
|
||||
meta,
|
||||
ContentBlockStartData(
|
||||
event="content-block-start",
|
||||
index=idx,
|
||||
content_block=start_block,
|
||||
),
|
||||
)
|
||||
# Then emit the first delta
|
||||
first_delta = _delta_block(
|
||||
_make_start_block(delta_block), state.blocks[idx]
|
||||
)
|
||||
if first_delta is not None:
|
||||
self._emit(
|
||||
meta,
|
||||
ContentBlockDeltaData(
|
||||
event="content-block-delta",
|
||||
index=idx,
|
||||
content_block=first_delta,
|
||||
),
|
||||
)
|
||||
else:
|
||||
# Existing block — compute delta, accumulate, emit
|
||||
previous = dict(state.blocks[idx])
|
||||
state.blocks[idx] = _accumulate_block(state.blocks[idx], delta_block)
|
||||
delta = _delta_block(previous, state.blocks[idx])
|
||||
if delta is not None:
|
||||
self._emit(
|
||||
meta,
|
||||
ContentBlockDeltaData(
|
||||
event="content-block-delta",
|
||||
index=idx,
|
||||
content_block=delta,
|
||||
),
|
||||
)
|
||||
|
||||
# Accumulate usage from chunk
|
||||
if msg.usage_metadata:
|
||||
state.usage = _accumulate_usage(state.usage, msg.usage_metadata)
|
||||
|
||||
def on_llm_end(
|
||||
self,
|
||||
response: LLMResult,
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: UUID | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
meta = self.metadata.pop(run_id, None)
|
||||
state = self.protocol_runs.pop(run_id, None)
|
||||
if meta is None or state is None:
|
||||
return
|
||||
|
||||
# Extract finish reason and usage from the final generation
|
||||
finish_reason: FinishReason = "stop"
|
||||
final_usage = state.usage
|
||||
|
||||
if response.generations and response.generations[0]:
|
||||
gen = response.generations[0][0]
|
||||
if isinstance(gen, ChatGeneration):
|
||||
final_msg = gen.message
|
||||
# Get finish reason from response_metadata
|
||||
rm = getattr(final_msg, "response_metadata", {}) or {}
|
||||
raw_reason = rm.get("finish_reason") or rm.get("stop_reason")
|
||||
if raw_reason:
|
||||
finish_reason = _normalize_finish_reason(raw_reason)
|
||||
# If we have tool calls in the final message, infer tool_use
|
||||
if (
|
||||
finish_reason == "stop"
|
||||
and hasattr(final_msg, "tool_calls")
|
||||
and final_msg.tool_calls
|
||||
):
|
||||
finish_reason = "tool_use"
|
||||
# Get usage from final message if not accumulated from chunks
|
||||
if final_usage is None and hasattr(final_msg, "usage_metadata"):
|
||||
final_usage = (
|
||||
dict(final_msg.usage_metadata)
|
||||
if final_msg.usage_metadata
|
||||
else None
|
||||
)
|
||||
|
||||
# If we never got streaming tokens (non-streamed model call),
|
||||
# emit the full message lifecycle now
|
||||
if not state.started:
|
||||
self._emit_full_message(meta, final_msg, finish_reason, final_usage)
|
||||
return
|
||||
|
||||
# Close out any open content blocks
|
||||
for idx in sorted(state.blocks):
|
||||
finalized = _finalize_block(state.blocks[idx])
|
||||
self._emit(
|
||||
meta,
|
||||
ContentBlockFinishData(
|
||||
event="content-block-finish",
|
||||
index=idx,
|
||||
content_block=finalized,
|
||||
),
|
||||
)
|
||||
|
||||
# Emit message-finish
|
||||
finish_data: dict[str, Any] = {
|
||||
"event": "message-finish",
|
||||
"reason": finish_reason,
|
||||
}
|
||||
usage_info = _to_protocol_usage(final_usage)
|
||||
if usage_info is not None:
|
||||
finish_data["usage"] = usage_info
|
||||
self._emit(meta, finish_data)
|
||||
|
||||
# Track the message as seen for dedup
|
||||
if state.message_id:
|
||||
self.seen.add(state.message_id)
|
||||
|
||||
def on_llm_error(
|
||||
self,
|
||||
error: BaseException,
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: UUID | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
meta = self.metadata.pop(run_id, None)
|
||||
state = self.protocol_runs.pop(run_id, None)
|
||||
self.stable_message_ids.pop(run_id, None)
|
||||
if meta is None or state is None:
|
||||
return
|
||||
if state.started:
|
||||
self._emit(
|
||||
meta,
|
||||
MessageErrorData(
|
||||
event="error",
|
||||
message=str(error),
|
||||
),
|
||||
)
|
||||
|
||||
# -- Chain callbacks (for node-level message dedup) ---------------------
|
||||
|
||||
def on_chain_start(
|
||||
self,
|
||||
serialized: dict[str, Any],
|
||||
inputs: dict[str, Any],
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: UUID | None = None,
|
||||
tags: list[str] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
if (
|
||||
metadata
|
||||
and kwargs.get("name") == metadata.get("langgraph_node")
|
||||
and (not tags or TAG_HIDDEN not in tags)
|
||||
):
|
||||
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
|
||||
:-1
|
||||
]
|
||||
if not self.subgraphs and len(ns) > 0:
|
||||
return
|
||||
self.metadata[run_id] = (ns, metadata)
|
||||
# Record input message IDs for deduplication
|
||||
self._record_seen_messages(inputs)
|
||||
|
||||
def on_chain_end(
|
||||
self,
|
||||
response: Any,
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: UUID | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
meta = self.metadata.pop(run_id, None)
|
||||
if meta is None:
|
||||
return
|
||||
# Emit protocol events for any new messages in the node's output
|
||||
self._emit_chain_messages(meta, response)
|
||||
|
||||
def on_chain_error(
|
||||
self,
|
||||
error: BaseException,
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: UUID | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
self.metadata.pop(run_id, None)
|
||||
|
||||
# -- Iterator taps (required by _StreamingCallbackHandler) ---------------
|
||||
|
||||
def tap_output_aiter(
|
||||
self, run_id: UUID, output: AsyncIterator[T]
|
||||
) -> AsyncIterator[T]:
|
||||
return output
|
||||
|
||||
def tap_output_iter(self, run_id: UUID, output: Iterator[T]) -> Iterator[T]:
|
||||
return output
|
||||
|
||||
# -- Internal helpers ---------------------------------------------------
|
||||
|
||||
def _normalize_message_id(self, msg: BaseMessage, run_id: UUID) -> str | None:
|
||||
"""Return a stable message ID for this run, creating one if needed."""
|
||||
msg_id = msg.id
|
||||
if msg_id is None:
|
||||
msg_id = self.stable_message_ids.get(run_id)
|
||||
if msg_id is None:
|
||||
msg_id = f"run-{run_id}"
|
||||
self.stable_message_ids[run_id] = msg_id
|
||||
# Mutate the message for consistency downstream
|
||||
if msg.id != msg_id:
|
||||
msg.id = msg_id
|
||||
return msg_id
|
||||
|
||||
def _emit_full_message(
|
||||
self,
|
||||
meta: Meta,
|
||||
msg: BaseMessage,
|
||||
finish_reason: FinishReason,
|
||||
usage: dict[str, Any] | None,
|
||||
role: str = "ai",
|
||||
) -> None:
|
||||
"""Emit a complete message lifecycle for a non-streamed model call."""
|
||||
message_id = msg.id or str(uuid4())
|
||||
if message_id in self.seen:
|
||||
return
|
||||
self.seen.add(message_id)
|
||||
|
||||
# message-start
|
||||
start_data = dict(
|
||||
MessageStartData(
|
||||
event="message-start",
|
||||
role=role,
|
||||
)
|
||||
)
|
||||
start_data["message_id"] = message_id
|
||||
self._emit(meta, start_data)
|
||||
|
||||
# Extract all blocks from the final message
|
||||
blocks = _extract_final_blocks(msg)
|
||||
for idx, block in blocks:
|
||||
# content-block-start with the full content
|
||||
self._emit(
|
||||
meta,
|
||||
ContentBlockStartData(
|
||||
event="content-block-start",
|
||||
index=idx,
|
||||
content_block=_make_start_block(block),
|
||||
),
|
||||
)
|
||||
# content-block-delta with the full content
|
||||
delta = _delta_block(_make_start_block(block), block)
|
||||
if delta is not None:
|
||||
self._emit(
|
||||
meta,
|
||||
ContentBlockDeltaData(
|
||||
event="content-block-delta",
|
||||
index=idx,
|
||||
content_block=delta,
|
||||
),
|
||||
)
|
||||
# content-block-finish
|
||||
finalized = _finalize_block(block)
|
||||
self._emit(
|
||||
meta,
|
||||
ContentBlockFinishData(
|
||||
event="content-block-finish",
|
||||
index=idx,
|
||||
content_block=finalized,
|
||||
),
|
||||
)
|
||||
|
||||
# message-finish
|
||||
finish_data: dict[str, Any] = {
|
||||
"event": "message-finish",
|
||||
"reason": finish_reason,
|
||||
}
|
||||
usage_info = _to_protocol_usage(usage)
|
||||
if usage_info is not None:
|
||||
finish_data["usage"] = usage_info
|
||||
self._emit(meta, finish_data)
|
||||
|
||||
def _record_seen_messages(self, obj: Any) -> None:
|
||||
"""Record message IDs from node inputs for deduplication."""
|
||||
if isinstance(obj, BaseMessage):
|
||||
if obj.id is not None:
|
||||
self.seen.add(obj.id)
|
||||
elif isinstance(obj, dict):
|
||||
for value in obj.values():
|
||||
self._record_seen_messages(value)
|
||||
elif isinstance(obj, Sequence) and not isinstance(obj, (str, bytes)):
|
||||
for item in obj:
|
||||
self._record_seen_messages(item)
|
||||
|
||||
def _emit_chain_messages(self, meta: Meta, response: Any) -> None:
|
||||
"""Emit protocol events for messages found in chain output."""
|
||||
from langgraph.types import Command
|
||||
|
||||
if isinstance(response, Command):
|
||||
self._emit_chain_messages(meta, response.update)
|
||||
elif isinstance(response, BaseMessage):
|
||||
self._emit_message_from_chain(meta, response)
|
||||
elif isinstance(response, Sequence) and not isinstance(response, (str, bytes)):
|
||||
for item in response:
|
||||
if isinstance(item, Command):
|
||||
self._emit_chain_messages(meta, item.update)
|
||||
elif isinstance(item, BaseMessage):
|
||||
self._emit_message_from_chain(meta, item)
|
||||
elif isinstance(response, dict):
|
||||
for value in response.values():
|
||||
if isinstance(value, BaseMessage):
|
||||
self._emit_message_from_chain(meta, value)
|
||||
elif isinstance(value, Sequence) and not isinstance(
|
||||
value, (str, bytes)
|
||||
):
|
||||
for item in value:
|
||||
if isinstance(item, BaseMessage):
|
||||
self._emit_message_from_chain(meta, item)
|
||||
|
||||
def _emit_message_from_chain(self, meta: Meta, msg: BaseMessage) -> None:
|
||||
"""Emit a full message lifecycle for a message from a chain output,
|
||||
deduplicating against previously-seen messages."""
|
||||
if msg.id is not None and msg.id in self.seen:
|
||||
return
|
||||
if msg.id is None:
|
||||
msg.id = str(uuid4())
|
||||
|
||||
# Determine role and finish reason
|
||||
role = "ai"
|
||||
if hasattr(msg, "type"):
|
||||
if msg.type == "human":
|
||||
role = "human"
|
||||
elif msg.type == "system":
|
||||
role = "system"
|
||||
|
||||
finish_reason: FinishReason = "stop"
|
||||
rm = getattr(msg, "response_metadata", {}) or {}
|
||||
raw_reason = rm.get("finish_reason") or rm.get("stop_reason")
|
||||
if raw_reason:
|
||||
finish_reason = _normalize_finish_reason(raw_reason)
|
||||
if finish_reason == "stop" and hasattr(msg, "tool_calls") and msg.tool_calls:
|
||||
finish_reason = "tool_use"
|
||||
|
||||
raw_usage = getattr(msg, "usage_metadata", None)
|
||||
usage = dict(raw_usage) if raw_usage else None
|
||||
|
||||
self._emit_full_message(meta, msg, finish_reason, usage, role=role)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Block extraction for finalized (non-streamed) messages
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _extract_final_blocks(msg: BaseMessage) -> list[tuple[int, CompatBlock]]:
|
||||
"""Extract ``(index, block)`` pairs from a finalized ``AIMessage``."""
|
||||
blocks: list[tuple[int, CompatBlock]] = []
|
||||
content = msg.content
|
||||
|
||||
if isinstance(content, str) and content:
|
||||
blocks.append((0, dict(TextBlock(type="text", text=content))))
|
||||
elif isinstance(content, list):
|
||||
for i, item in enumerate(content):
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
ctype = item.get("type", "")
|
||||
if ctype == "text" and item.get("text"):
|
||||
blocks.append((i, dict(TextBlock(type="text", text=item["text"]))))
|
||||
elif ctype in ("reasoning_content", "reasoning", "thinking"):
|
||||
reasoning_text = (
|
||||
item.get("reasoning_content")
|
||||
or item.get("reasoning")
|
||||
or item.get("thinking", "")
|
||||
)
|
||||
if reasoning_text:
|
||||
blocks.append(
|
||||
(
|
||||
i,
|
||||
dict(
|
||||
ReasoningBlock(
|
||||
type="reasoning", reasoning=reasoning_text
|
||||
)
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# Finalized tool calls (already parsed, not chunks)
|
||||
for tc in getattr(msg, "tool_calls", None) or []:
|
||||
idx = len(blocks)
|
||||
blocks.append(
|
||||
(
|
||||
idx,
|
||||
dict(
|
||||
ToolCallBlock(
|
||||
type="tool_call",
|
||||
id=tc.get("id", ""),
|
||||
name=tc.get("name", ""),
|
||||
args=tc.get("args", {}),
|
||||
)
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
return blocks
|
||||
|
||||
|
||||
def _make_start_block(block: CompatBlock) -> CompatBlock:
|
||||
"""Create an empty start placeholder for a content block."""
|
||||
btype = block.get("type", "text")
|
||||
if btype == "text":
|
||||
return TextBlock(type="text", text="")
|
||||
elif btype == "reasoning":
|
||||
return ReasoningBlock(type="reasoning", reasoning="")
|
||||
elif btype == "tool_call_chunk":
|
||||
result: CompatBlock = {"type": "tool_call_chunk", "args": ""}
|
||||
if "id" in block:
|
||||
result["id"] = block["id"]
|
||||
if "name" in block:
|
||||
result["name"] = block["name"]
|
||||
return result
|
||||
elif btype == "tool_call":
|
||||
# Already finalized — return as-is for start event
|
||||
return ToolCallBlock(
|
||||
type="tool_call",
|
||||
id=block.get("id", ""),
|
||||
name=block.get("name", ""),
|
||||
args=block.get("args", {}),
|
||||
)
|
||||
return dict(block)
|
||||
|
||||
|
||||
__all__ = ["PROTOCOL_MESSAGES_STREAM_KEY", "StreamProtocolMessagesHandler"]
|
||||
@@ -128,6 +128,10 @@ from langgraph.pregel._loop import (
|
||||
SyncPregelLoop,
|
||||
)
|
||||
from langgraph.pregel._messages import StreamMessagesHandler
|
||||
from langgraph.pregel._messages_v2 import (
|
||||
PROTOCOL_MESSAGES_STREAM_KEY,
|
||||
StreamProtocolMessagesHandler,
|
||||
)
|
||||
from langgraph.pregel._read import DEFAULT_BOUND, PregelNode
|
||||
from langgraph.pregel._retry import RetryPolicy
|
||||
from langgraph.pregel._runner import PregelRunner
|
||||
@@ -2616,8 +2620,15 @@ class Pregel(
|
||||
# set up messages stream mode
|
||||
if "messages" in stream_modes:
|
||||
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||
_msg_cls = (
|
||||
StreamProtocolMessagesHandler
|
||||
if config.get("configurable", {}).get(
|
||||
PROTOCOL_MESSAGES_STREAM_KEY, False
|
||||
)
|
||||
else StreamMessagesHandler
|
||||
)
|
||||
run_manager.inheritable_handlers.append(
|
||||
StreamMessagesHandler(
|
||||
_msg_cls(
|
||||
stream.put,
|
||||
subgraphs,
|
||||
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
|
||||
@@ -2935,7 +2946,10 @@ class Pregel(
|
||||
True
|
||||
for h in run_manager.handlers
|
||||
if isinstance(h, _StreamingCallbackHandler)
|
||||
and not isinstance(h, StreamMessagesHandler)
|
||||
and not isinstance(
|
||||
h,
|
||||
(StreamMessagesHandler, StreamProtocolMessagesHandler),
|
||||
)
|
||||
),
|
||||
False,
|
||||
)
|
||||
@@ -2972,10 +2986,16 @@ class Pregel(
|
||||
config[CONF][CONFIG_KEY_CHECKPOINT_NS] = recast_checkpoint_ns(ns)
|
||||
# set up messages stream mode
|
||||
if "messages" in stream_modes:
|
||||
# namespace can be None in a root level graph?
|
||||
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||
_msg_cls = (
|
||||
StreamProtocolMessagesHandler
|
||||
if config.get("configurable", {}).get(
|
||||
PROTOCOL_MESSAGES_STREAM_KEY, False
|
||||
)
|
||||
else StreamMessagesHandler
|
||||
)
|
||||
run_manager.inheritable_handlers.append(
|
||||
StreamMessagesHandler(
|
||||
_msg_cls(
|
||||
stream_put,
|
||||
subgraphs,
|
||||
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
"""Stream protocol types and infrastructure for LangGraph."""
|
||||
|
||||
from langgraph.stream._convert import STREAM_V2_MODES, convert_to_protocol_event
|
||||
from langgraph.stream._mux import AsyncStreamMux, StreamMux
|
||||
from langgraph.stream._types import (
|
||||
InterruptPayload,
|
||||
ProtocolEvent,
|
||||
StreamTransformer,
|
||||
)
|
||||
from langgraph.stream.chat_model_stream import AsyncChatModelStream, ChatModelStream
|
||||
from langgraph.stream.run_stream import (
|
||||
AsyncGraphRunStream,
|
||||
AsyncSubgraphRunStream,
|
||||
GraphRunStream,
|
||||
create_async_graph_run_stream,
|
||||
create_graph_run_stream,
|
||||
)
|
||||
from langgraph.stream.stream_channel import StreamChannel, is_stream_channel
|
||||
from langgraph.stream.streaming_handler import StreamingHandler
|
||||
from langgraph.stream.transformers import (
|
||||
MessagesTransformer,
|
||||
ValuesTransformer,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"STREAM_V2_MODES",
|
||||
"AsyncChatModelStream",
|
||||
"AsyncGraphRunStream",
|
||||
"AsyncStreamMux",
|
||||
"AsyncSubgraphRunStream",
|
||||
"ChatModelStream",
|
||||
"GraphRunStream",
|
||||
"InterruptPayload",
|
||||
"MessagesTransformer",
|
||||
"ProtocolEvent",
|
||||
"StreamChannel",
|
||||
"StreamMux",
|
||||
"StreamTransformer",
|
||||
"StreamingHandler",
|
||||
"ValuesTransformer",
|
||||
"convert_to_protocol_event",
|
||||
"create_async_graph_run_stream",
|
||||
"create_graph_run_stream",
|
||||
"is_stream_channel",
|
||||
]
|
||||
@@ -0,0 +1,67 @@
|
||||
"""Convert raw ``StreamChunk`` tuples to ``ProtocolEvent`` envelopes.
|
||||
|
||||
Each ``StreamMode`` is mapped to a ``ProtocolEvent`` whose ``method``
|
||||
field matches the mode name and whose ``params.data`` wraps the
|
||||
original payload.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from langgraph.stream._types import ProtocolEvent, _ProtocolEventParams
|
||||
from langgraph.types import StreamMode
|
||||
|
||||
#: All stream modes requested by ``StreamingHandler`` when calling the
|
||||
#: underlying ``stream()`` / ``astream()``.
|
||||
STREAM_V2_MODES: list[StreamMode] = [
|
||||
"values",
|
||||
"updates",
|
||||
"messages",
|
||||
"custom",
|
||||
"checkpoints",
|
||||
"tasks",
|
||||
"debug",
|
||||
]
|
||||
|
||||
_SUPPORTED_MODES: set[str] = set(STREAM_V2_MODES)
|
||||
|
||||
|
||||
def convert_to_protocol_event(
|
||||
ns: tuple[str, ...],
|
||||
mode: str,
|
||||
payload: Any,
|
||||
*,
|
||||
node: str | None = None,
|
||||
) -> ProtocolEvent | None:
|
||||
"""Convert a ``StreamChunk`` to a ``ProtocolEvent``.
|
||||
|
||||
Returns ``None`` for unsupported or unknown modes.
|
||||
|
||||
The ``seq`` field is left as ``0`` here; the :class:`StreamMux` is
|
||||
the sole seq assigner and overwrites it inside ``push()``.
|
||||
|
||||
Args:
|
||||
ns: Namespace tuple from the ``StreamChunk``.
|
||||
mode: Stream mode string (``"values"``, ``"updates"``, etc.).
|
||||
payload: The raw payload from the stream.
|
||||
node: Optional node name for provenance.
|
||||
"""
|
||||
if mode not in _SUPPORTED_MODES:
|
||||
return None
|
||||
|
||||
params: _ProtocolEventParams = {
|
||||
"namespace": list(ns),
|
||||
"data": payload,
|
||||
}
|
||||
if node is not None:
|
||||
params["node"] = node
|
||||
|
||||
return ProtocolEvent(
|
||||
type="event",
|
||||
method=mode,
|
||||
params=params,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["STREAM_V2_MODES", "convert_to_protocol_event"]
|
||||
@@ -0,0 +1,385 @@
|
||||
"""Central event dispatcher with transformer pipeline for StreamingHandler.
|
||||
|
||||
``StreamMux`` is the synchronous core: it holds the main
|
||||
event log (a plain list), tracks discovered namespaces for subgraph stream
|
||||
creation, and pipes every event through the registered
|
||||
:class:`StreamTransformer` pipeline before appending it to the log.
|
||||
|
||||
``AsyncStreamMux`` extends ``StreamMux`` with async consumer APIs
|
||||
(output futures, async event subscriptions, subgraph discovery).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
from langgraph.stream._types import InterruptPayload, ProtocolEvent, StreamTransformer
|
||||
from langgraph.stream.stream_channel import StreamChannel, is_stream_channel
|
||||
|
||||
|
||||
class StreamMux:
|
||||
"""Synchronous event dispatcher for the StreamingHandler infrastructure.
|
||||
|
||||
The mux owns the main event log, applies the transformer pipeline to
|
||||
every incoming event, and tracks namespace discovery and latest values.
|
||||
|
||||
For async consumer APIs (output futures, async event subscriptions,
|
||||
subgraph discovery) use :class:`AsyncStreamMux`.
|
||||
"""
|
||||
|
||||
def __init__(self, transformers: list[StreamTransformer] | None = None) -> None:
|
||||
self._event_log: list[ProtocolEvent] = []
|
||||
self._transformers: list[StreamTransformer] = list(transformers or [])
|
||||
self._current_namespace: list[str] = []
|
||||
self._next_emit_seq: int = 0
|
||||
|
||||
# Namespace discovery: maps top-level ns segment → True
|
||||
self._discovered_ns: dict[str, bool] = {}
|
||||
|
||||
# Latest values per namespace (list-of-strings key)
|
||||
self._latest_values: dict[str, Any] = {}
|
||||
|
||||
# Interrupt tracking
|
||||
self._interrupts: list[InterruptPayload] = []
|
||||
self._interrupted = False
|
||||
|
||||
# Closed state
|
||||
self._closed = False
|
||||
self._error: BaseException | None = None
|
||||
|
||||
# -- Producer API -------------------------------------------------------
|
||||
|
||||
def push(self, event: ProtocolEvent) -> None:
|
||||
"""Push an event through the transformer pipeline and into the log.
|
||||
|
||||
Each registered transformer's ``process()`` is called in order.
|
||||
If any transformer returns ``False``, the event is suppressed
|
||||
(not appended to the main log).
|
||||
"""
|
||||
if self._closed:
|
||||
return
|
||||
|
||||
# Mux is the sole seq assigner — ensures all events in the log
|
||||
# (including those from StreamChannel forwarders) share a single
|
||||
# monotonically increasing counter.
|
||||
event["seq"] = self._next_emit_seq
|
||||
self._next_emit_seq += 1
|
||||
|
||||
# Track namespace
|
||||
ns = event["params"].get("namespace", [])
|
||||
if ns:
|
||||
top_segment = ns[0]
|
||||
if top_segment not in self._discovered_ns:
|
||||
self._discovered_ns[top_segment] = True
|
||||
|
||||
# Track values
|
||||
if event["method"] == "values":
|
||||
ns_key = _ns_key(ns)
|
||||
self._latest_values[ns_key] = event["params"]["data"]
|
||||
|
||||
# Track interrupts from values events
|
||||
if event["method"] == "values":
|
||||
data = event["params"]["data"]
|
||||
if isinstance(data, dict) and "__interrupt__" in data:
|
||||
interrupt_info = data["__interrupt__"]
|
||||
if isinstance(interrupt_info, (list, tuple)):
|
||||
for item in interrupt_info:
|
||||
iid = getattr(item, "id", None) or str(id(item))
|
||||
self._interrupts.append(
|
||||
InterruptPayload(
|
||||
interrupt_id=iid,
|
||||
payload=item,
|
||||
)
|
||||
)
|
||||
self._interrupted = True
|
||||
|
||||
# Run transformer pipeline
|
||||
self._current_namespace = ns
|
||||
keep = True
|
||||
for transformer in self._transformers:
|
||||
result = transformer.process(event)
|
||||
if result is False:
|
||||
keep = False
|
||||
self._current_namespace = []
|
||||
|
||||
# Append to main log if not suppressed
|
||||
if keep:
|
||||
self._event_log.append(event)
|
||||
|
||||
def close(self, output: Any = None) -> None:
|
||||
"""Close the mux and finalize all transformers."""
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
|
||||
for transformer in self._transformers:
|
||||
transformer.finalize()
|
||||
|
||||
def fail(self, error: BaseException) -> None:
|
||||
"""Fail the mux and propagate the error to all consumers."""
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
self._error = error
|
||||
|
||||
for transformer in self._transformers:
|
||||
transformer.fail(error)
|
||||
|
||||
# -- Inspection ---------------------------------------------------------
|
||||
|
||||
@property
|
||||
def interrupted(self) -> bool:
|
||||
return self._interrupted
|
||||
|
||||
@property
|
||||
def interrupts(self) -> list[InterruptPayload]:
|
||||
return list(self._interrupts)
|
||||
|
||||
@property
|
||||
def event_log(self) -> list[ProtocolEvent]:
|
||||
return self._event_log
|
||||
|
||||
def get_latest_values(self, ns: list[str] | None = None) -> Any:
|
||||
"""Return the most recent values for a namespace."""
|
||||
return self._latest_values.get(_ns_key(ns or []))
|
||||
|
||||
# -- Internal -----------------------------------------------------------
|
||||
|
||||
def register_transformer(self, transformer: StreamTransformer) -> None:
|
||||
"""Register a new transformer and replay all buffered events through it.
|
||||
|
||||
This is the safe way to add a late-arriving transformer after the mux
|
||||
has already started processing events. The sequence is:
|
||||
|
||||
1. Snapshot the current log length.
|
||||
2. Append the transformer so future ``push()`` calls reach it.
|
||||
3. Replay events ``[0, snapshot)`` through the transformer.
|
||||
4. If the mux is already closed, call ``finalize()`` immediately so
|
||||
the transformer's log/channel terminates cleanly.
|
||||
|
||||
No namespace filtering is applied — all buffered events are
|
||||
replayed. Transformers that need namespace filtering should do
|
||||
so inside their ``process()`` implementation.
|
||||
"""
|
||||
snapshot = len(self._event_log)
|
||||
self._transformers.append(transformer)
|
||||
for i in range(snapshot):
|
||||
transformer.process(self._event_log[i])
|
||||
if self._closed:
|
||||
transformer.finalize()
|
||||
|
||||
def wire_channels(self, projection: Any) -> None:
|
||||
"""Scan *projection* for :class:`StreamChannel` instances and wire them.
|
||||
|
||||
For each ``StreamChannel`` found, registers a push callback that
|
||||
appends a :class:`ProtocolEvent` directly to the main event log
|
||||
with ``method`` set to the channel's name.
|
||||
|
||||
Channel events bypass the transformer pipeline (matching the JS
|
||||
implementation). They are visible to raw event iteration and
|
||||
remote SDK clients but not to other transformers' ``process()``.
|
||||
"""
|
||||
if projection is None:
|
||||
return
|
||||
items: dict[str, Any] = {}
|
||||
if isinstance(projection, dict):
|
||||
items = projection
|
||||
elif hasattr(projection, "__dict__"):
|
||||
items = vars(projection)
|
||||
for _key, value in items.items():
|
||||
if is_stream_channel(value):
|
||||
channel: StreamChannel[Any] = value
|
||||
def _make_forwarder(ch: StreamChannel[Any]) -> Any:
|
||||
def _forward(item: Any) -> None:
|
||||
if self._closed:
|
||||
return
|
||||
# Append directly to the event log, bypassing
|
||||
# the transformer pipeline. This matches the JS
|
||||
# implementation and avoids re-entrancy bugs
|
||||
# (namespace clobbering, infinite recursion).
|
||||
self._event_log.append(
|
||||
ProtocolEvent(
|
||||
type="event",
|
||||
seq=self._next_emit_seq,
|
||||
method=ch.channel_name,
|
||||
params={
|
||||
"namespace": list(self._current_namespace),
|
||||
"data": item,
|
||||
},
|
||||
)
|
||||
)
|
||||
self._next_emit_seq += 1
|
||||
|
||||
return _forward
|
||||
|
||||
channel._wire(_make_forwarder(channel))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# AsyncStreamMux — async consumer APIs on top of the sync core
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class AsyncStreamMux(StreamMux):
|
||||
"""Async extension of :class:`StreamMux`.
|
||||
|
||||
Adds output futures, async event subscriptions, and subgraph
|
||||
discovery on top of the synchronous producer/transformer core.
|
||||
"""
|
||||
|
||||
def __init__(self, transformers: list[StreamTransformer] | None = None) -> None:
|
||||
super().__init__(transformers=transformers)
|
||||
# Notification event — set on every push/close/fail to wake async consumers
|
||||
self._notify: asyncio.Event = asyncio.Event()
|
||||
# Waiters for new namespace discovery
|
||||
self._ns_waiters: list[asyncio.Future[None]] = []
|
||||
# Output promise tracking
|
||||
self._output_futures: dict[str, asyncio.Future[Any]] = {}
|
||||
|
||||
# -- Producer overrides (extend to resolve async primitives) -------------
|
||||
|
||||
def push(self, event: ProtocolEvent) -> None:
|
||||
# Peek at namespace before super().push() so we can detect new
|
||||
# discoveries and wake waiters.
|
||||
ns = event["params"].get("namespace", [])
|
||||
is_new_ns = bool(ns) and ns[0] not in self._discovered_ns
|
||||
super().push(event)
|
||||
if is_new_ns and ns[0] in self._discovered_ns:
|
||||
self._wake_ns_waiters()
|
||||
self._notify.set()
|
||||
|
||||
def close(self, output: Any = None) -> None:
|
||||
super().close(output)
|
||||
self._notify.set()
|
||||
# Resolve output futures
|
||||
for ns_key, fut in self._output_futures.items():
|
||||
if not fut.done():
|
||||
value = self._latest_values.get(ns_key)
|
||||
try:
|
||||
fut.get_loop().call_soon_threadsafe(fut.set_result, value)
|
||||
except RuntimeError:
|
||||
pass
|
||||
# Wake namespace waiters
|
||||
self._wake_ns_waiters()
|
||||
|
||||
def fail(self, error: BaseException) -> None:
|
||||
super().fail(error)
|
||||
self._notify.set()
|
||||
# Reject output futures
|
||||
for fut in self._output_futures.values():
|
||||
if not fut.done():
|
||||
try:
|
||||
fut.get_loop().call_soon_threadsafe(fut.set_exception, error)
|
||||
except RuntimeError:
|
||||
pass
|
||||
# Wake namespace waiters
|
||||
self._wake_ns_waiters()
|
||||
|
||||
# -- Async consumer API -------------------------------------------------
|
||||
|
||||
async def subscribe_events(
|
||||
self, path: list[str] | None = None, offset: int = 0
|
||||
) -> AsyncIterator[ProtocolEvent]:
|
||||
"""Async iterate over events matching *path*.
|
||||
|
||||
If *path* is ``None`` or empty, all events are yielded.
|
||||
Otherwise, only events whose namespace starts with *path*
|
||||
are yielded.
|
||||
|
||||
Uses the list + ``asyncio.Event`` notification pattern: poll
|
||||
the event log, yield what's new, await the notify event for more.
|
||||
"""
|
||||
cursor = offset
|
||||
while True:
|
||||
while cursor < len(self._event_log):
|
||||
event = self._event_log[cursor]
|
||||
cursor += 1
|
||||
if not path or _ns_starts_with(
|
||||
event["params"].get("namespace", []), path
|
||||
):
|
||||
yield event
|
||||
if self._closed:
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
return
|
||||
self._notify.clear()
|
||||
await self._notify.wait()
|
||||
|
||||
async def subscribe_subgraphs(
|
||||
self, path: list[str] | None = None, offset: int = 0
|
||||
) -> AsyncIterator[str]:
|
||||
"""Yield top-level namespace segments as they are discovered.
|
||||
|
||||
Each yielded value is the first namespace segment of a newly
|
||||
discovered subgraph (e.g. ``"agent:0"``).
|
||||
"""
|
||||
yielded: set[str] = set()
|
||||
while True:
|
||||
# Yield any newly discovered namespaces
|
||||
for ns_segment in list(self._discovered_ns):
|
||||
if ns_segment not in yielded:
|
||||
# Filter by path prefix if specified
|
||||
if path:
|
||||
if not ns_segment.startswith(path[0]):
|
||||
continue
|
||||
yielded.add(ns_segment)
|
||||
yield ns_segment
|
||||
|
||||
if self._closed:
|
||||
return
|
||||
|
||||
# Wait for new namespaces
|
||||
loop = asyncio.get_running_loop()
|
||||
fut: asyncio.Future[None] = loop.create_future()
|
||||
self._ns_waiters.append(fut)
|
||||
await fut
|
||||
|
||||
def get_output_future(self, ns: list[str] | None = None) -> asyncio.Future[Any]:
|
||||
"""Get or create an output future for a namespace.
|
||||
|
||||
The future resolves to the latest ``values`` event data when
|
||||
the mux is closed.
|
||||
"""
|
||||
ns_key = _ns_key(ns or [])
|
||||
if ns_key not in self._output_futures:
|
||||
loop = asyncio.get_running_loop()
|
||||
self._output_futures[ns_key] = loop.create_future()
|
||||
|
||||
# If already closed, resolve immediately
|
||||
if self._closed:
|
||||
value = self._latest_values.get(ns_key)
|
||||
if self._error is not None:
|
||||
self._output_futures[ns_key].set_exception(self._error)
|
||||
else:
|
||||
self._output_futures[ns_key].set_result(value)
|
||||
|
||||
return self._output_futures[ns_key]
|
||||
|
||||
# -- Internal -----------------------------------------------------------
|
||||
|
||||
def _wake_ns_waiters(self) -> None:
|
||||
for fut in self._ns_waiters:
|
||||
if not fut.done():
|
||||
try:
|
||||
fut.get_loop().call_soon_threadsafe(fut.set_result, None)
|
||||
except RuntimeError:
|
||||
pass
|
||||
self._ns_waiters.clear()
|
||||
|
||||
|
||||
|
||||
def _ns_key(ns: list[str] | tuple[str, ...]) -> str:
|
||||
"""Convert a namespace list to a hashable key."""
|
||||
return "|".join(ns)
|
||||
|
||||
|
||||
def _ns_starts_with(ns: list[str], prefix: list[str]) -> bool:
|
||||
"""Check if *ns* starts with *prefix*."""
|
||||
if len(ns) < len(prefix):
|
||||
return False
|
||||
return ns[: len(prefix)] == prefix
|
||||
|
||||
|
||||
__all__ = ["AsyncStreamMux", "StreamMux"]
|
||||
@@ -0,0 +1,166 @@
|
||||
"""Protocol types for StreamingHandler.
|
||||
|
||||
Re-exports CDDL-derived types from ``langchain-protocol`` and defines
|
||||
in-process-only types needed by the LangGraph streaming infrastructure.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Re-exports from langchain-protocol (CDDL-derived)
|
||||
# ---------------------------------------------------------------------------
|
||||
# Primitives
|
||||
# Content blocks
|
||||
# Messages data
|
||||
# Tools data
|
||||
from langchain_protocol import (
|
||||
Annotation,
|
||||
Citation,
|
||||
ContentBlock,
|
||||
ContentBlockDeltaData,
|
||||
ContentBlockFinishData,
|
||||
ContentBlockStartData,
|
||||
FinalizedContentBlock,
|
||||
FinishReason,
|
||||
InvalidToolCallBlock,
|
||||
MessageErrorData,
|
||||
MessageFinishData,
|
||||
MessageMetadata,
|
||||
MessageRole,
|
||||
MessagesData,
|
||||
MessageStartData,
|
||||
MetadataScalar,
|
||||
Namespace,
|
||||
ReasoningBlock,
|
||||
TextBlock,
|
||||
ToolCallBlock,
|
||||
ToolCallChunkBlock,
|
||||
ToolErrorData,
|
||||
ToolFinishedData,
|
||||
ToolOutputDeltaData,
|
||||
ToolsData,
|
||||
ToolStartedData,
|
||||
UsageInfo,
|
||||
)
|
||||
from typing_extensions import NotRequired, TypedDict
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# In-process types (not in the CDDL spec)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _ProtocolEventParams(TypedDict):
|
||||
"""Payload envelope for a :class:`ProtocolEvent`."""
|
||||
|
||||
namespace: Namespace
|
||||
node: NotRequired[str]
|
||||
data: Any
|
||||
|
||||
|
||||
class ProtocolEvent(TypedDict):
|
||||
"""A single protocol event emitted by the StreamingHandler infrastructure.
|
||||
|
||||
``method`` corresponds to a
|
||||
:pydata:`~langgraph.types.StreamMode` value (``"messages"``,
|
||||
``"updates"``, etc.).
|
||||
"""
|
||||
|
||||
type: str # always "event"
|
||||
seq: NotRequired[int] # assigned by StreamMux.push(); absent before push()
|
||||
method: str # StreamMode value
|
||||
params: _ProtocolEventParams
|
||||
|
||||
|
||||
class StreamTransformer(ABC):
|
||||
"""Extension point for custom stream projections.
|
||||
|
||||
Implementations are registered with ``StreamingHandler`` and receive every
|
||||
:class:`ProtocolEvent` before it is appended to the event log.
|
||||
|
||||
Any :class:`~langgraph.stream.stream_channel.StreamChannel` instances
|
||||
returned by ``init()`` are automatically wired to the protocol event
|
||||
stream by the mux.
|
||||
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def init(self) -> Any:
|
||||
"""Return the initial projection value.
|
||||
|
||||
Called once before the run. Any
|
||||
:class:`~langgraph.stream.stream_channel.StreamChannel` instances
|
||||
in the return value are automatically wired by the mux.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
"""Process an event.
|
||||
|
||||
Return ``True`` to keep the event in the log, ``False`` to suppress
|
||||
it.
|
||||
"""
|
||||
...
|
||||
|
||||
def finalize(self) -> None:
|
||||
"""Called once when the run completes successfully.
|
||||
|
||||
Optional — the mux auto-closes any :class:`StreamChannel` instances,
|
||||
so transformers that only use channels can omit this.
|
||||
"""
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
"""Called once when the run fails.
|
||||
|
||||
Optional — the mux auto-fails any :class:`StreamChannel` instances,
|
||||
so transformers that only use channels can omit this.
|
||||
"""
|
||||
|
||||
|
||||
class InterruptPayload(TypedDict):
|
||||
"""An interrupt produced during a StreamingHandler run."""
|
||||
|
||||
interrupt_id: str
|
||||
payload: Any
|
||||
|
||||
|
||||
__all__ = [
|
||||
# Primitives (re-exported)
|
||||
"Namespace",
|
||||
"MessageRole",
|
||||
"MessageMetadata",
|
||||
"MetadataScalar",
|
||||
# Content blocks (re-exported)
|
||||
"TextBlock",
|
||||
"ReasoningBlock",
|
||||
"ToolCallBlock",
|
||||
"ToolCallChunkBlock",
|
||||
"InvalidToolCallBlock",
|
||||
"ContentBlock",
|
||||
"FinalizedContentBlock",
|
||||
"Annotation",
|
||||
"Citation",
|
||||
# Messages data (re-exported)
|
||||
"MessagesData",
|
||||
"MessageStartData",
|
||||
"ContentBlockStartData",
|
||||
"ContentBlockDeltaData",
|
||||
"ContentBlockFinishData",
|
||||
"MessageFinishData",
|
||||
"MessageErrorData",
|
||||
"FinishReason",
|
||||
"UsageInfo",
|
||||
# Tools data (re-exported)
|
||||
"ToolsData",
|
||||
"ToolStartedData",
|
||||
"ToolOutputDeltaData",
|
||||
"ToolFinishedData",
|
||||
"ToolErrorData",
|
||||
# In-process types
|
||||
"ProtocolEvent",
|
||||
"StreamTransformer",
|
||||
"InterruptPayload",
|
||||
]
|
||||
@@ -0,0 +1,324 @@
|
||||
"""Per-message streaming objects for StreamingHandler.
|
||||
|
||||
``ChatModelStream`` is the synchronous variant returned by
|
||||
``GraphRunStream.messages``. Properties (``.text``, ``.reasoning``,
|
||||
``.usage``) return final accumulated values.
|
||||
|
||||
``AsyncChatModelStream`` is the asynchronous variant returned by
|
||||
``AsyncGraphRunStream.messages``. Projections are dual
|
||||
async-iterable + awaitable (e.g. ``async for delta in msg.text``
|
||||
or ``full = await msg.text``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Generator
|
||||
from typing import Any
|
||||
|
||||
from langgraph.stream._types import UsageInfo
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sync variant
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ChatModelStream:
|
||||
"""Synchronous per-message object for a single LLM response.
|
||||
|
||||
Created by :class:`~langgraph.stream.transformers.MessagesTransformer`
|
||||
and yielded by ``GraphRunStream.messages``. By the time the sync
|
||||
iterator yields a ``ChatModelStream``, the message lifecycle is
|
||||
complete and all properties contain their final values.
|
||||
|
||||
Projections:
|
||||
|
||||
- ``.text`` — accumulated text content (``str``)
|
||||
- ``.reasoning`` — accumulated reasoning content (``str``)
|
||||
- ``.usage`` — :class:`UsageInfo` or ``None``
|
||||
- ``.namespace`` / ``.node`` — provenance metadata
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
namespace: list[str] | None = None,
|
||||
node: str | None = None,
|
||||
message_id: str | None = None,
|
||||
) -> None:
|
||||
self._namespace = namespace or []
|
||||
self._node = node
|
||||
self._message_id = message_id
|
||||
|
||||
# Accumulated state
|
||||
self._text_acc = ""
|
||||
self._reasoning_acc = ""
|
||||
self._usage_value: UsageInfo | None = None
|
||||
self._done = False
|
||||
|
||||
# -- Public projections ------------------------------------------------
|
||||
|
||||
@property
|
||||
def text(self) -> str:
|
||||
"""Accumulated text content."""
|
||||
return self._text_acc
|
||||
|
||||
@property
|
||||
def reasoning(self) -> str:
|
||||
"""Accumulated reasoning content."""
|
||||
return self._reasoning_acc
|
||||
|
||||
@property
|
||||
def usage(self) -> UsageInfo | None:
|
||||
"""Usage info, available after the message finishes."""
|
||||
return self._usage_value
|
||||
|
||||
@property
|
||||
def namespace(self) -> list[str]:
|
||||
return self._namespace
|
||||
|
||||
@property
|
||||
def node(self) -> str | None:
|
||||
return self._node
|
||||
|
||||
@property
|
||||
def message_id(self) -> str | None:
|
||||
return self._message_id
|
||||
|
||||
@property
|
||||
def done(self) -> bool:
|
||||
return self._done
|
||||
|
||||
# -- Internal API (called by MessagesTransformer) ----------------------
|
||||
|
||||
def _push_content_block_delta(self, data: dict[str, Any]) -> None:
|
||||
"""Process a ``content-block-delta`` event."""
|
||||
block = data.get("content_block", {})
|
||||
btype = block.get("type", "")
|
||||
|
||||
if btype == "text":
|
||||
delta_text = block.get("text", "")
|
||||
if delta_text:
|
||||
self._text_acc += delta_text
|
||||
elif btype == "reasoning":
|
||||
delta_r = block.get("reasoning", "")
|
||||
if delta_r:
|
||||
self._reasoning_acc += delta_r
|
||||
|
||||
def _push_content_block_finish(self, data: dict[str, Any]) -> None:
|
||||
"""Process a ``content-block-finish`` event."""
|
||||
block = data.get("content_block", {})
|
||||
btype = block.get("type", "")
|
||||
|
||||
if btype == "text":
|
||||
full_text = block.get("text", "")
|
||||
if full_text and full_text != self._text_acc:
|
||||
self._text_acc = full_text
|
||||
elif btype == "reasoning":
|
||||
full_r = block.get("reasoning", "")
|
||||
if full_r and full_r != self._reasoning_acc:
|
||||
self._reasoning_acc = full_r
|
||||
|
||||
def _finish(self, data: dict[str, Any]) -> None:
|
||||
"""Process a ``message-finish`` event."""
|
||||
self._done = True
|
||||
self._usage_value = data.get("usage")
|
||||
|
||||
def _fail(self, error: BaseException) -> None:
|
||||
"""Process a ``message-error`` event."""
|
||||
self._done = True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dual-projection helpers — sync data container + async notification layer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _DualProjection:
|
||||
"""Sync data container for incremental deltas and a final value.
|
||||
|
||||
Stores deltas as they arrive and tracks the final accumulated value.
|
||||
No async primitives — see :class:`_AsyncDualProjection` for the
|
||||
async-iterable + awaitable extension.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._deltas: list[Any] = []
|
||||
self._done = False
|
||||
self._error: BaseException | None = None
|
||||
self._final_value: Any = None
|
||||
self._final_set = False
|
||||
|
||||
# -- Producer API (called by AsyncChatModelStream) ---------------------
|
||||
|
||||
def _push(self, delta: Any) -> None:
|
||||
"""Add a new delta value."""
|
||||
self._deltas.append(delta)
|
||||
|
||||
def _finish(self, accumulated: Any) -> None:
|
||||
"""Set the final accumulated value and mark as done."""
|
||||
self._final_value = accumulated
|
||||
self._final_set = True
|
||||
self._done = True
|
||||
|
||||
def _fail(self, error: BaseException) -> None:
|
||||
self._error = error
|
||||
self._done = True
|
||||
|
||||
|
||||
class _AsyncDualProjection(_DualProjection):
|
||||
"""Async extension of :class:`_DualProjection`.
|
||||
|
||||
Async iterable of deltas that is also awaitable for the final value.
|
||||
Uses an ``asyncio.Event`` to notify async consumers when new data
|
||||
arrives — the same pattern as :class:`AsyncStreamMux`.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self._notify: asyncio.Event = asyncio.Event()
|
||||
|
||||
# -- Producer overrides (extend to notify) -----------------------------
|
||||
|
||||
def _push(self, delta: Any) -> None:
|
||||
super()._push(delta)
|
||||
self._notify.set()
|
||||
|
||||
def _finish(self, accumulated: Any) -> None:
|
||||
super()._finish(accumulated)
|
||||
self._notify.set()
|
||||
|
||||
def _fail(self, error: BaseException) -> None:
|
||||
super()._fail(error)
|
||||
self._notify.set()
|
||||
|
||||
# -- Async iterable (yields deltas) ------------------------------------
|
||||
|
||||
def __aiter__(self) -> _AsyncDualProjectionIterator:
|
||||
return _AsyncDualProjectionIterator(self)
|
||||
|
||||
# -- Awaitable (returns final value) -----------------------------------
|
||||
|
||||
def __await__(self) -> Generator[Any, None, Any]:
|
||||
return self._await_impl().__await__()
|
||||
|
||||
async def _await_impl(self) -> Any:
|
||||
while not self._final_set:
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
self._notify.clear()
|
||||
await self._notify.wait()
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
return self._final_value
|
||||
|
||||
|
||||
class _AsyncDualProjectionIterator:
|
||||
"""Async iterator over an :class:`_AsyncDualProjection`'s deltas."""
|
||||
|
||||
__slots__ = ("_proj", "_offset")
|
||||
|
||||
def __init__(self, proj: _AsyncDualProjection) -> None:
|
||||
self._proj = proj
|
||||
self._offset = 0
|
||||
|
||||
def __aiter__(self) -> _AsyncDualProjectionIterator:
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> Any:
|
||||
while True:
|
||||
if self._offset < len(self._proj._deltas):
|
||||
item = self._proj._deltas[self._offset]
|
||||
self._offset += 1
|
||||
return item
|
||||
if self._proj._error is not None:
|
||||
raise self._proj._error
|
||||
if self._proj._done:
|
||||
raise StopAsyncIteration
|
||||
self._proj._notify.clear()
|
||||
await self._proj._notify.wait()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Async variant
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class AsyncChatModelStream(ChatModelStream):
|
||||
"""Asynchronous per-message streaming object for a single LLM response.
|
||||
|
||||
Created by :class:`~langgraph.stream.transformers.MessagesTransformer`
|
||||
and yielded by ``AsyncGraphRunStream.messages``. Content-block events
|
||||
are fed into this object until ``message-finish``.
|
||||
|
||||
Projections:
|
||||
|
||||
- ``.text`` — async iterable of text deltas; awaitable for full text
|
||||
- ``.reasoning`` — async iterable of reasoning deltas; awaitable for
|
||||
full reasoning text
|
||||
- ``.usage`` — awaitable for :class:`UsageInfo`
|
||||
- ``.namespace`` / ``.node`` — provenance metadata
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
namespace: list[str] | None = None,
|
||||
node: str | None = None,
|
||||
message_id: str | None = None,
|
||||
) -> None:
|
||||
super().__init__(namespace=namespace, node=node, message_id=message_id)
|
||||
self._text_proj = _AsyncDualProjection()
|
||||
self._reasoning_proj = _AsyncDualProjection()
|
||||
self._usage_proj = _AsyncDualProjection()
|
||||
|
||||
# -- Public projections (override sync properties) ---------------------
|
||||
|
||||
@property
|
||||
def text(self) -> _AsyncDualProjection:
|
||||
"""Text content — async iterable of deltas, awaitable for full text."""
|
||||
return self._text_proj
|
||||
|
||||
@property
|
||||
def reasoning(self) -> _AsyncDualProjection:
|
||||
"""Reasoning content — async iterable of deltas, awaitable for full text."""
|
||||
return self._reasoning_proj
|
||||
|
||||
@property
|
||||
def usage(self) -> _AsyncDualProjection:
|
||||
"""Usage info — awaitable for :class:`UsageInfo`."""
|
||||
return self._usage_proj
|
||||
|
||||
# -- Internal API (extend base to also drive projections) --------------
|
||||
|
||||
def _push_content_block_delta(self, data: dict[str, Any]) -> None:
|
||||
"""Process a ``content-block-delta`` event."""
|
||||
super()._push_content_block_delta(data)
|
||||
block = data.get("content_block", {})
|
||||
btype = block.get("type", "")
|
||||
|
||||
if btype == "text":
|
||||
delta_text = block.get("text", "")
|
||||
if delta_text:
|
||||
self._text_proj._push(delta_text)
|
||||
elif btype == "reasoning":
|
||||
delta_r = block.get("reasoning", "")
|
||||
if delta_r:
|
||||
self._reasoning_proj._push(delta_r)
|
||||
|
||||
def _finish(self, data: dict[str, Any]) -> None:
|
||||
"""Process a ``message-finish`` event."""
|
||||
super()._finish(data)
|
||||
self._text_proj._finish(self._text_acc)
|
||||
self._reasoning_proj._finish(self._reasoning_acc)
|
||||
self._usage_proj._finish(self._usage_value)
|
||||
|
||||
def _fail(self, error: BaseException) -> None:
|
||||
"""Process a ``message-error`` event."""
|
||||
super()._fail(error)
|
||||
self._text_proj._fail(error)
|
||||
self._reasoning_proj._fail(error)
|
||||
self._usage_proj._fail(error)
|
||||
|
||||
|
||||
__all__ = ["AsyncChatModelStream", "ChatModelStream"]
|
||||
@@ -0,0 +1,586 @@
|
||||
"""GraphRunStream and AsyncGraphRunStream for StreamingHandler.
|
||||
|
||||
These are the top-level objects returned by
|
||||
``StreamingHandler.stream()`` / ``StreamingHandler.astream()``.
|
||||
|
||||
``AsyncGraphRunStream`` wraps an :class:`AsyncStreamMux` and exposes
|
||||
``.values``, ``.messages``, ``.subgraphs``, ``.output``, and
|
||||
``.messages_from()``.
|
||||
|
||||
``GraphRunStream`` wraps a :class:`StreamMux` and exposes the sync
|
||||
equivalents: ``.values``, ``.messages``, and ``.output``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator, Callable, Iterator
|
||||
from typing import Any
|
||||
|
||||
from langgraph.stream._convert import convert_to_protocol_event
|
||||
from langgraph.stream._mux import AsyncStreamMux, StreamMux
|
||||
from langgraph.stream._types import InterruptPayload, ProtocolEvent, StreamTransformer
|
||||
from langgraph.stream.chat_model_stream import AsyncChatModelStream, ChatModelStream
|
||||
from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Values projection — dual async-iterable + awaitable
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _ValuesProjection:
|
||||
"""Async iterable of intermediate state snapshots; awaitable for final."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mux: AsyncStreamMux,
|
||||
values_transformer: ValuesTransformer,
|
||||
ns: list[str],
|
||||
mapper: Callable[[Any], Any] | None = None,
|
||||
) -> None:
|
||||
self._mux = mux
|
||||
self._values_transformer = values_transformer
|
||||
self._ns = ns
|
||||
self._mapper = mapper
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[Any]:
|
||||
log = self._values_transformer.values_log
|
||||
cursor = 0
|
||||
while True:
|
||||
while cursor < len(log):
|
||||
item = log[cursor]
|
||||
cursor += 1
|
||||
if item.get("namespace", []) == self._ns:
|
||||
data = item["data"]
|
||||
if data is not None and self._mapper is not None:
|
||||
yield self._mapper(data)
|
||||
else:
|
||||
yield data
|
||||
if self._mux._closed:
|
||||
return
|
||||
self._mux._notify.clear()
|
||||
await self._mux._notify.wait()
|
||||
|
||||
def __await__(self) -> Any:
|
||||
return self._await_impl().__await__()
|
||||
|
||||
async def _await_impl(self) -> Any:
|
||||
value = await self._mux.get_output_future(self._ns)
|
||||
if value is not None and self._mapper is not None:
|
||||
return self._mapper(value)
|
||||
return value
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Messages projection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _MessagesProjection:
|
||||
"""Async iterable of :class:`AsyncChatModelStream` instances."""
|
||||
|
||||
def __init__(self, mux: AsyncStreamMux, messages_transformer: MessagesTransformer) -> None:
|
||||
self._mux = mux
|
||||
self._transformer = messages_transformer
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[AsyncChatModelStream]:
|
||||
log = self._transformer.messages_log
|
||||
cursor = 0
|
||||
while True:
|
||||
while cursor < len(log):
|
||||
yield log[cursor]
|
||||
cursor += 1
|
||||
if self._mux._closed:
|
||||
return
|
||||
self._mux._notify.clear()
|
||||
await self._mux._notify.wait()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Subgraphs projection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _SubgraphsProjection:
|
||||
"""Async iterable yielding :class:`AsyncSubgraphRunStream` for each discovered subgraph."""
|
||||
|
||||
def __init__(self, mux: AsyncStreamMux, ns: list[str]) -> None:
|
||||
self._mux = mux
|
||||
self._ns = ns
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[AsyncSubgraphRunStream]:
|
||||
async for segment in self._mux.subscribe_subgraphs(self._ns):
|
||||
child_ns = self._ns + [segment]
|
||||
child_transformers: list[StreamTransformer] = [
|
||||
ValuesTransformer(),
|
||||
MessagesTransformer(
|
||||
namespace=child_ns, stream_cls=AsyncChatModelStream
|
||||
),
|
||||
]
|
||||
for t in child_transformers:
|
||||
t.init()
|
||||
self._mux.register_transformer(t)
|
||||
|
||||
yield AsyncSubgraphRunStream(
|
||||
mux=self._mux,
|
||||
namespace=child_ns,
|
||||
transformers=child_transformers,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# AsyncGraphRunStream
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class AsyncGraphRunStream:
|
||||
"""The async run stream returned by ``StreamingHandler.astream()``.
|
||||
|
||||
Async-iterable over all :class:`ProtocolEvent` instances. Named
|
||||
projections provide ergonomic access to values, messages, subgraphs,
|
||||
and output.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
mux: AsyncStreamMux,
|
||||
namespace: list[str] | None = None,
|
||||
transformers: list[StreamTransformer],
|
||||
abort_event: asyncio.Event | None = None,
|
||||
output_mapper: Callable[[Any], Any] | None = None,
|
||||
) -> None:
|
||||
self._mux = mux
|
||||
self._ns = namespace or []
|
||||
self._transformers = transformers
|
||||
self._abort_event = abort_event or asyncio.Event()
|
||||
self._output_mapper = output_mapper
|
||||
|
||||
# -- Transformer lookup -------------------------------------------------
|
||||
|
||||
def _find_transformer(self, name: str) -> StreamTransformer | None:
|
||||
for t in self._transformers:
|
||||
if getattr(t, "name", None) == name:
|
||||
return t
|
||||
return None
|
||||
|
||||
# -- Raw event iteration ------------------------------------------------
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[ProtocolEvent]:
|
||||
return self._mux.subscribe_events(self._ns)
|
||||
|
||||
# -- Named projections --------------------------------------------------
|
||||
|
||||
@property
|
||||
def values(self) -> _ValuesProjection:
|
||||
"""Async iterable of state snapshots; awaitable for final state."""
|
||||
t = self._find_transformer("values")
|
||||
return _ValuesProjection(self._mux, t, self._ns, self._output_mapper)
|
||||
|
||||
@property
|
||||
def output(self) -> _ValuesProjection:
|
||||
"""Awaitable for the final output state."""
|
||||
t = self._find_transformer("values")
|
||||
return _ValuesProjection(self._mux, t, self._ns, self._output_mapper)
|
||||
|
||||
@property
|
||||
def messages(self) -> _MessagesProjection:
|
||||
"""Async iterable of :class:`AsyncChatModelStream` instances."""
|
||||
t = self._find_transformer("messages")
|
||||
return _MessagesProjection(self._mux, t)
|
||||
|
||||
def messages_from(self, node: str) -> _MessagesProjection:
|
||||
"""Async iterable of messages from a specific node."""
|
||||
filtered = MessagesTransformer(
|
||||
namespace=self._ns,
|
||||
node_filter=node,
|
||||
stream_cls=AsyncChatModelStream,
|
||||
)
|
||||
self._mux.register_transformer(filtered)
|
||||
return _MessagesProjection(self._mux, filtered)
|
||||
|
||||
@property
|
||||
def subgraphs(self) -> _SubgraphsProjection:
|
||||
"""Async iterable of :class:`AsyncSubgraphRunStream` for child graphs."""
|
||||
return _SubgraphsProjection(self._mux, self._ns)
|
||||
|
||||
# -- State --------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def interrupted(self) -> bool:
|
||||
return self._mux.interrupted
|
||||
|
||||
@property
|
||||
def interrupts(self) -> list[InterruptPayload]:
|
||||
return self._mux.interrupts
|
||||
|
||||
# -- Cancellation -------------------------------------------------------
|
||||
|
||||
def abort(self, reason: str | None = None) -> None:
|
||||
"""Signal cancellation of the run."""
|
||||
self._abort_event.set()
|
||||
|
||||
@property
|
||||
def signal(self) -> asyncio.Event:
|
||||
"""The underlying cancellation event."""
|
||||
return self._abort_event
|
||||
|
||||
# -- Extensions ---------------------------------------------------------
|
||||
|
||||
@property
|
||||
def extensions(self) -> dict[str, Any]:
|
||||
"""All transformer projections."""
|
||||
result: dict[str, Any] = {}
|
||||
for t in self._transformers:
|
||||
name = getattr(t, "name", None)
|
||||
value = getattr(t, "value", None)
|
||||
if name is not None and value is not None:
|
||||
result[name] = value
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# AsyncSubgraphRunStream
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class AsyncSubgraphRunStream(AsyncGraphRunStream):
|
||||
"""An :class:`AsyncGraphRunStream` for a child subgraph.
|
||||
|
||||
Adds ``.name`` and ``.index`` parsed from the last namespace segment
|
||||
(e.g. ``"researcher:2"`` → ``name="researcher"``, ``index=2``).
|
||||
"""
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
if self._ns:
|
||||
segment = self._ns[-1]
|
||||
return segment.split(":")[0] if ":" in segment else segment
|
||||
return ""
|
||||
|
||||
@property
|
||||
def index(self) -> int:
|
||||
if self._ns:
|
||||
segment = self._ns[-1]
|
||||
if ":" in segment:
|
||||
try:
|
||||
return int(segment.split(":")[-1])
|
||||
except ValueError:
|
||||
pass
|
||||
return 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Async factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def create_async_graph_run_stream(
|
||||
source: AsyncIterator[tuple[tuple[str, ...], str, Any]],
|
||||
*,
|
||||
transformers: list[StreamTransformer] | None = None,
|
||||
abort_event: asyncio.Event | None = None,
|
||||
output_mapper: Callable[[Any], Any] | None = None,
|
||||
) -> AsyncGraphRunStream:
|
||||
"""Create an :class:`AsyncGraphRunStream` from a raw async stream source.
|
||||
|
||||
1. Creates a :class:`StreamMux`
|
||||
2. Registers built-in ``ValuesTransformer`` and ``MessagesTransformer``
|
||||
3. Registers user-supplied transformers
|
||||
4. Creates the root ``AsyncGraphRunStream``
|
||||
5. Starts a background pump task that reads from *source*,
|
||||
converts each chunk to a ``ProtocolEvent``, and pushes it
|
||||
through the mux
|
||||
6. Returns the ``AsyncGraphRunStream``
|
||||
"""
|
||||
abort = abort_event or asyncio.Event()
|
||||
|
||||
# Built-in transformers first, then user-supplied
|
||||
all_transformers: list[StreamTransformer] = [
|
||||
ValuesTransformer(),
|
||||
MessagesTransformer(stream_cls=AsyncChatModelStream),
|
||||
]
|
||||
all_transformers.extend(transformers or [])
|
||||
|
||||
# Initialize transformers, collecting projections to wire after mux creation
|
||||
projections: list[Any] = []
|
||||
for t in all_transformers:
|
||||
projection = t.init()
|
||||
if projection is not None:
|
||||
projections.append(projection)
|
||||
|
||||
mux = AsyncStreamMux(transformers=all_transformers)
|
||||
|
||||
# Wire any StreamChannel instances found in transformer projections
|
||||
for projection in projections:
|
||||
mux.wire_channels(projection)
|
||||
|
||||
# Create the root stream
|
||||
run_stream = AsyncGraphRunStream(
|
||||
mux=mux,
|
||||
transformers=all_transformers,
|
||||
abort_event=abort,
|
||||
output_mapper=output_mapper,
|
||||
)
|
||||
|
||||
# Start the pump task
|
||||
async def pump() -> None:
|
||||
try:
|
||||
async for ns, mode, payload in source:
|
||||
if abort.is_set():
|
||||
break
|
||||
# Extract node name embedded by StreamProtocolMessagesHandler.
|
||||
node: str | None = None
|
||||
if (
|
||||
mode == "messages"
|
||||
and isinstance(payload, dict)
|
||||
and "__node__" in payload
|
||||
):
|
||||
payload = dict(payload)
|
||||
node = payload.pop("__node__")
|
||||
event = convert_to_protocol_event(ns, mode, payload, node=node)
|
||||
if event is not None:
|
||||
mux.push(event)
|
||||
mux.close()
|
||||
except Exception as exc:
|
||||
mux.fail(exc)
|
||||
|
||||
asyncio.get_running_loop().create_task(pump())
|
||||
|
||||
return run_stream
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GraphRunStream — returned by StreamingHandler.stream()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _PumpDrivenLog:
|
||||
"""Wraps a list so that iteration drives the sync pump.
|
||||
|
||||
Used by all :class:`GraphRunStream` projections (``__iter__``,
|
||||
``.values``, ``.messages``, ``.extensions``) so that iterating
|
||||
any projection lazily consumes the source.
|
||||
"""
|
||||
|
||||
__slots__ = ("_log", "_pump_one")
|
||||
|
||||
def __init__(self, log: list, pump_one: Callable[[], bool]) -> None:
|
||||
self._log = log
|
||||
self._pump_one = pump_one
|
||||
|
||||
def __iter__(self) -> Iterator[Any]:
|
||||
cursor = 0
|
||||
while True:
|
||||
if cursor < len(self._log):
|
||||
yield self._log[cursor]
|
||||
cursor += 1
|
||||
elif not self._pump_one():
|
||||
return
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._log)
|
||||
|
||||
def __getitem__(self, index: int) -> Any:
|
||||
return self._log[index]
|
||||
|
||||
|
||||
class GraphRunStream:
|
||||
"""Synchronous run stream returned by ``StreamingHandler.stream()``.
|
||||
|
||||
All projections are blocking / sync-iterable. Internally uses
|
||||
the same ``StreamMux`` and transformer pipeline, but without an
|
||||
async event loop.
|
||||
|
||||
The source iterator is consumed lazily: each projection pulls
|
||||
events from the source on demand rather than eagerly buffering
|
||||
everything upfront. This means callers see events as soon as
|
||||
they are produced by the underlying ``stream()`` call.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
mux: StreamMux,
|
||||
source: Iterator[tuple[tuple[str, ...], str, Any]],
|
||||
namespace: list[str] | None = None,
|
||||
transformers: list[StreamTransformer],
|
||||
output_mapper: Callable[[Any], Any] | None = None,
|
||||
) -> None:
|
||||
self._mux = mux
|
||||
self._source = source
|
||||
self._source_exhausted = False
|
||||
self._ns = namespace or []
|
||||
self._transformers = transformers
|
||||
self._output_mapper = output_mapper
|
||||
|
||||
# -- Transformer lookup -------------------------------------------------
|
||||
|
||||
def _find_transformer(self, name: str) -> StreamTransformer | None:
|
||||
for t in self._transformers:
|
||||
if getattr(t, "name", None) == name:
|
||||
return t
|
||||
return None
|
||||
|
||||
# -- Lazy pump ----------------------------------------------------------
|
||||
|
||||
def _pump_one(self) -> bool:
|
||||
"""Pull one item from the source, convert it, and push through the mux.
|
||||
|
||||
Returns ``True`` if an item was consumed, ``False`` if the source
|
||||
is exhausted (or was already exhausted).
|
||||
"""
|
||||
if self._source_exhausted:
|
||||
return False
|
||||
try:
|
||||
ns, mode, payload = next(self._source)
|
||||
except StopIteration:
|
||||
self._source_exhausted = True
|
||||
self._mux.close()
|
||||
return False
|
||||
except Exception as exc:
|
||||
self._source_exhausted = True
|
||||
self._mux.fail(exc)
|
||||
return False
|
||||
|
||||
node: str | None = None
|
||||
if mode == "messages" and isinstance(payload, dict) and "__node__" in payload:
|
||||
payload = dict(payload)
|
||||
node = payload.pop("__node__")
|
||||
event = convert_to_protocol_event(ns, mode, payload, node=node)
|
||||
if event is not None:
|
||||
self._mux.push(event)
|
||||
return True
|
||||
|
||||
def _pump_all(self) -> None:
|
||||
"""Drain the source iterator completely."""
|
||||
while self._pump_one():
|
||||
pass
|
||||
|
||||
# -- Helpers ------------------------------------------------------------
|
||||
|
||||
def _map(self, value: Any) -> Any:
|
||||
if value is not None and self._output_mapper is not None:
|
||||
return self._output_mapper(value)
|
||||
return value
|
||||
|
||||
# -- Raw event iteration (sync) -----------------------------------------
|
||||
|
||||
def __iter__(self) -> Iterator[ProtocolEvent]:
|
||||
for event in _PumpDrivenLog(self._mux.event_log, self._pump_one):
|
||||
ns = event["params"].get("namespace", [])
|
||||
if not self._ns or ns[: len(self._ns)] == self._ns:
|
||||
yield event
|
||||
|
||||
# -- Named projections (sync) -------------------------------------------
|
||||
|
||||
@property
|
||||
def output(self) -> Any:
|
||||
"""The final output state (blocking). Drains the source."""
|
||||
self._pump_all()
|
||||
return self._map(self._mux.get_latest_values(self._ns))
|
||||
|
||||
@property
|
||||
def values(self) -> Iterator[Any]:
|
||||
"""Sync iterable of intermediate state snapshots."""
|
||||
t = self._find_transformer("values")
|
||||
if t is None:
|
||||
return
|
||||
for item in _PumpDrivenLog(t.value, self._pump_one):
|
||||
if item.get("namespace", []) == self._ns:
|
||||
yield self._map(item["data"])
|
||||
|
||||
@property
|
||||
def messages(self) -> Iterator[ChatModelStream]:
|
||||
"""Sync iterable of :class:`ChatModelStream` instances.
|
||||
|
||||
Each yielded ``ChatModelStream`` is fully populated (``done=True``)
|
||||
so that sync consumers can read ``.text``, ``.reasoning``, and
|
||||
``.usage`` immediately.
|
||||
"""
|
||||
t = self._find_transformer("messages")
|
||||
if t is None:
|
||||
return
|
||||
for msg in _PumpDrivenLog(t.value, self._pump_one):
|
||||
# Pump until this message is complete so sync consumers
|
||||
# get a fully populated ChatModelStream.
|
||||
while not msg.done:
|
||||
if not self._pump_one():
|
||||
break
|
||||
yield msg
|
||||
|
||||
# -- State --------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def interrupted(self) -> bool:
|
||||
return self._mux.interrupted
|
||||
|
||||
@property
|
||||
def interrupts(self) -> list[InterruptPayload]:
|
||||
return self._mux.interrupts
|
||||
|
||||
# -- Extensions ---------------------------------------------------------
|
||||
|
||||
@property
|
||||
def extensions(self) -> dict[str, Any]:
|
||||
"""All transformer projections as pump-driven iterables."""
|
||||
result: dict[str, Any] = {}
|
||||
for t in self._transformers:
|
||||
name = getattr(t, "name", None)
|
||||
value = getattr(t, "value", None)
|
||||
if name is not None and value is not None:
|
||||
if isinstance(value, list):
|
||||
result[name] = _PumpDrivenLog(value, self._pump_one)
|
||||
else:
|
||||
result[name] = value
|
||||
return result
|
||||
|
||||
|
||||
def create_graph_run_stream(
|
||||
source: Iterator[tuple[tuple[str, ...], str, Any]],
|
||||
*,
|
||||
transformers: list[StreamTransformer] | None = None,
|
||||
output_mapper: Callable[[Any], Any] | None = None,
|
||||
) -> GraphRunStream:
|
||||
"""Create a :class:`GraphRunStream` from a sync stream source.
|
||||
|
||||
The source iterator is stored on the returned stream and consumed
|
||||
lazily as projections are iterated.
|
||||
|
||||
Built-in transformers (values, messages) are always registered first
|
||||
so that user-supplied transformers see events after built-in
|
||||
processing.
|
||||
"""
|
||||
# Built-in transformers first, then user-supplied
|
||||
all_transformers: list[StreamTransformer] = [
|
||||
ValuesTransformer(),
|
||||
MessagesTransformer(),
|
||||
]
|
||||
all_transformers.extend(transformers or [])
|
||||
|
||||
projections: list[Any] = []
|
||||
for t in all_transformers:
|
||||
projection = t.init()
|
||||
if projection is not None:
|
||||
projections.append(projection)
|
||||
|
||||
mux = StreamMux(transformers=all_transformers)
|
||||
|
||||
for projection in projections:
|
||||
mux.wire_channels(projection)
|
||||
|
||||
return GraphRunStream(
|
||||
mux=mux,
|
||||
source=source,
|
||||
transformers=all_transformers,
|
||||
output_mapper=output_mapper,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AsyncGraphRunStream",
|
||||
"AsyncSubgraphRunStream",
|
||||
"GraphRunStream",
|
||||
"create_async_graph_run_stream",
|
||||
"create_graph_run_stream",
|
||||
]
|
||||
@@ -0,0 +1,49 @@
|
||||
"""StreamChannel — typed push-based channel for StreamTransformer projections.
|
||||
|
||||
A ``StreamChannel`` wraps a list and declares a protocol channel name.
|
||||
When the :class:`StreamMux` detects a ``StreamChannel`` in a transformer's
|
||||
``init()`` return, it wires every ``push()`` call to inject a
|
||||
:class:`ProtocolEvent` into the main event stream using the channel's
|
||||
name as the ``method``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any, Generic, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class StreamChannel(Generic[T]):
|
||||
"""A typed push-based channel that integrates with the mux.
|
||||
|
||||
Transformer authors create a ``StreamChannel`` in ``init()`` and
|
||||
call ``push()`` inside ``process()`` to emit domain objects. The
|
||||
mux auto-wires pushes to protocol events.
|
||||
"""
|
||||
|
||||
__slots__ = ("channel_name", "_items", "_on_push")
|
||||
|
||||
def __init__(self, name: str) -> None:
|
||||
self.channel_name = name
|
||||
self._items: list[T] = []
|
||||
self._on_push: Callable[[Any], None] | None = None
|
||||
|
||||
def push(self, item: T) -> None:
|
||||
"""Push an item to the channel."""
|
||||
self._items.append(item)
|
||||
if self._on_push is not None:
|
||||
self._on_push(item)
|
||||
|
||||
def _wire(self, fn: Callable[[Any], None]) -> None:
|
||||
"""Wire a callback invoked on every ``push()``. Called by the mux."""
|
||||
self._on_push = fn
|
||||
|
||||
|
||||
def is_stream_channel(value: object) -> bool:
|
||||
"""Check if *value* is a :class:`StreamChannel` instance."""
|
||||
return isinstance(value, StreamChannel)
|
||||
|
||||
|
||||
__all__ = ["StreamChannel", "is_stream_channel"]
|
||||
@@ -0,0 +1,168 @@
|
||||
"""Experimental streaming wrapper for CompiledGraph.
|
||||
|
||||
``StreamingHandler`` wraps a compiled graph and exposes the new streaming
|
||||
API without adding methods to the ``CompiledGraph`` class itself.
|
||||
|
||||
Usage::
|
||||
|
||||
from langgraph.stream import StreamingHandler
|
||||
|
||||
s = StreamingHandler(graph)
|
||||
|
||||
# async
|
||||
run = await s.astream(input)
|
||||
async for msg in run.messages:
|
||||
...
|
||||
|
||||
# sync
|
||||
run = s.stream(input)
|
||||
for event in run:
|
||||
...
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from langgraph._internal._config import patch_configurable
|
||||
from langgraph.stream._convert import STREAM_V2_MODES
|
||||
from langgraph.stream._types import StreamTransformer
|
||||
from langgraph.stream.run_stream import (
|
||||
AsyncGraphRunStream,
|
||||
GraphRunStream,
|
||||
create_async_graph_run_stream,
|
||||
create_graph_run_stream,
|
||||
)
|
||||
from langgraph.types import All
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langgraph.pregel import Pregel
|
||||
|
||||
#: Config key that activates the protocol messages handler.
|
||||
#: Duplicated here to avoid a circular import with ``pregel._messages_v2``.
|
||||
PROTOCOL_MESSAGES_STREAM_KEY = "__protocol_messages_stream"
|
||||
|
||||
|
||||
class StreamingHandler:
|
||||
"""Experimental streaming wrapper around a compiled graph.
|
||||
|
||||
Provides ``.stream()`` and ``.astream()`` returning
|
||||
:class:`GraphRunStream` / :class:`AsyncGraphRunStream` with
|
||||
ergonomic projections (``run.values``, ``run.messages``,
|
||||
``run.subgraphs``, ``run.output``).
|
||||
|
||||
Args:
|
||||
graph: A compiled LangGraph (``Pregel`` instance).
|
||||
"""
|
||||
|
||||
def __init__(self, graph: Pregel) -> None:
|
||||
self._graph = graph
|
||||
|
||||
async def astream(
|
||||
self,
|
||||
input: Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
context: Any | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
debug: bool | None = None,
|
||||
transformers: list[StreamTransformer] | None = None,
|
||||
) -> AsyncGraphRunStream:
|
||||
"""Stream graph execution, returning an
|
||||
:class:`~langgraph.stream.run_stream.AsyncGraphRunStream`.
|
||||
|
||||
The returned stream provides ergonomic projections:
|
||||
|
||||
- ``await run.output`` -- final state
|
||||
- ``async for v in run.values`` -- intermediate state snapshots
|
||||
- ``async for msg in run.messages`` -- per-message
|
||||
:class:`~langgraph.stream.chat_model_stream.AsyncChatModelStream`
|
||||
objects
|
||||
- ``async for sub in run.subgraphs`` -- child
|
||||
:class:`~langgraph.stream.run_stream.AsyncSubgraphRunStream`
|
||||
instances
|
||||
- ``async for event in run`` -- raw
|
||||
:class:`~langgraph.stream._types.ProtocolEvent` objects
|
||||
|
||||
Args:
|
||||
input: The input to the graph.
|
||||
config: The configuration to use for the run.
|
||||
context: The static context to use for the run.
|
||||
interrupt_before: Nodes to interrupt before.
|
||||
interrupt_after: Nodes to interrupt after.
|
||||
debug: Whether to emit debug events.
|
||||
transformers: Optional user-supplied
|
||||
:class:`~langgraph.stream._types.StreamTransformer` instances
|
||||
for custom projections (available on ``run.extensions``).
|
||||
|
||||
Returns:
|
||||
An :class:`~langgraph.stream.run_stream.AsyncGraphRunStream`.
|
||||
"""
|
||||
merged_config = patch_configurable(config, {PROTOCOL_MESSAGES_STREAM_KEY: True})
|
||||
|
||||
source = cast(
|
||||
AsyncIterator[tuple[tuple[str, ...], str, Any]],
|
||||
self._graph.astream(
|
||||
input,
|
||||
merged_config,
|
||||
context=context,
|
||||
stream_mode=STREAM_V2_MODES,
|
||||
subgraphs=True,
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
debug=debug,
|
||||
version="v1",
|
||||
),
|
||||
)
|
||||
|
||||
return await create_async_graph_run_stream(
|
||||
source,
|
||||
transformers=transformers,
|
||||
output_mapper=self._graph._output_mapper,
|
||||
)
|
||||
|
||||
def stream(
|
||||
self,
|
||||
input: Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
context: Any | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
debug: bool | None = None,
|
||||
transformers: list[StreamTransformer] | None = None,
|
||||
) -> GraphRunStream:
|
||||
"""Synchronous variant of :meth:`astream`.
|
||||
|
||||
Returns a :class:`~langgraph.stream.run_stream.GraphRunStream`
|
||||
immediately. The underlying source is consumed lazily as
|
||||
projections are iterated.
|
||||
|
||||
See :meth:`astream` for full documentation.
|
||||
"""
|
||||
merged_config = patch_configurable(config, {PROTOCOL_MESSAGES_STREAM_KEY: True})
|
||||
|
||||
source = cast(
|
||||
Iterator[tuple[tuple[str, ...], str, Any]],
|
||||
self._graph.stream(
|
||||
input,
|
||||
merged_config,
|
||||
context=context,
|
||||
stream_mode=STREAM_V2_MODES,
|
||||
subgraphs=True,
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
debug=debug,
|
||||
version="v1",
|
||||
),
|
||||
)
|
||||
|
||||
return create_graph_run_stream(
|
||||
source,
|
||||
transformers=transformers,
|
||||
output_mapper=self._graph._output_mapper,
|
||||
)
|
||||
@@ -0,0 +1,172 @@
|
||||
"""Built-in stream transformers for StreamingHandler.
|
||||
|
||||
``ValuesTransformer`` extracts ``values`` events and maintains the latest
|
||||
state per namespace. ``MessagesTransformer`` groups ``messages`` events
|
||||
into :class:`ChatModelStream` instances.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from langgraph.stream._types import ProtocolEvent, StreamTransformer
|
||||
from langgraph.stream.chat_model_stream import ChatModelStream
|
||||
|
||||
# Type alias for the stream class constructor signature
|
||||
_StreamCls = type[ChatModelStream]
|
||||
|
||||
|
||||
class ValuesTransformer(StreamTransformer):
|
||||
"""Extracts ``values`` events and populates a values log.
|
||||
|
||||
Maintains the latest state per namespace and provides a separate
|
||||
log that :class:`AsyncGraphRunStream` / :class:`GraphRunStream` uses for ``.values``
|
||||
iteration.
|
||||
"""
|
||||
|
||||
name = "values"
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._values_log: list[dict[str, Any]] = []
|
||||
self._latest: dict[str, Any] = {}
|
||||
|
||||
@property
|
||||
def value(self) -> list[dict[str, Any]]:
|
||||
return self._values_log
|
||||
|
||||
@property
|
||||
def values_log(self) -> list[dict[str, Any]]:
|
||||
return self._values_log
|
||||
|
||||
def get_latest(self, ns_key: str = "") -> Any:
|
||||
return self._latest.get(ns_key)
|
||||
|
||||
def init(self) -> Any:
|
||||
return None
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
if event["method"] != "values":
|
||||
return True
|
||||
|
||||
ns = event["params"].get("namespace", [])
|
||||
data = event["params"]["data"]
|
||||
ns_key = "|".join(ns) if ns else ""
|
||||
self._latest[ns_key] = data
|
||||
|
||||
# Append to the values log for iteration
|
||||
self._values_log.append({"namespace": ns, "data": data})
|
||||
return True
|
||||
|
||||
def finalize(self) -> None:
|
||||
pass
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class MessagesTransformer(StreamTransformer):
|
||||
"""Groups ``messages`` events into :class:`ChatModelStream` instances.
|
||||
|
||||
One ``ChatModelStream`` is created per ``message-start`` event.
|
||||
Content-block events are routed to the active stream until
|
||||
``message-finish`` or ``message-error`` closes it.
|
||||
"""
|
||||
|
||||
name = "messages"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
namespace: list[str] | None = None,
|
||||
node_filter: str | None = None,
|
||||
stream_cls: _StreamCls | None = None,
|
||||
) -> None:
|
||||
self._namespace = namespace
|
||||
self._node_filter = node_filter
|
||||
self._stream_cls: _StreamCls = stream_cls or ChatModelStream
|
||||
|
||||
# Message log for .messages iteration
|
||||
self._messages_log: list[ChatModelStream] = []
|
||||
|
||||
# Current active stream per namespace key
|
||||
self._active: dict[str, ChatModelStream] = {}
|
||||
|
||||
@property
|
||||
def value(self) -> list[ChatModelStream]:
|
||||
return self._messages_log
|
||||
|
||||
@property
|
||||
def messages_log(self) -> list[ChatModelStream]:
|
||||
return self._messages_log
|
||||
|
||||
def init(self) -> Any:
|
||||
return None
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
if event["method"] != "messages":
|
||||
return True
|
||||
|
||||
ns = event["params"].get("namespace", [])
|
||||
node = event["params"].get("node")
|
||||
data = event["params"]["data"]
|
||||
|
||||
# Apply namespace filter
|
||||
if self._namespace is not None:
|
||||
if ns[: len(self._namespace)] != self._namespace:
|
||||
return True
|
||||
|
||||
# Apply node filter
|
||||
if self._node_filter is not None and node != self._node_filter:
|
||||
return True
|
||||
|
||||
ns_key = "|".join(ns) if ns else ""
|
||||
event_type = data.get("event") if isinstance(data, dict) else None
|
||||
|
||||
if event_type == "message-start":
|
||||
stream = self._stream_cls(
|
||||
namespace=ns,
|
||||
node=node,
|
||||
message_id=data.get("message_id"),
|
||||
)
|
||||
self._active[ns_key] = stream
|
||||
self._messages_log.append(stream)
|
||||
|
||||
elif event_type in ("content-block-delta", "content-block-start"):
|
||||
active = self._active.get(ns_key)
|
||||
if active is not None and event_type == "content-block-delta":
|
||||
active._push_content_block_delta(data)
|
||||
|
||||
elif event_type == "content-block-finish":
|
||||
active = self._active.get(ns_key)
|
||||
if active is not None:
|
||||
active._push_content_block_finish(data)
|
||||
|
||||
elif event_type == "message-finish":
|
||||
active = self._active.pop(ns_key, None)
|
||||
if active is not None:
|
||||
active._finish(data)
|
||||
|
||||
elif event_type == "error":
|
||||
active = self._active.pop(ns_key, None)
|
||||
if active is not None:
|
||||
msg = data.get("message", "Unknown error")
|
||||
active._fail(RuntimeError(msg))
|
||||
|
||||
return True
|
||||
|
||||
def finalize(self) -> None:
|
||||
# Finish any remaining active streams
|
||||
for stream in self._active.values():
|
||||
stream._finish({"reason": "stop"})
|
||||
self._active.clear()
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
for stream in self._active.values():
|
||||
stream._fail(err)
|
||||
self._active.clear()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MessagesTransformer",
|
||||
"ValuesTransformer",
|
||||
]
|
||||
@@ -0,0 +1,623 @@
|
||||
import asyncio
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.config import get_stream_writer
|
||||
from langgraph.graph import END, START, MessagesState, StateGraph
|
||||
from langgraph.stream import AsyncChatModelStream, StreamingHandler
|
||||
from langgraph.stream._types import ProtocolEvent, StreamTransformer
|
||||
from tests.fake_chat import FakeChatModel
|
||||
|
||||
|
||||
class State(TypedDict):
|
||||
value: str
|
||||
items: Annotated[list[str], lambda a, b: a + b]
|
||||
|
||||
|
||||
def make_simple_graph():
|
||||
def node_a(state):
|
||||
return {"value": state["value"] + "_a", "items": ["a"]}
|
||||
|
||||
def node_b(state):
|
||||
return {"value": state["value"] + "_b", "items": ["b"]}
|
||||
|
||||
graph = StateGraph(State)
|
||||
graph.add_node("node_a", node_a)
|
||||
graph.add_node("node_b", node_b)
|
||||
graph.add_edge(START, "node_a")
|
||||
graph.add_edge("node_a", "node_b")
|
||||
graph.add_edge("node_b", END)
|
||||
return graph.compile()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_output():
|
||||
graph = make_simple_graph()
|
||||
run = await StreamingHandler(graph).astream({"value": "x", "items": []})
|
||||
await asyncio.sleep(0.1)
|
||||
output = await run.output
|
||||
assert output == {"value": "x_a_b", "items": ["a", "b"]}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_values_iteration():
|
||||
graph = make_simple_graph()
|
||||
run = await StreamingHandler(graph).astream({"value": "x", "items": []})
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
snapshots = []
|
||||
async for v in run.values:
|
||||
snapshots.append(v)
|
||||
|
||||
assert len(snapshots) == 3
|
||||
assert snapshots[0]["value"] == "x"
|
||||
assert snapshots[1]["value"] == "x_a"
|
||||
assert snapshots[2]["value"] == "x_a_b"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_updates_in_raw_events():
|
||||
graph = make_simple_graph()
|
||||
run = await StreamingHandler(graph).astream({"value": "x", "items": []})
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
updates = []
|
||||
async for event in run:
|
||||
if event["method"] == "updates":
|
||||
updates.append(event["params"]["data"])
|
||||
|
||||
assert len(updates) == 2
|
||||
assert "node_a" in updates[0]
|
||||
assert "node_b" in updates[1]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_messages_with_chat_model():
|
||||
model = FakeChatModel(messages=[AIMessage(content="Hello world")])
|
||||
|
||||
def agent(state):
|
||||
return {"messages": [model.invoke(state["messages"])]}
|
||||
|
||||
graph = StateGraph(MessagesState)
|
||||
graph.add_node("agent", agent)
|
||||
graph.add_edge(START, "agent")
|
||||
graph.add_edge("agent", END)
|
||||
compiled = graph.compile()
|
||||
|
||||
run = await StreamingHandler(compiled).astream(
|
||||
{"messages": [HumanMessage(content="hi")]}
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
messages_seen = []
|
||||
async for msg in run.messages:
|
||||
messages_seen.append(msg)
|
||||
|
||||
assert len(messages_seen) >= 1
|
||||
msg = messages_seen[0]
|
||||
assert isinstance(msg, AsyncChatModelStream)
|
||||
text = await msg.text
|
||||
assert text == "Hello world"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_custom_events():
|
||||
def node(state):
|
||||
writer = get_stream_writer()
|
||||
writer("hello")
|
||||
writer(42)
|
||||
return {"value": state["value"] + "_a", "items": ["a"]}
|
||||
|
||||
graph = StateGraph(State)
|
||||
graph.add_node("node_a", node)
|
||||
graph.add_edge(START, "node_a")
|
||||
graph.add_edge("node_a", END)
|
||||
compiled = graph.compile()
|
||||
|
||||
run = await StreamingHandler(compiled).astream({"value": "x", "items": []})
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
custom_payloads = []
|
||||
async for event in run:
|
||||
if event["method"] == "custom":
|
||||
custom_payloads.append(event["params"]["data"])
|
||||
|
||||
assert "hello" in custom_payloads
|
||||
assert 42 in custom_payloads
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_multiple_modes_present():
|
||||
graph = make_simple_graph()
|
||||
run = await StreamingHandler(graph).astream({"value": "x", "items": []})
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
methods = set()
|
||||
async for event in run:
|
||||
methods.add(event["method"])
|
||||
|
||||
assert {"values", "updates", "tasks", "debug"} <= methods
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_interrupted_false():
|
||||
graph = make_simple_graph()
|
||||
run = await StreamingHandler(graph).astream({"value": "x", "items": []})
|
||||
await asyncio.sleep(0.1)
|
||||
async for _ in run:
|
||||
pass
|
||||
assert run.interrupted is False
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_regression_v1_stream_unchanged():
|
||||
graph = make_simple_graph()
|
||||
chunks = []
|
||||
async for chunk in graph.astream(
|
||||
{"value": "x", "items": []}, stream_mode="values", version="v1"
|
||||
):
|
||||
chunks.append(chunk)
|
||||
for chunk in chunks:
|
||||
assert isinstance(chunk, dict)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_regression_v2_stream_unchanged():
|
||||
graph = make_simple_graph()
|
||||
chunks = []
|
||||
async for chunk in graph.astream(
|
||||
{"value": "x", "items": []}, stream_mode="values", version="v2"
|
||||
):
|
||||
chunks.append(chunk)
|
||||
assert len(chunks) >= 1
|
||||
for chunk in chunks:
|
||||
assert isinstance(chunk, dict)
|
||||
assert "type" in chunk
|
||||
assert chunk["type"] == "values"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_regression_invoke_unchanged():
|
||||
graph = make_simple_graph()
|
||||
result = await graph.ainvoke({"value": "x", "items": []})
|
||||
assert result == {"value": "x_a_b", "items": ["a", "b"]}
|
||||
|
||||
|
||||
def test_sync_stream_output():
|
||||
graph = make_simple_graph()
|
||||
run = StreamingHandler(graph).stream({"value": "x", "items": []})
|
||||
assert run.output == {"value": "x_a_b", "items": ["a", "b"]}
|
||||
|
||||
|
||||
def test_sync_stream_values():
|
||||
graph = make_simple_graph()
|
||||
run = StreamingHandler(graph).stream({"value": "x", "items": []})
|
||||
snapshots = list(run.values)
|
||||
assert len(snapshots) == 3
|
||||
assert snapshots[0]["value"] == "x"
|
||||
assert snapshots[2]["value"] == "x_a_b"
|
||||
|
||||
|
||||
def test_sync_stream_raw_events():
|
||||
graph = make_simple_graph()
|
||||
run = StreamingHandler(graph).stream({"value": "x", "items": []})
|
||||
methods = {e["method"] for e in run}
|
||||
assert {"values", "updates", "tasks", "debug"} <= methods
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Typed output (pydantic)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ModelState(BaseModel):
|
||||
value: str
|
||||
items: Annotated[list[str], lambda a, b: a + b]
|
||||
|
||||
|
||||
def _make_model_state_graph():
|
||||
def node_a(state):
|
||||
return {"value": state.value + "_a", "items": ["a"]}
|
||||
|
||||
graph = StateGraph(ModelState)
|
||||
graph.add_node("node_a", node_a)
|
||||
graph.add_edge(START, "node_a")
|
||||
graph.add_edge("node_a", END)
|
||||
return graph.compile()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pydantic_output():
|
||||
graph = _make_model_state_graph()
|
||||
run = await StreamingHandler(graph).astream(ModelState(value="x", items=[]))
|
||||
await asyncio.sleep(0.1)
|
||||
output = await run.output
|
||||
assert isinstance(output, ModelState)
|
||||
assert output.value == "x_a"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pydantic_values():
|
||||
graph = _make_model_state_graph()
|
||||
run = await StreamingHandler(graph).astream(ModelState(value="x", items=[]))
|
||||
await asyncio.sleep(0.1)
|
||||
snapshots = []
|
||||
async for v in run.values:
|
||||
snapshots.append(v)
|
||||
for v in snapshots:
|
||||
assert isinstance(v, ModelState)
|
||||
|
||||
|
||||
def test_sync_pydantic_output():
|
||||
graph = _make_model_state_graph()
|
||||
run = StreamingHandler(graph).stream(ModelState(value="x", items=[]))
|
||||
assert isinstance(run.output, ModelState)
|
||||
assert run.output.value == "x_a"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Interrupts
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_interrupts():
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
|
||||
from langgraph.types import interrupt
|
||||
|
||||
def ask_human(state: State):
|
||||
answer = interrupt("what do you want?")
|
||||
return {"value": state["value"] + f"_{answer}", "items": [answer]}
|
||||
|
||||
graph = StateGraph(State)
|
||||
graph.add_node("ask", ask_human)
|
||||
graph.add_edge(START, "ask")
|
||||
graph.add_edge("ask", END)
|
||||
compiled = graph.compile(checkpointer=MemorySaver())
|
||||
|
||||
config = {"configurable": {"thread_id": "t1"}}
|
||||
run = await StreamingHandler(compiled).astream(
|
||||
{"value": "x", "items": []}, config=config
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
# Drain events
|
||||
async for _ in run:
|
||||
pass
|
||||
assert run.interrupted is True
|
||||
assert len(run.interrupts) > 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# messages_from(node)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_messages_from_node():
|
||||
model = FakeChatModel(messages=[AIMessage(content="from agent")])
|
||||
|
||||
def agent(state):
|
||||
return {"messages": [model.invoke(state["messages"])]}
|
||||
|
||||
def postprocess(state):
|
||||
return {"messages": state["messages"]}
|
||||
|
||||
graph = StateGraph(MessagesState)
|
||||
graph.add_node("agent", agent)
|
||||
graph.add_node("postprocess", postprocess)
|
||||
graph.add_edge(START, "agent")
|
||||
graph.add_edge("agent", "postprocess")
|
||||
graph.add_edge("postprocess", END)
|
||||
compiled = graph.compile()
|
||||
|
||||
run = await StreamingHandler(compiled).astream(
|
||||
{"messages": [HumanMessage(content="hi")]}
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# All messages
|
||||
all_msgs = []
|
||||
async for m in run.messages:
|
||||
all_msgs.append(m)
|
||||
assert len(all_msgs) >= 1
|
||||
# Node provenance should be set
|
||||
assert all_msgs[0].node == "agent"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Subgraph child stream
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_subgraph_child_output():
|
||||
"""AsyncSubgraphRunStream.output should contain the child graph's final state."""
|
||||
|
||||
class ChildState(TypedDict):
|
||||
value: str
|
||||
|
||||
class ParentState(TypedDict):
|
||||
value: str
|
||||
|
||||
def child_node(state):
|
||||
return {"value": state["value"] + "_child"}
|
||||
|
||||
child_graph = StateGraph(ChildState)
|
||||
child_graph.add_node("child_node", child_node)
|
||||
child_graph.add_edge(START, "child_node")
|
||||
child_graph.add_edge("child_node", END)
|
||||
# Add the compiled child as a node — this triggers LangGraph's
|
||||
# subgraph streaming mechanism and emits child namespace events.
|
||||
child_compiled = child_graph.compile()
|
||||
|
||||
parent_graph = StateGraph(ParentState)
|
||||
parent_graph.add_node("child_node", child_compiled)
|
||||
parent_graph.add_edge(START, "child_node")
|
||||
parent_graph.add_edge("child_node", END)
|
||||
parent_compiled = parent_graph.compile()
|
||||
|
||||
run = await StreamingHandler(parent_compiled).astream({"value": "x"})
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
subgraph_streams = []
|
||||
async for sub in run.subgraphs:
|
||||
subgraph_streams.append(sub)
|
||||
|
||||
assert len(subgraph_streams) >= 1
|
||||
child_output = await subgraph_streams[0].output
|
||||
assert child_output is not None
|
||||
assert child_output["value"] == "x_child"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Custom reducers / .extensions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _CountTransformer(StreamTransformer):
|
||||
"""Counts events. Exposes count via .value for extensions."""
|
||||
|
||||
name = "event_count"
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.value = 0
|
||||
|
||||
def init(self) -> Any:
|
||||
return None
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
self.value += 1
|
||||
return True
|
||||
|
||||
def finalize(self) -> None:
|
||||
pass
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_custom_reducer_extensions():
|
||||
graph = make_simple_graph()
|
||||
counter = _CountTransformer()
|
||||
run = await StreamingHandler(graph).astream(
|
||||
{"value": "x", "items": []}, transformers=[counter]
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
async for _ in run:
|
||||
pass
|
||||
assert counter.value > 0
|
||||
assert run.extensions["event_count"] == counter.value
|
||||
|
||||
|
||||
def test_sync_custom_reducer_extensions():
|
||||
graph = make_simple_graph()
|
||||
counter = _CountTransformer()
|
||||
run = StreamingHandler(graph).stream(
|
||||
{"value": "x", "items": []}, transformers=[counter]
|
||||
)
|
||||
for _ in run:
|
||||
pass
|
||||
assert counter.value > 0
|
||||
assert run.extensions["event_count"] == counter.value
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Double iteration over .values
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_values_double_iteration():
|
||||
"""Iterating over run.values twice should yield the same snapshots both times."""
|
||||
graph = make_simple_graph()
|
||||
run = await StreamingHandler(graph).astream({"value": "x", "items": []})
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
first = []
|
||||
async for v in run.values:
|
||||
first.append(v)
|
||||
|
||||
second = []
|
||||
async for v in run.values:
|
||||
second.append(v)
|
||||
|
||||
assert len(first) == 3
|
||||
assert first == second
|
||||
|
||||
|
||||
def test_sync_values_double_iteration():
|
||||
"""Iterating over run.values twice should yield the same snapshots both times."""
|
||||
graph = make_simple_graph()
|
||||
run = StreamingHandler(graph).stream({"value": "x", "items": []})
|
||||
|
||||
first = list(run.values)
|
||||
second = list(run.values)
|
||||
|
||||
assert len(first) == 3
|
||||
assert first == second
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_raw_events_double_iteration():
|
||||
"""Iterating over the raw event stream twice should yield the same events."""
|
||||
graph = make_simple_graph()
|
||||
run = await StreamingHandler(graph).astream({"value": "x", "items": []})
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
first = []
|
||||
async for event in run:
|
||||
first.append(event)
|
||||
|
||||
second = []
|
||||
async for event in run:
|
||||
second.append(event)
|
||||
|
||||
assert len(first) > 0
|
||||
assert first == second
|
||||
|
||||
|
||||
def test_sync_raw_events_double_iteration():
|
||||
"""Iterating over the raw event stream twice should yield the same events."""
|
||||
graph = make_simple_graph()
|
||||
run = StreamingHandler(graph).stream({"value": "x", "items": []})
|
||||
|
||||
first = list(run)
|
||||
second = list(run)
|
||||
|
||||
assert len(first) > 0
|
||||
assert first == second
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tool transformer via extensions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _ToolExecution:
|
||||
def __init__(self, tool_call_id: str, tool_name: str, input: Any, output: Any):
|
||||
self.tool_call_id = tool_call_id
|
||||
self.tool_name = tool_name
|
||||
self.input = input
|
||||
self.output = output
|
||||
|
||||
|
||||
class _ToolsTransformer(StreamTransformer):
|
||||
"""Groups tool-started/tool-finished custom events into _ToolExecution objects."""
|
||||
|
||||
name = "tools"
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._log: list[_ToolExecution] = []
|
||||
self._pending: dict[str, dict] = {}
|
||||
self.value = self._log
|
||||
|
||||
def init(self) -> Any:
|
||||
return None
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
if event["method"] != "custom":
|
||||
return True
|
||||
data = event["params"]["data"]
|
||||
if not isinstance(data, dict) or "event" not in data:
|
||||
return True
|
||||
|
||||
tool_call_id = data.get("tool_call_id")
|
||||
if tool_call_id is None:
|
||||
return True
|
||||
|
||||
if data["event"] == "tool-started":
|
||||
self._pending[tool_call_id] = data
|
||||
return False
|
||||
|
||||
if data["event"] == "tool-finished":
|
||||
started = self._pending.pop(tool_call_id, {})
|
||||
self._log.append(_ToolExecution(
|
||||
tool_call_id=tool_call_id,
|
||||
tool_name=started.get("tool_name", ""),
|
||||
input=started.get("input"),
|
||||
output=data["output"],
|
||||
))
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def finalize(self) -> None:
|
||||
pass
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _make_tool_graph():
|
||||
"""Graph: agent emits a tool call, custom_tools executes it with writer events."""
|
||||
from langgraph.types import StreamWriter
|
||||
|
||||
def agent(state):
|
||||
return {
|
||||
"value": "called",
|
||||
"items": ["agent"],
|
||||
}
|
||||
|
||||
def custom_tools(state, *, writer: StreamWriter):
|
||||
writer({
|
||||
"event": "tool-started",
|
||||
"tool_call_id": "call_1",
|
||||
"tool_name": "get_weather",
|
||||
"input": {"city": "SF"},
|
||||
})
|
||||
writer({
|
||||
"event": "tool-finished",
|
||||
"tool_call_id": "call_1",
|
||||
"output": {"temp_f": 64},
|
||||
})
|
||||
return {"value": "done", "items": ["tools"]}
|
||||
|
||||
graph = StateGraph(State)
|
||||
graph.add_node("agent", agent)
|
||||
graph.add_node("custom_tools", custom_tools)
|
||||
graph.add_edge(START, "agent")
|
||||
graph.add_edge("agent", "custom_tools")
|
||||
graph.add_edge("custom_tools", END)
|
||||
return graph.compile()
|
||||
|
||||
|
||||
def test_sync_tool_transformer_via_extensions():
|
||||
"""Tool events flow through extensions and are iterable without draining raw events."""
|
||||
graph = _make_tool_graph()
|
||||
run = StreamingHandler(graph).stream(
|
||||
{"value": "", "items": []},
|
||||
transformers=[_ToolsTransformer()],
|
||||
)
|
||||
|
||||
# Iterating extensions drives the pump — no need to drain raw events first
|
||||
executions = list(run.extensions["tools"])
|
||||
assert len(executions) == 1
|
||||
assert executions[0].tool_name == "get_weather"
|
||||
assert executions[0].input == {"city": "SF"}
|
||||
assert executions[0].output == {"temp_f": 64}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_tool_transformer_via_extensions():
|
||||
"""Tool events flow through extensions in async mode."""
|
||||
graph = _make_tool_graph()
|
||||
run = await StreamingHandler(graph).astream(
|
||||
{"value": "", "items": []},
|
||||
transformers=[_ToolsTransformer()],
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Drain main stream so transformer processes all events
|
||||
async for _ in run:
|
||||
pass
|
||||
|
||||
tools_log = run.extensions["tools"]
|
||||
assert len(tools_log) == 1
|
||||
assert tools_log[0].tool_name == "get_weather"
|
||||
assert tools_log[0].output == {"temp_f": 64}
|
||||
@@ -0,0 +1,531 @@
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage
|
||||
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult
|
||||
|
||||
from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
|
||||
from langgraph.pregel._messages_v2 import StreamProtocolMessagesHandler
|
||||
from langgraph.types import Command
|
||||
|
||||
META = {"langgraph_checkpoint_ns": "root:", "langgraph_node": "agent"}
|
||||
|
||||
|
||||
def make_handler(subgraphs=True):
|
||||
events = []
|
||||
handler = StreamProtocolMessagesHandler(events.append, subgraphs)
|
||||
return handler, events
|
||||
|
||||
|
||||
def test_streamed_text():
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chat_model_start(
|
||||
serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[]
|
||||
)
|
||||
|
||||
for token_text in ("Hello", " ", "world"):
|
||||
chunk = ChatGenerationChunk(
|
||||
message=AIMessageChunk(content=token_text, id=f"run-{run_id}")
|
||||
)
|
||||
handler.on_llm_new_token(token_text, chunk=chunk, run_id=run_id)
|
||||
|
||||
final_msg = AIMessage(content="Hello world", id=f"run-{run_id}")
|
||||
handler.on_llm_end(
|
||||
LLMResult(generations=[[ChatGeneration(message=final_msg)]]),
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
data_events = [e[2] for e in events]
|
||||
assert data_events[0]["event"] == "message-start"
|
||||
assert data_events[1]["event"] == "content-block-start"
|
||||
assert data_events[1]["index"] == 0
|
||||
|
||||
deltas = [d for d in data_events if d["event"] == "content-block-delta"]
|
||||
assert len(deltas) == 3
|
||||
assert deltas[0]["content_block"]["text"] == "Hello"
|
||||
assert deltas[1]["content_block"]["text"] == " "
|
||||
assert deltas[2]["content_block"]["text"] == "world"
|
||||
|
||||
finish_blocks = [d for d in data_events if d["event"] == "content-block-finish"]
|
||||
assert len(finish_blocks) == 1
|
||||
assert finish_blocks[0]["content_block"]["text"] == "Hello world"
|
||||
|
||||
assert data_events[-1]["event"] == "message-finish"
|
||||
assert data_events[-1]["reason"] == "stop"
|
||||
|
||||
|
||||
def test_tool_calls():
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chat_model_start(
|
||||
serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[]
|
||||
)
|
||||
|
||||
chunk1 = ChatGenerationChunk(
|
||||
message=AIMessageChunk(
|
||||
content="",
|
||||
tool_call_chunks=[
|
||||
{"name": "search", "args": '{"q', "id": "call_1", "index": 0}
|
||||
],
|
||||
id=f"run-{run_id}",
|
||||
)
|
||||
)
|
||||
handler.on_llm_new_token("", chunk=chunk1, run_id=run_id)
|
||||
|
||||
chunk2 = ChatGenerationChunk(
|
||||
message=AIMessageChunk(
|
||||
content="",
|
||||
tool_call_chunks=[
|
||||
{"name": None, "args": 'uery":"hi"}', "id": None, "index": 0}
|
||||
],
|
||||
id=f"run-{run_id}",
|
||||
)
|
||||
)
|
||||
handler.on_llm_new_token("", chunk=chunk2, run_id=run_id)
|
||||
|
||||
final_msg = AIMessage(
|
||||
content="",
|
||||
tool_calls=[{"name": "search", "args": {"query": "hi"}, "id": "call_1"}],
|
||||
id=f"run-{run_id}",
|
||||
)
|
||||
handler.on_llm_end(
|
||||
LLMResult(generations=[[ChatGeneration(message=final_msg)]]),
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
data_events = [e[2] for e in events]
|
||||
finish_blocks = [d for d in data_events if d["event"] == "content-block-finish"]
|
||||
assert len(finish_blocks) == 1
|
||||
fb = finish_blocks[0]["content_block"]
|
||||
assert fb["type"] == "tool_call"
|
||||
assert fb["args"] == {"query": "hi"}
|
||||
assert fb["name"] == "search"
|
||||
assert fb["id"] == "call_1"
|
||||
|
||||
|
||||
def test_invalid_tool_call_json():
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chat_model_start(
|
||||
serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[]
|
||||
)
|
||||
|
||||
chunk = ChatGenerationChunk(
|
||||
message=AIMessageChunk(
|
||||
content="",
|
||||
tool_call_chunks=[
|
||||
{
|
||||
"name": "search",
|
||||
"args": "{not valid json",
|
||||
"id": "call_2",
|
||||
"index": 0,
|
||||
}
|
||||
],
|
||||
id=f"run-{run_id}",
|
||||
)
|
||||
)
|
||||
handler.on_llm_new_token("", chunk=chunk, run_id=run_id)
|
||||
|
||||
final_msg = AIMessage(content="", id=f"run-{run_id}")
|
||||
handler.on_llm_end(
|
||||
LLMResult(generations=[[ChatGeneration(message=final_msg)]]),
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
data_events = [e[2] for e in events]
|
||||
finish_blocks = [d for d in data_events if d["event"] == "content-block-finish"]
|
||||
assert len(finish_blocks) == 1
|
||||
fb = finish_blocks[0]["content_block"]
|
||||
assert fb["type"] == "invalid_tool_call"
|
||||
assert "Failed to parse" in fb["error"]
|
||||
|
||||
|
||||
def test_reasoning_blocks():
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chat_model_start(
|
||||
serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[]
|
||||
)
|
||||
|
||||
chunk = ChatGenerationChunk(
|
||||
message=AIMessageChunk(
|
||||
content=[{"type": "reasoning_content", "reasoning_content": "thinking..."}],
|
||||
id=f"run-{run_id}",
|
||||
)
|
||||
)
|
||||
handler.on_llm_new_token("", chunk=chunk, run_id=run_id)
|
||||
|
||||
final_msg = AIMessage(content="", id=f"run-{run_id}")
|
||||
handler.on_llm_end(
|
||||
LLMResult(generations=[[ChatGeneration(message=final_msg)]]),
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
data_events = [e[2] for e in events]
|
||||
block_starts = [d for d in data_events if d["event"] == "content-block-start"]
|
||||
assert len(block_starts) == 1
|
||||
assert block_starts[0]["content_block"]["type"] == "reasoning"
|
||||
|
||||
deltas = [d for d in data_events if d["event"] == "content-block-delta"]
|
||||
assert len(deltas) == 1
|
||||
assert deltas[0]["content_block"]["reasoning"] == "thinking..."
|
||||
|
||||
|
||||
def test_multiple_content_blocks():
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chat_model_start(
|
||||
serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[]
|
||||
)
|
||||
|
||||
chunk1 = ChatGenerationChunk(
|
||||
message=AIMessageChunk(content="hello", id=f"run-{run_id}")
|
||||
)
|
||||
handler.on_llm_new_token("hello", chunk=chunk1, run_id=run_id)
|
||||
|
||||
chunk2 = ChatGenerationChunk(
|
||||
message=AIMessageChunk(
|
||||
content="",
|
||||
tool_call_chunks=[
|
||||
{"name": "lookup", "args": '{"x":1}', "id": "call_3", "index": 1}
|
||||
],
|
||||
id=f"run-{run_id}",
|
||||
)
|
||||
)
|
||||
handler.on_llm_new_token("", chunk=chunk2, run_id=run_id)
|
||||
|
||||
final_msg = AIMessage(
|
||||
content="hello",
|
||||
tool_calls=[{"name": "lookup", "args": {"x": 1}, "id": "call_3"}],
|
||||
id=f"run-{run_id}",
|
||||
)
|
||||
handler.on_llm_end(
|
||||
LLMResult(generations=[[ChatGeneration(message=final_msg)]]),
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
data_events = [e[2] for e in events]
|
||||
finish_blocks = [d for d in data_events if d["event"] == "content-block-finish"]
|
||||
assert len(finish_blocks) == 2
|
||||
|
||||
|
||||
def test_usage_metadata():
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chat_model_start(
|
||||
serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[]
|
||||
)
|
||||
|
||||
chunk = ChatGenerationChunk(
|
||||
message=AIMessageChunk(content="hi", id=f"run-{run_id}")
|
||||
)
|
||||
handler.on_llm_new_token("hi", chunk=chunk, run_id=run_id)
|
||||
|
||||
final_msg = AIMessage(
|
||||
content="hi",
|
||||
id=f"run-{run_id}",
|
||||
usage_metadata={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
||||
)
|
||||
handler.on_llm_end(
|
||||
LLMResult(generations=[[ChatGeneration(message=final_msg)]]),
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
data_events = [e[2] for e in events]
|
||||
finish_event = [d for d in data_events if d["event"] == "message-finish"][0]
|
||||
assert "usage" in finish_event
|
||||
assert finish_event["usage"]["input_tokens"] == 10
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"raw_reason,expected",
|
||||
[
|
||||
("stop", "stop"),
|
||||
("tool_calls", "tool_use"),
|
||||
("length", "length"),
|
||||
("content_filter", "content_filter"),
|
||||
("end_turn", "stop"),
|
||||
],
|
||||
)
|
||||
def test_finish_reason_normalization(raw_reason, expected):
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chat_model_start(
|
||||
serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[]
|
||||
)
|
||||
chunk = ChatGenerationChunk(message=AIMessageChunk(content="x", id=f"run-{run_id}"))
|
||||
handler.on_llm_new_token("x", chunk=chunk, run_id=run_id)
|
||||
|
||||
final_msg = AIMessage(
|
||||
content="x",
|
||||
id=f"run-{run_id}",
|
||||
response_metadata={"finish_reason": raw_reason},
|
||||
)
|
||||
handler.on_llm_end(
|
||||
LLMResult(generations=[[ChatGeneration(message=final_msg)]]),
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
data_events = [e[2] for e in events]
|
||||
finish_event = [d for d in data_events if d["event"] == "message-finish"][0]
|
||||
assert finish_event["reason"] == expected
|
||||
|
||||
|
||||
def test_tag_nostream():
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chat_model_start(
|
||||
serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[TAG_NOSTREAM]
|
||||
)
|
||||
chunk = ChatGenerationChunk(
|
||||
message=AIMessageChunk(content="secret", id=f"run-{run_id}")
|
||||
)
|
||||
handler.on_llm_new_token("secret", chunk=chunk, run_id=run_id)
|
||||
|
||||
final_msg = AIMessage(content="secret", id=f"run-{run_id}")
|
||||
handler.on_llm_end(
|
||||
LLMResult(generations=[[ChatGeneration(message=final_msg)]]),
|
||||
run_id=run_id,
|
||||
)
|
||||
assert events == []
|
||||
|
||||
|
||||
def test_tag_hidden_chain():
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chain_start(
|
||||
serialized={},
|
||||
inputs={},
|
||||
run_id=run_id,
|
||||
metadata=META,
|
||||
tags=[TAG_HIDDEN],
|
||||
name="agent",
|
||||
)
|
||||
handler.on_chain_end(
|
||||
{"messages": [AIMessage(content="hidden", id="msg-1")]},
|
||||
run_id=run_id,
|
||||
)
|
||||
assert events == []
|
||||
|
||||
|
||||
def test_subgraph_filtering():
|
||||
handler, events = make_handler(subgraphs=False)
|
||||
run_id = uuid4()
|
||||
|
||||
subgraph_meta = {
|
||||
"langgraph_checkpoint_ns": "root:|child:",
|
||||
"langgraph_node": "agent",
|
||||
}
|
||||
handler.on_chat_model_start(
|
||||
serialized={}, messages=[[]], run_id=run_id, metadata=subgraph_meta, tags=[]
|
||||
)
|
||||
chunk = ChatGenerationChunk(
|
||||
message=AIMessageChunk(content="sub", id=f"run-{run_id}")
|
||||
)
|
||||
handler.on_llm_new_token("sub", chunk=chunk, run_id=run_id)
|
||||
|
||||
final_msg = AIMessage(content="sub", id=f"run-{run_id}")
|
||||
handler.on_llm_end(
|
||||
LLMResult(generations=[[ChatGeneration(message=final_msg)]]),
|
||||
run_id=run_id,
|
||||
)
|
||||
assert events == []
|
||||
|
||||
|
||||
def test_chain_emits_messages():
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chain_start(
|
||||
serialized={}, inputs={}, run_id=run_id, metadata=META, tags=[], name="agent"
|
||||
)
|
||||
handler.on_chain_end(
|
||||
{"messages": [AIMessage(content="hello", id="msg-chain-1")]},
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
data_events = [e[2] for e in events]
|
||||
assert len(data_events) > 0
|
||||
assert data_events[0]["event"] == "message-start"
|
||||
assert data_events[-1]["event"] == "message-finish"
|
||||
|
||||
|
||||
def test_llm_error_after_start():
|
||||
"""on_llm_error should emit a message-error event for a started stream."""
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chat_model_start(
|
||||
serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[]
|
||||
)
|
||||
|
||||
chunk = ChatGenerationChunk(
|
||||
message=AIMessageChunk(content="partial", id=f"run-{run_id}")
|
||||
)
|
||||
handler.on_llm_new_token("partial", chunk=chunk, run_id=run_id)
|
||||
|
||||
handler.on_llm_error(RuntimeError("connection lost"), run_id=run_id)
|
||||
|
||||
data_events = [e[2] for e in events]
|
||||
assert data_events[0]["event"] == "message-start"
|
||||
error_events = [d for d in data_events if d["event"] == "error"]
|
||||
assert len(error_events) == 1
|
||||
assert "connection lost" in error_events[0]["message"]
|
||||
|
||||
|
||||
def test_llm_error_before_start_no_emit():
|
||||
"""on_llm_error before any tokens should not emit error events."""
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chat_model_start(
|
||||
serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[]
|
||||
)
|
||||
|
||||
# Error before any token — state.started is False
|
||||
handler.on_llm_error(RuntimeError("immediate fail"), run_id=run_id)
|
||||
|
||||
data_events = [e[2] for e in events]
|
||||
error_events = [d for d in data_events if d.get("event") == "error"]
|
||||
assert len(error_events) == 0
|
||||
|
||||
|
||||
def test_non_streamed_model():
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chat_model_start(
|
||||
serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[]
|
||||
)
|
||||
|
||||
final_msg = AIMessage(
|
||||
content="full response",
|
||||
id=f"run-{run_id}",
|
||||
response_metadata={"finish_reason": "stop"},
|
||||
)
|
||||
handler.on_llm_end(
|
||||
LLMResult(generations=[[ChatGeneration(message=final_msg)]]),
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
data_events = [e[2] for e in events]
|
||||
assert len(data_events) > 0
|
||||
assert data_events[0]["event"] == "message-start"
|
||||
|
||||
deltas = [d for d in data_events if d["event"] == "content-block-delta"]
|
||||
assert len(deltas) == 1
|
||||
assert deltas[0]["content_block"]["text"] == "full response"
|
||||
|
||||
assert data_events[-1]["event"] == "message-finish"
|
||||
assert data_events[-1]["reason"] == "stop"
|
||||
|
||||
|
||||
def test_chain_emits_command_with_message():
|
||||
"""on_chain_end should emit protocol events for messages inside a Command."""
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chain_start(
|
||||
serialized={}, inputs={}, run_id=run_id, metadata=META, tags=[], name="agent"
|
||||
)
|
||||
handler.on_chain_end(
|
||||
Command(update={"messages": [AIMessage(content="from command", id="cmd-1")]}),
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
data_events = [e[2] for e in events]
|
||||
assert len(data_events) > 0
|
||||
assert data_events[0]["event"] == "message-start"
|
||||
deltas = [d for d in data_events if d["event"] == "content-block-delta"]
|
||||
assert len(deltas) == 1
|
||||
assert deltas[0]["content_block"]["text"] == "from command"
|
||||
assert data_events[-1]["event"] == "message-finish"
|
||||
|
||||
|
||||
def test_chain_emits_command_in_list():
|
||||
"""on_chain_end should handle a list containing Command objects."""
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chain_start(
|
||||
serialized={}, inputs={}, run_id=run_id, metadata=META, tags=[], name="agent"
|
||||
)
|
||||
handler.on_chain_end(
|
||||
[Command(update={"messages": [AIMessage(content="listed", id="cmd-2")]})],
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
data_events = [e[2] for e in events]
|
||||
starts = [d for d in data_events if d["event"] == "message-start"]
|
||||
assert len(starts) == 1
|
||||
|
||||
|
||||
def test_chain_deduplicates_seen_messages():
|
||||
"""Messages already seen from LLM streaming should not be re-emitted by chain end."""
|
||||
handler, events = make_handler()
|
||||
run_id_llm = uuid4()
|
||||
run_id_chain = uuid4()
|
||||
msg_id = f"run-{run_id_llm}"
|
||||
|
||||
# Simulate LLM streaming
|
||||
handler.on_chat_model_start(
|
||||
serialized={}, messages=[[]], run_id=run_id_llm, metadata=META, tags=[]
|
||||
)
|
||||
chunk = ChatGenerationChunk(message=AIMessageChunk(content="hello", id=msg_id))
|
||||
handler.on_llm_new_token("hello", chunk=chunk, run_id=run_id_llm)
|
||||
|
||||
final_msg = AIMessage(content="hello", id=msg_id)
|
||||
handler.on_llm_end(
|
||||
LLMResult(generations=[[ChatGeneration(message=final_msg)]]),
|
||||
run_id=run_id_llm,
|
||||
)
|
||||
|
||||
events_before = len(events)
|
||||
|
||||
# Now chain end with the same message ID
|
||||
handler.on_chain_start(
|
||||
serialized={},
|
||||
inputs={},
|
||||
run_id=run_id_chain,
|
||||
metadata=META,
|
||||
tags=[],
|
||||
name="agent",
|
||||
)
|
||||
handler.on_chain_end(
|
||||
{"messages": [AIMessage(content="hello", id=msg_id)]},
|
||||
run_id=run_id_chain,
|
||||
)
|
||||
|
||||
# No new events should have been emitted for the duplicate
|
||||
data_events_after = [e[2] for e in events[events_before:]]
|
||||
starts = [d for d in data_events_after if d.get("event") == "message-start"]
|
||||
assert len(starts) == 0
|
||||
|
||||
|
||||
def test_chain_emits_human_message_role():
|
||||
"""Non-AI messages from chain output should have the correct role."""
|
||||
handler, events = make_handler()
|
||||
run_id = uuid4()
|
||||
|
||||
handler.on_chain_start(
|
||||
serialized={}, inputs={}, run_id=run_id, metadata=META, tags=[], name="agent"
|
||||
)
|
||||
handler.on_chain_end(
|
||||
{"messages": [HumanMessage(content="user msg", id="hmsg-1")]},
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
data_events = [e[2] for e in events]
|
||||
starts = [d for d in data_events if d["event"] == "message-start"]
|
||||
assert len(starts) == 1
|
||||
assert starts[0]["role"] == "human"
|
||||
@@ -0,0 +1,245 @@
|
||||
import pytest
|
||||
|
||||
from langgraph.stream.chat_model_stream import AsyncChatModelStream, ChatModelStream
|
||||
|
||||
|
||||
def _text_delta(text: str) -> dict:
|
||||
return {"content_block": {"type": "text", "text": text}}
|
||||
|
||||
|
||||
def _reasoning_delta(text: str) -> dict:
|
||||
return {"content_block": {"type": "reasoning", "reasoning": text}}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sync ChatModelStream tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_text_accumulates():
|
||||
stream = ChatModelStream()
|
||||
stream._push_content_block_delta(_text_delta("Hello"))
|
||||
stream._push_content_block_delta(_text_delta(", world"))
|
||||
stream._finish({"reason": "stop"})
|
||||
|
||||
assert stream.text == "Hello, world"
|
||||
assert isinstance(stream.text, str)
|
||||
|
||||
|
||||
def test_sync_reasoning_accumulates():
|
||||
stream = ChatModelStream()
|
||||
stream._push_content_block_delta(_reasoning_delta("step 1"))
|
||||
stream._push_content_block_delta(_reasoning_delta(" -> step 2"))
|
||||
stream._finish({"reason": "stop"})
|
||||
|
||||
assert stream.reasoning == "step 1 -> step 2"
|
||||
assert isinstance(stream.reasoning, str)
|
||||
|
||||
|
||||
def test_sync_usage():
|
||||
stream = ChatModelStream()
|
||||
usage = {"input_tokens": 10, "output_tokens": 5}
|
||||
stream._finish({"reason": "stop", "usage": usage})
|
||||
assert stream.usage == usage
|
||||
|
||||
|
||||
def test_sync_mixed_blocks():
|
||||
stream = ChatModelStream()
|
||||
stream._push_content_block_delta(_text_delta("answer"))
|
||||
stream._push_content_block_delta(
|
||||
{"content_block": {"type": "tool_call", "name": "search"}}
|
||||
)
|
||||
stream._push_content_block_delta(_text_delta(" here"))
|
||||
stream._finish({"reason": "stop"})
|
||||
|
||||
assert stream.text == "answer here"
|
||||
|
||||
|
||||
def test_sync_tool_call_only_text_empty():
|
||||
stream = ChatModelStream()
|
||||
stream._push_content_block_delta(
|
||||
{"content_block": {"type": "tool_call", "name": "search"}}
|
||||
)
|
||||
stream._finish({"reason": "stop"})
|
||||
assert stream.text == ""
|
||||
|
||||
|
||||
def test_sync_fail_marks_done():
|
||||
stream = ChatModelStream()
|
||||
assert not stream.done
|
||||
stream._fail(RuntimeError("err"))
|
||||
assert stream.done
|
||||
|
||||
|
||||
def test_sync_namespace_and_node():
|
||||
stream = ChatModelStream(
|
||||
namespace=["agent:0", "tools:1"],
|
||||
node="chat_model",
|
||||
message_id="msg-123",
|
||||
)
|
||||
assert stream.namespace == ["agent:0", "tools:1"]
|
||||
assert stream.node == "chat_model"
|
||||
assert stream.message_id == "msg-123"
|
||||
|
||||
|
||||
def test_sync_content_block_finish_authoritative():
|
||||
"""content-block-finish with authoritative text overrides accumulated."""
|
||||
stream = ChatModelStream()
|
||||
stream._push_content_block_delta(_text_delta("partial"))
|
||||
stream._push_content_block_finish(
|
||||
{"content_block": {"type": "text", "text": "full text"}}
|
||||
)
|
||||
assert stream.text == "full text"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Async ChatModelStream tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_text_iterable_yields_deltas():
|
||||
stream = AsyncChatModelStream()
|
||||
stream._push_content_block_delta(_text_delta("Hello"))
|
||||
stream._push_content_block_delta(_text_delta(", world"))
|
||||
stream._finish({"reason": "stop"})
|
||||
|
||||
collected = []
|
||||
async for delta in stream.text:
|
||||
collected.append(delta)
|
||||
assert collected == ["Hello", ", world"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_text_awaitable_returns_full():
|
||||
stream = AsyncChatModelStream()
|
||||
stream._push_content_block_delta(_text_delta("Hello"))
|
||||
stream._push_content_block_delta(_text_delta(", world"))
|
||||
stream._finish({"reason": "stop"})
|
||||
|
||||
result = await stream.text
|
||||
assert result == "Hello, world"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_reasoning_dual_pattern():
|
||||
stream = AsyncChatModelStream()
|
||||
stream._push_content_block_delta(_reasoning_delta("step 1"))
|
||||
stream._push_content_block_delta(_reasoning_delta(" -> step 2"))
|
||||
stream._finish({"reason": "stop"})
|
||||
|
||||
collected = []
|
||||
async for delta in stream.reasoning:
|
||||
collected.append(delta)
|
||||
assert collected == ["step 1", " -> step 2"]
|
||||
|
||||
stream2 = AsyncChatModelStream()
|
||||
stream2._push_content_block_delta(_reasoning_delta("thinking"))
|
||||
stream2._finish({"reason": "stop"})
|
||||
full = await stream2.reasoning
|
||||
assert full == "thinking"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_usage_resolves():
|
||||
stream = AsyncChatModelStream()
|
||||
usage = {"input_tokens": 10, "output_tokens": 5}
|
||||
stream._finish({"reason": "stop", "usage": usage})
|
||||
result = await stream.usage
|
||||
assert result == usage
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_mixed_blocks_text_only():
|
||||
stream = AsyncChatModelStream()
|
||||
stream._push_content_block_delta(_text_delta("answer"))
|
||||
stream._push_content_block_delta(
|
||||
{"content_block": {"type": "tool_call", "name": "search"}}
|
||||
)
|
||||
stream._push_content_block_delta(_text_delta(" here"))
|
||||
stream._finish({"reason": "stop"})
|
||||
|
||||
collected = []
|
||||
async for delta in stream.text:
|
||||
collected.append(delta)
|
||||
assert collected == ["answer", " here"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_tool_call_only_text_empty():
|
||||
stream = AsyncChatModelStream()
|
||||
stream._push_content_block_delta(
|
||||
{"content_block": {"type": "tool_call", "name": "search"}}
|
||||
)
|
||||
stream._finish({"reason": "stop"})
|
||||
result = await stream.text
|
||||
assert result == ""
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_fail_raises_on_text_await():
|
||||
stream = AsyncChatModelStream()
|
||||
stream._push_content_block_delta(_text_delta("partial"))
|
||||
stream._fail(RuntimeError("model error"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="model error"):
|
||||
await stream.text
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_fail_raises_on_reasoning_await():
|
||||
stream = AsyncChatModelStream()
|
||||
stream._push_content_block_delta(_reasoning_delta("thinking"))
|
||||
stream._fail(RuntimeError("model error"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="model error"):
|
||||
await stream.reasoning
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_fail_raises_on_usage_await():
|
||||
stream = AsyncChatModelStream()
|
||||
stream._fail(RuntimeError("model error"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="model error"):
|
||||
await stream.usage
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_fail_raises_during_text_iteration():
|
||||
stream = AsyncChatModelStream()
|
||||
stream._push_content_block_delta(_text_delta("partial"))
|
||||
stream._fail(RuntimeError("model error"))
|
||||
|
||||
collected = []
|
||||
with pytest.raises(RuntimeError, match="model error"):
|
||||
async for delta in stream.text:
|
||||
collected.append(delta)
|
||||
assert collected == ["partial"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_fail_marks_done():
|
||||
stream = AsyncChatModelStream()
|
||||
assert not stream.done
|
||||
stream._fail(RuntimeError("err"))
|
||||
assert stream.done
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_namespace_and_node():
|
||||
stream = AsyncChatModelStream(
|
||||
namespace=["agent:0", "tools:1"],
|
||||
node="chat_model",
|
||||
message_id="msg-123",
|
||||
)
|
||||
assert stream.namespace == ["agent:0", "tools:1"]
|
||||
assert stream.node == "chat_model"
|
||||
assert stream.message_id == "msg-123"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_inherits_from_sync():
|
||||
"""AsyncChatModelStream is a subclass of ChatModelStream."""
|
||||
stream = AsyncChatModelStream()
|
||||
assert isinstance(stream, ChatModelStream)
|
||||
@@ -0,0 +1,79 @@
|
||||
from langgraph.stream._convert import STREAM_V2_MODES, convert_to_protocol_event
|
||||
|
||||
|
||||
def test_values_mode():
|
||||
evt = convert_to_protocol_event((), "values", {"x": 1})
|
||||
assert evt is not None
|
||||
assert evt["method"] == "values"
|
||||
assert evt["params"]["data"] == {"x": 1}
|
||||
|
||||
|
||||
def test_updates_mode():
|
||||
evt = convert_to_protocol_event((), "updates", {"node": "out"})
|
||||
assert evt is not None
|
||||
assert evt["method"] == "updates"
|
||||
|
||||
|
||||
def test_messages_mode():
|
||||
evt = convert_to_protocol_event((), "messages", {"event": "msg"})
|
||||
assert evt is not None
|
||||
assert evt["method"] == "messages"
|
||||
|
||||
|
||||
def test_custom_mode():
|
||||
evt = convert_to_protocol_event((), "custom", "hello")
|
||||
assert evt is not None
|
||||
assert evt["method"] == "custom"
|
||||
assert evt["params"]["data"] == "hello"
|
||||
|
||||
|
||||
def test_debug_mode():
|
||||
evt = convert_to_protocol_event((), "debug", {})
|
||||
assert evt is not None
|
||||
assert evt["method"] == "debug"
|
||||
|
||||
|
||||
def test_checkpoints_mode():
|
||||
evt = convert_to_protocol_event((), "checkpoints", {})
|
||||
assert evt is not None
|
||||
assert evt["method"] == "checkpoints"
|
||||
|
||||
|
||||
def test_tasks_mode():
|
||||
evt = convert_to_protocol_event((), "tasks", {})
|
||||
assert evt is not None
|
||||
assert evt["method"] == "tasks"
|
||||
|
||||
|
||||
def test_namespace_passthrough():
|
||||
evt = convert_to_protocol_event(("agent", "0"), "values", {})
|
||||
assert evt is not None
|
||||
assert evt["params"]["namespace"] == ["agent", "0"]
|
||||
|
||||
|
||||
def test_unknown_mode_returns_none():
|
||||
assert convert_to_protocol_event((), "unknown_mode", {}) is None
|
||||
|
||||
|
||||
def test_node_parameter():
|
||||
evt = convert_to_protocol_event((), "values", {}, node="agent")
|
||||
assert evt is not None
|
||||
assert evt["params"]["node"] == "agent"
|
||||
|
||||
|
||||
def test_type_is_event():
|
||||
evt = convert_to_protocol_event((), "values", {})
|
||||
assert evt is not None
|
||||
assert evt["type"] == "event"
|
||||
|
||||
|
||||
def test_stream_v2_modes_complete():
|
||||
assert set(STREAM_V2_MODES) == {
|
||||
"values",
|
||||
"updates",
|
||||
"messages",
|
||||
"custom",
|
||||
"checkpoints",
|
||||
"tasks",
|
||||
"debug",
|
||||
}
|
||||
@@ -0,0 +1,293 @@
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from langgraph.stream._convert import convert_to_protocol_event
|
||||
from langgraph.stream._mux import AsyncStreamMux, StreamMux
|
||||
from langgraph.stream._types import ProtocolEvent, StreamTransformer
|
||||
from langgraph.stream.stream_channel import StreamChannel
|
||||
|
||||
|
||||
def _event(mode: str, data: Any, ns: list[str] | None = None) -> ProtocolEvent:
|
||||
ev = convert_to_protocol_event(tuple(ns or []), mode, data)
|
||||
assert ev is not None
|
||||
return ev
|
||||
|
||||
|
||||
class _MockTransformer(StreamTransformer):
|
||||
def __init__(self, *, suppress: bool = False):
|
||||
self.calls: list[ProtocolEvent] = []
|
||||
self._suppress = suppress
|
||||
|
||||
def init(self) -> Any:
|
||||
return None
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
self.calls.append(event)
|
||||
return not self._suppress
|
||||
|
||||
def finalize(self) -> None:
|
||||
pass
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_events_through_reducer_pipeline():
|
||||
reducer = _MockTransformer()
|
||||
mux = StreamMux(transformers=[reducer])
|
||||
event = _event("values", {"key": "val"})
|
||||
mux.push(event)
|
||||
assert len(reducer.calls) == 1
|
||||
assert reducer.calls[0] is event
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_reducer_suppresses_event():
|
||||
reducer = _MockTransformer(suppress=True)
|
||||
mux = StreamMux(transformers=[reducer])
|
||||
mux.push(_event("values", {"x": 1}))
|
||||
mux.close()
|
||||
assert len(reducer.calls) == 1
|
||||
assert len(mux.event_log) == 0
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_namespace_discovery():
|
||||
mux = StreamMux()
|
||||
mux.push(_event("values", {"a": 1}, ns=["child:0"]))
|
||||
assert "child:0" in mux._discovered_ns
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_top_level_ns_only():
|
||||
mux = StreamMux()
|
||||
mux.push(_event("values", {"a": 1}, ns=["agent:0", "tools:1"]))
|
||||
assert "agent:0" in mux._discovered_ns
|
||||
assert "tools:1" not in mux._discovered_ns
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_subscribe_events_filter():
|
||||
mux = AsyncStreamMux()
|
||||
mux.push(_event("values", {"a": 1}, ns=["child:0"]))
|
||||
mux.push(_event("values", {"b": 2}, ns=["other:1"]))
|
||||
mux.push(_event("values", {"c": 3}, ns=["child:0"]))
|
||||
mux.close()
|
||||
|
||||
collected = []
|
||||
async for ev in mux.subscribe_events(["child:0"]):
|
||||
collected.append(ev)
|
||||
assert len(collected) == 2
|
||||
assert collected[0]["params"]["data"] == {"a": 1}
|
||||
assert collected[1]["params"]["data"] == {"c": 3}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_close_resolves_output():
|
||||
mux = AsyncStreamMux()
|
||||
fut = mux.get_output_future()
|
||||
mux.push(_event("values", {"v": 1}))
|
||||
mux.push(_event("values", {"v": 2}))
|
||||
mux.close()
|
||||
result = await fut
|
||||
assert result == {"v": 2}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_fail_rejects_output():
|
||||
mux = AsyncStreamMux()
|
||||
fut = mux.get_output_future()
|
||||
mux.fail(ValueError("boom"))
|
||||
with pytest.raises(ValueError, match="boom"):
|
||||
await fut
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_latest_values_tracked():
|
||||
mux = StreamMux()
|
||||
mux.push(_event("values", {"v": 1}, ns=["child:0"]))
|
||||
mux.push(_event("values", {"v": 2}, ns=["child:0"]))
|
||||
assert mux.get_latest_values(["child:0"]) == {"v": 2}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_interrupt_tracking():
|
||||
"""StreamMux should track __interrupt__ payloads in values events."""
|
||||
|
||||
class _FakeInterrupt:
|
||||
def __init__(self, id: str, payload: Any):
|
||||
self.id = id
|
||||
self.payload = payload
|
||||
|
||||
mux = StreamMux()
|
||||
interrupt_obj = _FakeInterrupt("int-1", "what do you want?")
|
||||
mux.push(
|
||||
_event(
|
||||
"values",
|
||||
{"__interrupt__": [interrupt_obj]},
|
||||
)
|
||||
)
|
||||
assert mux.interrupted is True
|
||||
assert len(mux.interrupts) == 1
|
||||
assert mux.interrupts[0]["interrupt_id"] == "int-1"
|
||||
assert mux.interrupts[0]["payload"] is interrupt_obj
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_no_interrupt_by_default():
|
||||
mux = StreamMux()
|
||||
mux.push(_event("values", {"x": 1}))
|
||||
mux.close()
|
||||
assert mux.interrupted is False
|
||||
assert mux.interrupts == []
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_push_after_close_ignored():
|
||||
mux = StreamMux()
|
||||
mux.push(_event("values", {"a": 1}))
|
||||
mux.close()
|
||||
mux.push(_event("values", {"b": 2}))
|
||||
assert len(mux.event_log) == 1
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_fail_rejects_all_futures():
|
||||
mux = AsyncStreamMux()
|
||||
fut1 = mux.get_output_future([])
|
||||
fut2 = mux.get_output_future(["child:0"])
|
||||
mux.fail(ValueError("boom"))
|
||||
with pytest.raises(ValueError, match="boom"):
|
||||
await fut1
|
||||
with pytest.raises(ValueError, match="boom"):
|
||||
await fut2
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_channel_events_bypass_transformer_pipeline():
|
||||
"""Events emitted via ``StreamChannel.push()`` are appended directly
|
||||
to the event log, bypassing the transformer pipeline. This matches
|
||||
the JS implementation and avoids re-entrancy bugs.
|
||||
"""
|
||||
mock = _MockTransformer()
|
||||
mux = AsyncStreamMux(transformers=[mock])
|
||||
|
||||
channel: StreamChannel[str] = StreamChannel("my_channel")
|
||||
mux.wire_channels({"ch": channel})
|
||||
|
||||
# Regular push — transformer sees it
|
||||
mux.push(_event("values", {"a": 1}))
|
||||
assert len(mock.calls) == 1
|
||||
|
||||
# Channel push — bypasses transformers, goes straight to event log
|
||||
channel.push("hello from channel")
|
||||
|
||||
assert len(mock.calls) == 1, (
|
||||
f"Transformer saw {len(mock.calls)} events (expected 1). "
|
||||
"Channel events should bypass the transformer pipeline."
|
||||
)
|
||||
|
||||
# But the event IS in the log
|
||||
mux.close()
|
||||
events = []
|
||||
async for ev in mux.subscribe_events():
|
||||
events.append(ev)
|
||||
assert len(events) == 2
|
||||
assert events[1]["method"] == "my_channel"
|
||||
assert events[1]["params"]["data"] == "hello from channel"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_event_log_has_monotonic_seq_numbers():
|
||||
"""All events in the event log should have strictly monotonically
|
||||
increasing seq numbers so consumers can reason about ordering.
|
||||
|
||||
Events from ``mux.push()`` carry seq numbers assigned by the pump
|
||||
while channel-emitted events use a separate counter
|
||||
(``_next_emit_seq``). When interleaved, seq numbers can duplicate.
|
||||
"""
|
||||
mux = AsyncStreamMux()
|
||||
channel: StreamChannel[str] = StreamChannel("test_ch")
|
||||
mux.wire_channels({"ch": channel})
|
||||
|
||||
mux.push(_event("values", {"a": 1})) # log seq: 0
|
||||
channel.push("from_channel") # log seq: 0 (from _next_emit_seq)
|
||||
mux.push(_event("values", {"b": 2})) # log seq: 1
|
||||
mux.close()
|
||||
|
||||
seqs: list[int] = []
|
||||
async for event in mux.subscribe_events():
|
||||
seqs.append(event["seq"])
|
||||
|
||||
assert len(seqs) == 3, f"Expected 3 events but got {len(seqs)}"
|
||||
|
||||
for i in range(1, len(seqs)):
|
||||
assert seqs[i] > seqs[i - 1], (
|
||||
f"Seq numbers not strictly monotonic: {seqs}. "
|
||||
f"seq[{i}]={seqs[i]} <= seq[{i - 1}]={seqs[i - 1]}. "
|
||||
"Channel events use a separate counter from push() events."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_channel_push_during_process_preserves_namespace():
|
||||
"""When two transformers both call channel.push() during the same
|
||||
outer mux.push(), the second transformer's channel event should
|
||||
still carry the original event's namespace.
|
||||
|
||||
Bug: the first channel.push() re-enters mux.push(), which resets
|
||||
``_current_namespace`` to ``[]`` on exit. The second transformer's
|
||||
channel.push() then reads the clobbered value and its event gets
|
||||
``namespace: []`` instead of the original.
|
||||
"""
|
||||
|
||||
class _ChannelTransformer(StreamTransformer):
|
||||
"""Pushes to its channel whenever it sees a ``values`` event."""
|
||||
|
||||
def __init__(self, name: str) -> None:
|
||||
self.name = name
|
||||
self.channel: StreamChannel[str] = StreamChannel(name)
|
||||
|
||||
def init(self) -> Any:
|
||||
return {self.name: self.channel}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
if event["method"] == "values":
|
||||
self.channel.push(f"from_{self.name}")
|
||||
return True
|
||||
|
||||
def finalize(self) -> None:
|
||||
pass
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
pass
|
||||
|
||||
t1 = _ChannelTransformer("first")
|
||||
t2 = _ChannelTransformer("second")
|
||||
mux = AsyncStreamMux(transformers=[t1, t2])
|
||||
mux.wire_channels({"first": t1.channel})
|
||||
mux.wire_channels({"second": t2.channel})
|
||||
|
||||
# Push a values event with a non-root namespace
|
||||
mux.push(_event("values", {"x": 1}, ns=["agent:0"]))
|
||||
mux.close()
|
||||
|
||||
# Collect channel events emitted by each transformer
|
||||
channel_events: list[ProtocolEvent] = []
|
||||
async for ev in mux.subscribe_events():
|
||||
if ev["method"] in ("first", "second"):
|
||||
channel_events.append(ev)
|
||||
|
||||
assert len(channel_events) == 2, (
|
||||
f"Expected 2 channel events but got {len(channel_events)}"
|
||||
)
|
||||
|
||||
for ev in channel_events:
|
||||
assert ev["params"]["namespace"] == ["agent:0"], (
|
||||
f"Channel event for method={ev['method']!r} has "
|
||||
f"namespace={ev['params']['namespace']!r}, expected ['agent:0']. "
|
||||
"The nested mux.push() from the first channel.push() clobbered "
|
||||
"_current_namespace before the second transformer ran."
|
||||
)
|
||||
@@ -0,0 +1,161 @@
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from langgraph.stream._convert import convert_to_protocol_event
|
||||
from langgraph.stream._types import ProtocolEvent
|
||||
from langgraph.stream.chat_model_stream import ChatModelStream
|
||||
from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer
|
||||
|
||||
|
||||
def _event(
|
||||
mode: str,
|
||||
data: Any,
|
||||
ns: list[str] | None = None,
|
||||
node: str | None = None,
|
||||
) -> ProtocolEvent:
|
||||
ev = convert_to_protocol_event(tuple(ns or []), mode, data, node=node)
|
||||
assert ev is not None
|
||||
return ev
|
||||
|
||||
|
||||
# -- ValuesTransformer ---------------------------------------------------------
|
||||
|
||||
|
||||
def test_values_captures_values_events():
|
||||
reducer = ValuesTransformer()
|
||||
reducer.init()
|
||||
reducer.process(_event("values", {"a": 1}))
|
||||
reducer.process(_event("values", {"b": 2}))
|
||||
reducer.finalize()
|
||||
|
||||
assert len(reducer.values_log) == 2
|
||||
assert reducer.values_log[0]["data"] == {"a": 1}
|
||||
assert reducer.values_log[1]["data"] == {"b": 2}
|
||||
|
||||
|
||||
def test_values_ignores_other_modes():
|
||||
reducer = ValuesTransformer()
|
||||
reducer.init()
|
||||
reducer.process(_event("updates", {"x": 1}))
|
||||
reducer.process(_event("messages", {"event": "message-start"}))
|
||||
reducer.finalize()
|
||||
|
||||
assert len(reducer.values_log) == 0
|
||||
|
||||
|
||||
def test_values_latest_per_namespace():
|
||||
reducer = ValuesTransformer()
|
||||
reducer.init()
|
||||
reducer.process(_event("values", {"v": 1}, ns=["child:0"]))
|
||||
reducer.process(_event("values", {"v": 2}, ns=["child:0"]))
|
||||
assert reducer.get_latest("child:0") == {"v": 2}
|
||||
|
||||
|
||||
# -- MessagesTransformer -------------------------------------------------------
|
||||
|
||||
|
||||
def _msg_start(ns=None, node=None, message_id="msg-1"):
|
||||
return _event(
|
||||
"messages",
|
||||
{"event": "message-start", "message_id": message_id},
|
||||
ns=ns,
|
||||
node=node,
|
||||
)
|
||||
|
||||
|
||||
def _content_delta(text, ns=None, node=None):
|
||||
return _event(
|
||||
"messages",
|
||||
{
|
||||
"event": "content-block-delta",
|
||||
"content_block": {"type": "text", "text": text},
|
||||
},
|
||||
ns=ns,
|
||||
node=node,
|
||||
)
|
||||
|
||||
|
||||
def _msg_finish(ns=None, node=None):
|
||||
return _event(
|
||||
"messages",
|
||||
{"event": "message-finish", "reason": "stop"},
|
||||
ns=ns,
|
||||
node=node,
|
||||
)
|
||||
|
||||
|
||||
def test_messages_groups_lifecycle():
|
||||
reducer = MessagesTransformer()
|
||||
reducer.init()
|
||||
reducer.process(_msg_start())
|
||||
reducer.process(_content_delta("hi"))
|
||||
reducer.process(_msg_finish())
|
||||
reducer.finalize()
|
||||
|
||||
assert len(reducer.messages_log) == 1
|
||||
assert isinstance(reducer.messages_log[0], ChatModelStream)
|
||||
assert reducer.messages_log[0].done
|
||||
|
||||
|
||||
def test_messages_multiple_sequential():
|
||||
reducer = MessagesTransformer()
|
||||
reducer.init()
|
||||
reducer.process(_msg_start(message_id="m1"))
|
||||
reducer.process(_msg_finish())
|
||||
reducer.process(_msg_start(message_id="m2"))
|
||||
reducer.process(_msg_finish())
|
||||
reducer.finalize()
|
||||
|
||||
assert len(reducer.messages_log) == 2
|
||||
|
||||
|
||||
def test_messages_namespace_filter():
|
||||
reducer = MessagesTransformer(namespace=["root"])
|
||||
reducer.init()
|
||||
reducer.process(_msg_start(ns=["root"]))
|
||||
reducer.process(_msg_finish(ns=["root"]))
|
||||
reducer.process(_msg_start(ns=["other"], message_id="m2"))
|
||||
reducer.process(_msg_finish(ns=["other"]))
|
||||
reducer.finalize()
|
||||
|
||||
assert len(reducer.messages_log) == 1
|
||||
|
||||
|
||||
def test_messages_node_filter():
|
||||
reducer = MessagesTransformer(node_filter="agent")
|
||||
reducer.init()
|
||||
reducer.process(_msg_start(node="agent"))
|
||||
reducer.process(_msg_finish(node="agent"))
|
||||
reducer.process(_msg_start(node="tools", message_id="m2"))
|
||||
reducer.process(_msg_finish(node="tools"))
|
||||
reducer.finalize()
|
||||
|
||||
assert len(reducer.messages_log) == 1
|
||||
|
||||
|
||||
def test_messages_error_event():
|
||||
"""An error event should fail the active ChatModelStream."""
|
||||
reducer = MessagesTransformer()
|
||||
reducer.init()
|
||||
reducer.process(_msg_start())
|
||||
reducer.process(_content_delta("partial"))
|
||||
reducer.process(
|
||||
_event("messages", {"event": "error", "message": "connection lost"}),
|
||||
)
|
||||
reducer.finalize()
|
||||
|
||||
assert len(reducer.messages_log) == 1
|
||||
assert reducer.messages_log[0].done
|
||||
|
||||
|
||||
def test_messages_fail_propagates_to_active():
|
||||
"""transformer.fail() should mark active streams as done."""
|
||||
reducer = MessagesTransformer()
|
||||
reducer.init()
|
||||
reducer.process(_msg_start())
|
||||
reducer.process(_content_delta("partial"))
|
||||
reducer.fail(RuntimeError("graph failed"))
|
||||
|
||||
assert len(reducer.messages_log) == 1
|
||||
assert reducer.messages_log[0].done
|
||||
@@ -0,0 +1,610 @@
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from langgraph.stream._mux import AsyncStreamMux, StreamMux
|
||||
from langgraph.stream._types import ProtocolEvent
|
||||
from langgraph.stream.chat_model_stream import ChatModelStream
|
||||
from langgraph.stream.run_stream import (
|
||||
AsyncGraphRunStream,
|
||||
AsyncSubgraphRunStream,
|
||||
create_async_graph_run_stream,
|
||||
create_graph_run_stream,
|
||||
)
|
||||
from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer
|
||||
|
||||
|
||||
async def _mock_source(
|
||||
chunks: list[tuple[tuple[str, ...], str, Any]],
|
||||
) -> AsyncIterator[tuple[tuple[str, ...], str, Any]]:
|
||||
for chunk in chunks:
|
||||
yield chunk
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_aiter_yields_all_events():
|
||||
chunks = [
|
||||
((), "values", {"step": 1}),
|
||||
((), "values", {"step": 2}),
|
||||
((), "updates", {"node": "a"}),
|
||||
]
|
||||
run = await create_async_graph_run_stream(_mock_source(chunks))
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
collected: list[ProtocolEvent] = []
|
||||
async for event in run:
|
||||
collected.append(event)
|
||||
assert len(collected) == 3
|
||||
assert collected[0]["method"] == "values"
|
||||
assert collected[2]["method"] == "updates"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_subgraph_name_and_index():
|
||||
vr, mr = ValuesTransformer(), MessagesTransformer()
|
||||
mux = AsyncStreamMux(transformers=[vr, mr])
|
||||
sub = AsyncSubgraphRunStream(
|
||||
mux=mux,
|
||||
namespace=["researcher:2"],
|
||||
transformers=[vr, mr],
|
||||
)
|
||||
assert sub.name == "researcher"
|
||||
assert sub.index == 2
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_subgraph_name_no_index():
|
||||
vr, mr = ValuesTransformer(), MessagesTransformer()
|
||||
mux = AsyncStreamMux(transformers=[vr, mr])
|
||||
sub = AsyncSubgraphRunStream(
|
||||
mux=mux, namespace=["agent"], transformers=[vr, mr]
|
||||
)
|
||||
assert sub.name == "agent"
|
||||
assert sub.index == 0
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_values_iterable():
|
||||
chunks = [
|
||||
((), "values", {"v": 1}),
|
||||
((), "values", {"v": 2}),
|
||||
]
|
||||
run = await create_async_graph_run_stream(_mock_source(chunks))
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
collected = []
|
||||
async for v in run.values:
|
||||
collected.append(v)
|
||||
assert len(collected) == 2
|
||||
assert collected[0] == {"v": 1}
|
||||
assert collected[1] == {"v": 2}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_values_awaitable():
|
||||
chunks = [
|
||||
((), "values", {"v": 1}),
|
||||
((), "values", {"v": 2}),
|
||||
]
|
||||
run = await create_async_graph_run_stream(_mock_source(chunks))
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
result = await run.values
|
||||
assert result == {"v": 2}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_output_resolves():
|
||||
chunks = [((), "values", {"final": True})]
|
||||
run = await create_async_graph_run_stream(_mock_source(chunks))
|
||||
await asyncio.sleep(0.05)
|
||||
result = await run.output
|
||||
assert result == {"final": True}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_messages_yields_streams():
|
||||
chunks = [
|
||||
((), "messages", {"event": "message-start", "message_id": "m1"}),
|
||||
(
|
||||
(),
|
||||
"messages",
|
||||
{
|
||||
"event": "content-block-delta",
|
||||
"content_block": {"type": "text", "text": "hi"},
|
||||
},
|
||||
),
|
||||
((), "messages", {"event": "message-finish", "reason": "stop"}),
|
||||
]
|
||||
run = await create_async_graph_run_stream(_mock_source(chunks))
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
collected: list[ChatModelStream] = []
|
||||
async for stream in run.messages:
|
||||
collected.append(stream)
|
||||
assert len(collected) == 1
|
||||
assert isinstance(collected[0], ChatModelStream)
|
||||
assert collected[0].done
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_interrupted_false_by_default():
|
||||
vr, mr = ValuesTransformer(), MessagesTransformer()
|
||||
mux = AsyncStreamMux(transformers=[vr, mr])
|
||||
run = AsyncGraphRunStream(mux=mux, transformers=[vr, mr])
|
||||
assert run.interrupted is False
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_abort_sets_signal():
|
||||
vr, mr = ValuesTransformer(), MessagesTransformer()
|
||||
mux = AsyncStreamMux(transformers=[vr, mr])
|
||||
run = AsyncGraphRunStream(mux=mux, transformers=[vr, mr])
|
||||
assert not run.signal.is_set()
|
||||
run.abort()
|
||||
assert run.signal.is_set()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_abort_stops_pump():
|
||||
"""Calling abort() should stop the pump from processing further chunks."""
|
||||
gate = asyncio.Event()
|
||||
|
||||
async def _gated_source():
|
||||
yield ((), "values", {"v": 1})
|
||||
yield ((), "values", {"v": 2})
|
||||
await gate.wait() # Block until released
|
||||
yield ((), "values", {"v": 3}) # Should not be processed
|
||||
|
||||
run = await create_async_graph_run_stream(_gated_source())
|
||||
await asyncio.sleep(0.05) # Let first two events through
|
||||
run.abort()
|
||||
gate.set() # Unblock the source so the pump can check abort and exit
|
||||
await asyncio.sleep(0.05) # Let pump close the mux
|
||||
|
||||
collected = []
|
||||
async for event in run:
|
||||
if event["method"] == "values":
|
||||
collected.append(event["params"]["data"])
|
||||
# v:3 should not have been processed because abort was set
|
||||
assert all(v.get("v") != 3 for v in collected)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_messages_from_filters_by_node():
|
||||
"""messages_from(node) should only yield messages from the specified node."""
|
||||
chunks = [
|
||||
(
|
||||
(),
|
||||
"messages",
|
||||
{"event": "message-start", "message_id": "m1", "__node__": "agent"},
|
||||
),
|
||||
(
|
||||
(),
|
||||
"messages",
|
||||
{
|
||||
"event": "content-block-delta",
|
||||
"content_block": {"type": "text", "text": "from agent"},
|
||||
"__node__": "agent",
|
||||
},
|
||||
),
|
||||
(
|
||||
(),
|
||||
"messages",
|
||||
{"event": "message-finish", "reason": "stop", "__node__": "agent"},
|
||||
),
|
||||
(
|
||||
(),
|
||||
"messages",
|
||||
{"event": "message-start", "message_id": "m2", "__node__": "tools"},
|
||||
),
|
||||
(
|
||||
(),
|
||||
"messages",
|
||||
{
|
||||
"event": "content-block-delta",
|
||||
"content_block": {"type": "text", "text": "from tools"},
|
||||
"__node__": "tools",
|
||||
},
|
||||
),
|
||||
(
|
||||
(),
|
||||
"messages",
|
||||
{"event": "message-finish", "reason": "stop", "__node__": "tools"},
|
||||
),
|
||||
]
|
||||
run = await create_async_graph_run_stream(_mock_source(chunks))
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
agent_msgs: list[ChatModelStream] = []
|
||||
async for stream in run.messages_from("agent"):
|
||||
agent_msgs.append(stream)
|
||||
assert len(agent_msgs) == 1
|
||||
assert agent_msgs[0].node == "agent"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GraphRunStream / create_graph_run_stream
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _sync_source(
|
||||
chunks: list[tuple[tuple[str, ...], str, Any]],
|
||||
) -> Iterator[tuple[tuple[str, ...], str, Any]]:
|
||||
yield from chunks
|
||||
|
||||
|
||||
def test_sync_create_yields_all_events():
|
||||
chunks = [
|
||||
((), "values", {"step": 1}),
|
||||
((), "values", {"step": 2}),
|
||||
((), "updates", {"node": "a"}),
|
||||
]
|
||||
run = create_graph_run_stream(_sync_source(chunks))
|
||||
collected = list(run)
|
||||
assert len(collected) == 3
|
||||
assert collected[0]["method"] == "values"
|
||||
assert collected[2]["method"] == "updates"
|
||||
|
||||
|
||||
def test_sync_output():
|
||||
chunks = [
|
||||
((), "values", {"v": 1}),
|
||||
((), "values", {"v": 2}),
|
||||
]
|
||||
run = create_graph_run_stream(_sync_source(chunks))
|
||||
assert run.output == {"v": 2}
|
||||
|
||||
|
||||
def test_sync_values_iteration():
|
||||
chunks = [
|
||||
((), "values", {"v": 1}),
|
||||
((), "values", {"v": 2}),
|
||||
]
|
||||
run = create_graph_run_stream(_sync_source(chunks))
|
||||
collected = list(run.values)
|
||||
assert len(collected) == 2
|
||||
assert collected[0] == {"v": 1}
|
||||
assert collected[1] == {"v": 2}
|
||||
|
||||
|
||||
def test_sync_messages():
|
||||
chunks = [
|
||||
((), "messages", {"event": "message-start", "message_id": "m1"}),
|
||||
(
|
||||
(),
|
||||
"messages",
|
||||
{
|
||||
"event": "content-block-delta",
|
||||
"content_block": {"type": "text", "text": "hi"},
|
||||
},
|
||||
),
|
||||
((), "messages", {"event": "message-finish", "reason": "stop"}),
|
||||
]
|
||||
run = create_graph_run_stream(_sync_source(chunks))
|
||||
collected = list(run.messages)
|
||||
assert len(collected) == 1
|
||||
assert isinstance(collected[0], ChatModelStream)
|
||||
assert collected[0].done
|
||||
|
||||
|
||||
def test_sync_messages_text_accessible():
|
||||
"""Sync consumers should be able to read ChatModelStream text content
|
||||
without an async event loop.
|
||||
"""
|
||||
chunks = [
|
||||
((), "messages", {"event": "message-start", "message_id": "m1"}),
|
||||
(
|
||||
(),
|
||||
"messages",
|
||||
{
|
||||
"event": "content-block-delta",
|
||||
"content_block": {"type": "text", "text": "Hello"},
|
||||
},
|
||||
),
|
||||
(
|
||||
(),
|
||||
"messages",
|
||||
{
|
||||
"event": "content-block-delta",
|
||||
"content_block": {"type": "text", "text": " world"},
|
||||
},
|
||||
),
|
||||
((), "messages", {"event": "message-finish", "reason": "stop"}),
|
||||
]
|
||||
run = create_graph_run_stream(_sync_source(chunks))
|
||||
|
||||
for msg in run.messages:
|
||||
assert isinstance(msg.text, str)
|
||||
assert msg.text == "Hello world"
|
||||
assert msg.done
|
||||
|
||||
|
||||
def test_sync_messages_content_populated_when_yielded():
|
||||
"""When sync run.messages yields a ChatModelStream, its content should
|
||||
be fully populated (done=True) with all text accumulated.
|
||||
"""
|
||||
chunks = [
|
||||
((), "messages", {"event": "message-start", "message_id": "m1"}),
|
||||
(
|
||||
(),
|
||||
"messages",
|
||||
{
|
||||
"event": "content-block-delta",
|
||||
"content_block": {"type": "text", "text": "answer"},
|
||||
},
|
||||
),
|
||||
((), "messages", {"event": "message-finish", "reason": "stop"}),
|
||||
((), "messages", {"event": "message-start", "message_id": "m2"}),
|
||||
(
|
||||
(),
|
||||
"messages",
|
||||
{
|
||||
"event": "content-block-delta",
|
||||
"content_block": {"type": "text", "text": "second"},
|
||||
},
|
||||
),
|
||||
((), "messages", {"event": "message-finish", "reason": "stop"}),
|
||||
]
|
||||
run = create_graph_run_stream(_sync_source(chunks))
|
||||
|
||||
messages = list(run.messages)
|
||||
assert len(messages) == 2
|
||||
assert messages[0].done
|
||||
assert messages[0].text == "answer"
|
||||
assert messages[1].done
|
||||
assert messages[1].text == "second"
|
||||
|
||||
|
||||
def test_sync_output_mapper():
|
||||
chunks = [((), "values", {"v": 1})]
|
||||
run = create_graph_run_stream(
|
||||
_sync_source(chunks), output_mapper=lambda x: {"mapped": x["v"]}
|
||||
)
|
||||
assert run.output == {"mapped": 1}
|
||||
|
||||
|
||||
def test_sync_interrupted_false():
|
||||
chunks = [((), "values", {"v": 1})]
|
||||
run = create_graph_run_stream(_sync_source(chunks))
|
||||
assert run.interrupted is False
|
||||
|
||||
|
||||
def test_sync_source_error():
|
||||
"""If the source raises, the mux should fail and the error should propagate."""
|
||||
|
||||
def _bad_source():
|
||||
yield ((), "values", {"v": 1})
|
||||
raise ValueError("source error")
|
||||
|
||||
run = create_graph_run_stream(_bad_source())
|
||||
collected = list(run)
|
||||
# Events before the error are still accessible
|
||||
assert len(collected) >= 1
|
||||
assert collected[0]["method"] == "values"
|
||||
# The mux recorded the failure
|
||||
assert run._mux._error is not None
|
||||
assert isinstance(run._mux._error, ValueError)
|
||||
assert "source error" in str(run._mux._error)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GraphRunStream — lazy consumption tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sync_lazy_not_consumed_on_creation():
|
||||
"""Source iterator should not be consumed when the stream is created."""
|
||||
consumed = 0
|
||||
|
||||
def counting_source():
|
||||
nonlocal consumed
|
||||
for chunk in [
|
||||
((), "values", {"v": 1}),
|
||||
((), "values", {"v": 2}),
|
||||
((), "values", {"v": 3}),
|
||||
]:
|
||||
consumed += 1
|
||||
yield chunk
|
||||
|
||||
create_graph_run_stream(counting_source())
|
||||
assert consumed == 0
|
||||
|
||||
|
||||
def test_sync_lazy_values_pull_incrementally():
|
||||
"""Iterating .values should pull from the source one event at a time."""
|
||||
consumed = 0
|
||||
|
||||
def counting_source():
|
||||
nonlocal consumed
|
||||
for chunk in [
|
||||
((), "values", {"v": 1}),
|
||||
((), "values", {"v": 2}),
|
||||
((), "values", {"v": 3}),
|
||||
]:
|
||||
consumed += 1
|
||||
yield chunk
|
||||
|
||||
run = create_graph_run_stream(counting_source())
|
||||
assert consumed == 0
|
||||
|
||||
it = iter(run.values)
|
||||
v = next(it)
|
||||
assert v == {"v": 1}
|
||||
assert consumed == 1
|
||||
|
||||
v = next(it)
|
||||
assert v == {"v": 2}
|
||||
assert consumed == 2
|
||||
|
||||
# Source not fully drained yet
|
||||
assert consumed < 3
|
||||
|
||||
|
||||
def test_sync_lazy_output_drains_all():
|
||||
"""Accessing .output should drain the entire source."""
|
||||
consumed = 0
|
||||
|
||||
def counting_source():
|
||||
nonlocal consumed
|
||||
for chunk in [((), "values", {"v": i}) for i in range(5)]:
|
||||
consumed += 1
|
||||
yield chunk
|
||||
|
||||
run = create_graph_run_stream(counting_source())
|
||||
assert consumed == 0
|
||||
assert run.output == {"v": 4}
|
||||
assert consumed == 5
|
||||
|
||||
|
||||
def test_sync_lazy_early_break():
|
||||
"""Breaking out of a projection early should leave the source partially consumed."""
|
||||
consumed = 0
|
||||
|
||||
def counting_source():
|
||||
nonlocal consumed
|
||||
for chunk in [((), "values", {"v": i}) for i in range(10)]:
|
||||
consumed += 1
|
||||
yield chunk
|
||||
|
||||
run = create_graph_run_stream(counting_source())
|
||||
for v in run.values:
|
||||
break # consume only the first value
|
||||
assert consumed == 1
|
||||
assert consumed < 10
|
||||
|
||||
|
||||
def test_sync_lazy_interleaved_projections():
|
||||
"""Switching between projections replays buffered items then resumes pumping."""
|
||||
consumed = 0
|
||||
|
||||
def counting_source():
|
||||
nonlocal consumed
|
||||
for chunk in [
|
||||
((), "values", {"v": 1}),
|
||||
((), "messages", {"event": "message-start", "message_id": "m1"}),
|
||||
((), "messages", {"event": "message-finish", "reason": "stop"}),
|
||||
((), "values", {"v": 2}),
|
||||
]:
|
||||
consumed += 1
|
||||
yield chunk
|
||||
|
||||
run = create_graph_run_stream(counting_source())
|
||||
|
||||
# Pull first value — consumes 1 source item
|
||||
vit = iter(run.values)
|
||||
assert next(vit) == {"v": 1}
|
||||
assert consumed == 1
|
||||
|
||||
# Pull first message — pumps until message-finish (item 3) so the
|
||||
# ChatModelStream is fully populated before yielding.
|
||||
mit = iter(run.messages)
|
||||
msg = next(mit)
|
||||
assert isinstance(msg, ChatModelStream)
|
||||
assert msg.done
|
||||
assert consumed == 3
|
||||
|
||||
# Pull second value — pumps values (item 4)
|
||||
assert next(vit) == {"v": 2}
|
||||
assert consumed == 4
|
||||
|
||||
|
||||
def test_sync_lazy_iter_pulls_incrementally():
|
||||
"""Raw __iter__ should pull from the source lazily."""
|
||||
consumed = 0
|
||||
|
||||
def counting_source():
|
||||
nonlocal consumed
|
||||
for chunk in [
|
||||
((), "values", {"v": 1}),
|
||||
((), "updates", {"node": "a"}),
|
||||
((), "values", {"v": 2}),
|
||||
]:
|
||||
consumed += 1
|
||||
yield chunk
|
||||
|
||||
run = create_graph_run_stream(counting_source())
|
||||
it = iter(run)
|
||||
event = next(it)
|
||||
assert event["method"] == "values"
|
||||
assert consumed == 1
|
||||
|
||||
event = next(it)
|
||||
assert event["method"] == "updates"
|
||||
assert consumed == 2
|
||||
|
||||
|
||||
def test_sync_lazy_source_error():
|
||||
"""If the source raises mid-stream, earlier events are still accessible."""
|
||||
consumed = 0
|
||||
|
||||
def bad_source():
|
||||
nonlocal consumed
|
||||
consumed += 1
|
||||
yield ((), "values", {"v": 1})
|
||||
raise ValueError("boom")
|
||||
|
||||
run = create_graph_run_stream(bad_source())
|
||||
collected = list(run)
|
||||
assert len(collected) >= 1
|
||||
assert collected[0]["method"] == "values"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_subgraph_child_values_receive_post_discovery_events():
|
||||
"""Child AsyncSubgraphRunStream.values iteration should include events
|
||||
that arrive AFTER the subgraph namespace is first discovered.
|
||||
|
||||
``_SubgraphsProjection`` creates a local ``ValuesTransformer`` for
|
||||
each child and replays existing events, but never registers the
|
||||
transformer with the mux. Events that arrive after discovery are
|
||||
not routed to it, and ``finalize()`` is not called (the mux wasn't
|
||||
closed at discovery time), so the child's values_log is never
|
||||
closed and iteration hangs.
|
||||
"""
|
||||
gate = asyncio.Event()
|
||||
|
||||
async def _source() -> AsyncIterator[tuple[tuple[str, ...], str, Any]]:
|
||||
# First event from child namespace — triggers discovery
|
||||
yield (("child:0",), "values", {"v": 1})
|
||||
await gate.wait()
|
||||
# Second event from same child — arrives after discovery
|
||||
yield (("child:0",), "values", {"v": 2})
|
||||
# Root event so the mux tracks output
|
||||
yield ((), "values", {"done": True})
|
||||
|
||||
run = await create_async_graph_run_stream(_source())
|
||||
await asyncio.sleep(0.05) # let pump process first event
|
||||
|
||||
# Get the first subgraph while the mux is still open
|
||||
sub = None
|
||||
async for s in run.subgraphs:
|
||||
sub = s
|
||||
break
|
||||
|
||||
assert sub is not None
|
||||
|
||||
# Release the gate so the pump finishes
|
||||
gate.set()
|
||||
await asyncio.sleep(0.05) # let pump close mux
|
||||
|
||||
# ``await sub.output`` uses the mux's output future — works fine
|
||||
output = await sub.output
|
||||
assert output == {"v": 2}, "await sub.output should reflect the latest value"
|
||||
|
||||
# But ``async for v in sub.values`` only gets the replayed event
|
||||
# and then hangs because the child's values_log is never closed.
|
||||
values: list[Any] = []
|
||||
try:
|
||||
async with asyncio.timeout(1.0):
|
||||
async for v in sub.values:
|
||||
values.append(v)
|
||||
except (asyncio.TimeoutError, TimeoutError):
|
||||
pass
|
||||
|
||||
assert len(values) == 2, (
|
||||
f"Expected 2 child value snapshots but got {len(values)}: {values}. "
|
||||
"Child transformer missed post-discovery events."
|
||||
)
|
||||
@@ -0,0 +1,420 @@
|
||||
"""Prove V1 and StreamingHandler APIs expose identical information.
|
||||
|
||||
Each test runs the same graph through both APIs and asserts data
|
||||
equivalence — same state snapshots, same messages, same custom events,
|
||||
same interrupts. Sync APIs are used where possible; async tests cover
|
||||
features without sync equivalents (subgraphs projection, messages_from).
|
||||
|
||||
Run with:
|
||||
TEST=tests/test_streaming_comparison.py make test
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.config import get_stream_writer
|
||||
from langgraph.graph import END, START, MessagesState, StateGraph
|
||||
from langgraph.stream import StreamingHandler
|
||||
from langgraph.stream._convert import STREAM_V2_MODES
|
||||
from langgraph.types import interrupt
|
||||
from tests.fake_chat import FakeChatModel
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Graph factories
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class State(TypedDict):
|
||||
value: str
|
||||
items: Annotated[list[str], lambda a, b: a + b]
|
||||
|
||||
|
||||
def _linear_graph(n_nodes: int = 3):
|
||||
"""Chain of *n_nodes* that concatenate strings."""
|
||||
g = StateGraph(State)
|
||||
names = [f"node_{i}" for i in range(n_nodes)]
|
||||
for name in names:
|
||||
|
||||
def make_fn(n):
|
||||
def fn(state: State) -> dict:
|
||||
return {"value": state["value"] + f"_{n}", "items": [n]}
|
||||
|
||||
return fn
|
||||
|
||||
g.add_node(name, make_fn(name))
|
||||
|
||||
g.add_edge(START, names[0])
|
||||
for i in range(len(names) - 1):
|
||||
g.add_edge(names[i], names[i + 1])
|
||||
g.add_edge(names[-1], END)
|
||||
return g.compile()
|
||||
|
||||
|
||||
def _chat_graph():
|
||||
"""Single agent node with a FakeChatModel."""
|
||||
model = FakeChatModel(messages=[AIMessage(content="Hello from agent")])
|
||||
|
||||
def agent(state: dict) -> dict:
|
||||
return {"messages": [model.invoke(state["messages"])]}
|
||||
|
||||
g = StateGraph(MessagesState)
|
||||
g.add_node("agent", agent)
|
||||
g.add_edge(START, "agent")
|
||||
g.add_edge("agent", END)
|
||||
return g.compile()
|
||||
|
||||
|
||||
def _multi_node_chat_graph():
|
||||
"""Two LLM nodes: agent -> reviewer."""
|
||||
agent_model = FakeChatModel(messages=[AIMessage(content="Agent reply")])
|
||||
reviewer_model = FakeChatModel(messages=[AIMessage(content="Reviewer reply")])
|
||||
|
||||
def agent(state: dict) -> dict:
|
||||
return {"messages": [agent_model.invoke(state["messages"])]}
|
||||
|
||||
def reviewer(state: dict) -> dict:
|
||||
return {"messages": [reviewer_model.invoke(state["messages"])]}
|
||||
|
||||
g = StateGraph(MessagesState)
|
||||
g.add_node("agent", agent)
|
||||
g.add_node("reviewer", reviewer)
|
||||
g.add_edge(START, "agent")
|
||||
g.add_edge("agent", "reviewer")
|
||||
g.add_edge("reviewer", END)
|
||||
return g.compile()
|
||||
|
||||
|
||||
def _custom_events_graph():
|
||||
"""Node that emits custom events via StreamWriter."""
|
||||
|
||||
def worker(state: State) -> dict:
|
||||
writer = get_stream_writer()
|
||||
writer({"step": 1, "msg": "started"})
|
||||
writer({"step": 2, "msg": "processing"})
|
||||
writer({"step": 3, "msg": "done"})
|
||||
return {"value": state["value"] + "_done", "items": ["done"]}
|
||||
|
||||
g = StateGraph(State)
|
||||
g.add_node("worker", worker)
|
||||
g.add_edge(START, "worker")
|
||||
g.add_edge("worker", END)
|
||||
return g.compile()
|
||||
|
||||
|
||||
def _interrupt_graph():
|
||||
"""Graph that interrupts for human input."""
|
||||
|
||||
def ask_human(state: State) -> dict:
|
||||
answer = interrupt("What next?")
|
||||
return {"value": state["value"] + f"_{answer}", "items": [answer]}
|
||||
|
||||
g = StateGraph(State)
|
||||
g.add_node("ask", ask_human)
|
||||
g.add_edge(START, "ask")
|
||||
g.add_edge("ask", END)
|
||||
return g.compile(checkpointer=MemorySaver())
|
||||
|
||||
|
||||
def _subgraph():
|
||||
"""Parent with a compiled child subgraph."""
|
||||
|
||||
class ChildState(TypedDict):
|
||||
value: str
|
||||
|
||||
class ParentState(TypedDict):
|
||||
value: str
|
||||
|
||||
def child_node(state: ChildState) -> dict:
|
||||
return {"value": state["value"] + "_child"}
|
||||
|
||||
child = StateGraph(ChildState)
|
||||
child.add_node("inner", child_node)
|
||||
child.add_edge(START, "inner")
|
||||
child.add_edge("inner", END)
|
||||
child_compiled = child.compile()
|
||||
|
||||
parent = StateGraph(ParentState)
|
||||
parent.add_node("child", child_compiled)
|
||||
parent.add_edge(START, "child")
|
||||
parent.add_edge("child", END)
|
||||
return parent.compile()
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# 1. Final output
|
||||
# ===================================================================
|
||||
|
||||
|
||||
def test_output():
|
||||
"""graph.invoke() produces the same result as StreamingHandler().stream().output."""
|
||||
graph = _linear_graph()
|
||||
inp = {"value": "x", "items": []}
|
||||
|
||||
v1 = graph.invoke(inp)
|
||||
|
||||
run = StreamingHandler(graph).stream(inp)
|
||||
v2 = run.output
|
||||
|
||||
assert v1 == v2
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# 2. Intermediate state snapshots (values mode)
|
||||
# ===================================================================
|
||||
|
||||
|
||||
def test_values():
|
||||
"""stream(mode='values') snapshots == StreamingHandler().stream().values snapshots."""
|
||||
graph = _linear_graph()
|
||||
inp = {"value": "x", "items": []}
|
||||
|
||||
v1 = list(graph.stream(inp, stream_mode="values"))
|
||||
|
||||
run = StreamingHandler(graph).stream(inp)
|
||||
v2 = list(run.values)
|
||||
|
||||
assert v1 == v2
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# 3. Per-node updates (updates mode)
|
||||
# ===================================================================
|
||||
|
||||
|
||||
def test_updates():
|
||||
"""stream(mode='updates') data == StreamingHandler raw events[method=updates]."""
|
||||
graph = _linear_graph()
|
||||
inp = {"value": "x", "items": []}
|
||||
|
||||
v1 = list(graph.stream(inp, stream_mode="updates"))
|
||||
|
||||
run = StreamingHandler(graph).stream(inp)
|
||||
v2 = [
|
||||
e["params"]["data"]
|
||||
for e in run
|
||||
if e["method"] == "updates" and not e["params"]["namespace"]
|
||||
]
|
||||
|
||||
assert v1 == v2
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# 4. Message text and node attribution
|
||||
# ===================================================================
|
||||
|
||||
|
||||
def test_messages():
|
||||
"""Reassembled V1 message text per node == V2 .messages text per node."""
|
||||
graph = _multi_node_chat_graph()
|
||||
inp = {"messages": [HumanMessage(content="hi")]}
|
||||
|
||||
# V1: collect (chunk, metadata) pairs, group text by node
|
||||
v1_text_by_node: dict[str, list[str]] = {}
|
||||
for chunk, metadata in graph.stream(inp, stream_mode="messages"):
|
||||
node = metadata["langgraph_node"]
|
||||
v1_text_by_node.setdefault(node, []).append(chunk.content)
|
||||
v1_text = {k: "".join(v) for k, v in v1_text_by_node.items()}
|
||||
|
||||
# V2: each ChatModelStream has .text and .node
|
||||
run = StreamingHandler(graph).stream(inp)
|
||||
v2_text: dict[str, str] = {}
|
||||
for msg in run.messages:
|
||||
assert msg.done is True
|
||||
v2_text[msg.node] = msg.text
|
||||
|
||||
assert v1_text == v2_text
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# 5. Custom events
|
||||
# ===================================================================
|
||||
|
||||
|
||||
def test_custom_events():
|
||||
"""stream(mode='custom') payloads == StreamingHandler raw events[method=custom]."""
|
||||
graph = _custom_events_graph()
|
||||
inp = {"value": "x", "items": []}
|
||||
|
||||
v1 = list(graph.stream(inp, stream_mode="custom"))
|
||||
|
||||
run = StreamingHandler(graph).stream(inp)
|
||||
v2 = [
|
||||
e["params"]["data"]
|
||||
for e in run
|
||||
if e["method"] == "custom" and not e["params"]["namespace"]
|
||||
]
|
||||
|
||||
assert v1 == v2
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# 6. Mode coverage
|
||||
# ===================================================================
|
||||
|
||||
|
||||
def test_mode_coverage():
|
||||
"""V2 produces events for the same set of modes as V1."""
|
||||
graph = _chat_graph()
|
||||
inp = {"messages": [HumanMessage(content="hi")]}
|
||||
|
||||
# V1: request all modes, collect which ones appear
|
||||
v1_modes: set[str] = set()
|
||||
for ns, mode, _ in graph.stream(
|
||||
inp, stream_mode=STREAM_V2_MODES, subgraphs=True, version="v1"
|
||||
):
|
||||
if not ns:
|
||||
v1_modes.add(mode)
|
||||
|
||||
# V2: iterate raw events, collect methods
|
||||
run = StreamingHandler(graph).stream(inp)
|
||||
v2_modes = {e["method"] for e in run if not e["params"]["namespace"]}
|
||||
|
||||
assert v1_modes == v2_modes
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# 7. Interrupt detection
|
||||
# ===================================================================
|
||||
|
||||
|
||||
def test_interrupts():
|
||||
"""V1 __interrupt__ value == V2 .interrupted and .interrupts payload."""
|
||||
graph = _interrupt_graph()
|
||||
inp = {"value": "x", "items": []}
|
||||
|
||||
# V1: detect __interrupt__ in values stream
|
||||
config1 = {"configurable": {"thread_id": "equiv-1"}}
|
||||
v1_interrupt_value = None
|
||||
for chunk in graph.stream(inp, config1, stream_mode="values"):
|
||||
if isinstance(chunk, dict) and "__interrupt__" in chunk:
|
||||
info = chunk["__interrupt__"]
|
||||
if info:
|
||||
v1_interrupt_value = info[0].value
|
||||
|
||||
assert v1_interrupt_value is not None
|
||||
|
||||
# V2: .interrupted and .interrupts (fresh thread)
|
||||
config2 = {"configurable": {"thread_id": "equiv-2"}}
|
||||
run = StreamingHandler(graph).stream(inp, config=config2)
|
||||
for _ in run:
|
||||
pass
|
||||
|
||||
assert run.interrupted is True
|
||||
assert len(run.interrupts) > 0
|
||||
v2_interrupt_value = run.interrupts[0]["payload"].value
|
||||
|
||||
assert v1_interrupt_value == v2_interrupt_value
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# 8. Subgraph state snapshots
|
||||
# ===================================================================
|
||||
|
||||
|
||||
def test_subgraph_values():
|
||||
"""V1 child namespace values == V2 child namespace values."""
|
||||
graph = _subgraph()
|
||||
inp = {"value": "x"}
|
||||
|
||||
# V1: stream with subgraphs=True, collect child values
|
||||
v1_child_values = []
|
||||
for ns, data in graph.stream(inp, stream_mode="values", subgraphs=True):
|
||||
if ns:
|
||||
v1_child_values.append(data)
|
||||
|
||||
# V2: filter raw events for child namespace + values mode
|
||||
run = StreamingHandler(graph).stream(inp)
|
||||
v2_child_values = [
|
||||
e["params"]["data"]
|
||||
for e in run
|
||||
if e["method"] == "values" and e["params"]["namespace"]
|
||||
]
|
||||
|
||||
assert v1_child_values == v2_child_values
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# 9. Node filtering on messages
|
||||
# ===================================================================
|
||||
|
||||
|
||||
def test_messages_node_filtering():
|
||||
"""V1 manual metadata filter == V2 .messages filtered by .node."""
|
||||
graph = _multi_node_chat_graph()
|
||||
inp = {"messages": [HumanMessage(content="hi")]}
|
||||
|
||||
# V1: manual filter for "agent" node only
|
||||
v1_agent_text: list[str] = []
|
||||
for chunk, metadata in graph.stream(inp, stream_mode="messages"):
|
||||
if metadata.get("langgraph_node") == "agent":
|
||||
v1_agent_text.append(chunk.content)
|
||||
v1_text = "".join(v1_agent_text)
|
||||
|
||||
# V2: filter .messages by .node
|
||||
run = StreamingHandler(graph).stream(inp)
|
||||
v2_agent_msgs = [msg for msg in run.messages if msg.node == "agent"]
|
||||
assert len(v2_agent_msgs) == 1
|
||||
v2_text = v2_agent_msgs[0].text
|
||||
|
||||
assert v1_text == v2_text
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# 10. Async: subgraphs projection
|
||||
# ===================================================================
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_subgraph_projection():
|
||||
"""V2 .subgraphs child output matches V1 child namespace output."""
|
||||
graph = _subgraph()
|
||||
inp = {"value": "x"}
|
||||
|
||||
# V1
|
||||
v1_child_output = None
|
||||
async for ns, data in graph.astream(inp, stream_mode="values", subgraphs=True):
|
||||
if ns:
|
||||
v1_child_output = data
|
||||
|
||||
# V2: .subgraphs yields typed child stream objects
|
||||
run = await StreamingHandler(graph).astream(inp)
|
||||
v2_child_output = None
|
||||
async for sub in run.subgraphs:
|
||||
v2_child_output = await sub.output
|
||||
|
||||
assert v1_child_output == v2_child_output
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# 11. Async: messages_from projection
|
||||
# ===================================================================
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_messages_from():
|
||||
"""V2 .messages_from('agent') text matches V1 filtered by metadata."""
|
||||
graph = _multi_node_chat_graph()
|
||||
inp = {"messages": [HumanMessage(content="hi")]}
|
||||
|
||||
# V1: manual filter for agent node
|
||||
v1_agent_text: list[str] = []
|
||||
async for chunk, metadata in graph.astream(inp, stream_mode="messages"):
|
||||
if metadata.get("langgraph_node") == "agent":
|
||||
v1_agent_text.append(chunk.content)
|
||||
v1_text = "".join(v1_agent_text)
|
||||
|
||||
# V2: declarative node filtering
|
||||
run = await StreamingHandler(graph).astream(inp)
|
||||
v2_texts: list[str] = []
|
||||
async for msg in run.messages_from("agent"):
|
||||
v2_texts.append(await msg.text)
|
||||
assert len(v2_texts) == 1
|
||||
v2_text = v2_texts[0]
|
||||
|
||||
assert v1_text == v2_text
|
||||
Reference in New Issue
Block a user