From dc09b134007c2a8c054b4db6ff09cf779366ebd0 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 27 Nov 2024 14:10:50 -0800
Subject: [PATCH] sdk-py: Fix SSE parsing to split lines only \n \r \r\n per
 SSE spec

---
 libs/sdk-py/langgraph_sdk/client.py |   5 +-
 libs/sdk-py/langgraph_sdk/sse.py    | 106 ++++++++++++++++++++++++++++
 2 files changed, 109 insertions(+), 2 deletions(-)
 create mode 100644 libs/sdk-py/langgraph_sdk/sse.py

diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py
index 6a3bb6c9c..a018bb3b6 100644
--- a/libs/sdk-py/langgraph_sdk/client.py
+++ b/libs/sdk-py/langgraph_sdk/client.py
@@ -59,6 +59,7 @@ from langgraph_sdk.schema import (
     ThreadStatus,
     ThreadUpdateStateResponse,
 )
+from langgraph_sdk.sse import EventSource
 
 logger = logging.getLogger(__name__)
 
@@ -292,7 +293,7 @@ class HttpClient:
                 else:
                     logger.error(f"Error from langgraph-api: {body}", exc_info=e)
                 raise e
-            async for event in sse.aiter_sse():
+            async for event in EventSource(sse.response).aiter_sse():
                 yield StreamPart(
                     event.event, orjson.loads(event.data) if event.data else None
                 )
@@ -2426,7 +2427,7 @@ class SyncHttpClient:
                 else:
                     logger.error(f"Error from langgraph-api: {body}", exc_info=e)
                 raise e
-            for event in sse.iter_sse():
+            for event in EventSource(sse.response).iter_sse():
                 yield StreamPart(
                     event.event, orjson.loads(event.data) if event.data else None
                 )
diff --git a/libs/sdk-py/langgraph_sdk/sse.py b/libs/sdk-py/langgraph_sdk/sse.py
new file mode 100644
index 000000000..b87788387
--- /dev/null
+++ b/libs/sdk-py/langgraph_sdk/sse.py
@@ -0,0 +1,106 @@
+"""Adapted from httpx_sse to split lines on \n, \r, \r\n per the SSE spec."""
+
+import io
+from typing import AsyncIterator, Iterator
+
+import httpx
+import httpx_sse
+import httpx_sse._decoders
+
+
+class BytesLineDecoder:
+    """
+    Handles incrementally reading lines from text.
+
+    Has the same behaviour as the stdllib bytes splitlines,
+    but handling the input iteratively.
+    """
+
+    def __init__(self) -> None:
+        self.buffer = io.BytesIO()
+        self.trailing_cr: bool = False
+
+    def decode(self, text: bytes) -> list[bytes]:
+        # See https://docs.python.org/3/glossary.html#term-universal-newlines
+        NEWLINE_CHARS = b"\n\r"
+
+        # We always push a trailing `\r` into the next decode iteration.
+        if self.trailing_cr:
+            text = b"\r" + text
+            self.trailing_cr = False
+        if text.endswith(b"\r"):
+            self.trailing_cr = True
+            text = text[:-1]
+
+        if not text:
+            # NOTE: the edge case input of empty text doesn't occur in practice,
+            # because other httpx internals filter out this value
+            return []  # pragma: no cover
+
+        trailing_newline = text[-1] in NEWLINE_CHARS
+        lines = text.splitlines()
+
+        if len(lines) == 1 and not trailing_newline:
+            # No new lines, buffer the input and continue.
+            self.buffer.append(lines[0])
+            return []
+
+        if self.buffer:
+            # Include any existing buffer in the first portion of the
+            # splitlines result.
+            lines = [self.buffer.getvalue() + lines[0]] + lines[1:]
+            self.buffer.truncate(0)
+
+        if not trailing_newline:
+            # If the last segment of splitlines is not newline terminated,
+            # then drop it from our output and start a new buffer.
+            self.buffer.write(lines.pop())
+
+        return lines
+
+    def flush(self) -> list[bytes]:
+        if not self.buffer and not self.trailing_cr:
+            return []
+
+        lines = [self.buffer.getvalue()] if self.buffer else []
+        self.buffer.truncate(0)
+        self.trailing_cr = False
+        return lines
+
+
+async def aiter_lines_raw(response: httpx.Response) -> AsyncIterator[bytes]:
+    decoder = BytesLineDecoder()
+    async for chunk in response.aiter_bytes():
+        for line in decoder.decode(chunk):
+            yield line
+    for line in decoder.flush():
+        yield line
+
+
+def iter_lines_raw(response: httpx.Response) -> Iterator[bytes]:
+    decoder = BytesLineDecoder()
+    for chunk in response.iter_bytes():
+        for line in decoder.decode(chunk):
+            yield line
+    for line in decoder.flush():
+        yield line
+
+
+class EventSource(httpx_sse.EventSource):
+    async def aiter_sse(self) -> AsyncIterator[httpx_sse.ServerSentEvent]:
+        self._check_content_type()
+        decoder = httpx_sse._decoders.SSEDecoder()
+        async for line in aiter_lines_raw(self._response):
+            line = line.rstrip(b"\n")
+            sse = decoder.decode(line.decode())
+            if sse is not None:
+                yield sse
+
+    def iter_sse(self) -> Iterator[httpx_sse.ServerSentEvent]:
+        self._check_content_type()
+        decoder = httpx_sse._decoders.SSEDecoder()
+        for line in iter_lines_raw(self._response):
+            line = line.rstrip(b"\n")
+            sse = decoder.decode(line.decode())
+            if sse is not None:
+                yield sse