import os
import json

from langchain.schema import messages_from_dict, messages_to_dict
from langchain.memory import (
    ConversationBufferMemory,
    ChatMessageHistory,
)


class YeagerAIContext:
    """Context for the @yeager.ai agent."""

    def __init__(self, username: str, session_id: str, session_path: str):
        self.username = username
        self.session_id = session_id
        self.session_path = session_path

        self.session_message_history = ChatMessageHistory()
        self.chat_buffer_memory = ConversationBufferMemory(
            memory_key="chat_history", input_key="input"
        )

    def load_session_message_history(self):
        try:
            with open(os.path.join(self.session_path, "session_history.txt"), "r") as f:
                dicts = json.loads(f.read())
                self.session_message_history.messages = messages_from_dict(dicts)
        except FileNotFoundError:
            os.makedirs(self.session_path, exist_ok=True)
            with open(os.path.join(self.session_path, "session_history.txt"), "w") as f:
                f.close()

    def save_session_message_history(self):
        dicts = messages_to_dict(self.session_message_history.messages)
        with open(os.path.join(self.session_path, "session_history.txt"), "w") as f:
            f.write(json.dumps(dicts))
            f.close()

    def create_shadow_clones(self):
        self.load_session_message_history()
        self.chat_buffer_memory.chat_memory = self.session_message_history

    def dispell_shadow_clones(self):
        self.session_message_history = self.chat_buffer_memory.chat_memory
        self.save_session_message_history()

langchain.schema.messages_from_dict langchain.memory.ConversationBufferMemory langchain.memory.ChatMessageHistory langchain.schema.messages_to_dict
[((472, 492), 'langchain.memory.ChatMessageHistory', 'ChatMessageHistory', ([], {}), '()\n', (490, 492), False, 'from langchain.memory import ConversationBufferMemory, ChatMessageHistory\n'), ((527, 597), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""', 'input_key': '"""input"""'}), "(memory_key='chat_history', input_key='input')\n", (551, 597), False, 'from langchain.memory import ConversationBufferMemory, ChatMessageHistory\n'), ((1162, 1217), 'langchain.schema.messages_to_dict', 'messages_to_dict', (['self.session_message_history.messages'], {}), '(self.session_message_history.messages)\n', (1178, 1217), False, 'from langchain.schema import messages_from_dict, messages_to_dict\n'), ((868, 893), 'langchain.schema.messages_from_dict', 'messages_from_dict', (['dicts'], {}), '(dicts)\n', (886, 893), False, 'from langchain.schema import messages_from_dict, messages_to_dict\n'), ((940, 985), 'os.makedirs', 'os.makedirs', (['self.session_path'], {'exist_ok': '(True)'}), '(self.session_path, exist_ok=True)\n', (951, 985), False, 'import os\n'), ((1236, 1290), 'os.path.join', 'os.path.join', (['self.session_path', '"""session_history.txt"""'], {}), "(self.session_path, 'session_history.txt')\n", (1248, 1290), False, 'import os\n'), ((1323, 1340), 'json.dumps', 'json.dumps', (['dicts'], {}), '(dicts)\n', (1333, 1340), False, 'import json\n'), ((700, 754), 'os.path.join', 'os.path.join', (['self.session_path', '"""session_history.txt"""'], {}), "(self.session_path, 'session_history.txt')\n", (712, 754), False, 'import os\n'), ((1008, 1062), 'os.path.join', 'os.path.join', (['self.session_path', '"""session_history.txt"""'], {}), "(self.session_path, 'session_history.txt')\n", (1020, 1062), False, 'import os\n')]