diff --git a/litellm/litellm_core_utils/dot_notation_indexing.py b/litellm/litellm_core_utils/dot_notation_indexing.py index f739e0238d7..2b305e3ff7d 100644 --- a/litellm/litellm_core_utils/dot_notation_indexing.py +++ b/litellm/litellm_core_utils/dot_notation_indexing.py @@ -102,8 +102,8 @@ def _parse_path_segments(path: str) -> list: import re # Match field names OR bracket expressions - # Pattern: field_name (anything except . or [) | [anything_in_brackets] - pattern: Final = r"[^\.\[]+|\[[^\]]*\]" + # Pattern: field_name (anything except . or [) | [anything_in_brackets except [ or ]] + pattern: Final = r"[^\.\[]+|\[[^\]\[]*\]" segments: Final = re.findall(pattern, path) return segments diff --git a/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py b/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py index 49ffae03c40..c7702d4fca1 100644 --- a/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py +++ b/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py @@ -5,13 +5,25 @@ from dataclasses import dataclass from typing import Final +@dataclass(frozen=True, slots=True) +class ToolChoiceTokens: + none: int + required: int + named: int + + +OPENAI_TOOL_CHOICE_TOKENS: Final = ToolChoiceTokens(none=1, required=0, named=7) +_XAI_GROK_TOOL_CHOICE_TOKENS: Final = ToolChoiceTokens(none=21, required=11, named=12) +_SERVED_BY_XAI: Final = re.compile(r"^(?:xai/|(?:openrouter/)?x-ai/)?[^/]+$", re.IGNORECASE) + + @dataclass(frozen=True, slots=True) class ChatTemplateToolSection: preamble_tokens: int implicit_system_turn_tokens: int render_tool_schemas: Callable[[Sequence[Mapping[str, object]]], str] - counts_tool_choice: bool = False renders_strict_only_when_true: bool = False + xai_tool_choice_tokens: ToolChoiceTokens | None = None def _json_line_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: @@ -43,13 +55,21 @@ _TOOL_SECTIONS: Final = ( ( re.compile(r"grok-4[-.]1(?![.\d])", re.IGNORECASE), ChatTemplateToolSection( - 0, 0, _compact_json_line_tool_schemas, counts_tool_choice=True, renders_strict_only_when_true=True + 0, + 0, + _compact_json_line_tool_schemas, + renders_strict_only_when_true=True, + xai_tool_choice_tokens=_XAI_GROK_TOOL_CHOICE_TOKENS, ), ), ( re.compile(r"grok-4[-.]3(?![.\d])", re.IGNORECASE), ChatTemplateToolSection( - 0, 0, _json_line_tool_schemas, counts_tool_choice=True, renders_strict_only_when_true=True + 0, + 0, + _json_line_tool_schemas, + renders_strict_only_when_true=True, + xai_tool_choice_tokens=_XAI_GROK_TOOL_CHOICE_TOKENS, ), ), ) @@ -59,3 +79,7 @@ def chat_template_tool_section(model: str | None) -> ChatTemplateToolSection | N if not model: return None return next((section for pattern, section in _TOOL_SECTIONS if pattern.search(model)), None) + + +def chat_template_tool_choice_tokens(model: str, section: ChatTemplateToolSection) -> ToolChoiceTokens | None: + return section.xai_tool_choice_tokens if _SERVED_BY_XAI.match(model) else None diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 11b02b9272b..5fb0a619744 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -50,7 +50,10 @@ from litellm.constants import ( from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping from litellm.litellm_core_utils.prompt_templates.chat_template_tools import ( + OPENAI_TOOL_CHOICE_TOKENS, ChatTemplateToolSection, + ToolChoiceTokens, + chat_template_tool_choice_tokens, chat_template_tool_section, ) from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding @@ -816,7 +819,7 @@ class _MessageCountParams: model: str, custom_tokenizer: dict | SelectTokenizerResponse | None, ): - from litellm.utils import print_verbose + from litellm.utils import huggingface_tokenizer_kind, print_verbose actual_model: Final = _fix_model_name(model) if uses_legacy_message_accounting(model): @@ -831,6 +834,12 @@ class _MessageCountParams: self.tokens_per_name = 1 self.count_function = _get_count_function(model, custom_tokenizer) self.tool_section = chat_template_tool_section(model) + self.tool_choice_tokens = _tool_choice_tokens( + model or "", + self.tool_section, + bills_no_tool_choice=huggingface_tokenizer_kind(model or "") == "deepseek" + or "gemini" in (model or "").lower(), + ) self._model = model @cached_property @@ -838,6 +847,16 @@ class _MessageCountParams: return pdf_page_pricing(self._model) +def _tool_choice_tokens( + model: str, tool_section: ChatTemplateToolSection | None, bills_no_tool_choice: bool +) -> ToolChoiceTokens | None: + if tool_section is not None: + return chat_template_tool_choice_tokens(model, tool_section) + if bills_no_tool_choice: + return None + return OPENAI_TOOL_CHOICE_TOKENS + + def _is_openai_chat_model(model: str) -> bool: return model in litellm.open_ai_chat_completion_models or model in litellm.azure_llms @@ -951,7 +970,12 @@ def token_counter( if count_response_tokens is False: includes_system_message: Final = any(message.get("role", None) == "system" for message in new_messages) num_tokens += _count_extra( - params.count_function, tools, tool_choice, includes_system_message, params.tool_section + params.count_function, + tools, + tool_choice, + includes_system_message, + params.tool_section, + params.tool_choice_tokens, ) else: @@ -987,7 +1011,9 @@ def messages_reach_token_count( ) running_totals: Final = accumulate( per_message_counts, - initial=_count_extra(params.count_function, tools, None, includes_system_message, params.tool_section), + initial=_count_extra( + params.count_function, tools, None, includes_system_message, params.tool_section, params.tool_choice_tokens + ), ) return any(total >= threshold for total in running_totals) @@ -1101,6 +1127,7 @@ def _count_extra( tool_choice: str | Mapping[str, object] | None, includes_system_message: bool, tool_section: ChatTemplateToolSection | None, + tool_choice_tokens: ToolChoiceTokens | None, ) -> int: """Count extra tokens for function definitions and tool choices. Args: @@ -1109,17 +1136,13 @@ def _count_extra( tool_choice (Optional[ChatCompletionNamedToolChoiceParam]): The tool choice. includes_system_message (bool): Whether the messages include a system message. tool_section (Optional[ChatTemplateToolSection]): How the model's chat template renders tools, if known. + tool_choice_tokens (Optional[ToolChoiceTokens]): What the model's chat template adds per tool choice, if any. """ reply_primer_tokens: Final = 3 - tool_choice_tokens: Final = ( - _count_tool_choice(count_function, tool_choice) - if tool_section is None or tool_section.counts_tool_choice - else 0 - ) return ( reply_primer_tokens + _count_tool_definitions(count_function, tools, includes_system_message, tool_section) - + tool_choice_tokens + + (_count_tool_choice(count_function, tool_choice, tool_choice_tokens) if tools else 0) ) @@ -1140,16 +1163,17 @@ def _count_tool_definitions( def _count_tool_choice( count_function: TokenCounterFunction, tool_choice: str | Mapping[str, object] | None, + tool_choice_tokens: ToolChoiceTokens | None, ) -> int: + if tool_choice_tokens is None: + return 0 match tool_choice: - case "none": - return 1 - case {"function": {"name": str() as name}}: - return 7 + count_function(name) - case {"type": "none"}: - return 1 - case {"name": str() as name}: - return 7 + count_function(name) + case "none" | {"type": "none"}: + return tool_choice_tokens.none + case "required" | {"type": "any"}: + return tool_choice_tokens.required + case {"function": {"name": str() as name}} | {"name": str() as name}: + return tool_choice_tokens.named + count_function(name) case _: return 0 diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index ddd83fd52c7..77394780b5e 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -5,7 +5,7 @@ This module handles the transformation of requests to Anthropic's CountTokens AP """ import copy -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass from functools import partial from itertools import chain @@ -18,7 +18,7 @@ from typing import ( cast, # noqa: TID251 # caller chat messages go to chat inference translation as-is ) -from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config +from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS @@ -48,6 +48,7 @@ if TYPE_CHECKING: from litellm.llms.anthropic.messages.native.transformation import AnthropicMessagesConfig _COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue]) +_COUNT_OPTION: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _COUNTED_MESSAGES: Final = TypeAdapter(list[dict[str, JsonValue]]) _TRANSFORMED_REQUEST: Final = TypeAdapter(dict[str, object]) _DROPPED_PATHS: Final = TypeAdapter(tuple[object, ...]) @@ -417,19 +418,26 @@ def _max_output_tokens(model: str, custom_llm_provider: str | None) -> int: ) +def _json_options(options: Mapping[str, object]) -> Iterator[tuple[str, JsonValue]]: + for name, value in options.items(): + try: + yield name, _COUNT_OPTION.validate_python(value) + except ValidationError: + continue + + def without_dropped_paths(options: Mapping[str, JsonValue], additional_drop_params: object) -> Mapping[str, JsonValue]: - return _COUNT_REQUEST.validate_python( - delete_nested_values( - options, - ( - path - for path in _DROPPED_PATHS.validate_python( - additional_drop_params if isinstance(additional_drop_params, (list, tuple)) else () - ) - if isinstance(path, str) - ), - ) + kept: Final = delete_nested_values( + options, + ( + path + for path in _DROPPED_PATHS.validate_python( + additional_drop_params if isinstance(additional_drop_params, (list, tuple)) else () + ) + if isinstance(path, str) + ), ) + return dict(_json_options(kept)) def messages_count_options( diff --git a/litellm/llms/bedrock/count_tokens/kimi_k3.py b/litellm/llms/bedrock/count_tokens/kimi_k3.py index 1d293bdc99c..0af3c89f6cf 100644 --- a/litellm/llms/bedrock/count_tokens/kimi_k3.py +++ b/litellm/llms/bedrock/count_tokens/kimi_k3.py @@ -246,6 +246,10 @@ class _ToolSpecBlock(_Converse): class _ToolConfig(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore") tools: tuple[_ToolSpecBlock | _CachePoint, ...] + tool_choice: Mapping[str, JsonValue] | None = Field(default=None, alias="toolChoice") + + def forces_a_tool(self) -> bool: + return self.tool_choice is not None and not self.tool_choice.keys().isdisjoint({"any", "tool"}) _Message: TypeAlias = Annotated[_UserMessage | _AssistantMessage, Field(discriminator="role")] @@ -303,6 +307,12 @@ def _message(attrs: Iterable[tuple[str, str]], body: Iterable[_Segment]) -> _Seg return (*_open("message", attrs), *body, *_close("message"), _Control.END_OF_MSG) +_FORCED_TOOL_CHOICE: Final = _message( + (("role", "system"), ("type", "tool-choice")), + ("The system is invoked with `tool_choice=required`.\nYou MUST call tools in the next message.",), +) + + def _tool_declare(specs: Sequence[_ToolSpec]) -> _Segments: tools: Final = json.dumps( [{"type": "function", "function": spec.declared()} for spec in specs], @@ -413,10 +423,12 @@ def _rendered(request: _ConverseRequest) -> _Segments: turns: Final = chain.from_iterable( _turn(previous, message) for previous, message in zip((None, *request.messages), request.messages, strict=False) ) + forces_a_tool: Final = request.tool_config is not None and request.tool_config.forces_a_tool() return ( *(_tool_declare(specs) if specs else ()), *(_message((("role", "system"),), (system,)) if system else ()), *turns, + *(_FORCED_TOOL_CHOICE if forces_a_tool else ()), *_open("message", (("role", "assistant"),)), *_open("think"), ) @@ -450,17 +462,29 @@ def _as_converse_input(tool: ChatCompletionToolParam) -> tuple[Mapping[str, obje def _converse_request( - model: str, messages: Sequence[AllMessageValues], tools: Sequence[ChatCompletionToolParam] | None + model: str, + messages: Sequence[AllMessageValues], + tools: Sequence[ChatCompletionToolParam] | None, + tool_choice: str | Mapping[str, object] | None, ) -> _ConverseRequest | None: if not all(_text_only(message) for message in messages): return None try: + config: Final = litellm.AmazonConverseConfig() chat_tools: Final = tuple(chain.from_iterable(_as_converse_input(tool) for tool in tools or ())) + choice: Final = ( + config.map_tool_choice_values(model, tool_choice, drop_params=True) + if chat_tools and tool_choice is not None + else None + ) return _ConverseRequest.model_validate( - litellm.AmazonConverseConfig().transform_request( + config.transform_request( model=model, messages=copy.deepcopy(list(messages)), - optional_params={"tools": copy.deepcopy(list(chat_tools))} if chat_tools else {}, + optional_params={ + **({"tools": copy.deepcopy(list(chat_tools))} if chat_tools else {}), + **({"tool_choice": choice} if choice is not None else {}), + }, litellm_params={}, headers={}, ) @@ -474,8 +498,9 @@ def kimi_k3_tokens( model: str, messages: Sequence[AllMessageValues], tools: Sequence[ChatCompletionToolParam] | None, + tool_choice: str | Mapping[str, object] | None = None, ) -> int | None: - request: Final = _converse_request(model, messages, tools) + request: Final = _converse_request(model, messages, tools, tool_choice) if request is None: return None return sum( @@ -493,6 +518,8 @@ def kimi_k3_anthropic_tokens(encoding: tiktoken.Encoding, model: str, fields: Ma except Exception: # noqa: BLE001 # a request the adapter cannot translate falls back to the estimate return None match translated: + case {"tools": tools, "tool_choice": tool_choice}: + return kimi_k3_tokens(encoding, model, translated["messages"], tools, tool_choice) case {"tools": tools}: return kimi_k3_tokens(encoding, model, translated["messages"], tools) case _: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 71d44b03f55..53601c1670f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -18,6 +18,9 @@ from pydantic import ( NonNegativeInt, PositiveInt, TypeAdapter, + ValidationError, + ValidatorFunctionWrapHandler, + WrapValidator, field_validator, model_validator, ) @@ -3617,6 +3620,13 @@ class LiteLLM_SpendLogs_ResponseObject(LiteLLMPydanticObjectBase): response: list[LiteLLM_SpendLogs | Any] | None = None +def _none_when_invalid(value: object, handler: ValidatorFunctionWrapHandler) -> object: + try: + return handler(value) + except ValidationError: + return None + + class TokenCountRequest(LiteLLMPydanticObjectBase): model: str prompt: str | None = None @@ -3631,10 +3641,11 @@ class TokenCountRequest(LiteLLMPydanticObjectBase): """ tools: list[dict] | None = None + tool_choice: Annotated[JsonValue, WrapValidator(_none_when_invalid)] = None + additional_drop_params: Annotated[JsonValue, WrapValidator(_none_when_invalid)] = None system: Any | None = None thinking: JsonValue = None reasoning_effort: JsonValue = None - tool_choice: JsonValue = None parallel_tool_calls: JsonValue = None response_format: JsonValue = None web_search_options: JsonValue = None diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 86092e4e27c..356eb0901c5 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -70,7 +70,7 @@ def _anthropic_error_json_response(exc: ProxyException, request: Request) -> JSO return anthropic_handler_error_json_response(exc, request) -_TOKEN_COUNT_FIELDS: Final = ("model", "messages", "tools", "system") +_TOKEN_COUNT_FIELDS: Final = ("model", "messages", "tools", "tool_choice", "system") def _token_count_request(fields: Mapping[str, object]) -> TokenCountRequest | ProxyException: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5c4b1f4f7e7..2ded84562fc 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -57,6 +57,7 @@ from typing import ( get_origin, get_type_hints, ) +from urllib.parse import urlparse import anyio import anyio.lowlevel @@ -108,12 +109,17 @@ from litellm.litellm_core_utils.bug_report import ( bug_report_notice, should_report_bug, ) +from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, initialize_sentry, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker -from litellm.litellm_core_utils.prompt_templates.chat_template_tools import chat_template_tool_section +from litellm.litellm_core_utils.prompt_templates.chat_template_tools import ( + ChatTemplateToolSection, + chat_template_tool_choice_tokens, + chat_template_tool_section, +) from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.routing_groups import parse_routing_groups from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -373,6 +379,7 @@ from litellm.constants import ( SPEND_CAPTURE_RATE_CHECK_JOB_ID, USER_SPEND_ALERTS_JOB_ID, WEEKLY_SPEND_REPORT_JOB_ID, + XAI_API_BASE, ) from litellm.exceptions import RejectedRequestError from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException @@ -15895,6 +15902,7 @@ async def _local_token_estimate( prompt: str | None, messages: Sequence[AllMessageValues] | None, tools: list[ChatCompletionToolParam] | None, + tool_choice: str | Mapping[str, object] | None, custom_tokenizer: CustomHuggingfaceTokenizer | None, provider_failure: ProviderCountFailure | None, ) -> _LocalTokenEstimate: @@ -15908,6 +15916,7 @@ async def _local_token_estimate( text=prompt, messages=messages, tools=tools, + tool_choice=tool_choice, custom_tokenizer=tokenizer, ) total_tokens: Final = await ( @@ -15927,17 +15936,19 @@ async def _kimi_k3_count( model: str, messages: Sequence[AllMessageValues], tools: Sequence[ChatCompletionToolParam] | None, + tool_choice: str | Mapping[str, object] | 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": + fields: Final = request.model_dump(include={"messages", "system", "tools"}, exclude_none=True) return await offload_token_count(kimi_k3_anthropic_tokens)( - encoding, model, request.model_dump(include={"messages", "system", "tools"}, exclude_none=True) + encoding, model, fields if tool_choice is None else {**fields, "tool_choice": tool_choice} ) case "openai": - return await offload_token_count(kimi_k3_tokens)(encoding, model, messages, tools) + return await offload_token_count(kimi_k3_tokens)(encoding, model, messages, tools, tool_choice) class _TokenCountDeployment(TypedDict, total=False): @@ -16180,15 +16191,59 @@ def _deployment_api_base( def _deployment_provider( deployment: _TokenCountDeployment, request_model: str, responses_request: Mapping[str, object] | None ) -> str | None: - params: Final = deployment.get("litellm_params", {}) + resolved: Final = _deployment_llm_provider(deployment, request_model, responses_request) + return None if resolved is None else resolved[1] + + +def _deployment_served_model( + deployment: _TokenCountDeployment, request_model: str, responses_request: Mapping[str, object] | None +) -> str | None: + resolved: Final = _deployment_llm_provider(deployment, request_model, responses_request) + if resolved is None: + return None + model, provider = resolved + served_by: Final = "xai" if _is_xai_api_base(_deployment_api_base(deployment, responses_request)) else provider + return f"{served_by}/{model}" + + +def _is_xai_api_base(api_base: str | None) -> bool: + if api_base is None: + return False try: - return litellm.get_llm_provider( - model=_non_empty_str(params.get("model")) or request_model, - custom_llm_provider=_non_empty_str(params.get("custom_llm_provider")), + return urlparse(api_base).hostname == urlparse(XAI_API_BASE).hostname + except ValueError: + return False + + +def _deployment_llm_provider( + deployment: _TokenCountDeployment, request_model: str, responses_request: Mapping[str, object] | None +) -> tuple[str, str] | None: + params: Final = deployment.get("litellm_params", {}) + model: Final = _non_empty_str(params.get("model")) or request_model + custom_llm_provider: Final = _non_empty_str(params.get("custom_llm_provider")) + declared: Final = declared_authenticating_provider(model, custom_llm_provider) + if declared is not None: + return model.removeprefix(f"{declared}/"), declared + try: + resolved_model, provider, _, _ = litellm.get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, api_base=_deployment_api_base(deployment, responses_request), - )[1] + ) except litellm.BadRequestError: return None + return resolved_model, provider + + +def _xai_tool_choice_off_xai( + template_model: str, tool_section: ChatTemplateToolSection | None, served_model: str | None +) -> bool: + return ( + tool_section is not None + and served_model is not None + and chat_template_tool_choice_tokens(template_model, tool_section) is not None + and chat_template_tool_choice_tokens(served_model, tool_section) is None + ) def _serves_native_responses( @@ -16253,7 +16308,16 @@ async def token_counter( """ call_id: Final = resolve_litellm_call_id(litellm_call_id) return await answered_with_litellm_call_id( - count_request_tokens(request=request, call_endpoint=call_endpoint, litellm_call_id=call_id), + count_request_tokens( + request=request, + call_endpoint=call_endpoint, + request_litellm_params=MappingProxyType( + {} + if request.additional_drop_params is None + else {"additional_drop_params": request.additional_drop_params} + ), + litellm_call_id=call_id, + ), call_id, http_request, fastapi_response, @@ -16578,6 +16642,17 @@ async def _count_request_tokens( openai_model: Final = deployment_openai_model(deployment_view) template_model: Final = openai_model or _deployment_base_model(deployment_view) or model_to_use tool_section: Final = chat_template_tool_section(template_model) + count_policy: Final = _count_drop_policy(deployment_view, {**request_litellm_params, **(responses_request or {})}) + counted_tool_choice: Final = ( + request.tool_choice + if counted_messages is not None + and isinstance(request.tool_choice, (str, Mapping)) + and "tool_choice" not in count_policy.additional_drop_params + and not _xai_tool_choice_off_xai( + template_model, tool_section, _deployment_served_model(deployment_view, request.model, responses_request) + ) + else None + ) counted_tools: Final = ( cast(list[ChatCompletionToolParam], list(_with_strict_by_default(sent_tools))) # cast-ok: strict added if sent_tools is not None @@ -16604,6 +16679,7 @@ async def _count_request_tokens( model_to_use, kimi_k3_messages, counted_tools, + counted_tool_choice, ) if kimi_k3_messages is not None and custom_tokenizer is None else None @@ -16622,6 +16698,7 @@ async def _count_request_tokens( prompt=prompt if counted_messages is None else None, messages=counted_messages, tools=counted_tools, + tool_choice=counted_tool_choice, custom_tokenizer=custom_tokenizer, provider_failure=ladder, ) @@ -16654,8 +16731,6 @@ async def supported_openai_params(model: str): --header "Authorization: Bearer $LITELLM_MASTER_KEY" ``` """ - from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider - global llm_router try: resolved_models: Final = ( diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 2963cf17b6c..2f0e3f2bcfb 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1518,6 +1518,7 @@ class _TokenCountPayload(TypedDict): model: ReadOnly[str] messages: ReadOnly[tuple[Mapping[str, object], ...]] tools: ReadOnly[object] + tool_choice: ReadOnly[object] class _TokenCounter(Protocol): @@ -1850,6 +1851,7 @@ async def _responses_input_tokens( "model": model_name, "messages": _responses_input_as_token_count_messages(input_value=input_value, request_params=data), "tools": token_fields.get("tools"), + "tool_choice": token_fields.get("tool_choice"), } token_request: Final = TokenCountRequest.model_validate(payload) except ValidationError as e: diff --git a/litellm/rust/crates/token-counter/src/counter.rs b/litellm/rust/crates/token-counter/src/counter.rs index f3c16f93b84..0bbe018bdb2 100644 --- a/litellm/rust/crates/token-counter/src/counter.rs +++ b/litellm/rust/crates/token-counter/src/counter.rs @@ -167,6 +167,7 @@ impl TokenCounter { definitions + TOOL_DEFINITIONS_TOKENS - discount }; let choice_tokens = match tool_choice { + _ if tools.is_empty() => 0, Some(ToolChoice::Mode(mode)) if mode == "none" => TOOL_CHOICE_NONE_TOKENS, Some(ToolChoice::Mode(_)) | None => 0, Some(ToolChoice::Named(named)) => { diff --git a/litellm/rust/crates/token-counter/tests/token_counter.rs b/litellm/rust/crates/token-counter/tests/token_counter.rs index 15b70d5a734..e7fcadf307d 100644 --- a/litellm/rust/crates/token-counter/tests/token_counter.rs +++ b/litellm/rust/crates/token-counter/tests/token_counter.rs @@ -122,17 +122,29 @@ mod json { count( r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"# ), - base + 1 + base ); assert_eq!( count( - r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"auto"}"# + r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":{"type":"function","function":{"name":"f"}}}"# ), base ); let with_tools = count( r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[{"name":"f"}]}"#, ); + assert_eq!( + count( + r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[{"name":"f"}],"tool_choice":"none"}"# + ), + with_tools + 1 + ); + assert_eq!( + count( + r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[{"name":"f"}],"tool_choice":"auto"}"# + ), + with_tools + ); let with_tools_and_system = count( r#"{"model":"m","messages":[{"role":"system","content":"hi"}],"tools":[{"name":"f"}]}"#, ); @@ -433,7 +445,7 @@ mod tiktoken { count( r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"# ), - base + 1 + base ); } diff --git a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts index af8aa12c8d4..4244073f72d 100644 --- a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -50864,6 +50864,7 @@ export interface components { }; /** TokenCountRequest */ TokenCountRequest: { + additional_drop_params?: components["schemas"]["JsonValue"]; /** Contents */ contents?: { [key: string]: unknown; diff --git a/litellm/utils.py b/litellm/utils.py index c015290726e..e84f933abec 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -241,7 +241,6 @@ from litellm.types.llms.openai import ( AllMessageValues, AllPromptValues, ChatCompletionAssistantToolCall, - ChatCompletionNamedToolChoiceParam, ChatCompletionToolParam, ChatCompletionToolParamFunctionChunk, OpenAITextCompletionUserMessage, @@ -3287,7 +3286,7 @@ def token_counter( messages: Sequence | None = None, count_response_tokens: bool | None = False, tools: list[ChatCompletionToolParam] | None = None, - tool_choice: ChatCompletionNamedToolChoiceParam | None = None, + tool_choice: str | Mapping[str, object] | None = None, use_default_image_token_count: bool | None = False, default_token_count: int | None = None, ) -> int: diff --git a/tests/benchmarks/test_count_tokens_benchmarks.py b/tests/benchmarks/test_count_tokens_benchmarks.py index fc1a7c615eb..3edcad5f31d 100644 --- a/tests/benchmarks/test_count_tokens_benchmarks.py +++ b/tests/benchmarks/test_count_tokens_benchmarks.py @@ -25,3 +25,10 @@ def test_without_dropped_paths_many_client_paths_over_a_large_schema() -> None: result: Final = without_dropped_paths(_LARGE_SCHEMA_OPTIONS, _MANY_DROP_PATHS) assert result == {"output_config": _LARGE_SCHEMA_OPTIONS["output_config"]} + + +@pytest.mark.benchmark +def test_without_dropped_paths_one_client_path_of_unclosed_brackets() -> None: + options: Final[dict[str, JsonValue]] = {"thinking": {"type": "enabled", "budget_tokens": 2048}} + + assert without_dropped_paths(options, ("[" * 16_000,)) == options diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index a6102a6e829..d80661ea8fd 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -705,26 +705,87 @@ def test_token_counter_counts_function_tools_like_chat_completions_tools_for_cha # Azure AI Foundry DeepSeek-V4-Flash-0731 billed 295 prompt tokens for the weather tool with each of -# these tool_choice values on 2026-10-17, the same as without one +# these tool_choice values on 2026-10-17, the same as without one. On 2026-10-29 aiand's +# deepseek-v4.1-flash billed 296 for each, and Fireworks' deepseek-v4p1-flash and OpenRouter's +# deepseek-v3.2 billed the same for "required" and the named tool as without one +@pytest.mark.parametrize( + "model", + [ + "DeepSeek-V4-Flash-0731", + "aiand/deepseek-ai/deepseek-v4.1-flash", + "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash", + "openrouter/deepseek/deepseek-v3.2", + ], +) @pytest.mark.parametrize("tool_choice", ["none", "required", {"type": "function", "function": {"name": "get_weather"}}]) -def test_token_counter_tool_choice_adds_nothing_for_deepseek_v4(tool_choice: object): +def test_token_counter_tool_choice_adds_nothing_for_deepseek_v3_and_v4(model: str, tool_choice: object): tools: Final = [_openai_tool(_WEATHER_FUNCTION)] assert token_counter_new( - model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools, tool_choice=tool_choice - ) == token_counter_new(model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools) + model=model, messages=list(_USER_MESSAGES), tools=tools, tool_choice=tool_choice + ) == token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) -# A request with a one-function tool billed 283 prompt tokens on xAI grok-4.3 without tool_choice and 297 with a -# named one, and 357 and 358 on Azure AI Foundry grok-4-1-fast-reasoning, on 2026-10-29 -@pytest.mark.parametrize("model", ["grok-4.3", "grok-4-1-fast-reasoning"]) -def test_token_counter_counts_a_named_tool_choice_for_grok(model: str): +# Gemini API gemini-3.8-flash, gemini-3.1-pro-preview and gemini-3.5-flash-lite each billed 51 prompt tokens for the +# weather tool with no toolConfig, NONE, ANY and ANY limited to the tool on 2026-11-11 +@pytest.mark.parametrize("model", ["gemini/gemini-3.8-flash", "gemini/gemini-3.1-pro-preview", "gemini-3.5-flash-lite"]) +@pytest.mark.parametrize("tool_choice", ["none", "required", {"type": "function", "function": {"name": "get_weather"}}]) +def test_token_counter_tool_choice_adds_nothing_for_gemini(model: str, tool_choice: object): tools: Final = [_openai_tool(_WEATHER_FUNCTION)] - named: Final = {"type": "function", "function": {"name": "get_weather"}} assert token_counter_new( - model=model, messages=list(_USER_MESSAGES), tools=tools, tool_choice=named - ) > token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) + model=model, messages=list(_USER_MESSAGES), tools=tools, tool_choice=tool_choice + ) == token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) + + +# xAI grok-4.3 and grok-4-1-fast-reasoning each billed 275 prompt tokens for the weather tool with no tool_choice, +# 296 with "none", 286 with "required" and 289 with the named tool on 2026-11-10, and OpenRouter, which served +# x-ai/grok-4.3 from xAI, billed the same on 2026-11-11 +@pytest.mark.parametrize("model", ["xai/grok-4.3", "grok-4-1-fast-reasoning", "openrouter/x-ai/grok-4.3"]) +@pytest.mark.parametrize( + ("tool_choice", "billed_over_unset"), + [ + pytest.param("none", 21, id="none"), + pytest.param({"type": "none"}, 21, id="anthropic-none"), + pytest.param("required", 11, id="required"), + pytest.param({"type": "any"}, 11, id="anthropic-any"), + pytest.param({"type": "function", "function": {"name": "get_weather"}}, 14, id="named"), + pytest.param({"type": "tool", "name": "get_weather"}, 14, id="anthropic-tool"), + pytest.param({"type": "function", "name": "get_weather"}, 14, id="responses-function"), + pytest.param("auto", 0, id="auto"), + ], +) +def test_token_counter_counts_each_grok_tool_choice_the_way_xai_bills_it( + model: str, tool_choice: str | Mapping[str, object], billed_over_unset: int +): + tools: Final = [_openai_tool(_WEATHER_FUNCTION)] + + assert ( + token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools, tool_choice=tool_choice) + - token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) + == billed_over_unset + ) + + +# Azure AI Foundry grok-4-1 billed 351 prompt tokens for the weather tool with no tool_choice, 163 with "none" and +# 352 with "required" or the named tool on 2026-11-10, so xAI's costs would count it 10 to 21 over +@pytest.mark.parametrize("model", ["azure_ai/grok-4-1-fast-reasoning", "azure_ai/grok-4.3", "vertex_ai/xai/grok-4.3"]) +@pytest.mark.parametrize("tool_choice", ["none", "required", {"type": "function", "function": {"name": "get_weather"}}]) +def test_token_counter_tool_choice_adds_nothing_for_grok_served_off_xai(model: str, tool_choice: object): + tools: Final = [_openai_tool(_WEATHER_FUNCTION)] + + assert token_counter_new( + model=model, messages=list(_USER_MESSAGES), tools=tools, tool_choice=tool_choice + ) == token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) + + +# xAI and OpenAI both answered 400 to a tool_choice sent without tools on 2026-11-11 +@pytest.mark.parametrize("model", ["xai/grok-4.3", "gpt-4.1"]) +@pytest.mark.parametrize("tool_choice", ["none", {"type": "function", "function": {"name": "get_weather"}}]) +def test_token_counter_tool_choice_without_tools_adds_nothing(model: str, tool_choice: object): + assert token_counter_new(model=model, messages=list(_USER_MESSAGES), tool_choice=tool_choice) == token_counter_new( + model=model, messages=list(_USER_MESSAGES) + ) def _counted_with_strict(model: str, **strict: bool) -> int: @@ -761,7 +822,7 @@ def test_token_counter_renders_a_false_strict_like_a_true_one_for_other_chat_tem ("tool_choice", "expected_tool_choice"), [ pytest.param({"type": "auto"}, None, id="anthropic-auto"), - pytest.param({"type": "any"}, None, id="anthropic-any"), + pytest.param({"type": "any"}, "required", id="anthropic-any"), pytest.param({"type": "none"}, "none", id="anthropic-none"), pytest.param( {"type": "tool", "name": "get_weather"}, diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py index 8cf69ba4cd3..432243c240f 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py @@ -413,6 +413,68 @@ def test_anthropic_request_counts_like_the_same_openai_request(pair_encoding: ti assert anthropic == openai +_FORCED_TOOL_CHOICE: Final = ( + '<|open|>message role="system" type="tool-choice"<|sep|>The system is invoked with `tool_choice=required`.\n' + "You MUST call tools in the next message.<|close|>message<|sep|><|end_of_msg|>" +) +_WEATHER_QUESTION: Final = [{"role": "user", "content": "Weather in Paris?"}] +_ANTHROPIC_WEATHER: Final = { + "name": "get_weather", + "description": "Get the weather in °C", + "input_schema": _WEATHER["function"]["parameters"], +} + + +# Bedrock billed us.moonshotai.kimi-k3 and kimi-k3.1 206 input tokens for an any or named tool choice and 170 for an +# auto or unset one on 2026-11-12, and the template above renders this message for tool_choice=required only +@pytest.mark.parametrize( + ("tools", "tool_choice", "added"), + [ + pytest.param([_WEATHER], "required", _FORCED_TOOL_CHOICE, id="required"), + pytest.param( + [_WEATHER], {"type": "function", "function": {"name": "get_weather"}}, _FORCED_TOOL_CHOICE, id="named" + ), + pytest.param([_WEATHER], "auto", "", id="auto"), + pytest.param([_WEATHER], "none", "", id="none_dropped_by_bedrock"), + pytest.param(None, "required", "", id="no_tools"), + ], +) +def test_tool_choice_adds_the_template_message_bedrock_bills_only_for_a_forced_tool( + pair_encoding: tiktoken.Encoding, + tools: list[dict[str, object]] | None, + tool_choice: str | dict[str, object], + added: str, +): + unset: Final = kimi_k3_tokens(pair_encoding, MODEL, _WEATHER_QUESTION, tools) + + assert unset is not None + assert kimi_k3_tokens(pair_encoding, MODEL, _WEATHER_QUESTION, tools, tool_choice) == unset + _template_count( + pair_encoding, added + ) + + +@pytest.mark.parametrize( + ("anthropic_choice", "openai_choice"), + [ + pytest.param({"type": "any"}, "required", id="any"), + pytest.param({"type": "tool", "name": "get_weather"}, "required", id="tool"), + pytest.param({"type": "auto"}, None, id="auto"), + pytest.param({"type": "none"}, None, id="none"), + ], +) +def test_anthropic_tool_choice_counts_like_the_openai_tool_choice_it_translates_to( + pair_encoding: tiktoken.Encoding, anthropic_choice: dict[str, object], openai_choice: str | None +): + anthropic: Final = kimi_k3_anthropic_tokens( + pair_encoding, + MODEL, + {"messages": _WEATHER_QUESTION, "tools": [_ANTHROPIC_WEATHER], "tool_choice": anthropic_choice}, + ) + + assert anthropic is not None + assert anthropic == kimi_k3_tokens(pair_encoding, MODEL, _WEATHER_QUESTION, [_WEATHER], openai_choice) + + @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 diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index d9f523b951b..0d2c5b06db0 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -44,7 +44,7 @@ from starlette.types import Message import litellm from litellm import Router -from litellm._logging import JsonFormatter, verbose_proxy_logger +from litellm._logging import JsonFormatter, verbose_logger, verbose_proxy_logger from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler from litellm.llms.azure_ai.anthropic.count_tokens import handler as azure_ai_count_tokens_http from litellm.llms.base_llm.base_utils import BaseTokenCounter @@ -83,6 +83,7 @@ from litellm.proxy.proxy_server import ( _provider_token_count, _try_provider_token_count, count_request_tokens, + token_counter, ) from litellm.responses.litellm_completion_transformation.session_handler import ResponsesSessionHandler from litellm.responses.utils import ResponsesAPIRequestUtils @@ -3454,6 +3455,78 @@ async def test_bedrock_kimi_k3_is_counted_with_its_tokenizer_plus_the_shipped_fr assert (anthropic.total_tokens, anthropic.tokenizer_type) == (openai.total_tokens, "huggingface_tokenizer") +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("messages_format", "tools", "forced", "unforced"), + [ + pytest.param( + "openai", [_KIMI_WEATHER], {"type": "function", "function": {"name": "get_weather"}}, "auto", id="openai" + ), + pytest.param( + "anthropic", + [ + { + "name": "get_weather", + "description": "Get the weather", + "input_schema": _KIMI_WEATHER["function"]["parameters"], + } + ], + {"type": "tool", "name": "get_weather"}, + {"type": "auto"}, + id="anthropic", + ), + ], +) +@pytest.mark.parametrize("drops_tool_choice", [False, True]) +async def test_bedrock_kimi_k3_count_adds_the_template_message_of_a_forced_tool_choice_inference_sends( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + messages_format: TokenCountMessagesFormat, + tools: list[dict[str, object]], + forced: dict[str, object], + unforced: str | dict[str, object], + drops_tool_choice: bool, +): + _framed_kimi(monkeypatch) + if drops_tool_choice: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "kimi", + "litellm_params": { + "model": "bedrock/us.moonshotai.kimi-k3", + "additional_drop_params": ["tool_choice"], + }, + } + ] + ), + ) + encoding: Final = _available_kimi_k3(monkeypatch, tmp_path) + question: Final = [{"role": "user", "content": "Weather in Paris?"}] + + async def count(tool_choice: str | dict[str, object] | None) -> int: + counted: Final = await count_request_tokens( + request=TokenCountRequest(model="kimi", messages=question, tools=tools, tool_choice=tool_choice), + call_endpoint=True, + messages_format=messages_format, + ) + assert counted.tokenizer_type == "huggingface_tokenizer" + return counted.total_tokens + + plain: Final = kimi_k3_tokens(encoding, "us.moonshotai.kimi-k3", question, [_KIMI_WEATHER]) + required: Final = kimi_k3_tokens(encoding, "us.moonshotai.kimi-k3", question, [_KIMI_WEATHER], "required") + unset: Final = await count(None) + + assert plain is not None and required is not None and required > plain + assert (await count(unforced), await count(forced)) == ( + unset, + unset if drops_tool_choice else unset + required - plain, + ) + + @pytest.mark.asyncio async def test_bedrock_kimi_k3_deployment_with_a_custom_tokenizer_keeps_counting_with_it( monkeypatch: pytest.MonkeyPatch, tmp_path: Path @@ -6361,3 +6434,330 @@ async def test_messages_count_tokens_drops_what_the_deployment_drops_unless_the_ sent: Final = json.loads(route.calls.last.request.content) assert ("tool_choice" in sent, (sent.get("thinking") or {}).get("type")) == sent_options + + +_TOOL_CHOICE_MODEL: Final = "qwen/qwen3.8-27b" +_TOOL_CHOICE_QUESTION: Final = "What is the weather in Paris?" +_CITY_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], +} +_CHAT_WEATHER_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": "get_weather", "description": "Get the weather for a city", "parameters": _CITY_SCHEMA}, +} + + +def _tool_choice_router() -> Router: + return Router( + model_list=[ + { + "model_name": "no-count-api", + "litellm_params": {"model": f"groq/{_TOOL_CHOICE_MODEL}", "api_key": "sk-test"}, + } + ] + ) + + +async def _utils_route_count(tool_choice: Mapping[str, JsonValue], model: str = "no-count-api") -> int: + response: Final = await token_counter( + request=TokenCountRequest( + model=model, + messages=[{"role": "user", "content": _TOOL_CHOICE_QUESTION}], + tools=[_CHAT_WEATHER_TOOL], + **tool_choice, + ), + http_request=_json_request({}, "python-httpx/0.28.1"), + fastapi_response=Response(), + ) + return response.total_tokens + + +async def _messages_route_count(tool_choice: Mapping[str, JsonValue]) -> int: + body: Final = { + "model": "no-count-api", + "messages": [{"role": "user", "content": _TOOL_CHOICE_QUESTION}], + "tools": [{"name": "get_weather", "description": "Get the weather for a city", "input_schema": _CITY_SCHEMA}], + **tool_choice, + } + response: Final = await anthropic_count_tokens( + _json_request(body, "python-httpx/0.28.1"), Response(), UserAPIKeyAuth() + ) + assert isinstance(response, dict) + return response["input_tokens"] + + +async def _responses_route_count(tool_choice: Mapping[str, JsonValue]) -> int: + from litellm.proxy.response_api_endpoints.endpoints import responses_input_tokens + + body: Final = { + "model": "no-count-api", + "input": _TOOL_CHOICE_QUESTION, + "tools": [ + { + "type": "function", + "name": "get_weather", + "description": "Get the weather for a city", + "parameters": _CITY_SCHEMA, + } + ], + **tool_choice, + } + response: Final = await responses_input_tokens( + request=_json_request(body, "python-httpx/0.28.1"), + fastapi_response=Response(), + general_settings=MappingProxyType({}), + token_counter=count_request_tokens, + llm_router=litellm.proxy.proxy_server.llm_router, + user_api_key_dict=UserAPIKeyAuth(), + response_id_security=None, + managed_files=None, + ) + assert isinstance(response, dict) + return response["input_tokens"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route_count", "tool_choice"), + [ + (_utils_route_count, {"type": "function", "function": {"name": "get_weather"}}), + (_utils_route_count, "none"), + (_messages_route_count, {"type": "tool", "name": "get_weather"}), + (_messages_route_count, {"type": "none"}), + (_responses_route_count, {"type": "function", "name": "get_weather"}), + (_responses_route_count, "none"), + ], + ids=["utils-named", "utils-none", "messages-named", "messages-none", "responses-named", "responses-none"], +) +async def test_count_routes_without_a_count_api_count_the_tool_choice_like_token_counter( + monkeypatch: pytest.MonkeyPatch, + route_count: Callable[[Mapping[str, JsonValue]], Coroutine[None, None, int]], + tool_choice: JsonValue, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _tool_choice_router()) + monkeypatch.setattr(litellm.proxy.proxy_server, "proxy_config", litellm.proxy.proxy_server.ProxyConfig()) + question: Final = [{"role": "user", "content": _TOOL_CHOICE_QUESTION}] + sdk_added: Final = litellm.token_counter( + model=_TOOL_CHOICE_MODEL, messages=question, tools=[_CHAT_WEATHER_TOOL], tool_choice=tool_choice + ) - litellm.token_counter(model=_TOOL_CHOICE_MODEL, messages=question, tools=[_CHAT_WEATHER_TOOL]) + + route_added: Final = await route_count({"tool_choice": tool_choice}) - await route_count({}) + + assert route_added == sdk_added > 0 + + +_AZURE_GROK_PARAMS: Final = {"api_base": "https://example.services.ai.azure.com", "api_key": "sk-test"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment", "served_by_xai"), + [ + ({"litellm_params": {"model": "xai/grok-4-1-fast-reasoning", "api_key": "sk-test"}}, True), + ({"litellm_params": {"model": "openrouter/x-ai/grok-4.3", "api_key": "sk-test"}}, True), + ( + {"litellm_params": {"model": "openai/grok-4.3", "api_base": "https://api.x.ai/v1", "api_key": "sk-test"}}, + True, + ), + ({"litellm_params": {"model": "openai/grok-4.3", "api_base": "http://[bad", "api_key": "sk-test"}}, False), + ( + { + "litellm_params": {"model": "azure_ai/grok-4-1", **_AZURE_GROK_PARAMS}, + "model_info": {"base_model": "azure_ai/grok-4-1-fast-reasoning"}, + }, + False, + ), + ({"litellm_params": {"model": "azure_ai/grok-4-1-fast-reasoning", **_AZURE_GROK_PARAMS}}, False), + ({"litellm_params": {"model": "grok-4.3", "custom_llm_provider": "azure_ai", **_AZURE_GROK_PARAMS}}, False), + ( + { + "litellm_params": { + "model": "vertex_ai/xai/grok-4.3", + "vertex_project": "test-project", + "vertex_location": "global", + } + }, + False, + ), + ], +) +async def test_utils_token_counter_counts_xai_tool_choice_costs_only_on_grok_that_xai_serves( + monkeypatch: pytest.MonkeyPatch, deployment: dict[str, dict[str, str]], served_by_xai: bool +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, "llm_router", Router(model_list=[{"model_name": "grok", **deployment}]) + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "proxy_config", litellm.proxy.proxy_server.ProxyConfig()) + named: Final = {"tool_choice": {"type": "function", "function": {"name": "get_weather"}}} + question: Final = [{"role": "user", "content": _TOOL_CHOICE_QUESTION}] + sdk_added: Final = litellm.token_counter( + model="xai/grok-4.3", messages=question, tools=[_CHAT_WEATHER_TOOL], **named + ) - litellm.token_counter(model="xai/grok-4.3", messages=question, tools=[_CHAT_WEATHER_TOOL]) + + route_added: Final = await _utils_route_count(named, "grok") - await _utils_route_count({}, "grok") + + assert route_added == (sdk_added if served_by_xai else 0) + assert sdk_added > 0 + + +def _litellm_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple(record.getMessage() for record in caplog.records if record.name == verbose_logger.name) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route_count", "tool_choice"), + [(_utils_route_count, "none"), (_messages_route_count, {"type": "none"}), (_responses_route_count, "none")], + ids=["utils", "messages", "responses"], +) +@pytest.mark.parametrize("model", ["chatgpt/gpt-6.1-sol", "github_copilot/kimi-k3"]) +async def test_count_routes_leave_the_provider_login_alone_when_counting_a_tool_choice( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + caplog: pytest.LogCaptureFixture, + route_count: Callable[[Mapping[str, JsonValue]], Coroutine[None, None, int]], + tool_choice: JsonValue, + model: str, +) -> None: + (tmp_path / "auth.json").write_text("not a saved login") + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "no-count-api", "litellm_params": {"model": model}}]), + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "proxy_config", litellm.proxy.proxy_server.ProxyConfig()) + + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + without_choice: Final = await route_count({}) + warnings_without_choice: Final = _litellm_warnings(caplog) + caplog.clear() + with_choice: Final = await route_count({"tool_choice": tool_choice}) + + assert with_choice > without_choice + assert _litellm_warnings(caplog) == warnings_without_choice + + +def _nested_json(depth: int) -> JsonValue: + return "none" if depth == 0 else {"nested": _nested_json(depth - 1)} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route_count", "field"), + [ + (_utils_route_count, "tool_choice"), + (_utils_route_count, "additional_drop_params"), + (_messages_route_count, "tool_choice"), + (_messages_route_count, "thinking"), + (_responses_route_count, "tool_choice"), + ], + ids=["utils-tool-choice", "utils-drop-params", "messages-tool-choice", "messages-thinking", "responses-tool-choice"], +) +async def test_count_routes_skip_a_value_nested_too_deep_to_validate( + monkeypatch: pytest.MonkeyPatch, + route_count: Callable[[Mapping[str, JsonValue]], Coroutine[None, None, int]], + field: str, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _tool_choice_router()) + monkeypatch.setattr(litellm.proxy.proxy_server, "proxy_config", litellm.proxy.proxy_server.ProxyConfig()) + + assert await route_count({field: _nested_json(300)}) == await route_count({}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route_count", "tool_choice"), + [(_utils_route_count, "none"), (_messages_route_count, {"type": "none"}), (_responses_route_count, "none")], + ids=["utils", "messages", "responses"], +) +@pytest.mark.parametrize( + ("drop_params", "counts_choice"), + [({"additional_drop_params": ["tool_choice"]}, False), ({"drop_params": True}, True)], + ids=["tool-choice-dropped", "only-unsupported-dropped"], +) +async def test_count_routes_count_no_tool_choice_the_deployment_drops( + monkeypatch: pytest.MonkeyPatch, + route_count: Callable[[Mapping[str, JsonValue]], Coroutine[None, None, int]], + tool_choice: JsonValue, + drop_params: dict[str, JsonValue], + counts_choice: bool, +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "no-count-api", + "litellm_params": {"model": f"groq/{_TOOL_CHOICE_MODEL}", "api_key": "sk-test", **drop_params}, + } + ] + ), + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "proxy_config", litellm.proxy.proxy_server.ProxyConfig()) + + route_added: Final = await route_count({"tool_choice": tool_choice}) - await route_count({}) + + assert (route_added > 0) == counts_choice + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route_count", "tool_choice"), + [(_utils_route_count, "none"), (_messages_route_count, {"type": "none"}), (_responses_route_count, "none")], + ids=["utils", "messages", "responses"], +) +@pytest.mark.parametrize( + ("dropped", "counts_choice"), + [(["tool_choice"], False), (["temperature"], True)], + ids=["tool-choice-dropped", "other-param-dropped"], +) +async def test_count_routes_count_no_tool_choice_the_request_drops( + monkeypatch: pytest.MonkeyPatch, + route_count: Callable[[Mapping[str, JsonValue]], Coroutine[None, None, int]], + tool_choice: JsonValue, + dropped: list[JsonValue], + counts_choice: bool, +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "no-count-api", + "litellm_params": {"model": f"groq/{_TOOL_CHOICE_MODEL}", "api_key": "sk-test"}, + } + ] + ), + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "proxy_config", litellm.proxy.proxy_server.ProxyConfig()) + drop_list: Final = {"additional_drop_params": dropped} + + route_added: Final = await route_count({"tool_choice": tool_choice, **drop_list}) - await route_count(drop_list) + + assert (route_added > 0) == counts_choice + + +@pytest.mark.asyncio +async def test_utils_token_counter_counts_a_bare_prompt_with_a_tool_choice_like_the_prompt_alone( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _tool_choice_router()) + + with_choice: Final = await token_counter( + request=TokenCountRequest(model="no-count-api", prompt=_TOOL_CHOICE_QUESTION, tool_choice="required"), + http_request=_json_request({}, "python-httpx/0.28.1"), + fastapi_response=Response(), + ) + prompt_alone: Final = await token_counter( + request=TokenCountRequest(model="no-count-api", prompt=_TOOL_CHOICE_QUESTION), + http_request=_json_request({}, "python-httpx/0.28.1"), + fastapi_response=Response(), + ) + + assert (with_choice.error, with_choice.total_tokens) == (False, prompt_alone.total_tokens) diff --git a/tests/unit/rust_bridge/test_token_counter.py b/tests/unit/rust_bridge/test_token_counter.py index ba9c0b90498..62ac68f75f8 100644 --- a/tests/unit/rust_bridge/test_token_counter.py +++ b/tests/unit/rust_bridge/test_token_counter.py @@ -240,6 +240,7 @@ def test_disabled_token_counter_declines_every_model(monkeypatch: pytest.MonkeyP PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = ( {"model": MODEL, "messages": [{"role": "user", "content": "Hello, how are you today?"}]}, + {"model": MODEL, "messages": [{"role": "user", "content": "Hello, how are you today?"}], "tool_choice": "none"}, { "model": MODEL, "messages": [