google-genai 0.8.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

google/genai/tunings.py

@@ -15,7 +15,8 @@
 
 # Code generated by the Google Gen AI SDK generator DO NOT EDIT.
 
-from typing import Optional, Union
+import time
+from typing import Any, Optional, Union
 from urllib.parse import urlencode
 from . import _api_module
 from . import _common
@@ -655,42 +656,50 @@ def _ListTuningJobsResponse_from_vertex(
   return to_object
 
 
-def _TuningJobOrOperation_from_mldev(
+def _Operation_from_mldev(
     api_client: ApiClient,
     from_object: Union[dict, object],
     parent_object: dict = None,
 ) -> dict:
   to_object = {}
-  if getv(from_object, ['_self']) is not None:
-    setv(
-        to_object,
-        ['tuning_job'],
-        _TuningJob_from_mldev(
-            api_client,
-            t.t_resolve_operation(api_client, getv(from_object, ['_self'])),
-            to_object,
-        ),
-    )
+  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'], getv(from_object, ['response']))
 
   return to_object
 
 
-def _TuningJobOrOperation_from_vertex(
+def _Operation_from_vertex(
     api_client: ApiClient,
     from_object: Union[dict, object],
     parent_object: dict = None,
 ) -> dict:
   to_object = {}
-  if getv(from_object, ['_self']) is not None:
-    setv(
-        to_object,
-        ['tuning_job'],
-        _TuningJob_from_vertex(
-            api_client,
-            t.t_resolve_operation(api_client, getv(from_object, ['_self'])),
-            to_object,
-        ),
-    )
+  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'], getv(from_object, ['response']))
 
   return to_object
 
@@ -823,7 +832,7 @@ class Tunings(_api_module.BaseModule):
       base_model: str,
       training_dataset: types.TuningDatasetOrDict,
       config: Optional[types.CreateTuningJobConfigOrDict] = None,
-  ) -> types.TuningJobOrOperation:
+  ) -> types.TuningJob:
     """Creates a supervised fine-tuning job.
 
     Args:
@@ -841,16 +850,76 @@ class Tunings(_api_module.BaseModule):
         config=config,
     )
 
-    if self._api_client.vertexai:
+    if not self._api_client.vertexai:
+      raise ValueError('This method is only supported in the Vertex AI client.')
+    else:
       request_dict = _CreateTuningJobParameters_to_vertex(
           self._api_client, parameter_model
       )
       path = 'tuningJobs'.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.
+    request_dict.pop('config', None)
+
+    http_options = None
+    if isinstance(config, dict):
+      http_options = config.get('http_options', None)
+    elif hasattr(config, 'http_options'):
+      http_options = config.http_options
+
+    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 = _TuningJob_from_vertex(self._api_client, response_dict)
+    else:
+      response_dict = _TuningJob_from_mldev(self._api_client, response_dict)
+
+    return_value = types.TuningJob._from_response(
+        response=response_dict, kwargs=parameter_model
+    )
+    self._api_client._verify_response(return_value)
+    return return_value
+
+  def _tune_mldev(
+      self,
+      *,
+      base_model: str,
+      training_dataset: types.TuningDatasetOrDict,
+      config: Optional[types.CreateTuningJobConfigOrDict] = None,
+  ) -> types.Operation:
+    """Creates a supervised fine-tuning job.
+
+    Args:
+      base_model: The name of the model to tune.
+      training_dataset: The training dataset to use.
+      config: The configuration to use for the tuning job.
+
+    Returns:
+      A TuningJob operation.
+    """
+
+    parameter_model = types._CreateTuningJobParameters(
+        base_model=base_model,
+        training_dataset=training_dataset,
+        config=config,
+    )
+
+    if self._api_client.vertexai:
+      raise ValueError('This method is only supported in the default client.')
     else:
       request_dict = _CreateTuningJobParameters_to_mldev(
           self._api_client, parameter_model
       )
       path = 'tunedModels'.format_map(request_dict.get('_url'))
+
     query_params = request_dict.get('_query')
     if query_params:
       path = f'{path}?{urlencode(query_params)}'
@@ -871,17 +940,13 @@ class Tunings(_api_module.BaseModule):
     )
 
     if self._api_client.vertexai:
-      response_dict = _TuningJobOrOperation_from_vertex(
-          self._api_client, response_dict
-      )
+      response_dict = _Operation_from_vertex(self._api_client, response_dict)
     else:
-      response_dict = _TuningJobOrOperation_from_mldev(
-          self._api_client, response_dict
-      )
+      response_dict = _Operation_from_mldev(self._api_client, response_dict)
 
-    return_value = types.TuningJobOrOperation._from_response(
+    return_value = types.Operation._from_response(
         response=response_dict, kwargs=parameter_model
-    ).tuning_job
+    )
     self._api_client._verify_response(return_value)
     return return_value
 
@@ -909,21 +974,106 @@ class Tunings(_api_module.BaseModule):
       )
     return job
 
+  @_common.experimental_warning(
+      "The SDK's tuning implementation is experimental, "
+      'and may change in future versions.',
+  )
   def tune(
       self,
       *,
       base_model: str,
       training_dataset: types.TuningDatasetOrDict,
       config: Optional[types.CreateTuningJobConfigOrDict] = None,
-  ) -> types.TuningJobOrOperation:
-    result = self._tune(
-        base_model=base_model,
-        training_dataset=training_dataset,
-        config=config,
-    )
-    if result.name and self._api_client.vertexai:
-      _IpythonUtils.display_model_tuning_button(tuning_job_resource=result.name)
-    return result
+  ) -> types.TuningJob:
+    if self._api_client.vertexai:
+      tuning_job = self._tune(
+          base_model=base_model,
+          training_dataset=training_dataset,
+          config=config,
+      )
+    else:
+      operation = self._tune_mldev(
+          base_model=base_model,
+          training_dataset=training_dataset,
+          config=config,
+      )
+      operation_dict = operation.to_json_dict()
+      tuned_model_dict = _resolve_operation(self._api_client, operation_dict)
+      tuning_job_dict = _TuningJob_from_mldev(
+          self._api_client, tuned_model_dict
+      )
+      tuning_job = types.TuningJob._from_response(tuning_job_dict, None)
+    if tuning_job.name and self._api_client.vertexai:
+      _IpythonUtils.display_model_tuning_button(
+          tuning_job_resource=tuning_job.name
+      )
+    return tuning_job
+
+
+_LRO_POLLING_INITIAL_DELAY_SECONDS = 1.0
+_LRO_POLLING_MAXIMUM_DELAY_SECONDS = 20.0
+_LRO_POLLING_TIMEOUT_SECONDS = 900.0
+_LRO_POLLING_MULTIPLIER = 1.5
+
+
+def _resolve_operation(api_client: ApiClient, struct: dict[str, Any]):
+  if (name := struct.get('name')) and '/operations/' in name:
+    operation: dict[str, Any] = struct
+    total_seconds = 0.0
+    delay_seconds = _LRO_POLLING_INITIAL_DELAY_SECONDS
+    while not operation.get('done'):
+      if total_seconds > _LRO_POLLING_TIMEOUT_SECONDS:
+        raise RuntimeError(f'Operation {name} timed out.\n{operation}')
+      # TODO(b/374433890): Replace with LRO module once it's available.
+      operation: dict[str, Any] = api_client.request(
+          http_method='GET', path=name, request_dict={}
+      )
+      if 'ReplayApiClient' not in type(api_client).__name__:
+        time.sleep(delay_seconds)
+      total_seconds += total_seconds
+      # Exponential backoff
+      delay_seconds = min(
+          delay_seconds * _LRO_POLLING_MULTIPLIER,
+          _LRO_POLLING_MAXIMUM_DELAY_SECONDS,
+      )
+    if error := operation.get('error'):
+      raise RuntimeError(
+          f'Operation {name} failed with error: {error}.\n{operation}'
+      )
+    return operation.get('response')
+  else:
+    return struct
+
+
+async def _resolve_operation_async(
+    api_client: ApiClient, struct: dict[str, Any]
+):
+  if (name := struct.get('name')) and '/operations/' in name:
+    operation: dict[str, Any] = struct
+    total_seconds = 0.0
+    delay_seconds = _LRO_POLLING_INITIAL_DELAY_SECONDS
+    while not operation.get('done'):
+      if total_seconds > _LRO_POLLING_TIMEOUT_SECONDS:
+        raise RuntimeError(f'Operation {name} timed out.\n{operation}')
+      # TODO(b/374433890): Replace with LRO module once it's available.
+      operation: dict[str, Any] = await api_client.async_request(
+          http_method='GET', path=name, request_dict={}
+      )
+      if 'ReplayApiClient' not in type(api_client).__name__:
+        time.sleep(delay_seconds)
+      total_seconds += total_seconds
+      # Exponential backoff
+      delay_seconds = min(
+          delay_seconds * _LRO_POLLING_MULTIPLIER,
+          _LRO_POLLING_MAXIMUM_DELAY_SECONDS,
+      )
+    if error := operation.get('error'):
+      raise RuntimeError(
+          f'Operation {name} failed with error: {error}.\n{operation}'
+      )
+    return operation.get('response')
+  else:
+    return struct
 
 
 class AsyncTunings(_api_module.BaseModule):
@@ -1054,7 +1204,7 @@ class AsyncTunings(_api_module.BaseModule):
       base_model: str,
       training_dataset: types.TuningDatasetOrDict,
       config: Optional[types.CreateTuningJobConfigOrDict] = None,
-  ) -> types.TuningJobOrOperation:
+  ) -> types.TuningJob:
     """Creates a supervised fine-tuning job.
 
     Args:
@@ -1072,16 +1222,76 @@ class AsyncTunings(_api_module.BaseModule):
         config=config,
     )
 
