mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(token-counter): count Anthropic server tool history (#51334)
* fix(token-counter): count Anthropic server tool history Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(token-counter): bound nested results and update proxy coverage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(token-counter): include readable web search snippets Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kenneth-berri
GitHub
parent
2845677822
commit
be623b5bde
5 changed files with 325 additions and 27 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue