Source code for langchain.llms.vertexai
from __future__ import annotations
import asyncio
from concurrent.futures import Executor, ThreadPoolExecutor
from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, List, Optional
from pydantic import BaseModel, root_validator
from langchain.callbacks.manager import (
    AsyncCallbackManagerForLLMRun,
    CallbackManagerForLLMRun,
)
from langchain.llms.base import LLM, create_base_retry_decorator
from langchain.llms.utils import enforce_stop_tokens
from langchain.utilities.vertexai import (
    init_vertexai,
    raise_vertex_import_error,
)
if TYPE_CHECKING:
    from vertexai.language_models._language_models import _LanguageModel
[docs]def is_codey_model(model_name: str) -> bool:
    """Returns True if the model name is a Codey model.
    Args:
        model_name: The model name to check.
    Returns: True if the model name is a Codey model.
    """
    return "code" in model_name
def _create_retry_decorator(llm: VertexAI) -> Callable[[Any], Any]:
    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,
    ]
    decorator = create_base_retry_decorator(
        error_types=errors, max_retries=llm.max_retries  # type: ignore
    )
    return decorator
[docs]def completion_with_retry(llm: VertexAI, *args: Any, **kwargs: Any) -> Any:
    """Use tenacity to retry the completion call."""
    retry_decorator = _create_retry_decorator(llm)
    @retry_decorator