google-genai 1.6.0__py3-none-any.whl → 1.8.0__py3-none-any.whl

google/genai/_common.py

@@ -20,7 +20,7 @@ import datetime
 import enum
 import functools
 import typing
-from typing import Union
+from typing import Any, Union
 import uuid
 import warnings
 
@@ -93,7 +93,7 @@ def set_value_by_path(data, keys, value):
     data[keys[-1]] = value
 
 
-def get_value_by_path(data: object, keys: list[str]):
+def get_value_by_path(data: Any, keys: list[str]):
   """Examples:
 
   get_value_by_path({'a': {'b': v}}, ['a', 'b'])
@@ -128,7 +128,7 @@ def get_value_by_path(data: object, keys: list[str]):
   return data
 
 
-def convert_to_dict(obj: dict[str, object]) -> dict[str, object]:
+def convert_to_dict(obj: object) -> Any:
   """Recursively converts a given object to a dictionary.
 
   If the object is a Pydantic model, it uses the model's `model_dump()` method.
@@ -137,7 +137,9 @@ def convert_to_dict(obj: dict[str, object]) -> dict[str, object]:
     obj: The object to convert.
 
   Returns:
-    A dictionary representation of the object.
+    A dictionary representation of the object, a list of objects if a list is
+    passed, or the object itself if it is not a dictionary, list, or Pydantic
+    model.
   """
   if isinstance(obj, pydantic.BaseModel):
     return obj.model_dump(exclude_none=True)
@@ -150,7 +152,7 @@ def convert_to_dict(obj: dict[str, object]) -> dict[str, object]:
 
 
 def _remove_extra_fields(
-    model: pydantic.BaseModel, response: dict[str, object]
+    model: Any, response: dict[str, object]
 ) -> None:
   """Removes extra fields from the response that are not in the model.
 

google/genai/_replay_api_client.py

@@ -119,7 +119,8 @@ def _redact_request_body(body: dict[str, object]):
 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)
+  if not isinstance(http_request.data, bytes):
+    _redact_request_body(http_request.data)
 
 
 def _current_file_path_and_line():
@@ -321,6 +322,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)
     )
@@ -342,7 +345,8 @@ 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
@@ -352,9 +356,11 @@ class ReplayApiClient(BaseApiClient):
         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)
@@ -373,6 +379,8 @@ class ReplayApiClient(BaseApiClient):
   def _verify_response(self, response_model: BaseModel):
     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():
@@ -453,7 +461,7 @@ 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) -> HttpResponse:
     if isinstance(file_path, io.IOBase):
       offset = file_path.tell()
       content = file_path.read()
@@ -474,21 +482,21 @@ class ReplayApiClient(BaseApiClient):
         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})]
+            dict(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)]))
+      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:
+  ) -> HttpResponse:
     if isinstance(file_path, io.IOBase):
       offset = file_path.tell()
       content = file_path.read()
@@ -504,37 +512,40 @@ class ReplayApiClient(BaseApiClient):
           method='POST', url='', data={'file_path': file_path}, headers={}
       )
     if self._should_call_api():
-      result: Union[str, HttpResponse]
+      result: HttpResponse
       try:
         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})]
+            dict(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)]))
+      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_request(self, request):
+  def download_file(self, path: str, http_options: HttpOptions):
     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_request(request)
+        result = super().download_file(path, http_options)
       except HTTPError as e:
         result = HttpResponse(
-            e.response.headers, [json.dumps({'reason': e.response.reason})]
+            dict(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)
+      return self._build_response_from_replay(request).byte_stream[0]
 
   async def async_download_file(self, path: str, http_options):
     self._initialize_replay_session_if_not_loaded()
@@ -546,7 +557,7 @@ class ReplayApiClient(BaseApiClient):
         result = await super().async_download_file(path, http_options)
       except HTTPError as e:
         result = HttpResponse(
-            e.response.headers, [json.dumps({'reason': e.response.reason})]
+            dict(e.response.headers), [json.dumps({'reason': e.response.reason})]
         )
         result.status_code = e.response.status_code
         raise e

google/genai/_transformers.py

@@ -181,17 +181,26 @@ def t_models_url(
 
 def t_extract_models(
     api_client: _api_client.BaseApiClient,
-    response: dict[str, list[types.ModelDict]],
-) -> Optional[list[types.ModelDict]]:
+    response: dict[str, Any],
+) -> list[dict[str, Any]]:
   if not response:
     return []
-  elif response.get('models') is not None:
-    return response.get('models')
-  elif response.get('tunedModels') is not None:
-    return response.get('tunedModels')
-  elif response.get('publisherModels') is not None:
-    return response.get('publisherModels')
-  elif (
+
+  models: Optional[list[dict[str, Any]]] = response.get('models')
+  if models is not None:
+    return models
+
+  tuned_models: Optional[list[dict[str, Any]]] = response.get('tunedModels')
+  if tuned_models is not None:
+    return tuned_models
+
+  publisher_models: Optional[list[dict[str, Any]]] = response.get(
+      'publisherModels'
+  )
+  if publisher_models is not None:
+    return publisher_models
+
+  if (
       response.get('httpHeaders') is not None
       and response.get('jsonPayload') is None
   ):
@@ -526,7 +535,6 @@ def process_schema(
 ):
   """Updates the schema and each sub-schema inplace to be API-compatible.
 
-  - Removes the `title` field from the schema if the client is not vertexai.
   - Inlines the $defs.
 
   Example of a schema before and after (with mldev):
@@ -570,21 +578,22 @@ def process_schema(
         'items': {
           'properties': {
             'continent': {
-                'type': 'string'
+              'title': 'Continent',
+              'type': 'string'
             },
             'gdp': {
-                'type': 'integer'}
+              'title': 'Gdp',
+              'type': 'integer'
             },
           }
           'required':['continent', 'gdp'],
+          'title': 'CountryInfo',
           'type': 'object'
         },
         'type': 'array'
     }
   """
   if not client.vertexai:
-    schema.pop('title', None)
-
     if schema.get('default') is not None:
       raise ValueError(
           'Default value is not supported in the response schema for the Gemini'