mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-22 15:42:25 +02:00
500 lines
18 KiB
Python
500 lines
18 KiB
Python
"""Compare token-level LLM streaming: v1 graph.stream() vs v2 StreamingHandler.
|
|
|
|
Runs the same single-node chatbot graph six ways:
|
|
1. v1 sync — graph.stream(stream_mode="messages")
|
|
2. v2 sync — StreamingHandler.stream() -> run.messages
|
|
3. v2 sync — same, with a custom TokenMetrics transformer
|
|
4. v1 async — graph.astream(stream_mode="messages")
|
|
5. v2 async — StreamingHandler.astream() -> run.messages
|
|
6. v2 async — same, with a custom TokenMetrics transformer
|
|
|
|
The TokenMetrics transformer demonstrates v2's extensibility: it tracks
|
|
first-token latency, token count, and throughput as a reusable component
|
|
that runs alongside the built-in projections without any changes to the
|
|
consumption loop. The v1 equivalent requires weaving all that bookkeeping
|
|
into your stream loop inline.
|
|
|
|
Requirements:
|
|
pip install langchain-openai
|
|
export OPENAI_API_KEY=...
|
|
|
|
Usage:
|
|
python streaming_comparison.py
|
|
python streaming_comparison.py --sync-only
|
|
python streaming_comparison.py --async-only
|
|
python streaming_comparison.py --question "your question here"
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Annotated, Any
|
|
|
|
from langchain_core.messages import BaseMessage, HumanMessage
|
|
from langchain_openai import ChatOpenAI
|
|
from typing_extensions import TypedDict
|
|
|
|
from langgraph.graph import StateGraph, add_messages
|
|
from langgraph.stream import StreamChannel, StreamingHandler, StreamTransformer
|
|
from langgraph.stream._types import ProtocolEvent
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Graph setup
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list[BaseMessage], add_messages]
|
|
|
|
|
|
_llm = None
|
|
|
|
|
|
def _get_llm() -> ChatOpenAI:
|
|
global _llm
|
|
if _llm is None:
|
|
_llm = ChatOpenAI(model="gpt-4o-mini", streaming=True)
|
|
return _llm
|
|
|
|
|
|
def chatbot(state: State) -> dict[str, Any]:
|
|
return {"messages": [_get_llm().invoke(state["messages"])]}
|
|
|
|
|
|
async def achatbot(state: State) -> dict[str, Any]:
|
|
return {"messages": [await _get_llm().ainvoke(state["messages"])]}
|
|
|
|
|
|
def build_graph(*, use_async_node: bool = False):
|
|
builder = StateGraph(State)
|
|
builder.add_node("chatbot", achatbot if use_async_node else chatbot)
|
|
builder.set_entry_point("chatbot")
|
|
builder.set_finish_point("chatbot")
|
|
return builder.compile()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Custom transformer: TokenMetrics
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class TokenSnapshot:
|
|
"""A throughput measurement pushed to the channel.
|
|
|
|
Periodic snapshots have `is_final=False`. The last item pushed in
|
|
`finalize()` has `is_final=True` and includes `first_token_latency`.
|
|
"""
|
|
|
|
token_index: int
|
|
elapsed: float
|
|
tokens_per_sec: float
|
|
is_final: bool = False
|
|
first_token_latency: float | None = None
|
|
|
|
|
|
class TokenMetricsTransformer(StreamTransformer):
|
|
"""Tracks token throughput and first-token latency.
|
|
|
|
Observes "messages" events and counts chunks with non-empty content.
|
|
Pushes a TokenSnapshot to its StreamChannel every N tokens so
|
|
consumers can monitor throughput in real time. A final snapshot with
|
|
`is_final=True` is pushed when the run completes.
|
|
|
|
All data is consumed through the stream via
|
|
`run.extensions["token_metrics"]` — no need to read transformer
|
|
internals.
|
|
|
|
This is the kind of cross-cutting concern that v2 transformers handle
|
|
cleanly: write it once, attach it to any graph, and it works alongside
|
|
all other projections without touching the consumption loop.
|
|
"""
|
|
|
|
def __init__(self, *, snapshot_every: int = 10) -> None:
|
|
self._channel: StreamChannel[TokenSnapshot] = StreamChannel("token_metrics")
|
|
self._snapshot_every = snapshot_every
|
|
self._t0: float | None = None
|
|
self._first_token_time: float | None = None
|
|
self._token_count = 0
|
|
|
|
def init(self) -> dict[str, Any]:
|
|
return {"token_metrics": self._channel}
|
|
|
|
def process(self, event: ProtocolEvent) -> bool:
|
|
if event["method"] != "messages":
|
|
return True
|
|
|
|
# Only count root-namespace events
|
|
if event["params"]["namespace"]:
|
|
return True
|
|
|
|
chunk = event["params"]["data"]
|
|
# messages data is (chunk, metadata) tuple
|
|
if not isinstance(chunk, tuple) or len(chunk) != 2:
|
|
return True
|
|
message, _metadata = chunk
|
|
if not hasattr(message, "content") or not message.content:
|
|
return True
|
|
|
|
now = time.time()
|
|
if self._t0 is None:
|
|
self._t0 = now
|
|
if self._first_token_time is None:
|
|
self._first_token_time = now
|
|
|
|
self._token_count += 1
|
|
|
|
# Push periodic snapshots
|
|
if self._token_count % self._snapshot_every == 0:
|
|
elapsed = now - self._t0
|
|
self._channel.push(
|
|
TokenSnapshot(
|
|
token_index=self._token_count,
|
|
elapsed=elapsed,
|
|
tokens_per_sec=self._token_count / elapsed if elapsed > 0 else 0,
|
|
)
|
|
)
|
|
|
|
return True
|
|
|
|
def finalize(self) -> None:
|
|
now = time.time()
|
|
elapsed = now - self._t0 if self._t0 else 0
|
|
self._channel.push(
|
|
TokenSnapshot(
|
|
token_index=self._token_count,
|
|
elapsed=elapsed,
|
|
tokens_per_sec=self._token_count / elapsed if elapsed > 0 else 0,
|
|
is_final=True,
|
|
first_token_latency=(
|
|
self._first_token_time - self._t0
|
|
if self._first_token_time and self._t0
|
|
else None
|
|
),
|
|
)
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def separator(title: str) -> None:
|
|
print(f"\n{'=' * 64}")
|
|
print(f" {title}")
|
|
print("=" * 64)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# V1 sync: graph.stream(stream_mode="messages")
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def run_v1_sync(question: str) -> None:
|
|
separator("V1 SYNC: graph.stream(stream_mode='messages')")
|
|
|
|
graph = build_graph()
|
|
input_data = {"messages": [HumanMessage(content=question)]}
|
|
|
|
t0 = time.time()
|
|
first_token_time = None
|
|
token_count = 0
|
|
|
|
for chunk, metadata in graph.stream(input_data, stream_mode="messages"):
|
|
if hasattr(chunk, "content") and chunk.content:
|
|
if first_token_time is None:
|
|
first_token_time = time.time()
|
|
token_count += 1
|
|
print(chunk.content, end="", flush=True)
|
|
|
|
elapsed = time.time() - t0
|
|
ftl = f"{first_token_time - t0:.3f}s" if first_token_time else "n/a"
|
|
tps = token_count / elapsed if elapsed > 0 else 0
|
|
print(f"\n {token_count} tokens | first: {ftl} | {tps:.1f} tok/s | {elapsed:.2f}s total")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# V1 sync with inline metrics (the v1 equivalent of the transformer)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def run_v1_sync_with_metrics(question: str) -> None:
|
|
separator("V1 SYNC + INLINE METRICS")
|
|
print(" (same bookkeeping the transformer does, but woven into the loop)")
|
|
|
|
graph = build_graph()
|
|
input_data = {"messages": [HumanMessage(content=question)]}
|
|
|
|
# All of this state lives in your consumption loop.
|
|
# Duplicate it everywhere you stream this graph.
|
|
t0 = time.time()
|
|
first_token_time = None
|
|
token_count = 0
|
|
snapshot_every = 10
|
|
snapshots: list[dict] = []
|
|
|
|
for chunk, metadata in graph.stream(input_data, stream_mode="messages"):
|
|
if hasattr(chunk, "content") and chunk.content:
|
|
now = time.time()
|
|
if first_token_time is None:
|
|
first_token_time = now
|
|
token_count += 1
|
|
print(chunk.content, end="", flush=True)
|
|
|
|
# Inline throughput tracking — mixed into your display logic
|
|
if token_count % snapshot_every == 0:
|
|
elapsed = now - t0
|
|
snapshots.append({
|
|
"token_index": token_count,
|
|
"elapsed": elapsed,
|
|
"tokens_per_sec": token_count / elapsed if elapsed > 0 else 0,
|
|
})
|
|
|
|
# Compute summary inline
|
|
elapsed = time.time() - t0
|
|
ftl = first_token_time - t0 if first_token_time else None
|
|
tps = token_count / elapsed if elapsed > 0 else 0
|
|
|
|
print(f"\n {token_count} tokens | first: {ftl:.3f}s | {tps:.1f} tok/s | {elapsed:.2f}s total")
|
|
if snapshots:
|
|
print(f" Throughput snapshots: {len(snapshots)}")
|
|
for s in snapshots:
|
|
print(f" token {s['token_index']:>4d}: {s['tokens_per_sec']:.1f} tok/s")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# V2 sync: StreamingHandler → run.messages
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def run_v2_sync(question: str) -> None:
|
|
separator("V2 SYNC: StreamingHandler.stream() -> run.messages")
|
|
|
|
graph = build_graph()
|
|
handler = StreamingHandler(graph)
|
|
input_data = {"messages": [HumanMessage(content=question)]}
|
|
|
|
run = handler.stream(input_data)
|
|
|
|
# run.messages yields (chunk, metadata) tuples as the LLM produces them
|
|
for chunk, metadata in run.messages:
|
|
if hasattr(chunk, "content") and chunk.content:
|
|
print(chunk.content, end="", flush=True)
|
|
|
|
output = run.output
|
|
assert output is not None
|
|
print(f"\n Final: {output['messages'][-1].content[:80]}...")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# V2 sync with custom transformer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def run_v2_sync_with_transformer(question: str) -> None:
|
|
separator("V2 SYNC + TokenMetrics TRANSFORMER")
|
|
print(" (metrics computed by a reusable transformer — loop stays clean)")
|
|
|
|
graph = build_graph()
|
|
handler = StreamingHandler(graph)
|
|
input_data = {"messages": [HumanMessage(content=question)]}
|
|
|
|
run = handler.stream(input_data, transformers=[TokenMetricsTransformer(snapshot_every=10)])
|
|
|
|
# The consumption loop is unchanged — just print tokens.
|
|
# The transformer silently tracks metrics in the background.
|
|
for chunk, metadata in run.messages:
|
|
if hasattr(chunk, "content") and chunk.content:
|
|
print(chunk.content, end="", flush=True)
|
|
|
|
_ = run.output
|
|
|
|
# All metrics are read from the stream, not the transformer instance.
|
|
# The channel contains periodic snapshots + a final summary snapshot.
|
|
snapshots = list(run.extensions["token_metrics"])
|
|
final = next((s for s in snapshots if s.is_final), None)
|
|
periodic = [s for s in snapshots if not s.is_final]
|
|
|
|
if final:
|
|
ftl = f"{final.first_token_latency:.3f}s" if final.first_token_latency is not None else "n/a"
|
|
print(f"\n {final.token_index} tokens | first: {ftl} | {final.tokens_per_sec:.1f} tok/s | {final.elapsed:.2f}s total")
|
|
if periodic:
|
|
print(f" Throughput snapshots: {len(periodic)}")
|
|
for snap in periodic:
|
|
print(f" token {snap.token_index:>4d}: {snap.tokens_per_sec:.1f} tok/s")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# V1 async: graph.astream(stream_mode="messages")
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def run_v1_async(question: str) -> None:
|
|
separator("V1 ASYNC: graph.astream(stream_mode='messages')")
|
|
|
|
graph = build_graph(use_async_node=True)
|
|
input_data = {"messages": [HumanMessage(content=question)]}
|
|
|
|
t0 = time.time()
|
|
first_token_time = None
|
|
token_count = 0
|
|
|
|
async for chunk, metadata in graph.astream(input_data, stream_mode="messages"):
|
|
if hasattr(chunk, "content") and chunk.content:
|
|
if first_token_time is None:
|
|
first_token_time = time.time()
|
|
token_count += 1
|
|
print(chunk.content, end="", flush=True)
|
|
|
|
elapsed = time.time() - t0
|
|
ftl = f"{first_token_time - t0:.3f}s" if first_token_time else "n/a"
|
|
tps = token_count / elapsed if elapsed > 0 else 0
|
|
print(f"\n {token_count} tokens | first: {ftl} | {tps:.1f} tok/s | {elapsed:.2f}s total")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# V1 async with inline metrics
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def run_v1_async_with_metrics(question: str) -> None:
|
|
separator("V1 ASYNC + INLINE METRICS")
|
|
print(" (same bookkeeping the transformer does, but woven into the loop)")
|
|
|
|
graph = build_graph(use_async_node=True)
|
|
input_data = {"messages": [HumanMessage(content=question)]}
|
|
|
|
t0 = time.time()
|
|
first_token_time = None
|
|
token_count = 0
|
|
snapshot_every = 10
|
|
snapshots: list[dict] = []
|
|
|
|
async for chunk, metadata in graph.astream(input_data, stream_mode="messages"):
|
|
if hasattr(chunk, "content") and chunk.content:
|
|
now = time.time()
|
|
if first_token_time is None:
|
|
first_token_time = now
|
|
token_count += 1
|
|
print(chunk.content, end="", flush=True)
|
|
|
|
if token_count % snapshot_every == 0:
|
|
elapsed = now - t0
|
|
snapshots.append({
|
|
"token_index": token_count,
|
|
"elapsed": elapsed,
|
|
"tokens_per_sec": token_count / elapsed if elapsed > 0 else 0,
|
|
})
|
|
|
|
elapsed = time.time() - t0
|
|
ftl = first_token_time - t0 if first_token_time else None
|
|
tps = token_count / elapsed if elapsed > 0 else 0
|
|
|
|
print(f"\n {token_count} tokens | first: {ftl:.3f}s | {tps:.1f} tok/s | {elapsed:.2f}s total")
|
|
if snapshots:
|
|
print(f" Throughput snapshots: {len(snapshots)}")
|
|
for s in snapshots:
|
|
print(f" token {s['token_index']:>4d}: {s['tokens_per_sec']:.1f} tok/s")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# V2 async: StreamingHandler → run.messages
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def run_v2_async(question: str) -> None:
|
|
separator("V2 ASYNC: StreamingHandler.astream() -> run.messages")
|
|
|
|
graph = build_graph(use_async_node=True)
|
|
handler = StreamingHandler(graph)
|
|
input_data = {"messages": [HumanMessage(content=question)]}
|
|
|
|
run = await handler.astream(input_data)
|
|
|
|
async for chunk, metadata in run.messages:
|
|
if hasattr(chunk, "content") and chunk.content:
|
|
print(chunk.content, end="", flush=True)
|
|
|
|
output = await run.output
|
|
assert output is not None
|
|
print(f"\n Final: {output['messages'][-1].content[:80]}...")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# V2 async with custom transformer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def run_v2_async_with_transformer(question: str) -> None:
|
|
separator("V2 ASYNC + TokenMetrics TRANSFORMER")
|
|
print(" (metrics computed by a reusable transformer — loop stays clean)")
|
|
|
|
graph = build_graph(use_async_node=True)
|
|
handler = StreamingHandler(graph)
|
|
input_data = {"messages": [HumanMessage(content=question)]}
|
|
|
|
run = await handler.astream(input_data, transformers=[TokenMetricsTransformer(snapshot_every=10)])
|
|
|
|
async for chunk, metadata in run.messages:
|
|
if hasattr(chunk, "content") and chunk.content:
|
|
print(chunk.content, end="", flush=True)
|
|
|
|
_ = await run.output
|
|
|
|
snapshots = [s async for s in run.extensions["token_metrics"]]
|
|
final = next((s for s in snapshots if s.is_final), None)
|
|
periodic = [s for s in snapshots if not s.is_final]
|
|
|
|
if final:
|
|
ftl = f"{final.first_token_latency:.3f}s" if final.first_token_latency is not None else "n/a"
|
|
print(f"\n {final.token_index} tokens | first: {ftl} | {final.tokens_per_sec:.1f} tok/s | {final.elapsed:.2f}s total")
|
|
if periodic:
|
|
print(f" Throughput snapshots: {len(periodic)}")
|
|
for snap in periodic:
|
|
print(f" token {snap.token_index:>4d}: {snap.tokens_per_sec:.1f} tok/s")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Token streaming comparison: v1 vs v2")
|
|
parser.add_argument("--sync-only", action="store_true")
|
|
parser.add_argument("--async-only", action="store_true")
|
|
parser.add_argument(
|
|
"--question",
|
|
default="Explain quantum entanglement in 2-3 sentences.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if not args.async_only:
|
|
# V1 baseline
|
|
run_v1_sync(args.question)
|
|
# V1 with inline metrics — the manual approach
|
|
run_v1_sync_with_metrics(args.question)
|
|
# V2 basic
|
|
run_v2_sync(args.question)
|
|
# V2 with transformer — same metrics, zero loop changes
|
|
run_v2_sync_with_transformer(args.question)
|
|
|
|
if not args.sync_only:
|
|
# V1 baseline
|
|
asyncio.run(run_v1_async(args.question))
|
|
# V1 with inline metrics
|
|
asyncio.run(run_v1_async_with_metrics(args.question))
|
|
# V2 basic
|
|
asyncio.run(run_v2_async(args.question))
|
|
# V2 with transformer
|
|
asyncio.run(run_v2_async_with_transformer(args.question))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|