diff --git a/litellm/llms/bedrock/count_tokens/kimi_k3.py b/litellm/llms/bedrock/count_tokens/kimi_k3.py index 5eb46944f9a..28c3d1ebe20 100644 --- a/litellm/llms/bedrock/count_tokens/kimi_k3.py +++ b/litellm/llms/bedrock/count_tokens/kimi_k3.py @@ -422,7 +422,7 @@ def _rendered(request: _ConverseRequest) -> _Segments: ) -def _text_tokens(encoding: tiktoken.Encoding, text: str) -> int: +def kimi_k3_text_tokens(encoding: tiktoken.Encoding, text: str) -> int: spans: Final = (match.span() for match in _CHUNK.finditer(text)) return sum(len(encoding.encode_ordinary(text[start:end])) for start, end in spans) @@ -481,7 +481,7 @@ def kimi_k3_tokens( if request is None: return None return sum( - 1 if isinstance(segment, _Control) else _text_tokens(encoding, segment) for segment in _rendered(request) + 1 if isinstance(segment, _Control) else kimi_k3_text_tokens(encoding, segment) for segment in _rendered(request) ) diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py index 4e3059cfe26..cb9c0ea9334 100644 --- a/litellm/proxy/common_utils/token_count_settings.py +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -3,13 +3,13 @@ from dataclasses import dataclass, replace from types import MappingProxyType from typing import Final, Literal -from pydantic import BaseModel, ConfigDict, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import ProxyException -from litellm.types.token_count import TokenCountProvider, TokenCountSourceFields, TokenCountTokenizer +from litellm.types.token_count import LocalTokenizer, TokenCountProvider, TokenCountSourceFields, TokenCountTokenizer from litellm.types.utils import TokenCountResponse CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404}) @@ -51,11 +51,17 @@ class TokenCountModel: configured: bool = False +@dataclass(frozen=True, slots=True) +class UnknownLocalTokenizer: + value: str + + @dataclass(frozen=True, slots=True) class TokenCountSettings: framing: int | None count_model: TokenCountModel | None tokenizer: TokenCountTokenizer | None = None + local_tokenizer: LocalTokenizer | UnknownLocalTokenizer | None = None tool_framing: int | None = None responses_framing: int | None = None responses_tool_framing: int | None = None @@ -197,6 +203,18 @@ def _tokenizer(value: object) -> TokenCountTokenizer | None: return None +_LOCAL_TOKENIZER: Final[TypeAdapter[LocalTokenizer]] = TypeAdapter(LocalTokenizer) + + +def _local_tokenizer(value: object) -> LocalTokenizer | UnknownLocalTokenizer | None: + if value is None or (isinstance(value, str) and not value.strip()): + return None + try: + return _LOCAL_TOKENIZER.validate_python(value, strict=True) + except ValidationError: + return UnknownLocalTokenizer(value=str(value)) + + def resolve_token_count_settings( configured: Mapping[str, object], default: Mapping[str, object], deployment_region: Callable[[], str] ) -> TokenCountSettings: @@ -206,6 +224,7 @@ def resolve_token_count_settings( configured.get("token_count_model"), default.get("token_count_model"), deployment_region ), tokenizer=_tokenizer(default.get("token_count_tokenizer")), + local_tokenizer=_local_tokenizer(configured.get("token_count_tokenizer")), tool_framing=_configured_framing(configured, default, "token_count_tool_framing"), responses_framing=_configured_framing(configured, default, "token_count_responses_framing"), responses_tool_framing=_configured_framing( @@ -224,5 +243,7 @@ def count_source_fields(response: TokenCountResponse) -> TokenCountSourceFields: return {} case "token_count_model": return {"litellm_count_source": "token_count_model", "litellm_count_model": response.model_used} + case "local_estimate" if response.count_tokenizer is not None: + return {"litellm_count_source": "local_estimate", "litellm_count_tokenizer": response.count_tokenizer} case source: return {"litellm_count_source": source} diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 698c1e6e752..41437d2253c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -58,6 +58,7 @@ from typing import ( import anyio import anyio.lowlevel +import tiktoken import websockets import websockets.exceptions from openai.types.audio import TranscriptionStreamEvent @@ -126,6 +127,7 @@ from litellm.llms.bedrock.count_tokens.kimi_k3 import ( KimiK3Encodings, kimi_k3_anthropic_tokens, kimi_k3_template_tokens, + kimi_k3_text_tokens, kimi_k3_tokens, ) from litellm.proxy._types import ( @@ -183,6 +185,7 @@ from litellm.proxy.common_utils.token_count_settings import ( TOKENIZER_LOAD_WAIT_SECONDS, ProviderCountFailure, TokenCountSettings, + UnknownLocalTokenizer, cost_map_entry, resolve_token_count_settings, ) @@ -206,6 +209,7 @@ from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( ) from litellm.rust_bridge.catalog import Target from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.token_count import LocalTokenizer from litellm.types.utils import ( PRICING_OVERRIDES_KEY, CallTypes, @@ -15084,6 +15088,7 @@ from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository from litellm.responses.litellm_completion_transformation.count_request import ( ChatCountBridgeError, + ChatCountRequest, chat_count_request, with_provider_file_ids, ) @@ -15366,7 +15371,7 @@ def _provider_error_after_failed_estimate( provider_failure: ProviderCountFailure, error: BaseException ) -> ProxyException: verbose_proxy_logger.warning( - "Local token estimate failed after the provider's count error %s, answering that error: %s", + "Local token estimate failed, answering count error %s instead: %s", provider_failure.status_code, type(error).__name__, ) @@ -15527,6 +15532,22 @@ async def _tokenizer_or_provider_error( raise _provider_error_after_failed_estimate(provider_failure, tokenizer) from tokenizer +def _tokenizer_count( + template_model: str, + prompt: str | None, + messages: Sequence[AllMessageValues] | None, + tools: Sequence[ChatCompletionToolParam] | None, + tokenizer: SelectTokenizerResponse, +) -> Awaitable[int]: + return offload_token_count(litellm.token_counter)( + model=template_model, + text=prompt, + messages=messages, + tools=None if tools is None else list(tools), + custom_tokenizer=tokenizer, + ) + + async def _local_token_estimate( model: str, template_model: str, @@ -15541,13 +15562,7 @@ async def _local_token_estimate( if provider_failure is None else await _tokenizer_or_provider_error(model, custom_tokenizer, provider_failure) ) - count: Final = offload_token_count(litellm.token_counter)( - model=template_model, - text=prompt, - messages=messages, - tools=tools, - custom_tokenizer=tokenizer, - ) + count: Final = _tokenizer_count(template_model, prompt, messages, tools, tokenizer) total_tokens: Final = await ( count if provider_failure is None else _estimate_or_provider_error(count, provider_failure) ) @@ -15569,6 +15584,17 @@ async def _kimi_k3_count( encoding: Final = await _KIMI_K3_ENCODINGS.encoding(TOKENIZER_LOAD_WAIT_SECONDS) if encoding is None: return None + return await _kimi_k3_encoded_count(encoding, request, messages_format, model, messages, tools) + + +async def _kimi_k3_encoded_count( + encoding: tiktoken.Encoding, + request: TokenCountRequest, + messages_format: TokenCountMessagesFormat, + model: str, + messages: Sequence[AllMessageValues], + tools: Sequence[ChatCompletionToolParam] | None, +) -> int | None: match messages_format: case "anthropic": return await offload_token_count(kimi_k3_anthropic_tokens)( @@ -15578,6 +15604,174 @@ async def _kimi_k3_count( return await offload_token_count(kimi_k3_tokens)(encoding, model, messages, tools) +def _unknown_tokenizer_failure(named: UnknownLocalTokenizer, request_model: str) -> ProviderCountFailure: + return ProviderCountFailure( + status_code=400, + message=f"{request_model} sets token_count_tokenizer to {named.value!r}, " + f"which is not one of {', '.join(get_args(LocalTokenizer))}", + rejected_before_provider_call=True, + ) + + +def _local_tokenizer_failure(tokenizer: str, model: str, reason: str) -> ProviderCountFailure: + return ProviderCountFailure( + status_code=400, + message=f"{model} counts tokens with the {tokenizer} tokenizer set in token_count_tokenizer, which {reason}", + rejected_before_provider_call=True, + ) + + +async def _count_or_failure( + count: Awaitable[int | None], unreadable: ProviderCountFailure +) -> int | ProviderCountFailure: + try: + counted: Final = await count + except BaseException as e: + if not _answers_provider_error(e): + raise + verbose_proxy_logger.warning("Local token count failed, answering 400 instead: %s", type(e).__name__) + return unreadable + return unreadable if counted is None else counted + + +async def _kimi_k3_local_tokens( + prompt: str | None, + count_messages: Callable[[tiktoken.Encoding], Awaitable[int | None]], + unreadable: ProviderCountFailure, + unavailable: ProviderCountFailure, +) -> int | ProviderCountFailure: + encoding: Final = await _KIMI_K3_ENCODINGS.encoding(TOKENIZER_LOAD_WAIT_SECONDS) + if encoding is None: + return unavailable + return await _count_or_failure( + offload_token_count(kimi_k3_text_tokens)(encoding, prompt) if prompt is not None else count_messages(encoding), + unreadable, + ) + + +async def _local_tokenizer_tokens( + tokenizer: LocalTokenizer, + request_model: str, + template_model: str, + prompt: str | None, + messages: Sequence[AllMessageValues] | None, + tools: Sequence[ChatCompletionToolParam] | None, + kimi_k3_count: Callable[[tiktoken.Encoding], Awaitable[int | None]], +) -> _LocalTokenEstimate | ProviderCountFailure: + unreadable: Final = _local_tokenizer_failure(tokenizer, request_model, "could not count this request") + unavailable: Final = _local_tokenizer_failure(tokenizer, request_model, "could not be loaded") + if tokenizer == "kimi_k3": + kimi_k3: Final = await _kimi_k3_local_tokens(prompt, kimi_k3_count, unreadable, unavailable) + return ( + kimi_k3 + if isinstance(kimi_k3, ProviderCountFailure) + else _LocalTokenEstimate(total_tokens=kimi_k3, tokenizer_type="huggingface_tokenizer") + ) + selected: Final = await asyncify(litellm.utils.select_named_tokenizer)(tokenizer) + if selected is None: + return unavailable + total_tokens: Final = await _count_or_failure( + _tokenizer_count(template_model, prompt, messages, tools, selected), unreadable + ) + return ( + total_tokens + if isinstance(total_tokens, ProviderCountFailure) + else _LocalTokenEstimate(total_tokens=total_tokens, tokenizer_type=selected["type"]) + ) + + +@dataclass(frozen=True, slots=True) +class _LocalTokenizerRequest: + request: TokenCountRequest + messages_format: TokenCountMessagesFormat + model: str + template_model: str + prompt: str | None + messages: Sequence[AllMessageValues] | None + tools: Sequence[ChatCompletionToolParam] | None + bridged: ChatCountRequest | None + responses_api: bool + contents_estimable: bool + stored_turns_unseen: bool + + +async def _local_tokenizer_count( + tokenizer: LocalTokenizer, counted: _LocalTokenizerRequest, settings: TokenCountSettings +) -> TokenCountResponse | ProviderCountFailure: + request: Final = counted.request + if counted.stored_turns_unseen: + return _local_tokenizer_failure( + tokenizer, request.model, "cannot see the earlier turns previous_response_id points to on the provider" + ) + if not counted.contents_estimable: + return _local_tokenizer_failure(tokenizer, request.model, "cannot read these contents") + bridged: Final = counted.bridged + messages: Final = ( + counted.messages + if bridged is None + else cast( # cast-ok: the chat body Responses inference sends, chat-shaped dicts token_counter normalizes + Sequence[AllMessageValues], [dict(message) for message in bridged.messages] + ) + ) + tools: Final = ( + counted.tools + if bridged is None + else cast( # cast-ok: the chat tools Responses inference sends + Sequence[ChatCompletionToolParam] | None, [dict(tool) for tool in bridged.tools or ()] or None + ) + ) + estimate: Final = await _local_tokenizer_tokens( + tokenizer, + request.model, + counted.template_model, + counted.prompt, + messages, + tools, + kimi_k3_count=partial( + _kimi_k3_encoded_count, + request=request, + messages_format=counted.messages_format if bridged is None else "openai", + model=counted.model, + messages=messages or (), + tools=tools, + ), + ) + if isinstance(estimate, ProviderCountFailure): + return estimate + framing: Final = ( + 0 + if messages is None + else settings.hidden_prompt_tokens( + carries_tools=sends_function_tools(tools), responses_api=counted.responses_api + ) + ) + return TokenCountResponse( + total_tokens=estimate.total_tokens + framing, + request_model=request.model, + model_used=counted.model, + tokenizer_type=estimate.tokenizer_type, + count_source="local_estimate", + count_tokenizer=tokenizer, + ) + + +async def _count_with_local_tokenizer( + named: LocalTokenizer | UnknownLocalTokenizer, counted: _LocalTokenizerRequest, settings: TokenCountSettings +) -> TokenCountResponse: + """Count with the tokenizer the deployment names in token_count_tokenizer, never asking the provider. + + A Responses request bridged to chat is counted as the chat body inference sends, previous_response_id turns + included.""" + outcome: Final = ( + _unknown_tokenizer_failure(named, counted.request.model) + if isinstance(named, UnknownLocalTokenizer) + else await _local_tokenizer_count(named, counted, settings) + ) + if isinstance(outcome, ProviderCountFailure): + raise outcome.to_proxy_exception() + return outcome + + class _TokenCountDeployment(TypedDict, total=False): litellm_params: ReadOnly[Mapping[str, object]] model_info: ReadOnly[Mapping[str, object]] @@ -15982,6 +16176,7 @@ async def _count_request_tokens( default=await asyncify(cost_map_entry)(*_token_count_settings_model(deployment_view, request.model)), deployment_region=partial(_deployment_aws_region, deployment_view), ) + local_tokenizer: Final = settings.local_tokenizer # Try provider-specific token counting first - only for non-direct requests (from provider endpoints) provider_counter: BaseTokenCounter | None = None @@ -16045,12 +16240,16 @@ async def _count_request_tokens( routed_request: Final = ( None if responses_request is None else with_provider_file_ids(responses_request, routed_deployment) ) + counting_provider: Final = provider_counter if local_tokenizer is None else None chat_count: Final = ( None - if routed_request is None or provider_counter is None or isinstance(provider_counter, OpenAITokenCounter) + if routed_request is None + or (local_tokenizer is None and (provider_counter is None or isinstance(provider_counter, OpenAITokenCounter))) else await chat_count_request( model_to_use, - custom_llm_provider, + custom_llm_provider + if local_tokenizer is None + else custom_llm_provider or _deployment_provider(deployment_view, request.model, routed_request), _deployment_api_base(deployment_view, routed_request), routed_request, _count_drop_policy(routed_deployment, routed_request), @@ -16060,10 +16259,10 @@ async def _count_request_tokens( raise ProxyException(message=chat_count.message, type="invalid_request_error", param=None, code=400) ladder: Final = ( None - if provider_counter is None + if counting_provider is None else await _provider_token_count( count_with=_native_responses_count( - provider_counter, + counting_provider, custom_llm_provider, routed_request, routed_deployment, @@ -16071,7 +16270,7 @@ async def _count_request_tokens( ) or partial( _try_provider_token_count, - provider_counter, + counting_provider, custom_llm_provider, messages=( [dict(message) for message in chat_count.messages] @@ -16134,12 +16333,11 @@ async def _count_request_tokens( local_system: Final = system if counted_view is None else counted_view.system # Default LiteLLM token counting - custom_tokenizer: CustomHuggingfaceTokenizer | None = None - if model_info is not None: - custom_tokenizer = cast( - CustomHuggingfaceTokenizer | None, - model_info.get("custom_tokenizer", None), - ) + custom_tokenizer: Final = ( + None + if model_info is None + else cast(CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None)) + ) system_message: Final = _system_message(local_system) prompt_turn: Final = prompt_only and not (system_message is None and not tools) typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes @@ -16174,6 +16372,26 @@ async def _count_request_tokens( carries_tools=sends_function_tools(counted_tools), responses_api=native_responses ) ) + if local_tokenizer is not None: + return await _count_with_local_tokenizer( + local_tokenizer, + _LocalTokenizerRequest( + request=request + if counted_view is None + else request.model_copy(update={"messages": local_messages, "system": local_system}), + messages_format=messages_format, + model=model_to_use, + template_model=template_model, + prompt=prompt if counted_messages is None else None, + messages=counted_messages, + tools=counted_tools, + bridged=chat_count, + responses_api=native_responses, + contents_estimable=contents_estimable, + stored_turns_unseen=chat_count is None and bool(_responses_routing_kwargs(routed_request)), + ), + settings, + ) kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None exact: Final = ( await _kimi_k3_count( diff --git a/litellm/types/token_count.py b/litellm/types/token_count.py index 90ccefea177..f89bcbfaa35 100644 --- a/litellm/types/token_count.py +++ b/litellm/types/token_count.py @@ -4,6 +4,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict TokenCountSource: TypeAlias = Literal["provider", "token_count_model", "local_estimate"] TokenCountTokenizer: TypeAlias = Literal["kimi_k3"] +LocalTokenizer: TypeAlias = Literal["openai", "llama3", "deepseek", "cohere", "kimi_k3"] TokenCountProvider: TypeAlias = Literal["bedrock", "bedrock_mantle"] @@ -16,3 +17,4 @@ class TokenCountModelInfo(TypedDict): class TokenCountSourceFields(TypedDict, total=False): litellm_count_source: ReadOnly[TokenCountSource] litellm_count_model: ReadOnly[str] + litellm_count_tokenizer: ReadOnly[LocalTokenizer] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ee80b63f93b..9bcc84a7760 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -67,7 +67,7 @@ from litellm.types.llms.base import ( LiteLLMPydanticObjectBase, ) from litellm.types.mcp import MCPServerCostInfo -from litellm.types.token_count import TokenCountModelInfo, TokenCountSource, TokenCountTokenizer +from litellm.types.token_count import LocalTokenizer, TokenCountModelInfo, TokenCountSource, TokenCountTokenizer from litellm.types.workload_identity import ( ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS, @@ -4876,6 +4876,7 @@ class TokenCountResponse(LiteLLMPydanticObjectBase): model_used: str tokenizer_type: str count_source: TokenCountSource | None = None + count_tokenizer: LocalTokenizer | None = Field(default=None, exclude_if=lambda tokenizer: tokenizer is None) original_response: dict | None = None """ Original Response from upstream API call - if an API call was made for token counting diff --git a/litellm/utils.py b/litellm/utils.py index ec7bb041fb5..0047de0e177 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2997,7 +2997,26 @@ def _return_huggingface_tokenizer( model: str, outcomes: _HuggingFaceTokenizerLoadOutcomes = _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES ) -> SelectTokenizerResponse | _HuggingFaceTokenizerLoadFailure | None: kind: Final = huggingface_tokenizer_kind(model) - if kind is None or (kind not in _BUNDLED_TOKENIZER_LOAD_LOCKS and is_offline_mode()): + return None if kind is None else _huggingface_tokenizer_of_kind(kind, model, outcomes) + + +def select_named_tokenizer( + name: Literal["openai"] | HuggingFaceTokenizerKind, + outcomes: _HuggingFaceTokenizerLoadOutcomes = _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, +) -> SelectTokenizerResponse | None: + """The tokenizer `name` stands for, never another one: None when it cannot be loaded.""" + if name == "openai": + return _return_openai_tokenizer(name) + if litellm.disable_hf_tokenizer_download is True and name not in _BUNDLED_TOKENIZER_LOAD_LOCKS: + return None + loaded: Final = _huggingface_tokenizer_of_kind(name, name, outcomes) + return None if loaded is None or isinstance(loaded, _HuggingFaceTokenizerLoadFailure) else loaded + + +def _huggingface_tokenizer_of_kind( + kind: HuggingFaceTokenizerKind, model: str, outcomes: _HuggingFaceTokenizerLoadOutcomes +) -> SelectTokenizerResponse | _HuggingFaceTokenizerLoadFailure | None: + if kind not in _BUNDLED_TOKENIZER_LOAD_LOCKS and is_offline_mode(): return None key: Final = (kind, _huggingface_tokenizer_backend()) with _BUNDLED_TOKENIZER_LOAD_LOCKS.get(kind, nullcontext()): diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py index 68910f52b02..274eb051d6c 100644 --- a/tests/unit/proxy/common_utils/test_token_count_settings.py +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -9,6 +9,7 @@ from litellm.proxy.common_utils.token_count_settings import ( ProviderCountFailure, TokenCountModel, TokenCountSettings, + UnknownLocalTokenizer, count_source_fields, resolve_token_count_settings, ) @@ -249,13 +250,40 @@ def test_only_a_tokenizer_litellm_ships_is_taken_from_the_cost_map(value: object assert settings.tokenizer == expected -def _counted(source: str | None, model_used: str = "us.anthropic.claude-opus-5-5") -> TokenCountResponse: +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ({"token_count_tokenizer": "openai"}, "openai"), + ({"token_count_tokenizer": "deepseek"}, "deepseek"), + ({"token_count_tokenizer": "kimi_k3"}, "kimi_k3"), + ({"token_count_tokenizer": "local"}, UnknownLocalTokenizer(value="local")), + ({"token_count_tokenizer": "OpenAI"}, UnknownLocalTokenizer(value="OpenAI")), + ({"token_count_tokenizer": 5}, UnknownLocalTokenizer(value="5")), + ({"token_count_tokenizer": " "}, None), + ({"token_count_tokenizer": None}, None), + ({}, None), + ], +) +def test_only_the_deployment_names_the_tokenizer_that_counts_without_the_provider( + configured: dict[str, object], expected: object +) -> None: + settings: Final = resolve_token_count_settings( + configured=configured, default={"token_count_tokenizer": "kimi_k3"}, deployment_region=_in_us_west_2 + ) + + assert (settings.local_tokenizer, settings.tokenizer) == (expected, "kimi_k3") + + +def _counted( + source: str | None, model_used: str = "us.anthropic.claude-opus-5-5", count_tokenizer: str | None = None +) -> TokenCountResponse: return TokenCountResponse( total_tokens=7, request_model="sonnet", model_used=model_used, tokenizer_type="bedrock_api", count_source=source, + count_tokenizer=count_tokenizer, ) @@ -277,6 +305,13 @@ def test_count_source_fields_name_the_counting_model_only_when_token_count_model assert count_source_fields(_counted(source)) == expected +def test_count_source_fields_name_the_tokenizer_a_local_estimate_was_counted_with() -> None: + assert count_source_fields(_counted("local_estimate", count_tokenizer="llama3")) == { + "litellm_count_source": "local_estimate", + "litellm_count_tokenizer": "llama3", + } + + def test_every_shipped_token_count_model_resolves_to_a_mapped_aws_model() -> None: cost_map: Final = json.loads(_COST_MAP.read_text()) aws_models: Final = frozenset( diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 2f85d1c7516..66f10a5ee34 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -4331,13 +4331,14 @@ async def test_messages_count_tokens_on_a_responses_api_bridge_rejects_an_invali class _CountTokensPeer: - def __init__(self, reply: Mapping[str, JsonValue]) -> None: + def __init__(self, reply: Mapping[str, JsonValue], status_code: int = 200) -> None: self._reply: Final = reply + self._status_code: Final = status_code self.requests: tuple[httpx.Request, ...] = () def __call__(self, request: httpx.Request) -> httpx.Response: self.requests = (*self.requests, request) - return httpx.Response(200, json=dict(self._reply), request=request) + return httpx.Response(self._status_code, json=dict(self._reply), request=request) class _BedrockInvokeModel(BaseModel): @@ -5359,3 +5360,381 @@ def test_utils_token_counter_line_stays_one_line_when_the_model_carries_a_newlin assert line["request_model"] == forged_model assert "\n" not in message assert message.startswith(f'token count - litellm_call_id="count log" request_model={json.dumps(forged_model)} ') + + +_NOVA_PRO: Final = "us.amazon.nova-pro-v1:0" +_AWS_CANNOT_COUNT: Final = "The provided model doesn't support counting tokens." +_MIXED_SCRIPT_TURNS: Final[list[dict[str, JsonValue]]] = [ + {"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20} +] +_DEEPSEEK_TOKENIZER_JSON: Final = ( + Path(litellm.__file__).parent / "litellm_core_utils/tokenizers/deepseek_tokenizer.json" +) + + +def _expected_local_count(model: str, tokenizer: str) -> int: + custom_tokenizer: Final = ( + litellm.utils.create_tokenizer(_DEEPSEEK_TOKENIZER_JSON.read_text(encoding="utf-8")) + if tokenizer == "deepseek" + else None + ) + return litellm.token_counter(model=model, messages=_MIXED_SCRIPT_TURNS, custom_tokenizer=custom_tokenizer) + + +def _bedrock_router(model: str, model_info: Mapping[str, JsonValue]) -> Router: + return Router( + model_list=[ + { + "model_name": "bedrock-agent", + "litellm_params": { + "model": f"bedrock/{model}", + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", + }, + "model_info": dict(model_info), + } + ] + ) + + +async def _bedrock_counts( + monkeypatch: pytest.MonkeyPatch, model: str, model_info: Mapping[str, JsonValue], body: Mapping[str, JsonValue] +) -> tuple[object, TokenCountResponse | ProxyException, tuple[httpx.Request, ...]]: + from litellm.llms.bedrock.count_tokens import handler as bedrock_count_tokens_handler + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + peer: Final = _CountTokensPeer({"message": _AWS_CANNOT_COUNT}, status_code=400) + router: Final = _bedrock_router(model, model_info) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + with patch.object( + bedrock_count_tokens_handler, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ): + messages_reply: Final = await _messages_endpoint_count(router, {**body, "model": "bedrock-agent"}) + try: + utils_reply: TokenCountResponse | ProxyException = await token_counter( + request=TokenCountRequest.model_validate({**body, "model": "bedrock-agent"}), call_endpoint=True + ) + except ProxyException as rejected: + utils_reply = rejected # rebind-ok: the utils endpoint answers its rejection as a raised ProxyException + return messages_reply, utils_reply, peer.requests + + +def _rejection(reply: object) -> tuple[int, str]: + assert isinstance(reply, JSONResponse), reply + return reply.status_code, json.loads(bytes(reply.body))["error"]["message"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tokenizer", ["openai", "deepseek"]) +async def test_a_named_token_count_tokenizer_counts_with_exactly_that_tokenizer_and_never_asks_the_provider( + monkeypatch: pytest.MonkeyPatch, tokenizer: str +) -> None: + counts: Final = {name: _expected_local_count(_NOVA_PRO, name) for name in ("openai", "deepseek")} + assert counts["openai"] != counts["deepseek"], counts + + messages_reply, utils_reply, sent = await _bedrock_counts( + monkeypatch, _NOVA_PRO, {"token_count_tokenizer": tokenizer}, {"messages": _MIXED_SCRIPT_TURNS} + ) + + assert sent == () + assert messages_reply == { + "input_tokens": counts[tokenizer], + "litellm_count_source": "local_estimate", + "litellm_count_tokenizer": tokenizer, + } + assert isinstance(utils_reply, TokenCountResponse), utils_reply + assert (utils_reply.total_tokens, utils_reply.count_source, utils_reply.count_tokenizer) == ( + counts[tokenizer], + "local_estimate", + tokenizer, + ) + + +@pytest.mark.asyncio +async def test_without_a_named_tokenizer_a_model_the_provider_cannot_count_answers_the_providers_400( + monkeypatch: pytest.MonkeyPatch, +) -> None: + messages_reply, utils_reply, sent = await _bedrock_counts( + monkeypatch, _NOVA_PRO, {}, {"messages": _MIXED_SCRIPT_TURNS} + ) + + status, message = _rejection(messages_reply) + assert (status, _AWS_CANNOT_COUNT in message) == (400, True), message + assert isinstance(utils_reply, ProxyException), utils_reply + assert (utils_reply.code, _AWS_CANNOT_COUNT in utils_reply.message) == ("400", True), utils_reply.message + assert [request.url.path.endswith("/count-tokens") for request in sent] == [True, True] + + +@pytest.mark.asyncio +async def test_the_deployments_named_tokenizer_beats_the_tokenizer_the_cost_map_ships( + monkeypatch: pytest.MonkeyPatch, +) -> None: + kimi: Final = "us.moonshotai.kimi-k3" + assert cost_map_entry(kimi).get("token_count_tokenizer") == "kimi_k3" + + messages_reply, _, sent = await _bedrock_counts( + monkeypatch, kimi, {"token_count_tokenizer": "deepseek"}, {"messages": _MIXED_SCRIPT_TURNS} + ) + + assert sent == () + assert messages_reply == { + "input_tokens": _expected_local_count(kimi, "deepseek") + + (cost_map_entry(kimi).get("token_count_framing") or 0), + "litellm_count_source": "local_estimate", + "litellm_count_tokenizer": "deepseek", + } + + +@pytest.mark.asyncio +async def test_a_named_tokenizer_count_adds_the_framing_the_deployment_sets_for_requests_with_tools( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_info: Final = {"token_count_tokenizer": "openai", "token_count_framing": 5, "token_count_tool_framing": 40} + + _, utils_reply, sent = await _bedrock_counts( + monkeypatch, _NOVA_PRO, model_info, {"messages": _MIXED_SCRIPT_TURNS, "tools": [_KIMI_WEATHER]} + ) + + assert sent == () + assert isinstance(utils_reply, TokenCountResponse), utils_reply + assert utils_reply.total_tokens == 40 + litellm.token_counter( + model=_NOVA_PRO, messages=_MIXED_SCRIPT_TURNS, tools=[_KIMI_WEATHER] + ) + + +@pytest.mark.asyncio +async def test_a_named_kimi_k3_tokenizer_counts_text_with_kimi_k3_and_answers_400_for_a_body_it_cannot_render( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + encoding: Final = _available_kimi_k3(monkeypatch, tmp_path) + image_turn: Final[dict[str, JsonValue]] = { + "role": "user", + "content": [{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0K"}}], + } + + counted, _, counted_sent = await _bedrock_counts( + monkeypatch, _NOVA_PRO, {"token_count_tokenizer": "kimi_k3"}, {"messages": _MIXED_SCRIPT_TURNS} + ) + rejected, _, rejected_sent = await _bedrock_counts( + monkeypatch, _NOVA_PRO, {"token_count_tokenizer": "kimi_k3"}, {"messages": [image_turn]} + ) + + assert counted == { + "input_tokens": kimi_k3_tokens(encoding, _NOVA_PRO, _MIXED_SCRIPT_TURNS, None), + "litellm_count_source": "local_estimate", + "litellm_count_tokenizer": "kimi_k3", + } + status, message = _rejection(rejected) + assert ( + status, + "kimi_k3 tokenizer set in token_count_tokenizer, which could not count this request" in message, + ) == ( + 400, + True, + ), message + assert counted_sent == rejected_sent == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_info", "reason"), + [ + ({"token_count_tokenizer": "local"}, "sets token_count_tokenizer to 'local', which is not one of openai, "), + ( + {"token_count_tokenizer": "kimi_k3"}, + "counts tokens with the kimi_k3 tokenizer set in token_count_tokenizer, which could not be loaded", + ), + ], + ids=["unknown_name", "tokenizer_that_cannot_load"], +) +async def test_a_named_tokenizer_litellm_cannot_count_with_answers_400_without_asking_the_provider( + monkeypatch: pytest.MonkeyPatch, model_info: dict[str, JsonValue], reason: str +) -> None: + messages_reply, utils_reply, sent = await _bedrock_counts( + monkeypatch, _NOVA_PRO, model_info, {"messages": _MIXED_SCRIPT_TURNS} + ) + + status, message = _rejection(messages_reply) + assert (status, reason in message) == (400, True), message + assert isinstance(utils_reply, ProxyException), utils_reply + assert (utils_reply.code, reason in utils_reply.message) == ("400", True), utils_reply.message + assert sent == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "unreadable_body", + [ + pytest.param( + { + "messages": [ + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {"query": "new"}, + }, + {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_01", "content": []}, + ], + }, + ] + }, + id="web_search_blocks", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "Describe this"}], + "contents": [ + { + "role": "user", + "parts": [ + {"text": "Describe this"}, + {"fileData": {"mimeType": "image/png", "fileUri": "gs://b/c.png"}}, + ], + } + ], + }, + id="google_contents_with_media_the_translation_dropped", + ), + ], +) +async def test_a_named_tokenizer_that_cannot_read_the_request_answers_400_never_500( + monkeypatch: pytest.MonkeyPatch, unreadable_body: dict[str, JsonValue] +) -> None: + _, utils_reply, sent = await _bedrock_counts( + monkeypatch, _NOVA_PRO, {"token_count_tokenizer": "openai"}, unreadable_body + ) + + assert isinstance(utils_reply, ProxyException), utils_reply + assert utils_reply.code == "400", utils_reply.message + assert "counts tokens with the openai tokenizer set in token_count_tokenizer, which " in utils_reply.message + assert sent == () + + +_STORED_1998_TURN: Final[dict[str, JsonValue]] = { + "request_id": "chatcmpl-stored", + "call_type": "aresponses", + "session_id": "session-1", + "proxy_server_request": {"input": "Who won in 1998?", "model": "bedrock-agent"}, + "response": { + "id": "chatcmpl-stored", + "object": "chat.completion", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "The Bulls."}, "finish_reason": "stop"}], + }, + "status": "success", +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("litellm_params", "served_model"), + [ + ( + { + "model": f"bedrock/{_NOVA_PRO}", + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", + }, + _NOVA_PRO, + ), + ({"model": "mistral/mistral-small-latest", "api_key": "sk-test"}, "mistral-small-latest"), + ], + ids=["provider_with_a_count_api", "provider_without_a_count_api"], +) +async def test_a_named_tokenizer_counts_the_turns_previous_response_id_replays( + monkeypatch: pytest.MonkeyPatch, litellm_params: dict[str, JsonValue], served_model: str +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "bedrock-agent", + "litellm_params": litellm_params, + "model_info": {"token_count_tokenizer": "openai"}, + } + ] + ), + ) + inline_turns: Final[list[dict[str, JsonValue]]] = [ + {"role": "user", "content": "Who won in 1998?"}, + {"role": "assistant", "content": "The Bulls."}, + {"role": "user", "content": "And in 1999?"}, + ] + + with patch.object( + ResponsesSessionHandler, + "get_all_spend_logs_for_previous_response_id", + new_callable=AsyncMock, + return_value=[_STORED_1998_TURN], + ): + replayed: Final = await count_request_tokens( + request=TokenCountRequest(model="bedrock-agent", messages=[{"role": "user", "content": "And in 1999?"}]), + call_endpoint=True, + responses_request={ + "model": "bedrock-agent", + "input": "And in 1999?", + "previous_response_id": "chatcmpl-stored", + }, + ) + inline: Final = await count_request_tokens( + request=TokenCountRequest(model="bedrock-agent", messages=inline_turns), + call_endpoint=True, + responses_request={"model": "bedrock-agent", "input": inline_turns}, + ) + + assert (replayed.count_source, replayed.count_tokenizer) == ("local_estimate", "openai") + assert ( + replayed.total_tokens == inline.total_tokens == litellm.token_counter(model=served_model, messages=inline_turns) + ) + + +@pytest.mark.asyncio +async def test_a_named_tokenizer_answers_400_for_previous_response_id_turns_only_the_provider_stores( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peer: Final = _CountTokensPeer({"input_tokens": 99}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "gpt-counted", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "test-key", + "api_base": "https://count-tokens.test/v1", + }, + "model_info": {"token_count_tokenizer": "openai"}, + } + ] + ), + ) + + with ( + patch.object( + openai_count_tokens_http, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ), + pytest.raises(ProxyException) as rejected, + ): + await count_request_tokens( + request=TokenCountRequest(model="gpt-counted", messages=[{"role": "user", "content": "again"}]), + call_endpoint=True, + responses_request={"model": "gpt-counted", "input": "again", "previous_response_id": "resp_original"}, + ) + + assert rejected.value.code == "400", rejected.value.message + assert "cannot see the earlier turns previous_response_id points to" in rejected.value.message + assert peer.requests == () diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e2e1210bc2a..ab61650973f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -47888,6 +47888,11 @@ export interface components { * @enum {string} */ litellm_count_source?: "provider" | "token_count_model" | "local_estimate"; + /** + * Litellm Count Tokenizer + * @enum {string} + */ + litellm_count_tokenizer?: "openai" | "llama3" | "deepseek" | "cohere" | "kimi_k3"; /** Prompttokensdetails */ promptTokensDetails: components["schemas"]["PromptTokensDetails"][]; /** Totaltokens */ @@ -47918,6 +47923,8 @@ export interface components { TokenCountResponse: { /** Count Source */ count_source?: ("provider" | "token_count_model" | "local_estimate") | null; + /** Count Tokenizer */ + count_tokenizer?: ("openai" | "llama3" | "deepseek" | "cohere" | "kimi_k3") | null; /** * Error * @default false