chore: GenAI - Enforce formatting in the vertexai module

Also, format all existing modules.

PiperOrigin-RevId: 637170526

Author: Ark-kun

noxfile.py

@@ -28,7 +28,7 @@
 FLAKE8_VERSION = "flake8==6.1.0"
 BLACK_VERSION = "black==22.3.0"
 ISORT_VERSION = "isort==5.10.1"
-LINT_PATHS = ["docs", "google", "tests", "noxfile.py", "setup.py"]
+LINT_PATHS = ["docs", "google", "vertexai", "tests", "noxfile.py", "setup.py"]
 
 DEFAULT_PYTHON_VERSION = "3.8"
 

vertexai/extensions/_extensions.py

@@ -95,22 +95,20 @@ def __init__(self, extension_name: str):
         self.execution_api_client = initializer.global_config.create_client(
             client_class=aip_utils.ExtensionExecutionClientWithOverride,
         )
-        self._gca_resource = self._get_gca_resource(
-            resource_name=extension_name
-        )
+        self._gca_resource = self._get_gca_resource(resource_name=extension_name)
         self._api_spec = None
         self._operation_schemas = None
 
     @classmethod
     def create(
-            cls,
-            manifest: Union[_utils.JsonDict, types.ExtensionManifest],
-            *,
-            extension_name: Optional[str] = None,
-            display_name: Optional[str] = None,
-            description: Optional[str] = None,
-            runtime_config: Optional[_RuntimeConfigOrJson] = None,
-        ):
+        cls,
+        manifest: Union[_utils.JsonDict, types.ExtensionManifest],
+        *,
+        extension_name: Optional[str] = None,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        runtime_config: Optional[_RuntimeConfigOrJson] = None,
+    ):
         """Creates a new Extension.
 
         Args:
@@ -150,7 +148,8 @@ def create(
         )
         if runtime_config:
             extension.runtime_config = _utils.to_proto(
-                runtime_config, types.RuntimeConfig(),
+                runtime_config,
+                types.RuntimeConfig(),
             )
         operation_future = sdk_resource.api_client.import_extension(
             parent=initializer.global_config.common_location_path(),
@@ -169,10 +168,8 @@ def create(
         sdk_resource._gca_resource = sdk_resource._get_gca_resource(
             resource_name=created_extension.name
         )
-        sdk_resource.execution_api_client = (
-            initializer.global_config.create_client(
-                client_class=aip_utils.ExtensionExecutionClientWithOverride,
-            )
+        sdk_resource.execution_api_client = initializer.global_config.create_client(
+            client_class=aip_utils.ExtensionExecutionClientWithOverride,
         )
         sdk_resource._api_spec = None
         sdk_resource._operation_schemas = None
@@ -186,9 +183,7 @@ def resource_name(self) -> str:
     def api_spec(self) -> _utils.JsonDict:
         """Returns the (Open)API Spec of the extension."""
         if self._api_spec is None:
-            self._api_spec = _load_api_spec(
-                self._gca_resource.manifest.api_spec
-            )
+            self._api_spec = _load_api_spec(self._gca_resource.manifest.api_spec)
         return self._api_spec
 
     def operation_schemas(self) -> Sequence[_utils.JsonDict]:
@@ -201,11 +196,11 @@ def operation_schemas(self) -> Sequence[_utils.JsonDict]:
         return self._operation_schemas
 
     def execute(
-            self,
-            operation_id: str,
-            operation_params: Optional[_StructOrJson] = None,
-            runtime_auth_config: Optional[_AuthConfigOrJson] = None,
-        ) -> Union[_utils.JsonDict, str]:
+        self,
+        operation_id: str,
+        operation_params: Optional[_StructOrJson] = None,
+        runtime_auth_config: Optional[_AuthConfigOrJson] = None,
+    ) -> Union[_utils.JsonDict, str]:
         """Executes an operation of the extension with the specified params.
 
         Args:
@@ -230,7 +225,8 @@ def execute(
         )
         if runtime_auth_config:
             request.runtime_auth_config = _utils.to_proto(
-                runtime_auth_config, types.AuthConfig(),
+                runtime_auth_config,
+                types.AuthConfig(),
             )
         response = self.execution_api_client.execute_extension(request)
         return _try_parse_execution_response(response)
@@ -263,7 +259,8 @@ def from_hub(
         """
         if runtime_config:
             runtime_config = _utils.to_proto(
-                runtime_config, types.RuntimeConfig(),
+                runtime_config,
+                types.RuntimeConfig(),
             )
         if name == "code_interpreter":
             if runtime_config and not getattr(
@@ -301,8 +298,8 @@ def from_hub(
 
 
 def _try_parse_execution_response(
-        response: types.ExecuteExtensionResponse
-    ) -> Union[_utils.JsonDict, str]:
+    response: types.ExecuteExtensionResponse,
+) -> Union[_utils.JsonDict, str]:
     content: str = response.content
     try:
         content = json.loads(response.content)

vertexai/generative_models/_function_calling_utils.py

@@ -71,24 +71,23 @@ def _generate_json_schema_from_function_using_pydantic(
         name: (
             # 1. We infer the argument type here: use Any rather than None so
             # it will not try to auto-infer the type based on the default value.
-            (
-                param.annotation if param.annotation != inspect.Parameter.empty
-                else Any
-            ),
+            (param.annotation if param.annotation != inspect.Parameter.empty else Any),
             pydantic.Field(
                 # 2. We do not support default values for now.
                 default=(
-                    param.default if param.default != inspect.Parameter.empty
+                    param.default
+                    if param.default != inspect.Parameter.empty
                     # ! Need to use Undefined instead of None
                     else pydantic_fields.Undefined
                 ),
                 # 3. We support user-provided descriptions.
                 description=parameter_descriptions.get(name, None),
-            )
+            ),
         )
         for name, param in defaults.items()
         # We do not support *args or **kwargs
-        if param.kind in (
+        if param.kind
+        in (
             inspect.Parameter.POSITIONAL_OR_KEYWORD,
             inspect.Parameter.KEYWORD_ONLY,
             inspect.Parameter.POSITIONAL_ONLY,
@@ -105,10 +104,9 @@ def _generate_json_schema_from_function_using_pydantic(
         #     * https://github.com/pydantic/pydantic/issues/1270
         #     * https://stackoverflow.com/a/58841311
         #     * https://github.com/pydantic/pydantic/discussions/4872
-        if (
-                typing.get_origin(annotation) is typing.Union
-                and type(None) in typing.get_args(annotation)
-            ):
+        if typing.get_origin(annotation) is typing.Union and type(
+            None
+        ) in typing.get_args(annotation):
             # for "typing.Optional" arguments, function_arg might be a
             # dictionary like
             #
@@ -121,9 +119,12 @@ def _generate_json_schema_from_function_using_pydantic(
             property_schema["nullable"] = True
     # 6. Annotate required fields.
     function_schema["required"] = [
-        k for k in defaults if (
+        k
+        for k in defaults
+        if (
             defaults[k].default == inspect.Parameter.empty
-            and defaults[k].kind in (
+            and defaults[k].kind
+            in (
                 inspect.Parameter.POSITIONAL_OR_KEYWORD,
                 inspect.Parameter.KEYWORD_ONLY,
                 inspect.Parameter.POSITIONAL_ONLY,

vertexai/generative_models/_generative_models.py

@@ -1660,8 +1660,7 @@ def text(self) -> str:
             raise ValueError(
                 "Response has no candidates (and thus no text)."
                 " The response is likely blocked by the safety filters.\n"
-                "Response:\n"
-                + _dict_to_pretty_string(self.to_dict())
+                "Response:\n" + _dict_to_pretty_string(self.to_dict())
             )
         try:
             return self.candidates[0].text
@@ -1671,8 +1670,7 @@ def text(self) -> str:
             raise ValueError(
                 "Cannot get the response text.\n"
                 f"{e}\n"
-                "Response:\n"
-                + _dict_to_pretty_string(self.to_dict())
+                "Response:\n" + _dict_to_pretty_string(self.to_dict())
             ) from e
 
     @property
@@ -1754,8 +1752,7 @@ def text(self) -> str:
             raise ValueError(
                 "Cannot get the Candidate text.\n"
                 f"{e}\n"
-                "Candidate:\n"
-                + _dict_to_pretty_string(self.to_dict())
+                "Candidate:\n" + _dict_to_pretty_string(self.to_dict())
             ) from e
 
     @property
@@ -1830,8 +1827,7 @@ def text(self) -> str:
             raise ValueError(
                 "Response candidate content has no parts (and thus no text)."
                 " The candidate is likely blocked by the safety filters.\n"
-                "Content:\n"
-                + _dict_to_pretty_string(self.to_dict())
+                "Content:\n" + _dict_to_pretty_string(self.to_dict())
             )
         return self.parts[0].text
 
@@ -1921,8 +1917,7 @@ def text(self) -> str:
         if "text" not in self._raw_part:
             raise AttributeError(
                 "Response candidate content part has no text.\n"
-                "Part:\n"
-                + _dict_to_pretty_string(self.to_dict())
+                "Part:\n" + _dict_to_pretty_string(self.to_dict())
             )
         return self._raw_part.text
 
@@ -2023,8 +2018,7 @@ class GoogleSearchRetrieval:
         """
 
         def __init__(self):
-            """Initializes a Google Search Retrieval tool.
-            """
+            """Initializes a Google Search Retrieval tool."""
             self._raw_google_search_retrieval = gapic_tool_types.GoogleSearchRetrieval()
 
 
@@ -2400,7 +2394,9 @@ def respond_to_model_response(
                     )
                 callable_function = None
                 for tool in tools:
-                    new_callable_function = tool._callable_functions.get(function_call.name)
+                    new_callable_function = tool._callable_functions.get(
+                        function_call.name
+                    )
                     if new_callable_function and callable_function:
                         raise ValueError(
                             "Multiple functions with the same name are not supported."

vertexai/language_models/_distillation.py

@@ -105,9 +105,7 @@ def submit_distillation_pipeline_job(
         pipeline_arguments["learning_rate_multiplier"] = learning_rate_multiplier
     if evaluation_spec is not None:
         pipeline_arguments["evaluation_data_uri"] = evaluation_spec.evaluation_data
-        pipeline_arguments[
-            "evaluation_interval"
-        ] = evaluation_spec.evaluation_interval
+        pipeline_arguments["evaluation_interval"] = evaluation_spec.evaluation_interval
         pipeline_arguments[
             "enable_early_stopping"
         ] = evaluation_spec.enable_early_stopping
@@ -126,8 +124,7 @@ def submit_distillation_pipeline_job(
         pipeline_arguments["max_context_length"] = max_context_length
     if model_display_name is None:
         model_display_name = (
-            f"{student_short_model_id}"
-            f" distilled from {teacher_short_model_id}"
+            f"{student_short_model_id} distilled from {teacher_short_model_id}"
         )
     pipeline_arguments["model_display_name"] = model_display_name
     # # Not exposing these parameters:

vertexai/preview/extensions.py

@@ -20,6 +20,4 @@
     Extension,
 )
 
-__all__ = (
-    "Extension",
-)
+__all__ = ("Extension",)

vertexai/preview/reasoning_engines/__init__.py

@@ -21,7 +21,7 @@
     ReasoningEngine,
 )
 from vertexai.preview.reasoning_engines.templates.langchain import (
-    LangchainAgent
+    LangchainAgent,
 )
 
 __all__ = (