diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 6d50e4b065e..f6e0ed660d2 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -453,6 +453,131 @@ class _MessageCountParams: self.tool_section = chat_template_tool_section(model) +def _tool_name_from_mapping(tool: Mapping[str, object]) -> str | None: + name: Final = tool.get("name") + if isinstance(name, str) and name: + return name + fn: Final = tool.get("function") + if isinstance(fn, Mapping): + fn_name: Final = fn.get("name") + if isinstance(fn_name, str) and fn_name: + return fn_name + defn: Final = tool.get("definition") + if isinstance(defn, Mapping): + defn_name: Final = defn.get("name") + if isinstance(defn_name, str) and defn_name: + return defn_name + return None + + +def _extract_added_tool(block: Mapping[str, object]) -> Mapping[str, object] | None: + tool_val: Final = block.get("tool") + if isinstance(tool_val, Mapping): + defn: Final = tool_val.get("definition") + if isinstance(defn, Mapping): + return defn + if tool_val.get("name") or tool_val.get("function"): + return tool_val + block_defn: Final = block.get("definition") + if isinstance(block_defn, Mapping): + return block_defn + if block.get("name"): + return block + return None + + +def _extract_removed_tool_name(block: Mapping[str, object]) -> str | None: + tool_name: Final = block.get("tool_name") + if isinstance(tool_name, str) and tool_name: + return tool_name + name: Final = block.get("name") + if isinstance(name, str) and name: + return name + tool_val: Final = block.get("tool") + if isinstance(tool_val, str) and tool_val: + return tool_val + if isinstance(tool_val, Mapping): + inner_name: Final = tool_val.get("name") + if isinstance(inner_name, str) and inner_name: + return inner_name + inner_tool_name: Final = tool_val.get("tool_name") + if isinstance(inner_tool_name, str) and inner_tool_name: + return inner_tool_name + inner_defn: Final = tool_val.get("definition") + if isinstance(inner_defn, Mapping): + inner_defn_name: Final = inner_defn.get("name") + if isinstance(inner_defn_name, str) and inner_defn_name: + return inner_defn_name + return None + + +def apply_mid_conversation_tool_changes( + messages: Sequence[AllMessageValues] | None, + tools: Sequence[object] | None, +) -> tuple[list[AllMessageValues] | None, list[object] | None]: + """Apply Anthropic mid-conversation tool changes (tool_addition and tool_removal blocks) + in message order to derive the effective tool set and strip the event blocks from messages. + """ + if messages is None: + return None, list(tools) if tools else None + + # Fast-path check: do any messages contain tool_addition or tool_removal blocks? + has_tool_changes = False + for m in messages: + c = m.get("content") if isinstance(m, Mapping) else getattr(m, "content", None) + if isinstance(c, list): + for b in c: + if isinstance(b, Mapping) and b.get("type") in ("tool_addition", "tool_removal"): + has_tool_changes = True + break + if has_tool_changes: + break + + if not has_tool_changes: + return list(messages), list(tools) if tools else None + + effective_tools: dict[str, object] = {} + if tools: + for t in tools: + if isinstance(t, Mapping): + t_name = _tool_name_from_mapping(t) + if t_name: + effective_tools[t_name] = t + else: + effective_tools[f"__unnamed_{id(t)}"] = t + else: + effective_tools[f"__unnamed_{id(t)}"] = t + + cleaned_messages: list[AllMessageValues] = [] + for m in messages: + c = m.get("content") if isinstance(m, Mapping) else getattr(m, "content", None) + if isinstance(c, list): + new_content = [] + for b in c: + if isinstance(b, Mapping): + btype = b.get("type") + if btype == "tool_addition": + added_tool = _extract_added_tool(b) + if added_tool is not None: + added_name = _tool_name_from_mapping(added_tool) + if added_name: + effective_tools[added_name] = added_tool + continue + elif btype == "tool_removal": + removed_name = _extract_removed_tool_name(b) + if removed_name and removed_name in effective_tools: + del effective_tools[removed_name] + continue + new_content.append(b) + new_m = dict(m) if isinstance(m, Mapping) else dict(m.__dict__) + new_m["content"] = new_content + cleaned_messages.append(cast(AllMessageValues, new_m)) + else: + cleaned_messages.append(m) + + res_tools = list(effective_tools.values()) if effective_tools else None + return cleaned_messages, res_tools + def token_counter( model="", custom_tokenizer: dict | SelectTokenizerResponse | None = None, @@ -510,12 +635,13 @@ def token_counter( elif messages is not None: new_messages: Final = cast(list[AllMessageValues], convert_list_message_to_dict(messages)) + cleaned_messages, updated_tools = apply_mid_conversation_tool_changes(new_messages, tools) params: Final = _MessageCountParams(model, custom_tokenizer) - num_tokens = _count_messages(params, new_messages, use_default_image_token_count, default_token_count) + num_tokens = _count_messages(params, cleaned_messages, use_default_image_token_count, default_token_count) if count_response_tokens is False: - includes_system_message: Final = any(message.get("role", None) == "system" for message in new_messages) + includes_system_message: Final = any(message.get("role", None) == "system" for message in cleaned_messages) num_tokens += _count_extra( - params.count_function, tools, tool_choice, includes_system_message, params.tool_section + params.count_function, updated_tools, tool_choice, includes_system_message, params.tool_section ) else: @@ -544,14 +670,15 @@ def messages_reach_token_count( new_messages: Final = cast( # cast-ok: convert_list_message_to_dict is untyped, same as token_counter list[AllMessageValues], convert_list_message_to_dict(messages) ) + cleaned_messages, updated_tools = apply_mid_conversation_tool_changes(new_messages, tools) params: Final = _MessageCountParams(model, None) - includes_system_message: Final = any(message.get("role", None) == "system" for message in new_messages) + includes_system_message: Final = any(message.get("role", None) == "system" for message in cleaned_messages) per_message_counts: Final = ( - _count_messages(params, [message], use_default_image_token_count, None) for message in new_messages + _count_messages(params, [message], use_default_image_token_count, None) for message in cleaned_messages ) running_totals: Final = accumulate( per_message_counts, - initial=_count_extra(params.count_function, tools, None, includes_system_message, params.tool_section), + initial=_count_extra(params.count_function, updated_tools, None, includes_system_message, params.tool_section), ) return any(total >= threshold for total in running_totals) @@ -1101,6 +1228,10 @@ def _count_content_list( tool_name = str(c.get("tool_name") or "") if tool_name: num_tokens += count_function(tool_name) + elif c["type"] in ("tool_addition", "tool_removal"): + # Mid-conversation tool changes: tool schemas are counted via `tools`, + # so these event blocks contribute 0 additional content tokens. + pass else: content_type = c.get("type", type(c).__name__) if isinstance(c, dict) else type(c).__name__ raise UnsupportedContentTypeError( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 04f84f1f7b5..dc7e61e1bc6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -100,7 +100,7 @@ from litellm.litellm_core_utils.litellm_logging import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.litellm_core_utils.token_counter import apply_mid_conversation_tool_changes, offload_token_count from litellm.llms.bedrock.count_tokens.kimi_k3 import ( HUGGING_FACE_SOURCE, KimiK3Encodings, @@ -15439,10 +15439,15 @@ async def count_request_tokens( counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None ) - framing: Final = settings.framing if counted_messages is not None and settings.framing is not None else 0 - kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None + resolved_messages, resolved_tools = ( + apply_mid_conversation_tool_changes(counted_messages, counted_tools) + if counted_messages is not None + else (counted_messages, counted_tools) + ) + framing: Final = settings.framing if resolved_messages is not None and settings.framing is not None else 0 + kimi_k3_messages: Final = resolved_messages if settings.tokenizer == "kimi_k3" else None exact: Final = ( - await _kimi_k3_count(request, messages_format, model_to_use, kimi_k3_messages, counted_tools) + await _kimi_k3_count(request, messages_format, model_to_use, kimi_k3_messages, resolved_tools) if kimi_k3_messages is not None and custom_tokenizer is None else None ) @@ -15456,14 +15461,14 @@ async def count_request_tokens( ) estimate: Final = await _local_token_estimate( model=model_to_use, - prompt=prompt if counted_messages is None else None, - messages=counted_messages, - tools=counted_tools, + prompt=prompt if resolved_messages is None else None, + messages=resolved_messages, + tools=resolved_tools, custom_tokenizer=custom_tokenizer, provider_failure=ladder, ) template: Final = ( - 0 if kimi_k3_messages is None else kimi_k3_template_tokens(len(kimi_k3_messages), bool(counted_tools)) + 0 if kimi_k3_messages is None else kimi_k3_template_tokens(len(kimi_k3_messages), bool(resolved_tools)) ) return TokenCountResponse( total_tokens=estimate.total_tokens + template + framing, diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 92f971f64e9..88c82e7fa87 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -2469,3 +2469,109 @@ def test_get_modified_max_tokens_leaves_a_gemini_4_long_decode_budget_alone(monk ) == 200_000 ) + + +def test_token_counter_anthropic_tool_addition_and_removal() -> None: + """Anthropic mid-conversation tool changes: tool_addition derives the effective tool set + and tool_removal drops tools, with the tool schemas counted via the tool definitions.""" + from litellm.litellm_core_utils.token_counter import apply_mid_conversation_tool_changes + + tool_def = { + "name": "get_weather", + "description": "Get current weather", + "input_schema": { + "type": "object", + "properties": {"location": {"type": "string", "description": "City name"}}, + "required": ["location"], + }, + } + + messages_with_addition = [ + { + "role": "system", + "content": [ + { + "type": "tool_addition", + "tool": { + "type": "tool_definition", + "definition": tool_def, + }, + } + ], + }, + {"role": "user", "content": "What is the weather in Paris?"}, + ] + + cleaned, effective_tools = apply_mid_conversation_tool_changes(messages_with_addition, None) + assert cleaned == [{"role": "system", "content": []}, {"role": "user", "content": "What is the weather in Paris?"}] + assert effective_tools == [tool_def] + + # Count via token_counter with mid-conversation addition + count_from_addition = token_counter_new(model="claude-3-5-sonnet-20241022", messages=messages_with_addition) + + # Direct equivalent: empty system message and tools passed explicitly + direct_count = token_counter_new( + model="claude-3-5-sonnet-20241022", + messages=[{"role": "system", "content": []}, {"role": "user", "content": "What is the weather in Paris?"}], + tools=[tool_def], + ) + assert count_from_addition == direct_count + + # Now add tool_removal + messages_with_removal = [ + *messages_with_addition, + { + "role": "system", + "content": [ + { + "type": "tool_removal", + "tool": {"type": "tool_reference", "name": "get_weather"}, + } + ], + }, + {"role": "user", "content": "Thanks!"}, + ] + + cleaned_rem, effective_tools_rem = apply_mid_conversation_tool_changes(messages_with_removal, None) + assert effective_tools_rem is None + count_with_removal = token_counter_new(model="claude-3-5-sonnet-20241022", messages=messages_with_removal) + direct_count_no_tools = token_counter_new( + model="claude-3-5-sonnet-20241022", + messages=[ + {"role": "system", "content": []}, + {"role": "user", "content": "What is the weather in Paris?"}, + {"role": "system", "content": []}, + {"role": "user", "content": "Thanks!"}, + ], + tools=None, + ) + assert count_with_removal == direct_count_no_tools + + +def test_token_counter_content_list_with_tool_addition_does_not_raise() -> None: + """_count_content_list should safely ignore tool_addition and tool_removal without raising ValueError.""" + from litellm.litellm_core_utils.token_counter import _count_content_list + + def count_fn(s: str) -> int: + return len(s.split()) + + tokens = _count_content_list( + count_function=count_fn, + content_list=[ + {"type": "text", "text": "Hello world"}, + { + "type": "tool_addition", + "tool": { + "type": "tool_definition", + "definition": {"name": "test_tool", "input_schema": {"type": "object"}}, + }, + }, + { + "type": "tool_removal", + "tool": {"type": "tool_reference", "name": "test_tool"}, + }, + ], + use_default_image_token_count=False, + default_token_count=None, + ) + assert tokens == 2 diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index ebe3c7407d8..77f9af19484 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -3897,3 +3897,50 @@ async def test_responses_count_on_claude_counts_the_turns_previous_response_id_r ] assert [sorted(message) for message in messages if isinstance(message, dict)] == [["content", "role"]] * 3 spend_logs.assert_awaited_once_with("chatcmpl-stored", ("chatcmpl-stored", "resp_client_form")) + + +@pytest.mark.asyncio +async def test_anthropic_count_tokens_local_fallback_handles_tool_addition_and_removal(): + """When Anthropic provider token count is unavailable, the local estimate should handle + tool_addition and tool_removal blocks without raising ValueError / 500, and include the added tool schemas.""" + import litellm.proxy.proxy_server as proxy_server + + tool_def = { + "name": "get_weather", + "description": "Get current weather", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}}, + } + messages = [ + { + "role": "system", + "content": [ + { + "type": "tool_addition", + "tool": { + "type": "tool_definition", + "definition": tool_def, + }, + } + ], + }, + {"role": "user", "content": "What is the weather in Paris?"}, + ] + + response = await proxy_server.count_request_tokens( + request=TokenCountRequest(model="claude-3-5-sonnet-20241022", messages=messages), + call_endpoint=False, + messages_format="anthropic", + ) + + assert response.count_source == "local_estimate" + assert response.total_tokens > 0 + + text_only_response = await proxy_server.count_request_tokens( + request=TokenCountRequest( + model="claude-3-5-sonnet-20241022", + messages=[{"role": "system", "content": []}, {"role": "user", "content": "What is the weather in Paris?"}], + ), + call_endpoint=False, + messages_format="anthropic", + ) + assert response.total_tokens > text_only_response.total_tokens