diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 800a0939cec..fc84cd4f807 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -47,11 +47,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.core_helpers import as_str_mapping, has_attributes 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: @@ -80,7 +83,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 class UnsupportedContentTypeError(ValueError): @@ -890,10 +893,11 @@ 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 (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 if key == "function_call": if not isinstance(value, Mapping): @@ -902,6 +906,24 @@ def _count_function_call_tokens( raise ValueError(f"Unexpected key {key!r}; expected 'tool_calls' or 'function_call'") +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", "")) + + def _count_messages( params: _MessageCountParams, messages: list[AllMessageValues], @@ -1801,6 +1823,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 "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): for declaration in declarations: @@ -1825,6 +1851,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 a45a0448326..4af64bb6969 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, replayable_responses_input -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") 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") or {}) + 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) @@ -249,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) @@ -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 697ac76b4e5..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,6 +391,146 @@ 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: 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_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=_tool_call_turn( + {"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": _CUSTOM_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( + ("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"), + ( + ( + {"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" + + +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, @@ -744,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 575183772f9..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( @@ -227,6 +243,79 @@ 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"}, + ] + + +@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", + "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.