From 78e7b5c3dbc94e4034820cff17df6a24767d2c1e Mon Sep 17 00:00:00 2001 From: Paul Gauthier Date: Sat, 22 Jul 2023 08:45:40 -0300 Subject: [PATCH] wip --- aider/coders/base_coder.py | 2 ++ aider/history.py | 70 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+) create mode 100644 aider/history.py diff --git a/aider/coders/base_coder.py b/aider/coders/base_coder.py index e54fb1087..9e245bbed 100755 --- a/aider/coders/base_coder.py +++ b/aider/coders/base_coder.py @@ -353,6 +353,8 @@ class Coder: def move_back_cur_messages(self, message): self.done_messages += self.cur_messages + self.done_messages = summarize_chat_history(self.done_messages) + if message: self.done_messages += [ dict(role="user", content=message), diff --git a/aider/history.py b/aider/history.py new file mode 100644 index 000000000..f38e98e86 --- /dev/null +++ b/aider/history.py @@ -0,0 +1,70 @@ +import json + +import tiktoken +from prompt_toolkit.completion import Completion + +from aider import prompts + +from .dump import dump # noqa: F401 + + +class ChatSummary: + def __init__(self, model, max_tokens=1024): + self.tokenizer = tiktoken.encoding_for_model(model) + self.max_tokens = max_tokens + + def summarize(self, messages): + num = len(messages) + if num < 2: + return messages + + total = 0 + sized = [] + for msg in messages: + tokens = len(self.tokenizer.encode(json.dumps(msg))) + sized.append((tokens, msg)) + total += tokens + + if total <= self.max_tokens: + return messages + + num = num // 2 + + # we want the head to end with an assistant msg + if messages[num]["role"] == "assistant": + num += 1 + + head = messages[:num] + tail = messages[num:] + + summary = self.summarize_all(head) + + tail_tokens = sum(tokens for tokens, msg in sized[num:]) + summary_tokens = len(self.tokenizer.encode(json.dumps(summary))) + + result = summary + tail + if summary_tokens + tail_tokens < self.max_tokens: + return result + + return self.summarize(result) + + def summarize_all(self, messages): + content = "" + for msg in messages: + role = msg["role"].upper() + if role not in ("USER", "ASSISTANT"): + continue + content += f"# {role}\n" + content += msg["content"] + if not content.endswith("\n"): + content += "\n" + + dump(content) + + messages = [ + dict(role="system", content=prompts.summarize), + dict(role="user", content=content), + ] + + summary = simple_send_with_retries(model=models.GPT35.name, messages=messages) + dump(summary)