Source code for langchain_community.utilities.vertexai

"""初始化 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
    from vertexai.preview.generative_models import Image


def create_retry_decorator(
    llm: BaseLLM,
    *,
    max_retries: int = 1,
    run_manager: Optional[
        Union[AsyncCallbackManagerForLLMRun, CallbackManagerForLLMRun]
    ] = None,
) -> Callable[[Any], Any]:
    """为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


def raise_vertex_import_error(minimum_expected_version: str = "1.38.0") -> None:
    """引发与Vertex SDK不可用相关的ImportError。

参数：
    minimum_expected_version：SDK的最低预期版本。
引发：
    ImportError：提到SDK所需版本的ImportError。
"""
    raise ImportError(
        "Please, install or upgrade the google-cloud-aiplatform library: "
        f"pip install google-cloud-aiplatform>={minimum_expected_version}"
    )


def init_vertexai(
    project: Optional[str] = None,
    location: Optional[str] = None,
    credentials: Optional["Credentials"] = None,
) -> None:
    """初始化 Vertex AI。

参数：
    project: 在进行 Vertex API 调用时要使用的默认 GCP 项目。
    location: 在进行 API 调用时要使用的默认位置。
    credentials: 在进行 API 调用时要使用的默认自定义凭据。如果未提供凭据，则将从环境中获取。

抛出：
    ImportError: 如果导入 vertexai SDK 失败。
"""
    try:
        import vertexai
    except ImportError:
        raise_vertex_import_error()

    vertexai.init(
        project=project,
        location=location,
        credentials=credentials,
    )


def get_client_info(module: Optional[str] = None) -> "ClientInfo":
    r"""返回一个自定义用户代理标头。

参数：
    module（可选[str]）：
        可选。用于自定义用户代理标头的模块。
返回：
    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}",
    )


def load_image_from_gcs(path: str, project: Optional[str] = None) -> "Image":
    """从Google Cloud Storage加载图像。"""
    try:
        from google.cloud import storage
    except ImportError:
        raise ImportError("Could not import google-cloud-storage python package.")
    from vertexai.preview.generative_models import Image

    gcs_client = storage.Client(project=project)
    pieces = path.split("/")
    blobs = list(gcs_client.list_blobs(pieces[2], prefix="/".join(pieces[3:])))
    if len(blobs) > 1:
        raise ValueError(f"Found more than one candidate for {path}!")
    return Image.from_bytes(blobs[0].download_as_bytes())