mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-23 05:34:41 +00:00
feat(count_tokens)!: configurable token_count_framing and token_count_model, count source in the response (#46777)
* feat(count_tokens): configurable token_count_framing and token_count_fallback_model, count source in the response
Count endpoints now say where each number came from in litellm_count_source
(provider, fallback_model, local_estimate) and name the fallback model in
litellm_count_fallback_model. When the provider answers 400 or 404 for a model
it cannot count, the proxy counts the same request with token_count_fallback_model.
token_count_framing adds a per-model constant to local estimates. Both settings
ship in the cost map for Bedrock Claude Sonnet/Haiku 5.5, Bedrock Kimi K3, and
Databricks Kimi K3 and DeepSeek V4 Flash, and can be overridden in model_info.
* test(count_tokens): classify the new cost map keys in the schema generator and expect count_source in exact-shape tests
* test(count_tokens): e2e matrix against real inference and integration wire tests for count source, fallback and framing
* fix(count_tokens): keep the provider error when local counting is disabled, name the fallback error, keep settings in cooldown
Moves the cost map lookup off the event loop
* ci(cost map): bound token_count_framing to non-negative and token_count_fallback_model to non-empty
* fix(count_tokens): keep a failed fallback's own status, ship fallbacks only where they count, add system.ai framing
A fallback model that answers 404 returns the requested model's error unchanged, and any other fallback
failure answers with its own status and message. The eu, apac, au and jp Sonnet and Haiku 5.5 entries
no longer ship the Opus 5.5 fallback, which AWS serves only in us-east-1. A configured empty
token_count_fallback_model turns the shipped fallback off. Config framing accepts whole-number floats and
rejects negatives. The Databricks system.ai Kimi K3 and DeepSeek V4 Flash names carry the same framing as
their serving endpoint names, and the e2e matrix runs on Bedrock Haiku 5.5 too.
* test(databricks): list the system.ai Kimi K3 and DeepSeek V4 Flash names with their published DBU rates
* fix(count_tokens): skip a token count fallback model for an hour after it answers 404 on that deployment
A us. or global. Sonnet 5.5 deployment outside us-east-1 paid two extra Bedrock
calls and about 200 ms on every count, since Opus 5.5 cannot count there either.
The Kimi K3 e2e also runs the large system prompt shape now.
* test(count_tokens): use the shared integration port allocator in the fallback wire test
* fix(cost map): keep the Databricks system.ai entries to framing only, so their spend stays as today
* test(count_tokens): cover the fallback 404 skip, inference profile base_model, shipped framing, and deployment-scoped token count settings
Adds integration cells for a fallback model that answers 404, an application inference profile ARN deployment whose base_model carries the shipped fallback, and the shipped Kimi K3 framing with token_count_framing 0 turning it off. Adds unit tests that keep token count settings off the shared backend cost map key and that store a token count setting on PATCH only when the operator changed it
* fix(count_tokens): keep the counter's own 400 and the provider's error when a framed estimate cannot read the body
A 400 the OpenAI or Gemini counter returns before calling the provider is the caller's answer, so it no longer tries the fallback model or a framed estimate. When the framed local estimate cannot read a content block, the client gets the provider's error instead of a 500. Adds /audit round 2's integration tests for the fallback leg's converted body, the Anthropic envelope on fallback failures, and both framing cases.
* fix(count_tokens): answer the provider's error whenever a framed estimate raises
A tool schema the local estimate cannot format, such as a boolean subschema or an array with no items, raised AttributeError or KeyError and answered a 500 on a framed deployment whose provider had already answered 400. Any estimate failure now answers the provider's own error.
* fix(count_tokens): load the custom tokenizer inside the guarded framed estimate
A framed deployment whose custom_tokenizer could not load answered a 500 or 503 where the provider had already answered 400, because the tokenizer was selected outside the estimate's error handling. Selecting it inside answers the provider's error, logs a warning with the estimate's error, and still raises the Rust extension's fork-guard errors.
* fix(count_tokens): look up framing under the deployment's custom_llm_provider
* fix(count_tokens): share one tokenizer download per worker and bound the wait after a provider error
A framed deployment whose custom_tokenizer is not cached locally made every
count wait out the Hugging Face hub's whole retry cycle when the hub could not
be reached, and a burst of those loads held the default thread limiter that
every count's cost map lookup also needs. Counts on a worker now share one
download per minute, a failed download is reused for the rest of that minute,
and after a provider error a count waits at most 5 s before answering the
provider's error. The swallowed-estimate warning logs only the error type.
* fix(count_tokens): keep one tokenizer download per key running and load the local route per request
* test(count_tokens): type the scripted counter and wait on events instead of wall time where possible
* fix(count_tokens): run framed tokenizer loads on their own capped threads and answer a tokenizer panic with the provider error
Framed loads used anyio's default 40 threads, so a hub that drops packets let about 40 hanging loads stall every count on the worker. Loads now run on a per-worker pool of 4, and a count that finds 4 loads running answers the provider's error at once. Any finished load is reused for a minute after it finishes, so a load that fell back to tiktoken is not retried on every count. A pyo3 PanicException from the tokenizer, while loading or counting, now answers the provider's error instead of a plain 500
* feat(proxy): count tokens with token_count_model in its pinned region
token_count_model {model, region} replaces token_count_fallback_model. Bedrock Sonnet and Haiku 5.5 count with Opus 5.5 in us-east-1, or Melbourne for au. profiles, and a failure names the region to enable
* fix(count_tokens): use the renamed count model constant in the e2e test and list token_count_model with the object keys
* fix(proxy): keep the upstream error envelope when adding the region hint
* test(proxy): annotate the token count setting echo test
* ci(cost map): require non-blank token_count_model strings and test the region rules
* fix(count_tokens): answer the token_count_model's error even when framing is set
* fix(count_tokens): keep the provider error for Google contents instead of a framed estimate
* fix(count_tokens): estimate Google contents the route translated to messages
* fix(count_tokens): keep the provider error for Google contents with non-text parts
* test(proxy): drive tokenizer wait and cost map tests by events and a skippable loop clock
* fix(count_tokens): check Google text-only parts without a KeysView comparison basedpyright flags
* feat(proxy): count Bedrock Kimi K3 tokens exactly with its pinned tokenizer
* test(bedrock): keep the Kimi K3 image check in-process
* fix(count_tokens): join the region hint to an upstream message that already ends with a period
Also pins that a blank token_count_model in config falls through to the cost map default
* fix(count_tokens): count Gemini tool-call turns and Responses tools, and stop tokenizer retries from stalling counts
The Google countTokens route estimates contents with functionCall, functionResponse and thoughtSignature parts again, and still answers the provider error for media, thought and unknown parts
The Kimi K3 exact count renders Responses API and Anthropic tools through the chat tool translation the Vertex count already used, now shared from the prompt factory, and falls back to the estimate when the Converse body cannot carry a tool
A failed Kimi K3 tokenizer load is retried in the background without making counts wait, and a load still running after the wait logs a warning
* fix(count_tokens): count Kimi K3 tools the way Bedrock inference sends them
A flat Responses API function or custom tool becomes the chat tool the
Responses route sends to inference, and every other tool goes to the
Converse body as is, so hosted tools Bedrock drops are left out of the
count instead of being counted or turning it into an estimate.
* test(count_tokens): annotate new token counter test locals with Final
* fix(count_tokens): keep shipped count regions inside the deployment's geography
A plain Bedrock Sonnet or Haiku 5.5 deployment in a non-US region no longer sends its count prompt to us-east-1. Global count models and configured token_count_model keep their pinned region. Also expect litellm_count_source in the Vertex Llama count test main added
* fix(count_tokens): resolve count settings off the event loop
This commit is contained in:
Mateo Wang
GitHub
parent
3fb95c8fe3
commit
75faaeab64
41 changed files with 4732 additions and 158 deletions
|
|
@ -6166,3 +6166,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.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)
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
|
|||
verbose_logger.debug("Processing CountTokens request for resolved model: %s", resolved_model)
|
||||
|
||||
# Get AWS region using existing LiteLLM function
|
||||
aws_region_name: Final = self._get_aws_region_name(
|
||||
aws_region_name: Final = self.get_aws_region_name(
|
||||
optional_params=litellm_params,
|
||||
model=resolved_model,
|
||||
model_id=None,
|
||||
|
|
|
|||
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.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)
|
||||
|
|
@ -52,7 +52,7 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM):
|
|||
) -> dict[str, JsonValue]:
|
||||
try:
|
||||
request: Final = _COUNT_REQUEST.validate_python(request_data)
|
||||
aws_region_name: Final = self._get_aws_region_name(
|
||||
aws_region_name: Final = self.get_aws_region_name(
|
||||
optional_params=litellm_params, model=resolved_model, model_id=None
|
||||
)
|
||||
endpoint_url: Final = self.get_mantle_count_tokens_endpoint(aws_region_name)
|
||||
|
|
|
|||
|
|
@ -124,6 +124,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
|
||||
|
|
@ -75,27 +76,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.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:
|
||||
|
|
@ -112,7 +92,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(
|
||||
|
|
@ -163,6 +143,7 @@ def count_context_error_response(
|
|||
error=True,
|
||||
error_message=error.error_message,
|
||||
status_code=error.status_code,
|
||||
rejected_before_provider_call=True,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -21144,6 +21144,7 @@
|
|||
"input_cost_per_token": 1.4e-07,
|
||||
"input_dbu_cost_per_token": 2e-06,
|
||||
"litellm_provider": "databricks",
|
||||
"token_count_framing": 97,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 393216,
|
||||
"max_tokens": 393216,
|
||||
|
|
@ -21166,6 +21167,11 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"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,
|
||||
"cache_read_input_token_cost": 1.3202e-07,
|
||||
|
|
@ -22134,6 +22140,7 @@
|
|||
"input_cost_per_token": 2.99999e-06,
|
||||
"input_dbu_cost_per_token": 4.2857e-05,
|
||||
"litellm_provider": "databricks",
|
||||
"token_count_framing": 100,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
|
|
@ -22157,6 +22164,11 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"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,
|
||||
"cache_read_input_token_cost": 5.0001e-07,
|
||||
|
|
@ -71540,6 +71552,8 @@
|
|||
"input_cost_per_token_flex": 1.65e-06,
|
||||
"input_cost_per_token_priority": 5.775e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_framing": 64,
|
||||
"token_count_tokenizer": "kimi_k3",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
|
|
@ -71570,6 +71584,8 @@
|
|||
"input_cost_per_token_flex": 1.5e-06,
|
||||
"input_cost_per_token_priority": 5.25e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_framing": 64,
|
||||
"token_count_tokenizer": "kimi_k3",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
|
|
@ -71600,6 +71616,8 @@
|
|||
"input_cost_per_token_flex": 1.65e-06,
|
||||
"input_cost_per_token_priority": 5.775e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_framing": 64,
|
||||
"token_count_tokenizer": "kimi_k3",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
|
|
@ -73100,6 +73118,10 @@
|
|||
"cache_read_input_token_cost": 1e-07,
|
||||
"input_cost_per_token": 2e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_model": {
|
||||
"model": "anthropic.claude-opus-5-5",
|
||||
"region": "us-east-1"
|
||||
},
|
||||
"supports_tool_search": true,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -73173,6 +73195,10 @@
|
|||
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
|
||||
},
|
||||
"au.anthropic.claude-sonnet-5-5": {
|
||||
"token_count_model": {
|
||||
"model": "au.anthropic.claude-opus-5-5",
|
||||
"region": "ap-southeast-4"
|
||||
},
|
||||
"bedrock_converse_supports_strict_tools": false,
|
||||
"cache_creation_input_token_cost": 2.75e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 4.4e-06,
|
||||
|
|
@ -73373,6 +73399,10 @@
|
|||
"cache_read_input_token_cost": 1e-07,
|
||||
"input_cost_per_token": 2e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_model": {
|
||||
"model": "global.anthropic.claude-opus-5-5",
|
||||
"region": "us-east-1"
|
||||
},
|
||||
"supports_tool_search": true,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -73582,6 +73612,10 @@
|
|||
"cache_read_input_token_cost": 1.1e-07,
|
||||
"input_cost_per_token": 2.2e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_model": {
|
||||
"model": "us.anthropic.claude-opus-5-5",
|
||||
"region": "us-east-1"
|
||||
},
|
||||
"supports_tool_search": true,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -74382,6 +74416,8 @@
|
|||
"input_cost_per_token_flex": 1.65e-06,
|
||||
"input_cost_per_token_priority": 5.775e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_framing": 64,
|
||||
"token_count_tokenizer": "kimi_k3",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
|
|
@ -74555,6 +74591,10 @@
|
|||
"cache_read_input_token_cost": 1e-08,
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_model": {
|
||||
"model": "anthropic.claude-opus-5-5",
|
||||
"region": "us-east-1"
|
||||
},
|
||||
"supports_tool_search": true,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -74633,6 +74673,10 @@
|
|||
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
|
||||
},
|
||||
"au.anthropic.claude-haiku-5-5": {
|
||||
"token_count_model": {
|
||||
"model": "au.anthropic.claude-opus-5-5",
|
||||
"region": "ap-southeast-4"
|
||||
},
|
||||
"input_cost_per_token_above_100k_tokens": 5.5e-07,
|
||||
"output_cost_per_token_above_100k_tokens": 2.75e-06,
|
||||
"cache_creation_input_token_cost_above_100k_tokens": 6.875e-07,
|
||||
|
|
@ -74863,6 +74907,10 @@
|
|||
"cache_read_input_token_cost": 1e-08,
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_model": {
|
||||
"model": "global.anthropic.claude-opus-5-5",
|
||||
"region": "us-east-1"
|
||||
},
|
||||
"supports_tool_search": true,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -75010,6 +75058,10 @@
|
|||
"cache_read_input_token_cost": 1.1e-08,
|
||||
"input_cost_per_token": 1.1e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_model": {
|
||||
"model": "us.anthropic.claude-opus-5-5",
|
||||
"region": "us-east-1"
|
||||
},
|
||||
"supports_tool_search": true,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
|
|||
|
|
@ -48,6 +48,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
|
||||
from litellm.proxy.litellm_pre_call_utils import should_auto_drop_params_for_agentic_cli
|
||||
from litellm.router import Router
|
||||
|
||||
|
|
@ -420,7 +421,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}
|
||||
|
|
@ -21,6 +21,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg
|
|||
_read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
read_request_body,
|
||||
)
|
||||
from litellm.proxy.common_utils.token_count_settings import count_source_fields
|
||||
from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
|
||||
|
||||
|
||||
|
|
@ -309,10 +310,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
|
||||
|
||||
|
||||
# ============================================================
|
||||
|
|
|
|||
|
|
@ -55,6 +55,7 @@ from typing import (
|
|||
)
|
||||
|
||||
import anyio
|
||||
import anyio.lowlevel
|
||||
import websockets
|
||||
import websockets.exceptions
|
||||
from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, ValidationError, with_config
|
||||
|
|
@ -99,6 +100,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,
|
||||
AwsErrorTypeProxyException,
|
||||
|
|
@ -149,6 +157,15 @@ from litellm.proxy.common_utils.realtime_utils import ( # noqa: F401, RUF100 #
|
|||
_realtime_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
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 (
|
||||
|
|
@ -304,6 +321,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,
|
||||
|
|
@ -337,6 +355,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
drop_params_flag,
|
||||
get_litellm_metadata_from_kwargs,
|
||||
get_parent_otel_span_from_kwargs,
|
||||
iter_items,
|
||||
)
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
|
@ -993,7 +1012,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
|
||||
|
||||
|
|
@ -14673,25 +14698,36 @@ 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, object] | 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,
|
||||
context_management: CountContextManagement | None = 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,
|
||||
|
|
@ -14700,24 +14736,325 @@ async def _try_provider_token_count(
|
|||
context_management=context_management,
|
||||
)
|
||||
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}
|
||||
)
|
||||
|
|
@ -14728,6 +15065,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 _router_default_litellm_params(router: Router) -> Mapping[str, object]:
|
||||
defaults: Final[Mapping[str, object] | None] = getattr(router, "default_litellm_params", None)
|
||||
if defaults is None:
|
||||
|
|
@ -14831,6 +15176,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
|
||||
|
|
@ -14850,24 +15203,33 @@ 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,
|
||||
context_management=_count_context_management(context_management, request.model, request_litellm_params),
|
||||
),
|
||||
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,
|
||||
context_management=_count_context_management(context_management, request.model, request_litellm_params),
|
||||
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:
|
||||
|
|
@ -14885,11 +15247,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
|
||||
|
|
@ -14902,18 +15259,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",
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -67,6 +67,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.hooks.proxy_track_cost_callback import ProxyDBLogger
|
||||
from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
|
|
@ -86,6 +87,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:
|
||||
|
|
@ -1373,7 +1375,7 @@ class _ResponsesApiErrorBody(TypedDict):
|
|||
error: ReadOnly[_ResponsesApiErrorDetail]
|
||||
|
||||
|
||||
class _ResponsesInputTokensResult(TypedDict):
|
||||
class _ResponsesInputTokensResult(TokenCountSourceFields):
|
||||
object: ReadOnly[str]
|
||||
input_tokens: ReadOnly[int]
|
||||
|
||||
|
|
@ -1610,6 +1612,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
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from typing_extensions import (
|
|||
)
|
||||
|
||||
from litellm.types.llms.openai import ChatCompletionFileObject, 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
|
||||
|
|
@ -320,12 +321,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]
|
||||
|
|
@ -58,6 +58,7 @@ from litellm.types.llms.base import (
|
|||
LiteLLMPydanticObjectBase,
|
||||
)
|
||||
from litellm.types.mcp import MCPServerCostInfo
|
||||
from litellm.types.token_count import TokenCountModelInfo, TokenCountSource, TokenCountTokenizer
|
||||
from litellm.types.workload_identity import (
|
||||
ANTHROPIC_WIF_KWARGS_KEYS,
|
||||
OPENAI_WIF_KWARGS_KEYS,
|
||||
|
|
@ -230,6 +231,9 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
|
|||
bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None
|
||||
bedrock_converse_supports_strict_tools: bool | None
|
||||
supports_regex_lookaround: ReadOnly[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):
|
||||
|
|
@ -4051,11 +4055,14 @@ DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset(
|
|||
{"supports_regex_lookaround", "supports_assistant_prefill", "supports_bedrock_native_structured_output"}
|
||||
)
|
||||
|
||||
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_CAPABILITY_FIELDS
|
||||
- DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -4492,6 +4499,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
|
||||
|
|
@ -4502,6 +4510,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):
|
||||
|
|
|
|||
|
|
@ -6694,6 +6694,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():
|
||||
|
|
|
|||
|
|
@ -21144,6 +21144,7 @@
|
|||
"input_cost_per_token": 1.4e-07,
|
||||
"input_dbu_cost_per_token": 2e-06,
|
||||
"litellm_provider": "databricks",
|
||||
"token_count_framing": 97,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 393216,
|
||||
"max_tokens": 393216,
|
||||
|
|
@ -21166,6 +21167,11 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"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,
|
||||
"cache_read_input_token_cost": 1.3202e-07,
|
||||
|
|
@ -22134,6 +22140,7 @@
|
|||
"input_cost_per_token": 2.99999e-06,
|
||||
"input_dbu_cost_per_token": 4.2857e-05,
|
||||
"litellm_provider": "databricks",
|
||||
"token_count_framing": 100,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
|
|
@ -22157,6 +22164,11 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"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,
|
||||
"cache_read_input_token_cost": 5.0001e-07,
|
||||
|
|
@ -71540,6 +71552,8 @@
|
|||
"input_cost_per_token_flex": 1.65e-06,
|
||||
"input_cost_per_token_priority": 5.775e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_framing": 64,
|
||||
"token_count_tokenizer": "kimi_k3",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
|
|
@ -71570,6 +71584,8 @@
|
|||
"input_cost_per_token_flex": 1.5e-06,
|
||||
"input_cost_per_token_priority": 5.25e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_framing": 64,
|
||||
"token_count_tokenizer": "kimi_k3",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
|
|
@ -71600,6 +71616,8 @@
|
|||
"input_cost_per_token_flex": 1.65e-06,
|
||||
"input_cost_per_token_priority": 5.775e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_framing": 64,
|
||||
"token_count_tokenizer": "kimi_k3",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
|
|
@ -73100,6 +73118,10 @@
|
|||
"cache_read_input_token_cost": 1e-07,
|
||||
"input_cost_per_token": 2e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_model": {
|
||||
"model": "anthropic.claude-opus-5-5",
|
||||
"region": "us-east-1"
|
||||
},
|
||||
"supports_tool_search": true,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -73173,6 +73195,10 @@
|
|||
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
|
||||
},
|
||||
"au.anthropic.claude-sonnet-5-5": {
|
||||
"token_count_model": {
|
||||
"model": "au.anthropic.claude-opus-5-5",
|
||||
"region": "ap-southeast-4"
|
||||
},
|
||||
"bedrock_converse_supports_strict_tools": false,
|
||||
"cache_creation_input_token_cost": 2.75e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 4.4e-06,
|
||||
|
|
@ -73373,6 +73399,10 @@
|
|||
"cache_read_input_token_cost": 1e-07,
|
||||
"input_cost_per_token": 2e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_model": {
|
||||
"model": "global.anthropic.claude-opus-5-5",
|
||||
"region": "us-east-1"
|
||||
},
|
||||
"supports_tool_search": true,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -73582,6 +73612,10 @@
|
|||
"cache_read_input_token_cost": 1.1e-07,
|
||||
"input_cost_per_token": 2.2e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_model": {
|
||||
"model": "us.anthropic.claude-opus-5-5",
|
||||
"region": "us-east-1"
|
||||
},
|
||||
"supports_tool_search": true,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -74382,6 +74416,8 @@
|
|||
"input_cost_per_token_flex": 1.65e-06,
|
||||
"input_cost_per_token_priority": 5.775e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_framing": 64,
|
||||
"token_count_tokenizer": "kimi_k3",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
|
|
@ -74555,6 +74591,10 @@
|
|||
"cache_read_input_token_cost": 1e-08,
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_model": {
|
||||
"model": "anthropic.claude-opus-5-5",
|
||||
"region": "us-east-1"
|
||||
},
|
||||
"supports_tool_search": true,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -74633,6 +74673,10 @@
|
|||
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
|
||||
},
|
||||
"au.anthropic.claude-haiku-5-5": {
|
||||
"token_count_model": {
|
||||
"model": "au.anthropic.claude-opus-5-5",
|
||||
"region": "ap-southeast-4"
|
||||
},
|
||||
"input_cost_per_token_above_100k_tokens": 5.5e-07,
|
||||
"output_cost_per_token_above_100k_tokens": 2.75e-06,
|
||||
"cache_creation_input_token_cost_above_100k_tokens": 6.875e-07,
|
||||
|
|
@ -74863,6 +74907,10 @@
|
|||
"cache_read_input_token_cost": 1e-08,
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_model": {
|
||||
"model": "global.anthropic.claude-opus-5-5",
|
||||
"region": "us-east-1"
|
||||
},
|
||||
"supports_tool_search": true,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -75010,6 +75058,10 @@
|
|||
"cache_read_input_token_cost": 1.1e-08,
|
||||
"input_cost_per_token": 1.1e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"token_count_model": {
|
||||
"model": "us.anthropic.claude-opus-5-5",
|
||||
"region": "us-east-1"
|
||||
},
|
||||
"supports_tool_search": true,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
|
|||
|
|
@ -1210,6 +1210,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,
|
||||
|
|
|
|||
362
tests/e2e/llm_translation/test_count_tokens_source_e2e.py
Normal file
362
tests/e2e/llm_translation/test_count_tokens_source_e2e.py
Normal file
|
|
@ -0,0 +1,362 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from anthropic import Anthropic, BadRequestError, Omit, omit
|
||||
from anthropic.types import DocumentBlockParam, MessageParam, TextBlockParam, ToolParam
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import unwrap
|
||||
from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta
|
||||
from lifecycle import ResourceManager
|
||||
from models import (
|
||||
ChatMessage,
|
||||
CountTokensResponse,
|
||||
LiteLLMParamsBody,
|
||||
ModelInfoBody,
|
||||
ModelNewBody,
|
||||
ResponsesInputTokensResponse,
|
||||
TokenCounterBody,
|
||||
TokenCountModelBody,
|
||||
)
|
||||
from proxy_client import ProxyClient
|
||||
from reportlab.lib.pagesizes import LETTER
|
||||
from reportlab.pdfgen.canvas import Canvas
|
||||
from sdk_clients import NO_PROXY_CACHE, SdkClients
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
BEDROCK_SONNET: Final = "bedrock/us.anthropic.claude-sonnet-5-5"
|
||||
BEDROCK_HAIKU: Final = "bedrock/us.anthropic.claude-haiku-5-5"
|
||||
COST_MAP_COUNT_MODEL: Final = "us.anthropic.claude-opus-5-5"
|
||||
CONFIGURED_COUNT_MODEL: Final = "us.anthropic.claude-sonnet-5"
|
||||
ANTHROPIC_SONNET: Final = "anthropic/claude-sonnet-5-5"
|
||||
BEDROCK_KIMI: Final = "bedrock/us.moonshotai.kimi-k3"
|
||||
BEDROCK_NOVA: Final = "bedrock/us.amazon.nova-pro-v1:0"
|
||||
OPENAI_GPT: Final = "openai/gpt-5.6"
|
||||
|
||||
COUNT_TOLERANCE: Final = 0.05
|
||||
MAX_TOKENS: Final = 16
|
||||
PDF_LINES_PER_PAGE: Final = 45
|
||||
|
||||
LARGE_SYSTEM: Final = " ".join(
|
||||
f"Rule {n}: answer in plain English, cite clause {n} when it applies, and never invent a figure."
|
||||
for n in range(1, 301)
|
||||
)
|
||||
LARGE_TOOL: Final[ToolParam] = {
|
||||
"name": "search_orders",
|
||||
"description": " ".join(
|
||||
f"Filter {n} narrows the order search by the attribute recorded in field_{n}." for n in range(1, 151)
|
||||
),
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
f"field_{n}": {"type": "string", "description": f"Order attribute number {n}, matched exactly."}
|
||||
for n in range(1, 61)
|
||||
},
|
||||
"required": ["field_1"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Served:
|
||||
model: str
|
||||
key: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PromptShape:
|
||||
name: str
|
||||
system: str | None = None
|
||||
tool: ToolParam | None = None
|
||||
pdf_pages: int = 0
|
||||
|
||||
|
||||
SHAPES: Final = (
|
||||
PromptShape(name="baseline"),
|
||||
PromptShape(name="large_system", system=LARGE_SYSTEM),
|
||||
PromptShape(name="large_tool", tool=LARGE_TOOL),
|
||||
PromptShape(name="system_and_tool", system=LARGE_SYSTEM, tool=LARGE_TOOL),
|
||||
PromptShape(name="pdf_1_page", pdf_pages=1),
|
||||
PromptShape(name="pdf_12_pages", pdf_pages=12),
|
||||
)
|
||||
SHAPE_IDS: Final = tuple(shape.name for shape in SHAPES)
|
||||
TEXT_SHAPES: Final = SHAPES[:2]
|
||||
TEXT_SHAPE_IDS: Final = SHAPE_IDS[:2]
|
||||
|
||||
|
||||
def _bedrock(model: str) -> LiteLLMParamsBody:
|
||||
return LiteLLMParamsBody(
|
||||
model=model,
|
||||
aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID",
|
||||
aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY",
|
||||
aws_region_name="os.environ/AWS_REGION",
|
||||
)
|
||||
|
||||
|
||||
def _serve(proxy: ProxyClient, backend: LiteLLMParamsBody, model_info: ModelInfoBody | None = None) -> Iterator[Served]:
|
||||
resources: Final = ResourceManager(client=proxy)
|
||||
try:
|
||||
model: Final = f"e2e-count-source-{unique_marker()}"
|
||||
model_id: Final = proxy.register_model(
|
||||
ModelNewBody(model_name=model, litellm_params=backend, model_info=model_info or ModelInfoBody())
|
||||
)
|
||||
resources.defer(lambda: proxy.delete_model(model_id))
|
||||
yield Served(model=model, key=resources.key())
|
||||
finally:
|
||||
resources.teardown()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def bedrock_sonnet(proxy: ProxyClient) -> Iterator[Served]:
|
||||
yield from _serve(proxy, _bedrock(BEDROCK_SONNET))
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def bedrock_haiku(proxy: ProxyClient) -> Iterator[Served]:
|
||||
yield from _serve(proxy, _bedrock(BEDROCK_HAIKU))
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def bedrock_sonnet_with_configured_count_model(proxy: ProxyClient) -> Iterator[Served]:
|
||||
yield from _serve(
|
||||
proxy,
|
||||
_bedrock(BEDROCK_SONNET),
|
||||
ModelInfoBody(token_count_model=TokenCountModelBody(model=CONFIGURED_COUNT_MODEL)),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def anthropic_sonnet(proxy: ProxyClient) -> Iterator[Served]:
|
||||
yield from _serve(proxy, LiteLLMParamsBody(model=ANTHROPIC_SONNET, api_key="os.environ/ANTHROPIC_API_KEY"))
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def bedrock_kimi(proxy: ProxyClient) -> Iterator[Served]:
|
||||
yield from _serve(proxy, _bedrock(BEDROCK_KIMI))
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def bedrock_nova(proxy: ProxyClient) -> Iterator[Served]:
|
||||
yield from _serve(proxy, _bedrock(BEDROCK_NOVA))
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def openai_gpt(proxy: ProxyClient) -> Iterator[Served]:
|
||||
yield from _serve(proxy, LiteLLMParamsBody(model=OPENAI_GPT, api_key="os.environ/OPENAI_API_KEY"))
|
||||
|
||||
|
||||
def _pdf(pages: int) -> str:
|
||||
buffer: Final = io.BytesIO()
|
||||
canvas: Final = Canvas(buffer, pagesize=LETTER, invariant=True)
|
||||
for page in range(1, pages + 1):
|
||||
for line in range(1, PDF_LINES_PER_PAGE + 1):
|
||||
canvas.drawString(
|
||||
72, 740 - 14 * line, f"Page {page}, line {line}: revenue, costs and headcount for region {line}."
|
||||
)
|
||||
canvas.showPage()
|
||||
canvas.save()
|
||||
return base64.b64encode(buffer.getvalue()).decode()
|
||||
|
||||
|
||||
def _user_turn(shape: PromptShape) -> MessageParam:
|
||||
question: Final[TextBlockParam] = {"type": "text", "text": f"Reply with one word. {unique_marker()}"}
|
||||
if shape.pdf_pages == 0:
|
||||
return {"role": "user", "content": [question]}
|
||||
document: Final[DocumentBlockParam] = {
|
||||
"type": "document",
|
||||
"source": {"type": "base64", "media_type": "application/pdf", "data": _pdf(shape.pdf_pages)},
|
||||
}
|
||||
return {"role": "user", "content": [document, question]}
|
||||
|
||||
|
||||
def _count_and_bill(client: Anthropic, model: str, shape: PromptShape) -> tuple[CountTokensResponse, int]:
|
||||
messages: Final = [_user_turn(shape)]
|
||||
system: Final[str | Omit] = omit if shape.system is None else shape.system
|
||||
tools: Final[list[ToolParam] | Omit] = omit if shape.tool is None else [shape.tool]
|
||||
counted: Final = CountTokensResponse.model_validate_json(
|
||||
client.messages.with_raw_response.count_tokens(model=model, messages=messages, system=system, tools=tools).text
|
||||
)
|
||||
usage: Final = client.messages.create(
|
||||
model=model,
|
||||
max_tokens=MAX_TOKENS,
|
||||
messages=messages,
|
||||
system=system,
|
||||
tools=tools,
|
||||
extra_body=NO_PROXY_CACHE,
|
||||
).usage
|
||||
billed: Final = usage.input_tokens + (usage.cache_creation_input_tokens or 0) + (usage.cache_read_input_tokens or 0)
|
||||
return counted, billed
|
||||
|
||||
|
||||
def _within_tolerance(counted: int, billed: int) -> bool:
|
||||
return abs(counted - billed) <= billed * COUNT_TOLERANCE
|
||||
|
||||
|
||||
def _assert_counted_by_the_cost_map_count_model(served: Served, sdk: SdkClients, shape: PromptShape) -> None:
|
||||
counted, billed = _count_and_bill(sdk.anthropic(served.key), served.model, shape)
|
||||
|
||||
assert counted.litellm_count_source == "token_count_model", counted
|
||||
assert counted.litellm_count_model == COST_MAP_COUNT_MODEL, counted
|
||||
assert _within_tolerance(counted.input_tokens, billed), (counted, billed)
|
||||
|
||||
|
||||
class TestBedrockClaudeCountsWithCostMapCountModel:
|
||||
@pytest.mark.parametrize("shape", SHAPES, ids=SHAPE_IDS)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.LLM_TRANSLATION,
|
||||
route=Route.COUNT_TOKENS,
|
||||
providers=(Provider.BEDROCK,),
|
||||
models=(BEDROCK_SONNET, COST_MAP_COUNT_MODEL),
|
||||
capabilities=(Capability.FUNCTION_CALLING, Capability.PDF_INPUT),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_sonnet_count_names_the_count_model_and_matches_billed_input(
|
||||
self, bedrock_sonnet: Served, sdk: SdkClients, shape: PromptShape
|
||||
) -> None:
|
||||
_assert_counted_by_the_cost_map_count_model(bedrock_sonnet, sdk, shape)
|
||||
|
||||
@pytest.mark.parametrize("shape", SHAPES, ids=SHAPE_IDS)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.LLM_TRANSLATION,
|
||||
route=Route.COUNT_TOKENS,
|
||||
providers=(Provider.BEDROCK,),
|
||||
models=(BEDROCK_HAIKU, COST_MAP_COUNT_MODEL),
|
||||
capabilities=(Capability.FUNCTION_CALLING, Capability.PDF_INPUT),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_haiku_count_names_the_count_model_and_matches_billed_input(
|
||||
self, bedrock_haiku: Served, sdk: SdkClients, shape: PromptShape
|
||||
) -> None:
|
||||
_assert_counted_by_the_cost_map_count_model(bedrock_haiku, sdk, shape)
|
||||
|
||||
|
||||
class TestAnthropicClaudeCountsWithProvider:
|
||||
@pytest.mark.parametrize("shape", SHAPES, ids=SHAPE_IDS)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.LLM_TRANSLATION,
|
||||
route=Route.COUNT_TOKENS,
|
||||
providers=(Provider.ANTHROPIC,),
|
||||
models=(ANTHROPIC_SONNET,),
|
||||
capabilities=(Capability.FUNCTION_CALLING, Capability.PDF_INPUT),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_count_comes_from_the_provider_and_matches_billed_input(
|
||||
self, anthropic_sonnet: Served, sdk: SdkClients, shape: PromptShape
|
||||
) -> None:
|
||||
counted, billed = _count_and_bill(sdk.anthropic(anthropic_sonnet.key), anthropic_sonnet.model, shape)
|
||||
|
||||
assert counted.litellm_count_source == "provider", counted
|
||||
assert counted.litellm_count_model is None, counted
|
||||
assert _within_tolerance(counted.input_tokens, billed), (counted, billed)
|
||||
|
||||
|
||||
class TestBedrockKimiCountsWithFramedLocalEstimate:
|
||||
@pytest.mark.parametrize("shape", TEXT_SHAPES, ids=TEXT_SHAPE_IDS)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.LLM_TRANSLATION,
|
||||
route=Route.COUNT_TOKENS,
|
||||
providers=(Provider.BEDROCK,),
|
||||
models=(BEDROCK_KIMI,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_text_estimate_matches_billed_input(self, bedrock_kimi: Served, sdk: SdkClients, shape: PromptShape) -> None:
|
||||
counted, billed = _count_and_bill(sdk.anthropic(bedrock_kimi.key), bedrock_kimi.model, shape)
|
||||
|
||||
assert counted.litellm_count_source == "local_estimate", counted
|
||||
assert _within_tolerance(counted.input_tokens, billed), (counted, billed)
|
||||
|
||||
|
||||
class TestConfiguredCountModelOverridesCostMap:
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.LLM_TRANSLATION,
|
||||
route=Route.COUNT_TOKENS,
|
||||
providers=(Provider.BEDROCK,),
|
||||
models=(BEDROCK_SONNET, CONFIGURED_COUNT_MODEL),
|
||||
)
|
||||
)
|
||||
def test_count_names_the_model_info_count_model(
|
||||
self, bedrock_sonnet_with_configured_count_model: Served, sdk: SdkClients
|
||||
) -> None:
|
||||
served: Final = bedrock_sonnet_with_configured_count_model
|
||||
raw: Final = sdk.anthropic(served.key).messages.with_raw_response.count_tokens(
|
||||
model=served.model, messages=[_user_turn(SHAPES[0])]
|
||||
)
|
||||
counted: Final = CountTokensResponse.model_validate_json(raw.text)
|
||||
|
||||
assert counted.litellm_count_source == "token_count_model", counted
|
||||
assert counted.litellm_count_model == CONFIGURED_COUNT_MODEL, counted
|
||||
assert counted.input_tokens > 0, counted
|
||||
|
||||
|
||||
class TestProviderCountErrorWithoutCountModel:
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.LLM_TRANSLATION,
|
||||
route=Route.COUNT_TOKENS,
|
||||
providers=(Provider.BEDROCK,),
|
||||
models=(BEDROCK_NOVA,),
|
||||
)
|
||||
)
|
||||
def test_count_keeps_the_provider_400_instead_of_estimating(self, bedrock_nova: Served, sdk: SdkClients) -> None:
|
||||
with pytest.raises(BadRequestError) as rejected:
|
||||
_ = sdk.anthropic(bedrock_nova.key).messages.count_tokens(
|
||||
model=bedrock_nova.model, messages=[_user_turn(SHAPES[0])]
|
||||
)
|
||||
|
||||
assert rejected.value.status_code == 400, rejected.value
|
||||
|
||||
|
||||
class TestCountSourceOnOtherCountEndpoints:
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.LLM_TRANSLATION,
|
||||
route=Route.COUNT_TOKENS,
|
||||
providers=(Provider.OPENAI,),
|
||||
models=(OPENAI_GPT,),
|
||||
)
|
||||
)
|
||||
def test_responses_input_tokens_says_the_provider_counted(self, openai_gpt: Served, sdk: SdkClients) -> None:
|
||||
raw: Final = sdk.openai(openai_gpt.key).responses.input_tokens.with_raw_response.count(
|
||||
model=openai_gpt.model, input=f"Reply with one word. {unique_marker()}"
|
||||
)
|
||||
counted: Final = ResponsesInputTokensResponse.model_validate_json(raw.text)
|
||||
|
||||
assert counted.litellm_count_source == "provider", counted
|
||||
assert counted.input_tokens > 0, counted
|
||||
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.LLM_TRANSLATION,
|
||||
route=Route.COUNT_TOKENS,
|
||||
providers=(Provider.BEDROCK,),
|
||||
models=(BEDROCK_KIMI,),
|
||||
)
|
||||
)
|
||||
def test_token_counter_says_kimi_was_estimated_locally(self, proxy: ProxyClient, bedrock_kimi: Served) -> None:
|
||||
counted: Final = unwrap(
|
||||
proxy.token_counter(
|
||||
bedrock_kimi.key,
|
||||
TokenCounterBody(
|
||||
model=bedrock_kimi.model,
|
||||
messages=[ChatMessage(role="user", content=f"Reply with one word. {unique_marker()}")],
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
assert counted.count_source == "local_estimate", counted
|
||||
assert counted.total_tokens > 0, counted
|
||||
|
|
@ -677,13 +677,39 @@ class AnthropicMessagesResponse(BaseModel):
|
|||
usage: Usage | None = None
|
||||
|
||||
|
||||
class CountTokensResponse(BaseModel):
|
||||
TokenCountSource = Literal["provider", "token_count_model", "local_estimate"]
|
||||
|
||||
|
||||
class LitellmCountSource(BaseModel, frozen=True):
|
||||
litellm_count_source: TokenCountSource | None = None
|
||||
litellm_count_model: str | None = None
|
||||
|
||||
|
||||
class CountTokensResponse(LitellmCountSource, frozen=True):
|
||||
"""`/v1/messages/count_tokens` answer. `input_tokens` is required so a 200
|
||||
whose body lacks it fails validation instead of passing vacuously."""
|
||||
|
||||
input_tokens: int
|
||||
|
||||
|
||||
class ResponsesInputTokensResponse(LitellmCountSource, frozen=True):
|
||||
object: str
|
||||
input_tokens: int
|
||||
|
||||
|
||||
class TokenCounterBody(BaseModel, frozen=True):
|
||||
model: str
|
||||
messages: list[ChatMessage]
|
||||
|
||||
|
||||
class TokenCounterResponse(BaseModel, frozen=True):
|
||||
total_tokens: int
|
||||
request_model: str
|
||||
model_used: str
|
||||
tokenizer_type: str
|
||||
count_source: TokenCountSource | None = None
|
||||
|
||||
|
||||
class AnthropicErrorBody(BaseModel):
|
||||
type: str
|
||||
message: str
|
||||
|
|
@ -1392,6 +1418,11 @@ class LiteLLMParamsBody(BaseModel):
|
|||
ModelMode = Literal["chat", "batch", "realtime", "image_generation"]
|
||||
|
||||
|
||||
class TokenCountModelBody(BaseModel, frozen=True):
|
||||
model: str
|
||||
region: str | None = None
|
||||
|
||||
|
||||
class ModelInfoBody(BaseModel):
|
||||
# id is left unset so the proxy assigns a unique model_id per deployment.
|
||||
# Pinning it to the model_name made re-registrations of a fixed-name model
|
||||
|
|
@ -1404,6 +1435,7 @@ class ModelInfoBody(BaseModel):
|
|||
team_id: str | None = None
|
||||
allowed_fails: int | None = None
|
||||
allowed_fails_policy: dict[str, int] | None = None
|
||||
token_count_model: TokenCountModelBody | None = None
|
||||
|
||||
|
||||
class ModelNewBody(BaseModel):
|
||||
|
|
|
|||
|
|
@ -100,6 +100,8 @@ from models import (
|
|||
TeamNewBody,
|
||||
TeamNewResponse,
|
||||
TeamUpdateBody,
|
||||
TokenCounterBody,
|
||||
TokenCounterResponse,
|
||||
ToolsetCreateBody,
|
||||
ToolsetRow,
|
||||
ToolsetUpdateBody,
|
||||
|
|
@ -1049,6 +1051,15 @@ class ProxyClient:
|
|||
response_type=CountTokensResponse,
|
||||
)
|
||||
|
||||
@step("Count tokens with /utils/token_counter for {body.model}")
|
||||
def token_counter(self, key: str, body: TokenCounterBody) -> Result[TokenCounterResponse]:
|
||||
return self.transport.post(
|
||||
"/utils/token_counter",
|
||||
headers=self.transport.bearer(key),
|
||||
json=body,
|
||||
response_type=TokenCounterResponse,
|
||||
)
|
||||
|
||||
@step("Send a /v1/messages request to {body.model}")
|
||||
def messages(
|
||||
self, key: str, body: AnthropicMessagesBody, *, session_id: str | None = None
|
||||
|
|
|
|||
|
|
@ -104,7 +104,7 @@ def test_anthropic_count_tokens_forwards_body_and_refuses_a_key_without_model_ac
|
|||
|
||||
client: Final = _client(gateway, gateway.key)
|
||||
result: Final = _sdk_count(client, model, spelling, _COUNT_BETA)
|
||||
assert result.model_dump() == {"input_tokens": _COUNTED_INPUT_TOKENS}
|
||||
assert result.model_dump() == {"input_tokens": _COUNTED_INPUT_TOKENS, "litellm_count_source": "provider"}
|
||||
requests: Final = wire.drain()
|
||||
assert len(requests) == 1, requests
|
||||
|
||||
|
|
@ -236,5 +236,5 @@ def test_anthropic_count_tokens_merges_the_client_beta_with_the_token_counting_b
|
|||
num_retries=0,
|
||||
)
|
||||
result: Final = _sdk_count(_client(gateway, gateway.key), model, spelling, _CLIENT_BETA)
|
||||
assert result.model_dump() == {"input_tokens": _COUNTED_INPUT_TOKENS}
|
||||
assert result.model_dump() == {"input_tokens": _COUNTED_INPUT_TOKENS, "litellm_count_source": "provider"}
|
||||
assert len(wire.drain()) == 1
|
||||
|
|
|
|||
|
|
@ -245,7 +245,10 @@ def test_enabled_gateway_count_tokens_alias_matches_native_count_tokens_route(en
|
|||
"/claude_code_gateway/v1/messages/count_tokens?beta=true", json=body, headers=headers
|
||||
)
|
||||
assert gateway_response.status_code == 200, gateway_response.text
|
||||
assert _JSON_OBJECT.validate_json(gateway_response.content) == {"input_tokens": 41}, gateway_response.text
|
||||
assert _JSON_OBJECT.validate_json(gateway_response.content) == {
|
||||
"input_tokens": 41,
|
||||
"litellm_count_source": "provider",
|
||||
}, gateway_response.text
|
||||
native_response: Final = enabled_gateway.client.post(
|
||||
"/v1/messages/count_tokens?beta=true", json=body, headers=headers
|
||||
)
|
||||
|
|
|
|||
|
|
@ -274,7 +274,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),)
|
||||
|
||||
|
||||
|
|
@ -329,7 +329,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,)
|
||||
|
||||
|
||||
|
|
@ -342,7 +346,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}]},
|
||||
)
|
||||
|
|
@ -373,7 +381,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},)
|
||||
|
||||
|
||||
|
|
@ -386,7 +394,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]},
|
||||
)
|
||||
|
|
@ -418,7 +426,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
|
||||
|
||||
|
||||
|
|
@ -432,7 +440,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,)
|
||||
|
||||
|
||||
|
|
@ -447,7 +455,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,)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -378,7 +378,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),)
|
||||
|
||||
|
|
@ -394,7 +394,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,)
|
||||
|
||||
|
|
@ -429,7 +429,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},)
|
||||
|
||||
|
|
@ -445,7 +449,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() == ()
|
||||
|
||||
|
|
@ -460,7 +468,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),)
|
||||
|
||||
|
|
@ -520,7 +528,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(
|
||||
|
|
@ -565,7 +577,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(
|
||||
|
|
@ -595,7 +611,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,)
|
||||
|
||||
|
|
@ -640,7 +660,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"}
|
||||
)
|
||||
|
|
@ -728,7 +751,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"}]),)
|
||||
|
||||
|
|
@ -745,7 +772,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),)
|
||||
|
||||
|
|
@ -858,7 +885,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}]),)
|
||||
|
||||
|
|
@ -874,7 +905,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),)
|
||||
|
||||
|
|
@ -894,7 +925,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]),
|
||||
|
|
@ -947,7 +978,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
|
||||
|
||||
|
|
@ -968,7 +999,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,)
|
||||
|
||||
|
|
@ -990,7 +1023,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,)
|
||||
|
||||
|
|
@ -1015,7 +1048,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
|
||||
|
||||
|
|
@ -1220,7 +1255,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),)
|
||||
|
||||
|
|
@ -1241,7 +1276,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,)
|
||||
|
|
@ -1259,7 +1294,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,)
|
||||
|
||||
|
|
@ -1275,7 +1310,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
|
||||
|
||||
|
|
@ -1304,9 +1339,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
|
||||
|
|
|
|||
719
tests/integration/providers/test_count_tokens_model_wire.py
Normal file
719
tests/integration/providers/test_count_tokens_model_wire.py
Normal file
|
|
@ -0,0 +1,719 @@
|
|||
import json
|
||||
import re
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from itertools import chain
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from uuid import uuid4
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import pytest
|
||||
from anthropic.types import MessageCountTokensToolParam, MessageTokensCount
|
||||
from integration._support.bedrock_runtime_peer import target_of
|
||||
from integration._support.client import Gateway, Scenario, gateway_from_environment
|
||||
from integration._support.ports import allocate_port
|
||||
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.vertex import service_account_json
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
_SONNET: Final = "us.anthropic.claude-sonnet-5-5"
|
||||
_SONNET_BASE: Final = "anthropic.claude-sonnet-5-5"
|
||||
_OPUS: Final = "us.anthropic.claude-opus-5-5"
|
||||
_OPUS_BASE: Final = "anthropic.claude-opus-5-5"
|
||||
_NOVA: Final = "amazon.nova-lite-v1:0"
|
||||
_KIMI: Final = "us.moonshotai.kimi-k3"
|
||||
_KIMI_BASE: Final = "moonshotai.kimi-k3"
|
||||
_PROFILE_ID: Final = "lit9563count"
|
||||
_PROFILE_ARN: Final = f"arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/{_PROFILE_ID}"
|
||||
_COST_MAP_PATH: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json"
|
||||
_REGION: Final = "us-west-2"
|
||||
_COUNT_REGION: Final = "us-east-1"
|
||||
_ACCESS_KEY: Final = "AKIAINTEGRATIONCOUNT"
|
||||
_SECRET_KEY: Final = "integration-count-secret"
|
||||
_PROVIDER_COUNT: Final = 1345
|
||||
_COUNT_MODEL_COUNT: Final = 2718
|
||||
_FRAMING: Final = 50
|
||||
_UNSUPPORTED: Final = "The provided model doesn't support counting tokens."
|
||||
_COUNT_TARGET: Final = re.compile(r"^/model/(.+)/count-tokens$")
|
||||
_SCRIPTED_STATUS: Final = re.compile(r"status=(\d{3})")
|
||||
_SCRIPTED_COUNT_MODEL_STATUS: Final = re.compile(r"count_model=(\d{3})")
|
||||
_CREDENTIAL: Final = re.compile(r"Credential=([^/]+)/\d+/([^/]+)/")
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_COST_MAP: Final = TypeAdapter(dict[str, dict[str, JsonValue]])
|
||||
|
||||
_SYSTEM: Final = "You are a terse assistant that answers in one sentence"
|
||||
_WEATHER_SCHEMA: Final[dict[str, JsonValue]] = {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string", "description": "City to look up"}},
|
||||
"required": ["city"],
|
||||
}
|
||||
_SDK_TOOLS: Final[list[MessageCountTokensToolParam]] = [
|
||||
{
|
||||
"name": "get_weather",
|
||||
"description": "Look up the current weather for a city",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string", "description": "City to look up"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
}
|
||||
]
|
||||
_COUNT_MODEL_SOURCE: Final = {"litellm_count_source": "token_count_model", "litellm_count_model": _OPUS}
|
||||
_OPENAI_PRIMARY: Final = "gpt-count-primary"
|
||||
_OPENAI_COUNT_MODEL: Final = "gpt-count-model"
|
||||
_OPENAI_COUNT_TARGET: Final = "/v1/responses/input_tokens"
|
||||
_COUNTED_FIELDS: Final = ("input", "instructions", "tools", "text")
|
||||
_ANSWER_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}}
|
||||
_ANTHROPIC_TOOL: Final[dict[str, JsonValue]] = {
|
||||
"name": "get_weather",
|
||||
"description": "Look up the current weather for a city",
|
||||
"input_schema": _WEATHER_SCHEMA,
|
||||
}
|
||||
_TOOL_TURNS: Final[list[JsonValue]] = [
|
||||
{"role": "user", "content": "What is the weather in Paris?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "tool_use", "id": "toolu_01", "name": "get_weather", "input": {"city": "Paris"}}],
|
||||
},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "Sunny"}]},
|
||||
]
|
||||
_LOCAL_SOURCE: Final = {"litellm_count_source": "local_estimate"}
|
||||
|
||||
|
||||
def _prompt(status: int, count_model_status: int | None = None) -> str:
|
||||
scripted: Final = "" if count_model_status is None else f" count_model={count_model_status}"
|
||||
return f"Count this message status={status}{scripted}"
|
||||
|
||||
|
||||
def _converse_body(prompt: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"input": {
|
||||
"converse": {
|
||||
"messages": [{"role": "user", "content": [{"text": prompt}]}],
|
||||
"system": [{"text": _SYSTEM}],
|
||||
"toolConfig": {
|
||||
"tools": [
|
||||
{
|
||||
"toolSpec": {
|
||||
"name": "get_weather",
|
||||
"description": "Look up the current weather for a city",
|
||||
"inputSchema": {"json": _WEATHER_SCHEMA},
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _json_reply(status: int, payload: Mapping[str, JsonValue]) -> Reply:
|
||||
return Reply(status=status, body=json.dumps(payload).encode())
|
||||
|
||||
|
||||
def _counted_model(request: Request) -> str:
|
||||
counted: Final = _COUNT_TARGET.match(target_of(request))
|
||||
return target_of(request) if counted is None else counted.group(1)
|
||||
|
||||
|
||||
def _scripted(pattern: re.Pattern[str], request: Request, default: int) -> int:
|
||||
scripted: Final = pattern.search(request.body.decode(errors="replace"))
|
||||
return default if scripted is None else int(scripted.group(1))
|
||||
|
||||
|
||||
def _runtime(request: Request) -> Reply:
|
||||
model: Final = _counted_model(request)
|
||||
if model == _OPUS_BASE:
|
||||
count_model_status: Final = _scripted(_SCRIPTED_COUNT_MODEL_STATUS, request, 200)
|
||||
if count_model_status == 200:
|
||||
return _json_reply(200, {"inputTokens": _COUNT_MODEL_COUNT})
|
||||
return _json_reply(count_model_status, {"message": f"scripted runtime {count_model_status} for {model}"})
|
||||
status: Final = _scripted(_SCRIPTED_STATUS, request, 500)
|
||||
if status == 200:
|
||||
return _json_reply(200, {"inputTokens": _PROVIDER_COUNT})
|
||||
return _json_reply(status, {"message": _UNSUPPORTED if status == 400 else f"scripted runtime {status}"})
|
||||
|
||||
|
||||
def _mantle(request: Request) -> Reply:
|
||||
model: Final = _JSON_OBJECT.validate_json(request.body).get("model")
|
||||
return _json_reply(
|
||||
404, {"type": "error", "error": {"type": "not_found_error", "message": f"scripted mantle 404 for {model}"}}
|
||||
)
|
||||
|
||||
|
||||
def _signer(request: Request) -> tuple[str, str]:
|
||||
credential: Final = _CREDENTIAL.search(request.headers.get("authorization", ""))
|
||||
return ("", "") if credential is None else (credential.group(1), credential.group(2))
|
||||
|
||||
|
||||
def _runtime_counts(wire: Wire) -> tuple[tuple[str, tuple[str, str], dict[str, JsonValue]], ...]:
|
||||
return tuple(
|
||||
(_counted_model(request), _signer(request), _JSON_OBJECT.validate_json(request.body))
|
||||
for request in wire.drain()
|
||||
)
|
||||
|
||||
|
||||
def _runtime_models(wire: Wire) -> tuple[str, ...]:
|
||||
return tuple(_counted_model(request) for request in wire.drain())
|
||||
|
||||
|
||||
def _mantle_models(wire: Wire) -> tuple[JsonValue, ...]:
|
||||
return tuple(_JSON_OBJECT.validate_json(request.body).get("model") for request in wire.drain())
|
||||
|
||||
|
||||
def _url(port: int) -> str:
|
||||
return f"http://127.0.0.1:{port}"
|
||||
|
||||
|
||||
def _litellm_params(model: str, api_base: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model": f"bedrock/{model}",
|
||||
"api_base": api_base,
|
||||
"aws_access_key_id": _ACCESS_KEY,
|
||||
"aws_secret_access_key": _SECRET_KEY,
|
||||
"aws_region_name": _REGION,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def runtime_port() -> int:
|
||||
return allocate_port()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def mantle_port() -> int:
|
||||
return allocate_port()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def counting_proxy(tmp_path_factory: pytest.TempPathFactory, mantle_port: int) -> Iterator[OwnedProxy]:
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path_factory.mktemp("count-model"),
|
||||
{
|
||||
"BEDROCK_MANTLE_API_BASE": _url(mantle_port),
|
||||
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
|
||||
"HF_HUB_OFFLINE": "1",
|
||||
},
|
||||
workers=1,
|
||||
remove_environment=("AWS_BEARER_TOKEN_BEDROCK",),
|
||||
) as owned,
|
||||
):
|
||||
yield owned
|
||||
|
||||
|
||||
def _deployment(
|
||||
scenario: Scenario, runtime_port: int, model: str = _SONNET, model_info: Mapping[str, JsonValue] | None = None
|
||||
) -> str:
|
||||
return scenario.model(model_info=model_info, api_key=None, **_litellm_params(model, _url(runtime_port)))
|
||||
|
||||
|
||||
def _count(gateway: Gateway, model: str, status: int, count_model_status: int | None = None) -> MessageTokensCount:
|
||||
with anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) as client:
|
||||
return client.messages.count_tokens(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": _prompt(status, count_model_status)}],
|
||||
system=_SYSTEM,
|
||||
tools=_SDK_TOOLS,
|
||||
)
|
||||
|
||||
|
||||
def _messages_estimate(gateway: Gateway, model: str) -> tuple[int, JsonValue]:
|
||||
counted: Final = _count(gateway, model, 400)
|
||||
return counted.input_tokens, (counted.model_extra or {}).get("litellm_count_source")
|
||||
|
||||
|
||||
def _utils_estimate(gateway: Gateway, model: str) -> tuple[int, JsonValue]:
|
||||
response: Final = gateway.request(
|
||||
"POST", "/utils/token_counter", {"model": model, "messages": [{"role": "user", "content": _prompt(400)}]}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = _JSON_OBJECT.validate_json(response.content)
|
||||
total: Final = body["total_tokens"]
|
||||
assert isinstance(total, int), body
|
||||
return total, body.get("count_source")
|
||||
|
||||
|
||||
def _shipped_framing(model: str) -> JsonValue:
|
||||
return _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes())[model].get("token_count_framing")
|
||||
|
||||
|
||||
def test_shipped_token_count_model_counts_the_request_in_its_region_without_trying_the_model(
|
||||
counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with (
|
||||
wire_server(_runtime, port=runtime_port) as runtime,
|
||||
wire_server(_mantle, port=mantle_port) as mantle,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
counted: Final = _count(gateway, _deployment(scenario, runtime_port), 404)
|
||||
assert (counted.input_tokens, counted.model_extra) == (_COUNT_MODEL_COUNT, _COUNT_MODEL_SOURCE), counted
|
||||
assert _runtime_counts(runtime) == ((_OPUS_BASE, (_ACCESS_KEY, _COUNT_REGION), _converse_body(_prompt(404))),)
|
||||
assert _mantle_models(mantle) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("count_model_status", "error", "message"),
|
||||
[
|
||||
(401, anthropic.AuthenticationError, f"scripted runtime 401 for {_OPUS_BASE}"),
|
||||
(429, anthropic.RateLimitError, f"scripted runtime 429 for {_OPUS_BASE}"),
|
||||
(404, anthropic.NotFoundError, f"scripted mantle 404 for {_OPUS_BASE}"),
|
||||
],
|
||||
)
|
||||
def test_token_count_model_failure_answers_its_error_and_names_the_region_to_enable(
|
||||
counting_proxy: OwnedProxy,
|
||||
runtime_port: int,
|
||||
mantle_port: int,
|
||||
count_model_status: int,
|
||||
error: type[anthropic.APIStatusError],
|
||||
message: str,
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with (
|
||||
wire_server(_runtime, port=runtime_port) as runtime,
|
||||
wire_server(_mantle, port=mantle_port) as mantle,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
with pytest.raises(error) as refused:
|
||||
_count(gateway, _deployment(scenario, runtime_port), 200, count_model_status=count_model_status)
|
||||
assert message in str(refused.value) and f"{_OPUS} in {_COUNT_REGION}" in str(refused.value), refused.value
|
||||
assert _runtime_models(runtime) == (_OPUS_BASE,)
|
||||
assert _mantle_models(mantle) == ((_OPUS_BASE,) if count_model_status == 404 else ())
|
||||
|
||||
|
||||
def test_framing_is_added_to_the_local_estimate_when_the_runtime_cannot_count(
|
||||
counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with (
|
||||
wire_server(_runtime, port=runtime_port) as runtime,
|
||||
wire_server(_mantle, port=mantle_port) as mantle,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
framed: Final = _count(
|
||||
gateway, _deployment(scenario, runtime_port, _NOVA, {"token_count_framing": _FRAMING}), 400
|
||||
)
|
||||
unframed: Final = _count(gateway, _deployment(scenario, runtime_port, _NOVA, {"token_count_framing": 0}), 400)
|
||||
assert (framed.model_extra, unframed.model_extra) == (_LOCAL_SOURCE, _LOCAL_SOURCE), (framed, unframed)
|
||||
assert framed.input_tokens - unframed.input_tokens == _FRAMING, (framed, unframed)
|
||||
assert _runtime_models(runtime) == (_NOVA, _NOVA)
|
||||
assert _mantle_models(mantle) == ()
|
||||
|
||||
|
||||
def test_framing_on_one_deployment_leaves_a_sibling_of_the_same_model_raising_the_400(
|
||||
counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with (
|
||||
wire_server(_runtime, port=runtime_port) as runtime,
|
||||
wire_server(_mantle, port=mantle_port) as mantle,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
framed: Final = _deployment(scenario, runtime_port, _NOVA, {"token_count_framing": _FRAMING})
|
||||
sibling: Final = _deployment(scenario, runtime_port, _NOVA)
|
||||
assert _count(gateway, framed, 400).model_extra == _LOCAL_SOURCE
|
||||
with pytest.raises(anthropic.BadRequestError, match=re.escape(_UNSUPPORTED)):
|
||||
_count(gateway, sibling, 400)
|
||||
assert _runtime_models(runtime) == (_NOVA, _NOVA)
|
||||
assert _mantle_models(mantle) == ()
|
||||
|
||||
|
||||
def test_empty_configured_token_count_model_counts_the_model_itself_in_the_deployment_region(
|
||||
counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with (
|
||||
wire_server(_runtime, port=runtime_port) as runtime,
|
||||
wire_server(_mantle, port=mantle_port) as mantle,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
counted: Final = _count(gateway, _deployment(scenario, runtime_port, model_info={"token_count_model": {}}), 200)
|
||||
assert (counted.input_tokens, counted.model_extra) == (_PROVIDER_COUNT, {"litellm_count_source": "provider"})
|
||||
assert _runtime_counts(runtime) == ((_SONNET_BASE, (_ACCESS_KEY, _REGION), _converse_body(_prompt(200))),)
|
||||
assert _mantle_models(mantle) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("base_model_in", ["model_info", "litellm_params"])
|
||||
def test_inference_profile_arn_counts_with_the_token_count_model_of_its_base_model(
|
||||
counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int, base_model_in: str
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
base_model: Final[dict[str, JsonValue]] = {"base_model": f"bedrock/{_SONNET}"}
|
||||
with (
|
||||
wire_server(_runtime, port=runtime_port) as runtime,
|
||||
wire_server(_mantle, port=mantle_port) as mantle,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(
|
||||
model_info=base_model if base_model_in == "model_info" else None,
|
||||
api_key=None,
|
||||
**_litellm_params(_PROFILE_ARN, _url(runtime_port)),
|
||||
**(base_model if base_model_in == "litellm_params" else {}),
|
||||
)
|
||||
counted: Final = _count(gateway, model, 404)
|
||||
assert (counted.input_tokens, counted.model_extra) == (_COUNT_MODEL_COUNT, _COUNT_MODEL_SOURCE), counted
|
||||
assert _runtime_models(runtime) == (_OPUS_BASE,)
|
||||
assert _mantle_models(mantle) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("estimate", "runtime_calls"),
|
||||
[(_messages_estimate, (_KIMI_BASE, _KIMI_BASE)), (_utils_estimate, ())],
|
||||
ids=["messages_count_tokens", "utils_token_counter"],
|
||||
)
|
||||
def test_shipped_framing_is_added_to_the_estimate_and_a_configured_zero_turns_it_off(
|
||||
counting_proxy: OwnedProxy,
|
||||
runtime_port: int,
|
||||
mantle_port: int,
|
||||
estimate: Callable[[Gateway, str], tuple[int, JsonValue]],
|
||||
runtime_calls: tuple[str, ...],
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with (
|
||||
wire_server(_runtime, port=runtime_port) as runtime,
|
||||
wire_server(_mantle, port=mantle_port) as mantle,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
shipped: Final = estimate(gateway, _deployment(scenario, runtime_port, _KIMI))
|
||||
unframed: Final = estimate(gateway, _deployment(scenario, runtime_port, _KIMI, {"token_count_framing": 0}))
|
||||
assert (shipped[1], unframed[1]) == ("local_estimate", "local_estimate"), (shipped, unframed)
|
||||
framing: Final = _shipped_framing(_KIMI)
|
||||
assert isinstance(framing, int) and framing > 0, framing
|
||||
assert shipped[0] - unframed[0] == framing, (shipped, unframed)
|
||||
assert _runtime_models(runtime) == runtime_calls
|
||||
assert _mantle_models(mantle) == ()
|
||||
|
||||
|
||||
def _openai_count(request: Request) -> Reply:
|
||||
model: Final = _JSON_OBJECT.validate_json(request.body).get("model")
|
||||
if request.target == _OPENAI_COUNT_TARGET and model == _OPENAI_COUNT_MODEL:
|
||||
return _json_reply(200, {"object": "response.input_tokens", "input_tokens": _COUNT_MODEL_COUNT})
|
||||
return _json_reply(
|
||||
404,
|
||||
{"error": {"message": f"scripted openai 404 for {model}", "type": "invalid_request_error", "code": None}},
|
||||
)
|
||||
|
||||
|
||||
def _counted_fields(body: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
return {key: body[key] for key in _COUNTED_FIELDS if key in body}
|
||||
|
||||
|
||||
def _openai_counts(wire: Wire) -> tuple[tuple[str, JsonValue, dict[str, JsonValue]], ...]:
|
||||
bodies: Final = tuple((request.target, _JSON_OBJECT.validate_json(request.body)) for request in wire.drain())
|
||||
return tuple((target, body.get("model"), _counted_fields(body)) for target, body in bodies)
|
||||
|
||||
|
||||
def _openai_deployment(scenario: Scenario, wire: Wire, model_info: Mapping[str, JsonValue]) -> str:
|
||||
return scenario.model(model=f"openai/{_OPENAI_PRIMARY}", api_base=f"{wire.url}/v1", model_info=model_info)
|
||||
|
||||
|
||||
def test_token_count_model_counts_the_anthropic_body_converted_like_the_provider_count(
|
||||
counting_proxy: OwnedProxy,
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with wire_server(_openai_count) as openai, gateway.scenario() as scenario:
|
||||
model: Final = _openai_deployment(scenario, openai, {"token_count_model": {"model": _OPENAI_COUNT_MODEL}})
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/messages/count_tokens",
|
||||
{
|
||||
"model": model,
|
||||
"system": _SYSTEM,
|
||||
"tools": [_ANTHROPIC_TOOL],
|
||||
"output_config": {"format": {"type": "json_schema", "schema": _ANSWER_SCHEMA}},
|
||||
"messages": _TOOL_TURNS,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _JSON_OBJECT.validate_json(response.content) == {
|
||||
"input_tokens": _COUNT_MODEL_COUNT,
|
||||
"litellm_count_source": "token_count_model",
|
||||
"litellm_count_model": _OPENAI_COUNT_MODEL,
|
||||
}
|
||||
counts: Final = _openai_counts(openai)
|
||||
assert tuple((target, counted) for target, counted, _ in counts) == (
|
||||
(_OPENAI_COUNT_TARGET, _OPENAI_COUNT_MODEL),
|
||||
), counts
|
||||
primary: Final = counts[0][2]
|
||||
assert primary["instructions"] == _SYSTEM, primary
|
||||
text: Final = primary.get("text")
|
||||
assert isinstance(text, dict) and isinstance(text.get("format"), dict), primary
|
||||
input_items: Final = primary["input"]
|
||||
assert isinstance(input_items, list), primary
|
||||
item_types: Final = {item.get("type") for item in input_items if isinstance(item, dict)}
|
||||
assert {"function_call", "function_call_output"} <= item_types, input_items
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("body", "named"),
|
||||
[
|
||||
(
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Summarize the notes."},
|
||||
{
|
||||
"type": "search_result",
|
||||
"source": "https://example.com/notes",
|
||||
"title": "Notes",
|
||||
"content": [],
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
},
|
||||
"search_result",
|
||||
),
|
||||
(
|
||||
{"messages": [{"role": "user", "content": "hello"}, {"content": "no role"}]},
|
||||
"messages.1.role: Field required",
|
||||
),
|
||||
],
|
||||
ids=["unconvertible_block", "missing_role"],
|
||||
)
|
||||
def test_framing_keeps_the_400_for_a_body_the_openai_count_rejects_before_calling_openai(
|
||||
counting_proxy: OwnedProxy, body: Mapping[str, JsonValue], named: str
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with wire_server(_openai_count) as openai, gateway.scenario() as scenario:
|
||||
model: Final = _openai_deployment(scenario, openai, {"token_count_framing": _FRAMING})
|
||||
rejected: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model, **body})
|
||||
assert rejected.status_code == 400, rejected.text
|
||||
error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"]
|
||||
assert isinstance(error, dict) and named in str(error.get("message")), error
|
||||
assert openai.drain() == ()
|
||||
estimated: Final = gateway.request(
|
||||
"POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": "hello"}]}
|
||||
)
|
||||
assert estimated.status_code == 200, estimated.text
|
||||
assert _JSON_OBJECT.validate_json(estimated.content).get("litellm_count_source") == "local_estimate"
|
||||
assert tuple(counted for _, counted, _ in _openai_counts(openai)) == (_OPENAI_PRIMARY,)
|
||||
|
||||
|
||||
def _runtime_cannot_count(_request: Request) -> Reply:
|
||||
return _json_reply(400, {"message": _UNSUPPORTED})
|
||||
|
||||
|
||||
_REPLAYED_WEB_SEARCH: Final[list[JsonValue]] = [
|
||||
{"role": "user", "content": "What is the newest LiteLLM release?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "LiteLLM"}},
|
||||
{
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_01",
|
||||
"content": [
|
||||
{
|
||||
"type": "web_search_result",
|
||||
"url": "https://docs.litellm.ai/release_notes",
|
||||
"title": "Release notes",
|
||||
"encrypted_content": "",
|
||||
}
|
||||
],
|
||||
},
|
||||
{"type": "text", "text": "The newest release is in the release notes."},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Thanks."},
|
||||
]
|
||||
_BOOLEAN_SUBSCHEMA_TOOL: Final[dict[str, JsonValue]] = {
|
||||
"name": "lookup",
|
||||
"description": "Look up a value",
|
||||
"input_schema": {"type": "object", "properties": {"query": True}},
|
||||
}
|
||||
_UNREADABLE_BODIES: Final[dict[str, dict[str, JsonValue]]] = {
|
||||
"web_search_replay": {"messages": _REPLAYED_WEB_SEARCH},
|
||||
"boolean_subschema_tool": {
|
||||
"messages": [{"role": "user", "content": "Look up the value."}],
|
||||
"tools": [_BOOLEAN_SUBSCHEMA_TOOL],
|
||||
},
|
||||
}
|
||||
_COUNT_ENDPOINTS: Final[dict[str, tuple[str, str]]] = {
|
||||
"messages_count_tokens": ("/v1/messages/count_tokens", "litellm_count_source"),
|
||||
"utils_token_counter": ("/utils/token_counter?call_endpoint=true", "count_source"),
|
||||
}
|
||||
|
||||
|
||||
def _count_answer(response: httpx.Response, source_key: str) -> tuple[int, str]:
|
||||
if response.status_code not in (200, 400):
|
||||
return response.status_code, response.text
|
||||
body: Final = _JSON_OBJECT.validate_json(response.content)
|
||||
error: Final = body.get("error")
|
||||
if response.status_code == 400 and isinstance(error, dict):
|
||||
return 400, str(error.get("message"))
|
||||
return response.status_code, str(body.get(source_key))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endpoint", sorted(_COUNT_ENDPOINTS))
|
||||
@pytest.mark.parametrize("unreadable", sorted(_UNREADABLE_BODIES))
|
||||
def test_shipped_framing_never_answers_a_500_for_a_body_the_local_count_cannot_read(
|
||||
counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int, unreadable: str, endpoint: str
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
path, source_key = _COUNT_ENDPOINTS[endpoint]
|
||||
with (
|
||||
wire_server(_runtime_cannot_count, port=runtime_port) as runtime,
|
||||
wire_server(_mantle, port=mantle_port) as mantle,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
response: Final = gateway.request(
|
||||
"POST", path, {"model": _deployment(scenario, runtime_port, _KIMI), **_UNREADABLE_BODIES[unreadable]}
|
||||
)
|
||||
status, answer = _count_answer(response, source_key)
|
||||
assert (status, answer) == (200, "local_estimate") or (status == 400 and _UNSUPPORTED in answer), (
|
||||
response.status_code,
|
||||
response.text,
|
||||
)
|
||||
assert _runtime_models(runtime) == (_KIMI_BASE,)
|
||||
assert _mantle_models(mantle) == ()
|
||||
|
||||
|
||||
_GEMINI_MODEL: Final = "gemini-4-argon"
|
||||
_GEMINI_COUNT_MODEL: Final = "gemini-3.8-flash"
|
||||
_VERTEX_PROJECT: Final = "scripted-project"
|
||||
_COMPACT_UNDER_THE_MINIMUM_TRIGGER: Final[dict[str, JsonValue]] = {
|
||||
"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}]
|
||||
}
|
||||
_CLEAR_OLDER_TOOL_RESULTS: Final[dict[str, JsonValue]] = {
|
||||
"edits": [
|
||||
{
|
||||
"type": "clear_tool_uses_20250919",
|
||||
"trigger": {"type": "input_tokens", "value": 500},
|
||||
"keep": {"type": "tool_uses", "value": 1},
|
||||
}
|
||||
]
|
||||
}
|
||||
_READ_FILE_TOOL: Final[dict[str, JsonValue]] = {
|
||||
"name": "read_file",
|
||||
"description": "Read a file",
|
||||
"input_schema": {"type": "object", "properties": {"path": {"type": "string"}}},
|
||||
}
|
||||
_FILE_TEXT: Final = "line of a long file that the agent read earlier. " * 180
|
||||
_FILES: Final = 3
|
||||
|
||||
|
||||
def _google_cannot_count(_request: Request) -> Reply:
|
||||
return _json_reply(400, {"error": {"code": 400, "message": "scripted count 400", "status": "INVALID_ARGUMENT"}})
|
||||
|
||||
|
||||
def _google_counts_only_the_count_model(request: Request) -> Reply:
|
||||
if request.target.endswith(f"/{_GEMINI_COUNT_MODEL}:countTokens"):
|
||||
return _json_reply(200, {"totalTokens": _COUNT_MODEL_COUNT})
|
||||
return _google_cannot_count(request)
|
||||
|
||||
|
||||
def _google_deployment(
|
||||
gateway: Gateway, scenario: Scenario, wire: Wire, provider: str, model_info: dict[str, JsonValue]
|
||||
) -> str:
|
||||
if provider == "gemini":
|
||||
return scenario.model(
|
||||
model=f"gemini/{_GEMINI_MODEL}", api_base=wire.url, api_key="scripted-gemini-key", model_info=model_info
|
||||
)
|
||||
return scenario.model(
|
||||
model=f"vertex_ai/{_GEMINI_MODEL}",
|
||||
api_base=wire.url,
|
||||
api_key=None,
|
||||
vertex_project=_VERTEX_PROJECT,
|
||||
vertex_location="us-east5",
|
||||
vertex_credentials=service_account_json(_VERTEX_PROJECT, gateway.upstream_url.rstrip("/")),
|
||||
model_info=model_info,
|
||||
)
|
||||
|
||||
|
||||
def _google_count_targets(wire: Wire) -> tuple[bool, ...]:
|
||||
return tuple(request.target.endswith(":countTokens") for request in wire.drain())
|
||||
|
||||
|
||||
def _read_file_round(marker: str, index: int) -> tuple[JsonValue, ...]:
|
||||
return (
|
||||
{"role": "user", "content": f"Read file {index}."},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}.txt"}}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": f"toolu_{index}",
|
||||
"content": f"<file {marker}-{index}>{_FILE_TEXT}",
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _agent_messages(marker: str) -> list[JsonValue]:
|
||||
rounds: Final = tuple(chain.from_iterable(_read_file_round(marker, index) for index in range(_FILES)))
|
||||
return [*rounds, {"role": "user", "content": "Summarize what you read."}]
|
||||
|
||||
|
||||
def _counted_files(request: Request, marker: str) -> tuple[bool, tuple[bool, ...]]:
|
||||
sent: Final = json.dumps(_JSON_OBJECT.validate_json(request.body))
|
||||
return (
|
||||
request.target.endswith(f"/{_GEMINI_COUNT_MODEL}:countTokens"),
|
||||
tuple(f"<file {marker}-{index}>" in sent for index in range(_FILES)),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ["gemini", "vertex_ai"])
|
||||
def test_framing_keeps_the_400_for_a_context_edit_the_gemini_count_rejects_before_calling_google(
|
||||
counting_proxy: OwnedProxy, provider: str
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
hello: Final[list[JsonValue]] = [{"role": "user", "content": "hello"}]
|
||||
with wire_server(_google_cannot_count) as google, gateway.scenario() as scenario:
|
||||
model: Final = _google_deployment(gateway, scenario, google, provider, {"token_count_framing": _FRAMING})
|
||||
rejected: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/messages/count_tokens",
|
||||
{"model": model, "messages": hello, "context_management": _COMPACT_UNDER_THE_MINIMUM_TRIGGER},
|
||||
)
|
||||
assert rejected.status_code == 400, rejected.text
|
||||
assert _google_count_targets(google) == ()
|
||||
estimated: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model, "messages": hello})
|
||||
assert estimated.status_code == 200, estimated.text
|
||||
assert _JSON_OBJECT.validate_json(estimated.content).get("litellm_count_source") == "local_estimate"
|
||||
assert _google_count_targets(google) == (True,)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ["gemini", "vertex_ai"])
|
||||
def test_token_count_model_counts_the_body_left_after_the_context_management_edits(
|
||||
counting_proxy: OwnedProxy, provider: str
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
marker: Final = uuid4().hex
|
||||
with wire_server(_google_counts_only_the_count_model) as google, gateway.scenario() as scenario:
|
||||
model: Final = _google_deployment(
|
||||
gateway, scenario, google, provider, {"token_count_model": {"model": _GEMINI_COUNT_MODEL}}
|
||||
)
|
||||
counted: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/messages/count_tokens",
|
||||
{
|
||||
"model": model,
|
||||
"tools": [_READ_FILE_TOOL],
|
||||
"messages": _agent_messages(marker),
|
||||
"context_management": _CLEAR_OLDER_TOOL_RESULTS,
|
||||
},
|
||||
)
|
||||
assert counted.status_code == 200, counted.text
|
||||
assert _JSON_OBJECT.validate_json(counted.content) == {
|
||||
"input_tokens": _COUNT_MODEL_COUNT,
|
||||
"litellm_count_source": "token_count_model",
|
||||
"litellm_count_model": _GEMINI_COUNT_MODEL,
|
||||
}
|
||||
assert tuple(_counted_files(request, marker) for request in google.drain()) == ((True, (False, False, True)),)
|
||||
|
|
@ -46,6 +46,7 @@ _CLAUDE_BACKEND: Final = "claude-haiku-4-5"
|
|||
_GEMINI_BACKEND: Final = "gemini-2.5-flash"
|
||||
_PROVIDER_KEY: Final = "scripted-malformed-tools-key"
|
||||
_COUNT: Final = 37
|
||||
_PROVIDER_SOURCE: Final[Mapping[str, JsonValue]] = MappingProxyType({"litellm_count_source": "provider"})
|
||||
_CONFIG_STORE: Final = "vs_integration_config_store"
|
||||
_OTHER_STORE: Final = "vs_malformed_tools_other_store"
|
||||
_CLAUDE_CODE_AGENT: Final = "claude-cli/2.1.0 (external, cli)"
|
||||
|
|
@ -821,7 +822,7 @@ def test_count_tokens_forwards_well_formed_tools(rig: _Rig) -> None:
|
|||
"POST", "/v1/messages/count_tokens", _count_body(rig, marker, tools=[MESSAGES.tool])
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT}, response.text
|
||||
assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT, **_PROVIDER_SOURCE}, response.text
|
||||
assert _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] == [MESSAGES.tool]
|
||||
|
||||
|
||||
|
|
@ -853,6 +854,7 @@ def test_responses_input_tokens_counts_well_formed_tools(rig: _Rig) -> None:
|
|||
assert _JSON_OBJECT.validate_json(response.content) == {
|
||||
"object": "response.input_tokens",
|
||||
"input_tokens": _COUNT,
|
||||
**_PROVIDER_SOURCE,
|
||||
}, response.text
|
||||
sent: Final = _one_delivery(rig, marker, "/v1/messages/count_tokens")
|
||||
assert "get_weather" in json.dumps(sent["tools"]), sent
|
||||
|
|
@ -1396,7 +1398,7 @@ def test_count_tokens_drops_the_bad_entry_a_deployment_list_names(rig: _Rig, fie
|
|||
{**_count_body(rig, marker, tools=[MESSAGES.tool, 5]), "model": field_groups.anthropic},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT}, response.text
|
||||
assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT, **_PROVIDER_SOURCE}, response.text
|
||||
assert _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] == [MESSAGES.tool]
|
||||
|
||||
|
||||
|
|
@ -1656,7 +1658,7 @@ _CHECKED: Final[Mapping[str, _Checked]] = MappingProxyType(
|
|||
named_as=MESSAGES,
|
||||
send=_count_tokens,
|
||||
target="/v1/messages/count_tokens",
|
||||
answered=_counted({"input_tokens": _COUNT}),
|
||||
answered=_counted({"input_tokens": _COUNT, **_PROVIDER_SOURCE}),
|
||||
kept_by_a_tools_list=("get_weather",),
|
||||
bad_entry=_MESSAGES_BAD_ENTRY,
|
||||
),
|
||||
|
|
@ -1666,7 +1668,7 @@ _CHECKED: Final[Mapping[str, _Checked]] = MappingProxyType(
|
|||
named_as=MESSAGES,
|
||||
send=_input_tokens,
|
||||
target="/v1/messages/count_tokens",
|
||||
answered=_counted({"object": "response.input_tokens", "input_tokens": _COUNT}),
|
||||
answered=_counted({"object": "response.input_tokens", "input_tokens": _COUNT, **_PROVIDER_SOURCE}),
|
||||
kept_by_a_tools_list=("get_weather",),
|
||||
bad_entry=_RESPONSES_BAD_ENTRY,
|
||||
),
|
||||
|
|
|
|||
|
|
@ -1118,6 +1118,7 @@ def test_input_tokens_counts_marked_items_without_their_marks(gateway: Gateway)
|
|||
assert _JSON_OBJECT.validate_json(response.text) == {
|
||||
"object": "response.input_tokens",
|
||||
"input_tokens": _COUNTED_TOKENS,
|
||||
"litellm_count_source": "provider",
|
||||
}
|
||||
assert counted.target.endswith("/count_tokens"), counted.target
|
||||
assert anthropic_labels(counted) == []
|
||||
|
|
@ -1199,6 +1200,7 @@ def test_input_tokens_reads_each_input_shape_without_item_marks(gateway: Gateway
|
|||
assert _JSON_OBJECT.validate_json(response.text) == {
|
||||
"object": "response.input_tokens",
|
||||
"input_tokens": _COUNTED_TOKENS,
|
||||
"litellm_count_source": "provider",
|
||||
}
|
||||
assert [anthropic_labels(request) for request in counted] == [[]]
|
||||
|
||||
|
|
|
|||
|
|
@ -396,7 +396,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},)
|
||||
|
||||
|
||||
|
|
@ -404,7 +404,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,)
|
||||
|
||||
|
||||
|
|
@ -443,7 +443,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,)
|
||||
|
||||
|
||||
|
|
@ -498,7 +502,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
|
||||
|
||||
|
|
@ -616,7 +620,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": []},)
|
||||
|
||||
|
||||
|
|
@ -624,7 +628,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": ""},)
|
||||
|
||||
|
||||
|
|
@ -632,7 +636,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,)
|
||||
|
||||
|
||||
|
|
@ -650,7 +654,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},)
|
||||
|
||||
|
||||
|
|
@ -663,7 +667,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
|
||||
|
|
@ -696,7 +700,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,)
|
||||
|
||||
|
|
@ -705,7 +709,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
|
||||
|
||||
|
||||
|
|
@ -722,7 +726,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}]}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -453,9 +453,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")
|
||||
|
|
|
|||
|
|
@ -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}",
|
||||
|
|
|
|||
209
tests/unit/proxy/common_utils/test_token_count_settings.py
Normal file
209
tests/unit/proxy/common_utils/test_token_count_settings.py
Normal file
|
|
@ -0,0 +1,209 @@
|
|||
import json
|
||||
from pathlib import Path
|
||||
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"},
|
||||
}
|
||||
_COST_MAP: Final = Path(__file__).parents[4] / "model_prices_and_context_window.json"
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
def test_every_shipped_token_count_model_is_another_mapped_model_of_the_same_provider_in_a_named_region() -> None:
|
||||
cost_map: Final = json.loads(_COST_MAP.read_text())
|
||||
count_models: Final = {
|
||||
name: entry["token_count_model"]
|
||||
for name, entry in cost_map.items()
|
||||
if isinstance(entry, dict) and "token_count_model" in entry
|
||||
}
|
||||
assert count_models
|
||||
assert {
|
||||
name: count_model
|
||||
for name, count_model in count_models.items()
|
||||
if cost_map.get(count_model["model"], {}).get("litellm_provider") != cost_map[name]["litellm_provider"]
|
||||
or count_model["model"] == name
|
||||
or not count_model.get("region")
|
||||
} == {}
|
||||
|
||||
|
||||
_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
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -4184,6 +4184,39 @@ class TestModelInfoCostMapEchoFilter:
|
|||
assert "litellm_provider" not in info
|
||||
assert "supported_openai_params" not in info
|
||||
|
||||
@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
|
||||
) -> 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
|
||||
|
||||
|
|
|
|||
|
|
@ -54,6 +54,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,
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -584,6 +584,40 @@ def test_capability_flag_stays_on_the_deployment_that_set_it(flag: str) -> None:
|
|||
_restore_model_cost_entries(model_keys)
|
||||
|
||||
|
||||
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.anthropic.claude-sonnet-5-5"
|
||||
deploy_id: Final = "sonnet-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": "sonnet-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
|
||||
|
|
@ -3884,9 +3918,7 @@ def test_router_capacity_widens_to_a_larger_base_model_window(
|
|||
)
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
router: Final = _capacity_test_router(route, base, base_location, explicit_limit)
|
||||
expected: Final = (
|
||||
explicit_limit if explicit_limit is not None else (1024 if base_location is None else 8192)
|
||||
)
|
||||
expected: Final = explicit_limit if explicit_limit is not None else (1024 if base_location is None else 8192)
|
||||
assert _effective_max_input_tokens(router) == expected
|
||||
assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0)
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
|
|
|||
|
|
@ -1138,6 +1138,13 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
|
|||
"supports_box_questions": {"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_bedrock_runtime_chat_completions_tools_with_reasoning": {"type": "boolean"},
|
||||
"supports_bedrock_runtime_chat_completions_response_format": {"type": "boolean"},
|
||||
"supports_bedrock_runtime_chat_completions_inline_reasoning": {"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
|
|
@ -46847,9 +46847,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 */
|
||||
|
|
@ -46878,6 +46886,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