mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
fix(token_counter): count OpenAI models newer than tiktoken's table with o200k_base (#48728)
* fix(token_counter): count OpenAI models newer than tiktoken's table with o200k_base Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(token_counter): keep tiktoken's encoding for suffixed variants of models it knows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(token_counter): move the o200k provenance note into the assert message Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(token_counter): strip openai/ and azure/ prefixes before picking the encoding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(rust_bridge): azure/gpt-4o now resolves to o200k_base Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(token_counter): check the provider prefix through _fix_model_name to keep the type gate flat Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(token_counter): check o200k on the gpt-6 models tiktoken 0.14 still does not know * fix(token_counter): count Azure deployment names and OpenAI alias groups with their OpenAI tokenizer * fix(token_counter): count mixed, alias, and base-model deployments with every served model's tokenizer A group that serves OpenAI and non-OpenAI models now reserves at least each served model's count, model_group_alias names count as their target group, custom Azure deployment names count with their base model, and azure_ai/, chatgpt/ and similar prefixes resolve to the listed OpenAI model. * fix(token_counter): keep mixed groups on OpenAI tokenizers, cache served models, and count context checks, failed rows and stream fallback with the served model * fix(token-counter): keep mixed groups at their old reservation and count new OpenAI release names with o200k A group that serves an OpenAI model next to another vendor now counts its own name the way tiktoken reads it as written, plus its OpenAI models, so it reserves what main did instead of the lower o200k count base_model picks the tokenizer for estimated usage only when it names an OpenAI model, now also in the Cursor stream counter and the output-limit response Unlisted gpt-5+ and o-series names count with o200k_base, and wildcard routes run the context-window check with the matched OpenAI model * fix(token-counter): count an alias to a group no OpenAI model serves by the alias as written Keeps the count main gives such an alias, so its budget reservation and rate limit match main and no other vendor's tokenizer loads for it. Also sorts the import the main merge left unsorted. * test: type the new token-counting tests' parameters * fix(token-counter): count a release-named group no OpenAI model serves as main does, and snapshot the OpenAI model set before iterating it * fix(token-counter): read served models through the router's allowlisted deployments_named, and coordinate the registration race test * perf(token-counter): classify each distinct served model once when walking a model group Budgeted requests walk every deployment in their model group, since the router's cached copy is outside the router surface the proxy may use. Classifying each distinct served model once instead of once per deployment cuts a 200-deployment group from 2.3-4.5 ms to about 1.4 ms with the same result. --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
rj-cairn
GitHub
parent
c068d1350b
commit
5eb1042bbf
27 changed files with 788 additions and 42 deletions
|
|
@ -109,6 +109,7 @@ Follow these coding conventions for new/updated code (a three-line fix in a lega
|
|||
- No file sprawl: deliberate file and folder structure
|
||||
- Standard over hand-rolled: use the official SDK or a library where one exists. Where none does, follow industry standards instead of inventing local conventions
|
||||
- API-fragmentation-aware: when logic must branch on which API surface produced or consumes data (e.g. chat completions vs Anthropic Messages vs Responses API shapes), proactively look for an existing shared helper (e.g. `litellm_core_utils/prompt_templates/factory.py`) before writing per-surface parsing in the new module. If none exists, add one there instead of duplicating the same format-detection logic in every new guardrail/integration
|
||||
- Current default for new model names: logic that picks behavior from a model-name table (tokenizer, message framing, supported params) gives a provider's model the table does not list yet that provider's current behavior, never its oldest, since providers ship names faster than any table is updated
|
||||
|
||||
Follow conventional commits for commit names and PR titles
|
||||
|
||||
|
|
|
|||
|
|
@ -591,7 +591,7 @@ vertex_ai_safety_settings: Optional[dict] = None
|
|||
####### COMPLETION MODELS ###################
|
||||
from typing import Set
|
||||
|
||||
open_ai_chat_completion_models: Set = set()
|
||||
open_ai_chat_completion_models: Set[str] = set()
|
||||
open_ai_text_completion_models: Set = set()
|
||||
cohere_models: Set = set()
|
||||
cohere_chat_models: Set = set()
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRAC
|
|||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_content_from_model_response, # pyright: ignore[reportUnknownVariableType] # response param accepts an untyped dict
|
||||
)
|
||||
from litellm.litellm_core_utils.token_counter import usage_tokenizer_model
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
ChatCompletionAssistantContentValue,
|
||||
|
|
@ -51,7 +52,7 @@ from litellm.types.utils import (
|
|||
UnreportedUsage,
|
||||
Usage,
|
||||
)
|
||||
from litellm.utils import print_verbose, token_counter
|
||||
from litellm.utils import get_base_model_from_metadata, print_verbose, token_counter
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from openai.types.completion_usage import CompletionUsage
|
||||
|
|
@ -1552,6 +1553,9 @@ def stream_chunk_builder(
|
|||
return stream_chunk_builder_text_completion(chunks=chunks, messages=messages)
|
||||
|
||||
model: Final = chunks[0]["model"]
|
||||
usage_model: Final = usage_tokenizer_model(
|
||||
model, get_base_model_from_metadata(logging_obj.model_call_details) if logging_obj is not None else None
|
||||
)
|
||||
# Initialize the response dictionary
|
||||
response: Final = processor.build_base_response(chunks)
|
||||
|
||||
|
|
@ -1575,7 +1579,7 @@ def stream_chunk_builder(
|
|||
completion_output = get_content_from_model_response(response)
|
||||
usage = processor.calculate_usage(
|
||||
chunks=chunks,
|
||||
model=model,
|
||||
model=usage_model,
|
||||
completion_output=completion_output,
|
||||
messages=messages,
|
||||
reasoning_tokens=0,
|
||||
|
|
@ -1766,7 +1770,7 @@ def stream_chunk_builder(
|
|||
|
||||
usage = processor.calculate_usage(
|
||||
chunks=chunks,
|
||||
model=model,
|
||||
model=usage_model,
|
||||
completion_output=completion_output,
|
||||
messages=messages,
|
||||
reasoning_tokens=reasoning_tokens,
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ import anyio
|
|||
import anyio.lowlevel
|
||||
import httpx
|
||||
import tiktoken
|
||||
from tiktoken.model import MODEL_TO_ENCODING
|
||||
from typing_extensions import ParamSpec, TypeVar
|
||||
|
||||
import litellm
|
||||
|
|
@ -49,6 +50,7 @@ if TYPE_CHECKING:
|
|||
from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace
|
||||
from litellm.litellm_core_utils.url_utils import safe_get
|
||||
from litellm.llms.custom_httpx.http_handler import get_httpx_client, header_value
|
||||
from litellm.router_utils.deployment_pricing import deployment_pricing_models
|
||||
from litellm.rust_bridge.tokenizer import get_encoding
|
||||
from litellm.types.llms.anthropic import (
|
||||
AnthropicContentParamSource,
|
||||
|
|
@ -1005,17 +1007,144 @@ def openai_tokenizer_encoding(model: str) -> "Encoding":
|
|||
|
||||
|
||||
def openai_tokenizer_encoding_name(model: str) -> str:
|
||||
"""The tiktoken encoding name for `model`, without loading the encoding."""
|
||||
from litellm.utils import print_verbose
|
||||
"""The tiktoken encoding name for `model`, without loading the encoding.
|
||||
|
||||
model_to_use: Final = _fix_model_name(model)
|
||||
A listed OpenAI or Azure model that tiktoken does not know is either a suffixed variant of a model it does
|
||||
know (`text-embedding-ada-002-v2`) or a newer release than its table, which gets o200k_base, the encoding every
|
||||
OpenAI model since gpt-4o uses. Provider and region prefixes (`azure/eu/`, `azure_ai/`) are dropped, a name
|
||||
listed only under Azure (`gpt-41`, `openai/o1-mini`) counts as listed through its `azure/` cost-map key, and an
|
||||
unlisted name tiktoken counts with o200k_base (`openai/gpt-5.9`) or named like a gpt-5 or o-series release
|
||||
(`openai/gpt-6.2`) gets it."""
|
||||
model_to_use: Final = _openai_model_name(model) or _fix_model_name(model)
|
||||
if "gpt-4o" in model_to_use:
|
||||
return "o200k_base"
|
||||
try:
|
||||
return tiktoken.encoding_name_for_model(model_to_use)
|
||||
except KeyError:
|
||||
print_verbose("Warning: model not found. Using cl100k_base encoding.")
|
||||
return "cl100k_base"
|
||||
return _suffixed_variant_encoding_name(model_to_use) or "o200k_base"
|
||||
|
||||
|
||||
def unresolved_tokenizer_model(model: str) -> str:
|
||||
"""`model`, or `gpt-3.5-turbo` when resolving the OpenAI model `model` names moves it off cl100k_base, the encoding
|
||||
tiktoken gives the name as written, so a count with the result never drops below that one."""
|
||||
as_written: Final = _fix_model_name(model)
|
||||
as_written_encoding: Final = (
|
||||
"o200k_base" if "gpt-4o" in as_written else (_tiktoken_encoding_name(as_written) or "cl100k_base")
|
||||
)
|
||||
moved_off_cl100k: Final = (
|
||||
as_written_encoding == "cl100k_base" and openai_tokenizer_encoding_name(model) != "cl100k_base"
|
||||
)
|
||||
return "gpt-3.5-turbo" if moved_off_cl100k else model
|
||||
|
||||
|
||||
def usage_tokenizer_model(model: str, base_model: str | None) -> str:
|
||||
"""The model to estimate missing usage with: a deployment's `base_model` when it names an OpenAI model, since a
|
||||
custom deployment name says nothing about its tokenizer, else `model`."""
|
||||
return base_model if base_model is not None and is_openai_model(base_model) else model
|
||||
|
||||
|
||||
def is_openai_model(model: str) -> bool:
|
||||
"""Whether `model` names an OpenAI or Azure OpenAI model, listed in the cost map or counted with o200k_base by
|
||||
tiktoken, with or without its prefixes."""
|
||||
return _openai_model_name(model) is not None
|
||||
|
||||
|
||||
def deployment_openai_model(deployment: object) -> str | None:
|
||||
"""The OpenAI model a router deployment serves, read from its base model first, since a custom deployment name
|
||||
(`azure/my-gpt-5`) says nothing about its tokenizer, or None when it serves no OpenAI model."""
|
||||
return _as_openai_model(_served_model(deployment))
|
||||
|
||||
|
||||
def served_openai_models(deployments: Iterable[object]) -> tuple[str | None, ...]:
|
||||
"""The distinct OpenAI models `deployments` serve, with None standing for any deployment serving another model."""
|
||||
served: Final = dict.fromkeys(_served_model(deployment) for deployment in deployments)
|
||||
return tuple(dict.fromkeys(_as_openai_model(model) for model in served))
|
||||
|
||||
|
||||
def _served_model(deployment: object) -> str | None:
|
||||
return next(iter(deployment_pricing_models(deployment)), None)
|
||||
|
||||
|
||||
def _as_openai_model(model: str | None) -> str | None:
|
||||
return model if model is not None and is_openai_model(model) else None
|
||||
|
||||
|
||||
def single_encoding_openai_model(served: Sequence[str | None]) -> str | None:
|
||||
"""An OpenAI model to count a request with, when every `served_openai_models` entry is an OpenAI model on the same
|
||||
encoding, else None."""
|
||||
if not served or None in served:
|
||||
return None
|
||||
encodings: Final = frozenset(openai_tokenizer_encoding_name(model) for model in served if model is not None)
|
||||
return served[0] if len(encodings) == 1 else None
|
||||
|
||||
|
||||
def _openai_model_name(model: str) -> str | None:
|
||||
return _listed_openai_model_name(model) or _tiktoken_o200k_model_name(model) or _versioned_openai_model_name(model)
|
||||
|
||||
|
||||
_VERSIONED_OPENAI_MODEL: Final = re.compile(r"(?:gpt-(?:[5-9]|[1-9]\d+)|o[1-9]\d*)(?:[.-]|$)")
|
||||
|
||||
|
||||
def _versioned_openai_model_name(model: str) -> str | None:
|
||||
name: Final = model.rsplit("/", 1)[-1]
|
||||
return name if _VERSIONED_OPENAI_MODEL.match(name) else None
|
||||
|
||||
|
||||
def _tiktoken_o200k_model_name(model: str) -> str | None:
|
||||
name: Final = model.rsplit("/", 1)[-1]
|
||||
return name if _tiktoken_encoding_name(name) == "o200k_base" else None
|
||||
|
||||
|
||||
def _listed_openai_model_name(model: str) -> str | None:
|
||||
segments: Final = model.split("/")
|
||||
suffixes: Final = ("/".join(segments[start:]) for start in range(len(segments)))
|
||||
listed: Final = any(
|
||||
_is_openai_cost_map_model(suffix) or _is_openai_cost_map_model(f"azure/{suffix}") for suffix in suffixes
|
||||
)
|
||||
return segments[-1] if listed else None
|
||||
|
||||
|
||||
def _is_openai_cost_map_model(model: str) -> bool:
|
||||
return (
|
||||
model in litellm.azure_llms
|
||||
or model in litellm.open_ai_chat_completion_models
|
||||
or _is_azure_openai_text_model(model)
|
||||
)
|
||||
|
||||
|
||||
_TEXT_MODEL_MODES: Final = frozenset({"chat", "completion", "responses", "realtime"})
|
||||
|
||||
|
||||
def _is_azure_openai_text_model(model: str) -> bool:
|
||||
cost_info: Final[Mapping[str, object]] = litellm.model_cost.get(model) or {}
|
||||
return (
|
||||
cost_info.get("litellm_provider") == "azure"
|
||||
and cost_info.get("mode") in _TEXT_MODEL_MODES
|
||||
and _is_openai_model_name(model.rsplit("/", 1)[-1])
|
||||
)
|
||||
|
||||
|
||||
def _is_openai_model_name(name: str) -> bool:
|
||||
return (
|
||||
name in litellm.open_ai_chat_completion_models
|
||||
or _tiktoken_encoding_name(name) is not None
|
||||
or any(
|
||||
name.startswith(f"{listed}-") or listed.startswith(f"{name}-")
|
||||
for listed in tuple(litellm.open_ai_chat_completion_models)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _tiktoken_encoding_name(model: str) -> str | None:
|
||||
try:
|
||||
return tiktoken.encoding_name_for_model(model)
|
||||
except KeyError:
|
||||
return None
|
||||
|
||||
|
||||
def _suffixed_variant_encoding_name(model: str) -> str | None:
|
||||
known: Final = tuple(name for name in MODEL_TO_ENCODING if model.startswith(f"{name}-"))
|
||||
return MODEL_TO_ENCODING[max(known, key=len)] if known else None
|
||||
|
||||
|
||||
def uses_legacy_message_accounting(model: str) -> bool:
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ from litellm.types.utils import (
|
|||
from litellm.utils import (
|
||||
CustomStreamWrapper,
|
||||
convert_to_model_response_object,
|
||||
get_base_model_from_metadata,
|
||||
)
|
||||
|
||||
from ...types.llms.openai import HttpxBinaryResponseContent
|
||||
|
|
@ -145,6 +146,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
azure_client: AzureOpenAI | OpenAI,
|
||||
data: dict,
|
||||
timeout: float | httpx.Timeout,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
):
|
||||
"""
|
||||
Helper to:
|
||||
|
|
@ -160,7 +162,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
except BadRequestError as e:
|
||||
if not is_output_token_limit_error(e):
|
||||
raise
|
||||
return build_output_token_limit_response(e=e, data=data, is_async=False)
|
||||
return build_output_token_limit_response(
|
||||
e=e,
|
||||
data=data,
|
||||
is_async=False,
|
||||
base_model=get_base_model_from_metadata(logging_obj.model_call_details),
|
||||
)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
|
@ -192,7 +199,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
except BadRequestError as e:
|
||||
if not is_output_token_limit_error(e):
|
||||
raise
|
||||
return build_output_token_limit_response(e=e, data=data, is_async=True)
|
||||
return build_output_token_limit_response(
|
||||
e=e,
|
||||
data=data,
|
||||
is_async=True,
|
||||
base_model=get_base_model_from_metadata(logging_obj.model_call_details),
|
||||
)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
|
@ -350,7 +362,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
)
|
||||
|
||||
headers, response = self.make_sync_azure_openai_chat_completion_request(
|
||||
azure_client=azure_client, data=data, timeout=timeout
|
||||
azure_client=azure_client, data=data, timeout=timeout, logging_obj=logging_obj
|
||||
)
|
||||
if not isinstance(response, BaseModel):
|
||||
raise AzureOpenAIError(
|
||||
|
|
@ -564,7 +576,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
},
|
||||
)
|
||||
headers, response = self.make_sync_azure_openai_chat_completion_request(
|
||||
azure_client=azure_client, data=data, timeout=timeout
|
||||
azure_client=azure_client, data=data, timeout=timeout, logging_obj=logging_obj
|
||||
)
|
||||
logging_obj.model_call_details["response_headers"] = headers
|
||||
streamwrapper: Final = CustomStreamWrapper(
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ if TYPE_CHECKING:
|
|||
import litellm
|
||||
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object
|
||||
from litellm.litellm_core_utils.token_counter import token_counter
|
||||
from litellm.litellm_core_utils.token_counter import token_counter, usage_tokenizer_model
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
|
||||
|
|
@ -242,7 +242,7 @@ async def _aiter_once(chunk: ChatCompletionChunk) -> AsyncIterator[ChatCompletio
|
|||
|
||||
|
||||
def build_output_token_limit_response(
|
||||
e: openai.BadRequestError, data: Mapping[str, object], is_async: bool
|
||||
e: openai.BadRequestError, data: Mapping[str, object], is_async: bool, base_model: str | None
|
||||
) -> tuple[httpx.Headers, ChatCompletion | Iterator[ChatCompletionChunk] | AsyncIterator[ChatCompletionChunk]]:
|
||||
"""Synthesize the length-truncated response the provider itself returns for slightly larger output budgets.
|
||||
|
||||
|
|
@ -251,7 +251,11 @@ def build_output_token_limit_response(
|
|||
"""
|
||||
model: Final[str] = str(data.get("model", ""))
|
||||
messages: Final = data.get("messages")
|
||||
prompt_tokens: Final = token_counter(model=model, messages=messages) if isinstance(messages, list) else 0
|
||||
prompt_tokens: Final = (
|
||||
token_counter(model=usage_tokenizer_model(model, base_model), messages=messages)
|
||||
if isinstance(messages, list)
|
||||
else 0
|
||||
)
|
||||
if not data.get("stream"):
|
||||
return e.response.headers, _output_token_limit_completion(model, prompt_tokens)
|
||||
chunk: Final = _output_token_limit_chunk(model)
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ from litellm.utils import (
|
|||
CustomStreamWrapper,
|
||||
ProviderConfigManager,
|
||||
convert_to_model_response_object,
|
||||
get_base_model_from_metadata,
|
||||
)
|
||||
|
||||
from ...types.llms.openai import *
|
||||
|
|
@ -557,7 +558,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
except openai.BadRequestError as e:
|
||||
if not is_output_token_limit_error(e):
|
||||
raise
|
||||
return build_output_token_limit_response(e=e, data=data, is_async=True)
|
||||
return build_output_token_limit_response(
|
||||
e=e,
|
||||
data=data,
|
||||
is_async=True,
|
||||
base_model=get_base_model_from_metadata(logging_obj.model_call_details),
|
||||
)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
|
@ -604,7 +610,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
except openai.BadRequestError as e:
|
||||
if not is_output_token_limit_error(e):
|
||||
raise
|
||||
return build_output_token_limit_response(e=e, data=data, is_async=False)
|
||||
return build_output_token_limit_response(
|
||||
e=e,
|
||||
data=data,
|
||||
is_async=False,
|
||||
base_model=get_base_model_from_metadata(logging_obj.model_call_details),
|
||||
)
|
||||
except Exception as e:
|
||||
if raw_response is not None:
|
||||
raise Exception(
|
||||
|
|
|
|||
|
|
@ -76,6 +76,7 @@ from litellm.proxy.hooks.batch_enqueued_tokens import (
|
|||
canonical_provider_batch_id,
|
||||
)
|
||||
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
|
||||
from litellm.proxy.spend_tracking.input_tokens import tokenizer_models
|
||||
from litellm.responses.litellm_completion_transformation.compaction_items import billable_input_since_newest_compaction
|
||||
from litellm.router_utils.add_retry_fallback_headers import (
|
||||
ensure_response_additional_headers,
|
||||
|
|
@ -159,6 +160,12 @@ def _resolve_model_group_alias_via_proxy_router(model: str) -> str | None:
|
|||
return resolve_model_group_alias(llm_router.model_group_alias, model)
|
||||
|
||||
|
||||
def _tokenizer_models_via_proxy_router(model: str) -> tuple[str, ...]:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
return tokenizer_models(model, llm_router)
|
||||
|
||||
|
||||
def _resolve_ptu_team_ceiling_via_proxy_router(team_id: str, model_group: str) -> PTUTeamCeiling | None:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
|
|
@ -884,6 +891,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
] = _resolve_ptu_team_ceiling_via_proxy_router,
|
||||
fail_closed_resolver: Callable[[], bool] = _fail_closed_rate_limit_enforcement_from_general_settings,
|
||||
force_hash_tag_grouping_resolver: Callable[[], bool] = _force_redis_hash_tag_grouping_from_litellm_settings,
|
||||
tokenizer_models_resolver: Callable[[str], tuple[str, ...]] = _tokenizer_models_via_proxy_router,
|
||||
):
|
||||
self.internal_usage_cache = internal_usage_cache
|
||||
self._time_provider = time_provider or datetime.now
|
||||
|
|
@ -893,6 +901,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
self._ptu_team_ceiling_resolver = ptu_team_ceiling_resolver
|
||||
self._fail_closed_resolver = fail_closed_resolver
|
||||
self._force_hash_tag_grouping_resolver = force_hash_tag_grouping_resolver
|
||||
self._tokenizer_models_resolver = tokenizer_models_resolver
|
||||
if self.internal_usage_cache.dual_cache.redis_cache is not None:
|
||||
self.batch_rate_limiter_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
|
||||
BATCH_RATE_LIMITER_SCRIPT
|
||||
|
|
@ -4009,15 +4018,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
try:
|
||||
estimate: Final = max(
|
||||
0,
|
||||
int(
|
||||
token_counter(
|
||||
model=model or "",
|
||||
messages=countable_messages,
|
||||
text=selected_text,
|
||||
tools=countable_tools,
|
||||
tool_choice=countable_tool_choice,
|
||||
use_default_image_token_count=True,
|
||||
*(
|
||||
int(
|
||||
token_counter(
|
||||
model=tokenizer_model,
|
||||
messages=countable_messages,
|
||||
text=selected_text,
|
||||
tools=countable_tools,
|
||||
tool_choice=countable_tool_choice,
|
||||
use_default_image_token_count=True,
|
||||
)
|
||||
)
|
||||
for tokenizer_model in self._tokenizer_models_resolver(model or "")
|
||||
),
|
||||
)
|
||||
return estimate + audio_token_estimate
|
||||
|
|
|
|||
|
|
@ -100,7 +100,11 @@ 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 UnsupportedContentTypeError, offload_token_count
|
||||
from litellm.litellm_core_utils.token_counter import (
|
||||
UnsupportedContentTypeError,
|
||||
deployment_openai_model,
|
||||
offload_token_count,
|
||||
)
|
||||
from litellm.llms.bedrock.count_tokens.kimi_k3 import (
|
||||
HUGGING_FACE_SOURCE,
|
||||
KimiK3Encodings,
|
||||
|
|
@ -15542,7 +15546,7 @@ async def _count_request_tokens(
|
|||
count_source="local_estimate",
|
||||
)
|
||||
estimate: Final = await _local_token_estimate(
|
||||
model=model_to_use,
|
||||
model=deployment_openai_model(deployment_view) or model_to_use,
|
||||
prompt=prompt if counted_messages is None else None,
|
||||
messages=counted_messages,
|
||||
tools=counted_tools,
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata
|
|||
from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
from litellm.litellm_core_utils.token_counter import usage_tokenizer_model
|
||||
from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, first_validation_error
|
||||
from litellm.llms.openai.responses.guardrail_translation.handler import (
|
||||
OpenAIResponsesHandler,
|
||||
|
|
@ -96,6 +97,7 @@ from litellm.types.llms.openai import (
|
|||
from litellm.types.responses.main import DeleteResponseResult
|
||||
from litellm.types.token_count import TokenCountSourceFields
|
||||
from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, TokenCountResponse
|
||||
from litellm.utils import get_base_model_from_metadata
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
|
@ -537,7 +539,10 @@ async def cursor_model_list(
|
|||
|
||||
def _dispatched_prompt_token_counter(logging_obj: LiteLLMLoggingObj) -> Callable[[], int]:
|
||||
return lambda: count_request_input_tokens(
|
||||
model=str(logging_obj.model_call_details.get("model") or ""),
|
||||
model=usage_tokenizer_model(
|
||||
str(logging_obj.model_call_details.get("model") or ""),
|
||||
get_base_model_from_metadata(logging_obj.model_call_details),
|
||||
),
|
||||
request_input=messages_since_newest_compaction(logging_obj.model_call_details.get("messages")),
|
||||
system_input=None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import time
|
|||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Final, NoReturn, SupportsFloat, SupportsIndex, SupportsInt, cast
|
||||
|
||||
|
|
@ -39,7 +40,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
tag_cache_key,
|
||||
team_membership_reservation_cache_key,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.input_tokens import count_input_tokens, count_input_tokens_for_model
|
||||
from litellm.proxy.spend_tracking.input_tokens import count_input_tokens, count_input_tokens_for_model, tokenizer_models
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement, spend_counter_batch_scope
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.router import Router
|
||||
|
|
@ -1599,7 +1600,19 @@ async def count_request_input_tokens(
|
|||
models: Final = _get_request_models(request_body=request_body, route=route, llm_router=llm_router)
|
||||
if not models:
|
||||
return MappingProxyType({})
|
||||
return await count_input_tokens(request_body=request_body, raw_body=raw_body, models=models)
|
||||
counted_as: Final = MappingProxyType({model: tokenizer_models(model, llm_router) for model in models})
|
||||
counts: Final = await count_input_tokens(
|
||||
request_body=request_body,
|
||||
raw_body=raw_body,
|
||||
models=tuple(dict.fromkeys(chain.from_iterable(counted_as.values()))),
|
||||
)
|
||||
return MappingProxyType(
|
||||
{
|
||||
model: count
|
||||
for model, names in counted_as.items()
|
||||
if (count := max((counts[name] for name in names if name in counts), default=None)) is not None
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _estimate_input_tokens(
|
||||
|
|
|
|||
|
|
@ -15,14 +15,16 @@ import asyncio
|
|||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.token_counter import served_openai_models, unresolved_tokenizer_model
|
||||
from litellm.responses.litellm_completion_transformation.compaction_items import (
|
||||
billable_input_since_newest_compaction,
|
||||
holds_compaction_summary,
|
||||
)
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
from litellm.rust_bridge import runtime
|
||||
from litellm.rust_bridge.catalog import Route, RouteContext
|
||||
from litellm.rust_bridge.token_counter import (
|
||||
|
|
@ -33,6 +35,9 @@ from litellm.rust_bridge.token_counter import (
|
|||
rust_tokenizer,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
||||
TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000
|
||||
|
||||
_INPUT_SIZE_FIELDS: Final = ("messages", "prompt", "input", "query", "documents", "tools", "tool_choice")
|
||||
|
|
@ -47,6 +52,25 @@ def _approximate_input_size(request_body: Mapping[str, object]) -> int:
|
|||
return sum(len(str(request_body.get(field, ""))) for field in _INPUT_SIZE_FIELDS)
|
||||
|
||||
|
||||
def tokenizer_models(model: str, llm_router: Router | None) -> tuple[str, ...]:
|
||||
"""The names to count a request for `model` with, keeping the largest count.
|
||||
|
||||
An alias resolves to its target group. A group served only by OpenAI models counts as those models, so it gets
|
||||
their tokenizer. Any other group counts as its OpenAI models, if any, and as the requested name with the encoding
|
||||
tiktoken gives that name as written, so its other deployments keep that count without loading their vendor's
|
||||
tokenizer. A name the router does not index counts as written."""
|
||||
if llm_router is None:
|
||||
return (model,)
|
||||
group: Final = resolve_model_group_alias(llm_router.model_group_alias, model) or model
|
||||
served: Final = served_openai_models(llm_router.deployments_named(group) or ())
|
||||
if not served:
|
||||
return (model,)
|
||||
openai_served: Final = tuple(served_model for served_model in served if served_model is not None)
|
||||
if len(openai_served) == len(served):
|
||||
return openai_served
|
||||
return (unresolved_tokenizer_model(model), *openai_served)
|
||||
|
||||
|
||||
async def count_input_tokens(
|
||||
request_body: dict,
|
||||
raw_body: bytes | None,
|
||||
|
|
|
|||
|
|
@ -168,7 +168,7 @@ from litellm.litellm_core_utils.served_output_texts import (
|
|||
record_served_output_texts,
|
||||
served_stream_output_texts,
|
||||
)
|
||||
from litellm.litellm_core_utils.token_counter import offload_token_count
|
||||
from litellm.litellm_core_utils.token_counter import offload_token_count, usage_tokenizer_model
|
||||
from litellm.llms import load_guardrail_translation_mappings
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.proxy._types import (
|
||||
|
|
@ -1139,7 +1139,9 @@ def _estimated_prompt_usage(
|
|||
None,
|
||||
)
|
||||
return _estimate_dispatched_failure_usage(
|
||||
model=str(model_call_details.get("model") or ""),
|
||||
model=usage_tokenizer_model(
|
||||
str(model_call_details.get("model") or ""), get_base_model_from_metadata(model_call_details)
|
||||
),
|
||||
request_input=(
|
||||
messages_since_newest_compaction(dispatched_input)
|
||||
if dispatched_call_type in _RESPONSES_CALL_TYPES
|
||||
|
|
|
|||
|
|
@ -162,7 +162,11 @@ from litellm.litellm_core_utils.sensitive_data_masker import (
|
|||
SensitiveDataMasker,
|
||||
mask_sensitive_structure,
|
||||
)
|
||||
from litellm.litellm_core_utils.token_counter import offload_token_count
|
||||
from litellm.litellm_core_utils.token_counter import (
|
||||
offload_token_count,
|
||||
served_openai_models,
|
||||
single_encoding_openai_model,
|
||||
)
|
||||
from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, drop_policy
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo, anthropic_error_frame_exception
|
||||
from litellm.llms.base_llm.passthrough.transformation import replace_path_segment
|
||||
|
|
@ -1533,6 +1537,7 @@ class Router:
|
|||
self._cached_deployment_tracks_usage = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)(
|
||||
self._deployment_tracks_usage
|
||||
)
|
||||
self.served_openai_models = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)(self._served_openai_models)
|
||||
self._discovered_model_info_cache: InMemoryCache = InMemoryCache(
|
||||
max_size_in_memory=max(len(model_list or ()), 1),
|
||||
default_ttl=2 * MODEL_INFO_REFRESH_SECONDS,
|
||||
|
|
@ -11890,6 +11895,9 @@ class Router:
|
|||
indices: Final = self.model_name_to_deployment_indices.get(model_name)
|
||||
return None if indices is None else tuple(self.model_list[idx] for idx in indices)
|
||||
|
||||
def _served_openai_models(self, model_name: str) -> tuple[str | None, ...]:
|
||||
return served_openai_models(self.deployments_named(model_name) or ())
|
||||
|
||||
def team_public_model_names_of(self, team_id: str) -> frozenset[str]:
|
||||
"""The public names `team_id` serves its own deployments under."""
|
||||
with self._deployment_index_lock:
|
||||
|
|
@ -14044,6 +14052,7 @@ class Router:
|
|||
self._ptu_share_teams.cache_clear()
|
||||
self._cached_team_ptu_ceiling.cache_clear()
|
||||
self._cached_deployment_tracks_usage.cache_clear()
|
||||
self.served_openai_models.cache_clear()
|
||||
self._zero_cost_cache.clear()
|
||||
self._routing_group_rows = None
|
||||
|
||||
|
|
@ -14303,6 +14312,7 @@ class Router:
|
|||
messages: Sequence[Mapping[str, object]] | None,
|
||||
input: str | list[object] | None,
|
||||
request_kwargs: Mapping[str, object] | None = None,
|
||||
tokenizer_model: str = "",
|
||||
) -> int:
|
||||
"""
|
||||
Count input tokens for context-window pre-call checks.
|
||||
|
|
@ -14333,7 +14343,7 @@ class Router:
|
|||
system_message: Final = anthropic_system_to_openai_message(extras.get("system"))
|
||||
if messages is not None:
|
||||
counted_messages: Final = (system_message, *messages) if system_message is not None else messages
|
||||
return litellm.token_counter(messages=counted_messages, tools=tools)
|
||||
return litellm.token_counter(model=tokenizer_model, messages=counted_messages, tools=tools)
|
||||
if input is not None:
|
||||
from openai.types.responses.response_create_params import ResponseInputParam
|
||||
|
||||
|
|
@ -14348,6 +14358,7 @@ class Router:
|
|||
only_since_compaction=True,
|
||||
)
|
||||
return litellm.token_counter(
|
||||
model=tokenizer_model,
|
||||
messages=cast(list, input_messages), # cast-ok: transformed chat messages
|
||||
tools=tools,
|
||||
)
|
||||
|
|
@ -14404,6 +14415,10 @@ class Router:
|
|||
messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter
|
||||
input=cast(str | list | None, input), # cast-ok: forwarded to the sync counter
|
||||
request_kwargs=request_kwargs,
|
||||
tokenizer_model=single_encoding_openai_model(
|
||||
self.served_openai_models(model) or served_openai_models(healthy_deployments)
|
||||
)
|
||||
or "",
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request
|
||||
verbose_router_logger.error(
|
||||
|
|
@ -14418,6 +14433,8 @@ class Router:
|
|||
request_kwargs: Mapping[str, object] | None,
|
||||
input_token_count: int | None,
|
||||
count_inline: bool,
|
||||
model: str,
|
||||
deployments: Sequence[Mapping[str, object]],
|
||||
) -> int | None:
|
||||
if input_token_count is not None:
|
||||
return input_token_count
|
||||
|
|
@ -14428,6 +14445,10 @@ class Router:
|
|||
messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter
|
||||
input=cast(str | list[object] | None, input), # cast-ok: forwarded to the sync counter
|
||||
request_kwargs=request_kwargs,
|
||||
tokenizer_model=single_encoding_openai_model(
|
||||
self.served_openai_models(model) or served_openai_models(deployments)
|
||||
)
|
||||
or "",
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_router_logger.error(
|
||||
|
|
@ -14622,6 +14643,8 @@ class Router:
|
|||
count_inline=has_countable_input
|
||||
and not skip_inline_token_count
|
||||
and any(limit is not None for limit in deployment_max_input_tokens),
|
||||
model=model,
|
||||
deployments=_returned_deployments,
|
||||
)
|
||||
request_settings: Final[Mapping[str, object]] = request_kwargs if request_kwargs is not None else {}
|
||||
requested_params: Final[frozenset[str]] = (
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ from litellm import token_counter
|
|||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.litellm_core_utils.token_counter import offload_token_count
|
||||
from litellm.litellm_core_utils.token_counter import deployment_openai_model, offload_token_count
|
||||
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
|
||||
from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET
|
||||
from litellm.types.router import RouterCacheEnum, RouterErrors
|
||||
|
|
@ -127,6 +127,12 @@ def _get_cache_keys(deployment: dict, current_minute: str) -> tuple[str, str] |
|
|||
return itpm_key, otpm_key
|
||||
|
||||
|
||||
def _tokenizer_model(deployment: object) -> str:
|
||||
params: Final = deployment.get("litellm_params") if is_str_keyed_mapping(deployment) else None
|
||||
served: Final = params.get("model") if is_str_keyed_mapping(params) else None
|
||||
return deployment_openai_model(deployment) or (served if isinstance(served, str) else "")
|
||||
|
||||
|
||||
def _estimate_input_tokens(request_kwargs: dict[str, Any] | None, model: str = "") -> int:
|
||||
if not request_kwargs:
|
||||
return 0
|
||||
|
|
@ -427,7 +433,7 @@ def io_token_pre_call_check(
|
|||
return deployment
|
||||
|
||||
request_kwargs: Final = get_io_token_rate_limit_request_kwargs()
|
||||
_model: Final = (deployment.get("litellm_params") or {}).get("model") or ""
|
||||
_model: Final = _tokenizer_model(deployment)
|
||||
estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model)
|
||||
max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment)
|
||||
|
||||
|
|
@ -491,7 +497,7 @@ async def async_io_token_pre_call_check(
|
|||
return deployment
|
||||
|
||||
request_kwargs: Final = get_io_token_rate_limit_request_kwargs()
|
||||
_model: Final = (deployment.get("litellm_params") or {}).get("model") or ""
|
||||
_model: Final = _tokenizer_model(deployment)
|
||||
estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model)
|
||||
max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment)
|
||||
|
||||
|
|
|
|||
|
|
@ -175,6 +175,7 @@ ignored_function_names = [
|
|||
"_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
|
||||
"_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
|
||||
"_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py
|
||||
"_served_openai_models", # Cached as served_openai_models, tested through the pre-call check token count tests in test_router.py
|
||||
"priority_header_update", # Tested through the priority header tests in test_common_request_processing.py
|
||||
"_drops_params", # Read by split_request_priority and priority_header_update, tested through both
|
||||
"_responses_event_announces_background_run", # Tested through the background Responses stream tests in test_router.py
|
||||
|
|
|
|||
|
|
@ -2519,3 +2519,39 @@ def test_stream_chunk_builder_carries_the_bedrock_guardrail_trace_and_records_it
|
|||
(entry,) = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"]
|
||||
assert entry["guardrail_usage"] == {"topicPolicyUnits": 2}
|
||||
assert entry["guardrail_cost"] == pytest.approx(2 * 0.00015)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("base_model", "counted_as"), [("azure/gpt-6.1-sol", "azure/gpt-6.1-sol"), ("claude-sonnet-4-5", "my-gpt-6")]
|
||||
)
|
||||
def test_stream_chunk_builder_counts_a_custom_named_deployments_missing_usage_with_its_base_model_only_when_openai(
|
||||
base_model: str, counted_as: str
|
||||
) -> None:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
messages: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}]
|
||||
chunks: Final = [
|
||||
ModelResponseStream(id="s", model="my-gpt-6", choices=[StreamingChoices(index=0, delta=Delta(content="hi"))]),
|
||||
ModelResponseStream(
|
||||
id="s", model="my-gpt-6", choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")]
|
||||
),
|
||||
]
|
||||
logging_obj: Final = Logging(
|
||||
model="my-gpt-6",
|
||||
messages=messages,
|
||||
stream=True,
|
||||
call_type="acompletion",
|
||||
start_time=None,
|
||||
litellm_call_id="call",
|
||||
function_id="fn",
|
||||
)
|
||||
logging_obj.model_call_details["litellm_params"] = {"metadata": {"model_info": {"base_model": base_model}}}
|
||||
|
||||
built: Final = stream_chunk_builder(chunks=chunks, messages=messages, logging_obj=logging_obj)
|
||||
|
||||
assert built is not None
|
||||
assert litellm.token_counter(model=base_model, messages=messages) != litellm.token_counter(
|
||||
model="my-gpt-6", messages=messages
|
||||
)
|
||||
assert built.usage.prompt_tokens == litellm.token_counter(model=counted_as, messages=messages)
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ from litellm.litellm_core_utils.token_counter import (
|
|||
calculate_img_tokens,
|
||||
estimate_image_tokens,
|
||||
high_detail_image_token_upper_bound,
|
||||
is_openai_model,
|
||||
offload_token_count,
|
||||
)
|
||||
from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
|
||||
|
|
@ -2600,3 +2601,95 @@ def test_get_modified_max_tokens_leaves_a_gemini_4_long_decode_budget_alone(monk
|
|||
)
|
||||
== 200_000
|
||||
)
|
||||
|
||||
|
||||
_MULTILINGUAL_SAMPLE: Final = "日本語のテキストとहिन्दी पाठ और русский текст. " * 20
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"gpt-6-sol",
|
||||
"gpt-6.1-sol",
|
||||
"gpt-6.1-terra",
|
||||
"azure/gpt-41",
|
||||
"azure/gpt-5.1",
|
||||
"openai/gpt-5.1",
|
||||
"openai/gpt-4o",
|
||||
"gpt-41",
|
||||
"eu/gpt-6.1-sol",
|
||||
"azure/eu/gpt-6.1-sol",
|
||||
"azure/gpt-5.1-chat",
|
||||
"azure/gpt-6.1-sol-2026-09-29",
|
||||
"openai/gpt-6.1-sol-2026-09-29",
|
||||
"azure_ai/gpt-6.1-sol",
|
||||
"chatgpt/gpt-6.1-sol",
|
||||
"perplexity/openai/gpt-5",
|
||||
],
|
||||
)
|
||||
def test_token_counter_counts_openai_models_newer_than_tiktokens_table_with_o200k(model: str):
|
||||
expected: Final = len(tiktoken.get_encoding("o200k_base").encode(_MULTILINGUAL_SAMPLE))
|
||||
|
||||
assert expected != len(tiktoken.get_encoding("cl100k_base").encode(_MULTILINGUAL_SAMPLE))
|
||||
assert litellm.token_counter(model=model, text=_MULTILINGUAL_SAMPLE) == expected, (
|
||||
f"{model} should count with o200k_base. OpenAI's /v1/responses/input_tokens gave gpt-5, gpt-5.1, gpt-6-sol, "
|
||||
"gpt-6-luna, gpt-6-astra, gpt-6.1-sol and gpt-6.1-terra the same count on this sample, checked 2026-10-28"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["openai/gpt-5.9", "chatgpt/gpt-5.3-instant", "azure/gpt-5.9-mini"])
|
||||
def test_token_counter_counts_unlisted_openai_names_with_the_encoding_tiktoken_gives_them(model: str):
|
||||
assert litellm.token_counter(model=model, text=_MULTILINGUAL_SAMPLE) == len(
|
||||
tiktoken.encoding_for_model(model.rsplit("/", 1)[-1]).encode(_MULTILINGUAL_SAMPLE)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["gpt-6.2", "openai/gpt-6.2", "azure/gpt-6.2-mini", "azure/eu/o9-mini", "gpt-10"])
|
||||
def test_token_counter_counts_unlisted_names_of_newer_openai_releases_with_o200k(model: str):
|
||||
assert litellm.token_counter(model=model, text=_MULTILINGUAL_SAMPLE) == len(
|
||||
tiktoken.get_encoding("o200k_base").encode(_MULTILINGUAL_SAMPLE)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"groq/qwen3-32b",
|
||||
"openai/my-llama",
|
||||
"text-embedding-ada-002-v2",
|
||||
"azure/gpt-4",
|
||||
"azure/my-deployment",
|
||||
"azure/ada",
|
||||
"azure/mistral-large-latest",
|
||||
"azure/command-r-plus",
|
||||
"azure_ai/mistral-large-latest",
|
||||
],
|
||||
)
|
||||
def test_token_counter_keeps_cl100k_for_non_openai_models_and_variants_of_cl100k_models(model: str):
|
||||
assert litellm.token_counter(model=model, text=_MULTILINGUAL_SAMPLE) == len(
|
||||
tiktoken.get_encoding("cl100k_base").encode(_MULTILINGUAL_SAMPLE)
|
||||
)
|
||||
|
||||
|
||||
def test_openai_model_check_holds_while_openai_models_are_registered_concurrently(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setitem(litellm.model_cost, "azure/azure-only-chat-model", {"litellm_provider": "azure", "mode": "chat"})
|
||||
registered: Final = tuple(f"concurrently-registered-model-{index}" for index in range(200))
|
||||
checked: Final = threading.Semaphore(0)
|
||||
|
||||
def register() -> None:
|
||||
for name in registered:
|
||||
if not checked.acquire(timeout=10):
|
||||
return
|
||||
litellm.open_ai_chat_completion_models.add(name)
|
||||
|
||||
try:
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
registering: Final = pool.submit(register)
|
||||
while not registering.done():
|
||||
assert not is_openai_model("azure-only-chat-model")
|
||||
checked.release()
|
||||
registering.result()
|
||||
finally:
|
||||
litellm.open_ai_chat_completion_models.difference_update(registered)
|
||||
|
|
|
|||
|
|
@ -322,6 +322,29 @@ def test_mapped_response_still_bills_the_prompt_the_provider_processed(provider)
|
|||
assert litellm.completion_cost(completion_response=response) > 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ["openai", "azure"])
|
||||
@pytest.mark.parametrize(
|
||||
("base_model", "counted_as"), [("azure/gpt-6.1-sol", "azure/gpt-6.1-sol"), ("claude-sonnet-4-5", "my-gpt-6")]
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_mapped_response_bills_a_custom_named_deployment_with_its_base_model_only_when_it_is_openai(
|
||||
provider: str, base_model: str, counted_as: str
|
||||
) -> None:
|
||||
messages: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}]
|
||||
overrides: Final = {"model": f"{provider}/my-gpt-6", "messages": messages, "base_model": base_model}
|
||||
|
||||
sync_response: Final = litellm.completion(
|
||||
**_completion_kwargs(provider, _sync_client_raising(provider, OUTPUT_LIMIT_400_MESSAGE), **overrides)
|
||||
)
|
||||
async_response: Final = await litellm.acompletion(
|
||||
**_completion_kwargs(provider, _async_client_raising(provider, OUTPUT_LIMIT_400_MESSAGE), **overrides)
|
||||
)
|
||||
|
||||
expected: Final = token_counter(model=counted_as, messages=messages)
|
||||
assert token_counter(model=base_model, messages=messages) != token_counter(model="my-gpt-6", messages=messages)
|
||||
assert (sync_response.usage.prompt_tokens, async_response.usage.prompt_tokens) == (expected, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ["openai", "azure"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_output_limit_400_maps_to_length_truncated_response(provider):
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import time
|
|||
from collections.abc import AsyncGenerator, Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from functools import partial
|
||||
from typing import Any, Dict, Final, List, Optional
|
||||
from unittest.mock import patch
|
||||
|
||||
|
|
@ -45,6 +46,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
|||
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
||||
PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.input_tokens import tokenizer_models
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
|
||||
from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token
|
||||
from litellm.router_utils.ptu_shares import PTUTeamCeiling
|
||||
|
|
@ -9172,3 +9174,42 @@ async def test_a_successful_extraction_counts_its_input_tokens_toward_the_key_tp
|
|||
call_type="aextraction",
|
||||
)
|
||||
assert refused.value.status_code == 429
|
||||
|
||||
|
||||
_MULTILINGUAL_REQUEST: Final[dict[str, object]] = {
|
||||
"messages": [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}]
|
||||
}
|
||||
|
||||
|
||||
def _project_itpm_estimate(router: Router, model: str) -> int:
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(DualCache()),
|
||||
tokenizer_models_resolver=partial(tokenizer_models, llm_router=router),
|
||||
)
|
||||
return handler._estimate_precise_input_tokens(_MULTILINGUAL_REQUEST, model, "acompletion")
|
||||
|
||||
|
||||
def test_project_itpm_counts_a_group_mixing_in_a_non_openai_model_by_its_name_as_written_and_its_openai_models():
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{"model_name": "gpt-6.1-sol", "litellm_params": {"model": model, "api_key": "sk-test"}}
|
||||
for model in ("openai/gpt-6.1-sol", "anthropic/claude-sonnet-5-5")
|
||||
]
|
||||
)
|
||||
messages: Final = _MULTILINGUAL_REQUEST["messages"]
|
||||
|
||||
assert _project_itpm_estimate(router, "gpt-6.1-sol") == max(
|
||||
litellm.token_counter(model="my-model", messages=messages),
|
||||
litellm.token_counter(model="openai/gpt-6.1-sol", messages=messages),
|
||||
)
|
||||
|
||||
|
||||
def test_project_itpm_counts_a_model_group_alias_as_its_target_group():
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{"model_name": "gpt-6.1-sol", "litellm_params": {"model": "openai/gpt-6.1-sol", "api_key": "sk-test"}}
|
||||
],
|
||||
model_group_alias={"prod": "gpt-6.1-sol"},
|
||||
)
|
||||
|
||||
assert _project_itpm_estimate(router, "prod") == _project_itpm_estimate(router, "gpt-6.1-sol")
|
||||
|
|
|
|||
|
|
@ -189,7 +189,11 @@ _TURN_AFTER_COMPACTION: Final = {"role": "user", "content": "what is the cat's n
|
|||
|
||||
|
||||
async def _cursor_stream_failure_prompt_tokens(
|
||||
monkeypatch: pytest.MonkeyPatch, request_input: list[Mapping[str, str]], fails_after_output: bool
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
request_input: list[Mapping[str, str]],
|
||||
fails_after_output: bool,
|
||||
model: str = "gpt-6-astra",
|
||||
model_info: Mapping[str, str] | None = None,
|
||||
) -> int:
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
|
@ -208,7 +212,6 @@ async def _cursor_stream_failure_prompt_tokens(
|
|||
usage: Final = request_data.get("combined_usage_object")
|
||||
self.prompt_tokens = usage.prompt_tokens if isinstance(usage, litellm.Usage) else None
|
||||
|
||||
model: Final = "gpt-6-astra"
|
||||
upstream_url: Final = "https://streaming.example/v1"
|
||||
created: Final = {
|
||||
"type": "response.created",
|
||||
|
|
@ -239,6 +242,7 @@ async def _cursor_stream_failure_prompt_tokens(
|
|||
{
|
||||
"model_name": model,
|
||||
"litellm_params": {"model": "openai/" + model, "api_base": upstream_url, "api_key": "fixture-key"},
|
||||
"model_info": dict(model_info or {}),
|
||||
}
|
||||
],
|
||||
num_retries=0,
|
||||
|
|
@ -276,6 +280,26 @@ async def test_a_failed_cursor_stream_counts_its_input_from_the_newest_compactio
|
|||
assert resent == dispatched_after_trim
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("fails_after_output", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
("base_model", "counted_as"), [("azure/gpt-6.1-sol", "azure/gpt-6.1-sol"), ("claude-sonnet-4-5", "my-gpt-6")]
|
||||
)
|
||||
async def test_a_failed_cursor_stream_counts_a_custom_named_deployment_with_its_base_model_only_when_it_is_openai(
|
||||
monkeypatch: pytest.MonkeyPatch, base_model: str, counted_as: str, fails_after_output: bool
|
||||
) -> None:
|
||||
from litellm.proxy.utils import count_request_input_tokens
|
||||
|
||||
turn: Final = {"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}
|
||||
|
||||
counted: Final = await _cursor_stream_failure_prompt_tokens(
|
||||
monkeypatch, [turn], fails_after_output, model="my-gpt-6", model_info={"base_model": base_model}
|
||||
)
|
||||
|
||||
assert count_request_input_tokens(base_model, [turn], None) != count_request_input_tokens("my-gpt-6", [turn], None)
|
||||
assert counted == count_request_input_tokens(counted_as, [turn], None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_api_background_polling_rejects_missing_input():
|
||||
from fastapi import Response as FastAPIResponse
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ from litellm.proxy.spend_tracking.budget_reservation import (
|
|||
estimate_request_max_cost,
|
||||
get_budget_window_start,
|
||||
release_unbound_budget_reservation,
|
||||
count_request_input_tokens,
|
||||
reserve_budget_for_request,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
|
@ -426,3 +427,118 @@ async def test_reservation_holds_cost_against_cached_org_and_project_budgets(spe
|
|||
)
|
||||
def test_budget_window_start_reads_json_encoded_windows(window: str, expected: datetime | None) -> None:
|
||||
assert get_budget_window_start(window) == expected
|
||||
|
||||
|
||||
_MULTILINGUAL_MESSAGES: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}]
|
||||
|
||||
|
||||
def _router_with_groups(*groups: tuple[str, str]) -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{"model_name": name, "litellm_params": {"model": model, "api_key": "sk-test"}} for name, model in groups
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
async def _input_tokens(router: Router, model: str) -> dict[str, int]:
|
||||
return dict(
|
||||
await count_request_input_tokens(
|
||||
request_body={"model": model, "messages": _MULTILINGUAL_MESSAGES},
|
||||
route="/v1/chat/completions",
|
||||
llm_router=router,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_alias_group_of_an_openai_model_reserves_on_that_models_token_count():
|
||||
router: Final = _router_with_groups(("gpt-6.1-sol", "openai/gpt-6.1-sol"), ("prod-gpt-6", "openai/gpt-6.1-sol"))
|
||||
|
||||
by_name: Final = await _input_tokens(router, "gpt-6.1-sol")
|
||||
|
||||
assert await _input_tokens(router, "prod-gpt-6") == {"prod-gpt-6": by_name["gpt-6.1-sol"]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_group_alias_of_an_openai_group_reserves_on_that_groups_token_count():
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{"model_name": "gpt-6.1-sol", "litellm_params": {"model": "openai/gpt-6.1-sol", "api_key": "sk-test"}}
|
||||
],
|
||||
model_group_alias={"prod": "gpt-6.1-sol"},
|
||||
)
|
||||
|
||||
by_name: Final = await _input_tokens(router, "gpt-6.1-sol")
|
||||
|
||||
assert await _input_tokens(router, "prod") == {"prod": by_name["gpt-6.1-sol"]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_group_alias_of_a_group_no_openai_model_serves_reserves_on_the_alias_as_written():
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "claude-sonnet-5-5",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": "sk-test"},
|
||||
}
|
||||
],
|
||||
model_group_alias={"claude-alias": "claude-sonnet-5-5"},
|
||||
)
|
||||
|
||||
assert await _input_tokens(router, "claude-alias") == {
|
||||
"claude-alias": litellm.token_counter(model="claude-alias", messages=_MULTILINGUAL_MESSAGES)
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("name", ["gpt-6-claude", "o4-alias"])
|
||||
async def test_release_named_group_or_alias_no_openai_model_serves_reserves_on_an_unknown_names_count(
|
||||
name: str,
|
||||
) -> None:
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{"model_name": group, "litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": "sk-test"}}
|
||||
for group in ("claude-sonnet-5-5", "gpt-6-claude")
|
||||
],
|
||||
model_group_alias={"o4-alias": "claude-sonnet-5-5"},
|
||||
)
|
||||
|
||||
assert await _input_tokens(router, name) == {
|
||||
name: litellm.token_counter(model="my-model", messages=_MULTILINGUAL_MESSAGES)
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("group", "counted_as"),
|
||||
[("gpt-6.1-sol", "my-model"), ("openai/gpt-6.1-sol", "my-model"), ("mixed", "my-model"), ("gpt-4o", "gpt-4o")],
|
||||
)
|
||||
async def test_group_mixing_in_a_non_openai_model_reserves_on_its_name_as_written_and_its_openai_models(
|
||||
group: str, counted_as: str
|
||||
):
|
||||
served: Final = ("openai/gpt-6.1-sol", "anthropic/claude-sonnet-5-5")
|
||||
router: Final = _router_with_groups(*((group, model) for model in served))
|
||||
|
||||
reserved: Final = (await _input_tokens(router, group))[group]
|
||||
|
||||
assert reserved == max(
|
||||
litellm.token_counter(model=counted_as, messages=_MULTILINGUAL_MESSAGES),
|
||||
litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_MESSAGES),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_group_reserves_on_its_current_deployments_after_another_vendors_deployment_joins():
|
||||
router: Final = _router_with_groups(("prod", "openai/gpt-6.1-sol"))
|
||||
openai_only: Final = (await _input_tokens(router, "prod"))["prod"]
|
||||
|
||||
router.add_deployment(
|
||||
Deployment(
|
||||
model_name="prod",
|
||||
litellm_params=LiteLLM_Params(model="anthropic/claude-sonnet-5-5", api_key="sk-test"),
|
||||
)
|
||||
)
|
||||
|
||||
by_own_name: Final = litellm.token_counter(model="prod", messages=_MULTILINGUAL_MESSAGES)
|
||||
assert openai_only != by_own_name
|
||||
assert (await _input_tokens(router, "prod"))["prod"] == by_own_name
|
||||
|
|
|
|||
|
|
@ -3989,6 +3989,32 @@ async def test_responses_count_on_claude_counts_the_turns_previous_response_id_r
|
|||
spend_logs.assert_awaited_once_with("chatcmpl-stored", ("chatcmpl-stored", "resp_client_form"))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_counter_counts_a_custom_azure_deployment_name_with_its_base_models_tokenizer(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
azure_params: Final = {"api_key": "sk-test", "api_base": "https://example.openai.azure.com"}
|
||||
monkeypatch.setattr(
|
||||
litellm.proxy.proxy_server,
|
||||
"llm_router",
|
||||
Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "named",
|
||||
"litellm_params": {"model": "azure/my-gpt-6", "base_model": "azure/gpt-6.1-sol", **azure_params},
|
||||
},
|
||||
{"model_name": "served", "litellm_params": {"model": "azure/gpt-6.1-sol", **azure_params}},
|
||||
]
|
||||
),
|
||||
)
|
||||
messages: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}]
|
||||
|
||||
named: Final = await token_counter(request=TokenCountRequest(model="named", messages=messages))
|
||||
served: Final = await token_counter(request=TokenCountRequest(model="served", messages=messages))
|
||||
|
||||
assert named.total_tokens == served.total_tokens
|
||||
|
||||
|
||||
_COUNTED_HELLO: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{"model": "count-log-gpt", "messages": [{"role": "user", "content": "Hello, how many tokens is this?"}]}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2308,6 +2308,30 @@ def test_a_failed_dispatch_is_estimated_as_input_only():
|
|||
assert usage.total_tokens == usage.prompt_tokens
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("base_model", "counted_as"), [("azure/gpt-6.1-sol", "azure/gpt-6.1-sol"), ("claude-sonnet-4-5", "my-gpt-6")]
|
||||
)
|
||||
def test_a_failed_dispatch_to_a_custom_named_deployment_is_counted_with_its_base_model_only_when_it_is_openai(
|
||||
base_model: str, counted_as: str
|
||||
):
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from litellm.proxy.utils import _estimated_prompt_usage, count_request_input_tokens
|
||||
|
||||
messages: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}]
|
||||
logging_obj = _dispatched_failure_logging_obj(
|
||||
"my-gpt-6", get_litellm_params(metadata={"model_info": {"base_model": base_model}})
|
||||
)
|
||||
logging_obj.model_call_details["messages"] = messages
|
||||
|
||||
usage: Final = _estimated_prompt_usage(logging_obj.model_call_details, {}, "acompletion")
|
||||
|
||||
assert usage is not None
|
||||
assert count_request_input_tokens(base_model, messages, None) != count_request_input_tokens(
|
||||
"my-gpt-6", messages, None
|
||||
)
|
||||
assert usage.prompt_tokens == count_request_input_tokens(counted_as, messages, None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("request_input", [[], object()])
|
||||
def test_nothing_is_estimated_when_there_is_nothing_to_count(request_input):
|
||||
from litellm.proxy.utils import _estimate_dispatched_failure_usage
|
||||
|
|
|
|||
|
|
@ -144,7 +144,7 @@ async def test_each_tokenizer_gets_its_own_cached_counter(fake_tokenizers: None)
|
|||
("gemini/gemini-2.5-pro", "cl100k_base"),
|
||||
("mistral/mistral-large-latest", "cl100k_base"),
|
||||
("my-router-alias", "cl100k_base"),
|
||||
("azure/gpt-4o", "cl100k_base"),
|
||||
("azure/gpt-4o", "o200k_base"),
|
||||
("command-r-plus", "cl100k_base"),
|
||||
("gpt-4o", "o200k_base"),
|
||||
("gpt-4o-mini", "o200k_base"),
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits).
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -1039,3 +1040,34 @@ class TestContextSlotRetention:
|
|||
assert deployment is not None
|
||||
router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs)
|
||||
assert get_io_token_rate_limit_request_kwargs() is kwargs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_azure_deployment_name_reserves_itpm_on_its_base_models_token_count() -> None:
|
||||
async def reserved_input(deployment: dict[str, object]) -> object:
|
||||
metadata: Final[dict[str, object]] = {}
|
||||
request_kwargs: Final[dict[str, object]] = {
|
||||
"messages": [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}],
|
||||
"max_tokens": 5,
|
||||
"metadata": metadata,
|
||||
}
|
||||
set_io_token_rate_limit_request_kwargs(request_kwargs)
|
||||
await async_io_token_pre_call_check(DualCache(), deployment)
|
||||
return metadata[ITPM_RESERVED_KEY]
|
||||
|
||||
named = await reserved_input(
|
||||
{
|
||||
"litellm_params": {"model": "azure/my-gpt-6", "base_model": "azure/gpt-6.1-sol", "itpm": 100_000},
|
||||
"model_info": {"id": "named-deployment"},
|
||||
"model_name": "named",
|
||||
}
|
||||
)
|
||||
served = await reserved_input(
|
||||
{
|
||||
"litellm_params": {"model": "azure/gpt-6.1-sol", "itpm": 100_000},
|
||||
"model_info": {"id": "served-deployment"},
|
||||
"model_name": "served",
|
||||
}
|
||||
)
|
||||
|
||||
assert named == served
|
||||
|
|
|
|||
|
|
@ -7938,6 +7938,86 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch
|
|||
assert calls == [1]
|
||||
|
||||
|
||||
_MULTILINGUAL_PROMPT: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}]
|
||||
|
||||
|
||||
def _context_limited_router(*served: str) -> litellm.Router:
|
||||
limit: Final = (
|
||||
litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT)
|
||||
+ litellm.token_counter(model="prod", messages=_MULTILINGUAL_PROMPT)
|
||||
) // 2
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "prod",
|
||||
"litellm_params": {"model": model, "api_key": "sk-test"},
|
||||
"model_info": {"max_input_tokens": limit},
|
||||
}
|
||||
for model in served
|
||||
],
|
||||
enable_pre_call_checks=True,
|
||||
)
|
||||
|
||||
|
||||
def test_pre_call_checks_count_a_group_of_one_openai_model_with_that_models_tokenizer():
|
||||
router: Final = _context_limited_router("openai/gpt-6.1-sol")
|
||||
|
||||
kept: Final = router._pre_call_checks(
|
||||
model="prod", healthy_deployments=router.get_model_list(model_name="prod"), messages=_MULTILINGUAL_PROMPT
|
||||
)
|
||||
|
||||
assert len(kept) == 1
|
||||
|
||||
|
||||
def test_pre_call_checks_count_a_group_mixing_in_another_vendor_by_its_own_name():
|
||||
router: Final = _context_limited_router("openai/gpt-6.1-sol", "anthropic/claude-sonnet-5-5")
|
||||
|
||||
with pytest.raises(litellm.ContextWindowExceededError):
|
||||
router._pre_call_checks(
|
||||
model="prod", healthy_deployments=router.get_model_list(model_name="prod"), messages=_MULTILINGUAL_PROMPT
|
||||
)
|
||||
|
||||
|
||||
async def test_acount_pre_call_check_tokens_counts_a_group_of_one_openai_model_with_that_models_tokenizer():
|
||||
router: Final = _context_limited_router("openai/gpt-6.1-sol")
|
||||
|
||||
count: Final = await router._acount_pre_call_check_tokens(
|
||||
model="prod",
|
||||
healthy_deployments=router.get_model_list(model_name="prod"),
|
||||
messages=_MULTILINGUAL_PROMPT,
|
||||
input=None,
|
||||
request_kwargs=None,
|
||||
)
|
||||
|
||||
assert count == litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT)
|
||||
|
||||
|
||||
async def test_a_wildcard_route_checks_the_context_window_with_the_requested_openai_models_tokenizer():
|
||||
limit: Final = (
|
||||
litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT)
|
||||
+ litellm.token_counter(model="my-model", messages=_MULTILINGUAL_PROMPT)
|
||||
) // 2
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "openai/*",
|
||||
"litellm_params": {"model": "openai/*", "api_key": "sk-test"},
|
||||
"model_info": {"max_input_tokens": limit},
|
||||
}
|
||||
],
|
||||
enable_pre_call_checks=True,
|
||||
)
|
||||
|
||||
sync_response: Final = router.completion(
|
||||
model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT, mock_response="ok"
|
||||
)
|
||||
async_response: Final = await router.acompletion(
|
||||
model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT, mock_response="ok"
|
||||
)
|
||||
|
||||
assert (sync_response.choices[0].message.content, async_response.choices[0].message.content) == ("ok", "ok")
|
||||
|
||||
|
||||
def test_pre_call_checks_uses_precounted_tokens(monkeypatch):
|
||||
"""
|
||||
An async caller counts off the event loop and passes the result in. _pre_call_checks
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue