From cb177dbd7a34bc5ef5a454158c716ee46adc960a Mon Sep 17 00:00:00 2001 From: Carlos Freund Date: Fri, 2 May 2025 08:34:37 +0200 Subject: [PATCH] Fix and rewrite of token_counter (#10409) * added tests messages_with_counts: Made tolerance explicit for each test. But they match the new implementation(which beats the old) * new token counter impl * compare old and new implementation in test * delete old token counter * moved tests to /tests/litellm/litellm_core_utils * use existing types * docstrings * warn about using default params on unknown model. * created type for the token_counter_function * check key == "content" * throw error on invalid detail-type, ignore type-warning. * fix imports --- litellm/litellm_core_utils/token_counter.py | 370 +++++++++++++++++- litellm/utils.py | 306 +-------------- .../messages_with_counts.py | 18 +- .../litellm_core_utils}/test_token_counter.py | 144 +++++-- .../test_token_counter_tool.py | 84 ++++ .../test_token_counter_tool_data.py | 166 ++++++++ tests/litellm_utils_tests/test_utils.py | 36 -- 7 files changed, 743 insertions(+), 381 deletions(-) rename tests/{local_testing => litellm/litellm_core_utils}/messages_with_counts.py (98%) rename tests/{local_testing => litellm/litellm_core_utils}/test_token_counter.py (82%) create mode 100644 tests/litellm/litellm_core_utils/test_token_counter_tool.py create mode 100644 tests/litellm/litellm_core_utils/test_token_counter_tool_data.py diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index afd5ab5ff4..10ee7e4210 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -3,7 +3,9 @@ import base64 import io import struct -from typing import Literal, Optional, Tuple, Union +from typing import Callable, List, Literal, Optional, Tuple, Union + +import tiktoken import litellm from litellm import verbose_logger @@ -16,13 +18,21 @@ from litellm.constants import ( MAX_TILE_HEIGHT, MAX_TILE_WIDTH, ) +from litellm.litellm_core_utils.default_encoding import encoding as default_encoding from litellm.llms.custom_httpx.http_handler import _get_httpx_client +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionNamedToolChoiceParam, + ChatCompletionToolParam, + OpenAIMessageContent, +) +from litellm.types.utils import SelectTokenizerResponse def get_modified_max_tokens( model: str, base_model: str, - messages: Optional[list], + messages: Optional[List[AllMessageValues]], user_max_tokens: Optional[int], buffer_perc: Optional[float], buffer_num: Optional[float], @@ -281,3 +291,359 @@ def calculate_img_tokens( tile_tokens = (base_tokens * 2) * tiles_needed_high_res total_tokens = base_tokens + tile_tokens return total_tokens + + +TokenCounterFunction = Callable[[str], int] +""" +Type for a function that counts tokens in a string. +""" + +class _MessageCountParams: + """ + A class to hold the parameters for counting tokens in messages. + """ + def __init__( + self, + model: str, + custom_tokenizer: Optional[Union[dict, SelectTokenizerResponse]], + ): + from litellm.utils import print_verbose + actual_model = _fix_model_name(model) + if actual_model == "gpt-3.5-turbo-0301": + self.tokens_per_message = ( + 4 # every message follows <|start|>{role/name}\n{content}<|end|>\n + ) + self.tokens_per_name = -1 # if there's a name, the role is omitted + elif actual_model in litellm.open_ai_chat_completion_models: + self.tokens_per_message = 3 + self.tokens_per_name = 1 + elif actual_model in litellm.azure_llms: + self.tokens_per_message = 3 + self.tokens_per_name = 1 + else: + print_verbose(f"Warning: unknown model {model}. Using default token params.") + self.tokens_per_message = 3 + self.tokens_per_name = 1 + self.count_function = _get_count_function(model, custom_tokenizer) + + +def token_counter( + model="", + custom_tokenizer: Optional[Union[dict, SelectTokenizerResponse]] = None, + text: Optional[Union[str, List[str]]] = None, + messages: Optional[List[AllMessageValues]] = None, + count_response_tokens: Optional[bool] = False, + tools: Optional[List[ChatCompletionToolParam]] = None, + tool_choice: Optional[ChatCompletionNamedToolChoiceParam] = None, + use_default_image_token_count: Optional[bool] = False, + default_token_count: Optional[int] = None, +) -> int: + """ + Count the number of tokens in a given text using a specified model. + + Args: + model (str): The name of the model to use for tokenization. Default is an empty string. + custom_tokenizer (Optional[dict]): A custom tokenizer created with the `create_pretrained_tokenizer` or `create_tokenizer` method. Must be a dictionary with a string value for `type` and Tokenizer for `tokenizer`. Default is None. + text (str): The raw text string to be passed to the model. Default is None. + messages (Optional[List[AllMessageValues]]): Alternative to passing in text. A list of dictionaries representing messages with "role" and "content" keys. Default is None. + count_response_tokens (Optional[bool]): set to True to indicate we are processing a stream response. + tools (Optional[List[ChatCompletionToolParam]]): The available tools. Default is None. + tool_choice (Optional[ChatCompletionNamedToolChoiceParam]): The tool choice. Default is None. + use_default_image_token_count (Optional[bool]): When True, will NOT make a GET request to the image URL and instead return the default image dimensions. Default is False. + default_token_count (Optional[int]): The default number of tokens to return for a message block, if an error occurs. Default is None. + + Returns: + int: The number of tokens in the text. + """ + if text is not None and messages is not None: + raise ValueError("text and messages cannot both be set") + if use_default_image_token_count is None: + use_default_image_token_count = False + + if text is not None: + if tools or tool_choice: + raise ValueError("tools or tool_choice cannot be set if using text") + if isinstance(text, List): + text_to_count = "".join(t for t in text if isinstance(t, str)) + elif isinstance(text, str): + text_to_count = text + count_function = _get_count_function(model, custom_tokenizer) + num_tokens = count_function(text_to_count) + + elif messages is not None: + params = _MessageCountParams(model, custom_tokenizer) + num_tokens = _count_messages( + params, messages, use_default_image_token_count, default_token_count + ) + if count_response_tokens is False: + includes_system_message = any( + [message.get("role", None) == "system" for message in messages] + ) + num_tokens += _count_extra( + params.count_function, tools, tool_choice, includes_system_message + ) + + else: + raise ValueError("Either text or messages must be provided") + + return num_tokens + + +def _count_messages( + params: _MessageCountParams, + messages: List[AllMessageValues], + use_default_image_token_count: bool, + default_token_count: Optional[int], +) -> int: + """ + Count the number of tokens in a list of messages. + + Args: + params (_MessageCountParams): The parameters for counting tokens. + messages (List[AllMessageValues]): The list of messages to count tokens in. + use_default_image_token_count (bool): When True, will NOT make a GET request to the image URL and instead return the default image dimensions. + default_token_count (Optional[int]): The default number of tokens to return for a message block, if an error occurs. + """ + num_tokens = 0 + for message in messages: + num_tokens += params.tokens_per_message + for key, value in message.items(): + if value is None: + pass + elif key == "tool_calls": + if isinstance(value, List): + for tool_call in value: + if "function" in tool_call: + function_arguments = tool_call["function"].get( + "arguments", [] + ) + num_tokens += params.count_function(str(function_arguments)) + else: + raise ValueError( + f"Unsupported tool call {tool_call} must contain a function key" + ) + else: + raise ValueError( + f"Unsupported type {type(value)} for key tool_calls in message {message}" + ) + elif isinstance(value, str): + num_tokens += params.count_function(value) + if key == "name": + num_tokens += params.tokens_per_name + elif key == 'content' and isinstance(value, List): + num_tokens += _count_content_list( + params.count_function, + value, + use_default_image_token_count, + default_token_count, + ) + else: + raise ValueError( + f"Unsupported type {type(value)} for key {key} in message {message}" + ) + return num_tokens + + +def _count_extra( + count_function: TokenCounterFunction, + tools: Optional[List[ChatCompletionToolParam]], + tool_choice: Optional[ChatCompletionNamedToolChoiceParam], + includes_system_message: bool, +) -> int: + """Count extra tokens for function definitions and tool choices. + Args: + count_function (TokenCounterFunction): The function to count tokens. + tools (Optional[List[ChatCompletionToolParam]]): The available tools. + tool_choice (Optional[ChatCompletionNamedToolChoiceParam]): The tool choice. + includes_system_message (bool): Whether the messages include a system message. + """ + + num_tokens = 3 # every reply is primed with <|start|>assistant<|message|> + + if tools: + num_tokens += count_function(_format_function_definitions(tools)) + num_tokens += 9 # Additional tokens for function definition of tools + # If there's a system message and tools are present, subtract four tokens + if tools and includes_system_message: + num_tokens -= 4 + # If tool_choice is 'none', add one token. + # If it's an object, add 4 + the number of tokens in the function name. + # If it's undefined or 'auto', don't add anything. + if tool_choice == "none": + num_tokens += 1 + elif isinstance(tool_choice, dict): + num_tokens += 7 + num_tokens += count_function(str(tool_choice["function"]["name"])) + + return num_tokens + + +def _get_count_function( + model: Optional[str], + custom_tokenizer: Optional[Union[dict, SelectTokenizerResponse]] = None, +) -> TokenCounterFunction: + """ + Get the function to count tokens based on the model and custom tokenizer.""" + from litellm.utils import _select_tokenizer, print_verbose + + if model is not None or custom_tokenizer is not None: + tokenizer_json = custom_tokenizer or _select_tokenizer(model) # type: ignore + if tokenizer_json["type"] == "huggingface_tokenizer": + + def count_tokens(text: str) -> int: + enc = tokenizer_json["tokenizer"].encode(text) + return len(enc.ids) + + elif tokenizer_json["type"] == "openai_tokenizer": + model_to_use = _fix_model_name(model) # type: ignore + try: + if "gpt-4o" in model_to_use: + encoding = tiktoken.get_encoding("o200k_base") + else: + encoding = tiktoken.encoding_for_model(model_to_use) + except KeyError: + print_verbose("Warning: model not found. Using cl100k_base encoding.") + encoding = tiktoken.get_encoding("cl100k_base") + + def count_tokens(text: str) -> int: + return len(encoding.encode(text)) + + else: + raise ValueError("Unsupported tokenizer type") + else: + + def count_tokens(text: str) -> int: + return len(default_encoding.encode(text, disallowed_special=())) + + return count_tokens + + +def _fix_model_name(model: str) -> str: + """We normalize some model names to others""" + if model in litellm.azure_llms: + # azure llms use gpt-35-turbo instead of gpt-3.5-turbo πŸ™ƒ + return model.replace("-35", "-3.5") + elif model in litellm.open_ai_chat_completion_models: + return model # type: ignore + else: + return "gpt-3.5-turbo" + + +def _count_content_list( + count_function: TokenCounterFunction, + content_list: OpenAIMessageContent, + use_default_image_token_count: bool, + default_token_count: Optional[int], +) -> int: + """ + Get the number of tokens from a list of content. + """ + try: + num_tokens = 0 + for c in content_list: + if isinstance(c, str): + num_tokens += count_function(c) + elif c["type"] == "text": + num_tokens += count_function(c["text"]) + elif c["type"] == "image_url": + if isinstance(c["image_url"], dict): + image_url_dict = c["image_url"] + detail = image_url_dict.get("detail", "auto") + if detail not in ["low", "high", "auto"]: + raise ValueError( + f"Invalid detail value: {detail}. Expected 'low', 'high', or 'auto'." + ) + url = image_url_dict.get("url") + num_tokens += calculate_img_tokens( + data=url, + mode=detail, # type: ignore + use_default_image_token_count=use_default_image_token_count, + ) + elif isinstance(c["image_url"], str): + image_url_str = c["image_url"] + num_tokens += calculate_img_tokens( + data=image_url_str, + mode="auto", + use_default_image_token_count=use_default_image_token_count, + ) + else: + raise ValueError( + f"Invalid image_url type: {type(c['image_url'])}. Expected str or dict." + ) + else: + raise ValueError( + f"Invalid content type: {type(c)}. Expected str or dict." + ) + return num_tokens + except Exception as e: + if default_token_count is not None: + return default_token_count + raise ValueError( + f"Error getting number of tokens from content list: {e}, default_token_count={default_token_count}" + ) + + +def _format_function_definitions(tools): + """Formats tool definitions in the format that OpenAI appears to use. + Based on https://github.com/forestwanglin/openai-java/blob/main/jtokkit/src/main/java/xyz/felh/openai/jtokkit/utils/TikTokenUtils.java + """ + lines = [] + lines.append("namespace functions {") + lines.append("") + for tool in tools: + function = tool.get("function") + if function_description := function.get("description"): + lines.append(f"// {function_description}") + function_name = function.get("name") + parameters = function.get("parameters", {}) + properties = parameters.get("properties") + if properties and properties.keys(): + lines.append(f"type {function_name} = (_: {{") + lines.append(_format_object_parameters(parameters, 0)) + lines.append("}) => any;") + else: + lines.append(f"type {function_name} = () => any;") + lines.append("") + lines.append("} // namespace functions") + return "\n".join(lines) + + +def _format_object_parameters(parameters, indent): + properties = parameters.get("properties") + if not properties: + return "" + required_params = parameters.get("required", []) + lines = [] + for key, props in properties.items(): + description = props.get("description") + if description: + lines.append(f"// {description}") + question = "?" + if required_params and key in required_params: + question = "" + lines.append(f"{key}{question}: {_format_type(props, indent)},") + return "\n".join([" " * max(0, indent) + line for line in lines]) + + +def _format_type(props, indent): + type = props.get("type") + if type == "string": + if "enum" in props: + return " | ".join([f'"{item}"' for item in props["enum"]]) + return "string" + elif type == "array": + # items is required, OpenAI throws an error if it's missing + return f"{_format_type(props['items'], indent)}[]" + elif type == "object": + return f"{{\n{_format_object_parameters(props, indent + 2)}\n}}" + elif type in ["integer", "number"]: + if "enum" in props: + return " | ".join([f'"{item}"' for item in props["enum"]]) + return "number" + elif type == "boolean": + return "boolean" + elif type == "null": + return "null" + else: + # This is a guess, as an empty string doesn't yield the expected token count + return "any" diff --git a/litellm/utils.py b/litellm/utils.py index 87012d5eb5..82e46d2e5e 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -122,7 +122,6 @@ from litellm.litellm_core_utils.redact_messages import ( from litellm.litellm_core_utils.rules import Rules from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.litellm_core_utils.token_counter import ( - calculate_img_tokens, get_modified_max_tokens, ) from litellm.llms.bedrock.common_utils import BedrockModelInfo @@ -266,6 +265,8 @@ from .types.llms.openai import ( ChatCompletionToolCallFunctionChunk, ) from .types.router import LiteLLM_Params +from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new + ####### ENVIRONMENT VARIABLES #################### # Adjust to your specific application needs / system capabilities. @@ -759,7 +760,7 @@ def function_setup( # noqa: PLR0915 messages = "default-message-value" stream = True if "stream" in kwargs and kwargs["stream"] is True else False logging_obj = LiteLLMLogging( - model=model, + model=model, # type: ignore messages=messages, stream=stream, litellm_call_id=kwargs["litellm_call_id"], @@ -1607,98 +1608,6 @@ def decode(model="", tokens: List[int] = [], custom_tokenizer: Optional[dict] = dec = tokenizer_json["tokenizer"].decode(tokens) return dec - -def openai_token_counter( # noqa: PLR0915 - messages: Optional[list] = None, - model="gpt-3.5-turbo-0613", - text: Optional[str] = None, - is_tool_call: Optional[bool] = False, - tools: Optional[List[ChatCompletionToolParam]] = None, - tool_choice: Optional[ChatCompletionNamedToolChoiceParam] = None, - count_response_tokens: Optional[ - bool - ] = False, # Flag passed from litellm.stream_chunk_builder, to indicate counting tokens for LLM Response. We need this because for LLM input we add +3 tokens per message - based on OpenAI's token counter - use_default_image_token_count: Optional[bool] = False, - default_token_count: Optional[int] = None, -): - """ - Return the number of tokens used by a list of messages. - - Borrowed from https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb. - """ - print_verbose(f"LiteLLM: Utils - Counting tokens for OpenAI model={model}") - try: - if "gpt-4o" in model: - encoding = tiktoken.get_encoding("o200k_base") - else: - encoding = tiktoken.encoding_for_model(model) - except KeyError: - print_verbose("Warning: model not found. Using cl100k_base encoding.") - encoding = tiktoken.get_encoding("cl100k_base") - if model == "gpt-3.5-turbo-0301": - tokens_per_message = ( - 4 # every message follows <|start|>{role/name}\n{content}<|end|>\n - ) - tokens_per_name = -1 # if there's a name, the role is omitted - elif model in litellm.open_ai_chat_completion_models: - tokens_per_message = 3 - tokens_per_name = 1 - elif model in litellm.azure_llms: - tokens_per_message = 3 - tokens_per_name = 1 - else: - raise NotImplementedError( - f"""num_tokens_from_messages() is not implemented for model {model}. See https://github.com/openai/openai-python/blob/main/chatml.md for information on how messages are converted to tokens.""" - ) - num_tokens = 0 - includes_system_message = False - - if is_tool_call and text is not None: - # if it's a tool call we assembled 'text' in token_counter() - num_tokens = len(encoding.encode(text, disallowed_special=())) - elif messages is not None: - for message in messages: - num_tokens += tokens_per_message - if message.get("role", None) == "system": - includes_system_message = True - for key, value in message.items(): - if isinstance(value, str): - num_tokens += len(encoding.encode(value, disallowed_special=())) - if key == "name": - num_tokens += tokens_per_name - elif isinstance(value, List): - text, num_tokens_from_list = _get_num_tokens_from_content_list( - content_list=value, - use_default_image_token_count=use_default_image_token_count, - default_token_count=default_token_count, - ) - num_tokens += num_tokens_from_list - elif text is not None and count_response_tokens is True: - # This is the case where we need to count tokens for a streamed response. We should NOT add +3 tokens per message in this branch - num_tokens = len(encoding.encode(text, disallowed_special=())) - return num_tokens - elif text is not None: - num_tokens = len(encoding.encode(text, disallowed_special=())) - num_tokens += 3 # every reply is primed with <|start|>assistant<|message|> - - if tools: - num_tokens += len(encoding.encode(_format_function_definitions(tools))) - num_tokens += 9 # Additional tokens for function definition of tools - # If there's a system message and tools are present, subtract four tokens - if tools and includes_system_message: - num_tokens -= 4 - # If tool_choice is 'none', add one token. - # If it's an object, add 4 + the number of tokens in the function name. - # If it's undefined or 'auto', don't add anything. - if tool_choice == "none": - num_tokens += 1 - elif isinstance(tool_choice, dict): - num_tokens += 7 - num_tokens += len(encoding.encode(tool_choice["function"]["name"])) - - return num_tokens - - def create_pretrained_tokenizer( identifier: str, revision="main", auth_token: Optional[str] = None ): @@ -1741,118 +1650,6 @@ def create_tokenizer(json: str): return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} -def _format_function_definitions(tools): - """Formats tool definitions in the format that OpenAI appears to use. - Based on https://github.com/forestwanglin/openai-java/blob/main/jtokkit/src/main/java/xyz/felh/openai/jtokkit/utils/TikTokenUtils.java - """ - lines = [] - lines.append("namespace functions {") - lines.append("") - for tool in tools: - function = tool.get("function") - if function_description := function.get("description"): - lines.append(f"// {function_description}") - function_name = function.get("name") - parameters = function.get("parameters", {}) - properties = parameters.get("properties") - if properties and properties.keys(): - lines.append(f"type {function_name} = (_: {{") - lines.append(_format_object_parameters(parameters, 0)) - lines.append("}) => any;") - else: - lines.append(f"type {function_name} = () => any;") - lines.append("") - lines.append("} // namespace functions") - return "\n".join(lines) - - -def _format_object_parameters(parameters, indent): - properties = parameters.get("properties") - if not properties: - return "" - required_params = parameters.get("required", []) - lines = [] - for key, props in properties.items(): - description = props.get("description") - if description: - lines.append(f"// {description}") - question = "?" - if required_params and key in required_params: - question = "" - lines.append(f"{key}{question}: {_format_type(props, indent)},") - return "\n".join([" " * max(0, indent) + line for line in lines]) - - -def _format_type(props, indent): - type = props.get("type") - if type == "string": - if "enum" in props: - return " | ".join([f'"{item}"' for item in props["enum"]]) - return "string" - elif type == "array": - # items is required, OpenAI throws an error if it's missing - return f"{_format_type(props['items'], indent)}[]" - elif type == "object": - return f"{{\n{_format_object_parameters(props, indent + 2)}\n}}" - elif type in ["integer", "number"]: - if "enum" in props: - return " | ".join([f'"{item}"' for item in props["enum"]]) - return "number" - elif type == "boolean": - return "boolean" - elif type == "null": - return "null" - else: - # This is a guess, as an empty string doesn't yield the expected token count - return "any" - - -def _get_num_tokens_from_content_list( - content_list: List[Dict[str, Any]], - use_default_image_token_count: Optional[bool] = False, - default_token_count: Optional[int] = None, -) -> Tuple[str, int]: - """ - Get the number of tokens from a list of content. - - Returns: - Tuple[str, int]: A tuple containing the text and the number of tokens. - """ - try: - num_tokens = 0 - text = "" - for c in content_list: - if c["type"] == "text": - text += c["text"] - num_tokens += len(encoding.encode(c["text"], disallowed_special=())) - elif c["type"] == "image_url": - if isinstance(c["image_url"], dict): - image_url_dict = c["image_url"] - detail = image_url_dict.get("detail", "auto") - url = image_url_dict.get("url") - num_tokens += calculate_img_tokens( - data=url, - mode=detail, - use_default_image_token_count=use_default_image_token_count - or False, - ) - elif isinstance(c["image_url"], str): - image_url_str = c["image_url"] - num_tokens += calculate_img_tokens( - data=image_url_str, - mode="auto", - use_default_image_token_count=use_default_image_token_count - or False, - ) - return text, num_tokens - except Exception as e: - if default_token_count is not None: - return "", default_token_count - raise ValueError( - f"Error getting number of tokens from content list: {e}, default_token_count={default_token_count}" - ) - - def token_counter( model="", custom_tokenizer: Optional[Union[dict, SelectTokenizerResponse]] = None, @@ -1865,100 +1662,11 @@ def token_counter( default_token_count: Optional[int] = None, ) -> int: """ - Count the number of tokens in a given text using a specified model. - - Args: - model (str): The name of the model to use for tokenization. Default is an empty string. - custom_tokenizer (Optional[dict]): A custom tokenizer created with the `create_pretrained_tokenizer` or `create_tokenizer` method. Must be a dictionary with a string value for `type` and Tokenizer for `tokenizer`. Default is None. - text (str): The raw text string to be passed to the model. Default is None. - messages (Optional[List[Dict[str, str]]]): Alternative to passing in text. A list of dictionaries representing messages with "role" and "content" keys. Default is None. - default_token_count (Optional[int]): The default number of tokens to return for a message block, if an error occurs. Default is None. - - Returns: - int: The number of tokens in the text. + The same as `litellm.litellm_core_utils.token_counter`. + + Kept for backwards compatibility. """ - # use tiktoken, anthropic, cohere, llama2, or llama3's tokenizer depending on the model - is_tool_call = False - num_tokens = 0 - if text is None: - if messages is not None: - print_verbose(f"token_counter messages received: {messages}") - text = "" - for message in messages: - if message.get("content", None) is not None: - content = message.get("content") - if isinstance(content, str): - text += message["content"] - elif isinstance(content, List): - text, num_tokens = _get_num_tokens_from_content_list( - content_list=content, - use_default_image_token_count=use_default_image_token_count, - default_token_count=default_token_count, - ) - if message.get("tool_calls"): - is_tool_call = True - for tool_call in message["tool_calls"]: - if "function" in tool_call: - function_arguments = tool_call["function"]["arguments"] - text = ( - text if isinstance(text, str) else "".join(text or []) - ) + (str(function_arguments) if function_arguments else "") - - else: - raise ValueError("text and messages cannot both be None") - elif isinstance(text, List): - text = "".join(t for t in text if isinstance(t, str)) - elif isinstance(text, str): - count_response_tokens = True # user just trying to count tokens for a text. don't add the chat_ml +3 tokens to this - - if model is not None or custom_tokenizer is not None: - tokenizer_json = custom_tokenizer or _select_tokenizer(model=model) - if tokenizer_json["type"] == "huggingface_tokenizer": - enc = tokenizer_json["tokenizer"].encode(text) - num_tokens = len(enc.ids) - elif tokenizer_json["type"] == "openai_tokenizer": - if ( - model in litellm.open_ai_chat_completion_models - or model in litellm.azure_llms - ): - if model in litellm.azure_llms: - # azure llms use gpt-35-turbo instead of gpt-3.5-turbo πŸ™ƒ - model = model.replace("-35", "-3.5") - - print_verbose( - f"Token Counter - using OpenAI token counter, for model={model}" - ) - num_tokens = openai_token_counter( - text=text, # type: ignore - model=model, - messages=messages, - is_tool_call=is_tool_call, - count_response_tokens=count_response_tokens, - tools=tools, - tool_choice=tool_choice, - use_default_image_token_count=use_default_image_token_count - or False, - default_token_count=default_token_count, - ) - else: - print_verbose( - f"Token Counter - using generic token counter, for model={model}" - ) - num_tokens = openai_token_counter( - text=text, # type: ignore - model="gpt-3.5-turbo", - messages=messages, - is_tool_call=is_tool_call, - count_response_tokens=count_response_tokens, - tools=tools, - tool_choice=tool_choice, - use_default_image_token_count=use_default_image_token_count - or False, - default_token_count=default_token_count, - ) - else: - num_tokens = len(encoding.encode(text, disallowed_special=())) # type: ignore - return num_tokens + return token_counter_new(model, custom_tokenizer, text, messages, count_response_tokens, tools, tool_choice, use_default_image_token_count, default_token_count) def supports_httpx_timeout(custom_llm_provider: str) -> bool: diff --git a/tests/local_testing/messages_with_counts.py b/tests/litellm/litellm_core_utils/messages_with_counts.py similarity index 98% rename from tests/local_testing/messages_with_counts.py rename to tests/litellm/litellm_core_utils/messages_with_counts.py index da27a97552..4595a94024 100644 --- a/tests/local_testing/messages_with_counts.py +++ b/tests/litellm/litellm_core_utils/messages_with_counts.py @@ -227,6 +227,7 @@ integer_enum = { ], "tool_choice": "none", "count": 54, + "count-tolerate": 56, # over by 2 } @@ -255,6 +256,7 @@ integer_enum_tool_choice_name = { "function": {"name": "data_demonstration"}, }, # 4 tokens for "data_demonstration" "count": 64, + "count-tolerate": 66, # over by 2 } no_parameters = { @@ -316,6 +318,7 @@ no_parameter_description_or_required = { ], "tool_choice": "auto", "count": 49, + "count-tolerate": 50, # over by 1 } no_parameter_description = { @@ -339,6 +342,7 @@ no_parameter_description = { ], "tool_choice": "auto", "count": 49, + "count-tolerate": 50, # over by 1 } string_enum = { @@ -406,7 +410,8 @@ inner_object = { } ], "tool_choice": "none", - "count": 65, # counted 67, over by 2 + "count": 65, + "count-tolerate" : 67 #over by 2 } """ namespace functions { @@ -453,7 +458,8 @@ inner_object_with_enum_only = { } ], "tool_choice": "none", - "count": 73, # counted 74, over by 1 + "count": 73, + "count-tolerate" : 74 #over by 1 } """ namespace functions { @@ -504,7 +510,8 @@ inner_object_with_enum = { } ], "tool_choice": "none", - "count": 89, # counted 92, over by 3 + "count": 89, + "count-tolerate" : 92, #over by 3 } """ namespace functions { @@ -561,7 +568,8 @@ inner_object_and_string = { } ], "tool_choice": "none", - "count": 103, # counted 106, over by 3 + "count": 103, + "count-tolerate" : 106, #over by 3 } """ namespace functions { @@ -606,7 +614,7 @@ boolean = { } ], "tool_choice": "none", - "count": 89, # over by 3 + "count": 89, } array = { diff --git a/tests/local_testing/test_token_counter.py b/tests/litellm/litellm_core_utils/test_token_counter.py similarity index 82% rename from tests/local_testing/test_token_counter.py rename to tests/litellm/litellm_core_utils/test_token_counter.py index e9445a5c73..f8ec9ed84f 100644 --- a/tests/local_testing/test_token_counter.py +++ b/tests/litellm/litellm_core_utils/test_token_counter.py @@ -1,5 +1,5 @@ #### What this tests #### -# This tests litellm.token_counter() function +# This tests litellm.token_counter.token_counter() function import traceback import os import sys @@ -8,6 +8,7 @@ from unittest.mock import MagicMock import pytest + sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path @@ -19,9 +20,10 @@ from litellm import ( decode, encode, get_modified_max_tokens, - token_counter, + token_counter as token_counter_old, ) -from large_text import text +from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new +from tests.large_text import text from messages_with_counts import ( MESSAGES_TEXT, MESSAGES_WITH_IMAGES, @@ -29,38 +31,48 @@ from messages_with_counts import ( ) +def token_counter_both_assert_same(**args): + new = token_counter_new(**args) + old = token_counter_old(**args) + assert new == old, f"New token counter {new} does not match old token counter {old}" + return new + +## Choose which token_counter the test will use. + +#token_counter = token_counter_new +#token_counter = token_counter_old +token_counter = token_counter_both_assert_same + + def test_token_counter_normal_plus_function_calling(): - try: - messages = [ - {"role": "system", "content": "System prompt"}, - {"role": "user", "content": "content1"}, - {"role": "assistant", "content": "content2"}, - {"role": "user", "content": "conten3"}, - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_E0lOb1h6qtmflUyok4L06TgY", - "function": { - "arguments": '{"query":"search query","domain":"google.ca","gl":"ca","hl":"en"}', - "name": "SearchInternet", - }, - "type": "function", - } - ], - }, - { - "tool_call_id": "call_E0lOb1h6qtmflUyok4L06TgY", - "role": "tool", - "name": "SearchInternet", - "content": "tool content", - }, - ] - tokens = token_counter(model="gpt-3.5-turbo", messages=messages) - print(f"tokens: {tokens}") - except Exception as e: - pytest.fail(f"An exception occurred - {str(e)}") + messages = [ + {"role": "system", "content": "System prompt"}, + {"role": "user", "content": "content1"}, + {"role": "assistant", "content": "content2"}, + {"role": "user", "content": "conten3"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_E0lOb1h6qtmflUyok4L06TgY", + "function": { + "arguments": '{"query":"search query","domain":"google.ca","gl":"ca","hl":"en"}', + "name": "SearchInternet", + }, + "type": "function", + } + ], + }, + { + "tool_call_id": "call_E0lOb1h6qtmflUyok4L06TgY", + "role": "tool", + "name": "SearchInternet", + "content": "tool content", + }, + ] + tokens = token_counter(model="gpt-3.5-turbo", messages=messages) + assert tokens == 80 # test_token_counter_normal_plus_function_calling() @@ -76,6 +88,18 @@ def test_token_counter_textonly(message_count_pair): ) assert counted_tokens == message_count_pair["count"] +@pytest.mark.parametrize( + "message_count_pair", + MESSAGES_TEXT, +) +def test_token_counter_count_response_tokens(message_count_pair): + counted_tokens = token_counter( + model="gpt-35-turbo", messages=[message_count_pair["message"]], count_response_tokens=True + ) + # 3 tokens are not added because of count_response_tokens=True + expected = message_count_pair["count"] - 3 + assert counted_tokens == expected + @pytest.mark.parametrize( "message_count_pair", @@ -100,11 +124,23 @@ def test_token_counter_with_tools(message_count_pair): tool_choice=message_count_pair["tool_choice"], ) expected_tokens = message_count_pair["count"] - diff = counted_tokens - expected_tokens - assert ( - diff >= 0 and diff <= 3 - ), f"Expected {expected_tokens} tokens, got {counted_tokens}. Counted tokens is only allowed to be off by 3 in the over-counting direction." + actual_diff = counted_tokens - expected_tokens + if "count-tolerate" in message_count_pair: + if message_count_pair["count-tolerate"] == counted_tokens: + pass # expected + else: + tolerated_diff = message_count_pair["count-tolerate"] - expected_tokens + assert actual_diff <= tolerated_diff, f"Expected {expected_tokens} tokens, got {counted_tokens}. Counted tokens is only allowed to be off by {tolerated_diff} in the over-counting direction." + if actual_diff != tolerated_diff: + raise NeedsToleranceUpdateError(f"SOMETHING BROKEN GOT FIXED! THIS is good! Adjust 'count-tolerate' from {message_count_pair['count-tolerate']} to {counted_tokens}") + + else: + assert expected_tokens == counted_tokens, f"Expected {expected_tokens} tokens, got {counted_tokens}." + +class NeedsToleranceUpdateError(Exception): + """Custom exception to mark tests that have improved""" + pass def test_tokenizers(): try: @@ -184,7 +220,7 @@ def test_encoding_and_decoding(): # llama2 encoding + decoding llama2_tokens = encode(model="meta-llama/Llama-2-7b-chat", text=sample_text) llama2_text = decode( - model="meta-llama/Llama-2-7b-chat", tokens=llama2_tokens.ids + model="meta-llama/Llama-2-7b-chat", tokens=llama2_tokens.ids # type: ignore ) assert llama2_text == sample_text @@ -350,6 +386,7 @@ def test_empty_tools(): print(result) +@pytest.mark.skip(reason="Skipping this test temporarily because it relies on a function being called that I am removing.") def test_gpt_4o_token_counter(): with patch.object( litellm.utils, "openai_token_counter", new=MagicMock() @@ -383,6 +420,35 @@ def test_img_url_token_counter(img_url): def test_token_encode_disallowed_special(): encode(model="gpt-3.5-turbo", text="Hello, world! <|endoftext|>") +def test_token_counter(): + try: + messages = [{"role": "user", "content": "hi how are you what time is it"}] + tokens = token_counter(model="gpt-3.5-turbo", messages=messages) + print("gpt-35-turbo") + print(tokens) + assert tokens > 0 + + tokens = token_counter(model="claude-2", messages=messages) + print("claude-2") + print(tokens) + assert tokens > 0 + + tokens = token_counter(model="gemini/chat-bison", messages=messages) + print("gemini/chat-bison") + print(tokens) + assert tokens > 0 + + tokens = token_counter(model="ollama/llama2", messages=messages) + print("ollama/llama2") + print(tokens) + assert tokens > 0 + + tokens = token_counter(model="anthropic.claude-instant-v1", messages=messages) + print("anthropic.claude-instant-v1") + print(tokens) + assert tokens > 0 + except Exception as e: + pytest.fail(f"Error occurred: {e}") import unittest from unittest.mock import patch, MagicMock diff --git a/tests/litellm/litellm_core_utils/test_token_counter_tool.py b/tests/litellm/litellm_core_utils/test_token_counter_tool.py new file mode 100644 index 0000000000..36cab0bc50 --- /dev/null +++ b/tests/litellm/litellm_core_utils/test_token_counter_tool.py @@ -0,0 +1,84 @@ +#### What this tests #### +# This tests litellm.token_counter() function +import os +import sys + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +#Use the same token_counter as the main test. +from test_token_counter import token_counter + +from test_token_counter_tool_data import * + + +@pytest.mark.parametrize( + "messages", + [ + SHORT, + CONTENT_AND_TOOL_CALL, + SYSTEM_LONG, + TOOL_CALL_CONTENT_ARRAY, + ], + ids=[ + "SHORT", + "CONTENT_AND_TOOL_CALL", + "SYSTEM_LONG", + "TOOL_CALL_CONTENT_ARRAY", + ], +) +def test_token_counter_tool_increases(messages): + conversation = [] + prev_tokens = 0 + for message in messages: + conversation.append(message) + tokens = token_counter(model="gpt-3.5-turbo", messages=conversation, tools=TOOLS) # type: ignore + print(f"tokens: {tokens}") + assert ( + tokens > prev_tokens + ), f"Token did not increase: {tokens} <= {prev_tokens}" + prev_tokens = tokens + + +# Reuse in multiple tests +@pytest.mark.parametrize("usermessage", USER_MESSAGES, ids=USER_MESSAGES_IDS) +@pytest.mark.parametrize("tool_call", TOOL_CALL_MESSAGES, ids=TOOL_CALL_MESSAGE_ids) +def test_grow(usermessage, tool_call): + assertGrow(usermessage, tool_call, False) + + +@pytest.mark.parametrize("usermessage", USER_MESSAGES, ids=USER_MESSAGES_IDS) +@pytest.mark.parametrize("tool_call", TOOL_CALL_MESSAGES, ids=TOOL_CALL_MESSAGE_ids) +def test_sum(usermessage, tool_call): + assertGrow(usermessage, tool_call, True) + + +def assertGrow(usermessage, tool_call, assertBiggerThanBoth=True): + print(usermessage, tool_call) + tokens_usermessage = token_counter( + model="anthropic.claude-instant-v1", + messages=[usermessage], + tools=TOOLS, # type: ignore + ) + assert tokens_usermessage > 0 + tokens_tool_call = token_counter( + model="anthropic.claude-instant-v1", + messages=[tool_call], + tools=TOOLS, # type: ignore + ) + assert tokens_tool_call > 0 + tokens_both = token_counter( + model="anthropic.claude-instant-v1", + messages=[usermessage, tool_call], + tools=TOOLS, # type: ignore + ) + assert tokens_both > tokens_usermessage + assert tokens_both > tokens_tool_call + if assertBiggerThanBoth: + assert abs(tokens_usermessage + tokens_tool_call - tokens_both) <= 61, ( + f"tokens_usermessage: {tokens_usermessage}, tokens_tool_call: {tokens_tool_call}, " + + f"tokens_both: {tokens_both} diff: {tokens_usermessage + tokens_tool_call - tokens_both}" + ) diff --git a/tests/litellm/litellm_core_utils/test_token_counter_tool_data.py b/tests/litellm/litellm_core_utils/test_token_counter_tool_data.py new file mode 100644 index 0000000000..bf030b673b --- /dev/null +++ b/tests/litellm/litellm_core_utils/test_token_counter_tool_data.py @@ -0,0 +1,166 @@ +SHORT = [ + {"role": "system", "content": "System prompt"}, + { + "tool_call_id": "call_E0lOb1h6qtmflUyok4L06TgY", + "role": "tool", + "name": "SearchInternet", + "content": "tool content", + }, +] + +CONTENT_AND_TOOL_CALL = [ + {"role": "system", "content": "System prompt"}, + { + "role": "assistant", + "content": "..", + "tool_calls": [ + { + "id": "toolu_01NZjJ3e6fzhhbYYcb5k1k4d", + "type": "function", + "function": { + "name": "str_replace_editor", + "arguments": '{"command": "view","path": "/workspace"}', + }, + } + ], + }, +] + +_OPENHANDS_SYSTEM_MESSAGE = { + "content": [ + { + "type": "text", + "text": "You are OpenHands agent, a helpful AI assistant that can " + "interact with a computer to solve tasks.\n\n\nYour primary " + "role is to assist users by executing commands, modifying code, and " + "solving technical problems effectively. You should be thorough, " + "methodical, and prioritize quality over speed.\n* If the user asks a " + "question, like 'why is X happening', don’t try to fix the problem. ", + } + ], + "role": "system", + } + +SYSTEM_LONG = [ + _OPENHANDS_SYSTEM_MESSAGE, + { + "role": "assistant", + "tool_calls": [ + { + "id": "toolu_01NZjJ3e6fzhhbYYcb5k1k4d", + "type": "function", + "function": { + "name": "str_replace_editor", + "arguments": '{"command": "view","path": "/workspace"}', + }, + } + ], + }, +] + +TOOL_CALL_CONTENT_ARRAY = [ + _OPENHANDS_SYSTEM_MESSAGE, + { + "role": "assistant", + "content": [{"type": "text", "text": "very short"}], + "tool_calls": [ + { + "id": "toolu_01NZjJ3e6fzhhbYYcb5k1k4d", + "type": "function", + "function": { + "name": "str_replace_editor", + "arguments": '{"command": "view","path": "/workspace"}', + }, + } + ], + }, +] + +USER_MESSAGES = [ + {"role": "system", "content": "System prompt"}, + {"content": [{"type": "text", "text": "2025-04-25 (UTC)"}], "role": "user"}, + { + "role": "system", + "content": "System prompt long " + ("long" * 100), + }, + { + "role": "system", + "content": ["System prompt long " + ("long" * 100)], + }, + { + "role": "system", + "content": [{"type": "text", "text": "System prompt long " + ("long" * 100)}], + }, +] +USER_MESSAGES_IDS = ["system", "user-arr", "long", "long-arr", "long-arr-obj"] + +TOOL_CALL_MESSAGES = [ + { + "role": "assistant", + "tool_calls": [ + { + "id": "x", + "type": "function", + "function": { + "arguments": '{"arg":"value"}', + "name": "foo", + }, + } + ], + }, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_E0lOb1h6qtmflUyok4L06TgY", + "function": { + "arguments": '{"query":"search query","domain":"google.ca","gl":"ca","hl":"en"}', + "name": "SearchInternet", + }, + "type": "function", + } + ], + }, + {"role": "user", "content": "Hi"}, + { + "role": "assistant", + "content": [{"type": "text", "text": "very short"}], + "tool_calls": [ + { + "id": "toolu_01NZjJ3e6fzhhbYYcb5k1k4d", + "type": "function", + "function": { + "name": "str_replace_editor", + "arguments": '{"command": "view","path": "/workspace"}', + }, + } + ], + }, +] + +TOOL_CALL_MESSAGE_ids = ["t_short", "t_normal", "usermsg", "t_content_array"] + + +TOOLS = [ + { + "type": "function", + "function": { + "name": "data_demonstration", + "description": "This is the main function description", + "parameters": { + "type": "object", + "properties": { + "object_1": { + "type": "object", + "description": "The object data type as a property", + "properties": { + "string1": {"type": "string"}, + }, + } + }, + "required": ["object_1"], + }, + }, + } +] diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index 398b69067a..ec779c17bb 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -33,7 +33,6 @@ from litellm.utils import ( get_supported_openai_params, get_token_count, get_valid_models, - token_counter, trim_messages, validate_environment, ) @@ -445,41 +444,6 @@ def test_function_to_dict(): # test_function_to_dict() - -def test_token_counter(): - try: - messages = [{"role": "user", "content": "hi how are you what time is it"}] - tokens = token_counter(model="gpt-3.5-turbo", messages=messages) - print("gpt-35-turbo") - print(tokens) - assert tokens > 0 - - tokens = token_counter(model="claude-2", messages=messages) - print("claude-2") - print(tokens) - assert tokens > 0 - - tokens = token_counter(model="gemini/chat-bison", messages=messages) - print("gemini/chat-bison") - print(tokens) - assert tokens > 0 - - tokens = token_counter(model="ollama/llama2", messages=messages) - print("ollama/llama2") - print(tokens) - assert tokens > 0 - - tokens = token_counter(model="anthropic.claude-instant-v1", messages=messages) - print("anthropic.claude-instant-v1") - print(tokens) - assert tokens > 0 - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - -# test_token_counter() - - @pytest.mark.parametrize( "model, expected_bool", [