mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(bedrock): inject tool-search beta header for every Claude model on Invoke
This commit is contained in:
parent
ff4a40f017
commit
3849fff187
6 changed files with 214 additions and 90 deletions
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue