fix(bedrock): inject tool-search beta header for every Claude model on Invoke

This commit is contained in:
mateo-berri 2026-07-15 16:58:01 -07:00
parent ff4a40f017
commit 3849fff187
6 changed files with 214 additions and 90 deletions

View file

@ -22,8 +22,8 @@ from litellm.llms.bedrock.common_utils import (
normalize_tool_input_schema_types_for_bedrock_invoke,
pop_bedrock_invoke_output_config_format,
remove_custom_field_from_tools,
swap_tool_search_beta_header_for_bedrock,
)
from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponse
from litellm.utils import _supports_factory
@ -273,10 +273,12 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
)
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")
beta_set = swap_tool_search_beta_header_for_bedrock(
beta_set=beta_set,
tool_search_used=tool_search_used,
programmatic_tool_calling_used=programmatic_tool_calling_used,
input_examples_used=input_examples_used,
)
auto_beta_list = filter_and_transform_beta_headers(
beta_headers=list(beta_set - user_beta_set),

View file

@ -36,6 +36,8 @@ from litellm.llms.base_llm.anthropic_messages.transformation import (
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret
from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER
from litellm.types.llms.anthropic_tool_search import TOOL_SEARCH_BETA_HEADER_BEDROCK
if TYPE_CHECKING:
from litellm.types.llms.openai import AllMessageValues
@ -1274,6 +1276,32 @@ class BedrockEventStreamDecoderBase:
return chunk.decode() # type: ignore[no-any-return]
def swap_tool_search_beta_header_for_bedrock(
beta_set: set[str],
tool_search_used: bool,
programmatic_tool_calling_used: bool,
input_examples_used: bool,
) -> set[str]:
"""
Replace the Anthropic-direct tool-search beta header with Bedrock's
``tool-search-tool-2025-10-19`` whenever a tool-search tool is present in
the request (and programmatic tool calling / input examples, which use the
Anthropic-direct header instead, are not).
Bedrock accepts the ``tool_search_tool_*_20251119`` tool types on every
live Claude model once this header is attached; without it the request
400s with "Input tag 'tool_search_tool_regex_20251119' found using 'type'
does not match any of the expected tags". There is deliberately no
per-model gate here: earlier hardcoded pattern lists went stale as AWS
rolled out support (Haiku 4.5 and Opus 4.7 were missing), and a model
that genuinely lacks support rejects the tool type with the same 400
whether or not the header is attached.
"""
if not tool_search_used or programmatic_tool_calling_used or input_examples_used:
return set(beta_set)
return (set(beta_set) - {ANTHROPIC_TOOL_SEARCH_BETA_HEADER}) | {TOOL_SEARCH_BETA_HEADER_BEDROCK}
def get_anthropic_beta_from_headers(headers: dict) -> List[str]:
"""
Extract anthropic-beta header values and convert them to a list.

View file

@ -45,11 +45,9 @@ from litellm.llms.bedrock.common_utils import (
normalize_tool_input_schema_types_for_bedrock_invoke,
pop_bedrock_invoke_output_config_format,
remove_custom_field_from_tools,
swap_tool_search_beta_header_for_bedrock,
)
from litellm.types.llms.anthropic import (
ANTHROPIC_BETA_HEADER_VALUES,
ANTHROPIC_TOOL_SEARCH_BETA_HEADER,
)
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
from litellm.types.llms.bedrock import BedrockInvokeAnthropicMessagesRequest
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
@ -437,84 +435,6 @@ class AmazonAnthropicClaudeMessagesConfig(
"""
return is_claude_4_5_on_bedrock(model)
def _supports_tool_search_on_bedrock(self, model: str) -> bool:
"""
Check if the model supports tool search on Bedrock.
On Amazon Bedrock, server-side tool search is supported on Claude Opus 4.5
and Claude Sonnet 4.5 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
Returns:
True if the model supports tool search on Bedrock
"""
model_lower = model.lower()
# Supported models for tool search on Bedrock
supported_patterns = [
# Opus 4.5
"opus-4.5",
"opus_4.5",
"opus-4-5",
"opus_4_5",
# Sonnet 4.5
"sonnet-4.5",
"sonnet_4.5",
"sonnet-4-5",
"sonnet_4_5",
# Opus 4.6
"opus-4.6",
"opus_4.6",
"opus-4-6",
"opus_4_6",
# sonnet 4.6
"sonnet-4.6",
"sonnet_4.6",
"sonnet-4-6",
"sonnet_4_6",
# NOTE: Opus 4.7 on Bedrock does not support server-side tool search
# as of launch (2026-04-16). Bedrock rejects the tool type with:
# "tool type 'tool_search_tool_..._20251119' is not supported for this model".
# Re-add the opus-4.7 patterns here once AWS announces support.
]
return any(pattern in model_lower for pattern in supported_patterns)
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 self._supports_tool_search_on_bedrock(model):
beta_set.add("tool-search-tool-2025-10-19")
# Bedrock-InvokeModel-supported ``context_management.edits`` types and the
# ``anthropic-beta`` header that each one requires. ``clear_thinking_20251015``
# is intentionally absent — it is LiteLLM-internal, consumed via
@ -612,12 +532,11 @@ class AmazonAnthropicClaudeMessagesConfig(
beta_set=beta_set,
)
self._get_tool_search_beta_header_for_bedrock(
model=model,
beta_set = swap_tool_search_beta_header_for_bedrock(
beta_set=beta_set,
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 "tool-search-tool-2025-10-19" in beta_set:

View file

@ -545,3 +545,42 @@ def test_output_format_removed_from_bedrock_invoke_request():
assert (
"output_format" not in result
), f"output_format should be removed for Bedrock Invoke, got keys: {result.keys()}"
@pytest.mark.parametrize(
"model",
[
"us.anthropic.claude-haiku-4-5-20251001-v1:0",
"us.anthropic.claude-sonnet-4-6",
"us.anthropic.claude-opus-4-7",
],
)
def test_transform_request_injects_tool_search_beta_for_every_claude_model(model):
"""
Regression: the tool-search beta header used to be injected only for
models matching an ``opus-4`` name pattern. Haiku/Sonnet requests carrying
a ``tool_search_tool_regex_20251119`` tool were forwarded without the
header and Bedrock rejected the tool type with a 400.
"""
config = AmazonAnthropicClaudeConfig()
result = config.transform_request(
model=model,
messages=[{"role": "user", "content": "Hello"}],
optional_params={
"max_tokens": 64,
"tools": [
{
"type": "tool_search_tool_regex_20251119",
"name": "tool_search_tool_regex",
}
],
},
litellm_params={},
headers={},
)
betas = result.get("anthropic_beta") or []
assert "tool-search-tool-2025-10-19" in betas, (
f"tool-search beta header must be injected for {model}; got {betas}"
)

View file

@ -1497,6 +1497,86 @@ def test_bedrock_messages_renames_user_provided_aliased_beta_header():
)
_TOOL_SEARCH_REQUEST_TOOLS = [
{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"},
{
"name": "add_numbers",
"description": "Add two integers",
"input_schema": {
"type": "object",
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
"required": ["a", "b"],
},
},
]
@pytest.mark.parametrize(
"model",
[
"us.anthropic.claude-haiku-4-5-20251001-v1:0",
"us.anthropic.claude-sonnet-4-6",
"us.anthropic.claude-opus-4-7",
],
)
def test_bedrock_messages_auto_injects_tool_search_beta_for_every_claude_model(model):
"""
Regression: the tool-search beta header used to be auto-injected only for a
hardcoded list of models (Opus/Sonnet 4.5/4.6). Requests carrying a
`tool_search_tool_regex_20251119` tool on any other Claude model (Haiku 4.5,
Opus 4.7) were forwarded without the header and Bedrock rejected them with
"Input tag 'tool_search_tool_regex_20251119' found using 'type' does not
match any of the expected tags". The header must be attached whenever a
tool-search tool is present, for every model.
"""
from litellm.types.router import GenericLiteLLMParams
cfg = AmazonAnthropicClaudeMessagesConfig()
result = cfg.transform_anthropic_messages_request(
model=model,
messages=[{"role": "user", "content": [{"type": "text", "text": "Hello"}]}],
anthropic_messages_optional_request_params={
"max_tokens": 64,
"tools": [dict(t) for t in _TOOL_SEARCH_REQUEST_TOOLS],
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
betas = result.get("anthropic_beta") or []
assert "tool-search-tool-2025-10-19" in betas, (
f"tool-search beta header must be auto-injected for {model}; got {betas}"
)
tool_types = [t.get("type") for t in result.get("tools") or []]
assert "tool_search_tool_regex_20251119" in tool_types, (
"the tool_search tool itself must survive the transformation"
)
def test_bedrock_messages_no_tool_search_beta_without_tool_search_tool():
"""The tool-search beta header must only appear when a tool-search tool is
actually in the request; injecting it unconditionally would advertise a
capability the request doesn't use."""
from litellm.types.router import GenericLiteLLMParams
cfg = AmazonAnthropicClaudeMessagesConfig()
result = cfg.transform_anthropic_messages_request(
model="us.anthropic.claude-haiku-4-5-20251001-v1:0",
messages=[{"role": "user", "content": [{"type": "text", "text": "Hello"}]}],
anthropic_messages_optional_request_params={
"max_tokens": 64,
"tools": [dict(_TOOL_SEARCH_REQUEST_TOOLS[1])],
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
betas = result.get("anthropic_beta") or []
assert "tool-search-tool-2025-10-19" not in betas
@pytest.mark.asyncio
async def test_promote_message_stop_usage_preserves_message_delta_output_tokens():
"""

View file

@ -473,3 +473,59 @@ def test_capability_lookups_fall_back_to_base_model_when_regional_entry_lacks_fi
assert is_claude_4_5_on_bedrock(regional) is True
assert bedrock_converse_supports_parallel_tool_use_config(regional) is True
class TestSwapToolSearchBetaHeaderForBedrock:
def test_swaps_anthropic_header_for_bedrock_header(self):
from litellm.llms.bedrock.common_utils import (
swap_tool_search_beta_header_for_bedrock,
)
result = swap_tool_search_beta_header_for_bedrock(
beta_set={"advanced-tool-use-2025-11-20", "context-1m-2025-08-07"},
tool_search_used=True,
programmatic_tool_calling_used=False,
input_examples_used=False,
)
assert result == {"tool-search-tool-2025-10-19", "context-1m-2025-08-07"}
def test_adds_bedrock_header_even_without_anthropic_header(self):
from litellm.llms.bedrock.common_utils import (
swap_tool_search_beta_header_for_bedrock,
)
result = swap_tool_search_beta_header_for_bedrock(
beta_set=set(),
tool_search_used=True,
programmatic_tool_calling_used=False,
input_examples_used=False,
)
assert result == {"tool-search-tool-2025-10-19"}
@pytest.mark.parametrize(
"tool_search_used, programmatic_tool_calling_used, input_examples_used",
[
(False, False, False),
(True, True, False),
(True, False, True),
],
)
def test_leaves_beta_set_unchanged_outside_the_plain_tool_search_case(
self, tool_search_used, programmatic_tool_calling_used, input_examples_used
):
from litellm.llms.bedrock.common_utils import (
swap_tool_search_beta_header_for_bedrock,
)
original = {"advanced-tool-use-2025-11-20"}
result = swap_tool_search_beta_header_for_bedrock(
beta_set=original,
tool_search_used=tool_search_used,
programmatic_tool_calling_used=programmatic_tool_calling_used,
input_examples_used=input_examples_used,
)
assert result == {"advanced-tool-use-2025-11-20"}
assert result is not original