Source code for langchain.llms.vertexai
from __future__ import annotations
from concurrent.futures import Executor, ThreadPoolExecutor
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Dict,
    Iterator,
    List,
    Optional,
    Union,
)
from langchain.callbacks.manager import (
    AsyncCallbackManagerForLLMRun,
    CallbackManagerForLLMRun,
)
from langchain.llms.base import BaseLLM, create_base_retry_decorator
from langchain.pydantic_v1 import BaseModel, root_validator
from langchain.schema import (
    Generation,
    LLMResult,
)
from langchain.schema.output import GenerationChunk
from langchain.utilities.vertexai import (
    init_vertexai,
    raise_vertex_import_error,
)
if TYPE_CHECKING:
    from google.cloud.aiplatform.gapic import (
        PredictionServiceAsyncClient,
        PredictionServiceClient,
    )
    from vertexai.language_models._language_models import (
        TextGenerationResponse,
        _LanguageModel,
    )
def _response_to_generation(
    response: TextGenerationResponse,
) -> GenerationChunk:
    """Convert a stream response to a generation chunk."""
    try:
        generation_info = {
            "is_blocked": response.is_blocked,
            "safety_attributes": response.safety_attributes,
        }
    except Exception:
        generation_info = None
    return GenerationChunk(text=response.text, generation_info=generation_info)
[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.