scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

Merge ee402443f3 into b4a8bc715f

Browse source
This commit is contained in:
ryan-crabbe-berri 2026-11-05 07:14:38 +00:00 • committed by GitHub
commit 2af3808bdb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 710 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/bedrock/count_tokens/kimi_k3.py
View file

@ -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)
@ -479,7 +479,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)
)

25
litellm/proxy/common_utils/token_count_settings.py
View file

@ -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}

256
litellm/proxy/proxy_server.py
View file

@ -59,6 +59,7 @@ from typing import (
import anyio
import anyio.lowlevel
import tiktoken
import websockets
import websockets.exceptions
from openai.types.audio import TranscriptionStreamEvent
@ -127,6 +128,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 (
@ -184,6 +186,7 @@ from litellm.proxy.common_utils.token_count_settings import (
TOKENIZER_LOAD_WAIT_SECONDS,
ProviderCountFailure,
TokenCountSettings,
UnknownLocalTokenizer,
cost_map_entry,
resolve_token_count_settings,
)
@ -207,6 +210,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,
@ -15117,6 +15121,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,
)
@ -15399,7 +15404,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__,
)
@ -15560,6 +15565,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,
@ -15574,13 +15595,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)
)
@ -15602,6 +15617,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)(
@ -15611,6 +15637,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]]
@ -16015,6 +16209,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
@ -16078,12 +16273,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),
@ -16093,10 +16292,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,
@ -16104,7 +16303,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]
@ -16167,12 +16366,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
@ -16207,6 +16405,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(

2
litellm/types/token_count.py
View file

@ -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]

3
litellm/types/utils.py
View file

@ -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,
@ -4882,6 +4882,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

21
litellm/utils.py
View file

@ -2990,7 +2990,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()):

37
tests/unit/proxy/common_utils/test_token_count_settings.py
View file

@ -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(

383
tests/unit/proxy/test_proxy_token_counter.py
View file

@ -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 == ()

7
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -48036,6 +48036,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 */
@ -48066,6 +48071,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