feat(count_tokens)!: configurable token_count_framing and token_count_model, count source in the response (#46777)

Backport of #46777 to stable/1.104.x.
Cherry-picked from 75faaeab64 (main).
This commit is contained in:
Mateo Wang 2026-10-19 11:14:29 -07:00 • committed by mateo-berri
parent 52de3b0a29
commit b6f1a033de
32 changed files with 3571 additions and 150 deletions

View file

@ -3,9 +3,9 @@
import copy
import logging
import re
from collections.abc import Collection, Iterable, Mapping
from collections.abc import Collection, Iterable, Iterator, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, runtime_checkable
import httpx
from pydantic import TypeAdapter, ValidationError
@ -51,6 +51,28 @@ def safe_divide_seconds(seconds: float, denominator: float, default: float | Non
_DROP_PARAMS_BOOL: Final = TypeAdapter(bool)
def as_str_mapping(value: object) -> Mapping[str, object] | None:
return value if isinstance(value, Mapping) else None
@runtime_checkable
class _SupportsIter(Protocol):
def __iter__(self) -> Iterator[object]: ...
@runtime_checkable
class _SupportsGetItem(Protocol):
def __getitem__(self, index: int, /) -> object: ...
def iter_items(value: object) -> Iterable[object]:
if isinstance(value, _SupportsIter):
return value
if isinstance(value, _SupportsGetItem):
return iter(value)
raise TypeError(f"'{type(value).__name__}' object is not iterable")
def normalize_drop_params(value: object) -> bool | None:
if value is None or isinstance(value, bool):
return value

View file

@ -5580,3 +5580,25 @@ def resolve_structured_messages(
if structured:
return [msg if isinstance(msg, dict) else msg.model_dump() for msg in structured]
return None
def is_flat_responses_tool(tool: Mapping[str, object]) -> bool:
tool_type: Final = tool.get("type")
return tool_type in ("function", "custom") and tool_type not in tool and "input_schema" not in tool
def chat_completion_tools_for(tool: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
if is_flat_responses_tool(tool):
chat_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(
[dict(tool)] # pyright: ignore[reportArgumentType] # a JSON dict in a flat Responses tool shape
)
return tuple(chat_tools)
anthropic_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(
[dict(tool)] # pyright: ignore[reportArgumentType] # a JSON dict in an Anthropic or chat tool shape
)
return tuple(anthropic_tools)

View file

@ -0,0 +1,502 @@
"""
Exact Bedrock token counts for Kimi K3, which has no CountTokens support.
Renders the Converse body litellm sends through Moonshot's Kimi K3 chat template and counts it with the
official tokenizer, downloaded once from Hugging Face at a pinned revision and checked against its SHA-256.
"""
import asyncio
import base64
import copy
import hashlib
import json
import threading
import time
from collections.abc import Callable, Iterable, Mapping, Sequence
from concurrent.futures import Future
from dataclasses import dataclass
from enum import Enum
from itertools import chain, groupby
from pathlib import Path
from typing import Annotated, Final, Literal, TypeAlias
import regex
import tiktoken
from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError, model_validator
from typing_extensions import Self
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.prompt_templates.factory import (
chat_completion_tools_for,
is_flat_responses_tool,
)
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
KIMI_K3_REPO: Final = "moonshotai/Kimi-K3"
KIMI_K3_REVISION: Final = "f831ab66814297da540d832a5235f8e904f29d06"
KIMI_K3_RANKS_FILE: Final = "tiktoken.model"
KIMI_K3_RANKS_SHA256: Final = "b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103"
_PATTERN: Final = "|".join(
(
r"[\p{Han}]+",
r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?",
r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?",
r"\p{N}{1,3}",
r" ?[^\s\p{L}\p{N}]+[\r\n]*",
r"\s*[\r\n]+",
r"\s+(?!\S)",
r"\s+",
)
)
# Ends chunks only where the pattern above always splits, so their counts sum to the whole text's count while BPE
# stays linear. Only a run of over 1024 characters with no such split, like 2000 spaces, is cut mid-piece.
_CHUNK: Final = regex.compile(
r"(?s).{1,1024}(?:(?<=\S)(?= \S)|(?<=\p{N})(?=\P{N})|(?<=[^\s\p{N}])(?=\p{N})|(?<=\p{L})(?=[^\p{L}\p{M}'])"
r"|(?<=[\r\n])(?=\S)|\Z)|.{1024}\p{N}*"
)
@dataclass(frozen=True, slots=True)
class KimiK3TokenizerSource:
ranks_path: Callable[[], str]
sha256: str
def _hugging_face_ranks_path() -> str:
import huggingface_hub
return huggingface_hub.hf_hub_download(
repo_id=KIMI_K3_REPO,
filename=KIMI_K3_RANKS_FILE,
revision=KIMI_K3_REVISION,
local_files_only=litellm.disable_hf_tokenizer_download is True,
)
HUGGING_FACE_SOURCE: Final = KimiK3TokenizerSource(ranks_path=_hugging_face_ranks_path, sha256=KIMI_K3_RANKS_SHA256)
def _loaded(source: KimiK3TokenizerSource) -> tiktoken.Encoding | str:
try:
data: Final = Path(source.ranks_path()).read_bytes()
except Exception as e: # noqa: BLE001 # a download or read failure of any kind falls back to the estimate
return type(e).__name__
if hashlib.sha256(data).hexdigest() != source.sha256:
return "SHA-256 mismatch"
rows: Final = (line.split() for line in data.splitlines() if line)
return tiktoken.Encoding(
name="kimi_k3",
pat_str=_PATTERN,
mergeable_ranks={base64.b64decode(token): int(rank) for token, rank in rows},
special_tokens={},
)
def load_kimi_k3_encoding(source: KimiK3TokenizerSource) -> tiktoken.Encoding | None:
loaded: Final = _loaded(source)
if isinstance(loaded, tiktoken.Encoding):
return loaded
verbose_logger.warning(
"Kimi K3 tokenizer unavailable (%s), so Kimi K3 token counts are estimated until a later count loads it. "
"Allow downloads of %s from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache",
loaded,
KIMI_K3_RANKS_FILE,
KIMI_K3_REPO,
KIMI_K3_REVISION,
)
return None
def _started_load(source: KimiK3TokenizerSource) -> "Future[tiktoken.Encoding | None]":
encoding: Final = Future[tiktoken.Encoding | None]()
threading.Thread(
target=lambda: encoding.set_result(load_kimi_k3_encoding(source)), name="litellm-kimi-k3-tokenizer", daemon=True
).start()
return encoding
@dataclass(frozen=True, slots=True)
class _EncodingLoad:
encoding: "Future[tiktoken.Encoding | None]"
started_at: float
is_retry: bool
def retry_due(self, now: float, retry_after_seconds: float) -> bool:
return self.encoding.done() and self.encoding.result() is None and now - self.started_at >= retry_after_seconds
@dataclass(slots=True)
class KimiK3Encodings:
source: KimiK3TokenizerSource
retry_after_seconds: float
clock: Callable[[], float] = time.monotonic
current: _EncodingLoad | None = None
def _load(self) -> _EncodingLoad:
now: Final = self.clock()
current: Final = self.current
if current is not None and not current.retry_due(now, self.retry_after_seconds):
return current
load: Final = _EncodingLoad(encoding=_started_load(self.source), started_at=now, is_retry=current is not None)
self.current = load
return load
async def encoding(self, wait_seconds: float) -> tiktoken.Encoding | None:
load: Final = self._load()
remaining: Final = 0 if load.is_retry else load.started_at + wait_seconds - self.clock()
if load.encoding.done() or remaining <= 0:
return load.encoding.result() if load.encoding.done() else None
await asyncio.wait((asyncio.wrap_future(load.encoding),), timeout=remaining)
if load.encoding.done():
return load.encoding.result()
verbose_logger.warning(
"Kimi K3 tokenizer still loading after %s s, so Kimi K3 token counts are estimated until it loads. "
"Check that %s can be downloaded from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache",
wait_seconds,
KIMI_K3_RANKS_FILE,
KIMI_K3_REPO,
KIMI_K3_REVISION,
)
return None
class _Converse(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
class _Text(_Converse):
text: str
class _CachePoint(_Converse):
cache_point: JsonValue = Field(alias="cachePoint")
class _ToolUse(_Converse):
tool_use_id: str = Field(alias="toolUseId")
name: str
arguments: Mapping[str, JsonValue] = Field(alias="input")
class _ToolUseBlock(_Converse):
tool_use: _ToolUse = Field(alias="toolUse")
class _ToolResult(_Converse):
tool_use_id: str = Field(alias="toolUseId")
content: tuple[_Text, ...]
status: str | None = None
class _ToolResultBlock(_Converse):
tool_result: _ToolResult = Field(alias="toolResult")
class _ReasoningText(_Converse):
text: str
signature: str | None = None
class _Reasoning(_Converse):
reasoning_text: _ReasoningText = Field(alias="reasoningText")
class _ReasoningBlock(_Converse):
reasoning_content: _Reasoning = Field(alias="reasoningContent")
class _UserMessage(_Converse):
role: Literal["user"]
content: tuple[_Text | _ToolResultBlock | _CachePoint, ...]
class _AssistantMessage(_Converse):
role: Literal["assistant"]
content: tuple[_Text | _ToolUseBlock | _ReasoningBlock | _CachePoint, ...]
def calls(self) -> tuple[_ToolUse, ...]:
return tuple(block.tool_use for block in self.content if isinstance(block, _ToolUseBlock))
class _InputSchema(_Converse):
definition: JsonValue = Field(alias="json")
class _ToolSpec(_Converse):
name: str
description: str | None = None
input_schema: _InputSchema = Field(alias="inputSchema")
def declared(self) -> Mapping[str, JsonValue]:
described: Final[Mapping[str, JsonValue]] = (
{} if self.description is None else {"description": self.description}
)
return {"name": self.name, **described, "parameters": self.input_schema.definition}
class _ToolSpecBlock(_Converse):
tool_spec: _ToolSpec = Field(alias="toolSpec")
class _ToolConfig(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
tools: tuple[_ToolSpecBlock | _CachePoint, ...]
_Message: TypeAlias = Annotated[_UserMessage | _AssistantMessage, Field(discriminator="role")]
def _result_ids(message: _Message) -> frozenset[str]:
if isinstance(message, _AssistantMessage):
return frozenset()
return frozenset(block.tool_result.tool_use_id for block in message.content if isinstance(block, _ToolResultBlock))
def _call_ids(message: _Message | None) -> frozenset[str]:
if not isinstance(message, _AssistantMessage):
return frozenset()
return frozenset(call.tool_use_id for call in message.calls())
class _ConverseRequest(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
messages: tuple[_Message, ...]
system: tuple[_Text | _CachePoint, ...] = ()
tool_config: _ToolConfig | None = Field(default=None, alias="toolConfig")
@model_validator(mode="after")
def _results_answer_the_previous_calls(self) -> Self:
pairs: Final = zip((None, *self.messages), self.messages, strict=False)
if not all(_result_ids(message) <= _call_ids(previous) for previous, message in pairs):
raise ValueError("a toolResult does not answer a toolUse of the previous assistant message")
return self
class _Control(Enum):
OPEN = "<|open|>"
CLOSE = "<|close|>"
SEP = "<|sep|>"
END_OF_MSG = "<|end_of_msg|>"
_Segment: TypeAlias = str | _Control
_Segments: TypeAlias = tuple[_Segment, ...]
def _open(tag: str, attrs: Iterable[tuple[str, str]] = ()) -> _Segments:
escaped: Final = (
(f" {key}", '="', value.replace("&", "&amp;").replace('"', "&quot;"), '"') for key, value in attrs
)
return (_Control.OPEN, tag, *chain.from_iterable(escaped), _Control.SEP)
def _close(tag: str) -> _Segments:
return (_Control.CLOSE, tag, _Control.SEP)
def _message(attrs: Iterable[tuple[str, str]], body: Iterable[_Segment]) -> _Segments:
return (*_open("message", attrs), *body, *_close("message"), _Control.END_OF_MSG)
def _tool_declare(specs: Sequence[_ToolSpec]) -> _Segments:
tools: Final = json.dumps(
[{"type": "function", "function": spec.declared()} for spec in specs],
ensure_ascii=False,
separators=(",", ":"),
sort_keys=True,
)
return _message(
(("role", "system"), ("type", "tool-declare")),
(f"# Tools\nHere are the available tools, described in JSONSchema.\n\n```json\n{tools}\n```",),
)
def _text(blocks: Iterable[_Converse]) -> str:
return "".join(block.text for block in blocks if isinstance(block, _Text))
def _json_type(value: JsonValue) -> str:
match value:
case bool():
return "boolean"
case None:
return "null"
case int() | float():
return "number"
case str():
return "string"
case dict():
return "object"
case _:
return "array"
def _argument(key: str, value: JsonValue) -> _Segments:
text: Final = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False)
return (*_open("argument", (("key", key), ("type", _json_type(value)))), text, *_close("argument"))
def _call(index: int, call: _ToolUse) -> _Segments:
arguments: Final = chain.from_iterable(_argument(key, value) for key, value in call.arguments.items())
return (*_open("call", (("tool", call.name), ("index", str(index)))), *arguments, *_close("call"))
def _tool_calls(calls: Sequence[_ToolUse]) -> _Segments:
if not calls:
return ()
rendered: Final = chain.from_iterable(_call(index, call) for index, call in enumerate(calls, start=1))
return (*_open("tools"), *rendered, *_close("tools"))
def _assistant_message(message: _AssistantMessage) -> _Segments:
reasoning: Final = "".join(
block.reasoning_content.reasoning_text.text for block in message.content if isinstance(block, _ReasoningBlock)
)
return _message(
(("role", "assistant"),),
(
*_open("think"),
reasoning,
*_close("think"),
*_open("response"),
_text(message.content),
*_close("response"),
*_tool_calls(message.calls()),
),
)
def _tool_messages(results: Sequence[_ToolResult], previous: _AssistantMessage) -> _Segments:
names: Final = {call.tool_use_id: call.name for call in previous.calls()}
return tuple(
chain.from_iterable(
_message(
(("role", "tool"), ("tool", names[result.tool_use_id]), ("index", str(index))),
(_text(result.content),),
)
for index, result in enumerate(results, start=1)
)
)
def _user_run(blocks: Sequence[_Text | _ToolResultBlock], previous: _Message | None) -> _Segments:
results: Final = tuple(block.tool_result for block in blocks if isinstance(block, _ToolResultBlock))
if results and isinstance(previous, _AssistantMessage):
return _tool_messages(results, previous)
return _message((("role", "user"),), (_text(blocks),))
def _user_messages(message: _UserMessage, previous: _Message | None) -> _Segments:
blocks: Final = tuple(block for block in message.content if not isinstance(block, _CachePoint))
runs: Final = groupby(blocks, key=lambda block: isinstance(block, _ToolResultBlock))
return tuple(chain.from_iterable(_user_run(tuple(run), previous) for _, run in runs))
def _turn(previous: _Message | None, message: _Message) -> _Segments:
if isinstance(message, _AssistantMessage):
return _assistant_message(message)
return _user_messages(message, previous)
def _rendered(request: _ConverseRequest) -> _Segments:
specs: Final = (
()
if request.tool_config is None
else tuple(block.tool_spec for block in request.tool_config.tools if isinstance(block, _ToolSpecBlock))
)
system: Final = _text(request.system)
turns: Final = chain.from_iterable(
_turn(previous, message) for previous, message in zip((None, *request.messages), request.messages, strict=False)
)
return (
*(_tool_declare(specs) if specs else ()),
*(_message((("role", "system"),), (system,)) if system else ()),
*turns,
*_open("message", (("role", "assistant"),)),
*_open("think"),
)
def _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)
class _TextPart(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
type: Literal["text"]
class _TextOnlyMessage(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
content: str | tuple[_TextPart, ...] | None = None
def _text_only(message: AllMessageValues) -> bool:
try:
_TextOnlyMessage.model_validate(message)
except ValidationError:
return False
return True
def _as_converse_input(tool: ChatCompletionToolParam) -> tuple[Mapping[str, object], ...]:
return chat_completion_tools_for(tool) if is_flat_responses_tool(tool) else (tool,)
def _converse_request(
model: str, messages: Sequence[AllMessageValues], tools: Sequence[ChatCompletionToolParam] | None
) -> _ConverseRequest | None:
if not all(_text_only(message) for message in messages):
return None
try:
chat_tools: Final = tuple(
chat_tool for tool in tools or () for chat_tool in _as_converse_input(tool)
) # comprehension-ok: flattens each tool's translation
return _ConverseRequest.model_validate(
litellm.AmazonConverseConfig().transform_request(
model=model,
messages=copy.deepcopy(list(messages)),
optional_params={"tools": copy.deepcopy(list(chat_tools))} if chat_tools else {},
litellm_params={},
headers={},
)
)
except Exception: # noqa: BLE001 # a body litellm cannot send to Bedrock, or cannot render, falls back to the estimate
return None
def kimi_k3_tokens(
encoding: tiktoken.Encoding,
model: str,
messages: Sequence[AllMessageValues],
tools: Sequence[ChatCompletionToolParam] | None,
) -> int | None:
request: Final = _converse_request(model, messages, tools)
if request is None:
return None
return sum(
1 if isinstance(segment, _Control) else _text_tokens(encoding, segment) for segment in _rendered(request)
)
def kimi_k3_anthropic_tokens(encoding: tiktoken.Encoding, model: str, fields: Mapping[str, object]) -> int | None:
body: Final = {**fields, "model": model, "max_tokens": 1}
try:
translated, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
copy.deepcopy(body), # pyright: ignore[reportArgumentType] # raw /v1/messages fields
custom_llm_provider="bedrock",
)
except Exception: # noqa: BLE001 # a request the adapter cannot translate falls back to the estimate
return None
match translated:
case {"tools": tools}:
return kimi_k3_tokens(encoding, model, translated["messages"], tools)
case _:
return kimi_k3_tokens(encoding, model, translated["messages"], None)
def kimi_k3_template_tokens(message_count: int, has_tools: bool) -> int:
return 14 + 13 * max(message_count - 1, 0) + (40 if has_tools else 0)

View file

@ -116,6 +116,7 @@ class OpenAITokenCounter(BaseTokenCounter):
error=True,
error_message=request_error,
status_code=400,
rejected_before_provider_call=True,
)
deployment = deployment or {}

View file

@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final, Literal
from pydantic import JsonValue, TypeAdapter, ValidationError
from typing_extensions import assert_never
from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for
from litellm.types.llms.anthropic import AnthropicMessagesTextParam
from litellm.types.llms.openai import AllMessageValues
from litellm.types.llms.vertex_ai import ContentType
@ -70,27 +71,6 @@ def _is_function_tool(tool: Mapping[str, JsonValue]) -> bool:
return isinstance(tool.get("name"), str) or "input_schema" in tool or "function" in tool
def _is_responses_function_tool(tool: Mapping[str, JsonValue]) -> bool:
return tool.get("type") == "function" and "function" not in tool and "input_schema" not in tool
def _chat_tool(tool: Mapping[str, JsonValue]) -> tuple[Mapping[str, object], ...]:
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
if _is_responses_function_tool(tool):
chat_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(
[dict(tool)] # pyright: ignore[reportArgumentType] # validated JSON dict in the Responses function tool shape
)
return tuple(chat_tools)
anthropic_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(
[dict(tool)] # pyright: ignore[reportArgumentType] # validated JSON dict in an Anthropic or chat tool shape
)
return tuple(anthropic_tools)
def _map_function_tools(
tools: object | None,
) -> tuple[Mapping[str, object], ...] | None | VertexGeminiCountContextError:
@ -107,7 +87,7 @@ def _map_function_tools(
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
chat_tools: Final = tuple(
chat_tool for tool in function_tools for chat_tool in _chat_tool(tool)
chat_tool for tool in function_tools for chat_tool in chat_completion_tools_for(tool)
) # comprehension-ok: flattens each tool's translation
try:
mapped_tools: Final = VertexGeminiConfig()._map_function( # pyright: ignore[reportPrivateUsage] # the tool mapping inference uses
@ -158,6 +138,7 @@ def count_context_error_response(
error=True,
error_message=error.error_message,
status_code=400,
rejected_before_provider_call=True,
)

View file

@ -19337,7 +19337,13 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
"supports_vision": false,
"token_count_framing": 97
},
"databricks/system.ai.deepseek-v4-flash-0731": {
"litellm_provider": "databricks",
"mode": "chat",
"token_count_framing": 97
},
"databricks/databricks-deepseek-v4-pro-0813": {
"cache_creation_input_token_cost": 1.31999e-06,
@ -20328,7 +20334,13 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"token_count_framing": 100
},
"databricks/system.ai.kimi-k3": {
"litellm_provider": "databricks",
"mode": "chat",
"token_count_framing": 100
},
"databricks/databricks-llama-4-maverick": {
"cache_creation_input_token_cost": 5.0001e-07,
@ -75642,7 +75654,9 @@
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"token_count_framing": 64,
"token_count_tokenizer": "kimi_k3"
},
"global.moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 3.75e-06,
@ -75662,7 +75676,9 @@
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"token_count_framing": 64,
"token_count_tokenizer": "kimi_k3"
},
"us.moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 4.125e-06,
@ -75682,7 +75698,9 @@
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"token_count_framing": 64,
"token_count_tokenizer": "kimi_k3"
},
"openrouter/xiaomi/mimo-v2.6-flash": {
"cache_read_input_token_cost": 2.8e-09,

View file

@ -41,6 +41,7 @@ from litellm.proxy.common_utils.openai_error_payload import (
openai_error_type,
with_litellm_call_id,
)
from litellm.proxy.common_utils.token_count_settings import count_source_fields
router: Final = APIRouter()
@ -366,7 +367,7 @@ async def count_tokens(
log_llm_api_exception(e, litellm_call_id)
return _anthropic_error_json_response(proxy_exception_from_route_error(e), request)
else:
return {"input_tokens": token_response.total_tokens}
return {"input_tokens": token_response.total_tokens, **count_source_fields(token_response)}
@router.post(

View file

@ -0,0 +1,164 @@
from collections.abc import Callable, Mapping
from dataclasses import dataclass, replace
from types import MappingProxyType
from typing import Final, Literal
from pydantic import BaseModel, ConfigDict, 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 TokenCountSourceFields, TokenCountTokenizer
from litellm.types.utils import TokenCountResponse
CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404})
TOKENIZER_LOAD_WAIT_SECONDS: Final = 5.0
TOKENIZER_LOAD_REUSE_SECONDS: Final = 60.0
TOKENIZER_LOAD_MAX_RUNNING: Final = 4
class _UpstreamErrorDetail(BaseModel):
model_config = ConfigDict(frozen=True, extra="allow")
message: str
class _UpstreamErrorEnvelope(BaseModel):
model_config = ConfigDict(frozen=True, extra="allow")
type: Literal["error"]
error: _UpstreamErrorDetail
def _joined(message: str, hint: str) -> str:
return f"{message.rstrip().rstrip('.')}. {hint}"
def _with_hint(message: str, hint: str) -> str:
try:
envelope: Final = _UpstreamErrorEnvelope.model_validate_json(message)
except ValidationError:
return _joined(AnthropicExceptionMapping.extract_error_message(message), hint)
detail: Final = envelope.error.model_copy(update={"message": _joined(envelope.error.message, hint)})
return envelope.model_copy(update={"error": detail}).model_dump_json()
@dataclass(frozen=True, slots=True)
class TokenCountModel:
model: str
region: str | None
@dataclass(frozen=True, slots=True)
class TokenCountSettings:
framing: int | None
count_model: TokenCountModel | None
tokenizer: TokenCountTokenizer | None = None
@dataclass(frozen=True, slots=True)
class ProviderCountFailure:
status_code: int
message: str
rejected_before_provider_call: bool = False
@property
def cannot_count_model(self) -> bool:
return not self.rejected_before_provider_call and self.status_code in CANNOT_COUNT_MODEL_STATUSES
def with_hint(self, hint: str) -> "ProviderCountFailure":
return replace(self, message=_with_hint(self.message, hint))
def to_proxy_exception(self) -> ProxyException:
return ProxyException(
message=self.message,
type="token_counting_error",
param="model",
code=self.status_code,
)
def cost_map_entry(model: str, custom_llm_provider: str | None = None) -> Mapping[str, object]:
try:
return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except Exception: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models
return MappingProxyType({})
def _framing(value: object) -> int | None:
match value:
case bool():
return None
case int() if value >= 0:
return value
case float() if value.is_integer() and value >= 0:
return int(value)
case _:
return None
def _non_blank(value: object) -> str | None:
return value if isinstance(value, str) and value.strip() else None
def _count_model(value: Mapping[str, object]) -> TokenCountModel | None:
model: Final = _non_blank(value.get("model"))
return None if model is None else TokenCountModel(model=model, region=_non_blank(value.get("region")))
def _geography(region: str) -> str:
return region.rsplit("-", 2)[0]
def _within_deployment_geography(
count_model: TokenCountModel | None, deployment_region: Callable[[], str]
) -> TokenCountModel | None:
if count_model is None or count_model.region is None or count_model.model.startswith("global."):
return count_model
if _geography(deployment_region()) == _geography(count_model.region):
return count_model
return replace(count_model, region=None)
def _resolved_count_model(
configured: object, default: object, deployment_region: Callable[[], str]
) -> TokenCountModel | None:
configured_mapping: Final = as_str_mapping(configured)
if configured_mapping is not None:
return _count_model(configured_mapping)
default_mapping: Final = as_str_mapping(default)
return (
None
if default_mapping is None
else _within_deployment_geography(_count_model(default_mapping), deployment_region)
)
def _tokenizer(value: object) -> TokenCountTokenizer | None:
match value:
case "kimi_k3":
return "kimi_k3"
case _:
return None
def resolve_token_count_settings(
configured: Mapping[str, object], default: Mapping[str, object], deployment_region: Callable[[], str]
) -> TokenCountSettings:
framing: Final = _framing(configured.get("token_count_framing"))
return TokenCountSettings(
framing=framing if framing is not None else _framing(default.get("token_count_framing")),
count_model=_resolved_count_model(
configured.get("token_count_model"), default.get("token_count_model"), deployment_region
),
tokenizer=_tokenizer(default.get("token_count_tokenizer")),
)
def count_source_fields(response: TokenCountResponse) -> TokenCountSourceFields:
match response.count_source:
case None:
return {}
case "token_count_model":
return {"litellm_count_source": "token_count_model", "litellm_count_model": response.model_used}
case source:
return {"litellm_count_source": source}

View file

@ -18,6 +18,7 @@ from litellm.proxy.common_request_processing import (
resolve_litellm_call_id,
)
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.common_utils.token_count_settings import count_source_fields
from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
@ -306,10 +307,12 @@ async def google_count_tokens(request: Request, model_name: str):
ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500),
)
original_response: Final[dict] = token_response.original_response or {}
return TokenCountDetailsResponse(
totalTokens=original_response.get("totalTokens") or token_response.total_tokens or 0,
promptTokensDetails=original_response.get("promptTokensDetails", []),
)
counted: Final[TokenCountDetailsResponse] = {
"totalTokens": original_response.get("totalTokens") or token_response.total_tokens or 0,
"promptTokensDetails": original_response.get("promptTokensDetails", []),
**count_source_fields(token_response),
}
return counted
# ============================================================

View file

@ -26,6 +26,7 @@ from collections.abc import (
MutableMapping,
Sequence,
)
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from itertools import chain
from types import MappingProxyType, UnionType
@ -47,6 +48,7 @@ from typing import (
)
import anyio
import anyio.lowlevel
import websockets
import websockets.exceptions
from pydantic import BaseModel, Json, JsonValue, TypeAdapter, ValidationError
@ -87,6 +89,13 @@ from litellm.litellm_core_utils.litellm_logging import (
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.llms.bedrock.count_tokens.kimi_k3 import (
HUGGING_FACE_SOURCE,
KimiK3Encodings,
kimi_k3_anthropic_tokens,
kimi_k3_template_tokens,
kimi_k3_tokens,
)
from litellm.proxy._types import (
UI_TEAM_ID,
CallbackDelete,
@ -133,6 +142,15 @@ from litellm.proxy.common_utils.callback_utils import (
strip_callback_config,
)
from litellm.proxy.common_utils.realtime_utils import _realtime_request_body
from litellm.proxy.common_utils.token_count_settings import (
TOKENIZER_LOAD_MAX_RUNNING,
TOKENIZER_LOAD_REUSE_SECONDS,
TOKENIZER_LOAD_WAIT_SECONDS,
ProviderCountFailure,
TokenCountSettings,
cost_map_entry,
resolve_token_count_settings,
)
from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog
from litellm.router_utils.access_windows import access_windows_config_error
from litellm.router_utils.add_retry_fallback_headers import (
@ -278,6 +296,7 @@ from litellm.constants import (
CLI_SSO_SESSION_TTL_SECONDS,
DAYS_IN_A_MONTH,
DEFAULT_HEALTH_CHECK_INTERVAL,
DEFAULT_MAX_LRU_CACHE_SIZE,
DEFAULT_MODEL_CREATED_AT_TIME,
GLOBAL_PROXY_SPEND_CACHE_KEY,
LITELLM_PROXY_ADMIN_NAME,
@ -308,8 +327,10 @@ from litellm.litellm_core_utils.agentic_loop_settings import (
from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type
from litellm.litellm_core_utils.core_helpers import (
_get_parent_otel_span_from_kwargs,
as_str_mapping,
drop_params_flag,
get_litellm_metadata_from_kwargs,
iter_items,
)
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -880,7 +901,13 @@ from litellm.types.secret_managers.main import (
KeyManagementSettings,
KeyManagementSystem,
)
from litellm.types.utils import CredentialItem, CustomHuggingfaceTokenizer, RawRequestTypedDict, StandardLoggingPayload
from litellm.types.utils import (
CredentialItem,
CustomHuggingfaceTokenizer,
RawRequestTypedDict,
SelectTokenizerResponse,
StandardLoggingPayload,
)
from litellm.types.utils import ModelInfo as ModelMapInfo
from litellm.utils import _add_custom_logger_callback_to_specific_event
@ -13614,24 +13641,35 @@ async def _try_provider_token_count(
provider_counter: "BaseTokenCounter",
custom_llm_provider: str | None,
model_to_use: str,
region: str | None,
messages: list | None,
contents: list | None,
deployment: dict[str, Any] | None,
deployment: Mapping[str, object] | None,
request_model: str,
tools: list | None = None,
system: str | None = None,
messages_format: TokenCountMessagesFormat = "openai",
output_format: JsonValue = None,
) -> Optional["TokenCountResponse"]:
"""Attempt provider-specific token counting. Returns result on success, None to fall through to local counting."""
) -> Union["TokenCountResponse", ProviderCountFailure, None]:
"""Attempt provider-specific token counting. Returns None to fall through to local counting."""
if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider):
return None
deployment_in_region: Final = (
None
if deployment is None
else dict(deployment)
if region is None
else {
**deployment,
"litellm_params": {**(as_str_mapping(deployment.get("litellm_params")) or {}), "aws_region_name": region},
}
)
try:
result: Final = await provider_counter.count_tokens(
model_to_use=model_to_use or "",
messages=messages,
contents=contents,
deployment=deployment,
deployment=deployment_in_region,
request_model=request_model,
tools=tools,
system=system,
@ -13639,24 +13677,325 @@ async def _try_provider_token_count(
output_format=output_format,
)
except httpx.HTTPStatusError as e:
error_message: Final = getattr(e, "message", None) or str(e)
status_code: Final = getattr(e, "status_code", None) or e.response.status_code
raise ProxyException(
message=error_message,
type="token_counting_error",
param="model",
code=status_code,
return ProviderCountFailure(
status_code=getattr(e, "status_code", None) or e.response.status_code,
message=getattr(e, "message", None) or str(e),
)
if result is not None and result.error is True:
raise ProxyException(
return ProviderCountFailure(
status_code=result.status_code or 500,
message=result.error_message or "Token counting failed",
type="token_counting_error",
param="model",
code=result.status_code or 500,
rejected_before_provider_call=result.rejected_before_provider_call,
)
return result
async def _provider_token_count(
count_with: Callable[[str, str | None], Awaitable[Union["TokenCountResponse", ProviderCountFailure, None]]],
model_to_use: str,
settings: TokenCountSettings,
local_estimate_allowed: bool,
) -> Union["TokenCountResponse", ProviderCountFailure, None]:
"""Count with the provider, using the token_count_model in its region when one is set.
Returns None when the count should be estimated locally, or the provider's failure when the framed local
estimate stands in for it and that failure is the answer if the estimate cannot read the request."""
count_model: Final = settings.count_model
counted: Final = await (
count_with(model_to_use, None) if count_model is None else count_with(count_model.model, count_model.region)
)
if counted is None:
return None
if isinstance(counted, TokenCountResponse):
if count_model is None:
return counted.model_copy(update={"count_source": "provider"})
return counted.model_copy(update={"count_source": "token_count_model", "model_used": count_model.model})
failure: Final = (
counted
if count_model is None or count_model.region is None
else counted.with_hint(
f"This model counts tokens with {count_model.model} in {count_model.region}. "
"If that region is not enabled for the AWS account, enable it or set token_count_model in model_info"
)
)
if count_model is None and failure.cannot_count_model and settings.framing is not None and local_estimate_allowed:
return failure
raise failure.to_proxy_exception()
@dataclass(frozen=True, slots=True)
class _LocalTokenEstimate:
total_tokens: int
tokenizer_type: str
def _is_fork_guard_error(error: BaseException) -> bool:
from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking
return isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking))
def _is_tokenizer_panic(error: BaseException) -> bool:
return type(error).__module__ == "pyo3_runtime" and type(error).__name__ == "PanicException"
def _answers_provider_error(error: BaseException) -> bool:
return (isinstance(error, Exception) or _is_tokenizer_panic(error)) and not _is_fork_guard_error(error)
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",
provider_failure.status_code,
type(error).__name__,
)
return provider_failure.to_proxy_exception()
async def _estimate_or_provider_error(estimate: Awaitable[int], provider_failure: ProviderCountFailure) -> int:
try:
return await estimate
except BaseException as e:
if not _answers_provider_error(e):
raise
raise _provider_error_after_failed_estimate(provider_failure, e) from e
@dataclass(frozen=True, slots=True)
class _CustomTokenizerKey:
identifier: str
revision: str
auth_token: str | None
class _TokenizerLoadsBusy(Exception):
pass
@dataclass(frozen=True, slots=True)
class _TokenizerLoadOutcome:
tokenizer: SelectTokenizerResponse | BaseException
finished_at: float
@dataclass(frozen=True, slots=True)
class _SharedTokenizerLoad:
outcome: asyncio.Task[_TokenizerLoadOutcome]
started_at: float
def reusable(self, now: float) -> bool:
return not self.outcome.done() or now - self.outcome.result().finished_at < TOKENIZER_LOAD_REUSE_SECONDS
async def _loaded_tokenizer(
model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter
) -> SelectTokenizerResponse | BaseException:
try:
return await asyncify(litellm.utils._select_tokenizer, limiter=threads)(
model=model, custom_tokenizer=custom_tokenizer
)
except asyncio.CancelledError:
raise
except BaseException as e: # noqa: BLE001 # every count sharing this load decides what it answers, pyo3 panics included
return e
async def _tokenizer_load_outcome(
model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter
) -> _TokenizerLoadOutcome:
tokenizer: Final = await _loaded_tokenizer(model, custom_tokenizer, threads)
return _TokenizerLoadOutcome(tokenizer=tokenizer, finished_at=asyncio.get_running_loop().time())
@dataclass(slots=True)
class _TokenizerLoads:
threads: anyio.CapacityLimiter
running: int = 0
def start(
self, model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, loop: asyncio.AbstractEventLoop
) -> _SharedTokenizerLoad | None:
if self.running >= self.threads.total_tokens:
return None
self.running += 1
outcome: Final = loop.create_task(_tokenizer_load_outcome(model, custom_tokenizer, self.threads))
outcome.add_done_callback(self._finished)
return _SharedTokenizerLoad(outcome=outcome, started_at=loop.time())
def _finished(self, _outcome: asyncio.Task[_TokenizerLoadOutcome]) -> None:
self.running -= 1
_TOKENIZER_LOADS: Final = anyio.lowlevel.RunVar[_TokenizerLoads]("litellm_proxy_tokenizer_loads")
def _tokenizer_loads_for_this_loop() -> _TokenizerLoads:
existing: Final = _TOKENIZER_LOADS.get(None)
if existing is not None:
return existing
created: Final = _TokenizerLoads(threads=anyio.CapacityLimiter(TOKENIZER_LOAD_MAX_RUNNING))
_TOKENIZER_LOADS.set(created)
return created
@dataclass(slots=True)
class _TokenizerLoadSlot:
model: str
custom_tokenizer: CustomHuggingfaceTokenizer | None
current: _SharedTokenizerLoad | None = None
def load(self, loads: _TokenizerLoads, loop: asyncio.AbstractEventLoop) -> _SharedTokenizerLoad | None:
current: Final = self.current
if current is not None and current.reusable(loop.time()):
return current
self.current = loads.start(self.model, self.custom_tokenizer, loop)
return self.current
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
def _tokenizer_load_slot(
model: str, custom_tokenizer: _CustomTokenizerKey | None, loop: asyncio.AbstractEventLoop
) -> _TokenizerLoadSlot:
return _TokenizerLoadSlot(
model=model,
custom_tokenizer=None
if custom_tokenizer is None
else CustomHuggingfaceTokenizer(
identifier=custom_tokenizer.identifier,
revision=custom_tokenizer.revision,
auth_token=custom_tokenizer.auth_token,
),
)
async def _tokenizer_after_provider_error(
model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None
) -> SelectTokenizerResponse | BaseException:
loop: Final = asyncio.get_running_loop()
try:
load: Final = _tokenizer_load_slot(
model,
None
if custom_tokenizer is None
else _CustomTokenizerKey(
identifier=custom_tokenizer["identifier"],
revision=custom_tokenizer["revision"],
auth_token=custom_tokenizer["auth_token"],
),
loop,
).load(_tokenizer_loads_for_this_loop(), loop)
if load is None:
return _TokenizerLoadsBusy()
outcome: Final = await asyncio.wait_for(
asyncio.shield(load.outcome),
timeout=max(load.started_at + TOKENIZER_LOAD_WAIT_SECONDS - loop.time(), 0),
)
except Exception as e: # noqa: BLE001 # a tokenizer that cannot be read, loaded in time, or loaded at all answers the provider's error
return e
return outcome.tokenizer
async def _tokenizer_or_provider_error(
model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, provider_failure: ProviderCountFailure
) -> SelectTokenizerResponse:
tokenizer: Final = await _tokenizer_after_provider_error(model, custom_tokenizer)
if not isinstance(tokenizer, BaseException):
return tokenizer
if not _answers_provider_error(tokenizer):
raise tokenizer.with_traceback(None)
raise _provider_error_after_failed_estimate(provider_failure, tokenizer) from tokenizer
async def _local_token_estimate(
model: str,
prompt: str | None,
messages: Sequence[AllMessageValues] | None,
tools: list[ChatCompletionToolParam] | None,
custom_tokenizer: CustomHuggingfaceTokenizer | None,
provider_failure: ProviderCountFailure | None,
) -> _LocalTokenEstimate:
tokenizer: Final = (
await asyncify(litellm.utils._select_tokenizer)(model=model, custom_tokenizer=custom_tokenizer)
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=model,
text=prompt,
messages=messages,
tools=tools,
custom_tokenizer=tokenizer,
)
total_tokens: Final = await (
count if provider_failure is None else _estimate_or_provider_error(count, provider_failure)
)
return _LocalTokenEstimate(total_tokens=total_tokens, tokenizer_type=str(tokenizer["type"]))
_KIMI_K3_ENCODINGS: Final = KimiK3Encodings(
source=HUGGING_FACE_SOURCE, retry_after_seconds=TOKENIZER_LOAD_REUSE_SECONDS
)
async def _kimi_k3_count(
request: TokenCountRequest,
messages_format: TokenCountMessagesFormat,
model: str,
messages: Sequence[AllMessageValues],
tools: Sequence[ChatCompletionToolParam] | None,
) -> int | None:
encoding: Final = await _KIMI_K3_ENCODINGS.encoding(TOKENIZER_LOAD_WAIT_SECONDS)
if encoding is None:
return None
match messages_format:
case "anthropic":
return await offload_token_count(kimi_k3_anthropic_tokens)(
encoding, model, request.model_dump(include={"messages", "system", "tools"}, exclude_none=True)
)
case "openai":
return await offload_token_count(kimi_k3_tokens)(encoding, model, messages, tools)
class _TokenCountDeployment(TypedDict, total=False):
litellm_params: ReadOnly[Mapping[str, object]]
model_info: ReadOnly[Mapping[str, object]]
_TOKEN_COUNT_DEPLOYMENT: Final = TypeAdapter(_TokenCountDeployment)
def _model_group_deployment(router: Router | None, model_group: str) -> Mapping[str, object] | None:
group: Final = () if router is None else router._get_all_deployments(model_name=model_group)
return group[0] if group else None
def _non_empty_str(value: object) -> str | None:
return value if isinstance(value, str) and value else None
def _deployment_aws_region(deployment: _TokenCountDeployment) -> str:
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
params: Final = deployment.get("litellm_params", {})
return BaseAWSLLM()._get_aws_region_name(
{"aws_region_name": params.get("aws_region_name")}, model=_non_empty_str(params.get("model"))
)
def _token_count_settings_model(deployment: _TokenCountDeployment, request_model: str) -> tuple[str, str | None]:
model_info: Final = deployment.get("model_info", {})
params: Final = deployment.get("litellm_params", {})
base_model: Final = _non_empty_str(model_info.get("base_model")) or _non_empty_str(params.get("base_model"))
if base_model is not None:
return base_model, None
model: Final = _non_empty_str(params.get("model"))
if model is None:
return request_model, None
return model, _non_empty_str(params.get("custom_llm_provider"))
_PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset(
{litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value}
)
@ -13667,6 +14006,14 @@ def _drops_params(deployment: Mapping[str, object] | None) -> bool:
return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True)
_ESTIMABLE_GOOGLE_PART_KEYS: Final = frozenset({"text", "functionCall", "functionResponse", "thoughtSignature"})
def _estimable_contents(contents: Sequence[Mapping[str, object]]) -> bool:
parts: Final = chain.from_iterable(iter_items(content.get("parts") or ()) for content in contents)
return all(set(as_str_mapping(part) or ()) <= _ESTIMABLE_GOOGLE_PART_KEYS for part in parts)
def _system_message(system: object) -> ChatCompletionSystemMessage | None:
if not isinstance(system, (str, list)) or not system:
return None
@ -13735,6 +14082,14 @@ async def count_request_tokens(
model_to_use: str = (
litellm_model_name if isinstance(litellm_model_name, str) and litellm_model_name else request.model
)
deployment_view: Final = _TOKEN_COUNT_DEPLOYMENT.validate_python(
deployment or _model_group_deployment(llm_router, request.model) or {}
)
settings: Final = await asyncify(resolve_token_count_settings)(
configured=deployment_view.get("model_info", {}),
default=await asyncify(cost_map_entry)(*_token_count_settings_model(deployment_view, request.model)),
deployment_region=partial(_deployment_aws_region, deployment_view),
)
# Try provider-specific token counting first - only for non-direct requests (from provider endpoints)
provider_counter: BaseTokenCounter | None = None
@ -13754,23 +14109,32 @@ async def count_request_tokens(
code=400,
)
prompt_only: Final = bool(prompt) and conversation_empty
if provider_counter is not None:
prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES
result: Final = await _try_provider_token_count(
provider_counter=provider_counter,
custom_llm_provider=custom_llm_provider,
contents_estimable: Final = not contents or (bool(messages) and _estimable_contents(contents))
prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES
ladder: Final = (
None
if provider_counter is None
else await _provider_token_count(
count_with=partial(
_try_provider_token_count,
provider_counter,
custom_llm_provider,
messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages,
contents=None if prompt_as_message else contents,
deployment=deployment,
request_model=request.model,
tools=tools,
system=system,
messages_format=messages_format,
output_format=output_format,
),
model_to_use=model_to_use,
messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages,
contents=None if prompt_as_message else contents,
deployment=deployment,
request_model=request.model,
tools=tools,
system=system,
messages_format=messages_format,
output_format=output_format,
settings=settings,
local_estimate_allowed=litellm.disable_token_counter is not True and contents_estimable,
)
if result is not None:
return result
)
if isinstance(ladder, TokenCountResponse):
return ladder
# Check if token counter is disabled before fallback
if litellm.disable_token_counter is True:
@ -13788,11 +14152,6 @@ async def count_request_tokens(
CustomHuggingfaceTokenizer | None,
model_info.get("custom_tokenizer", None),
)
_tokenizer_used: Final = await asyncify(litellm.utils._select_tokenizer)(
model=model_to_use, custom_tokenizer=custom_tokenizer
)
tokenizer_used: Final = str(_tokenizer_used["type"])
system_message: Final = _system_message(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
@ -13805,18 +14164,38 @@ async def count_request_tokens(
counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats
list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None
)
total_tokens: Final = await offload_token_count(litellm.token_counter)(
framing: Final = settings.framing if counted_messages is not None and settings.framing is not None else 0
kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None
exact: Final = (
await _kimi_k3_count(request, messages_format, model_to_use, kimi_k3_messages, counted_tools)
if kimi_k3_messages is not None and custom_tokenizer is None
else None
)
if exact is not None:
return TokenCountResponse(
total_tokens=exact + framing,
request_model=request.model,
model_used=model_to_use,
tokenizer_type="huggingface_tokenizer",
count_source="local_estimate",
)
estimate: Final = await _local_token_estimate(
model=model_to_use,
text=prompt if counted_messages is None else None,
prompt=prompt if counted_messages is None else None,
messages=counted_messages,
tools=counted_tools,
custom_tokenizer=_tokenizer_used,
custom_tokenizer=custom_tokenizer,
provider_failure=ladder,
)
template: Final = (
0 if kimi_k3_messages is None else kimi_k3_template_tokens(len(kimi_k3_messages), bool(counted_tools))
)
return TokenCountResponse(
total_tokens=total_tokens,
total_tokens=estimate.total_tokens + template + framing,
request_model=request.model,
model_used=model_to_use,
tokenizer_type=tokenizer_used,
tokenizer_type=estimate.tokenizer_type,
count_source="local_estimate",
)

View file

@ -42,6 +42,7 @@ from litellm.proxy.common_utils.openai_error_payload import (
error_status_code,
openai_error_type_for_status,
)
from litellm.proxy.common_utils.token_count_settings import count_source_fields
from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
from litellm.types.llms.openai import (
REASONING_EFFORT,
@ -49,6 +50,7 @@ from litellm.types.llms.openai import (
ResponsesAPIResponse,
)
from litellm.types.responses.main import DeleteResponseResult
from litellm.types.token_count import TokenCountSourceFields
from litellm.types.utils import TokenCountResponse
if TYPE_CHECKING:
@ -1074,7 +1076,7 @@ class _ResponsesApiErrorBody(TypedDict):
error: ReadOnly[_ResponsesApiErrorDetail]
class _ResponsesInputTokensResult(TypedDict):
class _ResponsesInputTokensResult(TokenCountSourceFields):
object: ReadOnly[str]
input_tokens: ReadOnly[int]
@ -1289,6 +1291,7 @@ async def responses_input_tokens(
result: Final[_ResponsesInputTokensResult] = {
"object": "response.input_tokens",
"input_tokens": token_response.total_tokens,
**count_source_fields(token_response),
}
return result

View file

@ -7,6 +7,7 @@ from typing_extensions import (
)
from litellm.types.llms.openai import EmbeddingInput
from litellm.types.token_count import TokenCountSourceFields
# Gemini supports nested-list inputs (e.g. [["text", "image"]]) as an explicit
# opt-in for combined embeddings — a provider-specific extension of the
@ -308,12 +309,13 @@ class UsageMetadata(TypedDict, total=False):
candidatesTokensDetails: list[PromptTokensDetails] # Alternative key name used in some responses
class TokenCountDetailsResponse(TypedDict):
class TokenCountDetailsResponse(TokenCountSourceFields):
"""
Response structure for token count details with modality breakdown.
Example:
{'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]}
{'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}],
'litellm_count_source': 'provider'}
"""
totalTokens: int

View file

@ -0,0 +1,16 @@
from typing import Literal, TypeAlias
from typing_extensions import NotRequired, ReadOnly, TypedDict
TokenCountSource: TypeAlias = Literal["provider", "token_count_model", "local_estimate"]
TokenCountTokenizer: TypeAlias = Literal["kimi_k3"]
class TokenCountModelInfo(TypedDict):
model: ReadOnly[str]
region: ReadOnly[NotRequired[str]]
class TokenCountSourceFields(TypedDict, total=False):
litellm_count_source: ReadOnly[TokenCountSource]
litellm_count_model: ReadOnly[str]

View file

@ -55,6 +55,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_core_utils.core_helpers import map_finish_reason, process_response_headers
from . import litellm_params as _litellm_params
@ -211,6 +212,9 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
vertex_ai_audio_api: ReadOnly[Literal["lyria_predict", "lyria_interactions"] | None]
bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None
bedrock_converse_supports_strict_tools: bool | None
token_count_framing: ReadOnly[int | None]
token_count_model: ReadOnly[TokenCountModelInfo | None]
token_count_tokenizer: ReadOnly[TokenCountTokenizer | None]
class SearchContextCostPerQuery(TypedDict, total=False):
@ -3826,10 +3830,13 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"})
DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset({"token_count_framing", "token_count_model"})
SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = (
frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__)
- frozenset(CustomPricingLiteLLMParams.model_fields)
- DEPLOYMENT_SCOPED_PRICING_FIELDS
- DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS
)
@ -4226,6 +4233,7 @@ class TokenCountResponse(LiteLLMPydanticObjectBase):
request_model: str
model_used: str
tokenizer_type: str
count_source: TokenCountSource | None = None
original_response: dict | None = None
"""
Original Response from upstream API call - if an API call was made for token counting
@ -4236,6 +4244,7 @@ class TokenCountResponse(LiteLLMPydanticObjectBase):
HTTP status code from the token counting API (e.g., 200 for success, 429 for rate limit, 400 for bad request)
"""
status_code: int | None = None
rejected_before_provider_call: bool = Field(default=False, exclude=True)
class CustomHuggingfaceTokenizer(TypedDict):

View file

@ -6327,6 +6327,9 @@ def _get_model_info_helper(
uses_embed_content=_model_info.get("uses_embed_content", None),
supports_image_size=_model_info.get("supports_image_size", None),
supported_audio_formats=_model_info.get("supported_audio_formats", None),
token_count_framing=_model_info.get("token_count_framing", None),
token_count_model=_model_info.get("token_count_model", None),
token_count_tokenizer=_model_info.get("token_count_tokenizer", None),
vertex_ai_audio_api=_model_info.get("vertex_ai_audio_api", None),
)
for cost_key, cost_value in _model_info.items():

View file

@ -19337,7 +19337,13 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
"supports_vision": false,
"token_count_framing": 97
},
"databricks/system.ai.deepseek-v4-flash-0731": {
"litellm_provider": "databricks",
"mode": "chat",
"token_count_framing": 97
},
"databricks/databricks-deepseek-v4-pro-0813": {
"cache_creation_input_token_cost": 1.31999e-06,
@ -20328,7 +20334,13 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"token_count_framing": 100
},
"databricks/system.ai.kimi-k3": {
"litellm_provider": "databricks",
"mode": "chat",
"token_count_framing": 100
},
"databricks/databricks-llama-4-maverick": {
"cache_creation_input_token_cost": 5.0001e-07,
@ -75642,7 +75654,9 @@
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"token_count_framing": 64,
"token_count_tokenizer": "kimi_k3"
},
"global.moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 3.75e-06,
@ -75662,7 +75676,9 @@
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"token_count_framing": 64,
"token_count_tokenizer": "kimi_k3"
},
"us.moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 4.125e-06,
@ -75682,7 +75698,9 @@
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"token_count_framing": 64,
"token_count_tokenizer": "kimi_k3"
},
"openrouter/xiaomi/mimo-v2.6-flash": {
"cache_read_input_token_cost": 2.8e-09,

View file

@ -1083,6 +1083,36 @@
"additionalProperties": false
}
},
"token_count_framing": {
"type": "integer",
"minimum": 0,
"description": "Tokens added once per request to a local count estimate, for the provider's chat framing."
},
"token_count_model": {
"type": "object",
"required": [
"model"
],
"additionalProperties": false,
"properties": {
"model": {
"type": "string",
"pattern": "\\S"
},
"region": {
"type": "string",
"pattern": "\\S"
}
},
"description": "Model that counts tokens for this one, and the AWS region its count API runs in."
},
"token_count_tokenizer": {
"type": "string",
"description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.",
"enum": [
"kimi_k3"
]
},
"tpm": {
"type": "integer",
"minimum": 0,

View file

@ -280,7 +280,7 @@ def test_messages_count_tokens_lifts_the_leading_system_run(
) -> None:
with _listening(deployment) as peer:
response: Final = _count(gateway, count_request(deployment.model, case))
assert _payload(response) == {"input_tokens": _COUNT}, response.text
assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text
assert _bodies(peer, deployment.provider) == (expected_count_body(_MODEL, case),)
@ -335,7 +335,11 @@ def test_responses_input_tokens_lifts_instructions(deployment: _Deployment, gate
"/v1/responses/input_tokens",
{"model": deployment.model, "input": USER_TEXT, "instructions": INSTRUCTION},
)
assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _COUNT}, response.text
assert _payload(response) == {
"object": "response.input_tokens",
"input_tokens": _COUNT,
"litellm_count_source": "provider",
}, response.text
assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
@ -348,7 +352,11 @@ def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_ite
"/v1/responses/input_tokens",
{"model": deployment.model, "input": [MID_SYSTEM, USER], "instructions": INSTRUCTION},
)
assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _COUNT}, response.text
assert _payload(response) == {
"object": "response.input_tokens",
"input_tokens": _COUNT,
"litellm_count_source": "provider",
}, response.text
assert _bodies(peer, deployment.provider) == (
{"model": _MODEL, "messages": [USER], "system": [LIFTED, {"type": "text", "text": REMINDER}]},
)
@ -379,7 +387,7 @@ def test_messages_count_tokens_forwards_tools_beside_the_lifted_system(
) -> None:
with _listening(deployment) as peer:
response: Final = _count(gateway, {**count_request(deployment.model, STRING_CASE), "tools": TOOLS})
assert _payload(response) == {"input_tokens": _COUNT}, response.text
assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text
assert _bodies(peer, deployment.provider) == ({**_STRING_BODY, "tools": TOOLS},)
@ -392,7 +400,7 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place(
"messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP],
}
response: Final = _count(gateway, body)
assert _payload(response) == {"input_tokens": _COUNT}, response.text
assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text
assert _bodies(peer, deployment.provider) == (
{"model": _MODEL, "messages": [USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], "system": [LIFTED]},
)
@ -424,7 +432,7 @@ def test_messages_count_tokens_repeated_request_lifts_each_time(deployment: _Dep
answers: Final = tuple(
_payload(_count(gateway, count_request(deployment.model, STRING_CASE))) for _ in range(2)
)
assert answers == ({"input_tokens": _COUNT},) * 2
assert answers == ({"input_tokens": _COUNT, "litellm_count_source": "provider"},) * 2
assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * 2
@ -438,7 +446,7 @@ def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_
assert "input_tokens" not in response.text, response.text
assert peer.drain() == ()
follow_up: Final = _count(gateway, count_request(deployment.model, STRING_CASE))
assert _payload(follow_up) == {"input_tokens": _COUNT}, follow_up.text
assert _payload(follow_up) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, follow_up.text
assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
@ -453,7 +461,7 @@ def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value(
content=f'{{"model": "{deployment.model}", "messages": {first}, "messages": {last}}}',
headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
)
assert _payload(response) == {"input_tokens": _COUNT}, response.text
assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text
assert _bodies(peer, deployment.provider) == (_STRING_BODY,)

View file

@ -385,7 +385,7 @@ def test_messages_count_tokens_counts_through_mantle_when_the_runtime_cannot(
):
model: Final = _deployment(scenario, runtime.url)
response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS})
assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",)
assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),)
@ -401,7 +401,7 @@ def test_messages_count_tokens_without_system_or_tools_sends_a_bare_body_to_mant
):
model: Final = _deployment(scenario, runtime.url)
response: Final = _count(gateway, _bare(model))
assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == (_MANTLE_BARE,)
@ -436,7 +436,11 @@ def test_gemini_count_tokens_route_reports_the_mantle_total(counting_proxy: Owne
):
model: Final = _deployment(scenario, runtime.url)
response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY)
assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text
assert _payload(response) == {
"totalTokens": _MANTLE_COUNT,
"promptTokensDetails": [],
"litellm_count_source": "provider",
}, response.text
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == ({"model": _OPUS_BASE, "messages": _GEMINI_MESSAGES},)
@ -452,7 +456,11 @@ def test_gemini_count_tokens_route_reports_the_runtime_total_for_a_model_bedrock
):
model: Final = _deployment(scenario, runtime.url, _SONNET)
response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY)
assert _payload(response) == {"totalTokens": _RUNTIME_COUNT, "promptTokensDetails": []}, response.text
assert _payload(response) == {
"totalTokens": _RUNTIME_COUNT,
"promptTokensDetails": [],
"litellm_count_source": "provider",
}, response.text
assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",)
assert mantle.drain() == ()
@ -467,7 +475,7 @@ def test_runtime_404_on_claude_falls_back_to_mantle(counting_proxy: OwnedProxy,
):
model: Final = _deployment(scenario, runtime.url)
response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS})
assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),)
@ -527,7 +535,11 @@ def test_google_count_tokens_forwards_system_and_tools_to_mantle(
):
model: Final = _deployment(scenario, runtime.url)
response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body)
assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text
assert _payload(response) == {
"totalTokens": _MANTLE_COUNT,
"promptTokensDetails": [],
"litellm_count_source": "provider",
}, response.text
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == (
_mantle_body(
@ -572,7 +584,11 @@ def test_google_count_tokens_normalizes_uppercase_parameters_for_mantle(
):
model: Final = _deployment(scenario, runtime.url)
response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body)
assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text
assert _payload(response) == {
"totalTokens": _MANTLE_COUNT,
"promptTokensDetails": [],
"litellm_count_source": "provider",
}, response.text
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == (
_mantle_body(
@ -602,7 +618,11 @@ def test_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy
response: Final = gateway.request(
"POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"}
)
assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text
assert _payload(response) == {
"object": "response.input_tokens",
"input_tokens": _MANTLE_COUNT,
"litellm_count_source": "provider",
}, response.text
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == (_MANTLE_BARE,)
@ -647,7 +667,10 @@ def test_model_the_runtime_counts_never_reaches_mantle(counting_proxy: OwnedProx
gateway.scenario() as scenario,
):
model: Final = _deployment(scenario, runtime.url, _SONNET)
assert _payload(_count(gateway, _full(model))) == {"input_tokens": _RUNTIME_COUNT}
assert _payload(_count(gateway, _full(model))) == {
"input_tokens": _RUNTIME_COUNT,
"litellm_count_source": "provider",
}
detailed: Final = gateway.request(
"POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"}
)
@ -735,7 +758,11 @@ def test_responses_input_tokens_with_instructions_counts_through_mantle(
"/v1/responses/input_tokens",
{"model": model, "input": "Count this message", "instructions": "Be terse"},
)
assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text
assert _payload(response) == {
"object": "response.input_tokens",
"input_tokens": _MANTLE_COUNT,
"litellm_count_source": "provider",
}, response.text
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == (_mantle_body(system=[{"type": "text", "text": "Be terse"}]),)
@ -752,7 +779,7 @@ def test_messages_count_tokens_lifts_the_leading_system_run_for_mantle(
):
model: Final = _deployment(scenario, runtime.url)
response: Final = _count(gateway, count_request(model, case))
assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == (expected_count_body(_OPUS_BASE, case),)
@ -865,7 +892,11 @@ def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_ite
"/v1/responses/input_tokens",
{"model": model, "input": [MID_SYSTEM, USER], "instructions": INSTRUCTION},
)
assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text
assert _payload(response) == {
"object": "response.input_tokens",
"input_tokens": _MANTLE_COUNT,
"litellm_count_source": "provider",
}, response.text
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == (_mantle_body(system=[LIFTED, {"type": "text", "text": REMINDER}]),)
@ -881,7 +912,7 @@ def test_messages_count_tokens_forwards_tools_beside_the_lifted_system_to_mantle
):
model: Final = _deployment(scenario, runtime.url)
response: Final = _count(gateway, {**count_request(model, STRING_CASE), "tools": _TOOLS})
assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == (_mantle_body(system=[LIFTED], tools=_TOOLS),)
@ -901,7 +932,7 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place_for_mant
"messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP],
}
response: Final = _count(gateway, body)
assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == (
_mantle_body(messages=[USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], system=[LIFTED]),
@ -954,7 +985,7 @@ def test_messages_count_tokens_repeated_request_lifts_the_leading_system_each_ti
):
model: Final = _deployment(scenario, runtime.url)
answers: Final = tuple(_payload(_count(gateway, count_request(model, STRING_CASE))) for _ in range(2))
assert answers == ({"input_tokens": _MANTLE_COUNT},) * 2
assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2
assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2
assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) * 2
@ -975,7 +1006,9 @@ def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_
assert "input_tokens" not in response.text, response.text
assert mantle.drain() == ()
follow_up: Final = _count(gateway, count_request(model, STRING_CASE))
assert _payload(follow_up) == {"input_tokens": _MANTLE_COUNT}, follow_up.text
assert _payload(follow_up) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, (
follow_up.text
)
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,)
@ -997,7 +1030,7 @@ def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value_for_m
content=f'{{"model": "{model}", "messages": {first}, "messages": {last}}}',
headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
)
assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,)
@ -1022,7 +1055,9 @@ def test_disabled_token_counter_counts_a_leading_system_through_mantle(
)
with wire_server(_mantle(_strict), port=mantle_port) as strict:
counted: Final = _count(owned.gateway, count_request(_OWNED_OPUS, STRING_CASE))
assert _payload(counted) == {"input_tokens": _MANTLE_COUNT}, counted.text
assert _payload(counted) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, (
counted.text
)
assert _mantle_bodies(strict) == (_MANTLE_LIFTED,)
assert len(_runtime_count_targets(runtime)) == 1
@ -1222,7 +1257,7 @@ def test_messages_count_tokens_forwards_a_5kb_system_verbatim_to_mantle(
):
model: Final = _deployment(scenario, runtime.url)
response: Final = _count(gateway, {**_bare(model), "system": system})
assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == (_mantle_body(system=system),)
@ -1243,7 +1278,7 @@ def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value
content=f'{{"model": "{model}", "messages": {json.dumps(_MESSAGES)}, {fields}, {fields}}}',
headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
)
assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
assert len(_runtime_count_targets(runtime)) == 1
(sent,) = mantle.drain()
assert _mantle_requests((sent,)) == (_MANTLE_FULL,)
@ -1261,7 +1296,7 @@ def test_messages_count_tokens_leaves_empty_tools_and_system_out_of_the_mantle_b
):
model: Final = _deployment(scenario, runtime.url)
response: Final = _count(gateway, {**_bare(model), "tools": [], "system": ""})
assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
assert len(_runtime_count_targets(runtime)) == 1
assert _mantle_bodies(mantle) == (_MANTLE_BARE,)
@ -1277,7 +1312,7 @@ def test_messages_count_tokens_repeated_request_reaches_both_peers_each_time(
):
model: Final = _deployment(scenario, runtime.url)
answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2))
assert answers == ({"input_tokens": _MANTLE_COUNT},) * 2
assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2
assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2
assert _mantle_bodies(mantle) == (_MANTLE_BARE,) * 2
@ -1306,9 +1341,15 @@ def test_disabled_token_counter_surfaces_the_mantle_error_instead_of_counting_lo
assert _REJECTION in refused.text and "input_tokens" not in refused.text, refused.text
assert _mantle_bodies(refusing) == (_MANTLE_FULL,)
with wire_server(_mantle(), port=mantle_port) as counting:
assert _payload(_count(owned.gateway, _full(_OWNED_OPUS))) == {"input_tokens": _MANTLE_COUNT}
assert _payload(_count(owned.gateway, _full(_OWNED_OPUS))) == {
"input_tokens": _MANTLE_COUNT,
"litellm_count_source": "provider",
}
assert _mantle_bodies(counting) == (_MANTLE_FULL,)
assert _payload(_count(owned.gateway, _full(_OWNED_SONNET))) == {"input_tokens": _RUNTIME_COUNT}
assert _payload(_count(owned.gateway, _full(_OWNED_SONNET))) == {
"input_tokens": _RUNTIME_COUNT,
"litellm_count_source": "provider",
}
rejected: Final = _count(owned.gateway, _full(_OWNED_NOVA))
assert rejected.status_code == 400, rejected.text
assert _UNSUPPORTED in rejected.text and "input_tokens" not in rejected.text, rejected.text

