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",