JSONL stream

sdk/core/corehttp/corehttp/streaming/__init__.py

@@ -0,0 +1,30 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+
+from ._stream import Stream, AsyncStream
+
+
+__all__ = ["Stream", "AsyncStream"]

sdk/core/corehttp/corehttp/streaming/_stream.py

@@ -0,0 +1,149 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+
+from types import TracebackType
+from typing import Iterator, AsyncIterator, TypeVar, Callable, Any, Optional, Type
+
+from typing_extensions import Self
+
+from ..rest import HttpResponse, AsyncHttpResponse
+from .decoders import StreamDecoder, AsyncStreamDecoder
+
+
+ReturnType = TypeVar("ReturnType")
+
+
+class Stream(Iterator[ReturnType]):
+    """Stream class for streaming JSONL or Server-Sent Events (SSE).
+
+    :keyword response: The response object.
+    :paramtype response: ~corehttp.rest.HttpResponse
+    :keyword decoder: A decoder to use for the stream.
+    :paramtype decoder: ~corehttp.streaming.decoders.StreamDecoder
+    :keyword deserialization_callback: A callback that takes JSON and returns a deserialized object.
+    :paramtype deserialization_callback: Callable[[Any], ReturnType]
+    :keyword terminal_event: A terminal event that indicates the end of the SSE stream.
+    :paramtype terminal_event: Optional[str]
+    """
+
+    def __init__(
+        self,
+        *,
+        response: HttpResponse,
+        decoder: StreamDecoder,
+        deserialization_callback: Callable[[Any], ReturnType],
+        terminal_event: Optional[str] = None,
+    ) -> None:
+        self._response = response
+        self._decoder = decoder
+        self._deserialization_callback = deserialization_callback
+        self._terminal_event = terminal_event
+        self._iterator = self._iter_results()
+
+    def __next__(self) -> ReturnType:
+        return self._iterator.__next__()
+
+    def __iter__(self) -> Iterator[ReturnType]:
+        yield from self._iterator
+
+    def _iter_results(self) -> Iterator[ReturnType]:
+        for event in self._decoder.iter_events(self._response.iter_bytes()):
+            if event.data == self._terminal_event:
+                break
+
+            result = self._deserialization_callback(event.json())
+            yield result
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]] = None,
+        exc_value: Optional[BaseException] = None,
+        traceback: Optional[TracebackType] = None,
+    ) -> None:
+        self.close()
+
+    def __enter__(self) -> Self:
+        return self
+
+    def close(self) -> None:
+        self._response.close()
+
+
+class AsyncStream(AsyncIterator[ReturnType]):
+    """AsyncStream class for asynchronously streaming JSONL or Server-Sent Events (SSE).
+
+    :keyword response: The response object.
+    :paramtype response: ~corehttp.rest.AsyncHttpResponse
+    :keyword decoder: A decoder to use for the stream.
+    :paramtype decoder: ~corehttp.streaming.decoders.AsyncStreamDecoder
+    :keyword deserialization_callback: A callback that takes JSON and returns a deserialized object.
+    :paramtype deserialization_callback: Callable[[Any], ReturnType]
+    :keyword terminal_event: A terminal event that indicates the end of the SSE stream.
+    :paramtype terminal_event: Optional[str]
+    """
+
+    def __init__(
+        self,
+        *,
+        response: AsyncHttpResponse,
+        decoder: AsyncStreamDecoder,
+        deserialization_callback: Callable[[Any], ReturnType],
+        terminal_event: Optional[str] = None,
+    ) -> None:
+        self._response = response
+        self._decoder = decoder
+        self._deserialization_callback = deserialization_callback
+        self._terminal_event = terminal_event
+        self._iterator = self._iter_results()
+
+    async def __anext__(self) -> ReturnType:
+        return await self._iterator.__anext__()
+
+    async def __aiter__(self) -> AsyncIterator[ReturnType]:  # pylint: disable=invalid-overridden-method
+        async for item in self._iterator:
+            yield item
+
+    async def _iter_results(self) -> AsyncIterator[ReturnType]:
+        async for event in self._decoder.aiter_events(self._response.iter_bytes()):
+            if event.data == self._terminal_event:
+                break
+
+            result = self._deserialization_callback(event.json())
+            yield result
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]] = None,
+        exc_value: Optional[BaseException] = None,
+        traceback: Optional[TracebackType] = None,
+    ) -> None:
+        await self.close()
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def close(self) -> None:
+        await self._response.close()

sdk/core/corehttp/corehttp/streaming/decoders.py

