Repository: oumi-ai/oumi
File: src/oumi/inference/gcp_inference_engine.py (Python)

Code before:

    @override
    def _set_required_fields_for_inference(self, remote_params: RemoteParams) -> None:
        """Set required fields for inference."""
        if (
            not remote_params.api_url
            and not self._remote_params.api_url
            and not self.base_url
        ):
            if self._project_id and self._region:
                project_id = self._project_id
                region = self._region
            elif os.getenv(self._project_id_env_key) and os.getenv(
                self._region_env_key
            ):
                project_id = os.getenv(self._project_id_env_key)
                region = os.getenv(self._region_env_key)
            else:
                raise ValueError(
                    "The GCP inference engine requires that either the `api_url` is "
                    "set or that both the `project_id` and `region` are set. "
                    "You can set the `api_url` in `RemoteParams`. "
                    "Alternatively, you can set the `project_id` and `region` when "
                    "constructing the inference engine (GoogleVertexInferenceEngine), "
                    f"or as environment variables: `{self._project_id_env_key}` and "
                    f"`{self._region_env_key}`, respectively."
                )

            remote_params.api_url = self._API_URL_TEMPLATE.format(
                project_id=project_id,
                region=region,
            )

        super()._set_required_fields_for_inference(remote_params)

    @override
    def _get_api_key(self, remote_params: RemoteParams) -> str:
        """Gets the authentication token for GCP."""
        try:
            from google.auth import default  # pyright: ignore[reportMissingImports]
            from google.auth.transport.requests import (  # pyright: ignore[reportMissingImports]
                Request,
            )
            from google.oauth2 import (  # pyright: ignore[reportMissingImports]
                service_account,
            )
        except ModuleNotFoundError:
            raise RuntimeError(
                "Google-auth is not installed. "
                "Please install oumi with GCP extra:`pip install oumi[gcp]`, "
                "or install google-auth with `pip install google-auth`."

Review comment (suggestion):
```suggestion
                    "This inference engine requires that either `api_url` is "
                    "set in `RemoteParams` or that both `project_id` and `region` are set. "
                    "You can set the `project_id` and `region` when "
                    "constructing a GoogleVertexInferenceEngine, "
                    f"or as environment variables: `{self._project_id_env_key}` and "
                    f"`{self._region_env_key}`."
```

Code after:

        self._project_id_env_key = (
            project_id_env_key or self._DEFAULT_PROJECT_ID_ENV_KEY
        )
        self._region_env_key = region_env_key or self._DEFAULT_REGION_ENV_KEY
        self._project_id = project_id
        self._region = region

    @override
    def _set_required_fields_for_inference(self, remote_params: RemoteParams) -> None:
        """Set required fields for inference."""
        if (
            not remote_params.api_url
            and not self._remote_params.api_url
            and not self.base_url
        ):
            if self._project_id and self._region:
                project_id = self._project_id
                region = self._region
            elif os.getenv(self._project_id_env_key) and os.getenv(
                self._region_env_key
            ):
                project_id = os.getenv(self._project_id_env_key)
                region = os.getenv(self._region_env_key)
            else:
                raise ValueError(
                    "This inference engine requires that either `api_url` is set in "
                    "`RemoteParams` or that both `project_id` and `region` are set. "
                    "You can set the `project_id` and `region` when "
                    "constructing a GoogleVertexInferenceEngine, "
                    f"or as environment variables: `{self._project_id_env_key}` and "
                    f"`{self._region_env_key}`."
                )

            remote_params.api_url = self._API_URL_TEMPLATE.format(
                project_id=project_id,
                region=region,
            )

        super()._set_required_fields_for_inference(remote_params)

    @override
    def _get_api_key(self, remote_params: RemoteParams) -> str:
        """Gets the authentication token for GCP."""
        try:
            from google.auth import default  # pyright: ignore[reportMissingImports]
            from google.auth.transport.requests import (  # pyright: ignore[reportMissingImports]
                Request,
            )
            from google.oauth2 import (  # pyright: ignore[reportMissingImports]
                service_account,