Merge d655641b6f into e2d78b69cd

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(

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,

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

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