diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index e56dd83a277..ce022aea5b4 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -31,6 +31,7 @@ from litellm.constants import ( DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_TOKEN_COUNT, DEFAULT_IMAGE_WIDTH, + DEFAULT_MAX_RECURSE_DEPTH, MAX_IMAGE_URL_DOWNLOAD_SIZE_MB, MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES, MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES, @@ -1758,31 +1759,72 @@ def _count_file_tokens( ) +_ANTHROPIC_SERVER_TOOL_FIELDS: Final = MappingProxyType( + { + "server_tool_use": ("name", "input"), + "mcp_tool_use": ("name", "input"), + "web_search_tool_result": ("content",), + "web_fetch_tool_result": ("content",), + "code_execution_tool_result": ("content",), + "bash_code_execution_tool_result": ("content",), + "text_editor_code_execution_tool_result": ("content",), + "mcp_tool_result": ("content",), + "web_search_result": ("url", "title", "snippet"), + "web_fetch_result": ("url", "content"), + "code_execution_result": ("stdout", "stderr", "content"), + "bash_code_execution_result": ("stdout", "stderr", "content"), + "text_editor_code_execution_view_result": ("content",), + "text_editor_code_execution_str_replace_result": ("lines",), + "text_editor_code_execution_create_result": (), + "encrypted_code_execution_result": ("stderr", "content"), + "code_execution_output": (), + "bash_code_execution_output": (), + "web_search_tool_result_error": ("error_code",), + "web_fetch_tool_result_error": ("error_code",), + "code_execution_tool_result_error": ("error_code",), + "bash_code_execution_tool_result_error": ("error_code",), + "text_editor_code_execution_tool_result_error": ("error_code", "error_message"), + } +) +_ANTHROPIC_DOCUMENT: Final = TypeAdapter(AnthropicMessagesDocumentParam) + + def _count_anthropic_content( content: Mapping[str, object], count_function: TokenCounterFunction, use_default_image_token_count: bool, default_token_count: int | None, pdf_pricing: PdfPagePricing = "claude", + *, + _content_depth: int = 0, ) -> int: """ Count tokens in Anthropic-specific content blocks (tool_use, tool_result, etc.). - Uses TypedDict definitions from litellm.types.llms.anthropic to determine - what fields to count and how to handle nested structures. - - Dynamically infers which fields to count based on the TypedDict definition, - avoiding hardcoded field names. + Uses readable server tool fields or the client tool TypedDict definitions. """ - typeddict_cls: Final = _validate_anthropic_content(content) - type_hints: Final[Mapping[str, object]] = getattr(typeddict_cls, "__annotations__", {}) + if _content_depth >= DEFAULT_MAX_RECURSE_DEPTH: + raise UnsupportedContentTypeError("Anthropic tool content nesting exceeds the token counter limit") + content_type: Final = content.get("type") + if content_type == "document": + return _count_document_tokens( + _ANTHROPIC_DOCUMENT.validate_python(content), + count_function, + use_default_image_token_count, + default_token_count, + pdf_pricing, + ) + server_fields: Final = _ANTHROPIC_SERVER_TOOL_FIELDS.get(content_type) if isinstance(content_type, str) else None + fields: Final[Mapping[str, object]] = ( + dict.fromkeys(server_fields) + if server_fields is not None + else getattr(_validate_anthropic_content(content), "__annotations__", {}) + ) tokens = 0 - # Fields to skip (metadata/identifiers that don't contribute to prompt tokens) skip_fields: Final = {"type", "id", "tool_use_id", "cache_control", "is_error"} - # Iterate over all fields defined in the TypedDict - for field_name, field_type in type_hints.items(): + for field_name in fields: if field_name in skip_fields: continue @@ -1801,7 +1843,21 @@ def _count_anthropic_content( pdf_pricing, ) elif isinstance(field_value, dict): - tokens += count_function(str(field_value)) + if ( + server_fields is not None + and field_name == "content" + and (field_mapping := as_str_mapping(content[field_name])) is not None + ): + tokens += _count_anthropic_content( + field_mapping, + count_function, + use_default_image_token_count, + default_token_count, + pdf_pricing, + _content_depth=_content_depth + 1, + ) + else: + tokens += count_function(str(field_value)) except Exception as e: if default_token_count is not None: return default_token_count @@ -1867,7 +1923,9 @@ def _count_content_list( use_default_image_token_count, pdf_pricing, ) - elif c["type"] in ("tool_use", "tool_result", "search_result"): + elif ( + c["type"] in ("tool_use", "tool_result", "search_result") or c["type"] in _ANTHROPIC_SERVER_TOOL_FIELDS + ): num_tokens += _count_anthropic_content( c, count_function, @@ -1892,7 +1950,7 @@ def _count_content_list( 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"compaction, tool_reference, search_result)." + f"compaction, tool_reference, search_result, or Anthropic server tool blocks)." ) return num_tokens except Exception as e: diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index b6044a68547..9ca064e308a 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -11,6 +11,7 @@ IGNORE_FUNCTIONS = [ "_json_cost", # bounded at depth 32 and consumed under byte/node limits. "detached", # router_config: copies only the plain containers of the router's own constructor arguments, once per build or changeset "_format_type", + "_count_anthropic_content", # Nested server result mappings are bounded by DEFAULT_MAX_RECURSE_DEPTH. "remove_additional_properties", "remove_strict_from_schema", "filter_schema_fields", diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index d0b837646b4..08ee9e765a2 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -2304,6 +2304,185 @@ def test_anthropic_image_block_with_empty_base64_data(): assert tokens > 0 +@pytest.mark.parametrize( + "block, expected", + ( + ( + {"type": "tool_result", "tool_use_id": "toolu_1", "content": {"status": "done"}}, + "{'status': 'done'}", + ), + ( + {"type": "server_tool_use", "id": "srv_1", "name": "web_search", "input": {"query": "release"}}, + "web_search{'query': 'release'}", + ), + ( + { + "type": "mcp_tool_use", + "id": "mcp_1", + "server_name": "notes", + "name": "lookup", + "input": {"q": "release"}, + }, + "lookup{'q': 'release'}", + ), + ( + { + "type": "web_search_tool_result", + "tool_use_id": "srv_1", + "content": [ + { + "type": "web_search_result", + "url": "https://example.invalid", + "title": "Release", + "snippet": "Streaming added.", + "encrypted_content": "opaque" * 1000, + "page_age": "yesterday", + } + ], + }, + "https://example.invalidReleaseStreaming added.", + ), + ( + { + "type": "web_fetch_tool_result", + "tool_use_id": "srv_1", + "content": { + "type": "web_fetch_result", + "url": "https://example.invalid", + "retrieved_at": "metadata", + "content": { + "type": "document", + "title": "Release", + "source": {"type": "text", "media_type": "text/plain", "data": "Streaming added."}, + }, + }, + }, + "https://example.invalidReleaseStreaming added.", + ), + ( + { + "type": "code_execution_tool_result", + "tool_use_id": "srv_1", + "content": { + "type": "code_execution_result", + "stdout": "42", + "stderr": "warning", + "return_code": 0, + "content": [{"type": "code_execution_output", "file_id": "opaque-file-id"}], + }, + }, + "42warning", + ), + ( + { + "type": "code_execution_tool_result", + "tool_use_id": "srv_1", + "content": { + "type": "encrypted_code_execution_result", + "encrypted_stdout": "opaque" * 1000, + "stderr": "warning", + "return_code": 0, + "content": [], + }, + }, + "warning", + ), + ( + { + "type": "mcp_tool_result", + "tool_use_id": "mcp_1", + "is_error": False, + "content": [{"type": "text", "text": "Streaming added."}], + }, + "Streaming added.", + ), + ( + { + "type": "bash_code_execution_tool_result", + "tool_use_id": "srv_1", + "content": { + "type": "bash_code_execution_result", + "stdout": "42", + "stderr": "warning", + "return_code": 0, + "content": [{"type": "bash_code_execution_output", "file_id": "opaque-file-id"}], + }, + }, + "42warning", + ), + ( + { + "type": "text_editor_code_execution_tool_result", + "tool_use_id": "srv_1", + "content": { + "type": "text_editor_code_execution_view_result", + "content": "print(42)", + "file_type": "text", + "num_lines": 1, + }, + }, + "print(42)", + ), + ( + { + "type": "text_editor_code_execution_tool_result", + "tool_use_id": "srv_1", + "content": { + "type": "text_editor_code_execution_str_replace_result", + "lines": ["first", "second"], + "old_lines": 1, + "new_lines": 2, + }, + }, + "firstsecond", + ), + ( + { + "type": "text_editor_code_execution_tool_result", + "tool_use_id": "srv_1", + "content": {"type": "text_editor_code_execution_create_result", "is_file_update": False}, + }, + "", + ), + ), +) +def test_token_counter_counts_anthropic_tool_readable_content(block: Mapping[str, object], expected: str) -> None: + from litellm.litellm_core_utils.token_counter import _count_content_list + + assert _count_content_list(len, [block], False, None) == len(expected) + + +def test_token_counter_bounds_nested_server_results_and_honors_the_fallback() -> None: + from functools import reduce + + from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH + from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError, _count_content_list + + def wrap(content: Mapping[str, object], _: int) -> Mapping[str, object]: + return {"type": "web_fetch_result", "content": content} + + leaf: Final[Mapping[str, object]] = {"type": "web_fetch_result", "url": "https://example.invalid"} + block: Final = reduce(wrap, range(DEFAULT_MAX_RECURSE_DEPTH + 1), leaf) + + with pytest.raises(UnsupportedContentTypeError, match="nesting exceeds the token counter limit"): + _count_content_list(len, [block], False, None) + assert _count_content_list(len, [block], False, 123) == 123 + + +@pytest.mark.parametrize( + "tool", ("web_search", "web_fetch", "code_execution", "bash_code_execution", "text_editor_code_execution") +) +def test_token_counter_counts_server_tool_errors(tool: str) -> None: + from litellm.litellm_core_utils.token_counter import _count_content_list + + block: Final = { + "type": f"{tool}_tool_result", + "tool_use_id": "srv_1", + "content": {"type": f"{tool}_tool_result_error", "error_code": "unavailable"}, + } + assert _count_content_list(len, [block], False, None) == len("unavailable") + + def test_token_counter_counts_anthropic_search_result_text() -> None: from litellm.litellm_core_utils.token_counter import _count_content_list @@ -2334,12 +2513,12 @@ def test_token_counter_counts_a_search_result_without_a_title() -> None: ) -@pytest.mark.parametrize("nested", (False, True)) -def test_unknown_content_type_remains_a_value_error_with_an_optional_fallback(nested: bool) -> None: +@pytest.mark.parametrize("container", (None, "tool_result", "web_search_tool_result")) +def test_unknown_content_type_remains_a_value_error_with_an_optional_fallback(container: str | None) -> 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] + content: Final = [{"type": container, "tool_use_id": "toolu_1", "content": [unknown]}] if container else [unknown] with pytest.raises(ValueError, match="Invalid content item type: future_block"): _count_content_list(len, content, False, None) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index bbd3ce8c78d..f2b52817a17 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2119,12 +2119,17 @@ async def test_framing_keeps_the_400_for_blocks_the_provider_counter_cannot_conv @pytest.mark.asyncio -@pytest.mark.parametrize( - "unreadable_body", - [ - pytest.param( - { - "messages": [ +async def test_framed_estimate_counts_server_tool_history_when_the_provider_counter_rejects_it( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _framed_kimi(monkeypatch) + _available_kimi_k3(monkeypatch, tmp_path) + + async def count(title: str) -> TokenCountResponse: + return await count_request_tokens( + request=TokenCountRequest( + model="kimi", + messages=[ {"role": "user", "content": "What is new?"}, { "role": "assistant", @@ -2135,13 +2140,31 @@ async def test_framing_keeps_the_400_for_blocks_the_provider_counter_cannot_conv "name": "web_search", "input": {"query": "new"}, }, - {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_01", "content": []}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + {"type": "web_search_result", "url": "https://example.invalid", "title": title} + ], + }, ], }, - ] - }, - id="web_search_blocks", - ), + ], + ), + call_endpoint=True, + ) + + short: Final = await count("Release") + long: Final = await count("Release notes " * 100) + + assert (short.count_source, long.count_source) == ("local_estimate", "local_estimate") + assert long.total_tokens > short.total_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "unreadable_body", + [ pytest.param( {"messages": [], "contents": [{"role": "user", "parts": [{"text": "What is new?"}]}]}, id="google_contents", diff --git a/tests/unit/test_router/core/test_construction_precall_characterization_precall.py b/tests/unit/test_router/core/test_construction_precall_characterization_precall.py index c9e04117f4d..ab53f01662f 100644 --- a/tests/unit/test_router/core/test_construction_precall_characterization_precall.py +++ b/tests/unit/test_router/core/test_construction_precall_characterization_precall.py @@ -168,6 +168,43 @@ async def test_select_without_messages_or_input_runs_no_context_window_check(mak assert await _picked(_small_and_big(make_core)) == {"small", "big"} +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("title", "snippet"), + [("release notes " * 100, ""), ("Release", "release notes " * 100)], + ids=["title", "snippet"], +) +async def test_server_tool_history_does_not_bypass_the_context_window_check( + make_core: MakeCore, title: str, snippet: str +) -> None: + core: Final = make_core(model_list=[windowed("g", "server-tools", 50)], enable_pre_call_checks=True) + messages: Final = [ + {"role": "user", "content": "Search release notes"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srv_1", "name": "web_search", "input": {"query": "release notes"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srv_1", + "content": [ + { + "type": "web_search_result", + "url": "https://example.invalid", + "title": title, + "snippet": snippet, + "encrypted_content": "opaque", + } + ], + }, + ], + }, + {"role": "user", "content": "Summarize"}, + ] + with pytest.raises(litellm.ContextWindowExceededError, match="Context Window exceeded"): + await core.select("g", messages=messages) + + @pytest.mark.asyncio async def test_when_no_context_window_fits_select_raises_context_window_exceeded(make_core: MakeCore) -> None: core: Final = make_core(model_list=[windowed("g", "small", 5)], enable_pre_call_checks=True)