diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 591d77bdb16..3ffb77c4940 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -795,10 +795,13 @@ def _count_function_call_tokens( raise ValueError(f"Unsupported type {type(value)} for key tool_calls in message {message}") total = 0 for tool_call in value: - if "function" not in tool_call: - raise ValueError(f"Unsupported tool call {tool_call} must contain a function key") - function_arguments = tool_call["function"].get("arguments", "") - total += count_function(str(function_arguments)) + if "function" in tool_call: + kind, field = "function", "arguments" + elif "custom" in tool_call: + kind, field = "custom", "input" + else: + raise ValueError(f"Unsupported tool call {tool_call} must contain a function or custom key") + total += count_function(str(tool_call[kind].get(field, ""))) return total if key == "function_call": if not isinstance(value, Mapping): diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index d97970753f8..3dd3e58b9ba 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -381,6 +381,33 @@ def test_token_counter_legacy_function_call_counts_arguments(): assert fc_tokens > 500, f"4000-char arguments payload must contribute real tokens, got {fc_tokens}" +def test_token_counter_custom_tool_call_counts_input_like_function_arguments(): + code = "print('hello world from a custom tool')\n" * 50 + + def assistant_turn(tool_call: dict[str, object]) -> list[dict[str, object]]: + return [ + {"role": "user", "content": "run the code"}, + {"role": "assistant", "content": None, "tool_calls": [tool_call]}, + {"role": "tool", "tool_call_id": "call_1", "content": "done"}, + ] + + custom_tokens = token_counter( + model="gpt-4o", + messages=assistant_turn({"id": "call_1", "type": "custom", "custom": {"name": "code_exec", "input": code}}), + ) + function_tokens = token_counter( + model="gpt-4o", + messages=assistant_turn( + {"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": code}} + ), + ) + + assert custom_tokens == function_tokens, ( + f"custom tool call input must count like function arguments, got {custom_tokens} vs {function_tokens}" + ) + assert custom_tokens > 300, f"custom tool call input must contribute real tokens, got {custom_tokens}" + + @pytest.mark.parametrize( "message_count_pair", MESSAGES_TEXT,