limit history to ~3000 tokens

.. and some small fixes for the last update

Author: andreashappe

history.py

@@ -1,3 +1,11 @@
+import tiktoken
+
+def num_tokens_from_string(string: str) -> int:
+    """Returns the number of tokens in a text string."""
+    encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
+    return len(encoding.encode(string))
+
+
 class ResultHistory:
     def __init__(self):
         self.data = []
@@ -8,6 +16,30 @@ def append(self, cmd, result):
             "result": result
         })
 
-    def dump(self):
+    def get_full_history(self):
         return self.data
 
+    # only retrieve recent parts. We need this as prompts only allow
+    # for maximum token length. We currently do this in a quite stupid
+    # whay which could be optimized in the future
+    def get_history(self, limit=3072):
+        result = []
+        rest = limit
+
+        for itm in reversed(self.data):
+            size_cmd = num_tokens_from_string(itm["cmd"])
+            size_result = num_tokens_from_string(itm["result"])
+            size = size_cmd + size_result
+
+            if size <= rest:
+                result.append(itm)
+                rest -= size
+            else:
+                # if theres a bit space left, fill that up with parts of the last item
+                if (rest - size_cmd) >= 200:
+                    result.append({
+                        "cmd" : itm["cmd"],
+                        "result" : itm["result"][:(rest-size_cmd-2)] + ".."
+                    })
+                return list(reversed(result))
+        return list(reversed(result))

prompt_helper.py

@@ -7,20 +7,20 @@
 from llms.openai import get_openai_response
 
 log = logging.getLogger()
-filename = datetime.now().strftime('logs/run_%H_%M_%d_%m_%Y.log')
+filename = datetime.now().strftime('logs/run_%Y%m%d%m-%H%M.log')
 log.addHandler(logging.FileHandler(filename))
 
-def output_log(self, kind, msg):
+def output_log(kind, msg):
     print("[" + Fore.RED + kind + Style.RESET_ALL +"]: " + msg)
-    self.log.warning("[" + kind + "] " + msg)
+    log.warning("[" + kind + "] " + msg)
 
 # helper for generating and executing LLM prompts from a template
 def create_and_ask_prompt(template_file, log_prefix, **params):
     global logs
 
     template = Template(filename='templates/' + template_file)
     prompt = template.render(**params)
-    logs.warning(log_prefix + "-prompt", prompt)
+    output_log(log_prefix + "-prompt", prompt)
     result = get_openai_response(prompt)
-    logs.warning(log_prefix + "-answer", result)
+    output_log(log_prefix + "-answer", result)
     return result

requirements.txt

@@ -20,7 +20,9 @@ paramiko==3.1.0
 pycparser==2.21
 PyNaCl==1.5.0
 python-dotenv==1.0.0
+regex==2023.3.23
 requests==2.28.2
+tiktoken==0.3.3
 tqdm==4.65.0
 urllib3==1.26.15
 yarl==1.9.2

wintermute.py

@@ -21,7 +21,7 @@
 
 while True:
 
-    next_cmd = create_and_ask_prompt('gpt_query.txt', "next-cmd", user=initial_user, history=cmd_history.dump())
+    next_cmd = create_and_ask_prompt('gpt_query.txt', "next-cmd", user=initial_user, history=cmd_history.get_history())
 
     # disable this for now, it's tragic because the AI won't tell me why it had chosen something
     # create_and_ask_prompt("why.txt", "why", user=initial_user, history=cmd_history.dump(), next_cmd=next_cmd)