View file

@ -399,7 +399,7 @@ def test_messages_count_tokens_forwards_system_and_tools(gateway: Gateway, syste
with wire_server(_peer()) as wire, gateway.scenario() as scenario:
model: Final = _deployment(gateway, scenario, wire.url)
response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS})
assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
assert _counted_bodies(wire) == ({**_PEER_BARE, "system": system, "tools": _TOOLS},)
@ -407,7 +407,7 @@ def test_messages_count_tokens_without_system_or_tools_sends_bare_body(gateway:
with wire_server(_peer()) as wire, gateway.scenario() as scenario:
model: Final = _deployment(gateway, scenario, wire.url)
response: Final = _count(gateway, _bare(model))
assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
assert _counted_bodies(wire) == (_PEER_BARE,)
@ -446,7 +446,11 @@ def test_responses_input_tokens_counts_through_the_partner_peer(gateway: Gateway
response: Final = gateway.request(
"POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"}
)
assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _PEER_COUNT}, response.text
assert _payload(response) == {
"object": "response.input_tokens",
"input_tokens": _PEER_COUNT,
"litellm_count_source": "provider",
}, response.text
assert _counted_bodies(wire) == (_PEER_BARE,)
@ -501,7 +505,7 @@ def test_llama_count_still_reaches_the_meta_count_route(gateway: Gateway) -> Non
with wire_server(_peer()) as wire, gateway.scenario() as scenario:
model: Final = _deployment(gateway, scenario, wire.url, model=f"vertex_ai/{_LLAMA_BACKEND}")
response: Final = _count(gateway, _bare(model))
assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
(sent,) = _count_bodies(wire.drain(), _LLAMA_COUNT_TARGET)
assert sent["messages"] == _MESSAGES, sent
@ -619,7 +623,7 @@ def test_messages_count_tokens_forwards_an_empty_tools_list(gateway: Gateway) ->
with wire_server(_peer()) as wire, gateway.scenario() as scenario:
model: Final = _deployment(gateway, scenario, wire.url)
response: Final = _count(gateway, {**_bare(model), "tools": []})
assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
assert _counted_bodies(wire) == ({**_PEER_BARE, "tools": []},)
@ -627,7 +631,7 @@ def test_messages_count_tokens_forwards_an_empty_system_string(gateway: Gateway)
with wire_server(_peer()) as wire, gateway.scenario() as scenario:
model: Final = _deployment(gateway, scenario, wire.url)
response: Final = _count(gateway, {**_bare(model), "system": ""})
assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
assert _counted_bodies(wire) == ({**_PEER_BARE, "system": ""},)
@ -635,7 +639,7 @@ def test_messages_count_tokens_leaves_null_system_and_tools_out(gateway: Gateway
with wire_server(_peer()) as wire, gateway.scenario() as scenario:
model: Final = _deployment(gateway, scenario, wire.url)
response: Final = _count(gateway, {**_bare(model), "system": None, "tools": None})
assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
assert _counted_bodies(wire) == (_PEER_BARE,)
@ -653,7 +657,7 @@ def test_messages_count_tokens_forwards_a_5kb_system_verbatim(gateway: Gateway)
with wire_server(_peer()) as wire, gateway.scenario() as scenario:
model: Final = _deployment(gateway, scenario, wire.url)
response: Final = _count(gateway, {**_bare(model), "system": system})
assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
assert _counted_bodies(wire) == ({**_PEER_BARE, "system": system},)
@ -666,7 +670,7 @@ def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value
content=f'{{"model": "{model}", "messages": {json.dumps(_MESSAGES)}, {fields}, {fields}}}',
headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
)
assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
(sent,) = _count_requests(wire.drain())
assert _count_bodies((sent,)) == (_PEER_FULL,)
assert (sent.body.count(b'"system"'), sent.body.count(b'"tools"')) == (1, 1), sent.body
@ -699,7 +703,7 @@ def test_count_tokens_location_override_targets_the_count_region(gateway: Gatewa
gateway, scenario, wire.url, vertex_location="global", vertex_count_tokens_location="europe-west1"
)
response: Final = _count(gateway, _bare(model))
assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
target: Final = _COUNT_TARGET.replace(f"/locations/{_LOCATION}/", "/locations/europe-west1/")
assert _count_bodies(wire.drain(), target) == (_PEER_BARE,)
@ -708,7 +712,7 @@ def test_messages_count_tokens_repeated_request_reaches_the_peer_each_time(gatew
with wire_server(_peer()) as wire, gateway.scenario() as scenario:
model: Final = _deployment(gateway, scenario, wire.url)
answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2))
assert answers == ({"input_tokens": _PEER_COUNT},) * 2
assert answers == ({"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"},) * 2
assert _counted_bodies(wire) == (_PEER_BARE,) * 2
@ -725,7 +729,7 @@ def test_disabled_token_counter_surfaces_provider_failures_instead_of_counting_l
)
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
counted: Final = _count(owned.gateway, _full(_OWNED_MODEL))
assert _payload(counted) == {"input_tokens": _PEER_COUNT}, counted.text
assert _payload(counted) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, counted.text
rejected: Final = _count(
owned.gateway, {**_full(_OWNED_MODEL), "messages": [{"role": "user", "content": _REJECT_TEXT}]}
)

