From b6f1a033deda54461feda0a47fa59532aee5c089 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 11:14:29 -0700 Subject: [PATCH] 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 75faaeab641e161a3d6bc38a42fcaa71c5602d90 (main). --- litellm/litellm_core_utils/core_helpers.py | 26 +- .../prompt_templates/factory.py | 22 + litellm/llms/bedrock/count_tokens/kimi_k3.py | 502 ++++++ .../responses/count_tokens/token_counter.py | 1 + .../llms/vertex_ai/count_tokens/context.py | 25 +- ...odel_prices_and_context_window_backup.json | 28 +- .../proxy/anthropic_endpoints/endpoints.py | 3 +- .../common_utils/token_count_settings.py | 164 ++ litellm/proxy/google_endpoints/endpoints.py | 11 +- litellm/proxy/proxy_server.py | 461 +++++- .../proxy/response_api_endpoints/endpoints.py | 5 +- litellm/types/llms/vertex_ai.py | 6 +- litellm/types/token_count.py | 16 + litellm/types/utils.py | 9 + litellm/utils.py | 3 + model_prices_and_context_window.json | 28 +- model_prices_and_context_window.schema.json | 30 + ...anthropic_count_tokens_system_lift_wire.py | 24 +- .../test_bedrock_mantle_count_tokens_wire.py | 89 +- .../test_vertex_partner_count_tokens_wire.py | 28 +- .../integration/spend/test_cache_and_quota.py | 8 +- .../test_google_api_endpoints.py | 42 +- .../test_model_management_endpoints.py | 48 + .../proxy/proxy_server/test_routes_utils.py | 1 + .../count_tokens/test_token_counter.py | 5 + .../count_tokens/test_bedrock_kimi_k3.py | 524 +++++++ .../count_tokens/test_token_counter.py | 5 + .../common_utils/test_token_count_settings.py | 190 +++ tests/unit/proxy/test_proxy_token_counter.py | 1364 ++++++++++++++++- .../unit/test_router_model_cost_isolation.py | 34 + tests/unit/test_utils.py | 7 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 +- 32 files changed, 3571 insertions(+), 150 deletions(-) create mode 100644 litellm/llms/bedrock/count_tokens/kimi_k3.py create mode 100644 litellm/proxy/common_utils/token_count_settings.py create mode 100644 litellm/types/token_count.py create mode 100644 tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py create mode 100644 tests/unit/proxy/common_utils/test_token_count_settings.py diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index b095b4b12c6..d74007dfccb 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -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 diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 7b12d1e939f..26d219a5ebc 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -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) diff --git a/litellm/llms/bedrock/count_tokens/kimi_k3.py b/litellm/llms/bedrock/count_tokens/kimi_k3.py new file mode 100644 index 00000000000..04ae9fae9c2 --- /dev/null +++ b/litellm/llms/bedrock/count_tokens/kimi_k3.py @@ -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("&", "&").replace('"', """), '"') 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) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index bdaea0b4104..b34b6fbce2f 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -116,6 +116,7 @@ class OpenAITokenCounter(BaseTokenCounter): error=True, error_message=request_error, status_code=400, + rejected_before_provider_call=True, ) deployment = deployment or {} diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index 3c9c30b4a29..62503bf2aa6 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -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, ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d89c8baf9b0..9635cf3ce72 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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, diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 210bd5a8bfd..98e1cf92026 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -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( diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py new file mode 100644 index 00000000000..6f805616820 --- /dev/null +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -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} diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 818fd1ad516..5fbe780379e 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -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 # ============================================================ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ee1e8647bf2..8369e79596c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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", ) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 021c741b2bc..d33d8a83bfa 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -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 diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index ce51e46ef15..28fdf3acf67 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -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 diff --git a/litellm/types/token_count.py b/litellm/types/token_count.py new file mode 100644 index 00000000000..3b21401066c --- /dev/null +++ b/litellm/types/token_count.py @@ -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] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d3f856697e3..2cce951ca60 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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): diff --git a/litellm/utils.py b/litellm/utils.py index e91460b5987..64ecfcac98c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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(): diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d89c8baf9b0..9635cf3ce72 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 35624045fdf..b99f8ee6f98 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -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, diff --git a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py index a4b5c15ad54..b34b0b6b526 100644 --- a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py +++ b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py @@ -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,) diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index bc70c5d3b21..e20a397a708 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -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 diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py index 1c4b1d1a6e5..d5c841d9335 100644 --- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -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}]} ) diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py index 1cc03838f8d..b5145ea49d4 100644 --- a/tests/integration/spend/test_cache_and_quota.py +++ b/tests/integration/spend/test_cache_and_quota.py @@ -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") diff --git a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py index 114bbb184c6..2777d18a399 100644 --- a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py @@ -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" diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 5f7807650e1..b7f6865ab7e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py index fc5ab9475eb..ee165d7b204 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py @@ -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, diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py index 86ad6c8d9d9..4913d2b0d1d 100644 --- a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -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}", diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py new file mode 100644 index 00000000000..ffddb09cca7 --- /dev/null +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py @@ -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&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” & 中文汉字,日本語のテキスト 😀🎉\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) diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py index 65a9f6f044d..a45773b93fc 100644 --- a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py +++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py @@ -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}", diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py new file mode 100644 index 00000000000..84a32a02a5c --- /dev/null +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -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 diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 869e7eb60d2..86556f1c7bd 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2,14 +2,32 @@ # 1. Generate a Key, and use it to make a call +import asyncio +import base64 +import hashlib +import itertools import json -from typing import Final +import logging +import threading +import time +import traceback +from collections.abc import Callable, Coroutine, Iterator, Mapping +from functools import partial +from pathlib import Path +from types import MappingProxyType +from typing import Final, NoReturn from unittest.mock import AsyncMock, MagicMock, patch +import anyio import httpx import pytest +import tiktoken from dotenv import load_dotenv from pydantic import BaseModel, ConfigDict, JsonValue +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from tokenizers.processors import TemplateProcessing load_dotenv() @@ -21,10 +39,20 @@ from fastapi.responses import JSONResponse import litellm from litellm import Router +from litellm._logging import verbose_proxy_logger from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler +from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler +from litellm.llms.bedrock.count_tokens.kimi_k3 import ( + KIMI_K3_RANKS_SHA256, + KimiK3Encodings, + KimiK3TokenizerSource, + kimi_k3_template_tokens, + kimi_k3_tokens, + load_kimi_k3_encoding, +) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import VertexAIPartnerModelsTokenCounter @@ -32,8 +60,59 @@ from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.anthropic_endpoints.endpoints import ( count_tokens as anthropic_count_tokens, ) -from litellm.proxy.proxy_server import token_counter -from litellm.types.utils import TokenCountResponse +from litellm.proxy.common_utils.token_count_settings import ( + ProviderCountFailure, + TokenCountModel, + TokenCountSettings, + cost_map_entry, +) +from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request +from litellm.proxy.proxy_server import ( + _estimate_or_provider_error, + _provider_token_count, + _try_provider_token_count, + count_request_tokens, + token_counter, +) +from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem +from litellm.types.utils import SelectTokenizerResponse, TokenCountMessagesFormat, TokenCountResponse + + +def _hugging_face_offline() -> NoReturn: + raise OSError("unit tests never download the Kimi K3 tokenizer") + + +@pytest.fixture(autouse=True) +def _kimi_k3_tokenizer_unavailable(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "_KIMI_K3_ENCODINGS", + KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=_hugging_face_offline, sha256=KIMI_K3_RANKS_SHA256), + retry_after_seconds=60, + ), + ) + + +@pytest.fixture +def bedrock_sonnet_5_5_mapped_as_on_main() -> Iterator[None]: + entries: Final = { + name: { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "token_count_model": {"model": count_model, "region": "us-east-1"}, + } + for name, count_model in ( + ("anthropic.claude-sonnet-5-5", "anthropic.claude-opus-5-5"), + ("us.anthropic.claude-sonnet-5-5", "us.anthropic.claude-opus-5-5"), + ) + } + with pytest.MonkeyPatch.context() as patched: + patched.setattr(litellm, "model_cost", {**litellm.model_cost, **entries}) + litellm.utils._invalidate_model_cost_lowercase_map() + yield + litellm.utils._invalidate_model_cost_lowercase_map() @pytest.mark.asyncio @@ -1179,6 +1258,1285 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): proxy_server.count_request_tokens = original_token_counter +class _ScriptedCounter(BaseTokenCounter): + def __init__(self, replies: dict[str, TokenCountResponse], burst: int = 1) -> None: + self._replies = replies + self._burst = burst + self._burst_arrived = asyncio.Event() + self.calls: list[tuple[str, object, object, object]] = [] + self.regions: list[object] = [] + + def should_use_token_counting_api(self, custom_llm_provider: str | None = None) -> bool: + return True + + async def count_tokens( + self, + model_to_use: str, + messages: object, + contents: object, + deployment: object = None, + request_model: str = "", + tools: object = None, + system: object = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, + context_management: object = None, + ) -> TokenCountResponse | None: + self.calls.append((model_to_use, messages, tools, system)) + self.regions.append(None if deployment is None else deployment["litellm_params"].get("aws_region_name")) + if len(self.calls) >= self._burst: + self._burst_arrived.set() + await self._burst_arrived.wait() + return self._replies[model_to_use] + + +def _provider_reply( + model: str, total_tokens: int = 0, status_code: int | None = None, rejected_before_provider_call: bool = False +) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=total_tokens, + request_model="sonnet-5-5", + model_used=model, + tokenizer_type="bedrock_api", + error=status_code is not None, + error_message=None if status_code is None else f"scripted {status_code} for {model}", + status_code=status_code, + rejected_before_provider_call=rejected_before_provider_call, + ) + + +_COUNT_REQUEST = TokenCountRequest( + model="sonnet-5-5", + messages=[{"role": "user", "content": "hello"}], + tools=[{"name": "get_quote", "input_schema": {"type": "object"}}], + system="be brief", +) + + +_SONNET_DEPLOYMENT: Final = MappingProxyType( + {"litellm_params": {"model": "bedrock/us.anthropic.claude-sonnet-5-5", "aws_region_name": "us-west-2"}} +) +_OPUS_IN_US_EAST_1: Final = TokenCountModel(model="us.anthropic.claude-opus-5-5", region="us-east-1") + + +async def _ladder( + counter: _ScriptedCounter, settings: TokenCountSettings, local_estimate_allowed: bool = True +) -> TokenCountResponse | None: + return await _provider_token_count( + count_with=partial( + _try_provider_token_count, + counter, + "bedrock", + messages=_COUNT_REQUEST.messages, + contents=_COUNT_REQUEST.contents, + deployment=_SONNET_DEPLOYMENT, + request_model=_COUNT_REQUEST.model, + tools=_COUNT_REQUEST.tools, + system=_COUNT_REQUEST.system, + ), + model_to_use="us.anthropic.claude-sonnet-5-5", + settings=settings, + local_estimate_allowed=local_estimate_allowed, + ) + + +@pytest.mark.asyncio +async def test_model_without_token_count_model_counts_itself_in_the_deployment_region(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=447)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=10, count_model=None)) + + assert result is not None + assert (result.total_tokens, result.count_source) == (447, "provider") + assert ([call[0] for call in counter.calls], counter.regions) == (["us.anthropic.claude-sonnet-5-5"], ["us-west-2"]) + + +@pytest.mark.parametrize( + ("count_model", "region"), + [ + (_OPUS_IN_US_EAST_1, "us-east-1"), + (TokenCountModel(model="us.anthropic.claude-opus-5-5", region=None), "us-west-2"), + ], +) +@pytest.mark.asyncio +async def test_token_count_model_counts_the_same_request_in_its_own_region_without_trying_the_model( + count_model: TokenCountModel, region: str +): + counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=None, count_model=count_model)) + + assert result is not None + assert (result.total_tokens, result.count_source, result.model_used) == ( + 447, + "token_count_model", + "us.anthropic.claude-opus-5-5", + ) + assert ([call[0] for call in counter.calls], counter.regions) == (["us.anthropic.claude-opus-5-5"], [region]) + assert counter.calls[0][1:] == (_COUNT_REQUEST.messages, _COUNT_REQUEST.tools, _COUNT_REQUEST.system) + assert _SONNET_DEPLOYMENT["litellm_params"]["aws_region_name"] == "us-west-2" + + +@pytest.mark.parametrize("status_code", [401, 403, 429, 500]) +@pytest.mark.asyncio +async def test_errors_other_than_cannot_count_raise_even_with_framing(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=status_code)} + ) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=10, count_model=None)) + + assert (raised.value.code, raised.value.message) == (str(status_code), f"scripted {status_code} for sonnet") + + +@pytest.mark.parametrize("status_code", [400, 403, 404, 500]) +@pytest.mark.asyncio +async def test_token_count_model_failure_names_its_region_so_the_user_can_enable_it(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=None, count_model=_OPUS_IN_US_EAST_1)) + + assert raised.value.code == str(status_code) + assert raised.value.message.startswith(f"scripted {status_code} for opus") + assert "us.anthropic.claude-opus-5-5 in us-east-1" in raised.value.message + assert len(counter.calls) == 1 + + +@pytest.mark.asyncio +async def test_token_count_model_without_a_region_answers_its_error_unchanged(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=404)}) + + with pytest.raises(ProxyException) as raised: + await _ladder( + counter, + TokenCountSettings( + framing=None, count_model=TokenCountModel(model="us.anthropic.claude-opus-5-5", region=None) + ), + ) + + assert (raised.value.code, raised.value.message) == ("404", "scripted 404 for opus") + + +@pytest.mark.asyncio +async def test_uncountable_model_without_framing_raises(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=404)}) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=None, count_model=None)) + + assert raised.value.code == "404" + + +@pytest.mark.asyncio +async def test_uncountable_model_with_framing_falls_through_to_the_local_estimate(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=400)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=0, count_model=None)) + + assert isinstance(result, ProviderCountFailure) + assert len(counter.calls) == 1 + + +@pytest.mark.parametrize("status_code", [400, 404]) +@pytest.mark.asyncio +async def test_token_count_model_error_stands_even_with_framing(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=78, count_model=_OPUS_IN_US_EAST_1)) + + assert raised.value.code == str(status_code) + assert "us.anthropic.claude-opus-5-5 in us-east-1" in raised.value.message + + +@pytest.mark.asyncio +async def test_uncountable_model_with_framing_raises_when_local_counting_is_disabled(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=400)}) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=10, count_model=None), local_estimate_allowed=False) + + assert raised.value.code == "400" + + +@pytest.mark.asyncio +async def test_request_the_counter_rejects_before_calling_the_provider_answers_400_without_framing(): + counter: Final = _ScriptedCounter( + { + "us.anthropic.claude-sonnet-5-5": _provider_reply( + "sonnet", status_code=400, rejected_before_provider_call=True + ) + } + ) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=100, count_model=None)) + + assert (raised.value.code, raised.value.message) == ("400", "scripted 400 for sonnet") + + +@pytest.mark.parametrize("model", ["openai/gpt-5.4-mini", "gemini/gemini-3.8-flash"]) +@pytest.mark.asyncio +async def test_framing_keeps_the_400_for_blocks_the_provider_counter_cannot_convert( + monkeypatch: pytest.MonkeyPatch, model: str +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "framed", + "litellm_params": {"model": model, "api_key": "sk-unused-framed"}, + "model_info": {"token_count_framing": 100}, + } + ] + ), + ) + search_result: Final = {"type": "search_result", "source": "https://example.com", "title": "Notes", "content": []} + + with pytest.raises(ProxyException) as raised: + await count_request_tokens( + request=TokenCountRequest( + model="framed", + messages=[{"role": "user", "content": [{"type": "text", "text": "Summarize."}, search_result]}], + ), + call_endpoint=True, + messages_format="anthropic", + ) + + assert raised.value.code == "400" + assert "search_result" in raised.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "unreadable_body", + [ + pytest.param( + { + "messages": [ + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {"query": "new"}, + }, + {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_01", "content": []}, + ], + }, + ] + }, + id="web_search_blocks", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "What is new?"}], + "tools": [{"name": "search", "input_schema": {"type": "object", "properties": {"query": True}}}], + }, + id="boolean_subschema_tool", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "What is new?"}], + "tools": [ + {"name": "search", "input_schema": {"type": "object", "properties": {"tags": {"type": "array"}}}} + ], + }, + id="array_schema_without_items", + ), + pytest.param( + {"messages": [], "contents": [{"role": "user", "parts": [{"text": "What is new?"}]}]}, + id="google_contents", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "Describe this"}], + "contents": [ + { + "role": "user", + "parts": [ + {"text": "Describe this"}, + {"fileData": {"mimeType": "image/png", "fileUri": "gs://b/c.png"}}, + ], + } + ], + }, + id="google_contents_with_media_the_translation_dropped", + ), + pytest.param( + { + "messages": [ + {"role": "user", "content": "What is new?"}, + {"role": "assistant", "content": "Weighing the release notes"}, + ], + "contents": [ + {"role": "user", "parts": [{"text": "What is new?"}]}, + {"role": "model", "parts": [{"text": "Weighing the release notes", "thought": True}]}, + ], + }, + id="google_contents_with_a_thought_the_translation_counts_as_text", + ), + ], +) +async def test_framed_estimate_that_cannot_read_the_body_answers_the_provider_error( + monkeypatch: pytest.MonkeyPatch, unreadable_body: dict[str, list[dict[str, object]]] +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + with pytest.raises(ProxyException) as raised: + await token_counter(request=TokenCountRequest(model="kimi", **unreadable_body), call_endpoint=True) + readable: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + + assert (raised.value.code, raised.value.message) == ("400", "scripted 400 for kimi") + assert readable.count_source == "local_estimate" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "contents", + [ + pytest.param([{"role": "user", "parts": [{"text": "What is new in the release notes?"}]}], id="text"), + pytest.param( + [ + {"role": "user", "parts": [{"text": "Weather in Paris?"}]}, + { + "role": "model", + "parts": [ + {"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}, "thoughtSignature": "c2ln"} + ], + }, + { + "role": "user", + "parts": [{"functionResponse": {"name": "get_weather", "response": {"forecast": "Sun"}}}], + }, + ], + id="sdk_agent_loop", + ), + pytest.param( + [ + {"role": "user", "parts": [{"text": "Hi"}]}, + {"role": "model", "parts": [{"text": "Hello!", "thoughtSignature": "c2ln"}]}, + {"role": "user", "parts": [{"text": "What is new?"}]}, + ], + id="sdk_chat_with_a_thought_signature", + ), + pytest.param([], id="empty"), + ], +) +async def test_framed_estimate_counts_google_contents_like_the_same_conversation_sent_as_messages( + monkeypatch: pytest.MonkeyPatch, contents: list[dict[str, object]] +): + _framed_kimi(monkeypatch) + google_request: Final = _google_count_tokens_request("kimi", {"contents": contents}) + + from_google: Final = await token_counter(request=google_request, call_endpoint=True) + from_messages: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=google_request.messages), call_endpoint=True + ) + + assert from_google.count_source == "local_estimate" + assert from_google.total_tokens == from_messages.total_tokens + + +@pytest.mark.asyncio +async def test_framed_deployment_whose_custom_tokenizer_cannot_load_answers_the_provider_error( + monkeypatch: pytest.MonkeyPatch, +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "kimi", + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": {"custom_tokenizer": {"identifier": "moonshotai/kimi-k3", "revision": "main"}}, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + with pytest.raises(ProxyException) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + + assert (raised.value.code, raised.value.message) == ("400", "scripted 400 for kimi") + + +def _framed_kimis_with_custom_tokenizers( + monkeypatch: pytest.MonkeyPatch, tokenizers: Mapping[str, tuple[str, str]], burst: int = 1 +) -> None: + counter: Final = _ScriptedCounter( + { + "us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400), + "us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=12), + }, + burst=burst, + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": {"model": "bedrock/us.anthropic.claude-sonnet-5-5"}, + "model_info": {"token_count_model": {}}, + }, + *( + { + "model_name": name, + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": { + "custom_tokenizer": {"identifier": identifier, "revision": revision, "auth_token": None} + }, + } + for name, (identifier, revision) in tokenizers.items() + ), + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + +def _framed_kimi_with_custom_tokenizer(monkeypatch: pytest.MonkeyPatch, identifier: str) -> None: + _framed_kimis_with_custom_tokenizers(monkeypatch, {"kimi": (identifier, "main")}) + + +async def _count_status(model: str = "kimi") -> str: + try: + response: Final = await token_counter( + request=TokenCountRequest(model=model, messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + except ProxyException as e: + return e.code + return response.count_source or "" + + +class _SkippableClockLoop(asyncio.SelectorEventLoop): + skipped_seconds: float = 0.0 + + def time(self) -> float: + return super().time() + self.skipped_seconds + + def skip(self, seconds: float) -> None: + self.skipped_seconds += seconds + + +def _run_on_skippable_clock( + scenario: Callable[[_SkippableClockLoop], Coroutine[object, object, tuple[str, ...]]], +) -> tuple[str, ...]: + loop: Final = _SkippableClockLoop() + try: + return loop.run_until_complete(scenario(loop)) + finally: + loop.run_until_complete(loop.shutdown_default_executor()) + loop.close() + + +async def _until_finished(tasks: tuple[asyncio.Future[str], ...], count: int) -> None: + for finished in itertools.islice(asyncio.as_completed(tasks), count): + await finished + + +def _tokenizer_that_panics_while_encoding() -> SelectTokenizerResponse: + template: Final = Tokenizer(WordLevel({"[UNK]": 0, "": 1}, unk_token="[UNK]")) + template.pre_tokenizer = Whitespace() + template.post_processor = TemplateProcessing(single=" $A", special_tokens=[("", 1)]) + spec: Final = json.loads(template.to_str()) + return litellm.utils.create_tokenizer( + json.dumps({**spec, "post_processor": {**spec["post_processor"], "special_tokens": {}}}) + ) + + +def test_hanging_tokenizer_download_answers_the_provider_error_after_the_wait_with_one_download_while_it_runs( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/hanging-tokenizer") + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + downloads: Final[list[str]] = [] + + def hanging_download(model: str, custom_tokenizer: object = None) -> NoReturn: + downloads.append(model) + download_started.set() + hub_released.wait(10) + raise OSError("hub unreachable") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hanging_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e-9) + + async def burst_past_the_wait_then_one_more(loop: _SkippableClockLoop) -> tuple[str, ...]: + try: + burst: Final = asyncio.gather(*(_count_status() for _ in range(60))) + await asyncio.to_thread(download_started.wait, 10) + loop.skip(31) + answered: Final = await asyncio.wait_for(burst, 8) + later: Final = await asyncio.wait_for(_count_status(), 8) + finally: + hub_released.set() + return (*answered, later) + + assert set(_run_on_skippable_clock(burst_past_the_wait_then_one_more)) == {"400"} + assert len(downloads) == 1 + + +@pytest.mark.asyncio +async def test_framed_count_waits_for_a_tokenizer_download_that_finishes_within_the_wait( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/slow-tokenizer") + real_select_tokenizer: Final = litellm.utils._select_tokenizer + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + + def slow_download(model: str, custom_tokenizer: object = None): + download_started.set() + hub_released.wait(10) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", slow_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + count: Final = asyncio.create_task(_count_status()) + try: + await asyncio.to_thread(download_started.wait, 10) + finally: + hub_released.set() + + assert await asyncio.wait_for(count, 10) == "local_estimate" + + +@pytest.mark.asyncio +async def test_failed_tokenizer_download_is_reused_until_the_retry_window_passes(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/flaky-tokenizer") + real_select_tokenizer: Final = litellm.utils._select_tokenizer + downloads: Final[list[str]] = [] + + def hub_down_once(model: str, custom_tokenizer: object = None): + downloads.append(model) + if len(downloads) == 1: + raise OSError("hub unreachable") + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hub_down_once) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e9) + within_window: Final = [await _count_status() for _ in range(3)] + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e-9) + after_window: Final = await _count_status() + + assert within_window == ["400", "400", "400"] + assert after_window == "local_estimate" + assert len(downloads) == 2 + + +def test_failed_tokenizer_load_is_reused_for_the_window_after_it_fails(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/slow-failing-tokenizer") + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + downloads: Final[list[str]] = [] + + def slow_failing_download(model: str, custom_tokenizer: object = None) -> NoReturn: + downloads.append(model) + download_started.set() + hub_released.wait(10) + raise OSError("hub unreachable") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", slow_failing_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 1000) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 60) + + async def fail_after_longer_than_the_window_then_retry(loop: _SkippableClockLoop) -> tuple[str, ...]: + first: Final = asyncio.ensure_future(_count_status()) + try: + await asyncio.to_thread(download_started.wait, 10) + loop.skip(120) + finally: + hub_released.set() + return await asyncio.wait_for(first, 8), await asyncio.wait_for(_count_status(), 8) + + assert _run_on_skippable_clock(fail_after_longer_than_the_window_then_retry) == ("400", "400") + assert len(downloads) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "second_tokenizer", + [("lit-tests/tokenizer-b", "main"), ("lit-tests/tokenizer-a", "v2")], + ids=["identifier", "revision"], +) +async def test_framed_deployments_with_different_custom_tokenizers_each_load_their_own( + monkeypatch: pytest.MonkeyPatch, second_tokenizer: tuple[str, str] +): + _framed_kimis_with_custom_tokenizers( + monkeypatch, {"kimi-a": ("lit-tests/tokenizer-a", "main"), "kimi-b": second_tokenizer} + ) + real_select_tokenizer: Final = litellm.utils._select_tokenizer + loaded: Final[list[object]] = [] + + def load(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + loaded.append(custom_tokenizer) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", load) + + statuses: Final = [await _count_status(name) for name in ("kimi-a", "kimi-b")] + + assert statuses == ["local_estimate", "local_estimate"] + assert loaded == [ + {"identifier": "lit-tests/tokenizer-a", "revision": "main", "auth_token": None}, + {"identifier": second_tokenizer[0], "revision": second_tokenizer[1], "auth_token": None}, + ] + + +@pytest.mark.asyncio +async def test_tokenizer_loads_past_the_cap_answer_the_provider_error_without_stalling_other_counts( + monkeypatch: pytest.MonkeyPatch, +): + tokenizers: Final = {f"kimi-{n}": (f"lit-tests/unreachable-tokenizer-{n}", "main") for n in range(5)} + _framed_kimis_with_custom_tokenizers(monkeypatch, tokenizers, burst=len(tokenizers)) + real_select_tokenizer: Final = litellm.utils._select_tokenizer + hub_released: Final = threading.Event() + downloads: Final[list[object]] = [] + + def held_download(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + downloads.append(custom_tokenizer) + hub_released.wait(10) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", held_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_MAX_RUNNING", 2) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + monkeypatch.setattr(anyio.to_thread.current_default_thread_limiter(), "total_tokens", 2) + framed: Final = tuple(asyncio.ensure_future(_count_status(name)) for name in tokenizers) + try: + await asyncio.wait_for(_until_finished(framed, 3), 8) + unrelated: Final = await asyncio.wait_for(_count_status("sonnet"), 5) + finally: + hub_released.set() + statuses: Final = await asyncio.wait_for(asyncio.gather(*framed), 8) + + assert sorted(statuses) == ["400"] * 3 + ["local_estimate"] * 2 + assert unrelated == "provider" + assert len(downloads) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("panics_while", ["loading", "counting"]) +async def test_framed_count_answers_the_provider_error_when_the_tokenizer_panics( + monkeypatch: pytest.MonkeyPatch, panics_while: str +): + _framed_kimi_with_custom_tokenizer(monkeypatch, f"lit-tests/panicking-tokenizer-{panics_while}") + panicking: Final = _tokenizer_that_panics_while_encoding() + + def load(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + if panics_while == "loading": + panicking["tokenizer"].encode("hello world") + return panicking + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", load) + + assert await _count_status() == "400" + + +def _framed_kimi(monkeypatch: pytest.MonkeyPatch) -> None: + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + +@pytest.mark.asyncio +async def test_local_count_loads_the_tokenizer_again_after_a_failed_load(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/local-route-tokenizer") + real_select_tokenizer: Final = litellm.utils._select_tokenizer + downloads: Final[list[str]] = [] + + def hub_down_once(model: str, custom_tokenizer: object = None): + downloads.append(model) + if len(downloads) == 1: + raise OSError("hub unreachable") + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hub_down_once) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e9) + request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]) + + with pytest.raises(OSError, match="hub unreachable"): + await token_counter(request=request) + retried: Final = await token_counter(request=request) + + assert retried.count_source == "local_estimate" and retried.total_tokens > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("reuse_seconds", "tokenizer_types"), + [(1e9, ("openai_tokenizer", "openai_tokenizer")), (1e-9, ("openai_tokenizer", "huggingface_tokenizer"))], + ids=["within_the_window", "after_the_window"], +) +async def test_framed_count_loads_the_tokenizer_again_only_after_the_window_of_a_load_that_fell_back( + monkeypatch: pytest.MonkeyPatch, reuse_seconds: float, tokenizer_types: tuple[str, str] +): + _framed_kimi(monkeypatch) + real_select_tokenizer: Final = litellm.utils._select_tokenizer + hub_tokenizer: Final = litellm.utils.create_tokenizer( + Tokenizer(WordLevel({"[UNK]": 0}, unk_token="[UNK]")).to_str() + ) + loads: Final[list[str]] = [] + + def hub_up_on_the_second_load(model: str, custom_tokenizer: object = None): + loads.append(model) + return real_select_tokenizer(model="gpt-5.4-mini") if len(loads) == 1 else hub_tokenizer + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hub_up_on_the_second_load) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", reuse_seconds) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_MAX_RUNNING", 1) + request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]) + + first: Final = await token_counter(request=request, call_endpoint=True) + second: Final = await token_counter(request=request, call_endpoint=True) + + assert (first.tokenizer_type, second.tokenizer_type) == tokenizer_types + + +async def _fork_guard_traceback_frames() -> int: + with pytest.raises(ForkedAfterNativeRuntimeStarted) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + return len(traceback.extract_tb(raised.value.__traceback__)) + + +@pytest.mark.asyncio +async def test_fork_guard_error_while_loading_the_tokenizer_is_raised_to_every_count_sharing_the_load( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi(monkeypatch) + + def forked_load(model: str, custom_tokenizer: object = None) -> NoReturn: + raise ForkedAfterNativeRuntimeStarted("native runtime started before fork") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", forked_load) + + frames: Final = [await _fork_guard_traceback_frames() for _ in range(3)] + + assert len(set(frames)) == 1 + + +@pytest.mark.asyncio +async def test_swallowed_estimate_error_is_logged_by_type_without_the_request_content( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + verbose_proxy_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name), pytest.raises(ProxyException): + await token_counter( + request=TokenCountRequest( + model="kimi", + messages=[ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "ok", "tool_calls": "private-marker-in-the-prompt"}, + ], + ), + call_endpoint=True, + ) + finally: + verbose_proxy_logger.removeHandler(caplog.handler) + + warnings: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert any("count error 400" in message and "ValueError" in message for message in warnings) + assert not any("private-marker-in-the-prompt" in message for message in warnings) + + +@pytest.mark.asyncio +async def test_fork_guard_error_in_the_estimate_is_never_answered_as_the_provider_error(): + async def forked_estimate() -> NoReturn: + raise ForkedAfterNativeRuntimeStarted("native runtime started before fork") + + with pytest.raises(ForkedAfterNativeRuntimeStarted): + await _estimate_or_provider_error( + forked_estimate(), provider_failure=ProviderCountFailure(status_code=400, message="scripted 400") + ) + + +@pytest.mark.asyncio +async def test_local_estimate_adds_the_configured_framing_once_per_request(monkeypatch: pytest.MonkeyPatch): + def deployment(name: str, model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": f"sk-unused-{name}"}, + "model_info": model_info, + } + + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[deployment("framed", {"token_count_framing": 100}), deployment("plain", {})]), + ) + messages: Final = [{"role": "system", "content": "be brief"}, {"role": "user", "content": "hello there"}] + + plain: Final = await token_counter(request=TokenCountRequest(model="plain", messages=messages)) + framed: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + framed_prompt: Final = await token_counter(request=TokenCountRequest(model="framed", prompt="hello there")) + + assert plain.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + assert framed.total_tokens == plain.total_tokens + 100 + assert (plain.count_source, framed.count_source) == ("local_estimate", "local_estimate") + assert framed_prompt.total_tokens == litellm.token_counter(model="gpt-5.4-mini", text="hello there") + + +@pytest.mark.asyncio +async def test_framing_comes_from_the_row_the_deployment_provider_or_base_model_names(monkeypatch: pytest.MonkeyPatch): + def deployment(name: str, litellm_params: dict[str, object], model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {**litellm_params, "api_key": "unused", "api_base": "https://unused.invalid"}, + "model_info": model_info, + } + + kimi: Final = "databricks/system.ai.kimi-k3" + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + deployment("prefixed", {"model": kimi}, {}), + deployment("split", {"model": "system.ai.kimi-k3", "custom_llm_provider": "databricks"}, {}), + deployment("aliased", {"model": "gpt-5.4-mini", "custom_llm_provider": "openai"}, {"base_model": kimi}), + ] + ), + ) + messages: Final = [{"role": "user", "content": "hello there"}] + framing: Final = cost_map_entry(kimi)["token_count_framing"] + + prefixed: Final = await token_counter(request=TokenCountRequest(model="prefixed", messages=messages)) + split: Final = await token_counter(request=TokenCountRequest(model="split", messages=messages)) + aliased: Final = await token_counter(request=TokenCountRequest(model="aliased", messages=messages)) + + assert prefixed.total_tokens == litellm.token_counter(model="system.ai.kimi-k3", messages=messages) + framing + assert split.total_tokens == prefixed.total_tokens + assert aliased.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + framing + + +_KIMI_WEATHER: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} + + +def _byte_level_kimi_k3(path: Path) -> KimiK3TokenizerSource: + path.write_bytes(b"\n".join(base64.b64encode(bytes((value,))) + b" %d" % value for value in range(256))) + return KimiK3TokenizerSource(ranks_path=lambda: str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest()) + + +def _available_kimi_k3(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> tiktoken.Encoding: + source: Final = _byte_level_kimi_k3(tmp_path / "tiktoken.model") + monkeypatch.setattr( + litellm.proxy.proxy_server, "_KIMI_K3_ENCODINGS", KimiK3Encodings(source=source, retry_after_seconds=60) + ) + encoding: Final = load_kimi_k3_encoding(source) + assert encoding is not None + return encoding + + +def _shipped_kimi_k3_framing() -> int: + framing: Final = cost_map_entry("bedrock/us.moonshotai.kimi-k3")["token_count_framing"] + assert isinstance(framing, int) + return framing + + +@pytest.mark.asyncio +async def test_bedrock_kimi_k3_is_counted_with_its_tokenizer_plus_the_shipped_framing_in_both_request_formats( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + _framed_kimi(monkeypatch) + encoding: Final = _available_kimi_k3(monkeypatch, tmp_path) + question: Final = {"role": "user", "content": "Weather in Paris?"} + openai_messages: Final = [ + {"role": "system", "content": "Be terse."}, + question, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}} + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "Sun"}, + ] + + openai: Final = await count_request_tokens( + request=TokenCountRequest(model="kimi", messages=openai_messages, tools=[_KIMI_WEATHER]), call_endpoint=True + ) + anthropic: Final = await count_request_tokens( + request=TokenCountRequest( + model="kimi", + system="Be terse.", + messages=[ + question, + { + "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", + "input_schema": _KIMI_WEATHER["function"]["parameters"], + } + ], + ), + call_endpoint=True, + messages_format="anthropic", + ) + exact: Final = kimi_k3_tokens(encoding, "us.moonshotai.kimi-k3", openai_messages, [_KIMI_WEATHER]) + + assert exact is not None + assert (openai.total_tokens, openai.tokenizer_type, openai.count_source) == ( + exact + _shipped_kimi_k3_framing(), + "huggingface_tokenizer", + "local_estimate", + ) + assert (anthropic.total_tokens, anthropic.tokenizer_type) == (openai.total_tokens, "huggingface_tokenizer") + + +@pytest.mark.asyncio +async def test_bedrock_kimi_k3_deployment_with_a_custom_tokenizer_keeps_counting_with_it( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/kimi-custom-tokenizer") + _available_kimi_k3(monkeypatch, tmp_path) + custom: Final = litellm.utils._select_tokenizer(model="gpt-5.4-mini") + monkeypatch.setattr(litellm.utils, "_select_tokenizer", lambda model, custom_tokenizer=None: custom) + messages: Final = [{"role": "user", "content": "What is new?"}] + + counted: Final = await token_counter(request=TokenCountRequest(model="kimi", messages=messages), call_endpoint=True) + + assert counted.total_tokens == ( + litellm.token_counter(model=counted.model_used, messages=messages, custom_tokenizer=custom) + + kimi_k3_template_tokens(len(messages), False) + + _shipped_kimi_k3_framing() + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("messages", "tools", "template_tokens"), + [ + pytest.param([{"role": "user", "content": "What is new?"}], None, 14, id="one_message"), + pytest.param( + [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Weather in Paris?"}, + {"role": "assistant", "content": "Sunny."}, + {"role": "user", "content": "And Oslo?"}, + ], + [_KIMI_WEATHER], + 14 + 3 * 13 + 40, + id="turns_and_tools", + ), + ], +) +async def test_bedrock_kimi_k3_without_its_tokenizer_adds_the_template_offsets_to_the_local_estimate( + monkeypatch: pytest.MonkeyPatch, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + template_tokens: int, +): + _framed_kimi(monkeypatch) + + counted: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=messages, tools=tools), call_endpoint=True + ) + estimate: Final = litellm.token_counter( + model=counted.model_used, + messages=messages, + tools=tools, + custom_tokenizer=litellm.utils._select_tokenizer(model=counted.model_used), + ) + + assert counted.total_tokens == estimate + template_tokens + _shipped_kimi_k3_framing() + + +@pytest.mark.asyncio +async def test_slow_cost_map_lookup_does_not_stall_other_requests_on_the_worker(monkeypatch: pytest.MonkeyPatch): + real_get_model_info: Final = litellm.get_model_info + lookup_started: Final = threading.Event() + other_request_ran: Final = threading.Event() + released_by_the_other_request: Final[list[bool]] = [] + + def slow_get_model_info(model: str, *args: object, **kwargs: object): + if model == "slow-lookup-model": + lookup_started.set() + released_by_the_other_request.append(other_request_ran.wait(5)) + return real_get_model_info(model, *args, **kwargs) + + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "get_model_info", slow_get_model_info) + counting: Final = asyncio.create_task( + token_counter( + request=TokenCountRequest(model="slow-lookup-model", messages=[{"role": "user", "content": "hi"}]) + ) + ) + await asyncio.to_thread(lookup_started.wait, 5) + other_request_ran.set() + await counting + + assert set(released_by_the_other_request) == {True} + + +@pytest.mark.asyncio +async def test_uncountable_model_keeps_the_provider_error_when_local_counting_is_disabled( + monkeypatch: pytest.MonkeyPatch, +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "kimi", + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": {"token_count_framing": 78}, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + monkeypatch.setattr(litellm, "disable_token_counter", True) + + with pytest.raises(ProxyException) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert raised.value.code == "400" + assert "scripted 400 for kimi" in raised.value.message + + +@pytest.mark.parametrize( + ("model_info", "counted_with"), + [ + ({}, ("us.anthropic.claude-opus-5-5", "us-east-1", 447)), + ( + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}}, + ("us.anthropic.claude-opus-5-5", "ap-southeast-4", 447), + ), + ({"token_count_model": {}}, ("us.anthropic.claude-sonnet-5-5", "us-west-2", 300)), + ], +) +@pytest.mark.usefixtures("bedrock_sonnet_5_5_mapped_as_on_main") +@pytest.mark.asyncio +async def test_configured_token_count_model_overrides_the_shipped_one_and_an_empty_one_counts_the_model_itself( + monkeypatch: pytest.MonkeyPatch, model_info: dict[str, object], counted_with: tuple[str, str, int] +): + counter: Final = _ScriptedCounter( + { + "us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=300), + "us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447), + } + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-sonnet-5-5", + "aws_region_name": "us-west-2", + }, + "model_info": model_info, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + counted: Final = await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert (counter.calls[0][0], counter.regions[0], counted.total_tokens) == counted_with + assert len(counter.calls) == 1 + + +@pytest.mark.parametrize( + ("aws_region_name", "counted_in"), + [("eu-central-1", "eu-central-1"), ("us-west-2", "us-east-1")], +) +@pytest.mark.usefixtures("bedrock_sonnet_5_5_mapped_as_on_main") +@pytest.mark.asyncio +async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geography( + monkeypatch: pytest.MonkeyPatch, aws_region_name: str, counted_in: str +): + counter: Final = _ScriptedCounter({"anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": { + "model": "bedrock/anthropic.claude-sonnet-5-5", + "aws_region_name": aws_region_name, + }, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert (counter.calls[0][0], counter.regions[0]) == ("anthropic.claude-opus-5-5", counted_in) + + +class _SlowVault: + def sync_read_secret(self, secret_name: str) -> None: + time.sleep(0.5) + + +async def _longest_stall(done: asyncio.Event) -> float: + if done.is_set(): + return 0.0 + started: Final = time.monotonic() + await asyncio.sleep(0.01) + return max(time.monotonic() - started, await _longest_stall(done)) + + +@pytest.mark.usefixtures("bedrock_sonnet_5_5_mapped_as_on_main") +@pytest.mark.asyncio +async def test_region_lookup_for_a_shipped_count_model_leaves_the_event_loop_free(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + + counter: Final = _ScriptedCounter({"anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + + def bedrock_counter(deployment: object, model: str) -> tuple[_ScriptedCounter, str, str]: + return counter, model, "bedrock" + + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[{"model_name": "sonnet", "litellm_params": {"model": "bedrock/anthropic.claude-sonnet-5-5"}}] + ), + ) + monkeypatch.setattr(proxy_server, "_get_provider_token_counter", bedrock_counter) + monkeypatch.setattr(litellm, "secret_manager_client", _SlowVault()) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.HASHICORP_VAULT) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + done: Final = asyncio.Event() + + async def count() -> TokenCountResponse: + try: + return await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), + call_endpoint=True, + ) + finally: + done.set() + + started: Final = time.monotonic() + stall, response = await asyncio.gather(_longest_stall(done), count()) + + assert time.monotonic() - started >= 1.0 + assert stall < 0.4 + assert response.total_tokens == 447 + + +@pytest.mark.asyncio +async def test_framing_still_applies_while_the_deployment_is_cooling_down(monkeypatch: pytest.MonkeyPatch): + router: Final = Router( + model_list=[ + { + "model_name": "framed", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-unused-framed"}, + "model_info": {"id": "framed-deployment", "token_count_framing": 100}, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + messages: Final = [{"role": "user", "content": "hello there"}] + warm: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + + router.cooldown_cache.add_deployment_to_cooldown( + model_id="framed-deployment", + original_exception=Exception("rate limited"), + exception_status=429, + cooldown_time=60, + ) + cooling: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + + assert warm.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + 100 + assert cooling.total_tokens - 100 == litellm.token_counter(model="framed", messages=messages) + assert cooling.count_source == "local_estimate" + + class _GeminiPart(BaseModel): model_config = ConfigDict(frozen=True) text: str = "" diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index d73f5efa96b..6a6e4c5b4d4 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -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 diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a934f57f3d3..e80ca2a5820 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -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"}, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 16eaa1e3630..42c306cc3b3 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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