[Feat] Add support for Tool Search on /messages API - Azure, Bedrock, Anthropic API

litellm/llms/anthropic/experimental_pass_through/messages/transformation.py

@@ -2,7 +2,8 @@
 
 import httpx
 
-from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj, verbose_logger
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.litellm_logging import verbose_logger
 from litellm.llms.base_llm.anthropic_messages.transformation import (
     BaseAnthropicMessagesConfig,
 )
@@ -13,9 +14,10 @@
 from litellm.types.llms.anthropic_messages.anthropic_response import (
     AnthropicMessagesResponse,
 )
+from litellm.types.llms.anthropic_tool_search import get_tool_search_beta_header
 from litellm.types.router import GenericLiteLLMParams
 
-from ...common_utils import AnthropicError
+from ...common_utils import AnthropicError, AnthropicModelInfo
 
 DEFAULT_ANTHROPIC_API_BASE = "https://api.anthropic.com"
 DEFAULT_ANTHROPIC_API_VERSION = "2023-06-01"
@@ -75,9 +77,9 @@ def validate_anthropic_messages_environment(
         if "content-type" not in headers:
             headers["content-type"] = "application/json"
 
-        headers = self._update_headers_with_optional_anthropic_beta(
+        headers = self._update_headers_with_anthropic_beta(
             headers=headers,
-            context_management=optional_params.get("context_management"),
+            optional_params=optional_params,
         )
 
         return headers, api_base
@@ -153,16 +155,44 @@ def get_async_streaming_response_iterator(
         )
 
     @staticmethod
-    def _update_headers_with_optional_anthropic_beta(
-        headers: dict, context_management: Optional[Dict]
+    def _update_headers_with_anthropic_beta(
+        headers: dict,
+        optional_params: dict,
+        custom_llm_provider: str = "anthropic",
     ) -> dict:
-        if context_management is None:
-            return headers
-
+        """
+        Auto-inject anthropic-beta headers based on features used.
+
+        Handles:
+        - context_management: adds 'context-management-2025-06-27'
+        - tool_search: adds provider-specific tool search header
+
+        Args:
+            headers: Request headers dict
+            optional_params: Optional parameters including tools, context_management
+            custom_llm_provider: Provider name for looking up correct tool search header
+        """
+        beta_values: set = set()
+
+        # Get existing beta headers if any
         existing_beta = headers.get("anthropic-beta")
-        beta_value = ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value
-        if existing_beta is None:
-            headers["anthropic-beta"] = beta_value
-        elif beta_value not in [beta.strip() for beta in existing_beta.split(",")]:
-            headers["anthropic-beta"] = f"{existing_beta}, {beta_value}"
+        if existing_beta:
+            beta_values.update(b.strip() for b in existing_beta.split(","))
+
+        # Check for context management
+        if optional_params.get("context_management") is not None:
+            beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value)
+
+        # Check for tool search tools
+        tools = optional_params.get("tools")
+        if tools:
+            anthropic_model_info = AnthropicModelInfo()
+            if anthropic_model_info.is_tool_search_used(tools):
+                # Use provider-specific tool search header
+                tool_search_header = get_tool_search_beta_header(custom_llm_provider)
+                beta_values.add(tool_search_header)
+
+        if beta_values:
+            headers["anthropic-beta"] = ",".join(sorted(beta_values))
+
         return headers

litellm/llms/azure_ai/anthropic/messages_transformation.py

@@ -62,10 +62,10 @@ def validate_anthropic_messages_environment(
         if "content-type" not in headers:
             headers["content-type"] = "application/json"
 
-        # Update headers with optional anthropic beta features
-        headers = self._update_headers_with_optional_anthropic_beta(
+        # Update headers with anthropic beta features (context management, tool search, etc.)
+        headers = self._update_headers_with_anthropic_beta(
             headers=headers,
-            context_management=optional_params.get("context_management"),
+            optional_params=optional_params,
         )
 
         return headers, api_base

litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py

@@ -129,6 +129,37 @@ def _remove_ttl_from_cache_control(
                                 if isinstance(cache_control, dict) and "ttl" in cache_control:
                                     cache_control.pop("ttl", None)
 
+    def _get_tool_search_beta_header_for_bedrock(
+        self,
+        model: str,
+        tool_search_used: bool,
+        programmatic_tool_calling_used: bool,
+        input_examples_used: bool,
+        beta_set: set,
+    ) -> None:
+        """
+        Adjust tool search beta header for Bedrock.
+
+        Bedrock requires a different beta header for tool search on Opus 4 models
+        when tool search is used without programmatic tool calling or input examples.
+
+        Note: On Amazon Bedrock, server-side tool search is only supported on Claude Opus 4
+        with the `tool-search-tool-2025-10-19` beta header.
+
+        Ref: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool
+
+        Args:
+            model: The model name
+            tool_search_used: Whether tool search is used
+            programmatic_tool_calling_used: Whether programmatic tool calling is used
+            input_examples_used: Whether input examples are used
+            beta_set: The set of beta headers to modify in-place
+        """
+        if tool_search_used and not (programmatic_tool_calling_used or input_examples_used):
+            beta_set.discard(ANTHROPIC_TOOL_SEARCH_BETA_HEADER)
+            if "opus-4" in model.lower() or "opus_4" in model.lower():
+                beta_set.add("tool-search-tool-2025-10-19")
+
     def transform_anthropic_messages_request(
         self,
         model: str,
@@ -189,13 +220,13 @@ def transform_anthropic_messages_request(
         )
         beta_set.update(auto_betas)
 
-        if (
-            tool_search_used
-            and not (programmatic_tool_calling_used or input_examples_used)
-        ):
-            beta_set.discard(ANTHROPIC_TOOL_SEARCH_BETA_HEADER)
-            if "opus-4" in model.lower() or "opus_4" in model.lower():
-                beta_set.add("tool-search-tool-2025-10-19")
+        self._get_tool_search_beta_header_for_bedrock(
+            model=model,
+            tool_search_used=tool_search_used,
+            programmatic_tool_calling_used=programmatic_tool_calling_used,
+            input_examples_used=input_examples_used,
+            beta_set=beta_set,
+        )
 
         if beta_set:
             anthropic_messages_request["anthropic_beta"] = list(beta_set)

litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py

@@ -1,11 +1,16 @@
 from typing import Any, Dict, List, Optional, Tuple
 
+from litellm.llms.anthropic.common_utils import AnthropicModelInfo
 from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
     AnthropicMessagesConfig,
 )
+from litellm.types.llms.anthropic import (
+    ANTHROPIC_BETA_HEADER_VALUES,
+    ANTHROPIC_HOSTED_TOOLS,
+)
+from litellm.types.llms.anthropic_tool_search import get_tool_search_beta_header
 from litellm.types.llms.vertex_ai import VertexPartnerProvider
 from litellm.types.router import GenericLiteLLMParams
-from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_HOSTED_TOOLS
 
 from ....vertex_llm_base import VertexBase
 
@@ -51,13 +56,28 @@ def validate_anthropic_messages_environment(
 
         headers["content-type"] = "application/json"
 
-        # Add web search beta header for Vertex AI only if not already set
-        if "anthropic-beta" not in headers:
-            tools = optional_params.get("tools", [])
-            for tool in tools:
-                if isinstance(tool, dict) and tool.get("type", "").startswith(ANTHROPIC_HOSTED_TOOLS.WEB_SEARCH.value):
-                    headers["anthropic-beta"] = ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value
-                    break
+        # Add beta headers for Vertex AI
+        tools = optional_params.get("tools", [])
+        beta_values: set[str] = set()
+
+        # Get existing beta headers if any
+        existing_beta = headers.get("anthropic-beta")
+        if existing_beta:
+            beta_values.update(b.strip() for b in existing_beta.split(","))
+
+        # Check for web search tool
+        for tool in tools:
+            if isinstance(tool, dict) and tool.get("type", "").startswith(ANTHROPIC_HOSTED_TOOLS.WEB_SEARCH.value):
+                beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value)
+                break
+
+        # Check for tool search tools - Vertex AI uses different beta header
+        anthropic_model_info = AnthropicModelInfo()
+        if anthropic_model_info.is_tool_search_used(tools):
+            beta_values.add(get_tool_search_beta_header("vertex_ai"))
+
+        if beta_values:
+            headers["anthropic-beta"] = ",".join(beta_values)
 
         return headers, api_base
 

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -1236,7 +1236,7 @@ def _create_pipeline_operations(
         return pipeline_operations
 
     def _get_total_tokens_from_usage(
-        self, usage: Any | None, rate_limit_type: Literal["output", "input", "total"]
+        self, usage: Optional[Any], rate_limit_type: Literal["output", "input", "total"]
     ) -> int:
         """
         Get total tokens from response usage for rate limiting.

litellm/types/llms/anthropic.py

@@ -636,8 +636,10 @@ class ANTHROPIC_BETA_HEADER_VALUES(str, Enum):
     ADVANCED_TOOL_USE_2025_11_20 = "advanced-tool-use-2025-11-20"
 
 
-# Tool search beta header constant
+# Tool search beta header constant (for Anthropic direct API and Microsoft Foundry)
 ANTHROPIC_TOOL_SEARCH_BETA_HEADER = "advanced-tool-use-2025-11-20"
 
 # Effort beta header constant
 ANTHROPIC_EFFORT_BETA_HEADER = "effort-2025-11-24"
+
+

litellm/types/llms/anthropic_tool_search.py

@@ -0,0 +1,36 @@
+"""
+Tool Search Beta Header Configuration
+
+Reference: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool
+"""
+
+from typing import Dict
+
+from litellm.types.utils import LlmProviders
+
+# Tool search beta header values
+TOOL_SEARCH_BETA_HEADER_ANTHROPIC = "advanced-tool-use-2025-11-20"
+TOOL_SEARCH_BETA_HEADER_VERTEX = "tool-search-tool-2025-10-19"
+TOOL_SEARCH_BETA_HEADER_BEDROCK = "tool-search-tool-2025-10-19"
+
+
+# Mapping of custom_llm_provider -> tool search beta header
+TOOL_SEARCH_BETA_HEADER_BY_PROVIDER: Dict[str, str] = {
+    LlmProviders.ANTHROPIC.value: TOOL_SEARCH_BETA_HEADER_ANTHROPIC,
+    LlmProviders.AZURE.value: TOOL_SEARCH_BETA_HEADER_ANTHROPIC,
+    LlmProviders.AZURE_AI.value: TOOL_SEARCH_BETA_HEADER_ANTHROPIC,
+    LlmProviders.VERTEX_AI.value: TOOL_SEARCH_BETA_HEADER_VERTEX,
+    LlmProviders.VERTEX_AI_BETA.value: TOOL_SEARCH_BETA_HEADER_VERTEX,
+    LlmProviders.BEDROCK.value: TOOL_SEARCH_BETA_HEADER_BEDROCK,
+}
+
+
+def get_tool_search_beta_header(custom_llm_provider: str) -> str:
+    """
+    Get the tool search beta header for a given provider.
+    """
+    return TOOL_SEARCH_BETA_HEADER_BY_PROVIDER.get(
+        custom_llm_provider,
+        TOOL_SEARCH_BETA_HEADER_ANTHROPIC
+    )
+