memory.py

class ChatMemory:
    def __init__(self, system_prompt="你是一個專業的助手，會用繁體中文回答問題。", max_len: int = 1000):
        self.max_len = max_len
        self.history = []
        self.initialize(system_prompt)
        
    def initialize(self, system_message: str):
        self.history = [
            {'role': 'system', 'content': system_message}
        ]
        self.check_system_message_length(system_message)
        
    def check_system_message_length(self, system_message: str): 
        if len(system_message) > self.max_len:
            raise ValueError("The system prompt is too long. Please provide a shorter system prompt.")

    def add_user_message(self, message: str):
        self.history.append({'role': 'user', 'content': message})
        self._trim_history()

    def add_system_response(self, message: str):
        self.history.append({'role': 'system', 'content': message})
        self._trim_history()
    
    def _trim_history(self):
        if len(self.history) <= 3:  
            return
        while self._get_total_length() > self.max_len:
            self._remove_oldest_pair()
       
    def _remove_oldest_pair(self):
        for i in range(1, len(self.history) - 1):  
            if self.history[i]['role'] == 'user' and self.history[i + 1]['role'] == 'system':
                del self.history[i:i + 2]  
                break
                
    def _get_total_length(self):
        return sum(len(msg["content"]) for msg in self.history)
    
    def get_history(self):
        return self.history
    
    
if __name__ == '__main__':
    memory = ChatMemory(max_len=70)

    memory.initialize("你是一個專業的助手")
    memory.add_system_response("你好！我是 AI 助手，請問有什麼可以幫助你的？")
    memory.add_user_message("請問今天的天氣如何？")
    memory.add_system_response("今天的天氣很好，適合外出！")
    memory.add_user_message("你喜歡下雨嗎？")
    memory.add_system_response("我沒有感覺，但下雨對植物很重要！")

    print(memory.get_history())