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

google/genai/operations.py

@@ -30,6 +30,24 @@ from ._common import set_value_by_path as setv
 logger = logging.getLogger('google_genai.operations')
 
 
+def _FetchPredictOperationParameters_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, ['operation_name']) is not None:
+    setv(to_object, ['operationName'], getv(from_object, ['operation_name']))
+
+  if getv(from_object, ['resource_name']) is not None:
+    setv(
+        to_object,
+        ['_url', 'resourceName'],
+        getv(from_object, ['resource_name']),
+    )
+
+  return to_object
+
+
 def _GetOperationParameters_to_mldev(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
@@ -42,9 +60,6 @@ def _GetOperationParameters_to_mldev(
         getv(from_object, ['operation_name']),
     )
 
-  if getv(from_object, ['config']) is not None:
-    setv(to_object, ['config'], getv(from_object, ['config']))
-
   return to_object
 
 
@@ -60,30 +75,19 @@ def _GetOperationParameters_to_vertex(
         getv(from_object, ['operation_name']),
     )
 
-  if getv(from_object, ['config']) is not None:
-    setv(to_object, ['config'], getv(from_object, ['config']))
-
   return to_object
 
 
-def _FetchPredictOperationParameters_to_vertex(
+def _GetProjectOperationParameters_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, ['operation_name']) is not None:
-    setv(to_object, ['operationName'], getv(from_object, ['operation_name']))
-
-  if getv(from_object, ['resource_name']) is not None:
+  if getv(from_object, ['operation_id']) is not None:
     setv(
-        to_object,
-        ['_url', 'resourceName'],
-        getv(from_object, ['resource_name']),
+        to_object, ['_url', 'operation_id'], getv(from_object, ['operation_id'])
     )
 
-  if getv(from_object, ['config']) is not None:
-    setv(to_object, ['config'], getv(from_object, ['config']))
-
   return to_object
 
 
@@ -134,7 +138,7 @@ class Operations(_api_module.BaseModule):
 
     response = self._api_client.request('get', path, request_dict, http_options)
 
-    response_dict = '' if not response.body else json.loads(response.body)
+    response_dict = {} if not response.body else json.loads(response.body)
 
     return response_dict
 
@@ -184,10 +188,59 @@ class Operations(_api_module.BaseModule):
         'post', path, request_dict, http_options
     )
 
-    response_dict = '' if not response.body else json.loads(response.body)
+    response_dict = {} if not response.body else json.loads(response.body)
 
     return response_dict
 
+  def _get(
+      self,
+      *,
+      operation_id: str,
+      config: Optional[types.GetOperationConfigOrDict] = None,
+  ) -> types.ProjectOperation:
+    parameter_model = types._GetProjectOperationParameters(
+        operation_id=operation_id,
+        config=config,
+    )
+
+    request_url_dict: Optional[dict[str, str]]
+    if not self._api_client.vertexai:
+      raise ValueError('This method is only supported in the Vertex AI client.')
+    else:
+      request_dict = _GetProjectOperationParameters_to_vertex(parameter_model)
+      request_url_dict = request_dict.get('_url')
+      if request_url_dict:
+        path = 'operations/{operation_id}'.format_map(request_url_dict)
+      else:
+        path = 'operations/{operation_id}'
+
+    query_params = request_dict.get('_query')
+    if query_params:
+      path = f'{path}?{urlencode(query_params)}'
+    # TODO: remove the hack that pops config.
+    request_dict.pop('config', None)
+
+    http_options: Optional[types.HttpOptions] = None
+    if (
+        parameter_model.config is not None
+        and parameter_model.config.http_options is not None
+    ):
+      http_options = parameter_model.config.http_options
+
+    request_dict = _common.convert_to_dict(request_dict)
+    request_dict = _common.encode_unserializable_types(request_dict)
+
+    response = self._api_client.request('get', path, request_dict, http_options)
+
+    response_dict = {} if not response.body else json.loads(response.body)
+
+    return_value = types.ProjectOperation._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(
@@ -229,6 +282,8 @@ class Operations(_api_module.BaseModule):
       response_operation = operation.from_api_response(
           response_dict, is_vertex_ai=True
       )
+
+      self._api_client._verify_response(response_operation)  # type: ignore [arg-type]
       return response_operation  # type: ignore[no-any-return]
     else:
       response_dict = self._get_videos_operation(
@@ -238,6 +293,8 @@ class Operations(_api_module.BaseModule):
       response_operation = operation.from_api_response(
           response_dict, is_vertex_ai=False
       )
+
+      self._api_client._verify_response(response_operation)  # type: ignore [arg-type]
       return response_operation  # type: ignore[no-any-return]
 
 
@@ -290,7 +347,7 @@ class AsyncOperations(_api_module.BaseModule):
         'get', path, request_dict, http_options
     )
 
-    response_dict = '' if not response.body else json.loads(response.body)
+    response_dict = {} if not response.body else json.loads(response.body)
 
     return response_dict
 
@@ -340,10 +397,61 @@ class AsyncOperations(_api_module.BaseModule):
         'post', path, request_dict, http_options
     )
 
