fix: fix the async_stream_query method. Restore yield_parsed_json method lost during vertex -> agentplatform migration

PiperOrigin-RevId: 939654762

Author: vertex-sdk-bot

agentplatform/_genai/_agent_engines_utils.py

@@ -53,6 +53,7 @@
 
 from google.api_core import exceptions
 from google.genai import types as google_genai_types
+from google.api import httpbody_pb2
 from google.protobuf import struct_pb2
 from google.protobuf import json_format
 
@@ -1915,6 +1916,50 @@ def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterat
             yield line
 
 
+def _yield_parsed_json_from_httpbody(body: httpbody_pb2.HttpBody) -> Iterator[Any]:
+    """Converts the contents of an `HttpBody` proto message to JSON format.
+
+    Unlike `_yield_parsed_json`, which parses a `google.genai.types.HttpResponse`
+    (with a `body` attribute), this helper parses the gRPC `httpbody_pb2.HttpBody`
+    protos yielded by `stream_query_reasoning_engine`, which expose
+    `content_type` and `data` instead.
+
+    Args:
+        body (httpbody_pb2.HttpBody):
+            Required. The httpbody proto to be converted to JSON object(s).
+
+    Yields:
+        Any: A JSON object, a line of the original body, or the original body if
+            it is not JSON, or None.
+    """
+    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:
+        yield body
+        return
+
+    try:
+        utf8_data = data.decode("utf-8")
+    except Exception as e:
+        logger.warning(f"Failed to decode data: {data!r}. Exception: {e}")
+        yield body
+        return
+
+    if not utf8_data:
+        yield None
+        return
+
+    # Handle the case of multiple dictionaries delimited by newlines.
+    for line in utf8_data.split("\n"):
+        if line:
+            try:
+                line = json.loads(line)
+            except Exception as e:
+                logger.warning(f"failed to parse json: {line}. Exception: {e}")
+            yield line
+
+
 def _validate_resource_limits_or_raise(resource_limits: dict[str, str]) -> None:
     """Validates the resource limits.
 

agentplatform/agent_engines/_agent_engines.py

@@ -1635,7 +1635,9 @@ def _method(self, **kwargs) -> Iterable[Any]:
             ),
         )
         for chunk in response:
-            for parsed_json in _agent_engines_utils._yield_parsed_json(chunk):
+            for parsed_json in _agent_engines_utils._yield_parsed_json_from_httpbody(
+                chunk
+            ):
                 if parsed_json is not None:
                     yield parsed_json
 
@@ -1669,7 +1671,9 @@ async def _method(self, **kwargs) -> AsyncIterable[Any]:
             ),
         )
         for chunk in response:
-            for parsed_json in _agent_engines_utils._yield_parsed_json(chunk):
+            for parsed_json in _agent_engines_utils._yield_parsed_json_from_httpbody(
+                chunk
+            ):
                 if parsed_json is not None:
                     yield parsed_json
 

tests/unit/agentplatform/genai/test_agent_engines.py

@@ -37,6 +37,7 @@
 from agentplatform._genai import types as _genai_types
 from google.genai import client as genai_client
 from google.genai import types as genai_types
+from google.api import httpbody_pb2
 import pytest
 
 
@@ -1767,6 +1768,49 @@ def test_to_parsed_json(self, obj, expected):
         for got, want in zip(_agent_engines_utils._yield_parsed_json(obj), expected):
             assert got == want
 
+    # pytest does not allow absl.testing.parameterized.named_parameters.
+    @pytest.mark.parametrize(
+        "obj, expected",
+        [
+            (
+                # "valid_json",
+                httpbody_pb2.HttpBody(
+                    content_type="application/json", data=b'{"a": 1, "b": "hello"}'
+                ),
+                [{"a": 1, "b": "hello"}],
+            ),
+            (
+                # "invalid_json",
+                httpbody_pb2.HttpBody(
+                    content_type="application/json", data=b'{"a": 1, "b": "hello"'
+                ),
+                ['{"a": 1, "b": "hello"'],  # returns the unparsed string
+            ),
+            (
+                # "empty_data",
+                httpbody_pb2.HttpBody(content_type="application/json", data=b""),
+                [None],
+            ),
+            (
+                # "multiline_json",
+                httpbody_pb2.HttpBody(
+                    content_type="application/json",
+                    data=b'{"a": {"b": 1}}\n{"a": {"b": 2}}',
+                ),
+                [{"a": {"b": 1}}, {"a": {"b": 2}}],
+            ),
+        ],
+    )
+    def test_yield_parsed_json_from_httpbody(self, obj, expected):
+        got = list(_agent_engines_utils._yield_parsed_json_from_httpbody(obj))
+        assert got == expected
+
+    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)) == [
+            body
+        ]
+
     def test_create_base64_encoded_tarball(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             test_file_path = os.path.join(tmpdir, "test_file.txt")