fix: parse Server-Sent Events in Agent Engine streaming responses

agentplatform/_genai/_agent_engines_utils.py

@@ -1888,6 +1888,39 @@ async def _method(self, **kwargs) -> Any:  # type: ignore[no-untyped-def]
     return _method  # type: ignore[return-value]
 
 
+_SSE_DATA_PREFIX = "data:"
+_STREAMABLE_CONTENT_TYPES = ("application/json", "text/event-stream")
+
+
+def _strip_sse_framing(line: str) -> str:
+    """Returns the payload of a Server-Sent Events `data:` line.
+
+    Streaming responses are newline-delimited JSON. A response may instead
+    arrive as Server-Sent Events, in which case each JSON object is wrapped in a
+    `data:` frame; removing that framing here lets both shapes be parsed the
+    same way
+    (https://github.com/googleapis/python-aiplatform/issues/5586).
+
+    A serialized JSON value never begins with `data:` -- it begins with `{`,
+    `[`, `"`, a digit, or one of `true`/`false`/`null` -- so stripping the
+    prefix unconditionally cannot corrupt a newline-delimited JSON response. A
+    chunk whose value is the string `data: hello` is serialized as
+    `"data: hello"`, with the quote first.
+
+    Args:
+        line: A single line of the response body.
+
+    Returns:
+        The line with any SSE `data:` framing removed.
+    """
+    line = line.rstrip("\r")
+    if not line.startswith(_SSE_DATA_PREFIX):
+        return line
+    # The single space after the colon is optional per the SSE specification.
+    payload = line[len(_SSE_DATA_PREFIX) :]
+    return payload[1:] if payload.startswith(" ") else payload
+
+
 def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterator[Any]:
     """Converts the body of the HTTP Response message to JSON format.
 
@@ -1904,6 +1937,9 @@ def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterat
 
     # Handle the case of multiple dictionaries delimited by newlines.
     for line in http_response.body.split("\n"):
+        # Strip before the emptiness check so the blank line that terminates an
+        # SSE frame, and a `data:` line with an empty payload, are both skipped.
+        line = _strip_sse_framing(line)
         if line:
             try:
                 line = json.loads(line)
@@ -1931,7 +1967,11 @@ def _yield_parsed_json_from_httpbody(body: httpbody_pb2.HttpBody) -> Iterator[An
     content_type = getattr(body, "content_type", None)
     data = getattr(body, "data", None)
 
-    if content_type is None or data is None or "application/json" not in content_type:
+    if (
+        content_type is None
+        or data is None
+        or not any(t in content_type for t in _STREAMABLE_CONTENT_TYPES)
+    ):
         yield body
         return
 
@@ -1948,6 +1988,9 @@ def _yield_parsed_json_from_httpbody(body: httpbody_pb2.HttpBody) -> Iterator[An
 
     # Handle the case of multiple dictionaries delimited by newlines.
     for line in utf8_data.split("\n"):
+        # Strip before the emptiness check so the blank line that terminates an
+        # SSE frame, and a `data:` line with an empty payload, are both skipped.
+        line = _strip_sse_framing(line)
         if line:
             try:
                 line = json.loads(line)

tests/unit/agentplatform/genai/test_agent_engines.py

@@ -1805,6 +1805,60 @@ def test_yield_parsed_json_from_httpbody(self, obj, expected):
         got = list(_agent_engines_utils._yield_parsed_json_from_httpbody(obj))
         assert got == expected
 
+
+    # pytest does not allow absl.testing.parameterized.named_parameters.
+    @pytest.mark.parametrize(
+        "obj, expected",
+        [
+            (
+                # "sse_single_event",
+                genai_types.HttpResponse(body='data: {"a": 1}\n\n'),
+                [{"a": 1}],
+            ),
+            (
+                # "sse_multiple_events",
+                genai_types.HttpResponse(
+                    body='data: {"a": 1}\n\ndata: {"a": 2}\n\n'
+                ),
+                [{"a": 1}, {"a": 2}],
+            ),
+            (
+                # "sse_no_space_after_colon",
+                genai_types.HttpResponse(body='data:{"a": 1}\n\n'),
+                [{"a": 1}],
+            ),
+            (
+                # "sse_crlf_line_endings",
+                genai_types.HttpResponse(body='data: {"a": 1}\r\n\r\n'),
+                [{"a": 1}],
+            ),
+            (
+                # "sse_empty_data_line_is_skipped",
+                genai_types.HttpResponse(body='data:\n\ndata: {"a": 1}\n\n'),
+                [{"a": 1}],
+            ),
+            (
+                # "json_string_value_beginning_with_data_is_untouched",
+                genai_types.HttpResponse(body='"data: hello"'),
+                ["data: hello"],
+            ),
+        ],
+    )
+    def test_to_parsed_json_server_sent_events(self, obj, expected):
+        """An SSE-framed response is parsed into the same objects as NDJSON."""
+        assert list(_agent_engines_utils._yield_parsed_json(obj)) == expected
+
+    def test_yield_parsed_json_from_httpbody_event_stream(self):
+        """The gRPC path parses SSE instead of yielding the raw proto."""
+        body = httpbody_pb2.HttpBody(
+            content_type="text/event-stream",
+            data=b'data: {"a": 1}\n\ndata: {"a": 2}\n\n',
+        )
+        assert list(_agent_engines_utils._yield_parsed_json_from_httpbody(body)) == [
+            {"a": 1},
+            {"a": 2},
+        ]
+
     def test_yield_parsed_json_from_httpbody_non_json_content_type(self):
         body = httpbody_pb2.HttpBody(content_type="text/plain", data=b"hello")
         assert list(_agent_engines_utils._yield_parsed_json_from_httpbody(body)) == [

vertexai/_genai/_agent_engines_utils.py

@@ -2007,6 +2007,39 @@ async def _method(self, **kwargs) -> Any:  # type: ignore[no-untyped-def]
     _wrap_a2a_operation = _wrap_a2a_operation_v03
 
 
+_SSE_DATA_PREFIX = "data:"
+_STREAMABLE_CONTENT_TYPES = ("application/json", "text/event-stream")
+
+
+def _strip_sse_framing(line: str) -> str:
+    """Returns the payload of a Server-Sent Events `data:` line.
+
+    Streaming responses are newline-delimited JSON. A response may instead
+    arrive as Server-Sent Events, in which case each JSON object is wrapped in a
+    `data:` frame; removing that framing here lets both shapes be parsed the
+    same way
+    (https://github.com/googleapis/python-aiplatform/issues/5586).
+
+    A serialized JSON value never begins with `data:` -- it begins with `{`,
+    `[`, `"`, a digit, or one of `true`/`false`/`null` -- so stripping the
+    prefix unconditionally cannot corrupt a newline-delimited JSON response. A
+    chunk whose value is the string `data: hello` is serialized as
+    `"data: hello"`, with the quote first.
+
+    Args:
+        line: A single line of the response body.
+
+    Returns:
+        The line with any SSE `data:` framing removed.
+    """
+    line = line.rstrip("\r")
+    if not line.startswith(_SSE_DATA_PREFIX):
+        return line
+    # The single space after the colon is optional per the SSE specification.
+    payload = line[len(_SSE_DATA_PREFIX) :]
+    return payload[1:] if payload.startswith(" ") else payload
+
+
 def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterator[Any]:
     """Converts the body of the HTTP Response message to JSON format.
 
@@ -2023,6 +2056,9 @@ def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterat
 
     # Handle the case of multiple dictionaries delimited by newlines.
     for line in http_response.body.split("\n"):
+        # Strip before the emptiness check so the blank line that terminates an
+        # SSE frame, and a `data:` line with an empty payload, are both skipped.
+        line = _strip_sse_framing(line)
         if line:
             try:
                 line = json.loads(line)

vertexai/agent_engines/_utils.py

@@ -305,6 +305,40 @@ def to_json_serializable_autogen_object(
     return _autogen_run_response_protocol_to_dict(obj)
 
 
+
+_SSE_DATA_PREFIX = "data:"
+_STREAMABLE_CONTENT_TYPES = ("application/json", "text/event-stream")
+
+
+def _strip_sse_framing(line: str) -> str:
+    """Returns the payload of a Server-Sent Events `data:` line.
+
+    Streaming responses are newline-delimited JSON. A response may instead
+    arrive as Server-Sent Events, in which case each JSON object is wrapped in a
+    `data:` frame; removing that framing here lets both shapes be parsed the
+    same way
+    (https://github.com/googleapis/python-aiplatform/issues/5586).
+
+    A serialized JSON value never begins with `data:` -- it begins with `{`,
+    `[`, `"`, a digit, or one of `true`/`false`/`null` -- so stripping the
+    prefix unconditionally cannot corrupt a newline-delimited JSON response. A
+    chunk whose value is the string `data: hello` is serialized as
+    `"data: hello"`, with the quote first.
+
+    Args:
+        line: A single line of the response body.
+
+    Returns:
+        The line with any SSE `data:` framing removed.
+    """
+    line = line.rstrip("\r")
+    if not line.startswith(_SSE_DATA_PREFIX):
+        return line
+    # The single space after the colon is optional per the SSE specification.
+    payload = line[len(_SSE_DATA_PREFIX) :]
+    return payload[1:] if payload.startswith(" ") else payload
+
+
 def yield_parsed_json(body: httpbody_pb2.HttpBody) -> Iterable[Any]:
     """Converts the contents of the httpbody message to JSON format.
 
@@ -318,7 +352,11 @@ def yield_parsed_json(body: httpbody_pb2.HttpBody) -> Iterable[Any]:
     content_type = getattr(body, "content_type", None)
     data = getattr(body, "data", None)
 
-    if content_type is None or data is None or "application/json" not in content_type:
+    if (
+        content_type is None
+        or data is None
+        or not any(t in content_type for t in _STREAMABLE_CONTENT_TYPES)
+    ):
         yield body
         return
 
@@ -335,6 +373,9 @@ def yield_parsed_json(body: httpbody_pb2.HttpBody) -> Iterable[Any]:
 
     # Handle the case of multiple dictionaries delimited by newlines.
     for line in utf8_data.split("\n"):
+        # Strip before the emptiness check so the blank line that terminates an
+        # SSE frame, and a `data:` line with an empty payload, are both skipped.
+        line = _strip_sse_framing(line)
         if line:
             try:
                 line = json.loads(line)

vertexai/reasoning_engines/_utils.py

@@ -162,6 +162,40 @@ def to_json_serializable_llama_index_object(
     return str(obj)
 
 
+
+_SSE_DATA_PREFIX = "data:"
+_STREAMABLE_CONTENT_TYPES = ("application/json", "text/event-stream")
+
+
+def _strip_sse_framing(line: str) -> str:
+    """Returns the payload of a Server-Sent Events `data:` line.
+
+    Streaming responses are newline-delimited JSON. A response may instead
+    arrive as Server-Sent Events, in which case each JSON object is wrapped in a
+    `data:` frame; removing that framing here lets both shapes be parsed the
+    same way
+    (https://github.com/googleapis/python-aiplatform/issues/5586).
+
+    A serialized JSON value never begins with `data:` -- it begins with `{`,
+    `[`, `"`, a digit, or one of `true`/`false`/`null` -- so stripping the
+    prefix unconditionally cannot corrupt a newline-delimited JSON response. A
+    chunk whose value is the string `data: hello` is serialized as
+    `"data: hello"`, with the quote first.
+
+    Args:
+        line: A single line of the response body.
+
+    Returns:
+        The line with any SSE `data:` framing removed.
+    """
+    line = line.rstrip("\r")
+    if not line.startswith(_SSE_DATA_PREFIX):
+        return line
+    # The single space after the colon is optional per the SSE specification.
+    payload = line[len(_SSE_DATA_PREFIX) :]
+    return payload[1:] if payload.startswith(" ") else payload
+
+
 def yield_parsed_json(body: httpbody_pb2.HttpBody) -> Iterable[Any]:
     """Converts the contents of the httpbody message to JSON format.
 
@@ -175,7 +209,11 @@ def yield_parsed_json(body: httpbody_pb2.HttpBody) -> Iterable[Any]:
     content_type = getattr(body, "content_type", None)
     data = getattr(body, "data", None)
 
-    if content_type is None or data is None or "application/json" not in content_type:
+    if (
+        content_type is None
+        or data is None
+        or not any(t in content_type for t in _STREAMABLE_CONTENT_TYPES)
+    ):
         yield body
         return
 
@@ -192,6 +230,9 @@ def yield_parsed_json(body: httpbody_pb2.HttpBody) -> Iterable[Any]:
 
     # Handle the case of multiple dictionaries delimited by newlines.
     for line in utf8_data.split("\n"):
+        # Strip before the emptiness check so the blank line that terminates an
+        # SSE frame, and a `data:` line with an empty payload, are both skipped.
+        line = _strip_sse_framing(line)
         if line:
             try:
                 line = json.loads(line)