1 files changed, 449 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/google/genai/_replay_api_client.py b/.venv/lib/python3.12/site-packages/google/genai/_replay_api_client.py
new file mode 100644
index 00000000..3af2336d
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/google/genai/_replay_api_client.py
@@ -0,0 +1,449 @@
+# 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 base64
+import copy
+import datetime
+import inspect
+import io
+import json
+import os
+import re
+from typing import Any, Literal, Optional, Union
+
+import google.auth
+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 ._common import BaseModel
+
+
+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'.*-aiplatform.googleapis.com/[^/]+/',
+      '{VERTEX_URL_PREFIX}/',
+      result,
+  )
+  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 _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]]
+  byte_segments: Optional[list[bytes]] = None
+  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 response 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(self.replay_session.model_dump_json(exclude_unset=True, indent=2))
+    self.replay_session = None
+
+  def _record_interaction(
+      self,
+      http_request: HttpRequest,
+      http_response: Union[HttpResponse, errors.APIError, bytes],
+  ):
+    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()),
+          byte_segments=[
+              seg[:100] + b'...' for seg in http_response.byte_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)
+
+    actual_request_body = [request_data_copy]
+    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
+        ],
+        byte_stream=interaction.response.byte_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 = 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}'
+    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: Union[str, io.IOBase], upload_url: str, upload_size: int):
+    if isinstance(file_path, io.IOBase):
+      offset = file_path.tell()
+      content = file_path.read()
+      file_path.seek(offset, os.SEEK_SET)
+      request = HttpRequest(
+          method='POST',
+          url='',
+          data={'bytes': base64.b64encode(content).decode('utf-8')},
+          headers={}
+      )
+    else:
+      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
+
+  def _download_file_request(self, request):
+    self._initialize_replay_session_if_not_loaded()
+    if self._should_call_api():
+      try:
+        result = super()._download_file_request(request)
+      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, result)
+      return result
+    else:
+      return self._build_response_from_replay(request)
+