Assemble multiline SSE data before decoding chat chunks

pinecone/_internal/sse.py

@@ -0,0 +1,48 @@
+"""Assemble SSE data fields before decoding application payloads."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterable, AsyncIterator, Iterable, Iterator
+
+
+class _EventData:
+    def __init__(self) -> None:
+        self.lines: list[str] = []
+
+    def feed(self, line: str) -> str | None:
+        if not line:
+            return self.finish()
+        field, separator, value = line.partition(":")
+        if field == "data":
+            self.lines.append(value.removeprefix(" ") if separator else "")
+        return None
+
+    def finish(self) -> str | None:
+        if not self.lines:
+            return None
+        data = "\n".join(self.lines)
+        self.lines.clear()
+        return data
+
+
+def iter_sse_data(lines: Iterable[str]) -> Iterator[str]:
+    event = _EventData()
+    for line in lines:
+        data = event.feed(line)
+        if data:
+            yield data
+    # Preserve support for the backend's final event without a blank separator.
+    data = event.finish()
+    if data:
+        yield data
+
+
+async def aiter_sse_data(lines: AsyncIterable[str]) -> AsyncIterator[str]:
+    event = _EventData()
+    async for line in lines:
+        data = event.feed(line)
+        if data:
+            yield data
+    data = event.finish()
+    if data:
+        yield data

pinecone/async_client/assistants.py

@@ -22,6 +22,7 @@
     ASSISTANT_EVALUATION_BASE_URL,
     DEFAULT_BASE_URL,
 )
+from pinecone._internal.sse import aiter_sse_data
 from pinecone._internal.validation import require_str
 from pinecone.async_client._assistants_legacy import AsyncAssistantsLegacyNamespaceMixin
 from pinecone.errors.exceptions import (
@@ -1815,14 +1816,7 @@ async def _chat_streaming(
             headers={"Content-Type": "application/json"},
             timeout=_stream_timeout(self._config.timeout, timeout),
         ) as response:
-            async for line in response.aiter_lines():
-                if not line:
-                    continue
-                if not line.startswith("data:"):
-                    continue
-                line = line[5:].lstrip()
-                if not line:
-                    continue
+            async for line in aiter_sse_data(response.aiter_lines()):
                 # Unreachable against the 2026-07 backend, which ends the stream
                 # by closing it. Kept because an intermediary that terminates the
                 # SSE conventionally would otherwise reach orjson.loads and raise.
@@ -2013,14 +2007,7 @@ async def _chat_completions_streaming(
             headers={"Content-Type": "application/json"},
             timeout=_stream_timeout(self._config.timeout, timeout),
         ) as response:
-            async for line in response.aiter_lines():
-                if not line:
-                    continue
-                if not line.startswith("data:"):
-                    continue
-                line = line[5:].lstrip()
-                if not line:
-                    continue
+            async for line in aiter_sse_data(response.aiter_lines()):
                 # Unreachable against the 2026-07 backend, which ends the stream
                 # by closing it. Kept because an intermediary that terminates the
                 # SSE conventionally would otherwise reach orjson.loads and raise.

pinecone/client/assistants.py

@@ -19,6 +19,7 @@
     ASSISTANT_EVALUATION_BASE_URL,
     DEFAULT_BASE_URL,
 )
+from pinecone._internal.sse import iter_sse_data
 from pinecone._internal.validation import require_str
 from pinecone.client._assistants_legacy import AssistantsLegacyNamespaceMixin
 from pinecone.errors.exceptions import (
@@ -1923,14 +1924,7 @@ def _chat_streaming(
             headers={"Content-Type": "application/json"},
             timeout=_stream_timeout(self._config.timeout, timeout),
         ) as response:
-            for line in response.iter_lines():
-                if not line:
-                    continue
-                if not line.startswith("data:"):
-                    continue
-                line = line[5:].lstrip()
-                if not line:
-                    continue
+            for line in iter_sse_data(response.iter_lines()):
                 # Unreachable against the 2026-07 backend, which ends the stream
                 # by closing it. Kept because an intermediary that terminates the
                 # SSE conventionally would otherwise reach orjson.loads and raise.
@@ -1981,14 +1975,7 @@ def _chat_completions_streaming(
             headers={"Content-Type": "application/json"},
             timeout=_stream_timeout(self._config.timeout, timeout),
         ) as response:
-            for line in response.iter_lines():
-                if not line:
-                    continue
-                if not line.startswith("data:"):
-                    continue
-                line = line[5:].lstrip()
-                if not line:
-                    continue
+            for line in iter_sse_data(response.iter_lines()):
                 # Unreachable against the 2026-07 backend, which ends the stream
                 # by closing it. Kept because an intermediary that terminates the
                 # SSE conventionally would otherwise reach orjson.loads and raise.

tests/unit/test_async_assistant_chat_parity.py

@@ -616,3 +616,42 @@ async def test_backend_400_error_parity(
     assert async_type is sync_type
     assert async_message == sync_message
     assert BACKEND_400_MESSAGE in async_message
+
+
+@pytest.mark.parametrize("method_name", _STREAMING_METHODS)
+@pytest.mark.parametrize("transport", ["sync", "async"])
+async def test_multiline_sse_data(
+    method_name: str,
+    transport: str,
+    sync_assistants: Assistants,
+    async_assistants: AsyncAssistants,
+    respx_mock: respx.MockRouter,
+) -> None:
+    body = CHAT_SSE if method_name == "chat" else COMPLETION_SSE
+    events = []
+    for line in body.decode().splitlines():
+        if not line.startswith("data:"):
+            continue
+        data = line[5:].strip()
+        if data == "[DONE]":
+            events.append("data: [DONE]\n\n")
+        else:
+            pretty = orjson.dumps(orjson.loads(data), option=orjson.OPT_INDENT_2).decode()
+            events.append(
+                ": heartbeat\nevent: message\n"
+                + "".join(f"data: {part}\n" for part in pretty.splitlines())
+                + "\n"
+            )
+    route = _register_routes(respx_mock)[method_name]
+    kwargs = {**_CALLS[method_name], "stream": True}
+
+    async def collect(content: str | bytes) -> list[Any]:
+        route.mock(return_value=httpx.Response(200, content=content))
+        if transport == "sync":
+            return list(getattr(sync_assistants, method_name)(**kwargs))
+        stream = await getattr(async_assistants, method_name)(**kwargs)
+        return [chunk async for chunk in stream]
+
+    expected = await collect(body)
+    assert expected
+    assert await collect("".join(events)) == expected

tests/unit/test_sse.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+
+import pytest
+
+from pinecone._internal.sse import aiter_sse_data, iter_sse_data
+
+
+@pytest.mark.parametrize(
+    ("lines", "expected"),
+    [
+        (["data: first", "data: second", "", "data: third", ""], ["first\nsecond", "third"]),
+        ([": heartbeat", "event: message", "id: 42", "retry: 1000", ""], []),
+        (["data:", "", "data", ""], []),
+        (["data:  spaced", ""], [" spaced"]),
+        (["data: final"], ["final"]),
+        (["data: first", "data: final"], ["first\nfinal"]),
+        (["", "", "data: one", "", "", "data: two", ""], ["one", "two"]),
+    ],
+)
+async def test_event_data_parity(lines: list[str], expected: list[str]) -> None:
+    async def async_lines() -> AsyncIterator[str]:
+        for line in lines:
+            yield line
+
+    assert list(iter_sse_data(lines)) == expected
+    assert [data async for data in aiter_sse_data(async_lines())] == expected
+
+
+def test_emits_completed_event_before_reading_next_event() -> None:
+    def lines():
+        yield "data: first"
+        yield ""
+        raise AssertionError("read beyond the completed event")
+
+    assert next(iter_sse_data(lines())) == "first"
+
+
+async def test_async_emits_completed_event_before_reading_next_event() -> None:
+    async def lines() -> AsyncIterator[str]:
+        yield "data: first"
+        yield ""
+        raise AssertionError("read beyond the completed event")
+
+    assert await anext(aiter_sse_data(lines())) == "first"