-    response_dict = '' if not response.body else json.loads(response.body)
+    response_dict = {} if not response.body else json.loads(response.body)
 
     return response_dict
 
+  async def _get(
+      self,
+      *,
+      operation_id: str,
+      config: Optional[types.GetOperationConfigOrDict] = None,
+  ) -> types.ProjectOperation:
+    parameter_model = types._GetProjectOperationParameters(
+        operation_id=operation_id,
+        config=config,
+    )
+
+    request_url_dict: Optional[dict[str, str]]
+    if not self._api_client.vertexai:
+      raise ValueError('This method is only supported in the Vertex AI client.')
+    else:
+      request_dict = _GetProjectOperationParameters_to_vertex(parameter_model)
+      request_url_dict = request_dict.get('_url')
+      if request_url_dict:
+        path = 'operations/{operation_id}'.format_map(request_url_dict)
+      else:
+        path = 'operations/{operation_id}'
+
+    query_params = request_dict.get('_query')
+    if query_params:
+      path = f'{path}?{urlencode(query_params)}'
+    # TODO: remove the hack that pops config.
+    request_dict.pop('config', None)
+
+    http_options: Optional[types.HttpOptions] = None
+    if (
+        parameter_model.config is not None
+        and parameter_model.config.http_options is not None
+    ):
+      http_options = parameter_model.config.http_options
+
+    request_dict = _common.convert_to_dict(request_dict)
+    request_dict = _common.encode_unserializable_types(request_dict)
+
+    response = await self._api_client.async_request(
+        'get', path, request_dict, http_options
+    )
+
+    response_dict = {} if not response.body else json.loads(response.body)
+
+    return_value = types.ProjectOperation._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(

google/genai/pagers.py

@@ -25,7 +25,13 @@ from . import types
 T = TypeVar('T')
 
 PagedItem = Literal[
-    'batch_jobs', 'models', 'tuning_jobs', 'files', 'cached_contents'
+    'batch_jobs',
+    'models',
+    'tuning_jobs',
+    'files',
+    'cached_contents',
+    'file_search_stores',
+    'documents',
 ]
 
 

google/genai/tokens.py

@@ -257,12 +257,7 @@ class Tokens(_api_module.BaseModule):
     response = self._api_client.request(
         'post', path, request_dict, http_options
     )
-    response_dict = '' if not response.body else json.loads(response.body)
-
-    if not self._api_client.vertexai:
-      response_dict = tokens_converters._AuthToken_from_mldev(
-          response_dict
-      )
+    response_dict = {} if not response.body else json.loads(response.body)
 
     return_value = types.AuthToken._from_response(
         response=response_dict, kwargs=parameter_model.model_dump()
@@ -358,12 +353,7 @@ class AsyncTokens(_api_module.BaseModule):
         request_dict,
         http_options=http_options,
     )
-    response_dict = '' if not response.body else json.loads(response.body)
-
-    if not self._api_client.vertexai:
-      response_dict = tokens_converters._AuthToken_from_mldev(
-          response_dict
-      )
+    response_dict = {} if not response.body else json.loads(response.body)
 
     return_value = types.AuthToken._from_response(
         response=response_dict, kwargs=parameter_model.model_dump()