from typing import List, Optional, Any, Dict

from langchain.callbacks.manager import Callbacks
from langchain.chat_models import ChatAnthropic
from langchain.schema import BaseMessage, LLMResult
from pydantic import root_validator

from core.llm.wrappers.anthropic_wrapper import handle_anthropic_exceptions


class StreamableChatAnthropic(ChatAnthropic):
    """
    Wrapper around Anthropic's large language model.
    """

    @root_validator()
    def prepare_params(cls, values: Dict) -> Dict:
        values['model_name'] = values.get('model')
        values['max_tokens'] = values.get('max_tokens_to_sample')
        return values

    @handle_anthropic_exceptions
    def generate(
            self,
            messages: List[List[BaseMessage]],
            stop: Optional[List[str]] = None,
            callbacks: Callbacks = None,
            *,
            tags: Optional[List[str]] = None,
            metadata: Optional[Dict[str, Any]] = None,
            **kwargs: Any,
    ) -> LLMResult:
        return super().generate(messages, stop, callbacks, tags=tags, metadata=metadata, **kwargs)

    @classmethod
    def get_kwargs_from_model_params(cls, params: dict):
        params['model'] = params.get('model_name')
        del params['model_name']

        params['max_tokens_to_sample'] = params.get('max_tokens')
        del params['max_tokens']

        del params['frequency_penalty']
        del params['presence_penalty']

        return params
