From 370b3ef264812bd9aa72ee9ac1a322606a875243 Mon Sep 17 00:00:00 2001 From: Eric allen Date: Sun, 8 Oct 2023 15:33:48 -0400 Subject: [PATCH 1/5] feat: add %tokens magic command that counts tokens via tiktoken --- .../terminal_interface/magic_commands.py | 10 +++++++ interpreter/utils/count_tokens.py | 26 +++++++++++++++++++ 2 files changed, 36 insertions(+) create mode 100644 interpreter/utils/count_tokens.py diff --git a/interpreter/terminal_interface/magic_commands.py b/interpreter/terminal_interface/magic_commands.py index 48780a9a92..879abc9d62 100644 --- a/interpreter/terminal_interface/magic_commands.py +++ b/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 in the current session.", "%help": "Show this help message.", } @@ -100,6 +102,13 @@ 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): + if len(self.messages) == 0: + display_markdown_message(f"> System Prompt Tokens: {count_messages_tokens(messages=[self.system_message], model=self.model)}") + else: + messages_including_system = [self.system_message] + self.messages + display_markdown_message(f"> Total Tokens Used: {count_messages_tokens(messages=messages_including_system, model=self.model)}") + def handle_magic_command(self, user_input): # split the command into the command and the arguments, by the first whitespace switch = { @@ -109,6 +118,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 `%` diff --git a/interpreter/utils/count_tokens.py b/interpreter/utils/count_tokens.py new file mode 100644 index 0000000000..4121677b88 --- /dev/null +++ b/interpreter/utils/count_tokens.py @@ -0,0 +1,26 @@ +import tiktoken + +def count_tokens(text="", model="gpt-4"): + """ + Count the number of tokens in a list of tokens + """ + + encoder = tiktoken.encoding_for_model(model) + + return len(encoder.encode(text)) + + +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) + + return tokens_used \ No newline at end of file From a1895e7a21e3e1452099b5fd03ca6db3e2cc327f Mon Sep 17 00:00:00 2001 From: Eric allen Date: Sun, 8 Oct 2023 15:52:05 -0400 Subject: [PATCH 2/5] feat: add estimated cost from litellm to token counter --- interpreter/terminal_interface/magic_commands.py | 6 ++++-- interpreter/utils/count_tokens.py | 16 ++++++++++++++-- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/interpreter/terminal_interface/magic_commands.py b/interpreter/terminal_interface/magic_commands.py index 879abc9d62..6c66fe7988 100644 --- a/interpreter/terminal_interface/magic_commands.py +++ b/interpreter/terminal_interface/magic_commands.py @@ -104,10 +104,12 @@ def handle_load_message(self, json_path): def handle_count_tokens(self, arguments): if len(self.messages) == 0: - display_markdown_message(f"> System Prompt Tokens: {count_messages_tokens(messages=[self.system_message], model=self.model)}") + (tokens, cost) = count_messages_tokens(messages=[self.system_message], model=self.model) + display_markdown_message(f"> System Prompt Tokens: {tokens} (${cost})") else: messages_including_system = [self.system_message] + self.messages - display_markdown_message(f"> Total Tokens Used: {count_messages_tokens(messages=messages_including_system, model=self.model)}") + (tokens, cost) = count_messages_tokens(messages=messages_including_system, model=self.model) + display_markdown_message(f"> Total Tokens Used: {tokens} (${cost})") def handle_magic_command(self, user_input): # split the command into the command and the arguments, by the first whitespace diff --git a/interpreter/utils/count_tokens.py b/interpreter/utils/count_tokens.py index 4121677b88..0130d14031 100644 --- a/interpreter/utils/count_tokens.py +++ b/interpreter/utils/count_tokens.py @@ -1,14 +1,23 @@ import tiktoken +from litellm import cost_per_token def count_tokens(text="", model="gpt-4"): """ - Count the number of tokens in a list of tokens + 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 prompt_cost def count_messages_tokens(messages=[], model=None): """ @@ -23,4 +32,7 @@ def count_messages_tokens(messages=[], model=None): elif "message" in message: tokens_used += count_tokens(message["message"], model=model) - return tokens_used \ No newline at end of file + prompt_cost = token_cost(tokens_used, model=model) + + return (tokens_used, prompt_cost) + From ba72cb7b6cfe32aed15ff9b68c019cdd4cc30996 Mon Sep 17 00:00:00 2001 From: Eric allen Date: Sun, 8 Oct 2023 15:56:38 -0400 Subject: [PATCH 3/5] fix: add note about only including current messages --- interpreter/terminal_interface/magic_commands.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/interpreter/terminal_interface/magic_commands.py b/interpreter/terminal_interface/magic_commands.py index 6c66fe7988..8047b19743 100644 --- a/interpreter/terminal_interface/magic_commands.py +++ b/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 in the current session.", + "%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.", } @@ -109,7 +109,7 @@ def handle_count_tokens(self, arguments): else: messages_including_system = [self.system_message] + self.messages (tokens, cost) = count_messages_tokens(messages=messages_including_system, model=self.model) - display_markdown_message(f"> Total Tokens Used: {tokens} (${cost})") + display_markdown_message(f"> Tokens in Current Conversation: {tokens} (${cost})") def handle_magic_command(self, user_input): # split the command into the command and the arguments, by the first whitespace From 1115cc606828dc7767e794131656a0c015d8c696 Mon Sep 17 00:00:00 2001 From: Eric allen Date: Sun, 8 Oct 2023 16:05:42 -0400 Subject: [PATCH 4/5] chore: add %tokens to README --- README.md | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index d24edb4f06..0883349fe4 100644 --- a/README.md +++ b/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 From eef23fd3fa493b97cde35a5f26a56df9ceb60d8c Mon Sep 17 00:00:00 2001 From: Eric allen Date: Sun, 8 Oct 2023 19:49:13 -0400 Subject: [PATCH 5/5] fix: include generated code in token count; round to 6 decimals --- interpreter/terminal_interface/magic_commands.py | 9 +++++---- interpreter/utils/count_tokens.py | 8 +++++++- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/interpreter/terminal_interface/magic_commands.py b/interpreter/terminal_interface/magic_commands.py index 8047b19743..d884222fde 100644 --- a/interpreter/terminal_interface/magic_commands.py +++ b/interpreter/terminal_interface/magic_commands.py @@ -103,13 +103,14 @@ 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=[self.system_message], model=self.model) + (tokens, cost) = count_messages_tokens(messages=messages, model=self.model) display_markdown_message(f"> System Prompt Tokens: {tokens} (${cost})") else: - messages_including_system = [self.system_message] + self.messages - (tokens, cost) = count_messages_tokens(messages=messages_including_system, model=self.model) - display_markdown_message(f"> Tokens in Current Conversation: {tokens} (${cost})") + (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 diff --git a/interpreter/utils/count_tokens.py b/interpreter/utils/count_tokens.py index 0130d14031..bda66a325b 100644 --- a/interpreter/utils/count_tokens.py +++ b/interpreter/utils/count_tokens.py @@ -17,7 +17,7 @@ def token_cost(tokens=0, model="gpt-4"): (prompt_cost, _) = cost_per_token(model=model, prompt_tokens=tokens) - return prompt_cost + return round(prompt_cost, 6) def count_messages_tokens(messages=[], model=None): """ @@ -32,6 +32,12 @@ def count_messages_tokens(messages=[], model=None): 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)