mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(langgraph): correct SSE protocol mismatch in LangGraphSSEStreamIterator
Fixes #24093. The parser was treating `data[0]` as the event type string, expecting frames shaped like `["messages", payload]`. LangGraph actually emits standard SSE where the event type is carried by the `event:` header line and the `data:` payload for a `messages` event is `[AIMessageChunk, metadata_dict]`. Because `data[0]` is a dict (not the string `"messages"`), every real streaming frame was silently dropped, producing an empty response. Changes: - Add `_current_event: Optional[str]` state to track the SSE `event:` field between lines. - `_parse_sse_line` now sets `_current_event` on `event:` lines and resets it on empty (event-boundary) lines. - `_process_data` dispatches on `event_type` (the caller-supplied value) rather than `data[0]`, treating `messages` data as `[message_object, metadata_object]`. - Add `_process_messages_payload` helper for the direct AIMessageChunk dispatch path. - Retain backward-compatible fallback for legacy `["event", payload]` tuple format (requires `data[0]` to be a string). Tests added (test_langgraph.py): - test_sse_event_line_sets_current_event - test_sse_empty_line_resets_current_event - test_sse_standard_messages_event_produces_chunk (regression test for #24093) - test_sse_data_without_event_header_is_safe - test_sse_metadata_event_via_header - test_sse_multiple_messages_frames_each_produce_chunk - test_sse_empty_content_ai_chunk_returns_none - test_sse_legacy_tuple_format_still_works - test_sse_full_stream_simulation Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
gambletan
parent
25ee2fb3f9
commit
3018d51a45
2 changed files with 291 additions and 24 deletions
|
|
@ -22,9 +22,22 @@ class LangGraphSSEStreamIterator:
|
|||
Iterator for LangGraph SSE streaming responses.
|
||||
Supports both sync and async iteration.
|
||||
|
||||
LangGraph stream format with stream_mode="messages-tuple":
|
||||
Each SSE event is a tuple: (event_type, data)
|
||||
Common event types: "messages", "metadata"
|
||||
LangGraph SSE wire format (stream_mode="messages-tuple"):
|
||||
|
||||
event: messages
|
||||
data: [<AIMessageChunk>, <metadata_dict>]
|
||||
|
||||
event: metadata
|
||||
data: {"run_id": "...", ...}
|
||||
|
||||
The event type is delivered via the SSE ``event:`` header line.
|
||||
The ``data:`` payload for a ``messages`` event is a two-element array
|
||||
``[message_object, metadata_object]``, NOT a tuple
|
||||
``[event_type_string, payload]``.
|
||||
|
||||
The previous implementation incorrectly treated ``data[0]`` as the event
|
||||
type, causing every real ``messages`` frame to be silently dropped because
|
||||
``data[0]`` is a dict (the AI message), not the string ``"messages"``.
|
||||
"""
|
||||
|
||||
def __init__(self, response: httpx.Response, model: str):
|
||||
|
|
@ -33,6 +46,8 @@ class LangGraphSSEStreamIterator:
|
|||
self.finished = False
|
||||
self.line_iterator = None
|
||||
self.async_line_iterator = None
|
||||
# Tracks the most recent ``event:`` field within the current SSE event block.
|
||||
self._current_event: Optional[str] = None
|
||||
|
||||
def __iter__(self):
|
||||
"""Initialize sync iteration."""
|
||||
|
|
@ -48,15 +63,25 @@ class LangGraphSSEStreamIterator:
|
|||
"""
|
||||
Parse a single SSE line and return a ModelResponse chunk if applicable.
|
||||
|
||||
LangGraph SSE format can vary:
|
||||
- data: [...] (tuple format)
|
||||
- event: ...\ndata: ...
|
||||
Standard SSE multi-line format::
|
||||
|
||||
event: messages
|
||||
data: [...]
|
||||
|
||||
An empty line signals the end of an event block and resets
|
||||
``_current_event`` to ``None``.
|
||||
"""
|
||||
line = line.strip()
|
||||
if not line:
|
||||
# SSE spec: an empty line dispatches the event; reset state.
|
||||
if not line.strip():
|
||||
self._current_event = None
|
||||
return None
|
||||
|
||||
# Handle SSE data lines
|
||||
# Track the event type from the ``event:`` field.
|
||||
if line.startswith("event:"):
|
||||
self._current_event = line[6:].strip()
|
||||
return None
|
||||
|
||||
# Handle SSE data lines.
|
||||
if line.startswith("data:"):
|
||||
json_str = line[5:].strip()
|
||||
if not json_str:
|
||||
|
|
@ -64,32 +89,53 @@ class LangGraphSSEStreamIterator:
|
|||
|
||||
try:
|
||||
data = json.loads(json_str)
|
||||
return self._process_data(data)
|
||||
return self._process_data(data, event_type=self._current_event)
|
||||
except json.JSONDecodeError:
|
||||
verbose_logger.debug(f"Skipping non-JSON SSE line: {line[:100]}")
|
||||
return None
|
||||
|
||||
return None
|
||||
|
||||
def _process_data(self, data) -> Optional[ModelResponseStream]:
|
||||
def _process_data(
|
||||
self, data, event_type: Optional[str] = None
|
||||
) -> Optional[ModelResponseStream]:
|
||||
"""
|
||||
Process parsed data from SSE stream.
|
||||
Process parsed data from an SSE ``data:`` line.
|
||||
|
||||
LangGraph uses tuple format: [event_type, payload]
|
||||
Dispatch is based on *event_type* (the value of the preceding
|
||||
``event:`` header), not on the contents of *data* itself.
|
||||
|
||||
LangGraph ``messages`` event payload::
|
||||
|
||||
[<AIMessageChunk dict>, <metadata dict>]
|
||||
|
||||
LangGraph ``metadata`` event payload::
|
||||
|
||||
{"run_id": "...", ...}
|
||||
"""
|
||||
# Handle tuple format: ["messages", ...]
|
||||
if isinstance(data, list) and len(data) >= 2:
|
||||
event_type = data[0]
|
||||
# --- Standard SSE protocol: event type from the ``event:`` header ---
|
||||
if event_type == "messages":
|
||||
# data = [message_object, metadata_object]
|
||||
if isinstance(data, list) and len(data) >= 1:
|
||||
return self._process_messages_payload(data[0])
|
||||
return None
|
||||
|
||||
if event_type == "metadata":
|
||||
return self._process_metadata_event(data)
|
||||
|
||||
# --- Fallback: legacy/non-standard tuple format [event_type, payload] ---
|
||||
# Retained for backward compatibility with any client that wraps the
|
||||
# data array as ["messages", payload] or ["metadata", payload].
|
||||
if isinstance(data, list) and len(data) >= 2 and isinstance(data[0], str):
|
||||
legacy_event = data[0]
|
||||
payload = data[1]
|
||||
|
||||
if event_type == "messages":
|
||||
if legacy_event == "messages":
|
||||
return self._process_messages_event(payload)
|
||||
elif event_type == "metadata":
|
||||
# Metadata event, might contain usage info
|
||||
elif legacy_event == "metadata":
|
||||
return self._process_metadata_event(payload)
|
||||
|
||||
# Handle dict format (alternative response format)
|
||||
elif isinstance(data, dict):
|
||||
# --- Dict format (alternative / non-streaming-style response) ---
|
||||
if isinstance(data, dict):
|
||||
if "content" in data:
|
||||
return self._create_content_chunk(data.get("content", ""))
|
||||
elif "messages" in data:
|
||||
|
|
@ -101,11 +147,28 @@ class LangGraphSSEStreamIterator:
|
|||
|
||||
return None
|
||||
|
||||
def _process_messages_payload(self, msg: object) -> Optional[ModelResponseStream]:
|
||||
"""
|
||||
Extract content from a single LangGraph AIMessageChunk dict.
|
||||
|
||||
This is the direct ``data[0]`` element from a ``messages`` SSE event.
|
||||
"""
|
||||
if not isinstance(msg, dict):
|
||||
return None
|
||||
|
||||
msg_type = msg.get("type", "")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if msg_type in ("ai", "AIMessageChunk") and content:
|
||||
return self._create_content_chunk(content)
|
||||
|
||||
return None
|
||||
|
||||
def _process_messages_event(self, payload) -> Optional[ModelResponseStream]:
|
||||
"""
|
||||
Process a messages event from the stream.
|
||||
Process a messages payload in the legacy tuple format.
|
||||
|
||||
payload format: [[message_object, metadata], ...]
|
||||
Legacy payload format: [[message_object, metadata], ...]
|
||||
"""
|
||||
if isinstance(payload, list):
|
||||
for item in payload:
|
||||
|
|
|
|||
|
|
@ -171,3 +171,207 @@ def test_langgraph_provider_detection():
|
|||
assert provider == "langgraph"
|
||||
assert model == "agent"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests for LangGraphSSEStreamIterator — issue #24093
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_iterator(model: str = "agent"):
|
||||
from litellm.llms.langgraph.chat.sse_iterator import LangGraphSSEStreamIterator
|
||||
|
||||
# Build a bare instance without needing a real httpx.Response.
|
||||
iterator = LangGraphSSEStreamIterator.__new__(LangGraphSSEStreamIterator)
|
||||
iterator.model = model
|
||||
iterator.finished = False
|
||||
iterator.line_iterator = None
|
||||
iterator.async_line_iterator = None
|
||||
iterator._current_event = None
|
||||
return iterator
|
||||
|
||||
|
||||
_AI_CHUNK = {
|
||||
"content": "Hello, world!",
|
||||
"additional_kwargs": {},
|
||||
"response_metadata": {"finish_reason": "stop"},
|
||||
"type": "AIMessageChunk",
|
||||
"name": None,
|
||||
"id": "lc_run--test",
|
||||
"tool_calls": [],
|
||||
"invalid_tool_calls": [],
|
||||
"usage_metadata": None,
|
||||
"tool_call_chunks": [],
|
||||
"chunk_position": None,
|
||||
}
|
||||
|
||||
_METADATA = {
|
||||
"created_by": "system",
|
||||
"run_id": "019d04b0-f4bf-7842-a109-bccee3c1bf33",
|
||||
"thread_id": "d475429f-e595-48dc-a7fa-066afaabf43c",
|
||||
}
|
||||
|
||||
|
||||
def test_sse_event_line_sets_current_event():
|
||||
"""``event:`` lines are tracked in _current_event."""
|
||||
it = _make_iterator()
|
||||
result = it._parse_sse_line("event: messages")
|
||||
assert result is None # event lines don't produce output
|
||||
assert it._current_event == "messages"
|
||||
|
||||
|
||||
def test_sse_empty_line_resets_current_event():
|
||||
"""An empty line (SSE event boundary) resets _current_event."""
|
||||
it = _make_iterator()
|
||||
it._current_event = "messages"
|
||||
it._parse_sse_line("")
|
||||
assert it._current_event is None
|
||||
|
||||
|
||||
def test_sse_standard_messages_event_produces_chunk():
|
||||
"""
|
||||
The standard LangGraph SSE frame::
|
||||
|
||||
event: messages
|
||||
data: [<AIMessageChunk>, <metadata>]
|
||||
|
||||
must produce a content chunk.
|
||||
|
||||
Regression (issue #24093): the data array is NOT ``[event_type, payload]``
|
||||
but ``[msg_dict, meta_dict]``. The old code treated ``data[0]`` as the
|
||||
event type and therefore never matched ``"messages"``, silently dropping
|
||||
every streaming frame.
|
||||
"""
|
||||
import json
|
||||
|
||||
it = _make_iterator()
|
||||
it._parse_sse_line("event: messages")
|
||||
assert it._current_event == "messages"
|
||||
|
||||
data_line = "data: " + json.dumps([_AI_CHUNK, _METADATA])
|
||||
chunk = it._parse_sse_line(data_line)
|
||||
|
||||
assert chunk is not None, (
|
||||
"Expected a content chunk from a standard SSE messages event, got None. "
|
||||
"This is the regression from issue #24093."
|
||||
)
|
||||
assert len(chunk.choices) == 1
|
||||
assert chunk.choices[0].delta.content == "Hello, world!"
|
||||
|
||||
|
||||
def test_sse_data_without_event_header_is_safe():
|
||||
"""
|
||||
A ``data:`` line whose payload is [dict, dict] but has no preceding
|
||||
``event:`` line must not crash and must not accidentally match the legacy
|
||||
tuple path (which requires data[0] to be a string).
|
||||
"""
|
||||
import json
|
||||
|
||||
it = _make_iterator()
|
||||
# No event: header — _current_event is None
|
||||
data_line = "data: " + json.dumps([_AI_CHUNK, _METADATA])
|
||||
chunk = it._parse_sse_line(data_line)
|
||||
assert chunk is None
|
||||
|
||||
|
||||
def test_sse_metadata_event_via_header():
|
||||
"""
|
||||
Standard SSE ``event: metadata`` frame with a dict payload containing
|
||||
``run_id`` should set finished=True and return a final chunk.
|
||||
"""
|
||||
import json
|
||||
|
||||
it = _make_iterator()
|
||||
it._parse_sse_line("event: metadata")
|
||||
chunk = it._parse_sse_line("data: " + json.dumps(_METADATA))
|
||||
|
||||
assert chunk is not None
|
||||
assert chunk.choices[0].finish_reason == "stop"
|
||||
assert it.finished is True
|
||||
|
||||
|
||||
def test_sse_multiple_messages_frames_each_produce_chunk():
|
||||
"""Consecutive SSE events should each produce an independent chunk."""
|
||||
import json
|
||||
|
||||
it = _make_iterator()
|
||||
collected = []
|
||||
|
||||
for text in ("The", " answer", " is 42."):
|
||||
msg = dict(_AI_CHUNK)
|
||||
msg["content"] = text
|
||||
it._parse_sse_line("event: messages")
|
||||
chunk = it._parse_sse_line("data: " + json.dumps([msg, _METADATA]))
|
||||
it._parse_sse_line("") # SSE event boundary
|
||||
if chunk is not None:
|
||||
collected.append(chunk)
|
||||
|
||||
assert len(collected) == 3
|
||||
assert [c.choices[0].delta.content for c in collected] == ["The", " answer", " is 42."]
|
||||
|
||||
|
||||
def test_sse_empty_content_ai_chunk_returns_none():
|
||||
"""AIMessageChunk with empty content should not produce a chunk."""
|
||||
import json
|
||||
|
||||
it = _make_iterator()
|
||||
msg = dict(_AI_CHUNK)
|
||||
msg["content"] = ""
|
||||
it._parse_sse_line("event: messages")
|
||||
chunk = it._parse_sse_line("data: " + json.dumps([msg, _METADATA]))
|
||||
assert chunk is None
|
||||
|
||||
|
||||
def test_sse_legacy_tuple_format_still_works():
|
||||
"""
|
||||
Backward-compat: some clients may use the legacy ``["messages", payload]``
|
||||
data format without a preceding ``event:`` line. The fallback path must
|
||||
continue to handle this.
|
||||
"""
|
||||
import json
|
||||
|
||||
it = _make_iterator()
|
||||
# No event: header — _current_event is None; data[0] is a string → legacy path
|
||||
legacy_payload = [_AI_CHUNK]
|
||||
data_line = "data: " + json.dumps(["messages", legacy_payload])
|
||||
chunk = it._parse_sse_line(data_line)
|
||||
|
||||
assert chunk is not None
|
||||
assert chunk.choices[0].delta.content == "Hello, world!"
|
||||
|
||||
|
||||
def test_sse_full_stream_simulation():
|
||||
"""
|
||||
Simulate a realistic multi-frame LangGraph SSE stream and verify that all
|
||||
content frames are collected and the final metadata chunk has finish_reason.
|
||||
"""
|
||||
import json
|
||||
|
||||
it = _make_iterator()
|
||||
collected = []
|
||||
|
||||
# Three content frames
|
||||
for text in ("The answer", " is", " 42."):
|
||||
msg = dict(_AI_CHUNK)
|
||||
msg["content"] = text
|
||||
for line in [
|
||||
"event: messages",
|
||||
"data: " + json.dumps([msg, _METADATA]),
|
||||
"",
|
||||
]:
|
||||
chunk = it._parse_sse_line(line)
|
||||
if chunk is not None:
|
||||
collected.append(chunk)
|
||||
|
||||
# Final metadata frame
|
||||
for line in ["event: metadata", "data: " + json.dumps(_METADATA), ""]:
|
||||
chunk = it._parse_sse_line(line)
|
||||
if chunk is not None:
|
||||
collected.append(chunk)
|
||||
|
||||
content_chunks = [c for c in collected if c.choices[0].finish_reason is None]
|
||||
final_chunks = [c for c in collected if c.choices[0].finish_reason == "stop"]
|
||||
|
||||
assert len(content_chunks) == 3
|
||||
assert "".join(c.choices[0].delta.content for c in content_chunks) == "The answer is 42."
|
||||
assert len(final_chunks) == 1
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue