from typing import AsyncGenerator, Optional, Tuple
from langchain import ConversationChain
import logging

from typing import Optional, Tuple
from pydantic.v1 import SecretStr
from vocode.streaming.agent.base_agent import RespondAgent

from vocode.streaming.agent.utils import get_sentence_from_buffer

from langchain import ConversationChain
from langchain.schema import ChatMessage, AIMessage, HumanMessage
from langchain_community.chat_models import ChatAnthropic
import logging
from vocode import getenv

from vocode.streaming.models.agent import ChatAnthropicAgentConfig


from langchain.prompts import (
    ChatPromptTemplate,
    MessagesPlaceholder,
    HumanMessagePromptTemplate,
)

from vocode import getenv
from vocode.streaming.models.agent import ChatAnthropicAgentConfig
from langchain.memory import ConversationBufferMemory

SENTENCE_ENDINGS = [".", "!", "?"]


class ChatAnthropicAgent(RespondAgent[ChatAnthropicAgentConfig]):
    def __init__(
        self,
        agent_config: ChatAnthropicAgentConfig,
        logger: Optional[logging.Logger] = None,
        anthropic_api_key: Optional[SecretStr] = None,
    ):
        super().__init__(agent_config=agent_config, logger=logger)
        import anthropic

        # Convert anthropic_api_key to SecretStr if it's not None and not already a SecretStr
        if anthropic_api_key is not None and not isinstance(
            anthropic_api_key, SecretStr
        ):
            anthropic_api_key = SecretStr(anthropic_api_key)
        else:
            # Retrieve anthropic_api_key from environment variable and convert to SecretStr
            env_key = getenv("ANTHROPIC_API_KEY")
            if env_key:
                anthropic_api_key = SecretStr(env_key)

        if not anthropic_api_key:
            raise ValueError(
                "ANTHROPIC_API_KEY must be set in environment or passed in as a SecretStr"
            )
        self.prompt = ChatPromptTemplate.from_messages(
            [
                MessagesPlaceholder(variable_name="history"),
                HumanMessagePromptTemplate.from_template("{input}"),
            ]
        )

        self.llm = ChatAnthropic(
            model_name=agent_config.model_name,
            anthropic_api_key=anthropic_api_key,
        )

        # streaming not well supported by langchain, so we will connect directly
        self.anthropic_client = (
            anthropic.AsyncAnthropic(api_key=str(anthropic_api_key))
            if agent_config.generate_responses
            else None
        )

        self.memory = ConversationBufferMemory(return_messages=True)
        self.memory.chat_memory.messages.append(
            HumanMessage(content=self.agent_config.prompt_preamble)
        )
        if agent_config.initial_message:
            self.memory.chat_memory.messages.append(
                AIMessage(content=agent_config.initial_message.text)
            )

        self.conversation = ConversationChain(
            memory=self.memory, prompt=self.prompt, llm=self.llm
        )

    async def respond(
        self,
        human_input,
        conversation_id: str,
        is_interrupt: bool = False,
    ) -> Tuple[str, bool]:
        text = await self.conversation.apredict(input=human_input)
        self.logger.debug(f"LLM response: {text}")
        return text, False

    async def generate_response(
        self,
        human_input,
        conversation_id: str,
        is_interrupt: bool = False,
    ) -> AsyncGenerator[Tuple[str, bool], None]:
        self.memory.chat_memory.messages.append(HumanMessage(content=human_input))

        bot_memory_message = AIMessage(content="")
        self.memory.chat_memory.messages.append(bot_memory_message)
        prompt = self.llm._convert_messages_to_prompt(self.memory.chat_memory.messages)

        if self.anthropic_client:
            streamed_response = await self.anthropic_client.completions.create(
                prompt=prompt,
                max_tokens_to_sample=self.agent_config.max_tokens_to_sample,
                model=self.agent_config.model_name,
                stream=True,
            )

            buffer = ""
            async for completion in streamed_response:
                buffer += completion.completion
                sentence, remainder = get_sentence_from_buffer(buffer)
                if sentence:
                    bot_memory_message.content = bot_memory_message.content + sentence
                    buffer = remainder
                    yield sentence, True
                continue

    def update_last_bot_message_on_cut_off(self, message: str):
        for memory_message in self.memory.chat_memory.messages[::-1]:
            if (
                isinstance(memory_message, ChatMessage)
                and memory_message.role == "assistant"
            ) or isinstance(memory_message, AIMessage):
                memory_message.content = message
                return

