fix: GenAI SDK client - parsing json body of HTTP response to yield JSON

PiperOrigin-RevId: 788922635

Author: yeesian

tests/unit/vertexai/genai/test_agent_engines.py

@@ -1023,6 +1023,46 @@ def test_scan_with_explicit_ignore_modules(self):
             "pydantic": "1.11.1",
         }
 
+    # pytest does not allow absl.testing.parameterized.named_parameters.
+    @pytest.mark.parametrize(
+        "obj, expected",
+        [
+            (
+                # "valid_json",
+                genai_types.HttpResponse(body='{"a": 1, "b": "hello"}'),
+                [{"a": 1, "b": "hello"}],
+            ),
+            (
+                # "invalid_json",
+                genai_types.HttpResponse(body='{"a": 1, "b": "hello"'),
+                ['{"a": 1, "b": "hello"'],  # returns the unparsed string
+            ),
+            (
+                # "nil_body",
+                genai_types.HttpResponse(body=None),
+                [None],
+            ),
+            (
+                # "empty_data",
+                genai_types.HttpResponse(body=""),
+                [None],
+            ),
+            (
+                # "nested_json",
+                genai_types.HttpResponse(body='{"a": {"b": 1}}'),
+                [{"a": {"b": 1}}],
+            ),
+            (
+                # "multiline_json",
+                genai_types.HttpResponse(body='{"a": {"b": 1}}\n{"a": {"b": 2}}'),
+                [{"a": {"b": 1}}, {"a": {"b": 2}}],
+            ),
+        ],
+    )
+    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.mark.usefixtures("google_auth_mock")
 class TestAgentEngine:

vertexai/_genai/_agent_engines_utils.py

@@ -45,6 +45,7 @@
 import proto
 
 from google.api_core import exceptions
+from google.genai import types as google_genai_types
 from google.protobuf import struct_pb2
 from google.protobuf import json_format
 
@@ -1369,15 +1370,17 @@ def _method(self: genai_types.AgentEngine, **kwargs) -> Iterator[Any]:  # type:
             raise ValueError("api_client is not initialized.")
         if not self.api_resource:
             raise ValueError("api_resource is not initialized.")
-        for response in self.api_client._stream_query(
+        for http_response in self.api_client._stream_query(
             name=self.api_resource.name,
             config={
                 "class_method": method_name,
                 "input": kwargs,
                 "include_all_fields": True,
             },
         ):
-            yield response
+            for line in _yield_parsed_json(http_response=http_response):
+                if line is not None:
+                    yield line
 
     return _method
 
@@ -1405,14 +1408,40 @@ async def _method(self: genai_types.AgentEngine, **kwargs) -> AsyncIterator[Any]
             raise ValueError("api_client is not initialized.")
         if not self.api_resource:
             raise ValueError("api_resource is not initialized.")
-        for response in self.api_client._stream_query(
+        for http_response in self.api_client._stream_query(
             name=self.api_resource.name,
             config={
                 "class_method": method_name,
                 "input": kwargs,
                 "include_all_fields": True,
             },
         ):
-            yield response
+            for line in _yield_parsed_json(http_response=http_response):
+                if line is not None:
+                    yield line
 
     return _method
+
+
+def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterator[Any]:
+    """Converts the body of the HTTP Response message to JSON format.
+
+    Args:
+        http_response (google.genai.types.HttpResponse):
+            Required. The httpbody body to be converted to JSON object(s).
+
+    Yields:
+        Any: A JSON object or line of the original body or None.
+    """
+    if not http_response.body:
+        yield None
+        return
+
+    # Handle the case of multiple dictionaries delimited by newlines.
+    for line in http_response.body.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