google-genai 1.25.0__py3-none-any.whl → 1.27.0__py3-none-any.whl

google/genai/operations.py

@@ -17,12 +17,11 @@
 
 import json
 import logging
-from typing import Any, Optional, Union
+from typing import Any, Optional, TypeVar, Union
 from urllib.parse import urlencode
 
 from . import _api_module
 from . import _common
-from . import _transformers as t
 from . import types
 from ._common import get_value_by_path as getv
 from ._common import set_value_by_path as setv
@@ -87,218 +86,6 @@ def _FetchPredictOperationParameters_to_vertex(
   return to_object
 
 
-def _Video_from_mldev(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-  if getv(from_object, ['video', 'uri']) is not None:
-    setv(to_object, ['uri'], getv(from_object, ['video', 'uri']))
-
-  if getv(from_object, ['video', 'encodedVideo']) is not None:
-    setv(
-        to_object,
-        ['video_bytes'],
-        t.t_bytes(getv(from_object, ['video', 'encodedVideo'])),
-    )
-
-  if getv(from_object, ['encoding']) is not None:
-    setv(to_object, ['mime_type'], getv(from_object, ['encoding']))
-
-  return to_object
-
-
-def _GeneratedVideo_from_mldev(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-  if getv(from_object, ['_self']) is not None:
-    setv(
-        to_object,
-        ['video'],
-        _Video_from_mldev(getv(from_object, ['_self']), to_object),
-    )
-
-  return to_object
-
-
-def _GenerateVideosResponse_from_mldev(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-  if getv(from_object, ['generatedSamples']) is not None:
-    setv(
-        to_object,
-        ['generated_videos'],
-        [
-            _GeneratedVideo_from_mldev(item, to_object)
-            for item in getv(from_object, ['generatedSamples'])
-        ],
-    )
-
-  if getv(from_object, ['raiMediaFilteredCount']) is not None:
-    setv(
-        to_object,
-        ['rai_media_filtered_count'],
-        getv(from_object, ['raiMediaFilteredCount']),
-    )
-
-  if getv(from_object, ['raiMediaFilteredReasons']) is not None:
-    setv(
-        to_object,
-        ['rai_media_filtered_reasons'],
-        getv(from_object, ['raiMediaFilteredReasons']),
-    )
-
-  return to_object
-
-
-def _GenerateVideosOperation_from_mldev(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-  if getv(from_object, ['name']) is not None:
-    setv(to_object, ['name'], getv(from_object, ['name']))
-
-  if getv(from_object, ['metadata']) is not None:
-    setv(to_object, ['metadata'], getv(from_object, ['metadata']))
-
-  if getv(from_object, ['done']) is not None:
-    setv(to_object, ['done'], getv(from_object, ['done']))
-
-  if getv(from_object, ['error']) is not None:
-    setv(to_object, ['error'], getv(from_object, ['error']))
-
-  if getv(from_object, ['response', 'generateVideoResponse']) is not None:
-    setv(
-        to_object,
-        ['response'],
-        _GenerateVideosResponse_from_mldev(
-            getv(from_object, ['response', 'generateVideoResponse']), to_object
-        ),
-    )
-
-  if getv(from_object, ['response', 'generateVideoResponse']) is not None:
-    setv(
-        to_object,
-        ['result'],
-        _GenerateVideosResponse_from_mldev(
-            getv(from_object, ['response', 'generateVideoResponse']), to_object
-        ),
-    )
-
-  return to_object
-
-
-def _Video_from_vertex(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-  if getv(from_object, ['gcsUri']) is not None:
-    setv(to_object, ['uri'], getv(from_object, ['gcsUri']))
-
-  if getv(from_object, ['bytesBase64Encoded']) is not None:
-    setv(
-        to_object,
-        ['video_bytes'],
-        t.t_bytes(getv(from_object, ['bytesBase64Encoded'])),
-    )
-
-  if getv(from_object, ['mimeType']) is not None:
-    setv(to_object, ['mime_type'], getv(from_object, ['mimeType']))
-
-  return to_object
-
-
-def _GeneratedVideo_from_vertex(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-  if getv(from_object, ['_self']) is not None:
-    setv(
-        to_object,
-        ['video'],
-        _Video_from_vertex(getv(from_object, ['_self']), to_object),
-    )
-
-  return to_object
-
-
-def _GenerateVideosResponse_from_vertex(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-  if getv(from_object, ['videos']) is not None:
-    setv(
-        to_object,
-        ['generated_videos'],
-        [
-            _GeneratedVideo_from_vertex(item, to_object)
-            for item in getv(from_object, ['videos'])
-        ],
-    )
-
-  if getv(from_object, ['raiMediaFilteredCount']) is not None:
-    setv(
-        to_object,
-        ['rai_media_filtered_count'],
-        getv(from_object, ['raiMediaFilteredCount']),
-    )
-
-  if getv(from_object, ['raiMediaFilteredReasons']) is not None:
-    setv(
-        to_object,
-        ['rai_media_filtered_reasons'],
-        getv(from_object, ['raiMediaFilteredReasons']),
-    )
-
-  return to_object
-
-
-def _GenerateVideosOperation_from_vertex(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-  if getv(from_object, ['name']) is not None:
-    setv(to_object, ['name'], getv(from_object, ['name']))
-
-  if getv(from_object, ['metadata']) is not None:
-    setv(to_object, ['metadata'], getv(from_object, ['metadata']))
-
-  if getv(from_object, ['done']) is not None:
-    setv(to_object, ['done'], getv(from_object, ['done']))
-
-  if getv(from_object, ['error']) is not None:
-    setv(to_object, ['error'], getv(from_object, ['error']))
-
-  if getv(from_object, ['response']) is not None:
-    setv(
-        to_object,
-        ['response'],
-        _GenerateVideosResponse_from_vertex(
-            getv(from_object, ['response']), to_object
-        ),
-    )
-
-  if getv(from_object, ['response']) is not None:
-    setv(
-        to_object,
-        ['result'],
-        _GenerateVideosResponse_from_vertex(
-            getv(from_object, ['response']), to_object
-        ),
-    )
-
-  return to_object
-
-
 class Operations(_api_module.BaseModule):
 
   def _get_videos_operation(
@@ -306,7 +93,7 @@ class Operations(_api_module.BaseModule):
       *,
       operation_name: str,
       config: Optional[types.GetOperationConfigOrDict] = None,
-  ) -> types.GenerateVideosOperation:
+  ) -> Any:
     parameter_model = types._GetOperationParameters(
         operation_name=operation_name,
         config=config,
@@ -348,18 +135,7 @@ class Operations(_api_module.BaseModule):
 
     response_dict = '' if not response.body else json.loads(response.body)
 
-    if self._api_client.vertexai:
-      response_dict = _GenerateVideosOperation_from_vertex(response_dict)
-
-    else:
-      response_dict = _GenerateVideosOperation_from_mldev(response_dict)
-
-    return_value = types.GenerateVideosOperation._from_response(
-        response=response_dict, kwargs=parameter_model.model_dump()
-    )
-
-    self._api_client._verify_response(return_value)
-    return return_value
+    return response_dict
 
   def _fetch_predict_videos_operation(
       self,
@@ -367,7 +143,7 @@ class Operations(_api_module.BaseModule):
       operation_name: str,
       resource_name: str,
       config: Optional[types.FetchPredictOperationConfigOrDict] = None,
-  ) -> types.GenerateVideosOperation:
+  ) -> Any:
     parameter_model = types._FetchPredictOperationParameters(
         operation_name=operation_name,
         resource_name=resource_name,
@@ -409,22 +185,16 @@ class Operations(_api_module.BaseModule):
 
     response_dict = '' if not response.body else json.loads(response.body)
 
-    if self._api_client.vertexai:
-      response_dict = _GenerateVideosOperation_from_vertex(response_dict)
+    return response_dict
 
-    return_value = types.GenerateVideosOperation._from_response(
-        response=response_dict, kwargs=parameter_model.model_dump()
-    )
-
-    self._api_client._verify_response(return_value)
-    return return_value
+  T = TypeVar('T', bound=types.Operation)
 
   def get(
       self,
-      operation: types.GenerateVideosOperation,
+      operation: T,
       *,
       config: Optional[types.GetOperationConfigOrDict] = None,
-  ) -> types.GenerateVideosOperation:
+  ) -> T:
     """Gets the status of an operation."""
     # Currently, only GenerateVideosOperation is supported.
     # TODO(b/398040607): Support short form names
@@ -449,16 +219,25 @@ class Operations(_api_module.BaseModule):
       fetch_operation_config = types.FetchPredictOperationConfig(
           http_options=http_options
       )
-      return self._fetch_predict_videos_operation(
+      response_dict = self._fetch_predict_videos_operation(
           operation_name=operation_name,
           resource_name=resource_name,
           config=fetch_operation_config,
       )
+
+      response_operation = operation.from_api_response(
+          response_dict, is_vertex_ai=True
+      )
+      return response_operation  # type: ignore[no-any-return]
     else:
-      return self._get_videos_operation(
+      response_dict = self._get_videos_operation(
           operation_name=operation_name,
           config=config,
       )
+      response_operation = operation.from_api_response(
+          response_dict, is_vertex_ai=False
+      )
+      return response_operation  # type: ignore[no-any-return]
 
 
 class AsyncOperations(_api_module.BaseModule):
@@ -468,7 +247,7 @@ class AsyncOperations(_api_module.BaseModule):
       *,
       operation_name: str,
       config: Optional[types.GetOperationConfigOrDict] = None,
-  ) -> types.GenerateVideosOperation:
+  ) -> Any:
     parameter_model = types._GetOperationParameters(
         operation_name=operation_name,
         config=config,
@@ -512,18 +291,7 @@ class AsyncOperations(_api_module.BaseModule):
 
     response_dict = '' if not response.body else json.loads(response.body)
 
-    if self._api_client.vertexai:
-      response_dict = _GenerateVideosOperation_from_vertex(response_dict)
-
-    else:
-      response_dict = _GenerateVideosOperation_from_mldev(response_dict)
-
-    return_value = types.GenerateVideosOperation._from_response(
-        response=response_dict, kwargs=parameter_model.model_dump()
-    )
-
-    self._api_client._verify_response(return_value)
-    return return_value
+    return response_dict
 
   async def _fetch_predict_videos_operation(
       self,
@@ -531,7 +299,7 @@ class AsyncOperations(_api_module.BaseModule):
       operation_name: str,
       resource_name: str,
       config: Optional[types.FetchPredictOperationConfigOrDict] = None,
-  ) -> types.GenerateVideosOperation:
+  ) -> Any:
     parameter_model = types._FetchPredictOperationParameters(
         operation_name=operation_name,
         resource_name=resource_name,
@@ -573,22 +341,16 @@ class AsyncOperations(_api_module.BaseModule):
 
     response_dict = '' if not response.body else json.loads(response.body)
 
-    if self._api_client.vertexai:
-      response_dict = _GenerateVideosOperation_from_vertex(response_dict)
+    return response_dict
 
-    return_value = types.GenerateVideosOperation._from_response(
-        response=response_dict, kwargs=parameter_model.model_dump()
-    )
-
-    self._api_client._verify_response(return_value)
-    return return_value
+  T = TypeVar('T', bound=types.Operation)
 
   async def get(
       self,
-      operation: types.GenerateVideosOperation,
+      operation: T,
       *,
       config: Optional[types.GetOperationConfigOrDict] = None,
-  ) -> types.GenerateVideosOperation:
+  ) -> T:
     """Gets the status of an operation."""
     # Currently, only GenerateVideosOperation is supported.
     operation_name = operation.name
@@ -611,13 +373,21 @@ class AsyncOperations(_api_module.BaseModule):
       fetch_operation_config = types.FetchPredictOperationConfig(
           http_options=http_options
       )
-      return await self._fetch_predict_videos_operation(
+      response_dict = await self._fetch_predict_videos_operation(
           operation_name=operation_name,
           resource_name=resource_name,
           config=fetch_operation_config,
       )
+      response_operation = operation.from_api_response(
+          response_dict, is_vertex_ai=True
+      )
+      return response_operation  # type: ignore[no-any-return]
     else:
-      return await self._get_videos_operation(
+      response_dict = await self._get_videos_operation(
           operation_name=operation_name,
           config=config,
       )
+      response_operation = operation.from_api_response(
+          response_dict, is_vertex_ai=False
+      )
+      return response_operation  # type: ignore[no-any-return]

google/genai/pagers.py

@@ -18,7 +18,8 @@
 # pylint: disable=protected-access
 
 import copy
-from typing import Any, AsyncIterator,Awaitable, Callable, Generic, Iterator, Literal, TypeVar
+from typing import Any, AsyncIterator,Awaitable, Callable, Generic, Iterator, Literal, TypeVar, Union
+from . import types
 
 T = TypeVar('T')
 
@@ -43,6 +44,8 @@ class _BasePager(Generic[T]):
     self._page = getattr(response, self._name) or []
     self._idx = 0
 
+    self._sdk_http_response = getattr(response, 'sdk_http_response', None)
+
     if not config:
       request_config = {}
     elif isinstance(config, dict):
@@ -110,6 +113,13 @@ class _BasePager(Generic[T]):
 
     return self._page_size
 
+  @property
+  def sdk_http_response(self) -> Union[types.HttpResponse, None]:
+    """Returns the http response of the API response.
+    """
+
+    return self._sdk_http_response
+
   @property
   def config(self) -> dict[str, Any]:
     """Returns the configuration when making the API request for the next page.

google/genai/tunings.py

@@ -126,22 +126,6 @@ def _TuningDataset_to_mldev(
   return to_object
 
 
-def _TuningValidationDataset_to_mldev(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-  if getv(from_object, ['gcs_uri']) is not None:
-    raise ValueError('gcs_uri parameter is not supported in Gemini API.')
-
-  if getv(from_object, ['vertex_dataset_resource']) is not None:
-    raise ValueError(
-        'vertex_dataset_resource parameter is not supported in Gemini API.'
-    )
-
-  return to_object
-
-
 def _CreateTuningJobConfig_to_mldev(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
@@ -286,20 +270,6 @@ def _ListTuningJobsParameters_to_vertex(
   return to_object
 
 
-def _TuningExample_to_vertex(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-  if getv(from_object, ['text_input']) is not None:
-    raise ValueError('text_input parameter is not supported in Vertex AI.')
-
-  if getv(from_object, ['output']) is not None:
-    raise ValueError('output parameter is not supported in Vertex AI.')
-
-  return to_object
-
-
 def _TuningDataset_to_vertex(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
@@ -434,15 +404,6 @@ def _CreateTuningJobParameters_to_vertex(
   return to_object
 
 
-def _TunedModelCheckpoint_from_mldev(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-
-  return to_object
-
-
 def _TunedModel_from_mldev(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
@@ -545,6 +506,11 @@ def _ListTuningJobsResponse_from_mldev(
     parent_object: Optional[dict[str, Any]] = None,
 ) -> dict[str, Any]:
   to_object: dict[str, Any] = {}
+  if getv(from_object, ['sdkHttpResponse']) is not None:
+    setv(
+        to_object, ['sdk_http_response'], getv(from_object, ['sdkHttpResponse'])
+    )
+
   if getv(from_object, ['nextPageToken']) is not None:
     setv(to_object, ['next_page_token'], getv(from_object, ['nextPageToken']))
 
@@ -561,7 +527,7 @@ def _ListTuningJobsResponse_from_mldev(
   return to_object
 
 
-def _Operation_from_mldev(
+def _TuningOperation_from_mldev(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
 ) -> dict[str, Any]:
@@ -730,6 +696,11 @@ def _ListTuningJobsResponse_from_vertex(
     parent_object: Optional[dict[str, Any]] = None,
 ) -> dict[str, Any]:
   to_object: dict[str, Any] = {}
+  if getv(from_object, ['sdkHttpResponse']) is not None:
+    setv(
+        to_object, ['sdk_http_response'], getv(from_object, ['sdkHttpResponse'])
+    )
+
   if getv(from_object, ['nextPageToken']) is not None:
     setv(to_object, ['next_page_token'], getv(from_object, ['nextPageToken']))
 
@@ -878,7 +849,9 @@ class Tunings(_api_module.BaseModule):
     return_value = types.ListTuningJobsResponse._from_response(
         response=response_dict, kwargs=parameter_model.model_dump()
     )
-
+    return_value.sdk_http_response = types.HttpResponse(
+        headers=response.headers
+    )
     self._api_client._verify_response(return_value)
     return return_value
 
@@ -955,7 +928,7 @@ class Tunings(_api_module.BaseModule):
       base_model: str,
       training_dataset: types.TuningDatasetOrDict,
       config: Optional[types.CreateTuningJobConfigOrDict] = None,
-  ) -> types.Operation:
+  ) -> types.TuningOperation:
     """Creates a supervised fine-tuning job and returns the TuningJob object.
 
     Args:
@@ -1009,9 +982,9 @@ class Tunings(_api_module.BaseModule):
     response_dict = '' if not response.body else json.loads(response.body)
 
     if not self._api_client.vertexai:
-      response_dict = _Operation_from_mldev(response_dict)
+      response_dict = _TuningOperation_from_mldev(response_dict)
 
-    return_value = types.Operation._from_response(
+    return_value = types.TuningOperation._from_response(
         response=response_dict, kwargs=parameter_model.model_dump()
     )
 
@@ -1222,7 +1195,9 @@ class AsyncTunings(_api_module.BaseModule):
     return_value = types.ListTuningJobsResponse._from_response(
         response=response_dict, kwargs=parameter_model.model_dump()
     )
-
+    return_value.sdk_http_response = types.HttpResponse(
+        headers=response.headers
+    )
     self._api_client._verify_response(return_value)
     return return_value
 
@@ -1299,7 +1274,7 @@ class AsyncTunings(_api_module.BaseModule):
       base_model: str,
       training_dataset: types.TuningDatasetOrDict,
       config: Optional[types.CreateTuningJobConfigOrDict] = None,
-  ) -> types.Operation:
+  ) -> types.TuningOperation:
     """Creates a supervised fine-tuning job and returns the TuningJob object.
 
     Args:
@@ -1353,9 +1328,9 @@ class AsyncTunings(_api_module.BaseModule):
     response_dict = '' if not response.body else json.loads(response.body)
 
     if not self._api_client.vertexai:
-      response_dict = _Operation_from_mldev(response_dict)
+      response_dict = _TuningOperation_from_mldev(response_dict)
 
-    return_value = types.Operation._from_response(
+    return_value = types.TuningOperation._from_response(
         response=response_dict, kwargs=parameter_model.model_dump()
     )