"""Wrapper around Google's PaLM Chat API."""
from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional

from pydantic import BaseModel, root_validator
from tenacity import (
    before_sleep_log,
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from langchain.callbacks.manager import (
    AsyncCallbackManagerForLLMRun,
    CallbackManagerForLLMRun,
)
from langchain.chat_models.base import BaseChatModel
from langchain.schema import (
    AIMessage,
    BaseMessage,
    ChatGeneration,
    ChatMessage,
    ChatResult,
    HumanMessage,
    SystemMessage,
)
from langchain.utils import get_from_dict_or_env

if TYPE_CHECKING:
    import google.generativeai as genai

logger = logging.getLogger(__name__)


class ChatGooglePalmError(Exception):
    pass


def _truncate_at_stop_tokens(
    text: str,
    stop: Optional[List[str]],
) -> str:
    """Truncates text at the earliest stop token found."""
    if stop is None:
        return text

    for stop_token in stop:
        stop_token_idx = text.find(stop_token)
        if stop_token_idx != -1:
            text = text[:stop_token_idx]
    return text


def _response_to_result(
    response: genai.types.ChatResponse,
    stop: Optional[List[str]],
) -> ChatResult:
    """Converts a PaLM API response into a LangChain ChatResult."""
    if not response.candidates:
        raise ChatGooglePalmError("ChatResponse must have at least one candidate.")

    generations: List[ChatGeneration] = []
    for candidate in response.candidates:
        author = candidate.get("author")
        if author is None:
            raise ChatGooglePalmError(f"ChatResponse must have an author: {candidate}")

        content = _truncate_at_stop_tokens(candidate.get("content", ""), stop)
        if content is None:
            raise ChatGooglePalmError(f"ChatResponse must have a content: {candidate}")

        if author == "ai":
            generations.append(
                ChatGeneration(text=content, message=AIMessage(content=content))
            )
        elif author == "human":
            generations.append(
                ChatGeneration(
                    text=content,
                    message=HumanMessage(content=content),
                )
            )
        else:
            generations.append(
                ChatGeneration(
                    text=content,
                    message=ChatMessage(role=author, content=content),
                )
            )

    return ChatResult(generations=generations)


def _messages_to_prompt_dict(
    input_messages: List[BaseMessage],
) -> genai.types.MessagePromptDict:
    """Converts a list of LangChain messages into a PaLM API MessagePrompt structure."""
    import google.generativeai as genai

    context: str = ""
    examples: List[genai.types.MessageDict] = []
    messages: List[genai.types.MessageDict] = []

    remaining = list(enumerate(input_messages))

    while remaining:
        index, input_message = remaining.pop(0)

        if isinstance(input_message, SystemMessage):
            if index != 0:
                raise ChatGooglePalmError("System message must be first input message.")
            context = input_message.content
        elif isinstance(input_message, HumanMessage) and input_message.example:
            if messages:
                raise ChatGooglePalmError(
                    "Message examples must come before other messages."
                )
            _, next_input_message = remaining.pop(0)
            if isinstance(next_input_message, AIMessage) and next_input_message.example:
                examples.extend(
                    [
                        genai.types.MessageDict(
                            author="human", content=input_message.content
                        ),
                        genai.types.MessageDict(
                            author="ai", content=next_input_message.content
                        ),
                    ]
                )
            else:
                raise ChatGooglePalmError(
                    "Human example message must be immediately followed by an "
                    " AI example response."
                )
        elif isinstance(input_message, AIMessage) and input_message.example:
            raise ChatGooglePalmError(
                "AI example message must be immediately preceded by a Human "
                "example message."
            )
        elif isinstance(input_message, AIMessage):
            messages.append(
                genai.types.MessageDict(author="ai", content=input_message.content)
            )
        elif isinstance(input_message, HumanMessage):
            messages.append(
                genai.types.MessageDict(author="human", content=input_message.content)
            )
        elif isinstance(input_message, ChatMessage):
            messages.append(
                genai.types.MessageDict(
                    author=input_message.role, content=input_message.content
                )
            )
        else:
            raise ChatGooglePalmError(
                "Messages without an explicit role not supported by PaLM API."
            )

    return genai.types.MessagePromptDict(
        context=context,
        examples=examples,
        messages=messages,
    )


def _create_retry_decorator() -> Callable[[Any], Any]:
    """Returns a tenacity retry decorator, preconfigured to handle PaLM exceptions"""
    import google.api_core.exceptions

    multiplier = 2
    min_seconds = 1
    max_seconds = 60
    max_retries = 10

    return retry(
        reraise=True,
        stop=stop_after_attempt(max_retries),
        wait=wait_exponential(multiplier=multiplier, min=min_seconds, max=max_seconds),
        retry=(
            retry_if_exception_type(google.api_core.exceptions.ResourceExhausted)
            | retry_if_exception_type(google.api_core.exceptions.ServiceUnavailable)
            | retry_if_exception_type(google.api_core.exceptions.GoogleAPIError)
        ),
        before_sleep=before_sleep_log(logger, logging.WARNING),
    )


def chat_with_retry(llm: ChatGooglePalm, **kwargs: Any) -> Any:
    """Use tenacity to retry the completion call."""
    retry_decorator = _create_retry_decorator()

    @retry_decorator
    def _chat_with_retry(**kwargs: Any) -> Any:
        return llm.client.chat(**kwargs)

    return _chat_with_retry(**kwargs)


async def achat_with_retry(llm: ChatGooglePalm, **kwargs: Any) -> Any:
    """Use tenacity to retry the async completion call."""
    retry_decorator = _create_retry_decorator()

    @retry_decorator
    async def _achat_with_retry(**kwargs: Any) -> Any:
        # Use OpenAI's async api https://github.com/openai/openai-python#async-api
        return await llm.client.chat_async(**kwargs)

    return await _achat_with_retry(**kwargs)


class ChatGooglePalm(BaseChatModel, BaseModel):
    """Wrapper around Google's PaLM Chat API.

    To use you must have the google.generativeai Python package installed and
    either:

        1. The ``GOOGLE_API_KEY``` environment varaible set with your API key, or
        2. Pass your API key using the google_api_key kwarg to the ChatGoogle
           constructor.

    Example:
        .. code-block:: python

            from langchain.chat_models import ChatGooglePalm
            chat = ChatGooglePalm()

    """

    client: Any  #: :meta private:
    model_name: str = "models/chat-bison-001"
    """Model name to use."""
    google_api_key: Optional[str] = None
    temperature: Optional[float] = None
    """Run inference with this temperature. Must by in the closed
       interval [0.0, 1.0]."""
    top_p: Optional[float] = None
    """Decode using nucleus sampling: consider the smallest set of tokens whose
       probability sum is at least top_p. Must be in the closed interval [0.0, 1.0]."""
    top_k: Optional[int] = None
    """Decode using top-k sampling: consider the set of top_k most probable tokens.
       Must be positive."""
    n: int = 1
    """Number of chat completions to generate for each prompt. Note that the API may
       not return the full n completions if duplicates are generated."""

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """Validate api key, python package exists, temperature, top_p, and top_k."""
        google_api_key = get_from_dict_or_env(
            values, "google_api_key", "GOOGLE_API_KEY"
        )
        try:
            import google.generativeai as genai

            genai.configure(api_key=google_api_key)
        except ImportError:
            raise ChatGooglePalmError(
                "Could not import google.generativeai python package. "
                "Please install it with `pip install google-generativeai`"
            )

        values["client"] = genai

        if values["temperature"] is not None and not 0 <= values["temperature"] <= 1:
            raise ValueError("temperature must be in the range [0.0, 1.0]")

        if values["top_p"] is not None and not 0 <= values["top_p"] <= 1:
            raise ValueError("top_p must be in the range [0.0, 1.0]")

        if values["top_k"] is not None and values["top_k"] <= 0:
            raise ValueError("top_k must be positive")

        return values

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
    ) -> ChatResult:
        prompt = _messages_to_prompt_dict(messages)

        response: genai.types.ChatResponse = chat_with_retry(
            self,
            model=self.model_name,
            prompt=prompt,
            temperature=self.temperature,
            top_p=self.top_p,
            top_k=self.top_k,
            candidate_count=self.n,
        )

        return _response_to_result(response, stop)

    async def _agenerate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
    ) -> ChatResult:
        prompt = _messages_to_prompt_dict(messages)

        response: genai.types.ChatResponse = await achat_with_retry(
            self,
            model=self.model_name,
            prompt=prompt,
            temperature=self.temperature,
            top_p=self.top_p,
            top_k=self.top_k,
            candidate_count=self.n,
        )

        return _response_to_result(response, stop)

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """Get the identifying parameters."""
        return {
            "model_name": self.model_name,
            "temperature": self.temperature,
            "top_p": self.top_p,
            "top_k": self.top_k,
            "n": self.n,
        }

    @property
    def _llm_type(self) -> str:
        return "google-palm-chat"

langchain.schema.ChatMessage langchain.utils.get_from_dict_or_env langchain.schema.ChatResult langchain.schema.HumanMessage langchain.schema.AIMessage
[((792, 819), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (809, 819), False, 'import logging\n'), ((2563, 2598), 'langchain.schema.ChatResult', 'ChatResult', ([], {'generations': 'generations'}), '(generations=generations)\n', (2573, 2598), False, 'from langchain.schema import AIMessage, BaseMessage, ChatGeneration, ChatMessage, ChatResult, HumanMessage, SystemMessage\n'), ((5264, 5353), 'google.generativeai.types.MessagePromptDict', 'genai.types.MessagePromptDict', ([], {'context': 'context', 'examples': 'examples', 'messages': 'messages'}), '(context=context, examples=examples, messages=\n    messages)\n', (5293, 5353), True, 'import google.generativeai as genai\n'), ((8257, 8273), 'pydantic.root_validator', 'root_validator', ([], {}), '()\n', (8271, 8273), False, 'from pydantic import BaseModel, root_validator\n'), ((8442, 8506), 'langchain.utils.get_from_dict_or_env', 'get_from_dict_or_env', (['values', '"""google_api_key"""', '"""GOOGLE_API_KEY"""'], {}), "(values, 'google_api_key', 'GOOGLE_API_KEY')\n", (8462, 8506), False, 'from langchain.utils import get_from_dict_or_env\n'), ((5697, 5728), 'tenacity.stop_after_attempt', 'stop_after_attempt', (['max_retries'], {}), '(max_retries)\n', (5715, 5728), False, 'from tenacity import before_sleep_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential\n'), ((5743, 5816), 'tenacity.wait_exponential', 'wait_exponential', ([], {'multiplier': 'multiplier', 'min': 'min_seconds', 'max': 'max_seconds'}), '(multiplier=multiplier, min=min_seconds, max=max_seconds)\n', (5759, 5816), False, 'from tenacity import before_sleep_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential\n'), ((6114, 6155), 'tenacity.before_sleep_log', 'before_sleep_log', (['logger', 'logging.WARNING'], {}), '(logger, logging.WARNING)\n', (6130, 6155), False, 'from tenacity import before_sleep_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential\n'), ((8603, 8642), 'google.generativeai.configure', 'genai.configure', ([], {'api_key': 'google_api_key'}), '(api_key=google_api_key)\n', (8618, 8642), True, 'import google.generativeai as genai\n'), ((6015, 6081), 'tenacity.retry_if_exception_type', 'retry_if_exception_type', (['google.api_core.exceptions.GoogleAPIError'], {}), '(google.api_core.exceptions.GoogleAPIError)\n', (6038, 6081), False, 'from tenacity import before_sleep_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential\n'), ((5846, 5915), 'tenacity.retry_if_exception_type', 'retry_if_exception_type', (['google.api_core.exceptions.ResourceExhausted'], {}), '(google.api_core.exceptions.ResourceExhausted)\n', (5869, 5915), False, 'from tenacity import before_sleep_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential\n'), ((5930, 6000), 'tenacity.retry_if_exception_type', 'retry_if_exception_type', (['google.api_core.exceptions.ServiceUnavailable'], {}), '(google.api_core.exceptions.ServiceUnavailable)\n', (5953, 6000), False, 'from tenacity import before_sleep_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential\n'), ((2073, 2099), 'langchain.schema.AIMessage', 'AIMessage', ([], {'content': 'content'}), '(content=content)\n', (2082, 2099), False, 'from langchain.schema import AIMessage, BaseMessage, ChatGeneration, ChatMessage, ChatResult, HumanMessage, SystemMessage\n'), ((2273, 2302), 'langchain.schema.HumanMessage', 'HumanMessage', ([], {'content': 'content'}), '(content=content)\n', (2285, 2302), False, 'from langchain.schema import AIMessage, BaseMessage, ChatGeneration, ChatMessage, ChatResult, HumanMessage, SystemMessage\n'), ((2476, 2517), 'langchain.schema.ChatMessage', 'ChatMessage', ([], {'role': 'author', 'content': 'content'}), '(role=author, content=content)\n', (2487, 2517), False, 'from langchain.schema import AIMessage, BaseMessage, ChatGeneration, ChatMessage, ChatResult, HumanMessage, SystemMessage\n'), ((3748, 3818), 'google.generativeai.types.MessageDict', 'genai.types.MessageDict', ([], {'author': '"""human"""', 'content': 'input_message.content'}), "(author='human', content=input_message.content)\n", (3771, 3818), True, 'import google.generativeai as genai\n'), ((3898, 3970), 'google.generativeai.types.MessageDict', 'genai.types.MessageDict', ([], {'author': '"""ai"""', 'content': 'next_input_message.content'}), "(author='ai', content=next_input_message.content)\n", (3921, 3970), True, 'import google.generativeai as genai\n'), ((4608, 4675), 'google.generativeai.types.MessageDict', 'genai.types.MessageDict', ([], {'author': '"""ai"""', 'content': 'input_message.content'}), "(author='ai', content=input_message.content)\n", (4631, 4675), True, 'import google.generativeai as genai\n'), ((4789, 4859), 'google.generativeai.types.MessageDict', 'genai.types.MessageDict', ([], {'author': '"""human"""', 'content': 'input_message.content'}), "(author='human', content=input_message.content)\n", (4812, 4859), True, 'import google.generativeai as genai\n'), ((4972, 5058), 'google.generativeai.types.MessageDict', 'genai.types.MessageDict', ([], {'author': 'input_message.role', 'content': 'input_message.content'}), '(author=input_message.role, content=input_message.\n    content)\n', (4995, 5058), True, 'import google.generativeai as genai\n')]