-    if self._api_client.vertexai:
+    if not self._api_client.vertexai:
+      raise ValueError('This method is only supported in the Vertex AI client.')
+    else:
       request_dict = _CreateTuningJobParameters_to_vertex(
           self._api_client, parameter_model
       )
       path = 'tuningJobs'.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.
+    request_dict.pop('config', None)
+
+    http_options = None
+    if isinstance(config, dict):
+      http_options = config.get('http_options', None)
+    elif hasattr(config, 'http_options'):
+      http_options = config.http_options
+
+    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 = _TuningJob_from_vertex(self._api_client, response_dict)
+    else:
+      response_dict = _TuningJob_from_mldev(self._api_client, response_dict)
+
+    return_value = types.TuningJob._from_response(
+        response=response_dict, kwargs=parameter_model
+    )
+    self._api_client._verify_response(return_value)
+    return return_value
+
+  async def _tune_mldev(
+      self,
+      *,
+      base_model: str,
+      training_dataset: types.TuningDatasetOrDict,
+      config: Optional[types.CreateTuningJobConfigOrDict] = None,
+  ) -> types.Operation:
+    """Creates a supervised fine-tuning job.
+
+    Args:
+      base_model: The name of the model to tune.
+      training_dataset: The training dataset to use.
+      config: The configuration to use for the tuning job.
+
+    Returns:
+      A TuningJob operation.
+    """
+
+    parameter_model = types._CreateTuningJobParameters(
+        base_model=base_model,
+        training_dataset=training_dataset,
+        config=config,
+    )
+
+    if self._api_client.vertexai:
+      raise ValueError('This method is only supported in the default client.')
     else:
       request_dict = _CreateTuningJobParameters_to_mldev(
           self._api_client, parameter_model
       )
       path = 'tunedModels'.format_map(request_dict.get('_url'))
