codekingpro/portable-devtools

1"""Adapted from httpx_sse to split lines on \n, \r, \r\n per the SSE spec."""2 3from __future__ import annotations4 5import contextlib6from collections.abc import AsyncIterator, Iterator7from typing import cast8 9import httpx10import orjson11 12from langgraph_sdk.schema import StreamPart13 14BytesLike = bytes | bytearray | memoryview15 16 17class BytesLineDecoder:18    """19    Handles incrementally reading lines from text.20 21    Has the same behaviour as the stdllib bytes splitlines,22    but handling the input iteratively.23    """24 25    def __init__(self) -> None:26        self.buffer = bytearray()27        self.trailing_cr: bool = False28 29    def decode(self, text: bytes) -> list[BytesLike]:30        # See https://docs.python.org/3/glossary.html#term-universal-newlines31        NEWLINE_CHARS = b"\n\r"32 33        # We always push a trailing `\r` into the next decode iteration.34        if self.trailing_cr:35            text = b"\r" + text36            self.trailing_cr = False37        if text.endswith(b"\r"):38            self.trailing_cr = True39            text = text[:-1]40 41        if not text:42            # NOTE: the edge case input of empty text doesn't occur in practice,43            # because other httpx internals filter out this value44            return []  # pragma: no cover45 46        trailing_newline = text[-1] in NEWLINE_CHARS47        lines = cast(list[BytesLike], text.splitlines())48 49        if len(lines) == 1 and not trailing_newline:50            # No new lines, buffer the input and continue.51            self.buffer.extend(lines[0])52            return []53 54        if self.buffer:55            # Include any existing buffer in the first portion of the56            # splitlines result.57            self.buffer.extend(lines[0])58            lines = cast(list[BytesLike], [self.buffer, *lines[1:]])59            self.buffer = bytearray()60 61        if not trailing_newline:62            # If the last segment of splitlines is not newline terminated,63            # then drop it from our output and start a new buffer.64            self.buffer.extend(lines.pop())65 66        return lines67 68    def flush(self) -> list[BytesLike]:69        if not self.buffer and not self.trailing_cr:70            return []71 72        lines = [self.buffer]73        self.buffer = bytearray()74        self.trailing_cr = False75        return lines76 77 78class SSEDecoder:79    def __init__(self) -> None:80        self._event = ""81        self._data = bytearray()82        self._last_event_id = ""83        self._retry: int | None = None84 85    @property86    def last_event_id(self) -> str | None:87        """Return the last event identifier that was seen."""88 89        return self._last_event_id or None90 91    def decode(self, line: bytes) -> StreamPart | None:92        # See: https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation93 94        if not line:95            if (96                not self._event97                and not self._data98                and not self._last_event_id99                and self._retry is None100            ):101                return None102 103            sse = StreamPart(104                event=self._event,105                data=orjson.loads(self._data) if self._data else None,  # type: ignore[invalid-argument-type]106                id=self.last_event_id,107            )108 109            # NOTE: as per the SSE spec, do not reset last_event_id.110            self._event = ""111            self._data = bytearray()112            self._retry = None113 114            return sse115 116        if line.startswith(b":"):117            return None118 119        fieldname, _, value = line.partition(b":")120 121        if value.startswith(b" "):122            value = value[1:]123 124        if fieldname == b"event":125            self._event = value.decode()126        elif fieldname == b"data":127            self._data.extend(value)128        elif fieldname == b"id":129            if b"\0" in value:130                pass131            else:132                self._last_event_id = value.decode()133        elif fieldname == b"retry":134            with contextlib.suppress(TypeError, ValueError):135                self._retry = int(value)136        else:137            pass  # Field is ignored.138 139        return None140 141 142async def aiter_lines_raw(response: httpx.Response) -> AsyncIterator[BytesLike]:143    decoder = BytesLineDecoder()144    async for chunk in response.aiter_bytes():145        for line in decoder.decode(chunk):146            yield line147    for line in decoder.flush():148        yield line149 150 151def iter_lines_raw(response: httpx.Response) -> Iterator[BytesLike]:152    decoder = BytesLineDecoder()153    for chunk in response.iter_bytes():154        for line in decoder.decode(chunk):155            yield line156    for line in decoder.flush():157        yield line158