Merge fb5414393e into 4212f1af27

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):

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,