@@ -0,0 +1,149 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+
+import codecs
+from typing import Iterator, AsyncIterator, Protocol
+
+from typing_extensions import runtime_checkable
+
+from .events import JSONLEvent, EventType
+
+
+@runtime_checkable
+class StreamDecoder(Protocol):
+    """Protocol for stream decoders."""
+
+    def iter_events(self, iter_bytes: Iterator[bytes]) -> Iterator[EventType]:
+        """Iterate over events from a byte iterator.
+
+        :param iter_bytes: An iterator of byte chunks.
+        :type iter_bytes: Iterator[bytes]
+        :return: An iterator of events.
+        """
+        ...
+
+
+@runtime_checkable
+class AsyncStreamDecoder(Protocol):
+    """Protocol for async stream decoders."""
+
+    # Why this isn't async def: https://mypy.readthedocs.io/en/stable/more_types.html#asynchronous-iterators
+    def aiter_events(self, iter_bytes: AsyncIterator[bytes]) -> AsyncIterator[EventType]:
+        """Asynchronously iterate over events from a byte iterator.
+
+        :param iter_bytes: An asynchronous iterator of byte chunks.
+        :type iter_bytes: AsyncIterator[bytes]
+        :return: An asynchronous iterator of events.
+        """
+        ...
+
+
+def iter_lines(iter_bytes: Iterator[bytes]) -> Iterator[str]:
+    """Iterate over lines from a byte iterator.
+
+    :param iter_bytes: An iterator of byte chunks.
+    :type iter_bytes: Iterator[bytes]
+    :rtype: Iterator[str]
+    :return: An iterator of lines.
+    """
+    decoder = codecs.getincrementaldecoder("utf-8")()
+
+    decoded = ""
+    for chunk in iter_bytes:
+        decoded += decoder.decode(chunk)
+        if decoded:
+            decoded_lines = decoded.splitlines()
+            if decoded.endswith(("\n", "\r\n")):
+                yield from decoded.splitlines()
+                decoded = ""
+            else:
+                yield from decoded_lines[:-1]
+                decoded = decoded_lines[-1]
+
+    decoded += decoder.decode(b"", final=True)
+    if decoded:
+        yield from decoded.splitlines()
+
+
+async def aiter_lines(iter_bytes: AsyncIterator[bytes]) -> AsyncIterator[str]:
+    """Iterate over lines from a byte iterator.
+
+    :param iter_bytes: An iterator of byte chunks.
+    :type iter_bytes: Iterator[bytes]
+    :rtype: Iterator[str]
+    :return: An iterator of lines.
+    """
+    decoder = codecs.getincrementaldecoder("utf-8")()
+
+    decoded = ""
+    async for chunk in iter_bytes:
+        decoded += decoder.decode(chunk)
+        if decoded:
+            decoded_lines = decoded.splitlines()
+            if decoded.endswith(("\n", "\r\n")):
+                for line in decoded.splitlines():
+                    yield line
+                decoded = ""
+            else:
+                for line in decoded_lines[:-1]:
+                    yield line
+                decoded = decoded_lines[-1]
+
+    decoded += decoder.decode(b"", final=True)
+    if decoded:
+        for line in decoded.splitlines():
+            yield line
+
+
+class JSONLDecoder:
+    """Decoder for JSON Lines (JSONL) format. https://jsonlines.org/"""
+
+    def iter_events(self, iter_bytes: Iterator[bytes]) -> Iterator[JSONLEvent]:
+        """Iterate over JSONL events from a byte iterator.
+
+        :param iter_bytes: An iterator of byte chunks.
+        :type iter_bytes: Iterator[bytes]
+        :rtype: Iterator[JSONLEvent]
+        :return: An iterator of JSONLEvent objects.
+        """
+
+        yield from (JSONLEvent(data=line) for line in iter_lines(iter_bytes))
+
+
+class AsyncJSONLDecoder:
+    """Asynchronous decoder for JSON Lines (JSONL) format. https://jsonlines.org/"""
+
+    async def aiter_events(self, iter_bytes: AsyncIterator[bytes]) -> AsyncIterator[JSONLEvent]:
+        """Asynchronously iterate over JSONL events from a byte iterator.
+
+        :param iter_bytes: An asynchronous iterator of byte chunks.
+        :type iter_bytes: AsyncIterator[bytes]
+        :rtype: AsyncIterator[JSONLEvent]
+        :return: An asynchronous iterator of JSONLEvent objects.
+        """
+
+        async for line in aiter_lines(iter_bytes):
+            yield JSONLEvent(data=line)

sdk/core/corehttp/corehttp/streaming/events.py

@@ -0,0 +1,67 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+
+import json
+from typing import Any, Protocol
+
+from typing_extensions import runtime_checkable
+
+
+@runtime_checkable
+class EventType(Protocol):
+    """Protocol for event types."""
+
+    data: str
+    """The event data."""
+
+    def json(self) -> Any:
+        """Parse the event data as JSON.
+
+        :return: The parsed JSON data.
+        """
+        ...
+
+
+class JSONLEvent:
+    """Class representing a JSONL Event."""
+
+    data: str
+    """The event data."""
+
+    def __init__(self, *, data: str) -> None:
+        """Create a new JSONL event.
+
+        :keyword str data: The event data.
+        """
+        self.data = data
+
+    def json(self) -> Any:
+        """Parse the event data as JSON.
+
+        :rtype: any
+        :return: The JSON data for the event
+        """
+        return json.loads(self.data)