from sparkdesk.chat_utils import limit_chat_history

from ..chat import BaseChat
from .completion import ChatCompletion


class Chat(BaseChat):
    """
    Class for creating chat
    """

    completion_cls = ChatCompletion
    chat_completion: ChatCompletion

    async def ask(self, question: str):
        """
        Appends the user's question to the chat context,
        generates responses using the chat_completion model,
        and adds the assistant's response to the chat context.

        Args:
            question (str): The user's question.

        Yields:
            dict: A dictionary representing each response generated by
                the chat_completion model.
            reference: https://www.xfyun.cn/doc/spark/Web.html
        """

        self.chat_history.append({"role": "user", "content": question})
        self.chat_history = limit_chat_history(self.chat_history)

        response_messages = []
        responses = self.chat_completion.create(self.chat_history)
        try:
            async for response in responses:
                response_messages.extend(
                    [e["content"] for e in response["payload"]["choices"]["text"]]
                )
                yield response
        finally:  # GeneratorExit
            self.chat_history.append(
                {"role": "assistant", "content": "".join(response_messages)}
            )