+
     query_params = request_dict.get('_query')
     if query_params:
       path = f'{path}?{urlencode(query_params)}'
@@ -1102,17 +1312,13 @@ class AsyncTunings(_api_module.BaseModule):
     )
 
     if self._api_client.vertexai:
-      response_dict = _TuningJobOrOperation_from_vertex(
-          self._api_client, response_dict
-      )
+      response_dict = _Operation_from_vertex(self._api_client, response_dict)
     else:
-      response_dict = _TuningJobOrOperation_from_mldev(
-          self._api_client, response_dict
-      )
+      response_dict = _Operation_from_mldev(self._api_client, response_dict)
 
-    return_value = types.TuningJobOrOperation._from_response(
+    return_value = types.Operation._from_response(
         response=response_dict, kwargs=parameter_model
-    ).tuning_job
+    )
     self._api_client._verify_response(return_value)
     return return_value
 
@@ -1140,21 +1346,42 @@ class AsyncTunings(_api_module.BaseModule):
       )
     return job
 
+  @_common.experimental_warning(
+      "The SDK's tuning implementation is experimental, "
+      'and may change in future versions.'
+  )
   async def tune(
       self,
       *,
       base_model: str,
       training_dataset: types.TuningDatasetOrDict,
       config: Optional[types.CreateTuningJobConfigOrDict] = None,
-  ) -> types.TuningJobOrOperation:
-    result = await self._tune(
-        base_model=base_model,
-        training_dataset=training_dataset,
-        config=config,
-    )
-    if result.name and self._api_client.vertexai:
-      _IpythonUtils.display_model_tuning_button(tuning_job_resource=result.name)
-    return result
+  ) -> types.TuningJob:
+    if self._api_client.vertexai:
+      tuning_job = await self._tune(
+          base_model=base_model,
+          training_dataset=training_dataset,
+          config=config,
+      )
+    else:
+      operation = await self._tune_mldev(
+          base_model=base_model,
+          training_dataset=training_dataset,
+          config=config,
+      )
+      operation_dict = operation.to_json_dict()
+      tuned_model_dict = await _resolve_operation_async(
+          self._api_client, operation_dict
+      )
+      tuning_job_dict = _TuningJob_from_mldev(
+          self._api_client, tuned_model_dict
+      )
+      tuning_job = types.TuningJob._from_response(tuning_job_dict, None)
+    if tuning_job.name and self._api_client.vertexai:
+      _IpythonUtils.display_model_tuning_button(
+          tuning_job_resource=tuning_job.name
+      )
+    return tuning_job
 
 
 class _IpythonUtils: