LLM/chat.py

class Chat:
 """
 Handles the chat using to avoid OOM issues.
 """

 def __init__(self, size):
 self.size = size
 self.init_chat_message = None
 # maxlen is necessary pair, since a each new step we add an prompt and assitant answer
 self.buffer = []

 def append(self, item):
 self.buffer.append(item)
 if len(self.buffer) == 2 * (self.size + 1):
 self.buffer.pop(0)
 self.buffer.pop(0)

 def init_chat(self, init_chat_message):
 self.init_chat_message = init_chat_message

 def to_list(self):
 if self.init_chat_message:
 return [self.init_chat_message] + self.buffer
 else:
 return self.buffer