langchain_community.chat_models.ChatAnthropic langchain.prompts.HumanMessagePromptTemplate.from_template langchain.memory.ConversationBufferMemory langchain.prompts.MessagesPlaceholder langchain.schema.HumanMessage langchain.schema.AIMessage langchain.ConversationChain
[((2147, 2238), 'langchain_community.chat_models.ChatAnthropic', 'ChatAnthropic', ([], {'model_name': 'agent_config.model_name', 'anthropic_api_key': 'anthropic_api_key'}), '(model_name=agent_config.model_name, anthropic_api_key=\n    anthropic_api_key)\n', (2160, 2238), False, 'from langchain_community.chat_models import ChatAnthropic\n'), ((2556, 2602), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'return_messages': '(True)'}), '(return_messages=True)\n', (2580, 2602), False, 'from langchain.memory import ConversationBufferMemory\n'), ((2936, 3007), 'langchain.ConversationChain', 'ConversationChain', ([], {'memory': 'self.memory', 'prompt': 'self.prompt', 'llm': 'self.llm'}), '(memory=self.memory, prompt=self.prompt, llm=self.llm)\n', (2953, 3007), False, 'from langchain import ConversationChain\n'), ((3624, 3645), 'langchain.schema.AIMessage', 'AIMessage', ([], {'content': '""""""'}), "(content='')\n", (3633, 3645), False, 'from langchain.schema import ChatMessage, AIMessage, HumanMessage\n'), ((1468, 1496), 'pydantic.v1.SecretStr', 'SecretStr', (['anthropic_api_key'], {}), '(anthropic_api_key)\n', (1477, 1496), False, 'from pydantic.v1 import SecretStr\n'), ((1625, 1652), 'vocode.getenv', 'getenv', (['"""ANTHROPIC_API_KEY"""'], {}), "('ANTHROPIC_API_KEY')\n", (1631, 1652), False, 'from vocode import getenv\n'), ((2664, 2719), 'langchain.schema.HumanMessage', 'HumanMessage', ([], {'content': 'self.agent_config.prompt_preamble'}), '(content=self.agent_config.prompt_preamble)\n', (2676, 2719), False, 'from langchain.schema import ChatMessage, AIMessage, HumanMessage\n'), ((3559, 3592), 'langchain.schema.HumanMessage', 'HumanMessage', ([], {'content': 'human_input'}), '(content=human_input)\n', (3571, 3592), False, 'from langchain.schema import ChatMessage, AIMessage, HumanMessage\n'), ((1713, 1731), 'pydantic.v1.SecretStr', 'SecretStr', (['env_key'], {}), '(env_key)\n', (1722, 1731), False, 'from pydantic.v1 import SecretStr\n'), ((1988, 2032), 'langchain.prompts.MessagesPlaceholder', 'MessagesPlaceholder', ([], {'variable_name': '"""history"""'}), "(variable_name='history')\n", (2007, 2032), False, 'from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder, HumanMessagePromptTemplate\n'), ((2050, 2101), 'langchain.prompts.HumanMessagePromptTemplate.from_template', 'HumanMessagePromptTemplate.from_template', (['"""{input}"""'], {}), "('{input}')\n", (2090, 2101), False, 'from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder, HumanMessagePromptTemplate\n'), ((2840, 2892), 'langchain.schema.AIMessage', 'AIMessage', ([], {'content': 'agent_config.initial_message.text'}), '(content=agent_config.initial_message.text)\n', (2849, 2892), False, 'from langchain.schema import ChatMessage, AIMessage, HumanMessage\n'), ((4286, 4318), 'vocode.streaming.agent.utils.get_sentence_from_buffer', 'get_sentence_from_buffer', (['buffer'], {}), '(buffer)\n', (4310, 4318), False, 'from vocode.streaming.agent.utils import get_sentence_from_buffer\n')]