google-genai 0.0.1__py3-none-any.whl

google/genai/_replay_api_client.py

@@ -0,0 +1,478 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Replay API client."""
+
+import copy
+import inspect
+import json
+import os
+import re
+import datetime
+from typing import Any, Literal, Optional, Union
+
+import google.auth
+from pydantic import BaseModel
+from requests.exceptions import HTTPError
+
+from . import errors
+from ._api_client import ApiClient
+from ._api_client import HttpOptions
+from ._api_client import HttpRequest
+from ._api_client import HttpResponse
+from ._api_client import RequestJsonEncoder
+
+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)
+
+
+def _redact_language_label(language_label: str) -> str:
+  """Removed because replay requests are used for all languages."""
+  return re.sub(r'gl-python/', '{LANGUAGE_LABEL}/', language_label)
+
+
+def _redact_request_headers(headers):
+  """Redacts headers that should not be recorded."""
+  redacted_headers = {}
+  for header_name, header_value in headers.items():
+    if header_name.lower() == 'x-goog-api-key':
+      redacted_headers[header_name] = '{REDACTED}'
+    elif header_name.lower() == 'user-agent':
+      redacted_headers[header_name] = _redact_language_label(
+          _redact_version_numbers(header_value)
+      )
+    elif header_name.lower() == 'x-goog-api-client':
+      redacted_headers[header_name] = _redact_language_label(
+          _redact_version_numbers(header_value)
+      )
+    else:
+      redacted_headers[header_name] = header_value
+  return redacted_headers
+
+
+def _redact_request_url(url: str) -> str:
+  # Redact all the url parts before the resource name, so the test can work
+  # against any project, location, version, or whether it's EasyGCP.
+  result = re.sub(
+      r'.*/projects/[^/]+/locations/[^/]+/',
+      '{VERTEX_URL_PREFIX}/',
+      url,
+  )
+  result = re.sub(
+      r'https://generativelanguage.googleapis.com/[^/]+',
+      '{MLDEV_URL_PREFIX}',
+      result,
+  )
+  return result
+
+
+def _redact_project_location_path(path: str) -> str:
+  # Redact a field in the request that is known to vary based on project and
+  # location.
+  if 'projects/' in path and 'locations/' in path:
+    result = re.sub(
+        r'projects/[^/]+/locations/[^/]+/',
+        '{PROJECT_AND_LOCATION_PATH}/',
+        path,
+    )
+    return result
+  else:
+    return path
+
+
+def _redact_request_body(body: dict[str, object]) -> dict[str, object]:
+  for key, value in body.items():
+    if isinstance(value, str):
+      body[key] = _redact_project_location_path(value)
+
+
+def redact_http_request(http_request: HttpRequest):
+  http_request.headers = _redact_request_headers(http_request.headers)
+  http_request.url = _redact_request_url(http_request.url)
+  _redact_request_body(http_request.data)
+
+
+def process_bytes_fields(data: dict[str, object]):
+  """Converts bytes fields to strings.
+
+  This function doesn't modify the content of data dict.
+  """
+  if not isinstance(data, dict):
+    return data
+  for key, value in data.items():
+    if isinstance(value, bytes):
+      data[key] = value.decode()
+    elif isinstance(value, dict):
+      process_bytes_fields(value)
+    elif isinstance(value, list):
+      if all(isinstance(v, bytes) for v in value):
+        data[key] = [v.decode() for v in value]
+      else:
+        data[key] = [process_bytes_fields(v) for v in value]
+    else:
+      data[key] = value
+  return data
+
+
+def _current_file_path_and_line():
+  """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}'
+
+
+def _debug_print(message: str):
+  print(
+      'DEBUG (test',
+      os.environ.get('PYTEST_CURRENT_TEST'),
+      ')',
+      _current_file_path_and_line(),
+      ':\n    ',
+      message,
+  )
+
+
+class ReplayRequest(BaseModel):
+  """Represents a single request in a replay."""
+
+  method: str
+  url: str
+  headers: dict[str, str]
+  body_segments: list[dict[str, object]]
+
+
+class ReplayResponse(BaseModel):
+  """Represents a single response in a replay."""
+
+  status_code: int = 200
+  headers: dict[str, str]
+  body_segments: list[dict[str, object]]
+  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)
+
+
+class ReplayInteraction(BaseModel):
+  """Represents a single interaction, request and response in a replay."""
+
+  request: ReplayRequest
+  response: ReplayResponse
+
+
+class ReplayFile(BaseModel):
+  """Represents a recorded session."""
+
+  replay_id: str
+  interactions: list[ReplayInteraction]
+
+
+class ReplayApiClient(ApiClient):
+  """For integration testing, send recorded responese or records a response."""
+
+  def __init__(
+      self,
+      mode: Literal['record', 'replay', 'auto', 'api'],
+      replay_id: str,
+      replays_directory: Optional[str] = None,
+      vertexai: bool = False,
+      api_key: Optional[str] = None,
+      credentials: Optional[google.auth.credentials.Credentials] = None,
+      project: Optional[str] = None,
+      location: Optional[str] = None,
+      http_options: Optional[HttpOptions] = None,
+  ):
+    super().__init__(
+        vertexai=vertexai,
+        api_key=api_key,
+        credentials=credentials,
+        project=project,
+        location=location,
+        http_options=http_options,
+    )
+    self.replays_directory = replays_directory
+    if not self.replays_directory:
+      self.replays_directory = os.environ.get(
+          'GOOGLE_GENAI_REPLAYS_DIRECTORY', None
+      )
+    # Valid replay modes are replay-only or record-and-replay.
+    self.replay_session = None
+    self._mode = mode
+    self._replay_id = replay_id
+
+  def initialize_replay_session(self, replay_id: str):
+    self._replay_id = replay_id
+    self._initialize_replay_session()
+
+  def _get_replay_file_path(self):
+    return self._generate_file_path_from_replay_id(
+        self.replays_directory, self._replay_id
+    )
+
+  def _should_call_api(self):
+    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):
+    return self._should_call_api() and self._mode != 'api'
+
+  def _initialize_replay_session_if_not_loaded(self):
+    if not self.replay_session:
+      self._initialize_replay_session()
+
+  def _initialize_replay_session(self):
+    _debug_print('Test is using replay id: ' + self._replay_id)
+    self._replay_index = 0
+    self._sdk_response_index = 0
+    replay_file_path = self._get_replay_file_path()
+    # This should not be triggered from the constructor.
+    replay_file_exists = os.path.isfile(replay_file_path)
+    if self._mode == 'replay' and not replay_file_exists:
+      raise ValueError(
+          'Replay files do not exist for replay id: ' + self._replay_id
+      )
+
+    if self._mode in ['replay', 'auto'] and replay_file_exists:
+      with open(replay_file_path, 'r') as f:
+        self.replay_session = ReplayFile.model_validate(json.loads(f.read()))
+
+    if self._should_update_replay():
+      self.replay_session = ReplayFile(
+          replay_id=self._replay_id, interactions=[]
+      )
+
+  def _generate_file_path_from_replay_id(self, replay_directory, replay_id):
+    session_parts = replay_id.split('/')
+    if len(session_parts) < 3:
+      raise ValueError(
+          f'{replay_id}: Session ID must be in the format of'
+          ' module/function/[vertex|mldev]'
+      )
+    if replay_directory is None:
+      path_parts = []
+    else:
+      path_parts = [replay_directory]
+    path_parts.extend(session_parts)
+    return os.path.join(*path_parts) + '.json'
+
+  def close(self):
+    if not self._should_update_replay() or not self.replay_session:
+      return
+    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:
+      f.write(
+          json.dumps(
+              self.replay_session.model_dump(), indent=2, cls=RequestJsonEncoder
+          )
+      )
+    self.replay_session = None
+
+  def _record_interaction(
+      self,
+      http_request: HttpRequest,
+      http_response: Union[HttpResponse, errors.APIError],
+  ):
+    if not self._should_update_replay():
+      return
+    redact_http_request(http_request)
+    request = ReplayRequest(
+        method=http_request.method,
+        url=http_request.url,
+        headers=http_request.headers,
+        body_segments=[http_request.data],
+    )
+    if isinstance(http_response, HttpResponse):
+      response = ReplayResponse(
+          headers=dict(http_response.headers),
+          body_segments=list(http_response.segments()),
+          status_code=http_response.status_code,
+          sdk_response_segments=[],
+      )
+    else:
+      response = ReplayResponse(
+          headers=dict(http_response.response.headers),
+          body_segments=[http_response._to_replay_record()],
+          status_code=http_response.code,
+          sdk_response_segments=[],
+      )
+    self.replay_session.interactions.append(
+        ReplayInteraction(request=request, response=response)
+    )
+
+  def _match_request(
+      self,
+      http_request: HttpRequest,
+      interaction: ReplayInteraction,
+  ):
+    assert http_request.url == interaction.request.url
+    assert http_request.headers == interaction.request.headers, (
+        'Request headers mismatch:\n'
+        f'Actual: {http_request.headers}\n'
+        f'Expected: {interaction.request.headers}'
+    )
+    assert http_request.method == interaction.request.method
+
+    # Sanitize the request body, rewrite any fields that vary.
+    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)
+
+    # Need to call dumps() and loads() to convert dict bytes values to strings.
+    # Because the expected_request_body dict never contains bytes values.
+    actual_request_body = [
+        json.loads(json.dumps(request_data_copy, cls=RequestJsonEncoder))
+    ]
+    expected_request_body = interaction.request.body_segments
+    assert 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):
+    redact_http_request(http_request)
+
+    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(
+        headers=interaction.response.headers,
+        response_stream=[
+            json.dumps(segment)
+            for segment in interaction.response.body_segments
+        ],
+    )
+
+  def _verify_response(self, response_model: BaseModel):
+    if self._mode == 'api':
+      return
+    # 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]
+      interaction.response.sdk_response_segments.append(
+          response_model.model_dump(exclude_none=True)
+      )
+      return
+
+    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,
+    )
+    assert (
+        actual == expected
+    ), f'SDK response mismatch:\nActual: {actual}\nExpected: {expected}'
+    self._sdk_response_index += 1
+
+  def _request(
+      self,
+      http_request: HttpRequest,
+      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
+      if stream:
+        result_segments = []
+        for segment in result.segments():
+          result_segments.append(json.dumps(segment))
+        result = HttpResponse(result.headers, result_segments)
+        self._record_interaction(http_request, result)
+        # Need to return a RecordedResponse that rebuilds the response
+        # segments since the stream has been consumed.
+      else:
+        self._record_interaction(http_request, result)
+      _debug_print('api mode result: %s' % result.text)
+      return result
+    else:
+      return self._build_response_from_replay(http_request)
+
+  def upload_file(self, file_path: str, upload_url: str, upload_size: int):
+    request = HttpRequest(
+        method='POST', url='', data={'file_path': file_path}, headers={}
+    )
+    if self._should_call_api():
+      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})]
+        )
+        result.status_code = e.response.status_code
+        raise e
+      self._record_interaction(request, HttpResponse({}, [json.dumps(result)]))
+      return result
+    else:
+      return self._build_response_from_replay(request).text
+
+
+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 error replace because response need to be serialized with bytes
+      # string, not base64 string. Otherwise, we cannot tell the response is
+      # already decoded from base64 or not from the replay file.
+      return o.decode(encoding='utf-8', errors='replace')
+    elif 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'):
+        return o.isoformat().replace('+00:00', 'Z')
+      else:
+        return o.isoformat()
+    else:
+      return super().default(o)