View file

@ -399,9 +399,11 @@ def test_in_flight_count_tokens_does_not_reserve_key_budget_away_from_a_completi
[{"role": "user", "content": prompt}]
]
assert count.status_code == 200, count.text
assert count.json() == {"totalTokens": 12, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}]}, (
count.text
)
assert count.json() == {
"totalTokens": 12,
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}],
"litellm_count_source": "provider",
}, count.text
provider_calls: Final = wire.drain()
assert [(call.method, call.target) for call in provider_calls] == [
("POST", "/v1beta/models/gemini-3.8-flash:countTokens")

View file

@ -4,10 +4,11 @@ Test to verify the Google GenAI proxy API endpoints
"""
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, patch
import pytest
from litellm.types.utils import TokenCountResponse
def _build_test_client():
@ -88,9 +89,7 @@ def test_google_stream_generate_content_endpoint():
assert init_kwargs["data"]["_litellm_raw_sse_stream"] is True
assert init_kwargs["data"]["_litellm_skip_openai_stream_done"] is True
assert init_kwargs["data"]["model"] == "test-model"
assert init_kwargs["data"]["contents"] == [
{"role": "user", "parts": [{"text": "Hello"}]}
]
assert init_kwargs["data"]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}]
def test_google_generate_content_data_flows_through_processor():
@ -123,9 +122,7 @@ def test_google_generate_content_data_flows_through_processor():
data = mock_init.call_args.kwargs["data"]
assert data["model"] == "test-model"
assert data["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}]
assert data["systemInstruction"] == {
"parts": [{"text": "Your name is Doodle."}]
}
assert data["systemInstruction"] == {"parts": [{"text": "Your name is Doodle."}]}
# generationConfig arrives intact here; the rename to `config` is
# done downstream in route_request (see test_route_llm_request).
assert data["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"]
@ -160,12 +157,17 @@ def test_google_count_tokens_unchanged():
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
fake_response = MagicMock()
fake_response.original_response = {
"totalTokens": 7,
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}],
}
fake_response.total_tokens = 7
fake_response = TokenCountResponse(
total_tokens=7,
request_model="test-model",
model_used="test-model",
tokenizer_type="gemini_api",
original_response={
"totalTokens": 7,
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}],
},
count_source="provider",
)
with patch(
"litellm.proxy.proxy_server.token_counter",
@ -180,6 +182,7 @@ def test_google_count_tokens_unchanged():
assert response.status_code == 200
body = response.json()
assert body["totalTokens"] == 7
assert body["litellm_count_source"] == "provider"
def test_google_count_tokens_unexpected_failure_uses_the_google_error_envelope():
@ -212,9 +215,14 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs()
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
fake_response: Final = MagicMock()
fake_response.original_response = {"input_tokens": 2167}
fake_response.total_tokens = 2167
fake_response: Final = TokenCountResponse(
total_tokens=2167,
request_model="claude-opus-4-8",
model_used="us.anthropic.claude-opus-5-5",
tokenizer_type="bedrock_api",
original_response={"input_tokens": 2167},
count_source="token_count_model",
)
with patch(
"litellm.proxy.proxy_server.token_counter",
@ -230,3 +238,5 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs()
body: Final = response.json()
assert body["totalTokens"] == 2167
assert body["promptTokensDetails"] == []
assert body["litellm_count_source"] == "token_count_model"
assert body["litellm_count_model"] == "us.anthropic.claude-opus-5-5"

View file

@ -4239,6 +4239,54 @@ class TestModelInfoCostMapEchoFilter:
assert "litellm_provider" not in info
assert "supported_openai_params" not in info
@pytest.fixture
def bedrock_sonnet_5_5_mapped_as_on_main(self) -> Iterator[None]:
import litellm
entry: Final = {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "us-east-1"},
}
with pytest.MonkeyPatch.context() as patched:
patched.setattr(litellm, "model_cost", {**litellm.model_cost, "us.anthropic.claude-sonnet-5-5": entry})
litellm.utils._invalidate_model_cost_lowercase_map()
yield
litellm.utils._invalidate_model_cost_lowercase_map()
@pytest.mark.parametrize(
("model", "field", "edited"),
[
("bedrock/us.anthropic.claude-sonnet-5-5", "token_count_model", {}),
("bedrock/us.moonshotai.kimi-k3", "token_count_framing", 0),
],
)
def test_token_count_setting_is_stored_only_when_the_operator_changed_it(
self, model: str, field: str, edited: object, bedrock_sonnet_5_5_mapped_as_on_main: None
) -> None:
import litellm
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
entry: Final = litellm.get_model_info(model)
assert entry[field] not in (None, edited), entry
db_model: Final = Deployment(
model_name="count-settings",
litellm_params=LiteLLM_Params(model=model),
model_info=ModelInfo(id="dep-echo-count"),
)
echoed: Final = update_db_model(
db_model=db_model,
updated_patch=updateDeployment(model_info=ModelInfo(**{**entry, "id": "dep-echo-count"})),
)
changed: Final = update_db_model(
db_model=db_model,
updated_patch=updateDeployment(model_info=ModelInfo(**{**entry, "id": "dep-echo-count", field: edited})),
)
assert field not in json.loads(echoed["model_info"])
assert json.loads(changed["model_info"])[field] == edited
def test_metadata_without_a_cost_map_key_is_persisted(self):
import litellm

View file

@ -52,6 +52,7 @@ def test_token_counter_happy_path(client, auth_as, patched_token_counter):
"request_model": "gpt-4",
"model_used": "gpt-4",
"tokenizer_type": "openai_tokenizer",
"count_source": "local_estimate",
"original_response": None,
"error": False,
"error_message": None,

View file

@ -54,6 +54,7 @@ async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload(
"request_model": "azure_ai/claude-sonnet-4-5",
"model_used": "claude-sonnet-4-5",
"tokenizer_type": "azure_ai_anthropic_api",
"count_source": None,
"original_response": payload,
"error": False,
"error_message": None,
@ -98,6 +99,7 @@ async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_v
"request_model": "azure_ai/claude-sonnet-4-5",
"model_used": "claude-sonnet-4-5",
"tokenizer_type": "azure_ai_anthropic_api",
"count_source": None,
"original_response": None,
"error": True,
"error_message": str(rejection.value),
@ -126,6 +128,7 @@ async def test_count_payload_that_is_not_an_object_becomes_an_error_response_nam
"request_model": "azure_ai/claude-sonnet-4-5",
"model_used": "claude-sonnet-4-5",
"tokenizer_type": "azure_ai_anthropic_api",
"count_source": None,
"original_response": None,
"error": True,
"error_message": str(failed_lookup.value),
@ -156,6 +159,7 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_
"request_model": "azure_ai/claude-sonnet-4-5",
"model_used": "claude-sonnet-4-5",
"tokenizer_type": "azure_ai_anthropic_api",
"count_source": None,
"original_response": None,
"error": True,
"error_message": upstream_text,
@ -180,6 +184,7 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_
"request_model": "azure_ai/claude-sonnet-4-5",
"model_used": "claude-sonnet-4-5",
"tokenizer_type": "azure_ai_anthropic_api",
"count_source": None,
"original_response": None,
"error": True,
"error_message": f"CountTokens processing error: {undecodable.value}",

View file

@ -0,0 +1,524 @@
import asyncio
import base64
import hashlib
import json
import logging
import re
import threading
from collections.abc import Iterable
from itertools import product
from pathlib import Path
from typing import Final
import httpx
import pytest
import respx
import tiktoken
import litellm
from litellm._logging import verbose_logger
from litellm.llms.bedrock.count_tokens.kimi_k3 import (
KimiK3Encodings,
KimiK3TokenizerSource,
kimi_k3_anthropic_tokens,
kimi_k3_tokens,
load_kimi_k3_encoding,
)
from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig
MODEL: Final = "us.moonshotai.kimi-k3"
_CONTROL: Final = re.compile(r"<\|(?:open|close|sep|end_of_msg)\|>")
_WEATHER: Final = {
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the weather in °C",
"parameters": {
"type": "object",
"properties": {"city": {"type": "string"}, "days": {"type": "integer"}},
"required": ["city"],
},
},
}
def _call(call_id: str, arguments: dict[str, object]) -> dict[str, object]:
return {"id": call_id, "type": "function", "function": {"name": "get_weather", "arguments": json.dumps(arguments)}}
def _written_ranks(path: Path, tokens: Iterable[bytes]) -> KimiK3TokenizerSource:
path.write_bytes(b"\n".join(base64.b64encode(token) + b" %d" % rank for rank, token in enumerate(tokens)))
return KimiK3TokenizerSource(ranks_path=lambda: str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest())
def _single_bytes() -> tuple[bytes, ...]:
return tuple(bytes((value,)) for value in range(256))
@pytest.fixture(scope="module")
def byte_source(tmp_path_factory: pytest.TempPathFactory) -> KimiK3TokenizerSource:
return _written_ranks(tmp_path_factory.mktemp("bytes") / "tiktoken.model", _single_bytes())
@pytest.fixture(scope="module")
def byte_encoding(byte_source: KimiK3TokenizerSource) -> tiktoken.Encoding:
loaded: Final = load_kimi_k3_encoding(byte_source)
assert loaded is not None
return loaded
@pytest.fixture(scope="module")
def pair_encoding(tmp_path_factory: pytest.TempPathFactory) -> tiktoken.Encoding:
pairs: Final = (bytes(pair) for pair in product(range(256), repeat=2))
source: Final = _written_ranks(tmp_path_factory.mktemp("pairs") / "tiktoken.model", (*_single_bytes(), *pairs))
loaded: Final = load_kimi_k3_encoding(source)
assert loaded is not None
return loaded
def _template_count(encoding: tiktoken.Encoding, rendered: str) -> int:
texts: Final = _CONTROL.split(rendered)
return len(_CONTROL.findall(rendered)) + sum(len(encoding.encode_ordinary(text)) for text in texts)
_GENERATION_PROMPT: Final = '<|open|>message role="assistant"<|sep|><|open|>think<|sep|>'
# Expected renders follow the Kimi K3 chat template in moonshotai/Kimi-K3 at revision f831ab6, which matched
# Bedrock's prompt_tokens for us.moonshotai.kimi-k3 in 19 live calls on 2026-10-16
@pytest.mark.parametrize(
("messages", "tools", "rendered"),
[
pytest.param(
[{"role": "user", "content": "Hello there"}],
None,
'<|open|>message role="user"<|sep|>Hello there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT,
id="user",
),
pytest.param(
[
{
"role": "system",
"content": [{"type": "text", "text": "Be terse."}, {"type": "text", "text": "Answer in French."}],
},
{"role": "user", "content": [{"type": "text", "text": "Hi"}, {"type": "text", "text": " there"}]},
],
None,
'<|open|>message role="system"<|sep|>Be terse.Answer in French.<|close|>message<|sep|><|end_of_msg|>'
'<|open|>message role="user"<|sep|>Hi there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT,
id="text_parts_join",
),
pytest.param(
[
{"role": "user", "content": "What is 2+2?"},
{"role": "assistant", "content": "4"},
{"role": "user", "content": "And 3+3?"},
],
None,
'<|open|>message role="user"<|sep|>What is 2+2?<|close|>message<|sep|><|end_of_msg|>'
'<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>'
"<|open|>response<|sep|>4<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|>"
'<|open|>message role="user"<|sep|>And 3+3?<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT,
id="assistant_turn",
),
pytest.param(
[
{"role": "user", "content": "Weather in Paris and Oslo?"},
{
"role": "assistant",
"content": None,
"tool_calls": [_call("c1", {"city": "Paris", "days": 3}), _call("c2", {"city": "Oslo"})],
},
{"role": "tool", "tool_call_id": "c2", "content": "Snow"},
{"role": "tool", "tool_call_id": "c1", "content": "Sun"},
],
[_WEATHER],
'<|open|>message role="system" type="tool-declare"<|sep|># Tools\n'
"Here are the available tools, described in JSONSchema.\n\n```json\n"
'[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":'
'{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},'
'"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>'
'<|open|>message role="user"<|sep|>Weather in Paris and Oslo?<|close|>message<|sep|><|end_of_msg|>'
'<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>'
"<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>"
'<|open|>call tool="get_weather" index="1"<|sep|>'
'<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|>'
'<|open|>argument key="days" type="number"<|sep|>3<|close|>argument<|sep|><|close|>call<|sep|>'
'<|open|>call tool="get_weather" index="2"<|sep|>'
'<|open|>argument key="city" type="string"<|sep|>Oslo<|close|>argument<|sep|><|close|>call<|sep|>'
"<|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>"
'<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>'
'<|open|>message role="tool" tool="get_weather" index="2"<|sep|>Snow<|close|>message<|sep|><|end_of_msg|>'
+ _GENERATION_PROMPT,
id="tool_results_out_of_call_order",
),
pytest.param(
[
{"role": "user", "content": "Weather?"},
{
"role": "assistant",
"content": None,
"tool_calls": [
_call(
"c1", {"a&b": 'x"y', "flag": True, "none": None, "list": [1, "zé"], "object": {"k": True}}
)
],
},
{"role": "tool", "tool_call_id": "c1", "content": "Rain"},
],
[_WEATHER],
'<|open|>message role="system" type="tool-declare"<|sep|># Tools\n'
"Here are the available tools, described in JSONSchema.\n\n```json\n"
'[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":'
'{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},'
'"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>'
'<|open|>message role="user"<|sep|>Weather?<|close|>message<|sep|><|end_of_msg|>'
'<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>'
"<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>"
'<|open|>call tool="get_weather" index="1"<|sep|>'
'<|open|>argument key="a&amp;b" type="string"<|sep|>x"y<|close|>argument<|sep|>'
'<|open|>argument key="flag" type="boolean"<|sep|>true<|close|>argument<|sep|>'
'<|open|>argument key="none" type="null"<|sep|>null<|close|>argument<|sep|>'
'<|open|>argument key="list" type="array"<|sep|>[1, "zé"]<|close|>argument<|sep|>'
'<|open|>argument key="object" type="object"<|sep|>{"k": true}<|close|>argument<|sep|>'
"<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>"
'<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Rain<|close|>message<|sep|><|end_of_msg|>'
+ _GENERATION_PROMPT,
id="argument_types_and_escaping",
),
],
)
@pytest.mark.parametrize("encoding_fixture", ["byte_encoding", "pair_encoding"])
def test_count_is_the_kimi_k3_chat_template_render_of_the_bedrock_request(
request: pytest.FixtureRequest,
encoding_fixture: str,
messages: list[dict[str, object]],
tools: list[dict[str, object]] | None,
rendered: str,
):
encoding: Final[tiktoken.Encoding] = request.getfixturevalue(encoding_fixture)
assert kimi_k3_tokens(encoding, MODEL, messages, tools) == _template_count(encoding, rendered)
def _parallel_calls(count: int) -> list[dict[str, object]]:
ids: Final = tuple(f"c{number}" for number in range(1, count + 1))
return [
{"role": "user", "content": "Weather?"},
{"role": "assistant", "content": None, "tool_calls": [_call(call_id, {"city": "Paris"}) for call_id in ids]},
*({"role": "tool", "tool_call_id": call_id, "content": "Sun"} for call_id in ids),
]
def test_tenth_parallel_call_and_its_result_are_numbered_ten(byte_encoding: tiktoken.Encoding):
nine: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(9), [_WEATHER])
ten: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(10), [_WEATHER])
assert nine is not None and ten is not None
assert ten - nine == _template_count(
byte_encoding,
'<|open|>call tool="get_weather" index="10"<|sep|>'
'<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|><|close|>call<|sep|>'
'<|open|>message role="tool" tool="get_weather" index="10"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>',
)
@pytest.mark.parametrize(
"text",
[
pytest.param("The quick brown fox, jumps over 12345 lazy dogs. " * 60, id="prose"),
pytest.param("a" + "1234567890" * 250 + "1", id="long_digit_run"),
pytest.param("--- " * 600, id="symbols"),
pytest.param("def f(x):\n return x # it's 42\n\n" * 80, id="code"),
pytest.param("word next\t\ttab\n \n" * 120, id="whitespace"),
pytest.param("naïve café, “quotes” & <tags> 中文汉字,日本語のテキスト 😀🎉\r\n" * 40, id="unicode"),
pytest.param("We're sure they'd say it's DON'T or I'M " * 60, id="contractions"),
pytest.param("abc123def4567ghi89 " * 120, id="letters_and_digits"),
],
)
def test_long_text_counts_what_the_tokenizer_counts_for_the_whole_text(pair_encoding: tiktoken.Encoding, text: str):
def counted(content: str) -> int | None:
return kimi_k3_tokens(pair_encoding, MODEL, [{"role": "user", "content": content}], None)
short: Final = counted("x")
assert short is not None
assert counted(text) == short - len(pair_encoding.encode_ordinary("x")) + len(pair_encoding.encode_ordinary(text))
@pytest.mark.parametrize(
("messages", "tools"),
[
pytest.param(
[
{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
],
}
],
None,
id="image",
),
pytest.param(
[{"role": "user", "content": "Weather?"}, {"role": "tool", "tool_call_id": "c1", "content": "Sun"}],
[_WEATHER],
id="tool_result_without_its_call",
),
],
)
def test_request_the_template_cannot_render_is_not_counted(
pair_encoding: tiktoken.Encoding, messages: list[dict[str, object]], tools: list[dict[str, object]] | None
):
assert kimi_k3_tokens(pair_encoding, MODEL, messages, tools) is None
@pytest.mark.parametrize(
"tool",
[
pytest.param(
{
"type": "function",
"name": "get_weather",
"description": "Get the weather in °C",
"parameters": _WEATHER["function"]["parameters"],
},
id="responses_api",
),
pytest.param(
{
"name": "get_weather",
"description": "Get the weather in °C",
"input_schema": _WEATHER["function"]["parameters"],
},
id="anthropic_messages",
),
],
)
def test_tool_in_another_api_shape_counts_like_the_same_chat_completions_tool(
pair_encoding: tiktoken.Encoding, tool: dict[str, object]
):
messages: Final = [{"role": "user", "content": "Weather in Paris?"}]
counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [tool])
assert counted is not None
assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER])
@pytest.mark.parametrize(
"tool",
[
pytest.param({"type": "web_search_20250305", "name": "web_search"}, id="anthropic_web_search"),
pytest.param({"type": "web_search"}, id="responses_web_search"),
pytest.param(
{
"type": "mcp",
"server_label": "docs",
"server_url": "https://mcp.example/mcp",
"require_approval": "never",
},
id="mcp",
),
],
)
def test_tool_bedrock_inference_drops_counts_like_the_function_tools_alone(
pair_encoding: tiktoken.Encoding, tool: dict[str, object]
):
messages: Final = [{"role": "user", "content": "Weather in Paris?"}]
counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, tool])
assert counted is not None
assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER])
def test_responses_custom_tool_counts_like_the_function_tool_inference_sends(pair_encoding: tiktoken.Encoding):
messages: Final = [{"role": "user", "content": "Weather in Paris?"}]
custom_tool: Final = {"type": "custom", "name": "grammar", "description": "Free-form SQL"}
sent_to_inference, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(
[dict(custom_tool)]
)
without_it: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER])
counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, custom_tool])
assert counted is not None and without_it is not None
assert counted > without_it
assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, *sent_to_inference])
def test_counting_a_request_with_an_image_url_never_downloads_the_image(
pair_encoding: tiktoken.Encoding, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
):
monkeypatch.setattr(litellm, "user_url_validation", False)
image: Final = respx_mock.get("https://images.example/cat.png").mock(return_value=httpx.Response(200))
messages: Final = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{"type": "image_url", "image_url": {"url": "https://images.example/cat.png"}},
],
}
]
assert kimi_k3_tokens(pair_encoding, MODEL, messages, None) is None
assert not image.called
def test_anthropic_request_counts_like_the_same_openai_request(pair_encoding: tiktoken.Encoding):
anthropic: Final = kimi_k3_anthropic_tokens(
pair_encoding,
MODEL,
{
"system": "Be terse.",
"messages": [
{"role": "user", "content": "Weather in Paris?"},
{
"role": "assistant",
"content": [{"type": "tool_use", "id": "c1", "name": "get_weather", "input": {"city": "Paris"}}],
},
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "Sun"}]},
],
"tools": [
{
"name": "get_weather",
"description": "Get the weather in °C",
"input_schema": _WEATHER["function"]["parameters"],
}
],
},
)
openai: Final = kimi_k3_tokens(
pair_encoding,
MODEL,
[
{"role": "system", "content": "Be terse."},
{"role": "user", "content": "Weather in Paris?"},
{"role": "assistant", "content": None, "tool_calls": [_call("c1", {"city": "Paris"})]},
{"role": "tool", "tool_call_id": "c1", "content": "Sun"},
],
[_WEATHER],
)
assert openai is not None
assert anthropic == openai
@pytest.mark.parametrize("failure", ["checksum", "missing_file"])
def test_tokenizer_that_fails_its_checksum_or_cannot_be_read_is_not_loaded(
byte_source: KimiK3TokenizerSource, tmp_path: Path, caplog: pytest.LogCaptureFixture, failure: str
):
source: Final = (
KimiK3TokenizerSource(ranks_path=byte_source.ranks_path, sha256=hashlib.sha256(b"other").hexdigest())
if failure == "checksum"
else KimiK3TokenizerSource(ranks_path=lambda: str(tmp_path / "missing"), sha256=byte_source.sha256)
)
verbose_logger.addHandler(caplog.handler)
try:
with caplog.at_level(logging.WARNING, logger=verbose_logger.name):
loaded: Final = load_kimi_k3_encoding(source)
finally:
verbose_logger.removeHandler(caplog.handler)
assert loaded is None
assert any("Kimi K3 tokenizer unavailable" in record.getMessage() for record in caplog.records)
@pytest.mark.asyncio
async def test_slow_tokenizer_load_is_not_waited_on_past_the_wait_and_is_shared_once_done(
byte_source: KimiK3TokenizerSource,
):
released: Final = threading.Event()
def slow_ranks_path() -> str:
released.wait(10)
return byte_source.ranks_path()
encodings: Final = KimiK3Encodings(
source=KimiK3TokenizerSource(ranks_path=slow_ranks_path, sha256=byte_source.sha256),
retry_after_seconds=60,
clock=lambda: 0.0,
)
assert await encodings.encoding(wait_seconds=0) is None
released.set()
loaded: Final = await encodings.encoding(wait_seconds=5)
assert loaded is not None
assert await encodings.encoding(wait_seconds=0) is loaded
class _Clock:
def __init__(self) -> None:
self.now = 0.0
def __call__(self) -> float:
return self.now
async def _loaded_eventually(encodings: KimiK3Encodings) -> tiktoken.Encoding | None:
for _ in range(500):
if (loaded := await encodings.encoding(wait_seconds=0)) is not None:
return loaded
await asyncio.sleep(0.01)
return None
@pytest.mark.asyncio
async def test_failed_tokenizer_load_is_retried_in_the_background_only_after_the_retry_window(
byte_source: KimiK3TokenizerSource, tmp_path: Path
):
clock: Final = _Clock()
reads: Final[list[float]] = []
released: Final = threading.Event()
def ranks_path() -> str:
reads.append(clock.now)
if len(reads) == 1:
return str(tmp_path / "missing")
released.wait(10)
return byte_source.ranks_path()
encodings: Final = KimiK3Encodings(
source=KimiK3TokenizerSource(ranks_path=ranks_path, sha256=byte_source.sha256),
retry_after_seconds=60,
clock=clock,
)
assert await encodings.encoding(wait_seconds=5) is None
clock.now = 59.0
assert await encodings.encoding(wait_seconds=5) is None
clock.now = 60.0
assert await asyncio.wait_for(encodings.encoding(wait_seconds=5), timeout=1) is None
released.set()
assert await _loaded_eventually(encodings) is not None
assert reads == [0.0, 60.0]
@pytest.mark.asyncio
async def test_tokenizer_load_still_running_after_the_wait_is_logged(
byte_source: KimiK3TokenizerSource, caplog: pytest.LogCaptureFixture
):
released: Final = threading.Event()
def hung_ranks_path() -> str:
released.wait(10)
return byte_source.ranks_path()
encodings: Final = KimiK3Encodings(
source=KimiK3TokenizerSource(ranks_path=hung_ranks_path, sha256=byte_source.sha256), retry_after_seconds=60
)
verbose_logger.addHandler(caplog.handler)
try:
with caplog.at_level(logging.WARNING, logger=verbose_logger.name):
loaded: Final = await encodings.encoding(wait_seconds=0.05)
finally:
verbose_logger.removeHandler(caplog.handler)
released.set()
assert loaded is None
assert any("Kimi K3 tokenizer still loading" in record.getMessage() for record in caplog.records)

View file

@ -54,6 +54,7 @@ async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload(
"request_model": "openai/gpt-4o",
"model_used": "gpt-4o",
"tokenizer_type": "openai_api",
"count_source": None,
"original_response": payload,
"error": False,
"error_message": None,
@ -98,6 +99,7 @@ async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_v
"request_model": "openai/gpt-4o",
"model_used": "gpt-4o",
"tokenizer_type": "openai_api",
"count_source": None,
"original_response": None,
"error": True,
"error_message": str(rejection.value),
@ -126,6 +128,7 @@ async def test_count_payload_that_is_not_an_object_becomes_an_error_response_nam
"request_model": "openai/gpt-4o",
"model_used": "gpt-4o",
"tokenizer_type": "openai_api",
"count_source": None,
"original_response": None,
"error": True,
"error_message": str(failed_lookup.value),
@ -156,6 +159,7 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_
"request_model": "openai/gpt-4o",
"model_used": "gpt-4o",
"tokenizer_type": "openai_api",
"count_source": None,
"original_response": None,
"error": True,
"error_message": upstream_text,
@ -382,6 +386,7 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_
"request_model": "openai/gpt-4o",
"model_used": "gpt-4o",
"tokenizer_type": "openai_api",
"count_source": None,
"original_response": None,
"error": True,
"error_message": f"CountTokens processing error: {undecodable.value}",

View file

@ -0,0 +1,190 @@
import json
from typing import Final
import pytest
from litellm.proxy.common_utils.token_count_settings import (
ProviderCountFailure,
TokenCountModel,
TokenCountSettings,
count_source_fields,
resolve_token_count_settings,
)
from litellm.types.utils import TokenCountResponse
_OPUS_US_EAST_1: Final = TokenCountModel(model="anthropic.claude-opus-5-5", region="us-east-1")
_DEFAULT: Final = {
"token_count_framing": 97,
"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"},
}
def _in_us_west_2() -> str:
return "us-west-2"
@pytest.mark.parametrize(
("configured", "expected"),
[
({}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
(
{"token_count_framing": 5, "token_count_model": {"model": "custom-model", "region": "ap-southeast-4"}},
TokenCountSettings(framing=5, count_model=TokenCountModel(model="custom-model", region="ap-southeast-4")),
),
(
{"token_count_model": {"model": "custom-model"}},
TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)),
),
(
{"token_count_model": {"model": "custom-model", "region": " "}},
TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)),
),
({"token_count_framing": 0}, TokenCountSettings(framing=0, count_model=_OPUS_US_EAST_1)),
(
{"token_count_framing": None, "token_count_model": None},
TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1),
),
(
{"token_count_framing": True, "token_count_model": "custom-model"},
TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1),
),
({"token_count_framing": 78.0}, TokenCountSettings(framing=78, count_model=_OPUS_US_EAST_1)),
({"token_count_framing": 1.5}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
({"token_count_framing": -1}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
({"token_count_framing": "78"}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
({"token_count_model": {}}, TokenCountSettings(framing=97, count_model=None)),
({"token_count_model": ""}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
(
{"token_count_model": {"model": " ", "region": "us-east-1"}},
TokenCountSettings(framing=97, count_model=None),
),
],
)
def test_config_overrides_the_cost_map_default_key_by_key(
configured: dict[str, object], expected: TokenCountSettings
) -> None:
assert (
resolve_token_count_settings(configured=configured, default=_DEFAULT, deployment_region=_in_us_west_2)
== expected
)
@pytest.mark.parametrize(
("deployment_region", "configured", "expected_region"),
[
("us-west-2", {}, "us-east-1"),
("eu-central-1", {}, None),
("ca-central-1", {}, None),
("us-gov-west-1", {}, None),
(
"eu-central-1",
{"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}},
"us-east-1",
),
],
)
def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geography(
deployment_region: str, configured: dict[str, object], expected_region: str | None
) -> None:
settings: Final = resolve_token_count_settings(
configured=configured, default=_DEFAULT, deployment_region=lambda: deployment_region
)
assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region)
def test_shipped_global_count_model_keeps_its_region_for_a_deployment_in_any_geography() -> None:
settings: Final = resolve_token_count_settings(
configured={},
default={"token_count_model": {"model": "global.anthropic.claude-opus-5-5", "region": "us-east-1"}},
deployment_region=lambda: "eu-west-1",
)
assert settings.count_model == TokenCountModel(model="global.anthropic.claude-opus-5-5", region="us-east-1")
def test_shipped_australian_count_region_stays_with_an_australian_deployment() -> None:
settings: Final = resolve_token_count_settings(
configured={},
default={"token_count_model": {"model": "au.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}},
deployment_region=lambda: "ap-southeast-2",
)
assert settings.count_model == TokenCountModel(model="au.anthropic.claude-opus-5-5", region="ap-southeast-4")
def test_no_setting_anywhere_means_the_model_counts_itself_without_framing() -> None:
assert resolve_token_count_settings(
configured={}, default={}, deployment_region=_in_us_west_2
) == TokenCountSettings(framing=None, count_model=None)
@pytest.mark.parametrize(
("value", "expected"),
[("kimi_k3", "kimi_k3"), ("cl100k_base", None), (None, None), (True, None)],
)
def test_only_a_tokenizer_litellm_ships_is_taken_from_the_cost_map(value: object, expected: str | None) -> None:
settings: Final = resolve_token_count_settings(
configured={}, default={"token_count_tokenizer": value}, deployment_region=_in_us_west_2
)
assert settings.tokenizer == expected
def _counted(source: str | None, model_used: str = "us.anthropic.claude-opus-5-5") -> TokenCountResponse:
return TokenCountResponse(
total_tokens=7,
request_model="sonnet",
model_used=model_used,
tokenizer_type="bedrock_api",
count_source=source,
)
@pytest.mark.parametrize(
("source", "expected"),
[
(None, {}),
("provider", {"litellm_count_source": "provider"}),
("local_estimate", {"litellm_count_source": "local_estimate"}),
(
"token_count_model",
{"litellm_count_source": "token_count_model", "litellm_count_model": "us.anthropic.claude-opus-5-5"},
),
],
)
def test_count_source_fields_name_the_counting_model_only_when_token_count_model_counted(
source: str | None, expected: dict[str, str]
) -> None:
assert count_source_fields(_counted(source)) == expected
_UPSTREAM_404: Final = {
"type": "error",
"request_id": "req_upstream",
"error": {"type": "not_found_error", "message": "The model 'anthropic.claude-opus-5-5' does not exist"},
}
@pytest.mark.parametrize(
("message", "expected"),
[
(
json.dumps(_UPSTREAM_404),
{
**_UPSTREAM_404,
"error": {
"type": "not_found_error",
"message": "The model 'anthropic.claude-opus-5-5' does not exist. Enable eu-west-1",
},
},
),
(json.dumps({"message": "Too many tokens"}), "Too many tokens. Enable eu-west-1"),
("scripted 404", "scripted 404. Enable eu-west-1"),
(
"Credential should be scoped to a valid region. ",
"Credential should be scoped to a valid region. Enable eu-west-1",
),
],
)
def test_hint_lands_in_the_upstream_error_message_and_keeps_its_envelope(message: str, expected: object) -> None:
hinted: Final = ProviderCountFailure(status_code=404, message=message).with_hint("Enable eu-west-1")
assert (json.loads(hinted.message) if isinstance(expected, dict) else hinted.message) == expected
assert hinted.status_code == 404

File diff suppressed because it is too large Load diff

View file

@ -513,6 +513,40 @@ def test_should_not_pollute_shared_key_with_custom_nonzero_pricing():
)
def test_token_count_settings_stay_on_the_deployment_that_set_them() -> None:
"""A deployment's ``token_count_model`` and ``token_count_framing`` must not land on
the shared ``{provider}/{model}`` key, or a sibling deployment of that model would count with them."""
from litellm.proxy.common_utils.token_count_settings import cost_map_entry
backend_model: Final = "bedrock/us.moonshotai.kimi-k3"
deploy_id: Final = "kimi-deploy-keep-token-count"
configured: Final = {
"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-west-2"},
"token_count_framing": 7,
}
builtin: Final = {field: cost_map_entry(backend_model).get(field) for field in configured}
assert all(builtin[field] != value for field, value in configured.items()), builtin
model_keys: Final = {
key: copy.deepcopy(litellm.model_cost.get(key))
for key in (deploy_id, backend_model, backend_model.removeprefix("bedrock/"))
}
try:
Router(
model_list=[
{
"model_name": "kimi-keep-token-count",
"litellm_params": {"model": backend_model},
"model_info": {"id": deploy_id, **configured},
}
],
)
assert {field: litellm.model_cost[deploy_id][field] for field in configured} == configured
assert {field: cost_map_entry(backend_model).get(field) for field in configured} == builtin
finally:
_restore_model_cost_entries(model_keys)
def test_should_store_full_pricing_under_deployment_model_id():
"""
Per-deployment pricing (including zero) should be stored and

View file

@ -936,6 +936,13 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"supports_video_input": {"type": "boolean"},
"supports_vision": {"type": "boolean"},
"supports_web_search": {"type": "boolean"},
"token_count_framing": {"type": "integer"},
"token_count_model": {
"type": "object",
"properties": {"model": {"type": "string"}, "region": {"type": "string"}},
"required": ["model"],
},
"token_count_tokenizer": {"type": "string", "enum": ["kimi_k3"]},
"supports_url_context": {"type": "boolean"},
"supports_multimodal": {"type": "boolean"},
"uses_embed_content": {"type": "boolean"},

View file

@ -43384,9 +43384,17 @@ export interface components {
* @description Response structure for token count details with modality breakdown.
*
* Example:
* {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]}
* {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}],
* 'litellm_count_source': 'provider'}
*/
TokenCountDetailsResponse: {
/** Litellm Count Model */
litellm_count_model?: string;
/**
* Litellm Count Source
* @enum {string}
*/
litellm_count_source?: "provider" | "token_count_model" | "local_estimate";
/** Prompttokensdetails */
promptTokensDetails: components["schemas"]["PromptTokensDetails"][];
/** Totaltokens */
@ -43415,6 +43423,8 @@ export interface components {
};
/** TokenCountResponse */
TokenCountResponse: {
/** Count Source */
count_source?: ("provider" | "token_count_model" | "local_estimate") | null;
/**
* Error
* @default false