Merge branch 'InterwebAlchemy-feature/estimate-tokens'

README.md

@@ -247,7 +247,7 @@ enters debug mode. With 'false', it exits debug mode.
 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`: 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`.
+ • `%tokens [prompt]`: Calculate the tokens used by the current conversation's messages and estimate their cost, and optionally calculate the tokens and estimated cost of a `prompt` if one is provided. Relies on [LiteLLM's `cost_per_token()` method](https://docs.litellm.ai/docs/completion/token_usage#2-cost_per_token) for estimated cost.
  • `%help`: Show the help message.
 
 ### Configuration

interpreter/terminal_interface/magic_commands.py

@@ -41,7 +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`.",
+      "%tokens [prompt]": "Calculate the tokens used by the current conversation's messages and estimate their cost and optionally calculate the tokens and estimated cost of a `prompt` if one is provided.",
       "%help": "Show this help message.",
     }
 
@@ -102,15 +102,24 @@ 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):
+def handle_count_tokens(self, prompt):
     messages = [{"role": "system", "message": self.system_message}] + self.messages
 
+    outputs = []
+
     if len(self.messages) == 0:
       (tokens, cost) = count_messages_tokens(messages=messages, model=self.model)
-      display_markdown_message(f"> System Prompt Tokens: {tokens} (${cost})")
+      outputs.append((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})")
+      outputs.append(f"> Conversation Tokens: {tokens} (${cost})")
+
+    if prompt and prompt != '':
+      (tokens, cost) = count_messages_tokens(messages=[prompt], model=self.model)
+      outputs.append(f"> Prompt Tokens: {tokens} (${cost})") 
+
+    display_markdown_message("\n".join(outputs))
+
 
 def handle_magic_command(self, user_input):
     # split the command into the command and the arguments, by the first whitespace

tests/test_interpreter.py

@@ -3,6 +3,13 @@
 import time
 import pytest
 import interpreter
+from interpreter.utils.count_tokens import count_tokens, count_messages_tokens
+import time
+
+interpreter.auto_run = True
+interpreter.model = "gpt-4"
+interpreter.temperature = 0
+
 
 # this function will run before each test
 # we're clearing out the messages Array so we can start fresh and reduce token usage
@@ -59,6 +66,28 @@ def test_reset():
     assert interpreter.messages == []
 
 
+def test_token_counter():
+    system_tokens = count_tokens(text=interpreter.system_message, model=interpreter.model)
+
+    prompt = "How many tokens is this?"
+
+    prompt_tokens = count_tokens(text=prompt, model=interpreter.model)
+
+    messages = [{"role": "system", "message": interpreter.system_message}] + interpreter.messages
+
+    system_token_test = count_messages_tokens(messages=messages, model=interpreter.model)
+
+    system_tokens_ok = system_tokens == system_token_test[0]
+
+    messages.append({"role": "user", "message": prompt})
+
+    prompt_token_test = count_messages_tokens(messages=messages, model=interpreter.model)
+
+    prompt_tokens_ok = system_tokens + prompt_tokens == prompt_token_test[0]
+
+    assert system_tokens_ok and prompt_tokens_ok
+
+
 def test_hello_world():
     hello_world_response = "Hello, World!"