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.