diff options
author | S. Solomon Darnell | 2025-03-28 21:52:21 -0500 |
---|---|---|
committer | S. Solomon Darnell | 2025-03-28 21:52:21 -0500 |
commit | 4a52a71956a8d46fcb7294ac71734504bb09bcc2 (patch) | |
tree | ee3dc5af3b6313e921cd920906356f5d4febc4ed /.venv/lib/python3.12/site-packages/google/genai/batches.py | |
parent | cc961e04ba734dd72309fb548a2f97d67d578813 (diff) | |
download | gn-ai-master.tar.gz |
Diffstat (limited to '.venv/lib/python3.12/site-packages/google/genai/batches.py')
-rw-r--r-- | .venv/lib/python3.12/site-packages/google/genai/batches.py | 1293 |
1 files changed, 1293 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/google/genai/batches.py b/.venv/lib/python3.12/site-packages/google/genai/batches.py new file mode 100644 index 00000000..6dd5b4c5 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/google/genai/batches.py @@ -0,0 +1,1293 @@ +# 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. +# + +# Code generated by the Google Gen AI SDK generator DO NOT EDIT. + +from typing import Optional, Union +from urllib.parse import urlencode +from . import _common +from . import _extra_utils +from . import _transformers as t +from . import types +from ._api_client import ApiClient +from ._common import get_value_by_path as getv +from ._common import set_value_by_path as setv +from .pagers import AsyncPager, Pager + + +def _BatchJobSource_to_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['format']) is not None: + raise ValueError('format parameter is not supported in Google AI.') + + if getv(from_object, ['gcs_uri']) is not None: + raise ValueError('gcs_uri parameter is not supported in Google AI.') + + if getv(from_object, ['bigquery_uri']) is not None: + raise ValueError('bigquery_uri parameter is not supported in Google AI.') + + return to_object + + +def _BatchJobSource_to_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['format']) is not None: + setv(to_object, ['instancesFormat'], getv(from_object, ['format'])) + + if getv(from_object, ['gcs_uri']) is not None: + setv(to_object, ['gcsSource', 'uris'], getv(from_object, ['gcs_uri'])) + + if getv(from_object, ['bigquery_uri']) is not None: + setv( + to_object, + ['bigquerySource', 'inputUri'], + getv(from_object, ['bigquery_uri']), + ) + + return to_object + + +def _BatchJobDestination_to_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['format']) is not None: + raise ValueError('format parameter is not supported in Google AI.') + + if getv(from_object, ['gcs_uri']) is not None: + raise ValueError('gcs_uri parameter is not supported in Google AI.') + + if getv(from_object, ['bigquery_uri']) is not None: + raise ValueError('bigquery_uri parameter is not supported in Google AI.') + + return to_object + + +def _BatchJobDestination_to_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['format']) is not None: + setv(to_object, ['predictionsFormat'], getv(from_object, ['format'])) + + if getv(from_object, ['gcs_uri']) is not None: + setv( + to_object, + ['gcsDestination', 'outputUriPrefix'], + getv(from_object, ['gcs_uri']), + ) + + if getv(from_object, ['bigquery_uri']) is not None: + setv( + to_object, + ['bigqueryDestination', 'outputUri'], + getv(from_object, ['bigquery_uri']), + ) + + return to_object + + +def _CreateBatchJobConfig_to_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['http_options']) is not None: + setv(to_object, ['httpOptions'], getv(from_object, ['http_options'])) + + if getv(from_object, ['display_name']) is not None: + setv(parent_object, ['displayName'], getv(from_object, ['display_name'])) + + if getv(from_object, ['dest']) is not None: + raise ValueError('dest parameter is not supported in Google AI.') + + return to_object + + +def _CreateBatchJobConfig_to_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['http_options']) is not None: + setv(to_object, ['httpOptions'], getv(from_object, ['http_options'])) + + if getv(from_object, ['display_name']) is not None: + setv(parent_object, ['displayName'], getv(from_object, ['display_name'])) + + if getv(from_object, ['dest']) is not None: + setv( + parent_object, + ['outputConfig'], + _BatchJobDestination_to_vertex( + api_client, + t.t_batch_job_destination(api_client, getv(from_object, ['dest'])), + to_object, + ), + ) + + return to_object + + +def _CreateBatchJobParameters_to_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['model']) is not None: + raise ValueError('model parameter is not supported in Google AI.') + + if getv(from_object, ['src']) is not None: + raise ValueError('src parameter is not supported in Google AI.') + + if getv(from_object, ['config']) is not None: + setv( + to_object, + ['config'], + _CreateBatchJobConfig_to_mldev( + api_client, getv(from_object, ['config']), to_object + ), + ) + + return to_object + + +def _CreateBatchJobParameters_to_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['model']) is not None: + setv( + to_object, + ['model'], + t.t_model(api_client, getv(from_object, ['model'])), + ) + + if getv(from_object, ['src']) is not None: + setv( + to_object, + ['inputConfig'], + _BatchJobSource_to_vertex( + api_client, + t.t_batch_job_source(api_client, getv(from_object, ['src'])), + to_object, + ), + ) + + if getv(from_object, ['config']) is not None: + setv( + to_object, + ['config'], + _CreateBatchJobConfig_to_vertex( + api_client, getv(from_object, ['config']), to_object + ), + ) + + return to_object + + +def _GetBatchJobConfig_to_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['http_options']) is not None: + setv(to_object, ['httpOptions'], getv(from_object, ['http_options'])) + + return to_object + + +def _GetBatchJobConfig_to_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['http_options']) is not None: + setv(to_object, ['httpOptions'], getv(from_object, ['http_options'])) + + return to_object + + +def _GetBatchJobParameters_to_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['name']) is not None: + raise ValueError('name parameter is not supported in Google AI.') + + if getv(from_object, ['config']) is not None: + setv( + to_object, + ['config'], + _GetBatchJobConfig_to_mldev( + api_client, getv(from_object, ['config']), to_object + ), + ) + + return to_object + + +def _GetBatchJobParameters_to_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['name']) is not None: + setv( + to_object, + ['_url', 'name'], + t.t_batch_job_name(api_client, getv(from_object, ['name'])), + ) + + if getv(from_object, ['config']) is not None: + setv( + to_object, + ['config'], + _GetBatchJobConfig_to_vertex( + api_client, getv(from_object, ['config']), to_object + ), + ) + + return to_object + + +def _CancelBatchJobConfig_to_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['http_options']) is not None: + setv(to_object, ['httpOptions'], getv(from_object, ['http_options'])) + + return to_object + + +def _CancelBatchJobConfig_to_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['http_options']) is not None: + setv(to_object, ['httpOptions'], getv(from_object, ['http_options'])) + + return to_object + + +def _CancelBatchJobParameters_to_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['name']) is not None: + raise ValueError('name parameter is not supported in Google AI.') + + if getv(from_object, ['config']) is not None: + setv( + to_object, + ['config'], + _CancelBatchJobConfig_to_mldev( + api_client, getv(from_object, ['config']), to_object + ), + ) + + return to_object + + +def _CancelBatchJobParameters_to_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['name']) is not None: + setv( + to_object, + ['_url', 'name'], + t.t_batch_job_name(api_client, getv(from_object, ['name'])), + ) + + if getv(from_object, ['config']) is not None: + setv( + to_object, + ['config'], + _CancelBatchJobConfig_to_vertex( + api_client, getv(from_object, ['config']), to_object + ), + ) + + return to_object + + +def _ListBatchJobConfig_to_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['http_options']) is not None: + setv(to_object, ['httpOptions'], getv(from_object, ['http_options'])) + + if getv(from_object, ['page_size']) is not None: + setv( + parent_object, ['_query', 'pageSize'], getv(from_object, ['page_size']) + ) + + if getv(from_object, ['page_token']) is not None: + setv( + parent_object, + ['_query', 'pageToken'], + getv(from_object, ['page_token']), + ) + + if getv(from_object, ['filter']) is not None: + raise ValueError('filter parameter is not supported in Google AI.') + + return to_object + + +def _ListBatchJobConfig_to_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['http_options']) is not None: + setv(to_object, ['httpOptions'], getv(from_object, ['http_options'])) + + if getv(from_object, ['page_size']) is not None: + setv( + parent_object, ['_query', 'pageSize'], getv(from_object, ['page_size']) + ) + + if getv(from_object, ['page_token']) is not None: + setv( + parent_object, + ['_query', 'pageToken'], + getv(from_object, ['page_token']), + ) + + if getv(from_object, ['filter']) is not None: + setv(parent_object, ['_query', 'filter'], getv(from_object, ['filter'])) + + return to_object + + +def _ListBatchJobParameters_to_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['config']) is not None: + raise ValueError('config parameter is not supported in Google AI.') + + return to_object + + +def _ListBatchJobParameters_to_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['config']) is not None: + setv( + to_object, + ['config'], + _ListBatchJobConfig_to_vertex( + api_client, getv(from_object, ['config']), to_object + ), + ) + + return to_object + + +def _DeleteBatchJobParameters_to_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['name']) is not None: + raise ValueError('name parameter is not supported in Google AI.') + + return to_object + + +def _DeleteBatchJobParameters_to_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['name']) is not None: + setv( + to_object, + ['_url', 'name'], + t.t_batch_job_name(api_client, getv(from_object, ['name'])), + ) + + return to_object + + +def _JobError_from_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + + return to_object + + +def _JobError_from_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['details']) is not None: + setv(to_object, ['details'], getv(from_object, ['details'])) + + if getv(from_object, ['code']) is not None: + setv(to_object, ['code'], getv(from_object, ['code'])) + + if getv(from_object, ['message']) is not None: + setv(to_object, ['message'], getv(from_object, ['message'])) + + return to_object + + +def _BatchJobSource_from_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + + return to_object + + +def _BatchJobSource_from_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['instancesFormat']) is not None: + setv(to_object, ['format'], getv(from_object, ['instancesFormat'])) + + if getv(from_object, ['gcsSource', 'uris']) is not None: + setv(to_object, ['gcs_uri'], getv(from_object, ['gcsSource', 'uris'])) + + if getv(from_object, ['bigquerySource', 'inputUri']) is not None: + setv( + to_object, + ['bigquery_uri'], + getv(from_object, ['bigquerySource', 'inputUri']), + ) + + return to_object + + +def _BatchJobDestination_from_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + + return to_object + + +def _BatchJobDestination_from_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['predictionsFormat']) is not None: + setv(to_object, ['format'], getv(from_object, ['predictionsFormat'])) + + if getv(from_object, ['gcsDestination', 'outputUriPrefix']) is not None: + setv( + to_object, + ['gcs_uri'], + getv(from_object, ['gcsDestination', 'outputUriPrefix']), + ) + + if getv(from_object, ['bigqueryDestination', 'outputUri']) is not None: + setv( + to_object, + ['bigquery_uri'], + getv(from_object, ['bigqueryDestination', 'outputUri']), + ) + + return to_object + + +def _BatchJob_from_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + + return to_object + + +def _BatchJob_from_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['name']) is not None: + setv(to_object, ['name'], getv(from_object, ['name'])) + + if getv(from_object, ['displayName']) is not None: + setv(to_object, ['display_name'], getv(from_object, ['displayName'])) + + if getv(from_object, ['state']) is not None: + setv(to_object, ['state'], getv(from_object, ['state'])) + + if getv(from_object, ['error']) is not None: + setv( + to_object, + ['error'], + _JobError_from_vertex( + api_client, getv(from_object, ['error']), to_object + ), + ) + + if getv(from_object, ['createTime']) is not None: + setv(to_object, ['create_time'], getv(from_object, ['createTime'])) + + if getv(from_object, ['startTime']) is not None: + setv(to_object, ['start_time'], getv(from_object, ['startTime'])) + + if getv(from_object, ['endTime']) is not None: + setv(to_object, ['end_time'], getv(from_object, ['endTime'])) + + if getv(from_object, ['updateTime']) is not None: + setv(to_object, ['update_time'], getv(from_object, ['updateTime'])) + + if getv(from_object, ['model']) is not None: + setv(to_object, ['model'], getv(from_object, ['model'])) + + if getv(from_object, ['inputConfig']) is not None: + setv( + to_object, + ['src'], + _BatchJobSource_from_vertex( + api_client, getv(from_object, ['inputConfig']), to_object + ), + ) + + if getv(from_object, ['outputConfig']) is not None: + setv( + to_object, + ['dest'], + _BatchJobDestination_from_vertex( + api_client, getv(from_object, ['outputConfig']), to_object + ), + ) + + return to_object + + +def _ListBatchJobResponse_from_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['nextPageToken']) is not None: + setv(to_object, ['next_page_token'], getv(from_object, ['nextPageToken'])) + + return to_object + + +def _ListBatchJobResponse_from_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['nextPageToken']) is not None: + setv(to_object, ['next_page_token'], getv(from_object, ['nextPageToken'])) + + if getv(from_object, ['batchPredictionJobs']) is not None: + setv( + to_object, + ['batch_jobs'], + [ + _BatchJob_from_vertex(api_client, item, to_object) + for item in getv(from_object, ['batchPredictionJobs']) + ], + ) + + return to_object + + +def _DeleteResourceJob_from_mldev( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + + return to_object + + +def _DeleteResourceJob_from_vertex( + api_client: ApiClient, + from_object: Union[dict, object], + parent_object: dict = None, +) -> dict: + to_object = {} + if getv(from_object, ['name']) is not None: + setv(to_object, ['name'], getv(from_object, ['name'])) + + 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'], + _JobError_from_vertex( + api_client, getv(from_object, ['error']), to_object + ), + ) + + return to_object + + +class Batches(_common.BaseModule): + + def _create( + self, + *, + model: str, + src: str, + config: Optional[types.CreateBatchJobConfigOrDict] = None, + ) -> types.BatchJob: + parameter_model = types._CreateBatchJobParameters( + model=model, + src=src, + config=config, + ) + + if not self._api_client.vertexai: + raise ValueError('This method is only supported in the Vertex AI client.') + else: + request_dict = _CreateBatchJobParameters_to_vertex( + self._api_client, parameter_model + ) + path = 'batchPredictionJobs'.format_map(request_dict.get('_url')) + + query_params = request_dict.get('_query') + if query_params: + path = f'{path}?{urlencode(query_params)}' + # TODO: remove the hack that pops config. + config = request_dict.pop('config', None) + http_options = config.pop('httpOptions', None) if config else None + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response_dict = self._api_client.request( + 'post', path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _BatchJob_from_vertex(self._api_client, response_dict) + else: + response_dict = _BatchJob_from_mldev(self._api_client, response_dict) + + return_value = types.BatchJob._from_response(response_dict, parameter_model) + self._api_client._verify_response(return_value) + return return_value + + def get( + self, *, name: str, config: Optional[types.GetBatchJobConfigOrDict] = None + ) -> types.BatchJob: + """Gets a batch job. + + Args: + name (str): A fully-qualified BatchJob resource name or ID. + Example: "projects/.../locations/.../batchPredictionJobs/456" or "456" + when project and location are initialized in the client. + + Returns: + A BatchJob object that contains details about the batch job. + + Usage: + + .. code-block:: python + + batch_job = client.batches.get(name='123456789') + print(f"Batch job: {batch_job.name}, state {batch_job.state}") + """ + + parameter_model = types._GetBatchJobParameters( + name=name, + config=config, + ) + + if not self._api_client.vertexai: + raise ValueError('This method is only supported in the Vertex AI client.') + else: + request_dict = _GetBatchJobParameters_to_vertex( + self._api_client, parameter_model + ) + path = 'batchPredictionJobs/{name}'.format_map(request_dict.get('_url')) + + query_params = request_dict.get('_query') + if query_params: + path = f'{path}?{urlencode(query_params)}' + # TODO: remove the hack that pops config. + config = request_dict.pop('config', None) + http_options = config.pop('httpOptions', None) if config else None + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response_dict = self._api_client.request( + 'get', path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _BatchJob_from_vertex(self._api_client, response_dict) + else: + response_dict = _BatchJob_from_mldev(self._api_client, response_dict) + + return_value = types.BatchJob._from_response(response_dict, parameter_model) + self._api_client._verify_response(return_value) + return return_value + + def cancel( + self, + *, + name: str, + config: Optional[types.CancelBatchJobConfigOrDict] = None, + ) -> None: + parameter_model = types._CancelBatchJobParameters( + name=name, + config=config, + ) + + if not self._api_client.vertexai: + raise ValueError('This method is only supported in the Vertex AI client.') + else: + request_dict = _CancelBatchJobParameters_to_vertex( + self._api_client, parameter_model + ) + path = 'batchPredictionJobs/{name}:cancel'.format_map( + request_dict.get('_url') + ) + + query_params = request_dict.get('_query') + if query_params: + path = f'{path}?{urlencode(query_params)}' + # TODO: remove the hack that pops config. + config = request_dict.pop('config', None) + http_options = config.pop('httpOptions', None) if config else None + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response_dict = self._api_client.request( + 'post', path, request_dict, http_options + ) + + def _list( + self, *, config: types.ListBatchJobConfigOrDict + ) -> types.ListBatchJobResponse: + parameter_model = types._ListBatchJobParameters( + config=config, + ) + + if not self._api_client.vertexai: + raise ValueError('This method is only supported in the Vertex AI client.') + else: + request_dict = _ListBatchJobParameters_to_vertex( + self._api_client, parameter_model + ) + path = 'batchPredictionJobs'.format_map(request_dict.get('_url')) + + query_params = request_dict.get('_query') + if query_params: + path = f'{path}?{urlencode(query_params)}' + # TODO: remove the hack that pops config. + config = request_dict.pop('config', None) + http_options = config.pop('httpOptions', None) if config else None + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response_dict = self._api_client.request( + 'get', path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _ListBatchJobResponse_from_vertex( + self._api_client, response_dict + ) + else: + response_dict = _ListBatchJobResponse_from_mldev( + self._api_client, response_dict + ) + + return_value = types.ListBatchJobResponse._from_response( + response_dict, parameter_model + ) + self._api_client._verify_response(return_value) + return return_value + + def delete(self, *, name: str) -> types.DeleteResourceJob: + """Deletes a batch job. + + Args: + name (str): A fully-qualified BatchJob resource name or ID. + Example: "projects/.../locations/.../batchPredictionJobs/456" or "456" + when project and location are initialized in the client. + + Returns: + A DeleteResourceJob object that shows the status of the deletion. + + Usage: + + .. code-block:: python + + client.batches.delete(name='123456789') + """ + + parameter_model = types._DeleteBatchJobParameters( + name=name, + ) + + if not self._api_client.vertexai: + raise ValueError('This method is only supported in the Vertex AI client.') + else: + request_dict = _DeleteBatchJobParameters_to_vertex( + self._api_client, parameter_model + ) + path = 'batchPredictionJobs/{name}'.format_map(request_dict.get('_url')) + + query_params = request_dict.get('_query') + if query_params: + path = f'{path}?{urlencode(query_params)}' + # TODO: remove the hack that pops config. + config = request_dict.pop('config', None) + http_options = config.pop('httpOptions', None) if config else None + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response_dict = self._api_client.request( + 'delete', path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _DeleteResourceJob_from_vertex( + self._api_client, response_dict + ) + else: + response_dict = _DeleteResourceJob_from_mldev( + self._api_client, response_dict + ) + + return_value = types.DeleteResourceJob._from_response( + response_dict, parameter_model + ) + self._api_client._verify_response(return_value) + return return_value + + def create( + self, + *, + model: str, + src: str, + config: Optional[types.CreateBatchJobConfigOrDict] = None, + ) -> types.BatchJob: + """Creates a batch job. + + Args: + model (str): The model to use for the batch job. + src (str): The source of the batch job. Currently supports GCS URI(-s) or + BigQuery URI. Example: "gs://path/to/input/data" or + "bq://projectId.bqDatasetId.bqTableId". + config (CreateBatchJobConfig): Optional configuration for the batch job. + + Returns: + A BatchJob object that contains details about the batch job. + + Usage: + + .. code-block:: python + + batch_job = client.batches.create( + model="gemini-1.5-flash", + src="gs://path/to/input/data", + ) + print(batch_job.state) + """ + config = _extra_utils.format_destination(src, config) + return self._create(model=model, src=src, config=config) + + def list( + self, *, config: Optional[types.ListBatchJobConfigOrDict] = None + ) -> Pager[types.BatchJob]: + """Lists batch jobs. + + Args: + config (ListBatchJobConfig): Optional configuration for the list request. + + Returns: + A Pager object that contains one page of batch jobs. When iterating over + the pager, it automatically fetches the next page if there are more. + + Usage: + + .. code-block:: python + + batch_jobs = client.batches.list(config={"page_size": 10}) + for batch_job in batch_jobs: + print(f"Batch job: {batch_job.name}, state {batch_job.state}") + """ + return Pager( + 'batch_jobs', + self._list, + self._list(config=config), + config, + ) + + +class AsyncBatches(_common.BaseModule): + + async def _create( + self, + *, + model: str, + src: str, + config: Optional[types.CreateBatchJobConfigOrDict] = None, + ) -> types.BatchJob: + parameter_model = types._CreateBatchJobParameters( + model=model, + src=src, + config=config, + ) + + if not self._api_client.vertexai: + raise ValueError('This method is only supported in the Vertex AI client.') + else: + request_dict = _CreateBatchJobParameters_to_vertex( + self._api_client, parameter_model + ) + path = 'batchPredictionJobs'.format_map(request_dict.get('_url')) + + query_params = request_dict.get('_query') + if query_params: + path = f'{path}?{urlencode(query_params)}' + # TODO: remove the hack that pops config. + config = request_dict.pop('config', None) + http_options = config.pop('httpOptions', None) if config else None + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response_dict = await self._api_client.async_request( + 'post', path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _BatchJob_from_vertex(self._api_client, response_dict) + else: + response_dict = _BatchJob_from_mldev(self._api_client, response_dict) + + return_value = types.BatchJob._from_response(response_dict, parameter_model) + self._api_client._verify_response(return_value) + return return_value + + async def get( + self, *, name: str, config: Optional[types.GetBatchJobConfigOrDict] = None + ) -> types.BatchJob: + """Gets a batch job. + + Args: + name (str): A fully-qualified BatchJob resource name or ID. + Example: "projects/.../locations/.../batchPredictionJobs/456" or "456" + when project and location are initialized in the client. + + Returns: + A BatchJob object that contains details about the batch job. + + Usage: + + .. code-block:: python + + batch_job = client.batches.get(name='123456789') + print(f"Batch job: {batch_job.name}, state {batch_job.state}") + """ + + parameter_model = types._GetBatchJobParameters( + name=name, + config=config, + ) + + if not self._api_client.vertexai: + raise ValueError('This method is only supported in the Vertex AI client.') + else: + request_dict = _GetBatchJobParameters_to_vertex( + self._api_client, parameter_model + ) + path = 'batchPredictionJobs/{name}'.format_map(request_dict.get('_url')) + + query_params = request_dict.get('_query') + if query_params: + path = f'{path}?{urlencode(query_params)}' + # TODO: remove the hack that pops config. + config = request_dict.pop('config', None) + http_options = config.pop('httpOptions', None) if config else None + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response_dict = await self._api_client.async_request( + 'get', path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _BatchJob_from_vertex(self._api_client, response_dict) + else: + response_dict = _BatchJob_from_mldev(self._api_client, response_dict) + + return_value = types.BatchJob._from_response(response_dict, parameter_model) + self._api_client._verify_response(return_value) + return return_value + + async def cancel( + self, + *, + name: str, + config: Optional[types.CancelBatchJobConfigOrDict] = None, + ) -> None: + parameter_model = types._CancelBatchJobParameters( + name=name, + config=config, + ) + + if not self._api_client.vertexai: + raise ValueError('This method is only supported in the Vertex AI client.') + else: + request_dict = _CancelBatchJobParameters_to_vertex( + self._api_client, parameter_model + ) + path = 'batchPredictionJobs/{name}:cancel'.format_map( + request_dict.get('_url') + ) + + query_params = request_dict.get('_query') + if query_params: + path = f'{path}?{urlencode(query_params)}' + # TODO: remove the hack that pops config. + config = request_dict.pop('config', None) + http_options = config.pop('httpOptions', None) if config else None + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response_dict = await self._api_client.async_request( + 'post', path, request_dict, http_options + ) + + async def _list( + self, *, config: types.ListBatchJobConfigOrDict + ) -> types.ListBatchJobResponse: + parameter_model = types._ListBatchJobParameters( + config=config, + ) + + if not self._api_client.vertexai: + raise ValueError('This method is only supported in the Vertex AI client.') + else: + request_dict = _ListBatchJobParameters_to_vertex( + self._api_client, parameter_model + ) + path = 'batchPredictionJobs'.format_map(request_dict.get('_url')) + + query_params = request_dict.get('_query') + if query_params: + path = f'{path}?{urlencode(query_params)}' + # TODO: remove the hack that pops config. + config = request_dict.pop('config', None) + http_options = config.pop('httpOptions', None) if config else None + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response_dict = await self._api_client.async_request( + 'get', path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _ListBatchJobResponse_from_vertex( + self._api_client, response_dict + ) + else: + response_dict = _ListBatchJobResponse_from_mldev( + self._api_client, response_dict + ) + + return_value = types.ListBatchJobResponse._from_response( + response_dict, parameter_model + ) + self._api_client._verify_response(return_value) + return return_value + + async def delete(self, *, name: str) -> types.DeleteResourceJob: + """Deletes a batch job. + + Args: + name (str): A fully-qualified BatchJob resource name or ID. + Example: "projects/.../locations/.../batchPredictionJobs/456" or "456" + when project and location are initialized in the client. + + Returns: + A DeleteResourceJob object that shows the status of the deletion. + + Usage: + + .. code-block:: python + + client.batches.delete(name='123456789') + """ + + parameter_model = types._DeleteBatchJobParameters( + name=name, + ) + + if not self._api_client.vertexai: + raise ValueError('This method is only supported in the Vertex AI client.') + else: + request_dict = _DeleteBatchJobParameters_to_vertex( + self._api_client, parameter_model + ) + path = 'batchPredictionJobs/{name}'.format_map(request_dict.get('_url')) + + query_params = request_dict.get('_query') + if query_params: + path = f'{path}?{urlencode(query_params)}' + # TODO: remove the hack that pops config. + config = request_dict.pop('config', None) + http_options = config.pop('httpOptions', None) if config else None + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response_dict = await self._api_client.async_request( + 'delete', path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _DeleteResourceJob_from_vertex( + self._api_client, response_dict + ) + else: + response_dict = _DeleteResourceJob_from_mldev( + self._api_client, response_dict + ) + + return_value = types.DeleteResourceJob._from_response( + response_dict, parameter_model + ) + self._api_client._verify_response(return_value) + return return_value + + async def create( + self, + *, + model: str, + src: str, + config: Optional[types.CreateBatchJobConfigOrDict] = None, + ) -> types.BatchJob: + """Creates a batch job asynchronously. + + Args: + model (str): The model to use for the batch job. + src (str): The source of the batch job. Currently supports GCS URI(-s) or + BigQuery URI. Example: "gs://path/to/input/data" or + "bq://projectId.bqDatasetId.bqTableId". + config (CreateBatchJobConfig): Optional configuration for the batch job. + + Returns: + A BatchJob object that contains details about the batch job. + + Usage: + + .. code-block:: python + + batch_job = await client.aio.batches.create( + model="gemini-1.5-flash", + src="gs://path/to/input/data", + ) + """ + config = _extra_utils.format_destination(src, config) + return await self._create(model=model, src=src, config=config) + + async def list( + self, *, config: Optional[types.ListBatchJobConfigOrDict] = None + ) -> AsyncPager[types.BatchJob]: + """Lists batch jobs asynchronously. + + Args: + config (ListBatchJobConfig): Optional configuration for the list request. + + Returns: + A Pager object that contains one page of batch jobs. When iterating over + the pager, it automatically fetches the next page if there are more. + + Usage: + + .. code-block:: python + + batch_jobs = await client.aio.batches.list(config={'page_size': 5}) + print(f"current page: {batch_jobs.page}") + await batch_jobs_pager.next_page() + print(f"next page: {batch_jobs_pager.page}") + """ + return AsyncPager( + 'batch_jobs', + self._list, + await self._list(config=config), + config, + ) |