feat: Use ser_json_byte val_json_bytes in bytes type public interface

PiperOrigin-RevId: 715040448

Author: qiaodev

google/genai/_api_client.py

@@ -574,13 +574,12 @@ def _verify_response(self, response_model: BaseModel):
     pass
 
 
+# TODO(b/389693448): Cleanup datetime hacks.
 class RequestJsonEncoder(json.JSONEncoder):
   """Encode bytes as strings without modify its content."""
 
   def default(self, o):
-    if isinstance(o, bytes):
-      return o.decode()
-    elif isinstance(o, datetime.datetime):
+    if isinstance(o, datetime.datetime):
       # This Zulu time format is used by the Vertex AI API and the test recorder
       # Using strftime works well, but we want to align with the replay encoder.
       # o.astimezone(datetime.timezone.utc).strftime('%Y-%m-%dT%H:%M:%S.%fZ')

google/genai/_common.py

@@ -189,6 +189,8 @@ class BaseModel(pydantic.BaseModel):
       extra='forbid',
       # This allows us to use arbitrary types in the model. E.g. PIL.Image.
       arbitrary_types_allowed=True,
+      ser_json_bytes='base64',
+      val_json_bytes='base64',
   )
 
   @classmethod
@@ -200,7 +202,10 @@ def _from_response(
     # We will provide another mechanism to allow users to access these fields.
     _remove_extra_fields(cls, response)
     validated_response = cls.model_validate(response)
-    return apply_base64_decoding_for_model(validated_response)
+    return validated_response
+
+  def to_json_dict(self) -> dict[str, object]:
+    return self.model_dump(exclude_none=True, mode='json')
 
 
 def timestamped_unique_name() -> str:
@@ -216,40 +221,21 @@ def timestamped_unique_name() -> str:
 
 def apply_base64_encoding(data: dict[str, object]) -> dict[str, object]:
   """Applies base64 encoding to bytes values in the given data."""
-  return process_bytes_fields(data, encode=True)
-
-
-def apply_base64_decoding(data: dict[str, object]) -> dict[str, object]:
-  """Applies base64 decoding to bytes values in the given data."""
-  return process_bytes_fields(data, encode=False)
-
-
-def apply_base64_decoding_for_model(data: BaseModel) -> BaseModel:
-  d = data.model_dump(exclude_none=True)
-  d = apply_base64_decoding(d)
-  return data.model_validate(d)
-
-
-def process_bytes_fields(data: dict[str, object], encode=True) -> dict[str, object]:
   processed_data = {}
   if not isinstance(data, dict):
     return data
   for key, value in data.items():
     if isinstance(value, bytes):
-      if encode:
-        processed_data[key] = base64.b64encode(value)
-      else:
-        processed_data[key] = base64.b64decode(value)
+      processed_data[key] = base64.urlsafe_b64encode(value).decode('ascii')
     elif isinstance(value, dict):
-      processed_data[key] = process_bytes_fields(value, encode)
+      processed_data[key] = apply_base64_encoding(value)
     elif isinstance(value, list):
-      if encode and all(isinstance(v, bytes) for v in value):
-        processed_data[key] = [base64.b64encode(v) for v in value]
-      elif all(isinstance(v, bytes) for v in value):
-        processed_data[key] = [base64.b64decode(v) for v in value]
+      if all(isinstance(v, bytes) for v in value):
+        processed_data[key] = [
+            base64.urlsafe_b64encode(v).decode('ascii') for v in value
+        ]
       else:
-        processed_data[key] = [process_bytes_fields(v, encode) for v in value]
+        processed_data[key] = [apply_base64_encoding(v) for v in value]
     else:
       processed_data[key] = value
   return processed_data
-

google/genai/_replay_api_client.py

@@ -25,7 +25,6 @@
 from typing import Any, Literal, Optional, Union
 
 import google.auth
-from pydantic import BaseModel
 from requests.exceptions import HTTPError
 
 from . import errors
@@ -34,6 +33,7 @@
 from ._api_client import HttpRequest
 from ._api_client import HttpResponse
 from ._api_client import RequestJsonEncoder
+from ._common import BaseModel
 
 def _redact_version_numbers(version_string: str) -> str:
   """Redacts version numbers in the form x.y.z from a string."""
@@ -264,18 +264,9 @@ def close(self):
     replay_file_path = self._get_replay_file_path()
     os.makedirs(os.path.dirname(replay_file_path), exist_ok=True)
     with open(replay_file_path, 'w') as f:
-      replay_session_dict = self.replay_session.model_dump()
-      # Use for non-utf-8 bytes in image/video... output.
-      for interaction in replay_session_dict['interactions']:
-        segments = []
-        for response in interaction['response']['sdk_response_segments']:
-          segments.append(json.loads(json.dumps(
-              response, cls=ResponseJsonEncoder
-          )))
-        interaction['response']['sdk_response_segments'] = segments
       f.write(
           json.dumps(
-              replay_session_dict, indent=2, cls=RequestJsonEncoder
+              self.replay_session.model_dump(mode='json'), indent=2, cls=ResponseJsonEncoder
           )
       )
     self.replay_session = None
@@ -376,15 +367,8 @@ def _verify_response(self, response_model: BaseModel):
     if isinstance(response_model, list):
       response_model = response_model[0]
     print('response_model: ', response_model.model_dump(exclude_none=True))
-    actual = json.dumps(
-        response_model.model_dump(exclude_none=True),
-        cls=ResponseJsonEncoder,
-        sort_keys=True,
-    )
-    expected = json.dumps(
-        interaction.response.sdk_response_segments[self._sdk_response_index],
-        sort_keys=True,
-    )
+    actual = response_model.model_dump(exclude_none=True, mode='json')
+    expected = interaction.response.sdk_response_segments[self._sdk_response_index]
     assert (
         actual == expected
     ), f'SDK response mismatch:\nActual: {actual}\nExpected: {expected}'
@@ -437,36 +421,12 @@ def upload_file(self, file_path: str, upload_url: str, upload_size: int):
       return self._build_response_from_replay(request).text
 
 
+# TODO(b/389693448): Cleanup datetime hacks.
 class ResponseJsonEncoder(json.JSONEncoder):
   """The replay test json encoder for response.
-
-  We need RequestJsonEncoder and ResponseJsonEncoder because:
-    1. In production, we only need RequestJsonEncoder to help json module
-    to convert non-stringable and stringable types to json string. Especially
-    for bytes type, the value of bytes field is encoded to base64 string so it
-    is always stringable and the RequestJsonEncoder doesn't have to deal with
-    utf-8 JSON broken issue.
-    2. In replay test, we also need ResponseJsonEncoder to help json module
-    convert non-stringable and stringable types to json string. But response
-    object returned from SDK method is different from the request api_client
-    sent to server. For the bytes type, there is no base64 string in response
-    anymore, because SDK handles it internally. So bytes type in Response is
-    non-stringable. The ResponseJsonEncoder uses different encoding
-    strategy than the RequestJsonEncoder to deal with utf-8 JSON broken issue.
   """
   def default(self, o):
-    if isinstance(o, bytes):
-      # Use base64.b64encode() to encode bytes to string so that the media bytes
-      # fields are serializable.
-      # o.decode(encoding='utf-8', errors='replace') doesn't work because it
-      # uses a fixed error string `\ufffd` for all non-utf-8 characters,
-      # which cannot be converted back to original bytes. And other languages
-      # only have the original bytes to compare with.
-      # Since we use base64.b64encoding() in replay test, a change that breaks
-      # native bytes can be captured by
-      # test_compute_tokens.py::test_token_bytes_deserialization.
-      return base64.b64encode(o).decode(encoding='utf-8')
-    elif isinstance(o, datetime.datetime):
+    if isinstance(o, datetime.datetime):
       # dt.isoformat() prints "2024-11-15T23:27:45.624657+00:00"
       # but replay files want "2024-11-15T23:27:45.624657Z"
       if o.isoformat().endswith('+00:00'):

google/genai/_transformers.py

@@ -476,3 +476,17 @@ def t_tuning_job_status(
     return 'JOB_STATE_FAILED'
   else:
     return status
+
+
+# Some fields don't accept url safe base64 encoding.
+# We shouldn't use this transformer if the backend adhere to Cloud Type
+# format https://cloud.google.com/docs/discovery/type-format.
+# TODO(b/389133914): Remove the hack after Vertex backend fix the issue.
+def t_bytes(api_client: _api_client.ApiClient, data: bytes) -> str:
+  if not isinstance(data, bytes):
+    return data
+  if api_client.vertexai:
+    return base64.b64encode(data).decode('ascii')
+  else:
+    return base64.urlsafe_encode(data).decode('ascii')
+

google/genai/models.py

@@ -1547,7 +1547,11 @@ def _Image_to_mldev(
     raise ValueError('gcs_uri parameter is not supported in Google AI.')
 
   if getv(from_object, ['image_bytes']) is not None:
-    setv(to_object, ['bytesBase64Encoded'], getv(from_object, ['image_bytes']))
+    setv(
+        to_object,
+        ['bytesBase64Encoded'],
+        t.t_bytes(api_client, getv(from_object, ['image_bytes'])),
+    )
 
   return to_object
 
@@ -1562,7 +1566,11 @@ def _Image_to_vertex(
     setv(to_object, ['gcsUri'], getv(from_object, ['gcs_uri']))
 
   if getv(from_object, ['image_bytes']) is not None:
-    setv(to_object, ['bytesBase64Encoded'], getv(from_object, ['image_bytes']))
+    setv(
+        to_object,
+        ['bytesBase64Encoded'],
+        t.t_bytes(api_client, getv(from_object, ['image_bytes'])),
+    )
 
   return to_object
 
@@ -3193,7 +3201,11 @@ def _Image_from_mldev(
   to_object = {}
 
   if getv(from_object, ['bytesBase64Encoded']) is not None:
-    setv(to_object, ['image_bytes'], getv(from_object, ['bytesBase64Encoded']))
+    setv(
+        to_object,
+        ['image_bytes'],
+        t.t_bytes(api_client, getv(from_object, ['bytesBase64Encoded'])),
+    )
 
   return to_object
 
@@ -3208,7 +3220,11 @@ def _Image_from_vertex(
     setv(to_object, ['gcs_uri'], getv(from_object, ['gcsUri']))
 
   if getv(from_object, ['bytesBase64Encoded']) is not None:
-    setv(to_object, ['image_bytes'], getv(from_object, ['bytesBase64Encoded']))
+    setv(
+        to_object,
+        ['image_bytes'],
+        t.t_bytes(api_client, getv(from_object, ['bytesBase64Encoded'])),
+    )
 
   return to_object
 

google/genai/tests/client/test_json_encoder.py

@@ -1,3 +1,4 @@
+import datetime
 # Copyright 2024 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -19,21 +20,26 @@
 
 
 def test_json_encoder():
-  assert json.dumps({'key': 'value'}, cls=RequestJsonEncoder) == '{"key": "value"}'
-  assert json.dumps({'key': b'value'}, cls=RequestJsonEncoder) == '{"key": "value"}'
+  date_value = datetime.datetime.fromtimestamp(
+      1736397612, tz=datetime.timezone.utc
+  )
+  assert (
+      json.dumps({'key': date_value}, cls=RequestJsonEncoder)
+      == '{"key": "2025-01-09T04:40:12Z"}'
+  )
   assert (
-      json.dumps({'nested': {'key': 'value'}}, cls=RequestJsonEncoder)
-      == '{"nested": {"key": "value"}}'
+      json.dumps({'nested': {'key': date_value}}, cls=RequestJsonEncoder)
+      == '{"nested": {"key": "2025-01-09T04:40:12Z"}}'
   )
   assert (
-      json.dumps({'nested': {'key': b'value'}}, cls=RequestJsonEncoder)
-      == '{"nested": {"key": "value"}}'
+      json.dumps({'nested': {'key': date_value}}, cls=RequestJsonEncoder)
+      == '{"nested": {"key": "2025-01-09T04:40:12Z"}}'
   )
   assert (
-      json.dumps({'list': ['value', 'value']}, cls=RequestJsonEncoder)
-      == '{"list": ["value", "value"]}'
+      json.dumps({'list': [date_value, date_value]}, cls=RequestJsonEncoder)
+      == '{"list": ["2025-01-09T04:40:12Z", "2025-01-09T04:40:12Z"]}'
   )
   assert (
-      json.dumps({'list': [b'value', b'value']}, cls=RequestJsonEncoder)
-      == '{"list": ["value", "value"]}'
+      json.dumps({'list': [date_value, date_value]}, cls=RequestJsonEncoder)
+      == '{"list": ["2025-01-09T04:40:12Z", "2025-01-09T04:40:12Z"]}'
   )

google/genai/tests/live/test_live.py

@@ -198,7 +198,7 @@ async def test_async_session_send_realtime_input(
       api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
   )
   realtime_input = types.LiveClientRealtimeInput(
-      media_chunks=[types.Blob(data='000000', mime_type='audio/pcm')]
+      media_chunks=[types.Blob(data='MDAwMDAw', mime_type='audio/pcm')]
   )
   await session.send(input=realtime_input)
   mock_websocket.send.assert_called_once()

google/genai/tests/models/test_generate_content_part.py

@@ -309,9 +309,6 @@
                 ),
             ],
         ),
-        # Base64 string is invalid input.
-        exception_if_vertex='400',
-        exception_if_mldev='400',
     ),
     pytest_helper.TestTableItem(
         name='test_union_none_part',
@@ -613,15 +610,14 @@ def test_from_function_call_response(client):
 
 @pytest.mark.asyncio
 async def test_image_base64_stream_async(client):
-  with pytest.raises(errors.ClientError):
-    async for part in client.aio.models.generate_content_stream(
-        model='gemini-1.5-flash-001',
-        contents=[
-            'What is this image about?',
-            {'inline_data': {'data': image_string, 'mimeType': 'image/png'}},
-        ],
-    ):
-      pass
+  async for part in client.aio.models.generate_content_stream(
+      model='gemini-1.5-flash-001',
+      contents=[
+          'What is this image about?',
+          {'inline_data': {'data': image_string, 'mimeType': 'image/png'}},
+      ],
+  ):
+    pass
 
 
 # function_call and function_response are tested in generate_content_tools.py

google/genai/tests/pytest_helper.py

@@ -167,7 +167,7 @@ def setup(
           # exclude_unset=True is needed to avoid warnings.
           # See https://github.com/pydantic/pydantic/issues/6467.
           json.dumps(
-              test_table_file.model_dump(exclude_unset=True, by_alias=True),
+              test_table_file.model_dump(exclude_unset=True, by_alias=True, mode='json'),
               indent=2,
               cls=ResponseJsonEncoder,
           )