Source code for langchain_community.utilities.vertexai
"""Utilities to init Vertex AI."""
from importlib import metadata
from typing import TYPE_CHECKING, Any, Callable, Optional, Union
from langchain_core.callbacks import (
AsyncCallbackManagerForLLMRun,
CallbackManagerForLLMRun,
)
from langchain_core.language_models.llms import BaseLLM, create_base_retry_decorator
if TYPE_CHECKING:
from google.api_core.gapic_v1.client_info import ClientInfo
from google.auth.credentials import Credentials
[docs]def create_retry_decorator(
llm: BaseLLM,
*,
max_retries: int = 1,
run_manager: Optional[
Union[AsyncCallbackManagerForLLMRun, CallbackManagerForLLMRun]
] = None,
) -> Callable[[Any], Any]:
"""Creates a retry decorator for Vertex / Palm LLMs."""
import google.api_core
errors = [
google.api_core.exceptions.ResourceExhausted,
google.api_core.exceptions.ServiceUnavailable,
google.api_core.exceptions.Aborted,
google.api_core.exceptions.DeadlineExceeded,
google.api_core.exceptions.GoogleAPIError,
]
decorator = create_base_retry_decorator(
error_types=errors, max_retries=max_retries, run_manager=run_manager
)
return decorator
[docs]def raise_vertex_import_error(minimum_expected_version: str = "1.36.0") -> None:
"""Raise ImportError related to Vertex SDK being not available.
Args:
minimum_expected_version: The lowest expected version of the SDK.
Raises:
ImportError: an ImportError that mentions a required version of the SDK.
"""
raise ImportError(
"Please, install or upgrade the google-cloud-aiplatform library: "
f"pip install google-cloud-aiplatform>={minimum_expected_version}"
)
[docs]def init_vertexai(
project: Optional[str] = None,
location: Optional[str] = None,
credentials: Optional["Credentials"] = None,
) -> None:
"""Init vertexai.
Args:
project: The default GCP project to use when making Vertex API calls.
location: The default location to use when making API calls.
credentials: The default custom
credentials to use when making API calls. If not provided credentials
will be ascertained from the environment.
Raises:
ImportError: If importing vertexai SDK did not succeed.
"""
try:
import vertexai
except ImportError:
raise_vertex_import_error()
vertexai.init(
project=project,
location=location,
credentials=credentials,
)
[docs]def get_client_info(module: Optional[str] = None) -> "ClientInfo":
r"""Returns a custom user agent header.
Args:
module (Optional[str]):
Optional. The module for a custom user agent header.
Returns:
google.api_core.gapic_v1.client_info.ClientInfo
"""
try:
from google.api_core.gapic_v1.client_info import ClientInfo
except ImportError as exc:
raise ImportError(
"Could not import ClientInfo. Please, install it with "
"pip install google-api-core"
) from exc
langchain_version = metadata.version("langchain")
client_library_version = (
f"{langchain_version}-{module}" if module else langchain_version
)
return ClientInfo(
client_library_version=client_library_version,
user_agent=f"langchain/{client_library_version}",
)