diff --git a/pinecone/_internal/sse.py b/pinecone/_internal/sse.py new file mode 100644 index 000000000..0fee7fc56 --- /dev/null +++ b/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 diff --git a/pinecone/async_client/assistants.py b/pinecone/async_client/assistants.py index 5c15db933..f5df523bc 100644 --- a/pinecone/async_client/assistants.py +++ b/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. diff --git a/pinecone/client/assistants.py b/pinecone/client/assistants.py index 004653483..154b8d2a1 100644 --- a/pinecone/client/assistants.py +++ b/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. diff --git a/tests/unit/test_async_assistant_chat_parity.py b/tests/unit/test_async_assistant_chat_parity.py index 6b06240a7..92d3a54c6 100644 --- a/tests/unit/test_async_assistant_chat_parity.py +++ b/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 diff --git a/tests/unit/test_sse.py b/tests/unit/test_sse.py new file mode 100644 index 000000000..8a5c530bd --- /dev/null +++ b/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"