google/genai/_test_api_client.py

@@ -0,0 +1,149 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import asyncio
+import time
+from unittest.mock import MagicMock, patch
+import pytest
+from .api_client import ApiClient
+
+
+@patch('genai.api_client.ApiClient._build_request')
+@patch('genai.api_client.ApiClient._request')
+def test_request_streamed_non_blocking(mock_request, mock_build_request):
+  api_client = ApiClient(api_key='test_api_key')
+  http_method = 'GET'
+  path = 'test/path'
+  request_dict = {'key': 'value'}
+
+  mock_http_request = MagicMock()
+  mock_build_request.return_value = mock_http_request
+
+  def delayed_segments():
+    chunks = ['{"chunk": 1}', '{"chunk": 2}', '{"chunk": 3}']
+    for chunk in chunks:
+      time.sleep(0.1)  # 100ms delay
+      yield chunk
+
+  mock_response = MagicMock()
+  mock_response.segments.side_effect = delayed_segments
+  mock_request.return_value = mock_response
+
+  chunks = []
+  start_time = time.time()
+  for chunk in api_client.request_streamed(http_method, path, request_dict):
+    chunks.append(chunk)
+    assert len(chunks) <= 3
+  end_time = time.time()
+
+  mock_build_request.assert_called_once_with(
+      http_method, path, request_dict, None
+  )
+  mock_request.assert_called_once_with(mock_http_request, stream=True)
+  assert chunks == ['{"chunk": 1}', '{"chunk": 2}', '{"chunk": 3}']
+  assert end_time - start_time > 0.3
+
+
+@patch('genai.api_client.ApiClient._build_request')
+@patch('genai.api_client.ApiClient._async_request')
+@pytest.mark.asyncio
+async def test_async_request(mock_async_request, mock_build_request):
+  api_client = ApiClient(api_key='test_api_key')
+  http_method = 'GET'
+  path = 'test/path'
+  request_dict = {'key': 'value'}
+
+  mock_http_request = MagicMock()
+  mock_build_request.return_value = mock_http_request
+
+  class MockResponse:
+
+    def __init__(self, text):
+      self.text = text
+
+  async def delayed_response(http_request, stream):
+    await asyncio.sleep(0.1)  # 100ms delay
+    return MockResponse('value')
+
+  mock_async_request.side_effect = delayed_response
+
+  async_coroutine1 = api_client.async_request(http_method, path, request_dict)
+  async_coroutine2 = api_client.async_request(http_method, path, request_dict)
+  async_coroutine3 = api_client.async_request(http_method, path, request_dict)
+
+  start_time = time.time()
+  results = await asyncio.gather(
+      async_coroutine1, async_coroutine2, async_coroutine3
+  )
+  end_time = time.time()
+
+  mock_build_request.assert_called_with(http_method, path, request_dict, None)
+  assert mock_build_request.call_count == 3
+  mock_async_request.assert_called_with(
+      http_request=mock_http_request, stream=False
+  )
+  assert mock_async_request.call_count == 3
+  assert results == ['value', 'value', 'value']
+  assert 0.1 <= end_time - start_time < 0.15
+
+
+@patch('genai.api_client.ApiClient._build_request')
+@patch('genai.api_client.ApiClient._async_request')
+@pytest.mark.asyncio
+async def test_async_request_streamed_non_blocking(
+    mock_async_request, mock_build_request
+):
+  api_client = ApiClient(api_key='test_api_key')
+  http_method = 'GET'
+  path = 'test/path'
+  request_dict = {'key': 'value'}
+
+  mock_http_request = MagicMock()
+  mock_build_request.return_value = mock_http_request
+
+  class MockResponse:
+
+    def __init__(self, segments):
+      self._segments = segments
+
+    # should mock async generator here but source code combines sync and async streaming in one segment method.
+    # TODO: fix the above
+    def segments(self):
+      for segment in self._segments:
+        time.sleep(0.1)  # 100ms delay
+        yield segment
+
+  async def delayed_response(http_request, stream):
+    return MockResponse(['{"chunk": 1}', '{"chunk": 2}', '{"chunk": 3}'])
+
+  mock_async_request.side_effect = delayed_response
+
+  chunks = []
+  start_time = time.time()
+  async for chunk in api_client.async_request_streamed(
+      http_method, path, request_dict
+  ):
+    chunks.append(chunk)
+    assert len(chunks) <= 3
+  end_time = time.time()
+
+  mock_build_request.assert_called_once_with(
+      http_method, path, request_dict, None
+  )
+  mock_async_request.assert_called_once_with(
+      http_request=mock_http_request, stream=True
+  )
+  assert chunks == ['{"chunk": 1}', '{"chunk": 2}', '{"chunk": 3}']
+  assert end_time - start_time > 0.3