fix(token_counter): count legacy function_call.arguments (VERIA-492) (#31741)

* fix(token_counter): count legacy function_call.arguments (VERIA-492)

token_counter handled the modern assistant tool_calls field but had no
branch for the legacy OpenAI function_call payload. The value is a dict,
so it skipped every special-cased branch in _count_messages and fell
through to the unsupported-key continue, letting arbitrary text in
function_call.arguments slip past the count.

Resolves VERIA-492

* refactor(token_counter): raise on unexpected key in _count_function_call_tokens

Address Greptile P2: the helper's fallback branch previously applied
function_call logic to any key that wasn't tool_calls. Make the contract
explicit so a future caller can't silently miscount.

litellm/litellm_core_utils/token_counter.py

@@ -407,6 +407,37 @@ def token_counter(
     return num_tokens
 
 
+def _count_function_call_tokens(
+    key: str,
+    value: Any,
+    message: Mapping[str, Any],
+    count_function: TokenCounterFunction,
+) -> int:
+    """
+    Count tokens contributed by an assistant message's tool/function call payload.
+
+    Handles both the modern `tool_calls` list and the legacy OpenAI
+    `function_call` dict. Only the `arguments` string is counted (matching the
+    existing tool_calls behavior); names are accounted for elsewhere via the
+    tool/function definitions and `tool_choice`.
+    """
+    if key == "tool_calls":
+        if not isinstance(value, List):
+            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))
+        return total
+    if key == "function_call":
+        if not isinstance(value, Mapping):
+            raise ValueError(f"Unsupported type {type(value)} for key function_call in message {message}")
+        return count_function(str(value.get("arguments", "")))
+    raise ValueError(f"Unexpected key {key!r}; expected 'tool_calls' or 'function_call'")
+
+
 def _count_messages(
     params: _MessageCountParams,
     messages: List[AllMessageValues],
@@ -430,16 +461,8 @@ def _count_messages(
         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 key in ("tool_calls", "function_call"):
+                num_tokens += _count_function_call_tokens(key, value, message, params.count_function)
             elif isinstance(value, str):
                 num_tokens += params.count_function(value)
                 if key == "name":

tests/test_litellm/litellm_core_utils/test_token_counter.py

@@ -97,6 +97,50 @@ def test_token_counter_normal_plus_function_calling():
 # test_token_counter_normal_plus_function_calling()
 
 
+def test_token_counter_legacy_function_call_counts_arguments():
+    """
+    Regression for VERIA-492 (Token-counter function_call bypass).
+
+    The legacy OpenAI assistant `function_call` field carries arbitrary text in
+    `arguments`. Before the fix, `_count_messages` had no branch for
+    `function_call` and fell through to the unsupported-key `continue`, so an
+    assistant turn could smuggle unlimited text past `token_counter` and the
+    proxy `/utils/token_counter` endpoint (and downstream pre-call budget /
+    `get_modified_max_tokens` math). After the fix it must be counted the
+    same as the equivalent `tool_calls` payload.
+    """
+    long_arg = "A" * 4000
+    fc_messages = [
+        {"role": "user", "content": "hi"},
+        {
+            "role": "assistant",
+            "content": None,
+            "function_call": {"name": "search", "arguments": long_arg},
+        },
+    ]
+    tc_messages = [
+        {"role": "user", "content": "hi"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "call_1",
+                    "type": "function",
+                    "function": {"name": "search", "arguments": long_arg},
+                }
+            ],
+        },
+    ]
+    fc_tokens = token_counter(model="gpt-3.5-turbo", messages=fc_messages)
+    tc_tokens = token_counter(model="gpt-3.5-turbo", messages=tc_messages)
+    assert fc_tokens == tc_tokens, (
+        f"function_call arguments must count like tool_calls arguments; "
+        f"got function_call={fc_tokens}, tool_calls={tc_tokens}"
+    )
+    assert fc_tokens > 500, f"4000-char arguments payload must contribute real tokens, got {fc_tokens}"
+
+
 @pytest.mark.parametrize(
     "message_count_pair",
     MESSAGES_TEXT,