Class MessageManage

utils/message_content.py:52–99

class MessageManage:
    def __init__(self, conversation_ids: list[str]):
        self.conversation_ids = conversation_ids

        self.conversation_map = {}
        for conversation_id in self.conversation_ids:
            self.conversation_map[conversation_id] = ConversationManager.query_by_idWithSort(conversation_id=conversation_id)

        self.memory_map = {}
        for conversation_id, conversations in self.conversation_map.items():
            memory = ConversationMemory()
            for conversation in conversations:
                memory.addQuestion(conversation.content)
                memory.addAnswer(conversation.content)
            self.memory_map[conversation_id] = memory

        self.llmChain_map = {}
        for conversation_id, conversations in self.conversation_map.items():
            template = """You are a chatbot having a conversation with a human.

            {chat_history}
            Human: {human_input}
            Chatbot:"""

            prompt = PromptTemplate(
                input_variables=["chat_history", "human_input"], template=template
            )
            memory = ConversationBufferMemory(memory_key="chat_history")
            llm = OpenAI(
                temperature=st.session_state["temperature"],
                openai_api_key=st.session_state["openai_api_key"])
            llm_chain = LLMChain(
                llm=llm,
                prompt=prompt,
                verbose=True,
                memory=memory,
            )
            self.llmChain_map[conversation_id] = llm_chain

    def add_conversation_id(self, conversation_id: str):
        conversations = ConversationManager.query_by_idWithSort(conversation_id=conversation_id)
        self.conversation_map[conversation_id] = conversations

    def add_conversation(self, conversation_id: str, role_name: str, content: str):
        conversation_id_res = ConversationManager.add(conversation_id=conversation_id, role_name=role_name, content=content)
        conversation_id_tmp = ConversationManager.query_latest_id(conversation_id=conversation_id)
        if conversation_id_res == conversation_id_tmp:
            self.conversation_map[conversation_id] = ConversationManager.query_by_idWithSort(conversation_id=conversation_id)