Merge pull request #607 from InterwebAlchemy/feature/token-count

README.md

@@ -240,13 +240,14 @@ In the interactive mode, you can use the below commands to enhance your experien
 
 **Available Commands:**  
  • `%debug [true/false]`: Toggle debug mode. Without arguments or with 'true', it
-enters debug mode. With 'false', it exits debug mode.  
- • `%reset`: Resets the current session.  
- • `%undo`: Remove previous messages and its response from the message history.  
+enters debug mode. With 'false', it exits debug mode.
+ • `%reset`: Resets the current session.
+ • `%undo`: Remove previous messages and its response from the message history.
  • `%save_message [path]`: Saves messages to a specified JSON path. If no path is
-provided, it defaults to 'messages.json'.  
+provided, it defaults to 'messages.json'.
  • `%load_message [path]`: Loads messages from a specified JSON path. If no path  
- is provided, it defaults to 'messages.json'.  
+ is provided, it defaults to 'messages.json'.
+ • `%tokens`: Counts the number of tokens used by the current conversation's messages and displays a cost estimate via [LiteLLM's `cost_per_token()` method](https://docs.litellm.ai/docs/completion/token_usage#2-cost_per_token). **Note**: This only accounts for messages currently in the conversation, not messages that have been sent or received and then removed via `%undo`.
  • `%help`: Show the help message.
 
 ### Configuration

interpreter/terminal_interface/magic_commands.py

@@ -1,4 +1,5 @@
 from ..utils.display_markdown_message import display_markdown_message
+from ..utils.count_tokens import count_messages_tokens
 import json
 import os
 
@@ -40,6 +41,7 @@ def handle_help(self, arguments):
       "%undo": "Remove previous messages and its response from the message history.",
       "%save_message [path]": "Saves messages to a specified JSON path. If no path is provided, it defaults to 'messages.json'.",
       "%load_message [path]": "Loads messages from a specified JSON path. If no path is provided, it defaults to 'messages.json'.",
+      "%tokens": "Show the tokens used by the current conversation's messages. **Note**: this will not take into account tokens that have already been used and then removed from the conversation with `%undo`.",
       "%help": "Show this help message.",
     }
 
@@ -100,6 +102,16 @@ def handle_load_message(self, json_path):
 
     display_markdown_message(f"> messages json loaded from {os.path.abspath(json_path)}")
 
+def handle_count_tokens(self, arguments):
+    messages = [{"role": "system", "message": self.system_message}] + self.messages
+
+    if len(self.messages) == 0:
+      (tokens, cost) = count_messages_tokens(messages=messages, model=self.model)
+      display_markdown_message(f"> System Prompt Tokens: {tokens} (${cost})")
+    else:
+      (tokens, cost) = count_messages_tokens(messages=messages, model=self.model)
+      display_markdown_message(f"> Conversation Tokens: {tokens} (${cost})")
+
 def handle_magic_command(self, user_input):
     # split the command into the command and the arguments, by the first whitespace
     switch = {
@@ -109,6 +121,7 @@ def handle_magic_command(self, user_input):
       "save_message": handle_save_message,
       "load_message": handle_load_message,
       "undo": handle_undo,
+      "tokens": handle_count_tokens,
     }
 
     user_input = user_input[1:].strip()  # Capture the part after the `%`

interpreter/utils/count_tokens.py

@@ -0,0 +1,44 @@
+import tiktoken
+from litellm import cost_per_token
+
+def count_tokens(text="", model="gpt-4"):
+    """
+    Count the number of tokens in a string
+    """
+
+    encoder = tiktoken.encoding_for_model(model)
+
+    return len(encoder.encode(text))
+
+def token_cost(tokens=0, model="gpt-4"):
+    """
+    Calculate the cost of the current number of tokens
+    """
+
+    (prompt_cost, _) = cost_per_token(model=model, prompt_tokens=tokens)
+
+    return round(prompt_cost, 6)
+
+def count_messages_tokens(messages=[], model=None):
+    """
+    Count the number of tokens in a list of messages
+    """
+
+    tokens_used = 0
+
+    for message in messages:
+        if isinstance(message, str):
+            tokens_used += count_tokens(message, model=model)
+        elif "message" in message:
+            tokens_used += count_tokens(message["message"], model=model)
+
+            if "code" in message:
+                tokens_used += count_tokens(message["code"], model=model)
+
+            if "output" in message:
+                tokens_used += count_tokens(message["output"], model=model)
+
+    prompt_cost = token_cost(tokens_used, model=model)
+
+    return (tokens_used, prompt_cost)
+