diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index ce022aea5b4..f19ea60fc76 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -1944,13 +1944,16 @@ def _count_content_list( tool_name = str(c.get("tool_name") or "") if tool_name: num_tokens += count_function(tool_name) + elif (change := as_str_mapping(c)) is not None and change.get("type") in ("tool_addition", "tool_removal"): + num_tokens += _tool_change_tokens(count_function, as_str_mapping(change.get("tool"))) else: content_type = c.get("type", type(c).__name__) if isinstance(c, dict) else type(c).__name__ raise UnsupportedContentTypeError( f"Invalid content item type: {content_type}. " f"Expected str or dict with 'type' field " f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, " - f"compaction, tool_reference, search_result, or Anthropic server tool blocks)." + f"compaction, tool_reference, tool_addition, tool_removal, search_result, " + f"or Anthropic server tool blocks)." ) return num_tokens except Exception as e: @@ -1962,6 +1965,24 @@ def _count_content_list( ) +def _tool_change_tokens(count_function: TokenCounterFunction, tool: Mapping[str, object] | None) -> int: + match tool: + case {"type": "tool_definition", "definition": definition} if ( + function := as_str_mapping(definition) + ) is not None: + return count_function(_format_function_definitions((function,))) + case {"type": "tool_reference", "name": str() as name}: + return count_function(name) + case {"type": "mcp_tool_reference", "server_name": str() as server_name, "name": str() as name}: + return count_function(server_name) + count_function(name) + case {"type": "mcp_toolset_reference", "server_name": str() as server_name}: + return count_function(server_name) + case _: + raise UnsupportedContentTypeError( + f"Invalid tool change target type: {None if tool is None else tool.get('type')}" + ) + + def _format_function_definitions(tools: Sequence[object]) -> str: """Formats tool definitions in the format that OpenAI appears to use. Based on https://github.com/forestwanglin/openai-java/blob/main/jtokkit/src/main/java/xyz/felh/openai/jtokkit/utils/TikTokenUtils.java diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 08ee9e765a2..7f42a39bc95 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -2,6 +2,7 @@ # This tests litellm.token_counter.token_counter() function import asyncio import base64 +import copy import importlib import io import json @@ -37,6 +38,8 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.token_counter import ( + UnsupportedContentTypeError, + _count_content_list, _encoding_count, _get_exact_count_function, _get_extrapolating_count_function, @@ -45,6 +48,7 @@ from litellm.litellm_core_utils.token_counter import ( estimate_image_tokens, high_detail_image_token_upper_bound, is_openai_model, + messages_reach_token_count, offload_token_count, ) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new @@ -2170,6 +2174,127 @@ def test_count_content_list_rejects_unknown_type(): assert "tool_reference" in message +_CLAUDE: Final = "anthropic/claude-sonnet-5-5" +_REMOVE_WEATHER: Final = {"type": "tool_removal", "tool": {"type": "tool_reference", "name": "get_weather"}} + + +def _weather_tool(description: str, *extra_properties: str) -> dict[str, object]: + return { + "name": "get_weather", + "description": description, + "input_schema": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + **{name: {"type": "string"} for name in extra_properties}, + }, + "required": ["city"], + }, + } + + +def _add_definition(definition: Mapping[str, object]) -> dict[str, object]: + return {"type": "tool_addition", "tool": {"type": "tool_definition", "definition": definition}} + + +def _conversation(*system_blocks: Mapping[str, object]) -> list[dict[str, object]]: + return [ + {"role": "user", "content": "What is the weather in Paris?"}, + {"role": "system", "content": [{"type": "text", "text": "Your tools changed."}, *system_blocks]}, + ] + + +def _then_remove_weather(messages: list[dict[str, object]]) -> list[dict[str, object]]: + return [*messages, {"role": "user", "content": "Never mind."}, {"role": "system", "content": [_REMOVE_WEATHER]}] + + +def test_tool_addition_counts_an_inline_definition_like_the_same_tool_declared_in_tools() -> None: + brief: Final = _weather_tool("Weather.") + detailed: Final = _weather_tool("Current weather, temperature and wind for a city.", "country", "date") + add_by_reference: Final = {"type": "tool_addition", "tool": {"type": "tool_reference", "name": "get_weather"}} + by_reference: Final = _conversation(add_by_reference) + + inline_delta: Final = token_counter_new(model=_CLAUDE, messages=_conversation(_add_definition(detailed))) - ( + token_counter_new(model=_CLAUDE, messages=_conversation(_add_definition(brief))) + ) + declared_delta: Final = token_counter_new(model=_CLAUDE, messages=by_reference, tools=[detailed]) - ( + token_counter_new(model=_CLAUDE, messages=by_reference, tools=[brief]) + ) + + assert declared_delta > 0 + assert inline_delta == declared_delta + + +@pytest.mark.parametrize("change", ["tool_addition", "tool_removal"]) +@pytest.mark.parametrize( + ("target", "referenced_names"), + [ + ({"type": "tool_reference", "name": "get_weather"}, ("get_weather",)), + ({"type": "mcp_tool_reference", "server_name": "calendar", "name": "list_events"}, ("calendar", "list_events")), + ({"type": "mcp_toolset_reference", "server_name": "calendar"}, ("calendar",)), + ], + ids=["tool_reference", "mcp_tool_reference", "mcp_toolset_reference"], +) +def test_tool_change_by_reference_counts_the_referenced_names( + change: str, target: dict[str, object], referenced_names: tuple[str, ...] +) -> None: + counted: Final = _count_content_list(len, [{"type": change, "tool": target}], False, None) + + assert counted == sum(map(len, referenced_names)) + + +def test_tool_removal_never_subtracts_an_earlier_inline_definition() -> None: + with_definition: Final = _conversation(_add_definition(_weather_tool("Current weather for a city."))) + without_definition: Final = _conversation() + + removal_with_definition: Final = token_counter_new( + model=_CLAUDE, messages=_then_remove_weather(with_definition) + ) - token_counter_new(model=_CLAUDE, messages=with_definition) + removal_without_definition: Final = token_counter_new( + model=_CLAUDE, messages=_then_remove_weather(without_definition) + ) - token_counter_new(model=_CLAUDE, messages=without_definition) + + assert removal_with_definition > 0 + assert removal_with_definition == removal_without_definition + + +def test_messages_reach_token_count_agrees_with_token_counter_across_tool_changes() -> None: + messages: Final = _then_remove_weather(_conversation(_add_definition(_weather_tool("Current weather for a city.")))) + total: Final = token_counter_new(model=_CLAUDE, messages=messages) + + reached: Final = ( + messages_reach_token_count(model=_CLAUDE, messages=messages, threshold=total), + messages_reach_token_count(model=_CLAUDE, messages=messages, threshold=total + 1), + ) + + assert reached == (True, False) + + +def test_counting_tool_changes_leaves_the_messages_unchanged() -> None: + messages: Final = _then_remove_weather(_conversation(_add_definition(_weather_tool("Current weather for a city.")))) + sent: Final = copy.deepcopy(messages) + + token_counter_new(model=_CLAUDE, messages=messages) + messages_reach_token_count(model=_CLAUDE, messages=messages, threshold=1) + + assert messages == sent + + +@pytest.mark.parametrize( + "tool_change", + [ + {"type": "tool_addition", "tool": {"type": "future_reference", "name": "get_weather"}}, + {"type": "tool_addition", "tool": {"type": "tool_definition", "definition": "get_weather"}}, + {"type": "tool_removal"}, + ], + ids=["unknown_target_type", "definition_not_an_object", "missing_target"], +) +def test_tool_change_with_an_unrecognized_target_stays_unsupported(tool_change: dict[str, object]) -> None: + with pytest.raises(UnsupportedContentTypeError, match="Invalid tool change target type"): + _count_content_list(len, [tool_change], False, None) + + @pytest.mark.parametrize( "source", [ diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index f2b52817a17..d4bdd0ee496 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -1682,6 +1682,75 @@ async def test_anthropic_count_tokens_on_a_claude_deployment_drops_a_malformed_t assert sent["messages"][2]["content"][0]["content"] == [{"type": "text", "text": "build ok"}] +def _weather_tool(description: str) -> dict[str, object]: + return { + "name": "get_weather", + "description": description, + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } + + +@pytest.mark.asyncio +async def test_anthropic_count_tokens_falls_back_to_a_local_estimate_that_counts_mid_conversation_tool_changes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "claude-framed", + "litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": "sk-ant-test"}, + "model_info": {"token_count_framing": 100}, + } + ] + ), + ) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + brief: Final = _weather_tool("Weather.") + detailed: Final = _weather_tool("Current weather, temperature and wind for a city in any country.") + + def conversation(tool: Mapping[str, object]) -> list[dict[str, object]]: + return [ + {"role": "user", "content": "What is the weather in Paris?"}, + {"role": "system", "content": [{"type": "tool_addition", "tool": tool}]}, + {"role": "user", "content": "Never mind."}, + { + "role": "system", + "content": [{"type": "tool_removal", "tool": {"type": "tool_reference", "name": "get_weather"}}], + }, + ] + + async def estimate(tool: Mapping[str, object], tools: list[dict[str, object]] | None) -> TokenCountResponse: + return await count_request_tokens( + request=TokenCountRequest(model="claude-framed", messages=conversation(tool), tools=tools), + call_endpoint=True, + messages_format="anthropic", + ) + + by_reference: Final = {"type": "tool_reference", "name": "get_weather"} + with respx.mock: + route: Final = respx.post("https://api.anthropic.com/v1/messages/count_tokens").mock( + return_value=httpx.Response( + 404, json={"type": "error", "error": {"type": "not_found_error", "message": "count unavailable"}} + ) + ) + inline_brief: Final = await estimate({"type": "tool_definition", "definition": brief}, None) + inline_detailed: Final = await estimate({"type": "tool_definition", "definition": detailed}, None) + declared_brief: Final = await estimate(by_reference, [brief]) + declared_detailed: Final = await estimate(by_reference, [detailed]) + + inline_delta: Final = inline_detailed.total_tokens - inline_brief.total_tokens + declared_delta: Final = declared_detailed.total_tokens - declared_brief.total_tokens + assert route.called + assert {inline_brief.count_source, inline_detailed.count_source} == {"local_estimate"} + assert declared_delta > 0 + assert inline_delta == declared_delta + + class _ScriptedCounter(BaseTokenCounter): def __init__(self, replies: dict[str, TokenCountResponse], burst: int = 1) -> None: self._replies = replies