mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
feat(count_tokens)!: configurable token_count_framing and token_count_model, count source in the response (#46777)
Backport of #46777 to stable/1.104.x.
Cherry-picked from 75faaeab64 (main).
This commit is contained in:
parent
52de3b0a29
commit
b6f1a033de
32 changed files with 3571 additions and 150 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
502
litellm/llms/bedrock/count_tokens/kimi_k3.py
Normal file
502
litellm/llms/bedrock/count_tokens/kimi_k3.py
Normal file
|
|
@ -0,0 +1,502 @@
|
|||
"""
|
||||
Exact Bedrock token counts for Kimi K3, which has no CountTokens support.
|
||||
|
||||
Renders the Converse body litellm sends through Moonshot's Kimi K3 chat template and counts it with the
|
||||
official tokenizer, downloaded once from Hugging Face at a pinned revision and checked against its SHA-256.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from concurrent.futures import Future
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from itertools import chain, groupby
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Final, Literal, TypeAlias
|
||||
|
||||
import regex
|
||||
import tiktoken
|
||||
from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError, model_validator
|
||||
from typing_extensions import Self
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
chat_completion_tools_for,
|
||||
is_flat_responses_tool,
|
||||
)
|
||||
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
|
||||
KIMI_K3_REPO: Final = "moonshotai/Kimi-K3"
|
||||
KIMI_K3_REVISION: Final = "f831ab66814297da540d832a5235f8e904f29d06"
|
||||
KIMI_K3_RANKS_FILE: Final = "tiktoken.model"
|
||||
KIMI_K3_RANKS_SHA256: Final = "b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103"
|
||||
_PATTERN: Final = "|".join(
|
||||
(
|
||||
r"[\p{Han}]+",
|
||||
r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?",
|
||||
r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?",
|
||||
r"\p{N}{1,3}",
|
||||
r" ?[^\s\p{L}\p{N}]+[\r\n]*",
|
||||
r"\s*[\r\n]+",
|
||||
r"\s+(?!\S)",
|
||||
r"\s+",
|
||||
)
|
||||
)
|
||||
# Ends chunks only where the pattern above always splits, so their counts sum to the whole text's count while BPE
|
||||
# stays linear. Only a run of over 1024 characters with no such split, like 2000 spaces, is cut mid-piece.
|
||||
_CHUNK: Final = regex.compile(
|
||||
r"(?s).{1,1024}(?:(?<=\S)(?= \S)|(?<=\p{N})(?=\P{N})|(?<=[^\s\p{N}])(?=\p{N})|(?<=\p{L})(?=[^\p{L}\p{M}'])"
|
||||
r"|(?<=[\r\n])(?=\S)|\Z)|.{1024}\p{N}*"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class KimiK3TokenizerSource:
|
||||
ranks_path: Callable[[], str]
|
||||
sha256: str
|
||||
|
||||
|
||||
def _hugging_face_ranks_path() -> str:
|
||||
import huggingface_hub
|
||||
|
||||
return huggingface_hub.hf_hub_download(
|
||||
repo_id=KIMI_K3_REPO,
|
||||
filename=KIMI_K3_RANKS_FILE,
|
||||
revision=KIMI_K3_REVISION,
|
||||
local_files_only=litellm.disable_hf_tokenizer_download is True,
|
||||
)
|
||||
|
||||
|
||||
HUGGING_FACE_SOURCE: Final = KimiK3TokenizerSource(ranks_path=_hugging_face_ranks_path, sha256=KIMI_K3_RANKS_SHA256)
|
||||
|
||||
|
||||
def _loaded(source: KimiK3TokenizerSource) -> tiktoken.Encoding | str:
|
||||
try:
|
||||
data: Final = Path(source.ranks_path()).read_bytes()
|
||||
except Exception as e: # noqa: BLE001 # a download or read failure of any kind falls back to the estimate
|
||||
return type(e).__name__
|
||||
if hashlib.sha256(data).hexdigest() != source.sha256:
|
||||
return "SHA-256 mismatch"
|
||||
rows: Final = (line.split() for line in data.splitlines() if line)
|
||||
return tiktoken.Encoding(
|
||||
name="kimi_k3",
|
||||
pat_str=_PATTERN,
|
||||
mergeable_ranks={base64.b64decode(token): int(rank) for token, rank in rows},
|
||||
special_tokens={},
|
||||
)
|
||||
|
||||
|
||||
def load_kimi_k3_encoding(source: KimiK3TokenizerSource) -> tiktoken.Encoding | None:
|
||||
loaded: Final = _loaded(source)
|
||||
if isinstance(loaded, tiktoken.Encoding):
|
||||
return loaded
|
||||
verbose_logger.warning(
|
||||
"Kimi K3 tokenizer unavailable (%s), so Kimi K3 token counts are estimated until a later count loads it. "
|
||||
"Allow downloads of %s from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache",
|
||||
loaded,
|
||||
KIMI_K3_RANKS_FILE,
|
||||
KIMI_K3_REPO,
|
||||
KIMI_K3_REVISION,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _started_load(source: KimiK3TokenizerSource) -> "Future[tiktoken.Encoding | None]":
|
||||
encoding: Final = Future[tiktoken.Encoding | None]()
|
||||
threading.Thread(
|
||||
target=lambda: encoding.set_result(load_kimi_k3_encoding(source)), name="litellm-kimi-k3-tokenizer", daemon=True
|
||||
).start()
|
||||
return encoding
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _EncodingLoad:
|
||||
encoding: "Future[tiktoken.Encoding | None]"
|
||||
started_at: float
|
||||
is_retry: bool
|
||||
|
||||
def retry_due(self, now: float, retry_after_seconds: float) -> bool:
|
||||
return self.encoding.done() and self.encoding.result() is None and now - self.started_at >= retry_after_seconds
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class KimiK3Encodings:
|
||||
source: KimiK3TokenizerSource
|
||||
retry_after_seconds: float
|
||||
clock: Callable[[], float] = time.monotonic
|
||||
current: _EncodingLoad | None = None
|
||||
|
||||
def _load(self) -> _EncodingLoad:
|
||||
now: Final = self.clock()
|
||||
current: Final = self.current
|
||||
if current is not None and not current.retry_due(now, self.retry_after_seconds):
|
||||
return current
|
||||
load: Final = _EncodingLoad(encoding=_started_load(self.source), started_at=now, is_retry=current is not None)
|
||||
self.current = load
|
||||
return load
|
||||
|
||||
async def encoding(self, wait_seconds: float) -> tiktoken.Encoding | None:
|
||||
load: Final = self._load()
|
||||
remaining: Final = 0 if load.is_retry else load.started_at + wait_seconds - self.clock()
|
||||
if load.encoding.done() or remaining <= 0:
|
||||
return load.encoding.result() if load.encoding.done() else None
|
||||
await asyncio.wait((asyncio.wrap_future(load.encoding),), timeout=remaining)
|
||||
if load.encoding.done():
|
||||
return load.encoding.result()
|
||||
verbose_logger.warning(
|
||||
"Kimi K3 tokenizer still loading after %s s, so Kimi K3 token counts are estimated until it loads. "
|
||||
"Check that %s can be downloaded from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache",
|
||||
wait_seconds,
|
||||
KIMI_K3_RANKS_FILE,
|
||||
KIMI_K3_REPO,
|
||||
KIMI_K3_REVISION,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class _Converse(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
|
||||
class _Text(_Converse):
|
||||
text: str
|
||||
|
||||
|
||||
class _CachePoint(_Converse):
|
||||
cache_point: JsonValue = Field(alias="cachePoint")
|
||||
|
||||
|
||||
class _ToolUse(_Converse):
|
||||
tool_use_id: str = Field(alias="toolUseId")
|
||||
name: str
|
||||
arguments: Mapping[str, JsonValue] = Field(alias="input")
|
||||
|
||||
|
||||
class _ToolUseBlock(_Converse):
|
||||
tool_use: _ToolUse = Field(alias="toolUse")
|
||||
|
||||
|
||||
class _ToolResult(_Converse):
|
||||
tool_use_id: str = Field(alias="toolUseId")
|
||||
content: tuple[_Text, ...]
|
||||
status: str | None = None
|
||||
|
||||
|
||||
class _ToolResultBlock(_Converse):
|
||||
tool_result: _ToolResult = Field(alias="toolResult")
|
||||
|
||||
|
||||
class _ReasoningText(_Converse):
|
||||
text: str
|
||||
signature: str | None = None
|
||||
|
||||
|
||||
class _Reasoning(_Converse):
|
||||
reasoning_text: _ReasoningText = Field(alias="reasoningText")
|
||||
|
||||
|
||||
class _ReasoningBlock(_Converse):
|
||||
reasoning_content: _Reasoning = Field(alias="reasoningContent")
|
||||
|
||||
|
||||
class _UserMessage(_Converse):
|
||||
role: Literal["user"]
|
||||
content: tuple[_Text | _ToolResultBlock | _CachePoint, ...]
|
||||
|
||||
|
||||
class _AssistantMessage(_Converse):
|
||||
role: Literal["assistant"]
|
||||
content: tuple[_Text | _ToolUseBlock | _ReasoningBlock | _CachePoint, ...]
|
||||
|
||||
def calls(self) -> tuple[_ToolUse, ...]:
|
||||
return tuple(block.tool_use for block in self.content if isinstance(block, _ToolUseBlock))
|
||||
|
||||
|
||||
class _InputSchema(_Converse):
|
||||
definition: JsonValue = Field(alias="json")
|
||||
|
||||
|
||||
class _ToolSpec(_Converse):
|
||||
name: str
|
||||
description: str | None = None
|
||||
input_schema: _InputSchema = Field(alias="inputSchema")
|
||||
|
||||
def declared(self) -> Mapping[str, JsonValue]:
|
||||
described: Final[Mapping[str, JsonValue]] = (
|
||||
{} if self.description is None else {"description": self.description}
|
||||
)
|
||||
return {"name": self.name, **described, "parameters": self.input_schema.definition}
|
||||
|
||||
|
||||
class _ToolSpecBlock(_Converse):
|
||||
tool_spec: _ToolSpec = Field(alias="toolSpec")
|
||||
|
||||
|
||||
class _ToolConfig(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
tools: tuple[_ToolSpecBlock | _CachePoint, ...]
|
||||
|
||||
|
||||
_Message: TypeAlias = Annotated[_UserMessage | _AssistantMessage, Field(discriminator="role")]
|
||||
|
||||
|
||||
def _result_ids(message: _Message) -> frozenset[str]:
|
||||
if isinstance(message, _AssistantMessage):
|
||||
return frozenset()
|
||||
return frozenset(block.tool_result.tool_use_id for block in message.content if isinstance(block, _ToolResultBlock))
|
||||
|
||||
|
||||
def _call_ids(message: _Message | None) -> frozenset[str]:
|
||||
if not isinstance(message, _AssistantMessage):
|
||||
return frozenset()
|
||||
return frozenset(call.tool_use_id for call in message.calls())
|
||||
|
||||
|
||||
class _ConverseRequest(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
messages: tuple[_Message, ...]
|
||||
system: tuple[_Text | _CachePoint, ...] = ()
|
||||
tool_config: _ToolConfig | None = Field(default=None, alias="toolConfig")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _results_answer_the_previous_calls(self) -> Self:
|
||||
pairs: Final = zip((None, *self.messages), self.messages, strict=False)
|
||||
if not all(_result_ids(message) <= _call_ids(previous) for previous, message in pairs):
|
||||
raise ValueError("a toolResult does not answer a toolUse of the previous assistant message")
|
||||
return self
|
||||
|
||||
|
||||
class _Control(Enum):
|
||||
OPEN = "<|open|>"
|
||||
CLOSE = "<|close|>"
|
||||
SEP = "<|sep|>"
|
||||
END_OF_MSG = "<|end_of_msg|>"
|
||||
|
||||
|
||||
_Segment: TypeAlias = str | _Control
|
||||
_Segments: TypeAlias = tuple[_Segment, ...]
|
||||
|
||||
|
||||
def _open(tag: str, attrs: Iterable[tuple[str, str]] = ()) -> _Segments:
|
||||
escaped: Final = (
|
||||
(f" {key}", '="', value.replace("&", "&").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)
|
||||
|
|
@ -116,6 +116,7 @@ class OpenAITokenCounter(BaseTokenCounter):
|
|||
error=True,
|
||||
error_message=request_error,
|
||||
status_code=400,
|
||||
rejected_before_provider_call=True,
|
||||
)
|
||||
|
||||
deployment = deployment or {}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
164
litellm/proxy/common_utils/token_count_settings.py
Normal file
164
litellm/proxy/common_utils/token_count_settings.py
Normal file
|
|
@ -0,0 +1,164 @@
|
|||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass, replace
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.types.token_count import TokenCountSourceFields, TokenCountTokenizer
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404})
|
||||
TOKENIZER_LOAD_WAIT_SECONDS: Final = 5.0
|
||||
TOKENIZER_LOAD_REUSE_SECONDS: Final = 60.0
|
||||
TOKENIZER_LOAD_MAX_RUNNING: Final = 4
|
||||
|
||||
|
||||
class _UpstreamErrorDetail(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="allow")
|
||||
message: str
|
||||
|
||||
|
||||
class _UpstreamErrorEnvelope(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="allow")
|
||||
type: Literal["error"]
|
||||
error: _UpstreamErrorDetail
|
||||
|
||||
|
||||
def _joined(message: str, hint: str) -> str:
|
||||
return f"{message.rstrip().rstrip('.')}. {hint}"
|
||||
|
||||
|
||||
def _with_hint(message: str, hint: str) -> str:
|
||||
try:
|
||||
envelope: Final = _UpstreamErrorEnvelope.model_validate_json(message)
|
||||
except ValidationError:
|
||||
return _joined(AnthropicExceptionMapping.extract_error_message(message), hint)
|
||||
detail: Final = envelope.error.model_copy(update={"message": _joined(envelope.error.message, hint)})
|
||||
return envelope.model_copy(update={"error": detail}).model_dump_json()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TokenCountModel:
|
||||
model: str
|
||||
region: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TokenCountSettings:
|
||||
framing: int | None
|
||||
count_model: TokenCountModel | None
|
||||
tokenizer: TokenCountTokenizer | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProviderCountFailure:
|
||||
status_code: int
|
||||
message: str
|
||||
rejected_before_provider_call: bool = False
|
||||
|
||||
@property
|
||||
def cannot_count_model(self) -> bool:
|
||||
return not self.rejected_before_provider_call and self.status_code in CANNOT_COUNT_MODEL_STATUSES
|
||||
|
||||
def with_hint(self, hint: str) -> "ProviderCountFailure":
|
||||
return replace(self, message=_with_hint(self.message, hint))
|
||||
|
||||
def to_proxy_exception(self) -> ProxyException:
|
||||
return ProxyException(
|
||||
message=self.message,
|
||||
type="token_counting_error",
|
||||
param="model",
|
||||
code=self.status_code,
|
||||
)
|
||||
|
||||
|
||||
def cost_map_entry(model: str, custom_llm_provider: str | None = None) -> Mapping[str, object]:
|
||||
try:
|
||||
return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
|
||||
except Exception: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models
|
||||
return MappingProxyType({})
|
||||
|
||||
|
||||
def _framing(value: object) -> int | None:
|
||||
match value:
|
||||
case bool():
|
||||
return None
|
||||
case int() if value >= 0:
|
||||
return value
|
||||
case float() if value.is_integer() and value >= 0:
|
||||
return int(value)
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _non_blank(value: object) -> str | None:
|
||||
return value if isinstance(value, str) and value.strip() else None
|
||||
|
||||
|
||||
def _count_model(value: Mapping[str, object]) -> TokenCountModel | None:
|
||||
model: Final = _non_blank(value.get("model"))
|
||||
return None if model is None else TokenCountModel(model=model, region=_non_blank(value.get("region")))
|
||||
|
||||
|
||||
def _geography(region: str) -> str:
|
||||
return region.rsplit("-", 2)[0]
|
||||
|
||||
|
||||
def _within_deployment_geography(
|
||||
count_model: TokenCountModel | None, deployment_region: Callable[[], str]
|
||||
) -> TokenCountModel | None:
|
||||
if count_model is None or count_model.region is None or count_model.model.startswith("global."):
|
||||
return count_model
|
||||
if _geography(deployment_region()) == _geography(count_model.region):
|
||||
return count_model
|
||||
return replace(count_model, region=None)
|
||||
|
||||
|
||||
def _resolved_count_model(
|
||||
configured: object, default: object, deployment_region: Callable[[], str]
|
||||
) -> TokenCountModel | None:
|
||||
configured_mapping: Final = as_str_mapping(configured)
|
||||
if configured_mapping is not None:
|
||||
return _count_model(configured_mapping)
|
||||
default_mapping: Final = as_str_mapping(default)
|
||||
return (
|
||||
None
|
||||
if default_mapping is None
|
||||
else _within_deployment_geography(_count_model(default_mapping), deployment_region)
|
||||
)
|
||||
|
||||
|
||||
def _tokenizer(value: object) -> TokenCountTokenizer | None:
|
||||
match value:
|
||||
case "kimi_k3":
|
||||
return "kimi_k3"
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def resolve_token_count_settings(
|
||||
configured: Mapping[str, object], default: Mapping[str, object], deployment_region: Callable[[], str]
|
||||
) -> TokenCountSettings:
|
||||
framing: Final = _framing(configured.get("token_count_framing"))
|
||||
return TokenCountSettings(
|
||||
framing=framing if framing is not None else _framing(default.get("token_count_framing")),
|
||||
count_model=_resolved_count_model(
|
||||
configured.get("token_count_model"), default.get("token_count_model"), deployment_region
|
||||
),
|
||||
tokenizer=_tokenizer(default.get("token_count_tokenizer")),
|
||||
)
|
||||
|
||||
|
||||
def count_source_fields(response: TokenCountResponse) -> TokenCountSourceFields:
|
||||
match response.count_source:
|
||||
case None:
|
||||
return {}
|
||||
case "token_count_model":
|
||||
return {"litellm_count_source": "token_count_model", "litellm_count_model": response.model_used}
|
||||
case source:
|
||||
return {"litellm_count_source": source}
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
# ============================================================
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
16
litellm/types/token_count.py
Normal file
16
litellm/types/token_count.py
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
from typing import Literal, TypeAlias
|
||||
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
TokenCountSource: TypeAlias = Literal["provider", "token_count_model", "local_estimate"]
|
||||
TokenCountTokenizer: TypeAlias = Literal["kimi_k3"]
|
||||
|
||||
|
||||
class TokenCountModelInfo(TypedDict):
|
||||
model: ReadOnly[str]
|
||||
region: ReadOnly[NotRequired[str]]
|
||||
|
||||
|
||||
class TokenCountSourceFields(TypedDict, total=False):
|
||||
litellm_count_source: ReadOnly[TokenCountSource]
|
||||
litellm_count_model: ReadOnly[str]
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}]}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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}",
|
||||
|
|
|
|||
524
tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py
Normal file
524
tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py
Normal file
|
|
@ -0,0 +1,524 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import threading
|
||||
from collections.abc import Iterable
|
||||
from itertools import product
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
import tiktoken
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.bedrock.count_tokens.kimi_k3 import (
|
||||
KimiK3Encodings,
|
||||
KimiK3TokenizerSource,
|
||||
kimi_k3_anthropic_tokens,
|
||||
kimi_k3_tokens,
|
||||
load_kimi_k3_encoding,
|
||||
)
|
||||
from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig
|
||||
|
||||
MODEL: Final = "us.moonshotai.kimi-k3"
|
||||
_CONTROL: Final = re.compile(r"<\|(?:open|close|sep|end_of_msg)\|>")
|
||||
_WEATHER: Final = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather in °C",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}, "days": {"type": "integer"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _call(call_id: str, arguments: dict[str, object]) -> dict[str, object]:
|
||||
return {"id": call_id, "type": "function", "function": {"name": "get_weather", "arguments": json.dumps(arguments)}}
|
||||
|
||||
|
||||
def _written_ranks(path: Path, tokens: Iterable[bytes]) -> KimiK3TokenizerSource:
|
||||
path.write_bytes(b"\n".join(base64.b64encode(token) + b" %d" % rank for rank, token in enumerate(tokens)))
|
||||
return KimiK3TokenizerSource(ranks_path=lambda: str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest())
|
||||
|
||||
|
||||
def _single_bytes() -> tuple[bytes, ...]:
|
||||
return tuple(bytes((value,)) for value in range(256))
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def byte_source(tmp_path_factory: pytest.TempPathFactory) -> KimiK3TokenizerSource:
|
||||
return _written_ranks(tmp_path_factory.mktemp("bytes") / "tiktoken.model", _single_bytes())
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def byte_encoding(byte_source: KimiK3TokenizerSource) -> tiktoken.Encoding:
|
||||
loaded: Final = load_kimi_k3_encoding(byte_source)
|
||||
assert loaded is not None
|
||||
return loaded
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def pair_encoding(tmp_path_factory: pytest.TempPathFactory) -> tiktoken.Encoding:
|
||||
pairs: Final = (bytes(pair) for pair in product(range(256), repeat=2))
|
||||
source: Final = _written_ranks(tmp_path_factory.mktemp("pairs") / "tiktoken.model", (*_single_bytes(), *pairs))
|
||||
loaded: Final = load_kimi_k3_encoding(source)
|
||||
assert loaded is not None
|
||||
return loaded
|
||||
|
||||
|
||||
def _template_count(encoding: tiktoken.Encoding, rendered: str) -> int:
|
||||
texts: Final = _CONTROL.split(rendered)
|
||||
return len(_CONTROL.findall(rendered)) + sum(len(encoding.encode_ordinary(text)) for text in texts)
|
||||
|
||||
|
||||
_GENERATION_PROMPT: Final = '<|open|>message role="assistant"<|sep|><|open|>think<|sep|>'
|
||||
|
||||
|
||||
# Expected renders follow the Kimi K3 chat template in moonshotai/Kimi-K3 at revision f831ab6, which matched
|
||||
# Bedrock's prompt_tokens for us.moonshotai.kimi-k3 in 19 live calls on 2026-10-16
|
||||
@pytest.mark.parametrize(
|
||||
("messages", "tools", "rendered"),
|
||||
[
|
||||
pytest.param(
|
||||
[{"role": "user", "content": "Hello there"}],
|
||||
None,
|
||||
'<|open|>message role="user"<|sep|>Hello there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT,
|
||||
id="user",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": "Be terse."}, {"type": "text", "text": "Answer in French."}],
|
||||
},
|
||||
{"role": "user", "content": [{"type": "text", "text": "Hi"}, {"type": "text", "text": " there"}]},
|
||||
],
|
||||
None,
|
||||
'<|open|>message role="system"<|sep|>Be terse.Answer in French.<|close|>message<|sep|><|end_of_msg|>'
|
||||
'<|open|>message role="user"<|sep|>Hi there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT,
|
||||
id="text_parts_join",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{"role": "assistant", "content": "4"},
|
||||
{"role": "user", "content": "And 3+3?"},
|
||||
],
|
||||
None,
|
||||
'<|open|>message role="user"<|sep|>What is 2+2?<|close|>message<|sep|><|end_of_msg|>'
|
||||
'<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>'
|
||||
"<|open|>response<|sep|>4<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|>"
|
||||
'<|open|>message role="user"<|sep|>And 3+3?<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT,
|
||||
id="assistant_turn",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
{"role": "user", "content": "Weather in Paris and Oslo?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [_call("c1", {"city": "Paris", "days": 3}), _call("c2", {"city": "Oslo"})],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "c2", "content": "Snow"},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "Sun"},
|
||||
],
|
||||
[_WEATHER],
|
||||
'<|open|>message role="system" type="tool-declare"<|sep|># Tools\n'
|
||||
"Here are the available tools, described in JSONSchema.\n\n```json\n"
|
||||
'[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":'
|
||||
'{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},'
|
||||
'"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>'
|
||||
'<|open|>message role="user"<|sep|>Weather in Paris and Oslo?<|close|>message<|sep|><|end_of_msg|>'
|
||||
'<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>'
|
||||
"<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>"
|
||||
'<|open|>call tool="get_weather" index="1"<|sep|>'
|
||||
'<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|>'
|
||||
'<|open|>argument key="days" type="number"<|sep|>3<|close|>argument<|sep|><|close|>call<|sep|>'
|
||||
'<|open|>call tool="get_weather" index="2"<|sep|>'
|
||||
'<|open|>argument key="city" type="string"<|sep|>Oslo<|close|>argument<|sep|><|close|>call<|sep|>'
|
||||
"<|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>"
|
||||
'<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>'
|
||||
'<|open|>message role="tool" tool="get_weather" index="2"<|sep|>Snow<|close|>message<|sep|><|end_of_msg|>'
|
||||
+ _GENERATION_PROMPT,
|
||||
id="tool_results_out_of_call_order",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
{"role": "user", "content": "Weather?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
_call(
|
||||
"c1", {"a&b": 'x"y', "flag": True, "none": None, "list": [1, "zé"], "object": {"k": True}}
|
||||
)
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "Rain"},
|
||||
],
|
||||
[_WEATHER],
|
||||
'<|open|>message role="system" type="tool-declare"<|sep|># Tools\n'
|
||||
"Here are the available tools, described in JSONSchema.\n\n```json\n"
|
||||
'[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":'
|
||||
'{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},'
|
||||
'"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>'
|
||||
'<|open|>message role="user"<|sep|>Weather?<|close|>message<|sep|><|end_of_msg|>'
|
||||
'<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>'
|
||||
"<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>"
|
||||
'<|open|>call tool="get_weather" index="1"<|sep|>'
|
||||
'<|open|>argument key="a&b" type="string"<|sep|>x"y<|close|>argument<|sep|>'
|
||||
'<|open|>argument key="flag" type="boolean"<|sep|>true<|close|>argument<|sep|>'
|
||||
'<|open|>argument key="none" type="null"<|sep|>null<|close|>argument<|sep|>'
|
||||
'<|open|>argument key="list" type="array"<|sep|>[1, "zé"]<|close|>argument<|sep|>'
|
||||
'<|open|>argument key="object" type="object"<|sep|>{"k": true}<|close|>argument<|sep|>'
|
||||
"<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>"
|
||||
'<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Rain<|close|>message<|sep|><|end_of_msg|>'
|
||||
+ _GENERATION_PROMPT,
|
||||
id="argument_types_and_escaping",
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("encoding_fixture", ["byte_encoding", "pair_encoding"])
|
||||
def test_count_is_the_kimi_k3_chat_template_render_of_the_bedrock_request(
|
||||
request: pytest.FixtureRequest,
|
||||
encoding_fixture: str,
|
||||
messages: list[dict[str, object]],
|
||||
tools: list[dict[str, object]] | None,
|
||||
rendered: str,
|
||||
):
|
||||
encoding: Final[tiktoken.Encoding] = request.getfixturevalue(encoding_fixture)
|
||||
|
||||
assert kimi_k3_tokens(encoding, MODEL, messages, tools) == _template_count(encoding, rendered)
|
||||
|
||||
|
||||
def _parallel_calls(count: int) -> list[dict[str, object]]:
|
||||
ids: Final = tuple(f"c{number}" for number in range(1, count + 1))
|
||||
return [
|
||||
{"role": "user", "content": "Weather?"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [_call(call_id, {"city": "Paris"}) for call_id in ids]},
|
||||
*({"role": "tool", "tool_call_id": call_id, "content": "Sun"} for call_id in ids),
|
||||
]
|
||||
|
||||
|
||||
def test_tenth_parallel_call_and_its_result_are_numbered_ten(byte_encoding: tiktoken.Encoding):
|
||||
nine: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(9), [_WEATHER])
|
||||
ten: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(10), [_WEATHER])
|
||||
assert nine is not None and ten is not None
|
||||
|
||||
assert ten - nine == _template_count(
|
||||
byte_encoding,
|
||||
'<|open|>call tool="get_weather" index="10"<|sep|>'
|
||||
'<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|><|close|>call<|sep|>'
|
||||
'<|open|>message role="tool" tool="get_weather" index="10"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>',
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text",
|
||||
[
|
||||
pytest.param("The quick brown fox, jumps over 12345 lazy dogs. " * 60, id="prose"),
|
||||
pytest.param("a" + "1234567890" * 250 + "1", id="long_digit_run"),
|
||||
pytest.param("--- " * 600, id="symbols"),
|
||||
pytest.param("def f(x):\n return x # it's 42\n\n" * 80, id="code"),
|
||||
pytest.param("word next\t\ttab\n \n" * 120, id="whitespace"),
|
||||
pytest.param("naïve café, “quotes” & <tags> 中文汉字,日本語のテキスト 😀🎉\r\n" * 40, id="unicode"),
|
||||
pytest.param("We're sure they'd say it's DON'T or I'M " * 60, id="contractions"),
|
||||
pytest.param("abc123def4567ghi89 " * 120, id="letters_and_digits"),
|
||||
],
|
||||
)
|
||||
def test_long_text_counts_what_the_tokenizer_counts_for_the_whole_text(pair_encoding: tiktoken.Encoding, text: str):
|
||||
def counted(content: str) -> int | None:
|
||||
return kimi_k3_tokens(pair_encoding, MODEL, [{"role": "user", "content": content}], None)
|
||||
|
||||
short: Final = counted("x")
|
||||
assert short is not None
|
||||
|
||||
assert counted(text) == short - len(pair_encoding.encode_ordinary("x")) + len(pair_encoding.encode_ordinary(text))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("messages", "tools"),
|
||||
[
|
||||
pytest.param(
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is this?"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
|
||||
],
|
||||
}
|
||||
],
|
||||
None,
|
||||
id="image",
|
||||
),
|
||||
pytest.param(
|
||||
[{"role": "user", "content": "Weather?"}, {"role": "tool", "tool_call_id": "c1", "content": "Sun"}],
|
||||
[_WEATHER],
|
||||
id="tool_result_without_its_call",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_request_the_template_cannot_render_is_not_counted(
|
||||
pair_encoding: tiktoken.Encoding, messages: list[dict[str, object]], tools: list[dict[str, object]] | None
|
||||
):
|
||||
assert kimi_k3_tokens(pair_encoding, MODEL, messages, tools) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool",
|
||||
[
|
||||
pytest.param(
|
||||
{
|
||||
"type": "function",
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather in °C",
|
||||
"parameters": _WEATHER["function"]["parameters"],
|
||||
},
|
||||
id="responses_api",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather in °C",
|
||||
"input_schema": _WEATHER["function"]["parameters"],
|
||||
},
|
||||
id="anthropic_messages",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_tool_in_another_api_shape_counts_like_the_same_chat_completions_tool(
|
||||
pair_encoding: tiktoken.Encoding, tool: dict[str, object]
|
||||
):
|
||||
messages: Final = [{"role": "user", "content": "Weather in Paris?"}]
|
||||
|
||||
counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [tool])
|
||||
|
||||
assert counted is not None
|
||||
assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool",
|
||||
[
|
||||
pytest.param({"type": "web_search_20250305", "name": "web_search"}, id="anthropic_web_search"),
|
||||
pytest.param({"type": "web_search"}, id="responses_web_search"),
|
||||
pytest.param(
|
||||
{
|
||||
"type": "mcp",
|
||||
"server_label": "docs",
|
||||
"server_url": "https://mcp.example/mcp",
|
||||
"require_approval": "never",
|
||||
},
|
||||
id="mcp",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_tool_bedrock_inference_drops_counts_like_the_function_tools_alone(
|
||||
pair_encoding: tiktoken.Encoding, tool: dict[str, object]
|
||||
):
|
||||
messages: Final = [{"role": "user", "content": "Weather in Paris?"}]
|
||||
|
||||
counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, tool])
|
||||
|
||||
assert counted is not None
|
||||
assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER])
|
||||
|
||||
|
||||
def test_responses_custom_tool_counts_like_the_function_tool_inference_sends(pair_encoding: tiktoken.Encoding):
|
||||
messages: Final = [{"role": "user", "content": "Weather in Paris?"}]
|
||||
custom_tool: Final = {"type": "custom", "name": "grammar", "description": "Free-form SQL"}
|
||||
sent_to_inference, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(
|
||||
[dict(custom_tool)]
|
||||
)
|
||||
|
||||
without_it: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER])
|
||||
|
||||
counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, custom_tool])
|
||||
|
||||
assert counted is not None and without_it is not None
|
||||
assert counted > without_it
|
||||
assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, *sent_to_inference])
|
||||
|
||||
|
||||
def test_counting_a_request_with_an_image_url_never_downloads_the_image(
|
||||
pair_encoding: tiktoken.Encoding, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
monkeypatch.setattr(litellm, "user_url_validation", False)
|
||||
image: Final = respx_mock.get("https://images.example/cat.png").mock(return_value=httpx.Response(200))
|
||||
messages: Final = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is this?"},
|
||||
{"type": "image_url", "image_url": {"url": "https://images.example/cat.png"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
assert kimi_k3_tokens(pair_encoding, MODEL, messages, None) is None
|
||||
assert not image.called
|
||||
|
||||
|
||||
def test_anthropic_request_counts_like_the_same_openai_request(pair_encoding: tiktoken.Encoding):
|
||||
anthropic: Final = kimi_k3_anthropic_tokens(
|
||||
pair_encoding,
|
||||
MODEL,
|
||||
{
|
||||
"system": "Be terse.",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Weather in Paris?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "tool_use", "id": "c1", "name": "get_weather", "input": {"city": "Paris"}}],
|
||||
},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "Sun"}]},
|
||||
],
|
||||
"tools": [
|
||||
{
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather in °C",
|
||||
"input_schema": _WEATHER["function"]["parameters"],
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
openai: Final = kimi_k3_tokens(
|
||||
pair_encoding,
|
||||
MODEL,
|
||||
[
|
||||
{"role": "system", "content": "Be terse."},
|
||||
{"role": "user", "content": "Weather in Paris?"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [_call("c1", {"city": "Paris"})]},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "Sun"},
|
||||
],
|
||||
[_WEATHER],
|
||||
)
|
||||
|
||||
assert openai is not None
|
||||
assert anthropic == openai
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["checksum", "missing_file"])
|
||||
def test_tokenizer_that_fails_its_checksum_or_cannot_be_read_is_not_loaded(
|
||||
byte_source: KimiK3TokenizerSource, tmp_path: Path, caplog: pytest.LogCaptureFixture, failure: str
|
||||
):
|
||||
source: Final = (
|
||||
KimiK3TokenizerSource(ranks_path=byte_source.ranks_path, sha256=hashlib.sha256(b"other").hexdigest())
|
||||
if failure == "checksum"
|
||||
else KimiK3TokenizerSource(ranks_path=lambda: str(tmp_path / "missing"), sha256=byte_source.sha256)
|
||||
)
|
||||
verbose_logger.addHandler(caplog.handler)
|
||||
try:
|
||||
with caplog.at_level(logging.WARNING, logger=verbose_logger.name):
|
||||
loaded: Final = load_kimi_k3_encoding(source)
|
||||
finally:
|
||||
verbose_logger.removeHandler(caplog.handler)
|
||||
|
||||
assert loaded is None
|
||||
assert any("Kimi K3 tokenizer unavailable" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_slow_tokenizer_load_is_not_waited_on_past_the_wait_and_is_shared_once_done(
|
||||
byte_source: KimiK3TokenizerSource,
|
||||
):
|
||||
released: Final = threading.Event()
|
||||
|
||||
def slow_ranks_path() -> str:
|
||||
released.wait(10)
|
||||
return byte_source.ranks_path()
|
||||
|
||||
encodings: Final = KimiK3Encodings(
|
||||
source=KimiK3TokenizerSource(ranks_path=slow_ranks_path, sha256=byte_source.sha256),
|
||||
retry_after_seconds=60,
|
||||
clock=lambda: 0.0,
|
||||
)
|
||||
|
||||
assert await encodings.encoding(wait_seconds=0) is None
|
||||
released.set()
|
||||
loaded: Final = await encodings.encoding(wait_seconds=5)
|
||||
assert loaded is not None
|
||||
assert await encodings.encoding(wait_seconds=0) is loaded
|
||||
|
||||
|
||||
class _Clock:
|
||||
def __init__(self) -> None:
|
||||
self.now = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
|
||||
async def _loaded_eventually(encodings: KimiK3Encodings) -> tiktoken.Encoding | None:
|
||||
for _ in range(500):
|
||||
if (loaded := await encodings.encoding(wait_seconds=0)) is not None:
|
||||
return loaded
|
||||
await asyncio.sleep(0.01)
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_tokenizer_load_is_retried_in_the_background_only_after_the_retry_window(
|
||||
byte_source: KimiK3TokenizerSource, tmp_path: Path
|
||||
):
|
||||
clock: Final = _Clock()
|
||||
reads: Final[list[float]] = []
|
||||
released: Final = threading.Event()
|
||||
|
||||
def ranks_path() -> str:
|
||||
reads.append(clock.now)
|
||||
if len(reads) == 1:
|
||||
return str(tmp_path / "missing")
|
||||
released.wait(10)
|
||||
return byte_source.ranks_path()
|
||||
|
||||
encodings: Final = KimiK3Encodings(
|
||||
source=KimiK3TokenizerSource(ranks_path=ranks_path, sha256=byte_source.sha256),
|
||||
retry_after_seconds=60,
|
||||
clock=clock,
|
||||
)
|
||||
|
||||
assert await encodings.encoding(wait_seconds=5) is None
|
||||
clock.now = 59.0
|
||||
assert await encodings.encoding(wait_seconds=5) is None
|
||||
clock.now = 60.0
|
||||
assert await asyncio.wait_for(encodings.encoding(wait_seconds=5), timeout=1) is None
|
||||
released.set()
|
||||
assert await _loaded_eventually(encodings) is not None
|
||||
assert reads == [0.0, 60.0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tokenizer_load_still_running_after_the_wait_is_logged(
|
||||
byte_source: KimiK3TokenizerSource, caplog: pytest.LogCaptureFixture
|
||||
):
|
||||
released: Final = threading.Event()
|
||||
|
||||
def hung_ranks_path() -> str:
|
||||
released.wait(10)
|
||||
return byte_source.ranks_path()
|
||||
|
||||
encodings: Final = KimiK3Encodings(
|
||||
source=KimiK3TokenizerSource(ranks_path=hung_ranks_path, sha256=byte_source.sha256), retry_after_seconds=60
|
||||
)
|
||||
verbose_logger.addHandler(caplog.handler)
|
||||
try:
|
||||
with caplog.at_level(logging.WARNING, logger=verbose_logger.name):
|
||||
loaded: Final = await encodings.encoding(wait_seconds=0.05)
|
||||
finally:
|
||||
verbose_logger.removeHandler(caplog.handler)
|
||||
released.set()
|
||||
|
||||
assert loaded is None
|
||||
assert any("Kimi K3 tokenizer still loading" in record.getMessage() for record in caplog.records)
|
||||
|
|
@ -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}",
|
||||
|
|
|
|||
190
tests/unit/proxy/common_utils/test_token_count_settings.py
Normal file
190
tests/unit/proxy/common_utils/test_token_count_settings.py
Normal file
|
|
@ -0,0 +1,190 @@
|
|||
import json
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.common_utils.token_count_settings import (
|
||||
ProviderCountFailure,
|
||||
TokenCountModel,
|
||||
TokenCountSettings,
|
||||
count_source_fields,
|
||||
resolve_token_count_settings,
|
||||
)
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
_OPUS_US_EAST_1: Final = TokenCountModel(model="anthropic.claude-opus-5-5", region="us-east-1")
|
||||
_DEFAULT: Final = {
|
||||
"token_count_framing": 97,
|
||||
"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"},
|
||||
}
|
||||
|
||||
|
||||
def _in_us_west_2() -> str:
|
||||
return "us-west-2"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("configured", "expected"),
|
||||
[
|
||||
({}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
|
||||
(
|
||||
{"token_count_framing": 5, "token_count_model": {"model": "custom-model", "region": "ap-southeast-4"}},
|
||||
TokenCountSettings(framing=5, count_model=TokenCountModel(model="custom-model", region="ap-southeast-4")),
|
||||
),
|
||||
(
|
||||
{"token_count_model": {"model": "custom-model"}},
|
||||
TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)),
|
||||
),
|
||||
(
|
||||
{"token_count_model": {"model": "custom-model", "region": " "}},
|
||||
TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)),
|
||||
),
|
||||
({"token_count_framing": 0}, TokenCountSettings(framing=0, count_model=_OPUS_US_EAST_1)),
|
||||
(
|
||||
{"token_count_framing": None, "token_count_model": None},
|
||||
TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1),
|
||||
),
|
||||
(
|
||||
{"token_count_framing": True, "token_count_model": "custom-model"},
|
||||
TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1),
|
||||
),
|
||||
({"token_count_framing": 78.0}, TokenCountSettings(framing=78, count_model=_OPUS_US_EAST_1)),
|
||||
({"token_count_framing": 1.5}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
|
||||
({"token_count_framing": -1}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
|
||||
({"token_count_framing": "78"}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
|
||||
({"token_count_model": {}}, TokenCountSettings(framing=97, count_model=None)),
|
||||
({"token_count_model": ""}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
|
||||
(
|
||||
{"token_count_model": {"model": " ", "region": "us-east-1"}},
|
||||
TokenCountSettings(framing=97, count_model=None),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_config_overrides_the_cost_map_default_key_by_key(
|
||||
configured: dict[str, object], expected: TokenCountSettings
|
||||
) -> None:
|
||||
assert (
|
||||
resolve_token_count_settings(configured=configured, default=_DEFAULT, deployment_region=_in_us_west_2)
|
||||
== expected
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("deployment_region", "configured", "expected_region"),
|
||||
[
|
||||
("us-west-2", {}, "us-east-1"),
|
||||
("eu-central-1", {}, None),
|
||||
("ca-central-1", {}, None),
|
||||
("us-gov-west-1", {}, None),
|
||||
(
|
||||
"eu-central-1",
|
||||
{"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}},
|
||||
"us-east-1",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geography(
|
||||
deployment_region: str, configured: dict[str, object], expected_region: str | None
|
||||
) -> None:
|
||||
settings: Final = resolve_token_count_settings(
|
||||
configured=configured, default=_DEFAULT, deployment_region=lambda: deployment_region
|
||||
)
|
||||
assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region)
|
||||
|
||||
|
||||
def test_shipped_global_count_model_keeps_its_region_for_a_deployment_in_any_geography() -> None:
|
||||
settings: Final = resolve_token_count_settings(
|
||||
configured={},
|
||||
default={"token_count_model": {"model": "global.anthropic.claude-opus-5-5", "region": "us-east-1"}},
|
||||
deployment_region=lambda: "eu-west-1",
|
||||
)
|
||||
assert settings.count_model == TokenCountModel(model="global.anthropic.claude-opus-5-5", region="us-east-1")
|
||||
|
||||
|
||||
def test_shipped_australian_count_region_stays_with_an_australian_deployment() -> None:
|
||||
settings: Final = resolve_token_count_settings(
|
||||
configured={},
|
||||
default={"token_count_model": {"model": "au.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}},
|
||||
deployment_region=lambda: "ap-southeast-2",
|
||||
)
|
||||
assert settings.count_model == TokenCountModel(model="au.anthropic.claude-opus-5-5", region="ap-southeast-4")
|
||||
|
||||
|
||||
def test_no_setting_anywhere_means_the_model_counts_itself_without_framing() -> None:
|
||||
assert resolve_token_count_settings(
|
||||
configured={}, default={}, deployment_region=_in_us_west_2
|
||||
) == TokenCountSettings(framing=None, count_model=None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("value", "expected"),
|
||||
[("kimi_k3", "kimi_k3"), ("cl100k_base", None), (None, None), (True, None)],
|
||||
)
|
||||
def test_only_a_tokenizer_litellm_ships_is_taken_from_the_cost_map(value: object, expected: str | None) -> None:
|
||||
settings: Final = resolve_token_count_settings(
|
||||
configured={}, default={"token_count_tokenizer": value}, deployment_region=_in_us_west_2
|
||||
)
|
||||
assert settings.tokenizer == expected
|
||||
|
||||
|
||||
def _counted(source: str | None, model_used: str = "us.anthropic.claude-opus-5-5") -> TokenCountResponse:
|
||||
return TokenCountResponse(
|
||||
total_tokens=7,
|
||||
request_model="sonnet",
|
||||
model_used=model_used,
|
||||
tokenizer_type="bedrock_api",
|
||||
count_source=source,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("source", "expected"),
|
||||
[
|
||||
(None, {}),
|
||||
("provider", {"litellm_count_source": "provider"}),
|
||||
("local_estimate", {"litellm_count_source": "local_estimate"}),
|
||||
(
|
||||
"token_count_model",
|
||||
{"litellm_count_source": "token_count_model", "litellm_count_model": "us.anthropic.claude-opus-5-5"},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_count_source_fields_name_the_counting_model_only_when_token_count_model_counted(
|
||||
source: str | None, expected: dict[str, str]
|
||||
) -> None:
|
||||
assert count_source_fields(_counted(source)) == expected
|
||||
|
||||
|
||||
_UPSTREAM_404: Final = {
|
||||
"type": "error",
|
||||
"request_id": "req_upstream",
|
||||
"error": {"type": "not_found_error", "message": "The model 'anthropic.claude-opus-5-5' does not exist"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("message", "expected"),
|
||||
[
|
||||
(
|
||||
json.dumps(_UPSTREAM_404),
|
||||
{
|
||||
**_UPSTREAM_404,
|
||||
"error": {
|
||||
"type": "not_found_error",
|
||||
"message": "The model 'anthropic.claude-opus-5-5' does not exist. Enable eu-west-1",
|
||||
},
|
||||
},
|
||||
),
|
||||
(json.dumps({"message": "Too many tokens"}), "Too many tokens. Enable eu-west-1"),
|
||||
("scripted 404", "scripted 404. Enable eu-west-1"),
|
||||
(
|
||||
"Credential should be scoped to a valid region. ",
|
||||
"Credential should be scoped to a valid region. Enable eu-west-1",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_hint_lands_in_the_upstream_error_message_and_keeps_its_envelope(message: str, expected: object) -> None:
|
||||
hinted: Final = ProviderCountFailure(status_code=404, message=message).with_hint("Enable eu-west-1")
|
||||
|
||||
assert (json.loads(hinted.message) if isinstance(expected, dict) else hinted.message) == expected
|
||||
assert hinted.status_code == 404
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"},
|
||||
|
|
|
|||
12
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
12
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue