google-genai 1.7.0__py3-none-any.whl → 1.53.0__py3-none-any.whl

google/genai/_replay_api_client.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Google LLC
+# Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,28 +17,82 @@
 
 import base64
 import copy
-import datetime
+import contextlib
+import enum
 import inspect
 import io
 import json
 import os
 import re
-from typing import Any, Literal, Optional, Union
+from typing import Any, Literal, Optional, Union, Iterator, AsyncIterator
 
 import google.auth
-from requests.exceptions import HTTPError
 
 from . import errors
 from ._api_client import BaseApiClient
-from ._api_client import HttpOptions
 from ._api_client import HttpRequest
 from ._api_client import HttpResponse
 from ._common import BaseModel
+from .types import HttpOptions, HttpOptionsOrDict
+
+
+def to_snake_case(name: str) -> str:
+  """Converts a string from camelCase or PascalCase to snake_case."""
+
+  if not isinstance(name, str):
+    name = str(name)
+  s1 = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', name)
+  return re.sub('([a-z0-9])([A-Z])', r'\1_\2', s1).lower()
+
+
+def _normalize_json_case(obj: Any) -> Any:
+  if isinstance(obj, dict):
+    return {
+        to_snake_case(k): _normalize_json_case(v)
+        for k, v in obj.items()
+    }
+  elif isinstance(obj, list):
+    return [_normalize_json_case(item) for item in obj]
+  elif isinstance(obj, enum.Enum):
+    return obj.value
+  elif isinstance(obj, str):
+    # Python >= 3.14 has a new division by zero error message.
+    if 'division by zero' in obj:
+      return obj.replace(
+          'division by zero', 'integer division or modulo by zero'
+      )
+  return obj
+
+
+def _equals_ignore_key_case(obj1: Any, obj2: Any) -> bool:
+  """Compares two Python objects for equality ignoring key casing.
+
+  Returns:
+      bool: True if the two objects are equal regardless of key casing
+  (camelCase vs. snake_case). For example, the following are considered equal:
+
+  {'my_key': 'my_value'}
+  {'myKey': 'my_value'}
+
+  This also considers enums and strings with the same value as equal.
+  For example, the following are considered equal:
+
+  {'type': <Type.STRING: 'STRING'>}}
+  {'type': 'STRING'}
+  """
+
+  normalized_obj_1 = _normalize_json_case(obj1)
+  normalized_obj_2 = _normalize_json_case(obj2)
+
+  if normalized_obj_1 == normalized_obj_2:
+    return True
+  else:
+    return False
 
 
 def _redact_version_numbers(version_string: str) -> str:
   """Redacts version numbers in the form x.y.z from a string."""
-  return re.sub(r'\d+\.\d+\.\d+', '{VERSION_NUMBER}', version_string)
+  return re.sub(r'\d+\.\d+\.\d+[a-zA-Z0-9]*', '{VERSION_NUMBER}', version_string)
 
 
 def _redact_language_label(language_label: str) -> str:
@@ -46,7 +100,7 @@ def _redact_language_label(language_label: str) -> str:
   return re.sub(r'gl-python/', '{LANGUAGE_LABEL}/', language_label)
 
 
-def _redact_request_headers(headers):
+def _redact_request_headers(headers: dict[str, str]) -> dict[str, str]:
   """Redacts headers that should not be recorded."""
   redacted_headers = {}
   for header_name, header_value in headers.items():
@@ -88,7 +142,7 @@ def _redact_request_url(url: str) -> str:
       result,
   )
   result = re.sub(
-      r'https://generativelanguage.googleapis.com/[^/]+',
+      r'.*generativelanguage.*.googleapis.com/[^/]+',
       '{MLDEV_URL_PREFIX}',
       result,
   )
@@ -109,28 +163,36 @@ def _redact_project_location_path(path: str) -> str:
     return path
 
 
-def _redact_request_body(body: dict[str, object]):
+def _redact_request_body(body: dict[str, object]) -> None:
   """Redacts fields in the request body in place."""
   for key, value in body.items():
     if isinstance(value, str):
       body[key] = _redact_project_location_path(value)
 
 
-def redact_http_request(http_request: HttpRequest):
+def redact_http_request(http_request: HttpRequest) -> None:
   http_request.headers = _redact_request_headers(http_request.headers)
   http_request.url = _redact_request_url(http_request.url)
-  _redact_request_body(http_request.data)
+  if not isinstance(http_request.data, bytes):
+    _redact_request_body(http_request.data)
 
 
-def _current_file_path_and_line():
+def _current_file_path_and_line() -> str:
   """Prints the current file path and line number."""
-  frame = inspect.currentframe().f_back.f_back
-  filepath = inspect.getfile(frame)
-  lineno = frame.f_lineno
-  return f'File: {filepath}, Line: {lineno}'
+  current_frame = inspect.currentframe()
+  if (
+      current_frame is not None
+      and current_frame.f_back is not None
+      and current_frame.f_back.f_back is not None
+  ):
+    frame = current_frame.f_back.f_back
+    filepath = inspect.getfile(frame)
+    lineno = frame.f_lineno
+    return f'File: {filepath}, Line: {lineno}'
+  return ''
 
 
-def _debug_print(message: str):
+def _debug_print(message: str) -> None:
   print(
       'DEBUG (test',
       os.environ.get('PYTEST_CURRENT_TEST'),
@@ -141,6 +203,28 @@ def _debug_print(message: str):
   )
 
 
+def pop_undeterministic_headers(headers: dict[str, str]) -> None:
+  """Remove headers that are not deterministic."""
+  headers.pop('Date', None)  # pytype: disable=attribute-error
+  headers.pop('Server-Timing', None)  # pytype: disable=attribute-error
+
+
+@contextlib.contextmanager
+def _record_on_api_error(client: 'ReplayApiClient', http_request: HttpRequest) -> Iterator[None]:
+  try:
+    yield
+  except errors.APIError as e:
+    client._record_interaction(http_request, e)
+    raise e
+
+@contextlib.asynccontextmanager
+async def _async_record_on_api_error(client: 'ReplayApiClient', http_request: HttpRequest) -> AsyncIterator[None]:
+  try:
+    yield
+  except errors.APIError as e:
+    client._record_interaction(http_request, e)
+    raise e
+
 class ReplayRequest(BaseModel):
   """Represents a single request in a replay."""
 
@@ -160,10 +244,7 @@ class ReplayResponse(BaseModel):
   sdk_response_segments: list[dict[str, object]]
 
   def model_post_init(self, __context: Any) -> None:
-    # Remove headers that are not deterministic so the replay files don't change
-    # every time they are recorded.
-    self.headers.pop('Date', None)
-    self.headers.pop('Server-Timing', None)
+    pop_undeterministic_headers(self.headers)
 
 
 class ReplayInteraction(BaseModel):
@@ -194,6 +275,7 @@ class ReplayApiClient(BaseApiClient):
       project: Optional[str] = None,
       location: Optional[str] = None,
       http_options: Optional[HttpOptions] = None,
+      private: bool = False,
   ):
     super().__init__(
         vertexai=vertexai,
@@ -209,33 +291,34 @@ class ReplayApiClient(BaseApiClient):
           'GOOGLE_GENAI_REPLAYS_DIRECTORY', None
       )
     # Valid replay modes are replay-only or record-and-replay.
-    self.replay_session = None
+    self.replay_session: Union[ReplayFile, None] = None
     self._mode = mode
     self._replay_id = replay_id
+    self._private = private
 
-  def initialize_replay_session(self, replay_id: str):
+  def initialize_replay_session(self, replay_id: str) -> None:
     self._replay_id = replay_id
     self._initialize_replay_session()
 
-  def _get_replay_file_path(self):
+  def _get_replay_file_path(self) -> str:
     return self._generate_file_path_from_replay_id(
         self.replays_directory, self._replay_id
     )
 
-  def _should_call_api(self):
+  def _should_call_api(self) -> bool:
     return self._mode in ['record', 'api'] or (
         self._mode == 'auto'
         and not os.path.isfile(self._get_replay_file_path())
     )
 
-  def _should_update_replay(self):
+  def _should_update_replay(self) -> bool:
     return self._should_call_api() and self._mode != 'api'
 
-  def _initialize_replay_session_if_not_loaded(self):
+  def _initialize_replay_session_if_not_loaded(self) -> None:
     if not self.replay_session:
       self._initialize_replay_session()
 
-  def _initialize_replay_session(self):
+  def _initialize_replay_session(self) -> None:
     _debug_print('Test is using replay id: ' + self._replay_id)
     self._replay_index = 0
     self._sdk_response_index = 0
@@ -256,7 +339,7 @@ class ReplayApiClient(BaseApiClient):
           replay_id=self._replay_id, interactions=[]
       )
 
-  def _generate_file_path_from_replay_id(self, replay_directory, replay_id):
+  def _generate_file_path_from_replay_id(self, replay_directory: Optional[str], replay_id: str) -> str:
     session_parts = replay_id.split('/')
     if len(session_parts) < 3:
       raise ValueError(
@@ -270,7 +353,7 @@ class ReplayApiClient(BaseApiClient):
     path_parts.extend(session_parts)
     return os.path.join(*path_parts) + '.json'
 
-  def close(self):
+  def close(self) -> None:
     if not self._should_update_replay() or not self.replay_session:
       return
     replay_file_path = self._get_replay_file_path()
@@ -283,7 +366,7 @@ class ReplayApiClient(BaseApiClient):
       self,
       http_request: HttpRequest,
       http_response: Union[HttpResponse, errors.APIError, bytes],
-  ):
+  ) -> None:
     if not self._should_update_replay():
       return
     redact_http_request(http_request)
@@ -321,6 +404,8 @@ class ReplayApiClient(BaseApiClient):
       raise ValueError(
           'Unsupported http_response type: ' + str(type(http_response))
       )
+    if self.replay_session is None:
+      raise ValueError('No replay session found.')
     self.replay_session.interactions.append(
         ReplayInteraction(request=request, response=response)
     )
@@ -329,7 +414,9 @@ class ReplayApiClient(BaseApiClient):
       self,
       http_request: HttpRequest,
       interaction: ReplayInteraction,
-  ):
+  ) -> None:
+    _debug_print(f'http_request.url: {http_request.url}')
+    _debug_print(f'interaction.request.url: {interaction.request.url}')
     assert http_request.url == interaction.request.url
     assert http_request.headers == interaction.request.headers, (
         'Request headers mismatch:\n'
@@ -342,26 +429,29 @@ class ReplayApiClient(BaseApiClient):
     request_data_copy = copy.deepcopy(http_request.data)
     # Both the request and recorded request must be redacted before comparing
     # so that the comparison is fair.
-    _redact_request_body(request_data_copy)
+    if not isinstance(request_data_copy, bytes):
+      _redact_request_body(request_data_copy)
 
     actual_request_body = [request_data_copy]
     expected_request_body = interaction.request.body_segments
-    assert actual_request_body == expected_request_body, (
+    assert _equals_ignore_key_case(actual_request_body, expected_request_body), (
         'Request body mismatch:\n'
         f'Actual: {actual_request_body}\n'
         f'Expected: {expected_request_body}'
     )
 
-  def _build_response_from_replay(self, http_request: HttpRequest):
+  def _build_response_from_replay(self, http_request: HttpRequest) -> HttpResponse:
     redact_http_request(http_request)
 
+    if self.replay_session is None:
+      raise ValueError('No replay session found.')
     interaction = self.replay_session.interactions[self._replay_index]
     # Replay is on the right side of the assert so the diff makes more sense.
     self._match_request(http_request, interaction)
     self._replay_index += 1
     self._sdk_response_index = 0
     errors.APIError.raise_for_response(interaction.response)
-    return HttpResponse(
+    http_response = HttpResponse(
         headers=interaction.response.headers,
         response_stream=[
             json.dumps(segment)
@@ -369,17 +459,29 @@ class ReplayApiClient(BaseApiClient):
         ],
         byte_stream=interaction.response.byte_segments,
     )
+    if http_response.response_stream == ['{}']:
+      http_response.response_stream = [""]
+    return http_response
 
-  def _verify_response(self, response_model: BaseModel):
+  def _verify_response(self, response_model: BaseModel) -> None:
     if self._mode == 'api':
       return
+    if not self.replay_session:
+      raise ValueError('No replay session found.')
     # replay_index is advanced in _build_response_from_replay, so we need to -1.
     interaction = self.replay_session.interactions[self._replay_index - 1]
     if self._should_update_replay():
       if isinstance(response_model, list):
         response_model = response_model[0]
-      if response_model and 'http_headers' in response_model.model_fields:
-        response_model.http_headers.pop('Date', None)
+      sdk_response_response = getattr(response_model, 'sdk_http_response', None)
+      if response_model and (
+          sdk_response_response is not None
+      ):
+        headers = getattr(
+            sdk_response_response, 'headers', None
+        )
+        if headers:
+          pop_undeterministic_headers(headers)
       interaction.response.sdk_response_segments.append(
           response_model.model_dump(exclude_none=True)
       )
@@ -387,29 +489,46 @@ class ReplayApiClient(BaseApiClient):
 
     if isinstance(response_model, list):
       response_model = response_model[0]
-    print('response_model: ', response_model.model_dump(exclude_none=True))
+    _debug_print(
+        f'response_model: {response_model.model_dump(exclude_none=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}'
+    # The sdk_http_response.body has format in the string, need to get rid of
+    # the format information before comparing.
+    if isinstance(expected, dict):
+      if 'sdk_http_response' in expected and isinstance(
+          expected['sdk_http_response'], dict
+      ):
+        if 'body' in expected['sdk_http_response']:
+          raw_body = expected['sdk_http_response']['body']
+          _debug_print(f'raw_body length: {len(raw_body)}')
+          _debug_print(f'raw_body: {raw_body}')
+          if isinstance(raw_body, str) and raw_body != '':
+            raw_body = json.loads(raw_body)
+            raw_body = json.dumps(raw_body)
+            expected['sdk_http_response']['body'] = raw_body
+    if not self._private:
+      assert (
+          actual == expected
+      ), f'SDK response mismatch:\nActual: {actual}\nExpected: {expected}'
+    else:
+      _debug_print(f'Expected SDK response mismatch:\nActual: {actual}\nExpected: {expected}')
     self._sdk_response_index += 1
 
   def _request(
       self,
       http_request: HttpRequest,
+      http_options: Optional[HttpOptionsOrDict] = None,
       stream: bool = False,
   ) -> HttpResponse:
     self._initialize_replay_session_if_not_loaded()
     if self._should_call_api():
       _debug_print('api mode request: %s' % http_request)
-      try:
-        result = super()._request(http_request, stream)
-      except errors.APIError as e:
-        self._record_interaction(http_request, e)
-        raise e
+      with _record_on_api_error(self, http_request):
+        result = super()._request(http_request, http_options, stream)
       if stream:
         result_segments = []
         for segment in result.segments():
@@ -428,16 +547,16 @@ class ReplayApiClient(BaseApiClient):
   async def _async_request(
       self,
       http_request: HttpRequest,
+      http_options: Optional[HttpOptionsOrDict] = None,
       stream: bool = False,
   ) -> HttpResponse:
     self._initialize_replay_session_if_not_loaded()
     if self._should_call_api():
       _debug_print('api mode request: %s' % http_request)
-      try:
-        result = await super()._async_request(http_request, stream)
-      except errors.APIError as e:
-        self._record_interaction(http_request, e)
-        raise e
+      async with _async_record_on_api_error(self, http_request):
+        result = await super()._async_request(
+            http_request, http_options, stream
+        )
       if stream:
         result_segments = []
         async for segment in result.async_segments():
@@ -453,7 +572,14 @@ class ReplayApiClient(BaseApiClient):
     else:
       return self._build_response_from_replay(http_request)
 
-  def upload_file(self, file_path: Union[str, io.IOBase], upload_url: str, upload_size: int):
+  def upload_file(
+      self,
+      file_path: Union[str, io.IOBase],
+      upload_url: str,
+      upload_size: int,
+      *,
+      http_options: Optional[HttpOptionsOrDict] = None,
+  ) -> HttpResponse:
     if isinstance(file_path, io.IOBase):
       offset = file_path.tell()
       content = file_path.read()
@@ -470,25 +596,23 @@ class ReplayApiClient(BaseApiClient):
       )
     if self._should_call_api():
       result: Union[str, HttpResponse]
-      try:
-        result = super().upload_file(file_path, upload_url, upload_size)
-      except HTTPError as e:
-        result = HttpResponse(
-            e.response.headers, [json.dumps({'reason': e.response.reason})]
+      with _record_on_api_error(self, request):
+        result = super().upload_file(
+            file_path, upload_url, upload_size, http_options=http_options
         )
-        result.status_code = e.response.status_code
-        raise e
-      self._record_interaction(request, HttpResponse({}, [json.dumps(result)]))
+      self._record_interaction(request, result)
       return result
     else:
-      return self._build_response_from_replay(request).json
+      return self._build_response_from_replay(request)
 
   async def async_upload_file(
       self,
       file_path: Union[str, io.IOBase],
       upload_url: str,
       upload_size: int,
-  ) -> str:
+      *,
+      http_options: Optional[HttpOptionsOrDict] = None,
+  ) -> HttpResponse:
     if isinstance(file_path, io.IOBase):
       offset = file_path.tell()
       content = file_path.read()
@@ -504,55 +628,43 @@ class ReplayApiClient(BaseApiClient):
           method='POST', url='', data={'file_path': file_path}, headers={}
       )
     if self._should_call_api():
-      result: Union[str, HttpResponse]
-      try:
+      result: HttpResponse
+      async with _async_record_on_api_error(self, request):
         result = await super().async_upload_file(
-            file_path, upload_url, upload_size
-        )
-      except HTTPError as e:
-        result = HttpResponse(
-            e.response.headers, [json.dumps({'reason': e.response.reason})]
+            file_path, upload_url, upload_size, http_options=http_options
         )
-        result.status_code = e.response.status_code
-        raise e
-      self._record_interaction(request, HttpResponse({}, [json.dumps(result)]))
+      self._record_interaction(request, result)
       return result
     else:
-      return self._build_response_from_replay(request).json
+      return self._build_response_from_replay(request)
 
-  def download_file(self, path: str, http_options: HttpOptions):
+  def download_file(
+      self, path: str, *, http_options: Optional[HttpOptionsOrDict] = None
+  ) -> Union[HttpResponse, bytes, Any]:
     self._initialize_replay_session_if_not_loaded()
     request = self._build_request(
         'get', path=path, request_dict={}, http_options=http_options
     )
     if self._should_call_api():
-      try:
-        result = super().download_file(path, http_options)
-      except HTTPError as e:
-        result = HttpResponse(
-            e.response.headers, [json.dumps({'reason': e.response.reason})]
-        )
-        result.status_code = e.response.status_code
-        raise e
+      with _record_on_api_error(self, request):
+        result = super().download_file(path, http_options=http_options)
       self._record_interaction(request, result)
       return result
     else:
       return self._build_response_from_replay(request).byte_stream[0]
 
-  async def async_download_file(self, path: str, http_options):
+  async def async_download_file(
+      self, path: str, *, http_options: Optional[HttpOptionsOrDict] = None
+  ) -> Any:
     self._initialize_replay_session_if_not_loaded()
     request = self._build_request(
         'get', path=path, request_dict={}, http_options=http_options
     )
     if self._should_call_api():
-      try:
-        result = await super().async_download_file(path, http_options)
-      except HTTPError as e:
-        result = HttpResponse(
-            e.response.headers, [json.dumps({'reason': e.response.reason})]
+      async with _async_record_on_api_error(self, request):
+        result = await super().async_download_file(
+            path, http_options=http_options
         )
-        result.status_code = e.response.status_code
-        raise e
       self._record_interaction(request, result)
       return result
     else:

google/genai/_test_api_client.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Google LLC
+# Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.