From 7acdf9f00d4abc6a72400bea80b93ce4dd7d7132 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 27 Oct 2026 19:33:38 -0700 Subject: [PATCH] fix(anthropic): count search results and reject unsupported token-count blocks (#48699) * fix(anthropic): count search results and reject unsupported token-count blocks Co-authored-by: Prasanna Jain * fix(anthropic): avoid a Final assignment inside the token-count loop Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Prasanna Jain --- litellm/litellm_core_utils/token_counter.py | 20 +++++-- .../proxy/anthropic_endpoints/endpoints.py | 5 ++ litellm/types/llms/anthropic.py | 9 ++++ .../test_count_tokens_error_shapes_wire.py | 54 +++++++++++++++++++ .../litellm_core_utils/test_token_counter.py | 28 ++++++++++ 5 files changed, 111 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index edb5ecf0f8d..6d50e4b065e 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -55,6 +55,7 @@ from litellm.types.llms.anthropic import ( AnthropicMessagesTextParam, AnthropicMessagesToolResultParam, AnthropicMessagesToolUseParam, + AnthropicSearchResultParam, ) from litellm.types.llms.openai import ( AllMessageValues, @@ -67,6 +68,10 @@ from litellm.types.llms.openai import ( from litellm.types.utils import Message, SelectTokenizerResponse +class UnsupportedContentTypeError(ValueError): + pass + + def get_modified_max_tokens( model: str, base_model: str, @@ -865,6 +870,7 @@ def _validate_anthropic_content(content: Mapping[str, object]) -> type: mapping: Final = { "tool_use": AnthropicMessagesToolUseParam, "tool_result": AnthropicMessagesToolResultParam, + "search_result": AnthropicSearchResultParam, } expected_cls: Final = mapping.get(content_type) if isinstance(content_type, str) else None @@ -1020,7 +1026,9 @@ def _count_anthropic_content( except Exception as e: if default_token_count is not None: return default_token_count - raise ValueError(f"Error counting field '{field_name}': {e}") + raise (UnsupportedContentTypeError if isinstance(e, UnsupportedContentTypeError) else ValueError)( + f"Error counting field '{field_name}': {e}" + ) return tokens @@ -1034,6 +1042,7 @@ def _count_content_list( | AnthropicMessagesTextParam | AnthropicMessagesImageParam | AnthropicMessagesDocumentParam + | AnthropicSearchResultParam ], use_default_image_token_count: bool, default_token_count: int | None, @@ -1068,7 +1077,7 @@ def _count_content_list( count_function, use_default_image_token_count, ) - elif c["type"] in ("tool_use", "tool_result"): + elif c["type"] in ("tool_use", "tool_result", "search_result"): num_tokens += _count_anthropic_content( c, count_function, @@ -1094,17 +1103,18 @@ def _count_content_list( num_tokens += count_function(tool_name) else: content_type = c.get("type", type(c).__name__) if isinstance(c, dict) else type(c).__name__ - raise ValueError( + raise UnsupportedContentTypeError( f"Invalid content item type: {content_type}. " f"Expected str or dict with 'type' field " f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, " - f"tool_reference)." + f"tool_reference, search_result)." ) return num_tokens except Exception as e: if default_token_count is not None: return default_token_count - raise ValueError( + error_type: Final = UnsupportedContentTypeError if isinstance(e, UnsupportedContentTypeError) else ValueError + raise error_type( f"Error getting number of tokens from content list: {e}, default_token_count={default_token_count}" ) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 3d2f89be15f..13a1844a6de 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -18,6 +18,7 @@ from litellm.anthropic_interface.exceptions import ( AnthropicExceptionMapping, ) from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError from litellm.litellm_core_utils.tool_fields import first_validation_error from litellm.llms.anthropic.messages.context_management import ( AnthropicContextManagementError, @@ -413,6 +414,10 @@ async def count_tokens( messages_format="anthropic", **_count_tokens_body_options(data, auto_drop_params), ) + except UnsupportedContentTypeError as e: + return _anthropic_error_json_response( + ProxyException(message=str(e), type="invalid_request_error", param="messages", code=400), request + ) except HTTPException as e: return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) except ProxyException as e: diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index ba422a0a901..de83e6ea66e 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -209,6 +209,13 @@ class AnthropicMessagesTextParam(TypedDict, total=False): prompt_cache_breakpoint: ReadOnly[PromptCacheBreakpoint] +class AnthropicSearchResultParam(TypedDict, total=False): + type: Required[ReadOnly[Literal["search_result"]]] + source: Required[ReadOnly[str]] + title: Required[ReadOnly[str]] + content: Required[ReadOnly[Iterable[AnthropicMessagesTextParam]]] + + class AnthropicMessagesToolUseParam(TypedDict, total=False): type: Required[Literal["tool_use"]] id: str @@ -357,6 +364,7 @@ class AnthropicMessagesToolResultParam(TypedDict, total=False): AnthropicMessagesToolResultContent | AnthropicMessagesImageParam | AnthropicMessagesDocumentParam + | AnthropicSearchResultParam | ToolReference ] ) @@ -368,6 +376,7 @@ AnthropicMessagesUserMessageValues = ( | AnthropicMessagesImageParam | AnthropicMessagesToolResultParam | AnthropicMessagesDocumentParam + | AnthropicSearchResultParam | AnthropicMessagesContainerUploadParam ) diff --git a/tests/integration/providers/test_count_tokens_error_shapes_wire.py b/tests/integration/providers/test_count_tokens_error_shapes_wire.py index 200f5e81673..a6c3a6dc62f 100644 --- a/tests/integration/providers/test_count_tokens_error_shapes_wire.py +++ b/tests/integration/providers/test_count_tokens_error_shapes_wire.py @@ -68,6 +68,60 @@ def _never_called(request: Request) -> Reply: raise AssertionError(f"count reached the provider without a key: {request.target}") +@pytest.mark.parametrize("nested", (False, True), ids=("user-content", "tool-result")) +def test_local_messages_count_accepts_search_result_text(gateway: Gateway, nested: bool) -> None: + search: Final[JsonValue] = { + "type": "search_result", + "source": "https://example.invalid/notes", + "title": "Project notes", + "content": [{"type": "text", "text": "The code word is 7f3a."}], + "citations": {"enabled": True}, + } + content: Final = [{"type": "tool_result", "tool_use_id": "toolu_1", "content": [search]}] if nested else [search] + text: Final[list[JsonValue]] = [ + {"type": "text", "text": "https://example.invalid/notes"}, + {"type": "text", "text": "Project notes"}, + {"type": "text", "text": "The code word is 7f3a."}, + ] + equivalent: Final = [{"type": "tool_result", "tool_use_id": "toolu_1", "content": text}] if nested else text + with gateway.scenario() as scenario: + model: Final = scenario.model(model="hosted_vllm/zai-org/GLM-5.3", api_base="http://127.0.0.1:9/v1") + counted: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": content}]} + ) + control: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": equivalent}]} + ) + assert counted.status_code == control.status_code == 200, (counted.text, control.text) + tokens: Final = _JSON_OBJECT.validate_json(counted.content)["input_tokens"] + assert tokens == _JSON_OBJECT.validate_json(control.content)["input_tokens"], (counted.text, control.text) + + +@pytest.mark.parametrize("nested", (False, True), ids=("user-content", "tool-result")) +@pytest.mark.parametrize( + ("block", "status", "error_type", "message"), + ( + ({"type": "future_block", "future_block": {"value": "ignore me"}}, 400, "invalid_request_error", "future_block"), + ({"type": "image"}, 500, "api_error", "source"), + ), + ids=("unknown-block", "unrelated-counting-error"), +) +def test_local_messages_count_classifies_only_unsupported_content_as_an_invalid_request( + gateway: Gateway, nested: bool, block: JsonValue, status: int, error_type: str, message: str +) -> None: + content: Final = [{"type": "tool_result", "tool_use_id": "toolu_1", "content": [block]}] if nested else [block] + with gateway.scenario() as scenario: + model: Final = scenario.model(model="hosted_vllm/zai-org/GLM-5.3", api_base="http://127.0.0.1:9/v1") + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": content}]} + ) + assert response.status_code == status, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + error: Final = payload["error"] + assert payload["type"] == "error" and isinstance(error, dict), payload + assert error["type"] == error_type and message in str(error["message"]), payload + + @pytest.mark.parametrize( ("deployment", "missing"), ( diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 9921bbe343a..92f971f64e9 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -1825,6 +1825,34 @@ def test_anthropic_image_block_with_empty_base64_data(): assert tokens > 0 +def test_token_counter_counts_anthropic_search_result_text() -> None: + from litellm.litellm_core_utils.token_counter import _count_content_list + + search_result: Final = { + "type": "search_result", + "source": "https://example.invalid/notes", + "title": "Project notes", + "content": [{"type": "text", "text": "The code word is 7f3a."}], + "citations": {"enabled": True}, + } + + assert _count_content_list(len, [search_result], False, None) == len( + "https://example.invalid/notesProject notesThe code word is 7f3a." + ) + + +@pytest.mark.parametrize("nested", (False, True)) +def test_unknown_content_type_remains_a_value_error_with_an_optional_fallback(nested: bool) -> None: + from litellm.litellm_core_utils.token_counter import _count_content_list + + unknown: Final = {"type": "future_block"} + content: Final = [{"type": "tool_result", "tool_use_id": "toolu_1", "content": [unknown]}] if nested else [unknown] + + with pytest.raises(ValueError, match="Invalid content item type: future_block"): + _count_content_list(len, content, False, None) + assert _count_content_list(len, content, False, 7) == 7 + + def test_anthropic_image_block_without_source_raises(): """An `image` block with no `source` raises, matching the OpenAI `image_url`-without-`url` behavior.""" from litellm.litellm_core_utils.token_counter import _count_content_list