Improved summarization split logic

This commit is contained in:
Paul Gauthier 2023-08-10 12:20:51 -03:00
parent ea87a350c2
commit 7eab7b54b8

View file

@ -26,14 +26,15 @@ class ChatSummary:
return sized
def summarize(self, messages, depth=0):
if len(messages) <= 4 or depth > 5:
return self.summarize_all(messages)
sized = self.tokenize(messages)
total = sum(tokens for tokens, _msg in sized)
if total <= self.max_tokens:
if total <= self.max_tokens and depth == 0:
return messages
min_split = 4
if len(messages) <= min_split or depth > 3:
return self.summarize_all(messages)
tail_tokens = 0
split_index = len(messages)
half_max_tokens = self.max_tokens // 2
@ -51,6 +52,9 @@ class ChatSummary:
while messages[split_index - 1]["role"] != "assistant" and split_index > 1:
split_index -= 1
if split_index <= min_split:
return self.summarize_all(messages)
head = messages[:split_index]
tail = messages[split_index:]