Source code for langchain.chat_models.vertexai
"""Wrapper around Google VertexAI chat-based models."""
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from pydantic import root_validator
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.chat_models.base import BaseChatModel
from langchain.llms.vertexai import _VertexAICommon, is_codey_model
from langchain.schema import (
    ChatGeneration,
    ChatResult,
)
from langchain.schema.messages import (
    AIMessage,
    BaseMessage,
    HumanMessage,
    SystemMessage,
)
from langchain.utilities.vertexai import raise_vertex_import_error
if TYPE_CHECKING:
    from vertexai.language_models import ChatMessage, InputOutputTextPair
@dataclass
class _ChatHistory:
    """Represents a context and a history of messages."""
    history: List["ChatMessage"] = field(default_factory=list)
    context: Optional[str] = None
def _parse_chat_history(history: List[BaseMessage]) -> _ChatHistory:
    """Parse a sequence of messages into history.
    Args:
        history: The list of messages to re-create the history of the chat.
    Returns:
        A parsed chat history.
    Raises:
        ValueError: If a sequence of message has a SystemMessage not at the
        first place.
    """
    from vertexai.language_models import ChatMessage
    vertex_messages, context = [], None
    for i, message in enumerate(history):
        if i == 0 and isinstance(message, SystemMessage):
            context = message.content
        elif isinstance(message, AIMessage):
            vertex_message = ChatMessage(content=message.content, author="bot")
            vertex_messages.append(vertex_message)