From 1ef9801ae0c2038d627f3bc73238704958a51f68 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 02:46:49 -0700 Subject: [PATCH] fix(token_counter): count null and SDK-shaped custom tool calls and map them for provider counts Pick the tool call payload with is_custom_tool_call_dict so a custom call that carries "function": null counts its input, and count a null payload as an empty call instead of failing the request. Tool calls with neither key still raise. Send custom tool calls, their results, and custom tool definitions to OpenAI's input_tokens endpoint as custom_tool_call, custom_tool_call_output, and flat custom tools, and count a custom tool definition's description and grammar locally. --- litellm/litellm_core_utils/token_counter.py | 38 ++++++-- .../responses/count_tokens/transformation.py | 79 +++++++++++++---- .../litellm_core_utils/test_token_counter.py | 87 ++++++++++++++++--- ...test_openai_count_tokens_transformation.py | 49 +++++++++++ 4 files changed, 214 insertions(+), 39 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 6c85a6606c6..70d1f249c3d 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -35,10 +35,14 @@ from litellm.constants import ( TOKEN_COUNTER_MAX_EXACT_CHARS, ) from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.chat_template_tools import ( ChatTemplateToolSection, chat_template_tool_section, ) +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_custom_tool_format_to_responses_shape, +) from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding if TYPE_CHECKING: @@ -64,7 +68,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolParam, OpenAIMessageContentListBlock, ) -from litellm.types.utils import Message, SelectTokenizerResponse +from litellm.types.utils import Message, SelectTokenizerResponse, is_custom_tool_call_dict def get_modified_max_tokens( @@ -570,13 +574,9 @@ 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" in tool_call: - kind, field = "function", "arguments" - elif "custom" in tool_call: - kind, field = "custom", "input" - else: + if (fields := as_str_mapping(tool_call)) is None or ("function" not in fields and "custom" not in fields): raise ValueError(f"Unsupported tool call {tool_call} must contain a function or custom key") - total += count_function(str(tool_call[kind].get(field, ""))) + total += count_function(_tool_call_text(fields)) return total if key == "function_call": if not isinstance(value, Mapping): @@ -585,6 +585,14 @@ def _count_function_call_tokens( raise ValueError(f"Unexpected key {key!r}; expected 'tool_calls' or 'function_call'") +def _tool_call_text(tool_call: Mapping[str, object]) -> str: + is_custom: Final = is_custom_tool_call_dict(tool_call) + payload: Final = as_str_mapping(tool_call.get("custom" if is_custom else "function")) + if payload is None: + return "" + return str(payload.get("input" if is_custom else "arguments", "")) + + def _count_messages( params: _MessageCountParams, messages: list[AllMessageValues], @@ -1149,6 +1157,10 @@ def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mappi if isinstance(function, Mapping): yield function return + custom: Final = as_str_mapping(tool.get("custom")) + if custom is not None: + yield _custom_tool_definition(custom) + return declarations: Final = tool.get("function_declarations") or tool.get("functionDeclarations") if isinstance(declarations, list): for declaration in declarations: @@ -1173,6 +1185,18 @@ def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mappi } +def _custom_tool_definition(custom: Mapping[str, object]) -> Mapping[str, object]: + format_obj: Final = as_str_mapping(custom.get("format")) + grammar: Final = ( + convert_custom_tool_format_to_responses_shape(format_obj).get("definition") if format_obj is not None else None + ) + description_parts: Final = tuple(str(part) for part in (custom.get("description"), grammar) if part) + return { + "name": custom.get("name"), + **({"description": "\n".join(description_parts)} if description_parts else {}), + } + + def _format_single_function_definition(function: Mapping[str, object]) -> tuple[str, ...]: function_name: Final = function.get("name") if not function_name: diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index b4e152ff406..f8c18795f61 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -5,14 +5,17 @@ This module handles the transformation of requests to OpenAI's /v1/responses/inp """ import copy -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from typing import Final, Literal, Protocol, TypeAlias from openai.types.responses.input_token_count_params import InputTokenCountParams from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict -from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes +from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items, require_attributes +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_custom_tool_format_to_responses_shape, +) from litellm.litellm_core_utils.tool_fields import DropPolicy from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.messages.convertible_blocks import malformed_tool_result_text_error @@ -23,7 +26,7 @@ from litellm.llms.anthropic.messages.responses_adapters.transformation import ( from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.anthropic import AnthropicMessagesRequest, AnthropicOutputSchema from litellm.types.llms.openai import ResponsesAPIRequestParams -from litellm.types.utils import all_litellm_params +from litellm.types.utils import all_litellm_params, is_custom_tool_call_dict class ResponsesInputTextPart(TypedDict): @@ -82,6 +85,47 @@ def _function_fields(value: object) -> _FunctionFields: return require_attributes(value, _FunctionFields, "get", "__contains__", "__getitem__") +def _is_custom_tool_call(tool_call: _FieldReader) -> bool: + return is_custom_tool_call_dict({"type": tool_call.get("type"), "custom": tool_call.get("custom")}) + + +def _assistant_tool_calls(messages: Sequence[Mapping[str, object]]) -> Iterable[_FieldReader]: + for msg in messages: + tool_calls = msg.get("tool_calls") + if msg.get("role") == "assistant" and tool_calls: + yield from (_field_reader(tool_call) for tool_call in iter_items(tool_calls)) + + +def _responses_tool_call_item(tool_call: _FieldReader) -> Mapping[str, object]: + call_id: Final = tool_call.get("id", "") + if _is_custom_tool_call(tool_call): + custom: Final = _field_reader(tool_call.get("custom")) + return { + "type": "custom_tool_call", + "call_id": call_id, + "name": custom.get("name", ""), + "input": custom.get("input", ""), + } + function: Final = _field_reader(tool_call.get("function", {})) + return { + "type": "function_call", + "call_id": call_id, + "name": function.get("name", ""), + "arguments": function.get("arguments", ""), + } + + +def _responses_custom_tool(custom: Mapping[str, object]) -> Mapping[str, object]: + description: Final = custom.get("description") + format_obj: Final = as_str_mapping(custom.get("format")) + return { + "type": "custom", + "name": custom.get("name", ""), + **({"description": description} if description is not None else {}), + **({"format": convert_custom_tool_format_to_responses_shape(format_obj)} if format_obj is not None else {}), + } + + _ANTHROPIC_OUTPUT_SCHEMA: Final = TypeAdapter(AnthropicOutputSchema) @@ -239,6 +283,8 @@ class OpenAICountTokensConfig: if "strict" in func: item["strict"] = func["strict"] transformed.append(item) + elif tool.get("type") == "custom" and (custom := as_str_mapping(tool.get("custom"))) is not None: + transformed.append(dict(_responses_custom_tool(custom))) elif "input_schema" in tool and "name" in tool and tool.get("type") in (None, "custom"): transformed.append( { @@ -267,6 +313,9 @@ class OpenAICountTokensConfig: """ input_items: Final[list[dict[str, object]]] = [] instructions_parts: Final[list[str]] = [] + custom_tool_call_ids: Final = frozenset( + tool_call.get("id") for tool_call in _assistant_tool_calls(messages) if _is_custom_tool_call(tool_call) + ) for msg in messages: role = msg.get("role", "") @@ -297,26 +346,20 @@ class OpenAICountTokensConfig: if content: input_items.append({"role": "assistant", "content": content}) if tool_calls: - for tc in iter_items(tool_calls): - tool_call = _field_reader(tc) - function = tool_call.get("function", {}) - call_id = tool_call.get("id", "") - function_fields = _field_reader(function) - input_items.append( - { - "type": "function_call", - "call_id": call_id, - "name": function_fields.get("name", ""), - "arguments": function_fields.get("arguments", ""), - } - ) + input_items.extend( + dict(_responses_tool_call_item(_field_reader(tool_call))) + for tool_call in iter_items(tool_calls) + ) elif not content: input_items.append({"role": "assistant", "content": content}) elif role == "tool": + tool_call_id = msg.get("tool_call_id", "") input_items.append( { - "type": "function_call_output", - "call_id": msg.get("tool_call_id", ""), + "type": "custom_tool_call_output" + if tool_call_id in custom_tool_call_ids + else "function_call_output", + "call_id": tool_call_id, "output": content if isinstance(content, str) else str(content), } ) diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 5d72b1463aa..456bfb4c594 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -377,24 +377,32 @@ 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 _tool_call_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"}, + ] - 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}}), - ) +_CUSTOM_CODE = "print('hello world from a custom tool')\n" * 50 + + +@pytest.mark.parametrize( + "custom_call", + ( + {"id": "call_1", "type": "custom", "custom": {"name": "code_exec", "input": _CUSTOM_CODE}}, + {"id": "call_1", "type": "custom", "function": None, "custom": {"name": "code_exec", "input": _CUSTOM_CODE}}, + {"id": "call_1", "custom": {"name": "code_exec", "input": _CUSTOM_CODE}}, + ), + ids=("typed", "sdk-shaped-null-function", "typeless"), +) +def test_token_counter_custom_tool_call_counts_input_like_function_arguments(custom_call: dict[str, object]): + custom_tokens = token_counter(model="gpt-4o", messages=_tool_call_turn(custom_call)) function_tokens = token_counter( model="gpt-4o", - messages=assistant_turn( - {"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": code}} + messages=_tool_call_turn( + {"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": _CUSTOM_CODE}} ), ) @@ -404,6 +412,57 @@ def test_token_counter_custom_tool_call_counts_input_like_function_arguments(): assert custom_tokens > 300, f"custom tool call input must contribute real tokens, got {custom_tokens}" +@pytest.mark.parametrize( + ("null_call", "empty_call"), + ( + ( + {"id": "call_1", "type": "custom", "custom": None}, + {"id": "call_1", "type": "custom", "custom": {"name": "code_exec", "input": ""}}, + ), + ( + {"id": "call_1", "type": "function", "function": None}, + {"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": ""}}, + ), + ), + ids=("custom", "function"), +) +def test_token_counter_null_tool_call_payload_counts_as_empty_call( + null_call: dict[str, object], empty_call: dict[str, object] +): + assert token_counter(model="gpt-4o", messages=_tool_call_turn(null_call)) == token_counter( + model="gpt-4o", messages=_tool_call_turn(empty_call) + ) + + +def test_token_counter_rejects_tool_call_without_function_or_custom_payload(): + with pytest.raises(ValueError, match="must contain a function or custom key"): + token_counter(model="gpt-4o", messages=_tool_call_turn({"id": "call_1", "type": "function"})) + + +def test_token_counter_counts_custom_tool_definition_description_and_grammar(): + messages = [{"role": "user", "content": "run the code"}] + grammar = "start: " + " | ".join(f'"command_{index}"' for index in range(100)) + + def custom_tool(**fields: object) -> dict[str, object]: + return {"type": "custom", "custom": {"name": "code_exec", **fields}} + + no_tools = token_counter(model="gpt-4o", messages=messages) + bare = token_counter(model="gpt-4o", messages=messages, tools=[custom_tool()]) + described = token_counter( + model="gpt-4o", + messages=messages, + tools=[ + custom_tool( + description="Runs one shell command", + format={"type": "grammar", "grammar": {"syntax": "lark", "definition": grammar}}, + ) + ], + ) + + assert bare > no_tools, f"a custom tool's name must count, got {bare} vs {no_tools} without tools" + assert described - bare > 300, f"description and grammar must count, got {described} vs {bare} bare" + + @pytest.mark.parametrize( "message_count_pair", MESSAGES_TEXT, diff --git a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py index d5e7bb4f93b..73a956579bd 100644 --- a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py @@ -225,6 +225,55 @@ def test_messages_to_responses_input_reads_a_tool_call_that_forwards_its_fields( } +def test_messages_to_responses_input_maps_custom_tool_calls_and_their_outputs(): + messages: Final = [ + {"role": "user", "content": "Run it"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_custom", "type": "custom", "custom": {"name": "code_exec", "input": "print(2+3)"}}, + {"id": "call_fn", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}}, + ], + }, + {"role": "tool", "tool_call_id": "call_custom", "content": "5"}, + {"role": "tool", "tool_call_id": "call_fn", "content": "sunny"}, + ] + + input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages) + + assert input_items[1:] == [ + {"type": "custom_tool_call", "call_id": "call_custom", "name": "code_exec", "input": "print(2+3)"}, + {"type": "function_call", "call_id": "call_fn", "name": "get_weather", "arguments": "{}"}, + {"type": "custom_tool_call_output", "call_id": "call_custom", "output": "5"}, + {"type": "function_call_output", "call_id": "call_fn", "output": "sunny"}, + ] + + +def test_transform_custom_tool_uses_the_flat_responses_shape(): + tool: Final = { + "type": "custom", + "custom": { + "name": "code_exec", + "description": "Run python", + "format": {"type": "grammar", "grammar": {"syntax": "lark", "definition": "start: /.+/"}}, + }, + } + + result: Final = OpenAICountTokensConfig().transform_request_to_count_tokens( + model="gpt-4o", input="Run it", tools=[tool] + ) + + assert result["tools"] == [ + { + "type": "custom", + "name": "code_exec", + "description": "Run python", + "format": {"type": "grammar", "syntax": "lark", "definition": "start: /.+/"}, + } + ] + + def test_messages_to_responses_input_preserves_images(): """An image block must survive the round trip, or OpenAI counts only the text.