From bcbf2b7a3d5dd397f148d244ca65661068eaa516 Mon Sep 17 00:00:00 2001 From: John Karam <168653085+johnkaram02@users.noreply.github.com> Date: Tue, 27 Oct 2026 23:53:19 +0000 Subject: [PATCH 1/4] fix(token_counter): count custom tool calls instead of raising token_counter raised "Unsupported tool call ... must contain a function key" for any assistant tool call without a function key. OpenAI custom tool calls carry their payload under custom.input instead, and LiteLLM returns them as ChatCompletionMessageCustomToolCall, so passing a conversation with one back to token_counter failed Count custom.input the same way function.arguments is counted --- litellm/litellm_core_utils/token_counter.py | 10 ++++--- .../litellm_core_utils/test_token_counter.py | 27 +++++++++++++++++++ 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index edb5ecf0f8d..5784e91f98c 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -570,10 +570,12 @@ 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: + total += count_function(str(tool_call["function"].get("arguments", ""))) + elif "custom" in tool_call: + total += count_function(str(tool_call["custom"].get("input", ""))) + else: + raise ValueError(f"Unsupported tool call {tool_call} must contain a function or custom key") 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 9921bbe343a..5d72b1463aa 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -377,6 +377,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, From fb5414393e63c065ee6b00ed8a4a14fa77c3b5a6 Mon Sep 17 00:00:00 2001 From: John Karam <168653085+johnkaram02@users.noreply.github.com> Date: Thu, 29 Oct 2026 13:19:34 +0000 Subject: [PATCH 2/4] refactor(token_counter): read function and custom tool call text through one lookup --- litellm/litellm_core_utils/token_counter.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 5784e91f98c..6c85a6606c6 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -571,11 +571,12 @@ def _count_function_call_tokens( total = 0 for tool_call in value: if "function" in tool_call: - total += count_function(str(tool_call["function"].get("arguments", ""))) + kind, field = "function", "arguments" elif "custom" in tool_call: - total += count_function(str(tool_call["custom"].get("input", ""))) + 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): 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 3/4] 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. From 41a857c34e4c878947c3de78fe7e8321007ad04e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 09:00:01 -0700 Subject: [PATCH 4/4] fix(token_counter): count SDK tool call objects and flat custom tools, keep Anthropic options tools as functions Reused ChatCompletionMessageToolCall and custom tool call objects are read through their get and contains methods, a null custom or function payload maps to an empty call on the provider count, flat Responses-shaped custom tools are counted like the chat shape, and Anthropic tools that carry a custom options object stay function tools. --- litellm/litellm_core_utils/token_counter.py | 26 ++++++-- .../responses/count_tokens/transformation.py | 8 +-- .../litellm_core_utils/test_token_counter.py | 65 ++++++++++++++++++- ...test_openai_count_tokens_transformation.py | 42 +++++++++++- 4 files changed, 127 insertions(+), 14 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 12ff7a04f90..fc84cd4f807 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -47,7 +47,7 @@ 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.core_helpers import as_str_mapping, has_attributes from litellm.litellm_core_utils.prompt_templates.chat_template_tools import ( ChatTemplateToolSection, chat_template_tool_section, @@ -893,7 +893,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 (fields := as_str_mapping(tool_call)) is None or ("function" not in fields and "custom" not in fields): + if (fields := _tool_call_fields(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(_tool_call_text(fields)) return total @@ -904,9 +906,19 @@ 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")) +class _ToolCallFields(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + def __contains__(self, key: object, /) -> bool: ... + + +def _tool_call_fields(value: object) -> _ToolCallFields | None: + return value if has_attributes(value, _ToolCallFields, "get", "__contains__") else None + + +def _tool_call_text(tool_call: _ToolCallFields) -> str: + is_custom: Final = is_custom_tool_call_dict({"type": tool_call.get("type"), "custom": tool_call.get("custom")}) + payload: Final = _tool_call_fields(tool_call.get("custom" if is_custom else "function")) if payload is None: return "" return str(payload.get("input" if is_custom else "arguments", "")) @@ -1812,8 +1824,8 @@ def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mappi yield function return custom: Final = as_str_mapping(tool.get("custom")) - if custom is not None: - yield _custom_tool_definition(custom) + if "input_schema" not in tool and (custom is not None or tool.get("type") == "custom"): + yield _custom_tool_definition(custom if custom is not None else tool) return declarations: Final = tool.get("function_declarations") or tool.get("functionDeclarations") if isinstance(declarations, list): diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 3141f57b74e..4af64bb6969 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -99,14 +99,14 @@ def _assistant_tool_calls(messages: Sequence[Mapping[str, object]]) -> Iterable[ 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")) + custom: Final = _field_reader(tool_call.get("custom") or {}) 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", {})) + function: Final = _field_reader(tool_call.get("function") or {}) return { "type": "function_call", "call_id": call_id, @@ -283,8 +283,6 @@ 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( { @@ -295,6 +293,8 @@ class OpenAICountTokensConfig: "strict": bool(tool.get("strict")), } ) + elif tool.get("type") == "custom" and (custom := as_str_mapping(tool.get("custom"))) is not None: + transformed.append(dict(_responses_custom_tool(custom))) else: # Pass through non-function tools (e.g., web_search, file_search) transformed.append(tool) diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index e228d299d2e..fc9b17106b8 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -50,7 +50,13 @@ from litellm.litellm_core_utils.token_counter import ( from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam -from litellm.types.utils import SelectTokenizerResponse +from litellm.types.utils import ( + ChatCompletionCustomToolCallPayload, + ChatCompletionMessageCustomToolCall, + ChatCompletionMessageToolCall, + Function, + SelectTokenizerResponse, +) from tests.large_text import text from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, @@ -385,7 +391,7 @@ 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 _tool_call_turn(tool_call: dict[str, object]) -> list[dict[str, object]]: +def _tool_call_turn(tool_call: object) -> list[dict[str, object]]: return [ {"role": "user", "content": "run the code"}, {"role": "assistant", "content": None, "tool_calls": [tool_call]}, @@ -420,6 +426,37 @@ def test_token_counter_custom_tool_call_counts_input_like_function_arguments(cus assert custom_tokens > 300, f"custom tool call input must contribute real tokens, got {custom_tokens}" +@pytest.mark.parametrize( + ("sdk_call", "dict_call"), + ( + ( + ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="code_exec", arguments=_CUSTOM_CODE) + ), + {"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": _CUSTOM_CODE}}, + ), + ( + ChatCompletionMessageCustomToolCall( + id="call_1", custom=ChatCompletionCustomToolCallPayload(name="code_exec", input=_CUSTOM_CODE) + ), + {"id": "call_1", "type": "custom", "custom": {"name": "code_exec", "input": _CUSTOM_CODE}}, + ), + ( + {"id": "call_1", "type": "function", "function": Function(name="code_exec", arguments=_CUSTOM_CODE)}, + {"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": _CUSTOM_CODE}}, + ), + ), + ids=("function", "custom", "dict-holding-function-object"), +) +def test_token_counter_counts_reused_sdk_tool_call_objects_like_their_dicts( + sdk_call: object, dict_call: dict[str, object] +): + sdk_tokens = token_counter(model="gpt-4o", messages=_tool_call_turn(sdk_call)) + + assert sdk_tokens == token_counter(model="gpt-4o", messages=_tool_call_turn(dict_call)) + assert sdk_tokens > 300, f"a reused SDK tool call must contribute its payload tokens, got {sdk_tokens}" + + @pytest.mark.parametrize( ("null_call", "empty_call"), ( @@ -471,6 +508,29 @@ def test_token_counter_counts_custom_tool_definition_description_and_grammar(): assert described - bare > 300, f"description and grammar must count, got {described} vs {bare} bare" +def test_token_counter_counts_responses_shaped_custom_tool_like_the_chat_shape(): + messages = [{"role": "user", "content": "run the code"}] + grammar = "start: " + " | ".join(f'"command_{index}"' for index in range(100)) + chat_tool = { + "type": "custom", + "custom": { + "name": "code_exec", + "description": "Runs one shell command", + "format": {"type": "grammar", "grammar": {"syntax": "lark", "definition": grammar}}, + }, + } + responses_tool = { + "type": "custom", + "name": "code_exec", + "description": "Runs one shell command", + "format": {"type": "grammar", "syntax": "lark", "definition": grammar}, + } + + assert token_counter(model="gpt-4o", messages=messages, tools=[responses_tool]) == token_counter( + model="gpt-4o", messages=messages, tools=[chat_tool] + ) + + @pytest.mark.parametrize( "message_count_pair", MESSAGES_TEXT, @@ -830,6 +890,7 @@ def test_token_counter_counts_anthropic_tool_choices_like_openai_choices( {**_anthropic_tool(_WEATHER_FUNCTION), "strict": True}, {**_anthropic_tool(_WEATHER_FUNCTION), "defer_loading": True, "input_examples": [{"location": "Paris"}]}, {**_anthropic_tool(_WEATHER_FUNCTION), "custom": {"eager_input_streaming": True}}, + {**_anthropic_tool(_WEATHER_FUNCTION), "type": "custom", "custom": {"eager_input_streaming": True}}, _anthropic_tool(_WEB_SEARCH_FUNCTION), ], ) 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 ca481e1b94f..b1697c1aba3 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 @@ -110,8 +110,24 @@ _SCHEMA: Final[dict[str, object]] = {"type": "object", "properties": {}} # muta "strict": False, }, ), + ( + { + "type": "custom", + "name": "lookup", + "description": "Find a record", + "input_schema": _SCHEMA, + "custom": {"eager_input_streaming": True}, + }, + { + "type": "function", + "name": "lookup", + "description": "Find a record", + "parameters": _SCHEMA, + "strict": False, + }, + ), ), - ids=("anthropic", "anthropic-strict", "anthropic-custom-type"), + ids=("anthropic", "anthropic-strict", "anthropic-custom-type", "anthropic-custom-options"), ) def test_transform_anthropic_tool_shapes(tool: dict[str, object], expected: dict[str, object]) -> None: result: Final = OpenAICountTokensConfig().transform_request_to_count_tokens( @@ -252,6 +268,30 @@ def test_messages_to_responses_input_maps_custom_tool_calls_and_their_outputs(): ] +@pytest.mark.parametrize( + ("tool_call", "expected"), + ( + ( + {"id": "call_1", "type": "custom", "custom": None}, + {"type": "custom_tool_call", "call_id": "call_1", "name": "", "input": ""}, + ), + ( + {"id": "call_1", "type": "function", "function": None}, + {"type": "function_call", "call_id": "call_1", "name": "", "arguments": ""}, + ), + ), + ids=("custom", "function"), +) +def test_messages_to_responses_input_maps_null_tool_call_payloads_to_empty_calls( + tool_call: dict[str, object], expected: dict[str, object] +) -> None: + input_items, _ = OpenAICountTokensConfig.messages_to_responses_input( + [{"role": "user", "content": "Run it"}, {"role": "assistant", "content": None, "tool_calls": [tool_call]}] + ) + + assert input_items[1:] == [expected] + + def test_transform_custom_tool_uses_the_flat_responses_shape(): tool: Final = { "type": "custom",