Implemented logging token usage (solves #322) (#438)

* Implemented logging token usage

Token usage is now tracked and logged into memory/logs/token_usage

* Step names are now inferred from function name

* Incorporated Anton's feedback

- Made LogUsage a dataclass
- For token logging, step name is now inferred via inspect module

* Formatted (black/ruff)

* Update gpt_engineer/ai.py

Co-authored-by: Anton Osika <anton.osika@gmail.com>

* formatting

---------

Co-authored-by: Anton Osika <anton.osika@gmail.com>

gpt_engineer/ai.py

@@ -2,25 +2,54 @@ from __future__ import annotations
 
 import logging
 
+from dataclasses import dataclass
 from typing import Dict, List
 
 import openai
+import tiktoken
 
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class TokenUsage:
+    step_name: str
+    in_step_prompt_tokens: int
+    in_step_completion_tokens: int
+    in_step_total_tokens: int
+    total_prompt_tokens: int
+    total_completion_tokens: int
+    total_tokens: int
+
+
 class AI:
     def __init__(self, model="gpt-4", temperature=0.1):
         self.temperature = temperature
         self.model = model
 
-    def start(self, system, user):
+        # initialize token usage log
+        self.cumulative_prompt_tokens = 0
+        self.cumulative_completion_tokens = 0
+        self.cumulative_total_tokens = 0
+        self.token_usage_log = []
+
+        try:
+            self.tokenizer = tiktoken.encoding_for_model(model)
+        except KeyError:
+            logger.debug(
+                f"Tiktoken encoder for model {model} not found. Using "
+                "cl100k_base encoder instead. The results may therefore be "
+                "inaccurate and should only be used as estimate."
+            )
+            self.tokenizer = tiktoken.get_encoding("cl100k_base")
+
+    def start(self, system, user, step_name):
         messages = [
             {"role": "system", "content": system},
             {"role": "user", "content": user},
         ]
 
-        return self.next(messages)
+        return self.next(messages, step_name=step_name)
 
     def fsystem(self, msg):
         return {"role": "system", "content": msg}
@@ -31,7 +60,7 @@ class AI:
     def fassistant(self, msg):
         return {"role": "assistant", "content": msg}
 
-    def next(self, messages: List[Dict[str, str]], prompt=None):
+    def next(self, messages: List[Dict[str, str]], prompt=None, *, step_name=None):
         if prompt:
             messages += [{"role": "user", "content": prompt}]
 
@@ -52,8 +81,65 @@ class AI:
         print()
         messages += [{"role": "assistant", "content": "".join(chat)}]
         logger.debug(f"Chat completion finished: {messages}")
+
+        self.update_token_usage_log(
+            messages=messages, answer="".join(chat), step_name=step_name
+        )
+
         return messages
 
+    def update_token_usage_log(self, messages, answer, step_name):
+        prompt_tokens = self.num_tokens_from_messages(messages)
+        completion_tokens = self.num_tokens(answer)
+        total_tokens = prompt_tokens + completion_tokens
+
+        self.cumulative_prompt_tokens += prompt_tokens
+        self.cumulative_completion_tokens += completion_tokens
+        self.cumulative_total_tokens += total_tokens
+
+        self.token_usage_log.append(
+            TokenUsage(
+                step_name=step_name,
+                in_step_prompt_tokens=prompt_tokens,
+                in_step_completion_tokens=completion_tokens,
+                in_step_total_tokens=total_tokens,
+                total_prompt_tokens=self.cumulative_prompt_tokens,
+                total_completion_tokens=self.cumulative_completion_tokens,
+                total_tokens=self.cumulative_total_tokens,
+            )
+        )
+
+    def format_token_usage_log(self):
+        result = "step_name,"
+        result += "prompt_tokens_in_step,completion_tokens_in_step,total_tokens_in_step"
+        result += ",total_prompt_tokens,total_completion_tokens,total_tokens\n"
+        for log in self.token_usage_log:
+            result += log.step_name + ","
+            result += str(log.in_step_prompt_tokens) + ","
+            result += str(log.in_step_completion_tokens) + ","
+            result += str(log.in_step_total_tokens) + ","
+            result += str(log.total_prompt_tokens) + ","
+            result += str(log.total_completion_tokens) + ","
+            result += str(log.total_tokens) + "\n"
+        return result
+
+    def num_tokens(self, txt):
+        return len(self.tokenizer.encode(txt))
+
+    def num_tokens_from_messages(self, messages):
+        """Returns the number of tokens used by a list of messages."""
+        n_tokens = 0
+        for message in messages:
+            n_tokens += (
+                4  # every message follows <im_start>{role/name}\n{content}<im_end>\n
+            )
+            for key, value in message.items():
+                n_tokens += self.num_tokens(value)
+                if key == "name":  # if there's a name, the role is omitted
+                    n_tokens += -1  # role is always required and always 1 token
+        n_tokens += 2  # every reply is primed with <im_start>assistant
+        return n_tokens
+
 
 def fallback_model(model: str) -> str:
     try: