import asyncio
import inspect
import warnings
from abc import ABC, abstractmethod
from functools import partial
from typing import Any, Dict, List, Mapping, Optional, Sequence

from pydantic import Field, root_validator

import langchain
from langchain.base_language import BaseLanguageModel
from langchain.callbacks.base import BaseCallbackManager
from langchain.callbacks.manager import (
    AsyncCallbackManager,
    AsyncCallbackManagerForLLMRun,
    CallbackManager,
    CallbackManagerForLLMRun,
    Callbacks,
)
from langchain.load.dump import dumpd, dumps
from langchain.schema import (
    ChatGeneration,
    ChatResult,
    LLMResult,
    PromptValue,
    RunInfo,
)
from langchain.schema.messages import AIMessage, BaseMessage, HumanMessage


def _get_verbosity() -> bool:
    return langchain.verbose


class BaseChatModel(BaseLanguageModel, ABC):
    cache: Optional[bool] = None
    verbose: bool = Field(default_factory=_get_verbosity)
    """Whether to print out response text."""
    callbacks: Callbacks = Field(default=None, exclude=True)
    callback_manager: Optional[BaseCallbackManager] = Field(default=None, exclude=True)
    tags: Optional[List[str]] = Field(default=None, exclude=True)
    """Tags to add to the run trace."""

    @root_validator()
    def raise_deprecation(cls, values: Dict) -> Dict:
        """Raise deprecation warning if callback_manager is used."""
        if values.get("callback_manager") is not None:
            warnings.warn(
                "callback_manager is deprecated. Please use callbacks instead.",
                DeprecationWarning,
            )
            values["callbacks"] = values.pop("callback_manager", None)
        return values

    class Config:
        """Configuration for this pydantic object."""

        arbitrary_types_allowed = True

    def _combine_llm_outputs(self, llm_outputs: List[Optional[dict]]) -> dict:
        return {}

    def _get_invocation_params(
        self,
        stop: Optional[List[str]] = None,
    ) -> dict:
        params = self.dict()
        params["stop"] = stop
        return params

    def _get_llm_string(self, stop: Optional[List[str]] = None, **kwargs: Any) -> str:
        if self.lc_serializable:
            params = {**kwargs, **{"stop": stop}}
            param_string = str(sorted([(k, v) for k, v in params.items()]))
            llm_string = dumps(self)
            return llm_string + "---" + param_string
        else:
            params = self._get_invocation_params(stop=stop)
            params = {**params, **kwargs}
            return str(sorted([(k, v) for k, v in params.items()]))

    def generate(
        self,
        messages: List[List[BaseMessage]],
        stop: Optional[List[str]] = None,
        callbacks: Callbacks = None,
        *,
        tags: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> LLMResult:
        """Top Level call"""
        params = self._get_invocation_params(stop=stop)
        options = {"stop": stop}

        callback_manager = CallbackManager.configure(
            callbacks,
            self.callbacks,
            self.verbose,
            tags,
            self.tags,
        )
        run_managers = callback_manager.on_chat_model_start(
            dumpd(self), messages, invocation_params=params, options=options
        )
        results = []
        for i, m in enumerate(messages):
            try:
                results.append(
                    self._generate_with_cache(
                        m,
                        stop=stop,
                        run_manager=run_managers[i] if run_managers else None,
                        **kwargs,
                    )
                )
            except (KeyboardInterrupt, Exception) as e:
                if run_managers:
                    run_managers[i].on_llm_error(e)
                raise e
        flattened_outputs = [
            LLMResult(generations=[res.generations], llm_output=res.llm_output)
            for res in results
        ]
        llm_output = self._combine_llm_outputs([res.llm_output for res in results])
        generations = [res.generations for res in results]
        output = LLMResult(generations=generations, llm_output=llm_output)
        if run_managers:
            run_infos = []
            for manager, flattened_output in zip(run_managers, flattened_outputs):
                manager.on_llm_end(flattened_output)
                run_infos.append(RunInfo(run_id=manager.run_id))
            output.run = run_infos
        return output

    async def agenerate(
        self,
        messages: List[List[BaseMessage]],
        stop: Optional[List[str]] = None,
        callbacks: Callbacks = None,
        *,
        tags: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> LLMResult:
        """Top Level call"""
        params = self._get_invocation_params(stop=stop)
        options = {"stop": stop}

        callback_manager = AsyncCallbackManager.configure(
            callbacks,
            self.callbacks,
            self.verbose,
            tags,
            self.tags,
        )

        run_managers = await callback_manager.on_chat_model_start(
            dumpd(self), messages, invocation_params=params, options=options
        )

        results = await asyncio.gather(
            *[
                self._agenerate_with_cache(
                    m,
                    stop=stop,
                    run_manager=run_managers[i] if run_managers else None,
                    **kwargs,
                )
                for i, m in enumerate(messages)
            ],
            return_exceptions=True,
        )
        exceptions = []
        for i, res in enumerate(results):
            if isinstance(res, Exception):
                if run_managers:
                    await run_managers[i].on_llm_error(res)
                exceptions.append(res)
        if exceptions:
            if run_managers:
                await asyncio.gather(
                    *[
                        run_manager.on_llm_end(
                            LLMResult(
                                generations=[res.generations], llm_output=res.llm_output
                            )
                        )
                        for run_manager, res in zip(run_managers, results)
                        if not isinstance(res, Exception)
                    ]
                )
            raise exceptions[0]
        flattened_outputs = [
            LLMResult(generations=[res.generations], llm_output=res.llm_output)
            for res in results
        ]
        llm_output = self._combine_llm_outputs([res.llm_output for res in results])
        generations = [res.generations for res in results]
        output = LLMResult(generations=generations, llm_output=llm_output)
        await asyncio.gather(
            *[
                run_manager.on_llm_end(flattened_output)
                for run_manager, flattened_output in zip(
                    run_managers, flattened_outputs
                )
            ]
        )
        if run_managers:
            output.run = [
                RunInfo(run_id=run_manager.run_id) for run_manager in run_managers
            ]
        return output

    def generate_prompt(
        self,
        prompts: List[PromptValue],
        stop: Optional[List[str]] = None,
        callbacks: Callbacks = None,
        **kwargs: Any,
    ) -> LLMResult:
        prompt_messages = [p.to_messages() for p in prompts]
        return self.generate(prompt_messages, stop=stop, callbacks=callbacks, **kwargs)

    async def agenerate_prompt(
        self,
        prompts: List[PromptValue],
        stop: Optional[List[str]] = None,
        callbacks: Callbacks = None,
        **kwargs: Any,
    ) -> LLMResult:
        prompt_messages = [p.to_messages() for p in prompts]
        return await self.agenerate(
            prompt_messages, stop=stop, callbacks=callbacks, **kwargs
        )

    def _generate_with_cache(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        new_arg_supported = inspect.signature(self._generate).parameters.get(
            "run_manager"
        )
        disregard_cache = self.cache is not None and not self.cache
        if langchain.llm_cache is None or disregard_cache:
            # This happens when langchain.cache is None, but self.cache is True
            if self.cache is not None and self.cache:
                raise ValueError(
                    "Asked to cache, but no cache found at `langchain.cache`."
                )
            if new_arg_supported:
                return self._generate(
                    messages, stop=stop, run_manager=run_manager, **kwargs
                )
            else:
                return self._generate(messages, stop=stop, **kwargs)
        else:
            llm_string = self._get_llm_string(stop=stop, **kwargs)
            prompt = dumps(messages)
            cache_val = langchain.llm_cache.lookup(prompt, llm_string)
            if isinstance(cache_val, list):
                return ChatResult(generations=cache_val)
            else:
                if new_arg_supported:
                    result = self._generate(
                        messages, stop=stop, run_manager=run_manager, **kwargs
                    )
                else:
                    result = self._generate(messages, stop=stop, **kwargs)
                langchain.llm_cache.update(prompt, llm_string, result.generations)
                return result

    async def _agenerate_with_cache(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        new_arg_supported = inspect.signature(self._agenerate).parameters.get(
            "run_manager"
        )
        disregard_cache = self.cache is not None and not self.cache
        if langchain.llm_cache is None or disregard_cache:
            # This happens when langchain.cache is None, but self.cache is True
            if self.cache is not None and self.cache:
                raise ValueError(
                    "Asked to cache, but no cache found at `langchain.cache`."
                )
            if new_arg_supported:
                return await self._agenerate(
                    messages, stop=stop, run_manager=run_manager, **kwargs
                )
            else:
                return await self._agenerate(messages, stop=stop, **kwargs)
        else:
            llm_string = self._get_llm_string(stop=stop, **kwargs)
            prompt = dumps(messages)
            cache_val = langchain.llm_cache.lookup(prompt, llm_string)
            if isinstance(cache_val, list):
                return ChatResult(generations=cache_val)
            else:
                if new_arg_supported:
                    result = await self._agenerate(
                        messages, stop=stop, run_manager=run_manager, **kwargs
                    )
                else:
                    result = await self._agenerate(messages, stop=stop, **kwargs)
                langchain.llm_cache.update(prompt, llm_string, result.generations)
                return result

    @abstractmethod
    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Top Level call"""

    @abstractmethod
    async def _agenerate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Top Level call"""

    def __call__(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        callbacks: Callbacks = None,
        **kwargs: Any,
    ) -> BaseMessage:
        generation = self.generate(
            [messages], stop=stop, callbacks=callbacks, **kwargs
        ).generations[0][0]
        if isinstance(generation, ChatGeneration):
            return generation.message
        else:
            raise ValueError("Unexpected generation type")

    async def _call_async(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        callbacks: Callbacks = None,
        **kwargs: Any,
    ) -> BaseMessage:
        result = await self.agenerate(
            [messages], stop=stop, callbacks=callbacks, **kwargs
        )
        generation = result.generations[0][0]
        if isinstance(generation, ChatGeneration):
            return generation.message
        else:
            raise ValueError("Unexpected generation type")

    def call_as_llm(
        self, message: str, stop: Optional[List[str]] = None, **kwargs: Any
    ) -> str:
        return self.predict(message, stop=stop, **kwargs)

    def predict(
        self, text: str, *, stop: Optional[Sequence[str]] = None, **kwargs: Any
    ) -> str:
        if stop is None:
            _stop = None
        else:
            _stop = list(stop)
        result = self([HumanMessage(content=text)], stop=_stop, **kwargs)
        return result.content

    def predict_messages(
        self,
        messages: List[BaseMessage],
        *,
        stop: Optional[Sequence[str]] = None,
        **kwargs: Any,
    ) -> BaseMessage:
        if stop is None:
            _stop = None
        else:
            _stop = list(stop)
        return self(messages, stop=_stop, **kwargs)

    async def apredict(
        self, text: str, *, stop: Optional[Sequence[str]] = None, **kwargs: Any
    ) -> str:
        if stop is None:
            _stop = None
        else:
            _stop = list(stop)
        result = await self._call_async(
            [HumanMessage(content=text)], stop=_stop, **kwargs
        )
        return result.content

    async def apredict_messages(
        self,
        messages: List[BaseMessage],
        *,
        stop: Optional[Sequence[str]] = None,
        **kwargs: Any,
    ) -> BaseMessage:
        if stop is None:
            _stop = None
        else:
            _stop = list(stop)
        return await self._call_async(messages, stop=_stop, **kwargs)

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

    @property
    @abstractmethod
    def _llm_type(self) -> str:
        """Return type of chat model."""

    def dict(self, **kwargs: Any) -> Dict:
        """Return a dictionary of the LLM."""
        starter_dict = dict(self._identifying_params)
        starter_dict["_type"] = self._llm_type
        return starter_dict


class SimpleChatModel(BaseChatModel):
    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        output_str = self._call(messages, stop=stop, run_manager=run_manager, **kwargs)
        message = AIMessage(content=output_str)
        generation = ChatGeneration(message=message)
        return ChatResult(generations=[generation])

    @abstractmethod
    def _call(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        """Simpler interface."""

    async def _agenerate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        func = partial(
            self._generate, messages, stop=stop, run_manager=run_manager, **kwargs
        )
        return await asyncio.get_event_loop().run_in_executor(None, func)

langchain.callbacks.manager.AsyncCallbackManager.configure langchain.schema.messages.AIMessage langchain.schema.ChatResult langchain.load.dump.dumps langchain.callbacks.manager.CallbackManager.configure langchain.load.dump.dumpd langchain.schema.RunInfo langchain.schema.messages.HumanMessage langchain.schema.ChatGeneration langchain.llm_cache.lookup langchain.llm_cache.update langchain.schema.LLMResult
[((915, 952), 'pydantic.Field', 'Field', ([], {'default_factory': '_get_verbosity'}), '(default_factory=_get_verbosity)\n', (920, 952), False, 'from pydantic import Field, root_validator\n'), ((1026, 1059), 'pydantic.Field', 'Field', ([], {'default': 'None', 'exclude': '(True)'}), '(default=None, exclude=True)\n', (1031, 1059), False, 'from pydantic import Field, root_validator\n'), ((1114, 1147), 'pydantic.Field', 'Field', ([], {'default': 'None', 'exclude': '(True)'}), '(default=None, exclude=True)\n', (1119, 1147), False, 'from pydantic import Field, root_validator\n'), ((1180, 1213), 'pydantic.Field', 'Field', ([], {'default': 'None', 'exclude': '(True)'}), '(default=None, exclude=True)\n', (1185, 1213), False, 'from pydantic import Field, root_validator\n'), ((1260, 1276), 'pydantic.root_validator', 'root_validator', ([], {}), '()\n', (1274, 1276), False, 'from pydantic import Field, root_validator\n'), ((3020, 3107), 'langchain.callbacks.manager.CallbackManager.configure', 'CallbackManager.configure', (['callbacks', 'self.callbacks', 'self.verbose', 'tags', 'self.tags'], {}), '(callbacks, self.callbacks, self.verbose, tags,\n    self.tags)\n', (3045, 3107), False, 'from langchain.callbacks.manager import AsyncCallbackManager, AsyncCallbackManagerForLLMRun, CallbackManager, CallbackManagerForLLMRun, Callbacks\n'), ((4172, 4229), 'langchain.schema.LLMResult', 'LLMResult', ([], {'generations': 'generations', 'llm_output': 'llm_output'}), '(generations=generations, llm_output=llm_output)\n', (4181, 4229), False, 'from langchain.schema import ChatGeneration, ChatResult, LLMResult, PromptValue, RunInfo\n'), ((4944, 5036), 'langchain.callbacks.manager.AsyncCallbackManager.configure', 'AsyncCallbackManager.configure', (['callbacks', 'self.callbacks', 'self.verbose', 'tags', 'self.tags'], {}), '(callbacks, self.callbacks, self.verbose,\n    tags, self.tags)\n', (4974, 5036), False, 'from langchain.callbacks.manager import AsyncCallbackManager, AsyncCallbackManagerForLLMRun, CallbackManager, CallbackManagerForLLMRun, Callbacks\n'), ((6747, 6804), 'langchain.schema.LLMResult', 'LLMResult', ([], {'generations': 'generations', 'llm_output': 'llm_output'}), '(generations=generations, llm_output=llm_output)\n', (6756, 6804), False, 'from langchain.schema import ChatGeneration, ChatResult, LLMResult, PromptValue, RunInfo\n'), ((15295, 15324), 'langchain.schema.messages.AIMessage', 'AIMessage', ([], {'content': 'output_str'}), '(content=output_str)\n', (15304, 15324), False, 'from langchain.schema.messages import AIMessage, BaseMessage, HumanMessage\n'), ((15346, 15377), 'langchain.schema.ChatGeneration', 'ChatGeneration', ([], {'message': 'message'}), '(message=message)\n', (15360, 15377), False, 'from langchain.schema import ChatGeneration, ChatResult, LLMResult, PromptValue, RunInfo\n'), ((15393, 15429), 'langchain.schema.ChatResult', 'ChatResult', ([], {'generations': '[generation]'}), '(generations=[generation])\n', (15403, 15429), False, 'from langchain.schema import ChatGeneration, ChatResult, LLMResult, PromptValue, RunInfo\n'), ((15941, 16020), 'functools.partial', 'partial', (['self._generate', 'messages'], {'stop': 'stop', 'run_manager': 'run_manager'}), '(self._generate, messages, stop=stop, run_manager=run_manager, **kwargs)\n', (15948, 16020), False, 'from functools import partial\n'), ((1467, 1569), 'warnings.warn', 'warnings.warn', (['"""callback_manager is deprecated. Please use callbacks instead."""', 'DeprecationWarning'], {}), "('callback_manager is deprecated. Please use callbacks instead.',\n    DeprecationWarning)\n", (1480, 1569), False, 'import warnings\n'), ((2374, 2385), 'langchain.load.dump.dumps', 'dumps', (['self'], {}), '(self)\n', (2379, 2385), False, 'from langchain.load.dump import dumpd, dumps\n'), ((3248, 3259), 'langchain.load.dump.dumpd', 'dumpd', (['self'], {}), '(self)\n', (3253, 3259), False, 'from langchain.load.dump import dumpd, dumps\n'), ((3903, 3970), 'langchain.schema.LLMResult', 'LLMResult', ([], {'generations': '[res.generations]', 'llm_output': 'res.llm_output'}), '(generations=[res.generations], llm_output=res.llm_output)\n', (3912, 3970), False, 'from langchain.schema import ChatGeneration, ChatResult, LLMResult, PromptValue, RunInfo\n'), ((6478, 6545), 'langchain.schema.LLMResult', 'LLMResult', ([], {'generations': '[res.generations]', 'llm_output': 'res.llm_output'}), '(generations=[res.generations], llm_output=res.llm_output)\n', (6487, 6545), False, 'from langchain.schema import ChatGeneration, ChatResult, LLMResult, PromptValue, RunInfo\n'), ((9053, 9068), 'langchain.load.dump.dumps', 'dumps', (['messages'], {}), '(messages)\n', (9058, 9068), False, 'from langchain.load.dump import dumpd, dumps\n'), ((9093, 9139), 'langchain.llm_cache.lookup', 'langchain.llm_cache.lookup', (['prompt', 'llm_string'], {}), '(prompt, llm_string)\n', (9119, 9139), False, 'import langchain\n'), ((10773, 10788), 'langchain.load.dump.dumps', 'dumps', (['messages'], {}), '(messages)\n', (10778, 10788), False, 'from langchain.load.dump import dumpd, dumps\n'), ((10813, 10859), 'langchain.llm_cache.lookup', 'langchain.llm_cache.lookup', (['prompt', 'llm_string'], {}), '(prompt, llm_string)\n', (10839, 10859), False, 'import langchain\n'), ((5184, 5195), 'langchain.load.dump.dumpd', 'dumpd', (['self'], {}), '(self)\n', (5189, 5195), False, 'from langchain.load.dump import dumpd, dumps\n'), ((7127, 7161), 'langchain.schema.RunInfo', 'RunInfo', ([], {'run_id': 'run_manager.run_id'}), '(run_id=run_manager.run_id)\n', (7134, 7161), False, 'from langchain.schema import ChatGeneration, ChatResult, LLMResult, PromptValue, RunInfo\n'), ((9207, 9240), 'langchain.schema.ChatResult', 'ChatResult', ([], {'generations': 'cache_val'}), '(generations=cache_val)\n', (9217, 9240), False, 'from langchain.schema import ChatGeneration, ChatResult, LLMResult, PromptValue, RunInfo\n'), ((9556, 9622), 'langchain.llm_cache.update', 'langchain.llm_cache.update', (['prompt', 'llm_string', 'result.generations'], {}), '(prompt, llm_string, result.generations)\n', (9582, 9622), False, 'import langchain\n'), ((10927, 10960), 'langchain.schema.ChatResult', 'ChatResult', ([], {'generations': 'cache_val'}), '(generations=cache_val)\n', (10937, 10960), False, 'from langchain.schema import ChatGeneration, ChatResult, LLMResult, PromptValue, RunInfo\n'), ((11290, 11356), 'langchain.llm_cache.update', 'langchain.llm_cache.update', (['prompt', 'llm_string', 'result.generations'], {}), '(prompt, llm_string, result.generations)\n', (11316, 11356), False, 'import langchain\n'), ((13349, 13375), 'langchain.schema.messages.HumanMessage', 'HumanMessage', ([], {'content': 'text'}), '(content=text)\n', (13361, 13375), False, 'from langchain.schema.messages import AIMessage, BaseMessage, HumanMessage\n'), ((4451, 4481), 'langchain.schema.RunInfo', 'RunInfo', ([], {'run_id': 'manager.run_id'}), '(run_id=manager.run_id)\n', (4458, 4481), False, 'from langchain.schema import ChatGeneration, ChatResult, LLMResult, PromptValue, RunInfo\n'), ((8220, 8253), 'inspect.signature', 'inspect.signature', (['self._generate'], {}), '(self._generate)\n', (8237, 8253), False, 'import inspect\n'), ((9925, 9959), 'inspect.signature', 'inspect.signature', (['self._agenerate'], {}), '(self._agenerate)\n', (9942, 9959), False, 'import inspect\n'), ((14025, 14051), 'langchain.schema.messages.HumanMessage', 'HumanMessage', ([], {'content': 'text'}), '(content=text)\n', (14037, 14051), False, 'from langchain.schema.messages import AIMessage, BaseMessage, HumanMessage\n'), ((16064, 16088), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (16086, 16088), False, 'import asyncio\n'), ((6075, 6142), 'langchain.schema.LLMResult', 'LLMResult', ([], {'generations': '[res.generations]', 'llm_output': 'res.llm_output'}), '(generations=[res.generations], llm_output=res.llm_output)\n', (6084, 6142), False, 'from langchain.schema import ChatGeneration, ChatResult, LLMResult, PromptValue, RunInfo\n')]