From 13cdf3a147f01ef50118e7c0bac5714c77d2dbfe Mon Sep 17 00:00:00 2001
From: Kayvan Zahiri <123409429+Kayvan-Zahiri@users.noreply.github.com>
Date: Sun, 18 Oct 2026 11:49:08 -0700
Subject: [PATCH] fix(llm_gateway): accept SSE data fields without a space

---
 assemblyai/llm_gateway/v1/api.py           |  7 +-
 assemblyai/llm_gateway/v1/async_api.py     |  7 +-
 tests/unit/test_llm_gateway_sse_spacing.py | 79 ++++++++++++++++++++++
 3 files changed, 89 insertions(+), 4 deletions(-)
 create mode 100644 tests/unit/test_llm_gateway_sse_spacing.py

diff --git a/assemblyai/llm_gateway/v1/api.py b/assemblyai/llm_gateway/v1/api.py
index db5f03aa..adf55f93 100644
--- a/assemblyai/llm_gateway/v1/api.py
+++ b/assemblyai/llm_gateway/v1/api.py
@@ -134,9 +134,12 @@ def stream_completion(
             raise _error_from_response(response)
 
         for line in response.iter_lines():
-            if not line.startswith("data: "):
+            if not line.startswith("data:"):
+                continue
+            # The space after an SSE field's colon is optional.
+            payload = line[len("data:") :].strip()
+            if not payload:
                 continue
-            payload = line[len("data: ") :].strip()
             if payload == "[DONE]":
                 break
             yield models.LLMGatewayCompletionChunk.parse_obj(json.loads(payload))
diff --git a/assemblyai/llm_gateway/v1/async_api.py b/assemblyai/llm_gateway/v1/async_api.py
index 3088ccb2..141556de 100644
--- a/assemblyai/llm_gateway/v1/async_api.py
+++ b/assemblyai/llm_gateway/v1/async_api.py
@@ -100,9 +100,12 @@ async def stream_completion(
             raise _error_from_response(response)
 
         async for line in response.aiter_lines():
-            if not line.startswith("data: "):
+            if not line.startswith("data:"):
+                continue
+            # The space after an SSE field's colon is optional.
+            payload = line[len("data:") :].strip()
+            if not payload:
                 continue
-            payload = line[len("data: ") :].strip()
             if payload == "[DONE]":
                 break
             yield models.LLMGatewayCompletionChunk.parse_obj(json.loads(payload))
diff --git a/tests/unit/test_llm_gateway_sse_spacing.py b/tests/unit/test_llm_gateway_sse_spacing.py
new file mode 100644
index 00000000..ca6475ba
--- /dev/null
+++ b/tests/unit/test_llm_gateway_sse_spacing.py
@@ -0,0 +1,79 @@
+import json
+
+import pytest
+from pytest_httpx import HTTPXMock, IteratorStream
+
+import assemblyai as aai
+
+
+def _response(data_prefix, done_prefix):
+    chunk = {
+        "id": "chatcmpl-test",
+        "object": "chat.completion.chunk",
+        "created": 1,
+        "model": "test-model",
+        "choices": [{"index": 0, "delta": {"content": "Hello"}, "finish_reason": None}],
+    }
+    # A following frame makes ignoring [DONE] observable, rather than relying
+    # on the response ending immediately after the sentinel.
+    after_done = {
+        **chunk,
+        "choices": [{"index": 0, "delta": {"content": "after done"}}],
+    }
+    return IteratorStream(
+        [
+            b": keepalive\r\nevent: message\r\n\r\ndata:\r\n\r\n",
+            (data_prefix + json.dumps(chunk) + "\r\n\r\n").encode(),
+            (done_prefix + "[DONE]\r\n\r\n").encode(),
+            ("data: " + json.dumps(after_done) + "\r\n\r\n").encode(),
+        ]
+    )
+
+
+@pytest.mark.parametrize("data_prefix", ["data:", "data: "])
+@pytest.mark.parametrize("done_prefix", ["data:", "data: "])
+def test_stream_accepts_optional_data_space(
+    httpx_mock: HTTPXMock, data_prefix, done_prefix
+):
+    client = aai.Client(api_key="test")
+    try:
+        httpx_mock.add_response(
+            method="POST",
+            url=client.settings.llm_gateway_base_url + "/v1/chat/completions",
+            headers={"content-type": "text/event-stream"},
+            stream=_response(data_prefix, done_prefix),
+        )
+        chunks = list(
+            aai.LLMGateway(client=client).chat.completions.create(
+                model="test-model",
+                messages=[{"role": "user", "content": "Hi"}],
+                stream=True,
+            )
+        )
+        assert [chunk.choices[0].delta.content for chunk in chunks] == ["Hello"]
+        assert client.last_response.is_closed
+    finally:
+        client.http_client.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("data_prefix", ["data:", "data: "])
+@pytest.mark.parametrize("done_prefix", ["data:", "data: "])
+async def test_async_stream_accepts_optional_data_space(
+    httpx_mock: HTTPXMock, data_prefix, done_prefix
+):
+    async with aai.AsyncClient(api_key="test") as client:
+        httpx_mock.add_response(
+            method="POST",
+            url=client.settings.llm_gateway_base_url + "/v1/chat/completions",
+            headers={"content-type": "text/event-stream"},
+            stream=_response(data_prefix, done_prefix),
+        )
+        stream = await aai.AsyncLLMGateway(client=client).chat.completions.create(
+            model="test-model",
+            messages=[{"role": "user", "content": "Hi"}],
+            stream=True,
+        )
+        chunks = [chunk async for chunk in stream]
+        assert [chunk.choices[0].delta.content for chunk in chunks] == ["Hello"]
+        assert client.last_response.is_closed