diff --git a/litellm/litellm_core_utils/dot_notation_indexing.py b/litellm/litellm_core_utils/dot_notation_indexing.py index 818cf1b9695..f739e0238d7 100644 --- a/litellm/litellm_core_utils/dot_notation_indexing.py +++ b/litellm/litellm_core_utils/dot_notation_indexing.py @@ -23,7 +23,8 @@ Used by JWT Auth to get the user role from the token, and by additional_drop_params to remove nested fields from optional parameters. """ -from collections.abc import Mapping +import copy +from collections.abc import Iterable, Mapping from typing import Any, Final, TypeVar T = TypeVar("T") @@ -212,8 +213,6 @@ def delete_nested_value( >>> delete_nested_value(data, "tools[*].input_examples") {"tools": [{"name": "t1"}]} """ - import copy - result: Final = copy.deepcopy(data if isinstance(data, dict) else dict(data)) try: @@ -233,6 +232,13 @@ def delete_nested_value( return result +def delete_nested_values(data: Mapping[str, object], paths: Iterable[str]) -> Mapping[str, object]: + result: Final = copy.deepcopy(dict(data)) + for path in dict.fromkeys(paths): + _delete_nested_value_custom(result, _parse_path_segments(path), 0) + return result + + def is_nested_path(path: str) -> bool: """ Check if path requires nested handling. diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index e32754adf03..b753fe03b5a 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -757,7 +757,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): case _: return False - def thinking_betas(self, thinking: AnthropicThinkingParam | None) -> tuple[str, ...]: + def thinking_betas(self, thinking: object) -> tuple[str, ...]: return ( *( (ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER,) diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index 46c90295198..10f15326186 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -83,7 +83,9 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): verbose_logger.debug("Making request to: %s", endpoint_url) # Get required headers - headers: Final = self.get_count_tokens_headers(auth_header, messages, context_management) + headers: Final = self.get_count_tokens_headers( + auth_header, messages, context_management, request_body.get("thinking") + ) # Use LiteLLM's async httpx client async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.ANTHROPIC) diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 755ccd129d8..aef687cc9ce 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -2,7 +2,7 @@ Anthropic Token Counter implementation using the CountTokens API. """ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import Any, Final @@ -11,14 +11,21 @@ from pydantic import JsonValue import litellm from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountFields, ChatCountTranslationError, anthropic_count_fields, + messages_count_options, +) +from litellm.llms.base_llm.base_utils import ( + BaseTokenCounter, + deployment_litellm_params, + deployment_timeout, + optional_str, ) -from litellm.llms.base_llm.base_utils import BaseTokenCounter, deployment_timeout from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests @@ -65,42 +72,71 @@ class AnthropicTokenCounter(BaseTokenCounter): Returns: TokenCountResponse with token count, or None if counting fails """ - from litellm.llms.anthropic.common_utils import AnthropicError, AnthropicModelInfo - if not messages: return None from litellm.llms.anthropic.chat.transformation import AnthropicConfig + from litellm.llms.anthropic.messages.native.transformation import AnthropicMessagesConfig + litellm_params: Final = deployment_litellm_params(deployment) match translated := await anthropic_count_fields( model_to_use, messages, system, tools, messages_format, AnthropicConfig(), chat_params ): case ChatCountTranslationError(): return translated.response(request_model, model_to_use, "anthropic_api") case AnthropicCountFields(): - return await self.count_tokens( + return await self._count( model_to_use=model_to_use, - messages=[dict(message) for message in translated.messages], - contents=contents, - deployment=deployment, + messages=translated.messages, + litellm_params=litellm_params, request_model=request_model, - tools=[dict(tool) for tool in translated.tools] if translated.tools is not None else None, + tools=translated.tools, system=translated.system, - messages_format="anthropic", - output_format=output_format, - context_management=context_management, count_options=translated.options, + context_management=context_management, ) case None: pass - deployment = deployment or {} - litellm_params: Final = deployment.get("litellm_params", {}) - api_base: Final = litellm_params.get("api_base") + match options := await asyncify(messages_count_options)( + AnthropicMessagesConfig(), + model_to_use, + messages, + count_options, + litellm_params, + ): + case ChatCountTranslationError(): + return options.response(request_model, model_to_use, "anthropic_api") + case _: + return await self._count( + model_to_use=model_to_use, + messages=messages, + litellm_params=litellm_params, + request_model=request_model, + tools=tools, + system=system, + count_options=options, + context_management=context_management, + ) + + async def _count( + self, + model_to_use: str, + messages: Sequence[Mapping[str, JsonValue]], + litellm_params: Mapping[str, object], + request_model: str, + tools: Sequence[Mapping[str, JsonValue]] | None, + system: JsonValue, + count_options: Mapping[str, JsonValue], + context_management: CountContextManagement | None, + ) -> TokenCountResponse | None: + from litellm.llms.anthropic.common_utils import AnthropicError, AnthropicModelInfo + + api_base: Final = optional_str(litellm_params.get("api_base")) try: configured_auth_header: Final = await AnthropicModelInfo.aget_auth_header( - api_key=litellm_params.get("api_key") or litellm.anthropic_key or litellm.api_key, + api_key=optional_str(litellm_params.get("api_key")) or litellm.anthropic_key or litellm.api_key, api_base=api_base, litellm_params=litellm_params, allow_workload_identity=True, @@ -126,10 +162,10 @@ class AnthropicTokenCounter(BaseTokenCounter): result: Final = await anthropic_count_tokens_handler.handle_count_tokens_request( model=model_to_use, - messages=messages, + messages=[dict(message) for message in messages], auth_header=auth_header, api_base=api_base, - tools=tools, + tools=None if tools is None else [dict(tool) for tool in tools], system=system, optional_params=count_options, context_management=context_management.spec if context_management is not None else None, diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index 3034927c612..f8b3f6f6c36 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -10,6 +10,7 @@ from dataclasses import dataclass from itertools import chain from types import MappingProxyType from typing import ( + TYPE_CHECKING, Final, Protocol, TypeAlias, @@ -19,14 +20,16 @@ from typing import ( from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict -from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION +from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_values from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( anthropic_system_blocks, split_leading_system_run, ) from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.common_utils import ( + AnthropicModelInfo, carries_signed_compaction_block, context_management_betas, context_management_with_edits, @@ -37,9 +40,16 @@ from litellm.llms.anthropic.common_utils import ( from litellm.llms.anthropic.wif import resolve_anthropic_base from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES from litellm.types.llms.openai import AllMessageValues, ChatCompletionImageObject, ValidUserMessageContentTypes +from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse +if TYPE_CHECKING: + from litellm.llms.anthropic.messages.native.transformation import AnthropicMessagesConfig + _COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue]) +_COUNTED_MESSAGES: Final = TypeAdapter(list[dict[str, JsonValue]]) +_TRANSFORMED_REQUEST: Final = TypeAdapter(dict[str, object]) +_DROPPED_PATHS: Final = TypeAdapter(tuple[object, ...]) _SYSTEM_BLOCKS: Final = TypeAdapter(list[JsonValue]) _IMAGE_BLOCK: Final = TypeAdapter(ChatCompletionImageObject) COUNT_TOKEN_OPTION_NAMES: Final = ("thinking", "tool_choice", "output_config") @@ -336,6 +346,75 @@ def anthropic_fields_for_chat_tool_use( ) +def _max_output_tokens(model: str, custom_llm_provider: str | None) -> int: + from litellm.utils import get_model_info + + try: + max_output_tokens: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider).get( + "max_output_tokens" + ) + except Exception: # noqa: BLE001 # get_model_info raises a plain Exception for a model the cost map lacks + return DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + return ( + max_output_tokens + if isinstance(max_output_tokens, int) and max_output_tokens > 0 + else DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + ) + + +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) + ), + ) + ) + + +def messages_count_options( + translator: "AnthropicMessagesConfig", + model: str, + messages: Sequence[Mapping[str, JsonValue]], + options: Mapping[str, JsonValue], + litellm_params: Mapping[str, object], +) -> Mapping[str, JsonValue] | ChatCountTranslationError: + """The thinking, tool_choice, and output_config the deployment's /v1/messages transform sends upstream + for these Anthropic Messages body options. The transform needs max_tokens, which a count request lacks, + so it gets the model's max output tokens: a thinking budget's size never changes a count, only whether + thinking stays on.""" + if not options: + return MappingProxyType({}) + from litellm.exceptions import BadRequestError + from litellm.llms.anthropic.common_utils import AnthropicError + + kept_options: Final = without_dropped_paths(options, litellm_params.get("additional_drop_params")) + thinking_betas: Final = AnthropicModelInfo().thinking_betas(kept_options.get("thinking")) + try: + request: Final = _TRANSFORMED_REQUEST.validate_python( + translator.transform_anthropic_messages_request( + model=model, + messages=_COUNTED_MESSAGES.validate_python(messages), + anthropic_messages_optional_request_params={ + **kept_options, + "max_tokens": _max_output_tokens(model, translator.custom_llm_provider), + }, + litellm_params=GenericLiteLLMParams.model_validate({"drop_params": litellm_params.get("drop_params")}), + headers={"anthropic-beta": ",".join(thinking_betas)} if thinking_betas else {}, + ) + ) + except (AnthropicError, BadRequestError) as e: + return ChatCountTranslationError(message=str(e.message)) + return _COUNT_REQUEST.validate_python( + {name: request[name] for name in COUNT_TOKEN_OPTION_NAMES if request.get(name) is not None} + ) + + async def anthropic_count_fields( model: str, messages: Sequence[Mapping[str, object]], @@ -433,6 +512,7 @@ class AnthropicCountTokensConfig: auth_header: Mapping[str, str], messages: Sequence[Mapping[str, JsonValue]], context_management: JsonValue = None, + thinking: JsonValue = None, ) -> dict[str, str]: """The count-tokens headers around a resolved Anthropic auth header (``AnthropicModelInfo.get_auth_header``): x-api-key for a static key, an Authorization @@ -449,6 +529,7 @@ class AnthropicCountTokensConfig: ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, *compaction_betas, *context_management_betas(context_management), + *AnthropicModelInfo().thinking_betas(thinking), ) return { "Content-Type": "application/json", diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 443a3a5edad..4d89c299bc5 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -15,7 +15,7 @@ from litellm._logging import verbose_logger from litellm.anthropic_beta_headers_manager import aload_beta_headers_config from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import drops_unsupported_params -from litellm.llms.anthropic.common_utils import AnthropicError +from litellm.llms.anthropic.common_utils import AnthropicError, AnthropicModelInfo from litellm.llms.azure.common_utils import EntraRejectionError from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( AzureAIAnthropicCountTokensConfig, @@ -86,7 +86,8 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): verbose_logger.debug("Making request to: %s", endpoint_url) - if context_management is not None: + thinking: Final = request_body.get("thinking") + if context_management is not None or AnthropicModelInfo().thinking_betas(thinking): await aload_beta_headers_config() # Get required headers with Azure authentication headers: Final = self.get_required_headers( @@ -94,6 +95,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): litellm_params=litellm_params, messages=messages, context_management=context_management, + thinking=thinking, ) if not {"x-api-key", "Authorization"} & headers.keys(): raise AnthropicError( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index 5031b69f83b..b25105c2236 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -3,7 +3,7 @@ Azure AI Anthropic Token Counter implementation using the CountTokens API. """ import os -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import Any, Final @@ -11,16 +11,23 @@ from pydantic import JsonValue import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountFields, ChatCountTranslationError, anthropic_count_fields, + messages_count_options, ) from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) -from litellm.llms.base_llm.base_utils import BaseTokenCounter, deployment_timeout +from litellm.llms.base_llm.base_utils import ( + BaseTokenCounter, + deployment_litellm_params, + deployment_timeout, + optional_str, +) from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests @@ -79,45 +86,69 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): Returns: TokenCountResponse with token count, or None if counting fails """ - from litellm.llms.anthropic.common_utils import AnthropicError - if not messages: return None + from litellm.llms.azure_ai.anthropic.messages_transformation import AzureAnthropicMessagesConfig from litellm.llms.azure_ai.anthropic.transformation import AzureAnthropicConfig + litellm_params: Final = deployment_litellm_params(deployment) match translated := await anthropic_count_fields( model_to_use, messages, system, tools, messages_format, AzureAnthropicConfig(), chat_params ): case ChatCountTranslationError(): return translated.response(request_model, model_to_use, "azure_ai_anthropic_api") case AnthropicCountFields(): - return await self.count_tokens( + return await self._count( model_to_use=model_to_use, - messages=[dict(message) for message in translated.messages], - contents=contents, - deployment=deployment, + messages=translated.messages, + litellm_params=litellm_params, request_model=request_model, - tools=[dict(tool) for tool in translated.tools] if translated.tools is not None else None, + tools=translated.tools, system=translated.system, - messages_format="anthropic", - output_format=output_format, - context_management=context_management, count_options=translated.options, + context_management=context_management, ) case None: pass - deployment = deployment or {} - litellm_params: Final = deployment.get("litellm_params", {}) + match options := await asyncify(messages_count_options)( + AzureAnthropicMessagesConfig(), + model_to_use, + messages, + count_options, + litellm_params, + ): + case ChatCountTranslationError(): + return options.response(request_model, model_to_use, "azure_ai_anthropic_api") + case _: + return await self._count( + model_to_use=model_to_use, + messages=messages, + litellm_params=litellm_params, + request_model=request_model, + tools=tools, + system=system, + count_options=options, + context_management=context_management, + ) - # Get Azure AI API key from deployment config or environment - api_key = litellm_params.get("api_key") - if not api_key: - api_key = os.getenv("AZURE_AI_API_KEY") + async def _count( + self, + model_to_use: str, + messages: Sequence[Mapping[str, JsonValue]], + litellm_params: Mapping[str, object], + request_model: str, + tools: Sequence[Mapping[str, JsonValue]] | None, + system: object | None, + count_options: Mapping[str, JsonValue], + context_management: CountContextManagement | None, + ) -> TokenCountResponse | None: + from litellm.llms.anthropic.common_utils import AnthropicError + api_key: Final = optional_str(litellm_params.get("api_key")) or os.getenv("AZURE_AI_API_KEY") api_base: Final = ( - litellm_params.get("api_base") + optional_str(litellm_params.get("api_base")) or os.getenv("AZURE_AI_API_BASE") or os.getenv("AZURE_API_BASE") or litellm.api_base @@ -129,11 +160,11 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): try: result: Final = await azure_ai_anthropic_count_tokens_handler.handle_count_tokens_request( model=model_to_use, - messages=messages, + messages=[dict(message) for message in messages], api_key=api_key, api_base=api_base, - litellm_params=litellm_params, - tools=tools, + litellm_params=dict(litellm_params), + tools=None if tools is None else [dict(tool) for tool in tools], system=system, optional_params=count_options, context_management=None if context_management is None else context_management.spec, diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index 91ed40d9271..f68620b18d6 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -14,6 +14,7 @@ from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_hea from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.common_utils import ( + AnthropicModelInfo, context_management_betas, count_tokens_feature_betas, merge_anthropic_beta_headers, @@ -39,6 +40,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): litellm_params: dict[str, object] | None = None, messages: Sequence[object] = (), context_management: JsonValue = None, + thinking: JsonValue = None, ) -> dict[str, str]: """ Get the required headers for the Azure AI Anthropic CountTokens API. @@ -51,6 +53,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): litellm_params: Optional LiteLLM parameters for additional auth config messages: The request messages, read for betas the counted features need context_management: The edit spec the count forwards, read for the betas its edits need + thinking: The thinking the count forwards, read for the betas inference adds for it Returns: Dictionary of required headers with both x-api-key and Azure authentication @@ -63,7 +66,10 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, ( *count_tokens_feature_betas("azure_ai", {}, messages), - *filter_and_transform_beta_headers(list(context_management_betas(context_management)), "azure_ai"), + *filter_and_transform_beta_headers( + [*context_management_betas(context_management), *AnthropicModelInfo().thinking_betas(thinking)], + "azure_ai", + ), ), ), **({"x-api-key": api_key} if api_key else {}), diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index 8b1ccda3a61..b5e617499b1 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -52,6 +52,15 @@ def deployment_timeout(litellm_params: Mapping[str, object]) -> float | httpx.Ti return None +def deployment_litellm_params(deployment: Mapping[str, object] | None) -> Mapping[str, object]: + params: Final = None if deployment is None else as_str_mapping(deployment.get("litellm_params")) + return MappingProxyType({}) if params is None else params + + +def optional_str(value: object) -> str | None: + return value if isinstance(value, str) else None + + class BaseTokenCounter(ABC): @abstractmethod async def count_tokens( @@ -67,6 +76,7 @@ class BaseTokenCounter(ABC): output_format: JsonValue = None, context_management: CountContextManagement | None = None, chat_params: Mapping[str, JsonValue] = MappingProxyType({}), + count_options: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse | None: pass diff --git a/litellm/llms/base_llm/files/batch_records.py b/litellm/llms/base_llm/files/batch_records.py index 87ac4bd2b93..b9328f4d0bf 100644 --- a/litellm/llms/base_llm/files/batch_records.py +++ b/litellm/llms/base_llm/files/batch_records.py @@ -11,7 +11,7 @@ def _frozen_mapping(items: Iterable[tuple[str, object]]) -> Mapping[str, object] @cache -def _responses_request_keys() -> frozenset[str]: +def responses_request_keys() -> frozenset[str]: return frozenset(get_type_hints(ResponsesAPIOptionalRequestParams)) @@ -46,7 +46,7 @@ def responses_batch_body_to_chat_body( chat_input: Final = cast(str | ResponseInputParam, responses_input) # cast-ok: forwarded as sent responses_request: Final = cast( # cast-ok: client-supplied fields forwarded verbatim, as real time does ResponsesAPIOptionalRequestParams, - _frozen_mapping((key, value) for key, value in openai_request_body.items() if key in _responses_request_keys()), + _frozen_mapping((key, value) for key, value in openai_request_body.items() if key in responses_request_keys()), ) return LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # transformer declares a bare dict return model=model if isinstance(model, str) else "", diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 3866240030b..78839108390 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -5,17 +5,22 @@ Bedrock Token Counter implementation using the CountTokens API. from collections.abc import Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.common_utils import context_management_betas, context_management_with_edits -from litellm.llms.anthropic.count_tokens.transformation import ChatCountTranslationError, anthropic_count_fields -from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.anthropic.count_tokens.transformation import ( + ChatCountTranslationError, + anthropic_count_fields, + messages_count_options, +) +from litellm.llms.base_llm.base_utils import BaseTokenCounter, deployment_litellm_params from litellm.llms.bedrock.common_utils import ( BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS, BedrockError, @@ -33,9 +38,14 @@ from litellm.types.utils import ( TokenCountResponse, ) +if TYPE_CHECKING: + from litellm.llms.anthropic.messages.native.transformation import AnthropicMessagesConfig + RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api" MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api" _RUNTIME_FALLBACK_STATUSES: Final = frozenset({400, 404}) +_NO_OPTIONS: Final[Mapping[str, JsonValue]] = MappingProxyType({}) +_COUNTED_MESSAGES: Final = TypeAdapter(list[dict[str, JsonValue]]) class _CountTokensReply(LiteLLMPydanticObjectBase): @@ -60,18 +70,40 @@ class CountTokensFailure: CountTokensOutcome = CountedTokens | CountTokensFailure -def _claude_count_translator(litellm_params: Mapping[str, object]) -> AnthropicConfig: +def _deployment_route(litellm_params: Mapping[str, object]) -> tuple[str, LlmProviders] | None: deployment_model: Final = litellm_params.get("model") if not isinstance(deployment_model, str): - return AnthropicConfig() + return None provider_hint: Final = litellm_params.get("custom_llm_provider") model, provider, _, _ = litellm.get_llm_provider( model=deployment_model, custom_llm_provider=provider_hint if isinstance(provider_hint, str) else None ) - chat_config: Final = litellm.ProviderConfigManager.get_provider_chat_config(model, LlmProviders(provider)) + return model, LlmProviders(provider) + + +def _claude_count_translator(litellm_params: Mapping[str, object]) -> AnthropicConfig: + route: Final = _deployment_route(litellm_params) + chat_config: Final = None if route is None else litellm.ProviderConfigManager.get_provider_chat_config(*route) return chat_config if isinstance(chat_config, AnthropicConfig) else AnthropicConfig() +def _claude_messages_translator(litellm_params: Mapping[str, object]) -> "AnthropicMessagesConfig": + from litellm.llms.anthropic.messages.native.transformation import AnthropicMessagesConfig + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + route: Final = _deployment_route(litellm_params) + messages_config: Final = ( + None if route is None else litellm.ProviderConfigManager.get_provider_anthropic_messages_config(*route) + ) + return ( + messages_config + if isinstance(messages_config, AnthropicMessagesConfig) + else AmazonAnthropicClaudeMessagesConfig() + ) + + def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool: return ( isinstance(outcome, CountTokensFailure) @@ -194,6 +226,7 @@ class BedrockTokenCounter(BaseTokenCounter): output_format: JsonValue = None, context_management: CountContextManagement | None = None, chat_params: Mapping[str, JsonValue] = MappingProxyType({}), + count_options: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse | None: if not messages: return None @@ -202,7 +235,7 @@ class BedrockTokenCounter(BaseTokenCounter): counted_model: Final = model_to_use.removeprefix("bedrock_mantle/") resolved_model: Final = get_bedrock_base_model(counted_model) primary_tokenizer: Final = MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE - litellm_params: Final = (deployment or {}).get("litellm_params", {}) + litellm_params: Final = dict(deployment_litellm_params(deployment)) translated: Final = ( await anthropic_count_fields( counted_model, @@ -218,6 +251,21 @@ class BedrockTokenCounter(BaseTokenCounter): ) if isinstance(translated, ChatCountTranslationError): return translated.response(request_model, model_to_use, primary_tokenizer) + options: Final[Mapping[str, JsonValue] | ChatCountTranslationError] = ( + translated.options + if translated is not None + else await asyncify(messages_count_options)( + _claude_messages_translator(litellm_params), + counted_model, + _COUNTED_MESSAGES.validate_python(messages), + count_options, + litellm_params, + ) + if is_mantle_claude_model(resolved_model) + else _NO_OPTIONS + ) + if isinstance(options, ChatCountTranslationError): + return options.response(request_model, model_to_use, primary_tokenizer) native_context_management: Final = _invoke_context_management( context_management if is_mantle_claude_model(resolved_model) else None ) @@ -231,7 +279,7 @@ class BedrockTokenCounter(BaseTokenCounter): ), **native_context_management, } - mantle_request: Final = runtime_request if translated is None else {**runtime_request, **translated.options} + mantle_request: Final = {**runtime_request, **options} primary: Final = ( await self._count_with( diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py index 0b589eda896..5c48ea97132 100644 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -8,7 +8,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import join_api_base -from litellm.llms.anthropic.common_utils import context_management_betas, tools_without_custom_field +from litellm.llms.anthropic.common_utils import AnthropicModelInfo, context_management_betas, tools_without_custom_field from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url @@ -82,7 +82,10 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM): ), context_management=request.get("context_management"), ) - betas: Final = context_management_betas(request.get("context_management")) + betas: Final = ( + *context_management_betas(request.get("context_management")), + *AnthropicModelInfo().thinking_betas(request_data.get("thinking")), + ) verbose_logger.debug("Making bedrock-mantle count_tokens request to: %s", endpoint_url) api_key: Final = litellm_params.get("api_key") signed_headers, signed_body = await run_aws_signing( diff --git a/litellm/llms/deepseek/chat/transformation.py b/litellm/llms/deepseek/chat/transformation.py index 872e153c91a..0def99c780d 100644 --- a/litellm/llms/deepseek/chat/transformation.py +++ b/litellm/llms/deepseek/chat/transformation.py @@ -3,6 +3,7 @@ Translates from OpenAI's `/v1/chat/completions` to DeepSeek's `/v1/chat/completi """ from collections.abc import Coroutine, Mapping, Sequence +from types import MappingProxyType from typing import Final, Literal, cast, overload import litellm @@ -33,6 +34,35 @@ def _warn_on_dropped_budget(thinking: object, model: str) -> None: ) +def _thinking_type(thinking: object, reasoning_effort: object) -> str | None: + requested: Final = as_str_mapping(thinking) + requested_type: Final = None if requested is None else requested.get("type") + if isinstance(requested_type, str) and requested_type in ("enabled", "disabled"): + return requested_type + if reasoning_effort is None: + return None + return "disabled" if reasoning_effort == "none" else "enabled" + + +def deepseek_reasoning_params(thinking: object, reasoning_effort: object) -> Mapping[str, object]: + """DeepSeek takes `{"type": "enabled"}` or `{"type": "disabled"}` with no budget_tokens. A thinking + param wins. Otherwise `reasoning_effort="none"` turns thinking off and any other effort turns it on, + and a supported effort is sent only while thinking is on. + + Reference: https://api-docs.deepseek.com/guides/thinking_mode + """ + thinking_type: Final = _thinking_type(thinking, reasoning_effort) + if thinking_type is None: + return MappingProxyType({}) + if ( + thinking_type == "enabled" + and isinstance(reasoning_effort, str) + and reasoning_effort in SUPPORTED_REASONING_EFFORTS + ): + return MappingProxyType({"thinking": {"type": thinking_type}, "reasoning_effort": reasoning_effort}) + return MappingProxyType({"thinking": {"type": thinking_type}}) + + class DeepSeekChatConfig(OpenAIGPTConfig): def _flattens_content_lists(self) -> bool: return True @@ -50,46 +80,16 @@ class DeepSeekChatConfig(OpenAIGPTConfig): model: str, drop_params: bool, ) -> dict[str, object]: - """ - Map OpenAI params to DeepSeek params. - - Handles `thinking` and `reasoning_effort` parameters for DeepSeek reasoner models. - DeepSeek supports `{"type": "enabled"}` and `{"type": "disabled"}` - no budget_tokens - like Anthropic. `reasoning_effort="none"` is the OpenAI-style way to ask for thinking - off, so it maps to `{"type": "disabled"}`; any other effort keeps thinking on. - - Reference: https://api-docs.deepseek.com/guides/thinking_mode - """ - # Let parent handle standard params first optional_params = super().map_openai_params( tool_params_without_hosted_mcp(non_default_params, model=model, route="deepseek", drop_params=drop_params), optional_params, model, drop_params, ) - - # Pop thinking/reasoning_effort from optional_params first (parent may have added them) - # Then re-add only if valid for DeepSeek thinking_value: Final = optional_params.pop("thinking", None) reasoning_effort: Final = optional_params.pop("reasoning_effort", None) _warn_on_dropped_budget(OPAQUE_VALUE.validate_python(thinking_value), model) - - # Handle thinking parameter - accept both enabled and disabled, ignore budget_tokens - if isinstance(thinking_value, dict) and thinking_value.get("type") in ("enabled", "disabled"): - optional_params["thinking"] = {"type": thinking_value["type"]} - - # Otherwise fall back to reasoning_effort: "none" disables, anything else enables - elif reasoning_effort is not None: - optional_params["thinking"] = {"type": "disabled" if reasoning_effort == "none" else "enabled"} - - if ( - isinstance(reasoning_effort, str) - and reasoning_effort in SUPPORTED_REASONING_EFFORTS - and optional_params["thinking"] == {"type": "enabled"} - ): - return {**optional_params, "reasoning_effort": reasoning_effort} - - return optional_params + return {**optional_params, **deepseek_reasoning_params(thinking_value, reasoning_effort)} def _fill_reasoning_content(self, messages: list[AllMessageValues]) -> list[AllMessageValues]: """ @@ -255,7 +255,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): ) @staticmethod - def _drop_unsupported_tools(optional_params: dict) -> dict: + def drop_unsupported_tools(optional_params: dict) -> dict: """ DeepSeek's /chat/completions only accepts tools of type "function". @@ -341,7 +341,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): (user explicitly enabled it), preventing spurious injection on models like deepseek-v3.2 that support thinking as opt-in but not always-on. """ - optional_params = self._drop_unsupported_tools(optional_params) + optional_params = self.drop_unsupported_tools(optional_params) if self._thinking_mode_active(model=model, optional_params=optional_params): messages = self._fill_reasoning_content(messages) return super().transform_request( @@ -364,7 +364,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): Async equivalent of transform_request — applies the same reasoning_content fix for multi-turn thinking-mode conversations. """ - optional_params = self._drop_unsupported_tools(optional_params) + optional_params = self.drop_unsupported_tools(optional_params) if self._thinking_mode_active(model=model, optional_params=optional_params): messages = self._fill_reasoning_content(messages) return await super().async_transform_request( diff --git a/litellm/llms/deepseek/common_utils.py b/litellm/llms/deepseek/common_utils.py new file mode 100644 index 00000000000..ad194694bca --- /dev/null +++ b/litellm/llms/deepseek/common_utils.py @@ -0,0 +1,46 @@ +from collections.abc import Mapping, Sequence +from typing import Final + +from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter +from litellm.llms.deepseek.messages.transformation import DeepSeekAnthropicMessagesConfig, deepseek_anthropic_base +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import LlmProviders + + +class DeepSeekModelInfo(BaseLLMModelInfo): + def get_models( + self, api_key: str | None = None, api_base: str | None = None + ) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override + from litellm.utils import bundled_provider_models + + return list(bundled_provider_models(LlmProviders.DEEPSEEK.value)) + + @staticmethod + def get_api_key(api_key: str | None = None) -> str | None: + return DeepSeekAnthropicMessagesConfig.get_api_key(api_key) + + @staticmethod + def get_api_base(api_base: str | None = None) -> str | None: + return deepseek_anthropic_base(api_base) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: matches validate_environment's dict contract shared by every provider override + resolved_api_key: Final = self.get_api_key(api_key) + return {**headers, **({} if resolved_api_key is None else {"x-api-key": resolved_api_key})} + + @staticmethod + def get_base_model(model: str) -> str | None: + return model.removeprefix(f"{LlmProviders.DEEPSEEK.value}/") + + def get_token_counter(self) -> BaseTokenCounter | None: + from litellm.llms.deepseek.count_tokens.token_counter import DeepSeekTokenCounter + + return DeepSeekTokenCounter() diff --git a/litellm/llms/deepseek/count_tokens/token_counter.py b/litellm/llms/deepseek/count_tokens/token_counter.py new file mode 100644 index 00000000000..1f2f1cbe7a1 --- /dev/null +++ b/litellm/llms/deepseek/count_tokens/token_counter.py @@ -0,0 +1,246 @@ +""" +DeepSeek token counting through the count_tokens route of its Anthropic-compatible API. +""" + +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final + +import httpx +from pydantic import JsonValue, TypeAdapter + +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify +from litellm.llms.anthropic.common_utils import AnthropicError +from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler +from litellm.llms.anthropic.count_tokens.transformation import ( + AnthropicCountFields, + ChatCountTranslationError, + anthropic_count_fields, + messages_count_options, +) +from litellm.llms.base_llm.base_utils import ( + BaseTokenCounter, + deployment_litellm_params, + deployment_timeout, + optional_str, +) +from litellm.llms.deepseek.chat.transformation import DeepSeekChatConfig, deepseek_reasoning_params +from litellm.llms.deepseek.common_utils import DeepSeekModelInfo +from litellm.llms.deepseek.messages.transformation import ( + DeepSeekAnthropicMessagesConfig, + deepseek_anthropic_base, + sanitize_tools_for_deepseek, +) +from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse + +DEEPSEEK_TOKENIZER_TYPE: Final = "deepseek_api" +_CHAT_REASONING_PARAMS: Final = ("thinking", "reasoning_effort") +DEEPSEEK_COUNT_TIMEOUT_SECONDS: Final = 10.0 +DEEPSEEK_API_HOST: Final = "api.deepseek.com" +_ESTIMATED_COUNT_STATUSES: Final = frozenset({404, 405, 408, 422, 429}) +_TOOLS: Final = TypeAdapter(list[dict[str, JsonValue]]) +_OPTIONS: Final = TypeAdapter(dict[str, JsonValue]) + +deepseek_count_tokens_handler: Final = AnthropicCountTokensHandler() + + +def _count_left_to_estimate(status_code: int) -> bool: + return status_code in _ESTIMATED_COUNT_STATUSES or status_code >= 500 + + +def _count_timeout(litellm_params: Mapping[str, object]) -> float | httpx.Timeout: + match timeout := deployment_timeout(litellm_params): + case None: + return DEEPSEEK_COUNT_TIMEOUT_SECONDS + case httpx.Timeout(): + return timeout + case _: + return min(timeout, DEEPSEEK_COUNT_TIMEOUT_SECONDS) + + +def _chat_request_inference_sends( + tools: Sequence[Mapping[str, JsonValue]] | None, chat_params: Mapping[str, JsonValue] +) -> tuple[list[dict[str, JsonValue]] | None, Mapping[str, JsonValue]]: + sent: Final = _OPTIONS.validate_python( + DeepSeekChatConfig.drop_unsupported_tools({**chat_params, **({} if tools is None else {"tools": list(tools)})}) + ) + sent_tools: Final = sent.get("tools") + return ( + None if sent_tools is None else _TOOLS.validate_python(sent_tools), + MappingProxyType({name: value for name, value in sent.items() if name != "tools"}), + ) + + +def _chat_reasoning_options(chat_params: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + mapped: Final = deepseek_reasoning_params(chat_params.get("thinking"), chat_params.get("reasoning_effort")) + effort: Final = mapped.get("reasoning_effort") + return _OPTIONS.validate_python( + { + **({} if mapped.get("thinking") is None else {"thinking": mapped["thinking"]}), + **({} if effort is None else {"output_config": {"effort": effort}}), + } + ) + + +def _with_chat_reasoning( + options: Mapping[str, JsonValue], reasoning: Mapping[str, JsonValue] +) -> Mapping[str, JsonValue]: + output_config: Final = options.get("output_config") + reasoning_output_config: Final = reasoning.get("output_config") + return _OPTIONS.validate_python( + { + **options, + **reasoning, + **( + {"output_config": {**output_config, **reasoning_output_config}} + if isinstance(output_config, dict) and isinstance(reasoning_output_config, dict) + else {} + ), + } + ) + + +class DeepSeekTokenCounter(BaseTokenCounter): + """Counts with the request DeepSeek bills: an Anthropic Messages body sent to the Anthropic-compatible + count_tokens route, with thinking and effort mapped the way the deployment's inference maps them.""" + + def __init__(self, handler: AnthropicCountTokensHandler = deepseek_count_tokens_handler) -> None: + self._handler: Final = handler + + def should_use_token_counting_api(self, custom_llm_provider: str | None = None) -> bool: + return custom_llm_provider == LlmProviders.DEEPSEEK.value + + async def count_tokens( + self, + model_to_use: str, + messages: Sequence[Mapping[str, JsonValue]] | None, + contents: Sequence[Mapping[str, object]] | None, + deployment: Mapping[str, object] | None = None, + request_model: str = "", + tools: Sequence[Mapping[str, JsonValue]] | None = None, + system: JsonValue = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, + context_management: CountContextManagement | None = None, + chat_params: Mapping[str, JsonValue] = MappingProxyType({}), + count_options: Mapping[str, JsonValue] = MappingProxyType({}), + ) -> TokenCountResponse | None: + litellm_params: Final = deployment_litellm_params(deployment) + api_base: Final = deepseek_anthropic_base(optional_str(litellm_params.get("api_base"))) + if not messages or httpx.URL(api_base).host != DEEPSEEK_API_HOST: + return None + + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + chat_tools, sent_chat_params = ( + _chat_request_inference_sends(tools, chat_params) if messages_format == "openai" else (tools, chat_params) + ) + match translated := await anthropic_count_fields( + model_to_use, + messages, + system, + chat_tools, + messages_format, + AnthropicConfig(), + {name: value for name, value in sent_chat_params.items() if name not in _CHAT_REASONING_PARAMS}, + ): + case ChatCountTranslationError(): + return translated.response(request_model, model_to_use, DEEPSEEK_TOKENIZER_TYPE) + case AnthropicCountFields(): + return await self._count( + model_to_use=model_to_use, + messages=translated.messages, + litellm_params=litellm_params, + api_base=api_base, + request_model=request_model, + tools=translated.tools, + system=translated.system, + count_options=_with_chat_reasoning(translated.options, _chat_reasoning_options(sent_chat_params)), + ) + case None: + pass + + if messages_format == "openai": + return await self._count( + model_to_use=model_to_use, + messages=messages, + litellm_params=litellm_params, + api_base=api_base, + request_model=request_model, + tools=chat_tools, + system=system, + count_options=_chat_reasoning_options(sent_chat_params), + ) + + match options := await asyncify(messages_count_options)( + DeepSeekAnthropicMessagesConfig(), model_to_use, messages, count_options, litellm_params + ): + case ChatCountTranslationError(): + return options.response(request_model, model_to_use, DEEPSEEK_TOKENIZER_TYPE) + case _: + return await self._count( + model_to_use=model_to_use, + messages=messages, + litellm_params=litellm_params, + api_base=api_base, + request_model=request_model, + tools=tools, + system=system, + count_options=options, + ) + + async def _count( + self, + model_to_use: str, + messages: Sequence[Mapping[str, JsonValue]], + litellm_params: Mapping[str, object], + api_base: str, + request_model: str, + tools: Sequence[Mapping[str, JsonValue]] | None, + system: JsonValue, + count_options: Mapping[str, JsonValue], + ) -> TokenCountResponse | None: + model_info: Final = DeepSeekModelInfo() + try: + result: Final = await self._handler.handle_count_tokens_request( + model=model_to_use, + messages=[dict(message) for message in messages], + auth_header=model_info.validate_environment( + headers={}, + model=model_to_use, + messages=(), + optional_params={}, + litellm_params=litellm_params, + api_key=optional_str(litellm_params.get("api_key")), + ), + api_base=api_base, + tools=None if tools is None else _TOOLS.validate_python(sanitize_tools_for_deepseek(tools)), + system=system, + optional_params=count_options, + timeout=_count_timeout(litellm_params), + ) + except AnthropicError as e: + if _count_left_to_estimate(e.status_code): + verbose_logger.warning( + "DeepSeek count_tokens route answered %s, so the count is estimated: %s", e.status_code, e.message + ) + return None + verbose_logger.warning("DeepSeek CountTokens error: status=%s, message=%s", e.status_code, e.message) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type=DEEPSEEK_TOKENIZER_TYPE, + error=True, + error_message=e.message, + status_code=e.status_code, + ) + input_tokens: Final = result.get("input_tokens") + return TokenCountResponse( + total_tokens=input_tokens if isinstance(input_tokens, int) else 0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type=DEEPSEEK_TOKENIZER_TYPE, + original_response=result, + ) diff --git a/litellm/llms/deepseek/messages/transformation.py b/litellm/llms/deepseek/messages/transformation.py index 3b2076b0085..9da594ee2ff 100644 --- a/litellm/llms/deepseek/messages/transformation.py +++ b/litellm/llms/deepseek/messages/transformation.py @@ -16,6 +16,7 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams _TOOL_LIST: Final = TypeAdapter(list[object]) +_REQUEST_FIELDS: Final = TypeAdapter(dict[str, object]) _TOOL_FIELDS: Final = TypeAdapter(dict[object, object]) @@ -34,6 +35,20 @@ def _sanitized_tool(tool: object) -> object: return tool if stripped is None else stripped +def sanitize_tools_for_deepseek(tools: Sequence[object]) -> Sequence[object]: + return [_sanitized_tool(tool) for tool in tools] + + +def deepseek_anthropic_base(api_base: str | None) -> str: + configured: Final = DeepSeekAnthropicMessagesConfig.get_api_base(api_base=api_base).rstrip("/") + if configured.endswith("/v1/messages") and "/anthropic/" in configured: + return configured.removesuffix("/v1/messages") + stripped: Final = configured.removesuffix("/v1/messages").removesuffix("/v1").removesuffix("/beta") + if stripped.endswith("/anthropic") or "/anthropic/" in stripped: + return stripped + return join_api_base(stripped, "anthropic") + + class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig): """ DeepSeek exposes an Anthropic-compatible Messages API at @@ -102,25 +117,7 @@ class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - base_url = self.get_api_base(api_base=api_base).rstrip("/") - - if base_url.endswith("/v1/messages") and "/anthropic/" in base_url: - return base_url - base_url = base_url.removesuffix("/v1/messages") - base_url = base_url.removesuffix("/v1") - base_url = base_url.removesuffix("/beta") - - if not base_url.endswith("/anthropic") and "/anthropic/" not in base_url: - base_url = join_api_base(base_url, "anthropic") - - return join_api_base(base_url, "v1/messages") - - @staticmethod - def _sanitize_tools_for_deepseek(tools: object) -> object: - if not isinstance(tools, list): - return tools - - return [_sanitized_tool(tool) for tool in _TOOL_LIST.validate_python(tools)] + return join_api_base(deepseek_anthropic_base(api_base), "v1/messages") def transform_anthropic_messages_request( self, @@ -130,13 +127,20 @@ class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> dict: - anthropic_messages_request: Final = super().transform_anthropic_messages_request( - model=model, - messages=messages, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, - litellm_params=litellm_params, - headers=headers, + anthropic_messages_request: Final = _REQUEST_FIELDS.validate_python( + super().transform_anthropic_messages_request( + model=model, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) ) - if "tools" in anthropic_messages_request: - anthropic_messages_request["tools"] = self._sanitize_tools_for_deepseek(anthropic_messages_request["tools"]) - return anthropic_messages_request + tools: Final = anthropic_messages_request.get("tools") + if not isinstance(tools, list): + return anthropic_messages_request + sanitized: Final = { + **anthropic_messages_request, + "tools": sanitize_tools_for_deepseek(_TOOL_LIST.validate_python(tools)), + } + return sanitized diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index b95219348e4..612a2d603bf 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -542,6 +542,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): output_format: JsonValue = None, context_management: CountContextManagement | None = None, chat_params: Mapping[str, JsonValue] = MappingProxyType({}), + count_options: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse | None: import copy diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index cd60bb1327f..787d1b840a2 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -180,6 +180,7 @@ class OpenAITokenCounter(BaseTokenCounter): output_format: JsonValue = None, context_management: CountContextManagement | None = None, chat_params: Mapping[str, JsonValue] = MappingProxyType({}), + count_options: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse | None: """ Count tokens using OpenAI's Responses API /input_tokens endpoint. diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index aa275f254b4..48bf955f707 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -7,7 +7,7 @@ from enum import Enum from functools import lru_cache from itertools import accumulate, chain, repeat, takewhile from types import MappingProxyType -from typing import Final, Literal, TypeVar, cast, get_type_hints +from typing import TYPE_CHECKING, Final, Literal, TypeVar, cast, get_type_hints import httpx from pydantic import JsonValue, TypeAdapter, ValidationError @@ -33,6 +33,9 @@ from litellm.types.llms.vertex_ai import ( from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages +if TYPE_CHECKING: + from litellm.llms.anthropic.count_tokens.transformation import AnthropicCountFields + GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType( {"fps": "fps", "start_offset": "startOffset", "end_offset": "endOffset"} ) @@ -97,6 +100,7 @@ class VertexAILyriaModelInfo(TypedDict): _VERTEX_AI_LYRIA_MODEL_INFO_ADAPTER: Final = TypeAdapter(VertexAILyriaModelInfo) +_JSON_MESSAGES: Final = TypeAdapter(list[dict[str, JsonValue]]) def _validate_vertex_ai_lyria_model_info(raw_model_info: object) -> VertexAILyriaModelInfo | None: @@ -1402,21 +1406,17 @@ class VertexAITokenCounter(BaseTokenCounter): return custom_llm_provider == LlmProviders.VERTEX_AI.value - async def _count_translated_claude_request( + async def _translated_claude_request( self, model_to_use: str, messages: Sequence[Mapping[str, object]] | None, - deployment: Mapping[str, object], request_model: str, tools: Sequence[Mapping[str, object]] | None, system: object | None, messages_format: TokenCountMessagesFormat, - output_format: JsonValue, - context_management: CountContextManagement | None, chat_params: Mapping[str, JsonValue], - ) -> TokenCountResponse | None: + ) -> "TokenCountResponse | AnthropicCountFields | None": from litellm.llms.anthropic.count_tokens.transformation import ( - AnthropicCountFields, ChatCountTranslationError, anthropic_count_fields, ) @@ -1426,27 +1426,43 @@ class VertexAITokenCounter(BaseTokenCounter): return None from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig - match translated := await anthropic_count_fields( + translated: Final = await anthropic_count_fields( model_to_use, messages or [], system, tools, messages_format, VertexAIAnthropicConfig(), chat_params - ): - case ChatCountTranslationError(): - return translated.response(request_model, model_to_use, "vertex_ai_partner_models") - case AnthropicCountFields(): - return await self.count_tokens( - model_to_use=model_to_use, - messages=[dict(message) for message in translated.messages], - contents=None, - deployment=dict(deployment), - request_model=request_model, - tools=[dict(tool) for tool in translated.tools] if translated.tools is not None else None, - system=translated.system, - messages_format="anthropic", - output_format=output_format, - context_management=context_management, - count_options=translated.options, - ) - case None: - return None + ) + if isinstance(translated, ChatCountTranslationError): + return translated.response(request_model, model_to_use, "vertex_ai_partner_models") + return translated + + async def _claude_count_options( + self, + model_to_use: str, + messages: Sequence[Mapping[str, object]], + request_model: str, + count_options: Mapping[str, JsonValue], + litellm_params: Mapping[str, object], + ) -> Mapping[str, JsonValue] | TokenCountResponse: + from litellm.litellm_core_utils.asyncify import asyncify + from litellm.llms.anthropic.count_tokens.transformation import ( + ChatCountTranslationError, + messages_count_options, + ) + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.messages.transformation import ( + VertexAIPartnerModelsAnthropicMessagesConfig, + ) + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_publisher + + if count_tokens_publisher(model_to_use) != "anthropic": + return count_options + options: Final = await asyncify(messages_count_options)( + VertexAIPartnerModelsAnthropicMessagesConfig(), + model_to_use, + _JSON_MESSAGES.validate_python(messages), + count_options, + litellm_params, + ) + if isinstance(options, ChatCountTranslationError): + return options.response(request_model, model_to_use, "vertex_ai_partner_models") + return options async def count_tokens( self, @@ -1502,26 +1518,32 @@ class VertexAITokenCounter(BaseTokenCounter): # Check if this is a partner model (Claude, Mistral, etc.) if VertexAIPartnerModels.is_vertex_partner_model(model_to_use): - translated_count: Final = await self._count_translated_claude_request( + translated: Final = await self._translated_claude_request( model_to_use=model_to_use, messages=messages, - deployment=resolved_deployment, request_model=request_model, tools=tools, system=system, messages_format=messages_format, - output_format=output_format, - context_management=context_management, chat_params=chat_params, ) - if translated_count is not None: - return translated_count + if isinstance(translated, TokenCountResponse): + return translated + partner_litellm_params: Final = dict(decoded_json_object(raw_litellm_params)) + partner_count_options: Final = ( + translated.options + if translated is not None + else await self._claude_count_options( + model_to_use, messages or [], request_model, count_options, partner_litellm_params + ) + ) + if isinstance(partner_count_options, TokenCountResponse): + return partner_count_options # Use partner models token counter partner_models_handler: Final = VertexAIPartnerModels() # Extract vertex-specific params from litellm_params - partner_litellm_params: Final = dict(decoded_json_object(raw_litellm_params)) vertex_project = partner_litellm_params.get("vertex_project") or partner_litellm_params.get( "vertex_ai_project" ) @@ -1540,14 +1562,22 @@ class VertexAITokenCounter(BaseTokenCounter): try: result = await partner_models_handler.count_tokens( model=model_to_use, - messages=messages or [], + messages=( + (messages or []) if translated is None else [dict(message) for message in translated.messages] + ), litellm_params=partner_litellm_params, vertex_project=vertex_project, vertex_location=vertex_location, vertex_credentials=vertex_credentials, - system=system, - tools=tools, - count_options=count_options, + system=system if translated is None else translated.system, + tools=( + tools + if translated is None + else None + if translated.tools is None + else [dict(tool) for tool in translated.tools] + ), + count_options=partner_count_options, context_management=( None if context_management is None else context_management_with_edits(context_management.spec) ), diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index 28d6321453f..20972a63edd 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -16,6 +16,7 @@ from litellm.litellm_core_utils.core_helpers import decoded_json_object, drops_u from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.common_utils import ( OBJECT_LIST_ADAPTER, + AnthropicModelInfo, context_management_betas, count_tokens_feature_betas, is_claude_model_target, @@ -64,7 +65,10 @@ def count_tokens_headers(access_token: str, request_data: Mapping[str, object]) "vertex_ai", request_data, _REQUEST_MESSAGES.validate_python(request_data["messages"]) ), *filter_and_transform_beta_headers( - list(context_management_betas(_CONTEXT_MANAGEMENT.validate_python(request_data.get("context_management")))), + [ + *context_management_betas(_CONTEXT_MANAGEMENT.validate_python(request_data.get("context_management"))), + *AnthropicModelInfo().thinking_betas(request_data.get("thinking")), + ], "vertex_ai", ), ) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index c298ca404e2..e905b97d29f 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -1,6 +1,6 @@ # What is this? ## API Handler for calling Vertex AI Partner Models -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from enum import Enum from types import MappingProxyType from typing import Final @@ -260,13 +260,13 @@ class VertexAIPartnerModels(VertexBase): async def count_tokens( self, model: str, - messages: list, - litellm_params: dict, - vertex_project=None, - vertex_location=None, - vertex_credentials=None, + messages: Sequence[Mapping[str, object]], + litellm_params: Mapping[str, object], + vertex_project: object = None, + vertex_location: object = None, + vertex_credentials: object = None, system: object | None = None, - tools: list[dict[str, object]] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, count_options: Mapping[str, object] = MappingProxyType({}), context_management: object = None, ): @@ -317,13 +317,18 @@ class VertexAIPartnerModels(VertexBase): } # Prepare litellm_params with credentials - _litellm_params: Final = litellm_params.copy() - if vertex_project: - _litellm_params["vertex_project"] = vertex_project - if vertex_location: - _litellm_params["vertex_location"] = vertex_location - if vertex_credentials: - _litellm_params["vertex_credentials"] = vertex_credentials + _litellm_params: Final = { + **litellm_params, + **{ + key: value + for key, value in ( + ("vertex_project", vertex_project), + ("vertex_location", vertex_location), + ("vertex_credentials", vertex_credentials), + ) + if value + }, + } # Call the token counter token_counter: Final = VertexAIPartnerModelsTokenCounter() diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4d832febe92..9ac4bb11854 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3631,6 +3631,12 @@ class TokenCountRequest(LiteLLMPydanticObjectBase): tools: list[dict] | None = 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 from litellm.types.integrations.slack_alerting import CallInfo as CallInfo # noqa: E402, PLC0414 # public re-export diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index ed3abbdc5b6..86092e4e27c 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -19,6 +19,7 @@ from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure from litellm.litellm_core_utils.tool_fields import first_validation_error from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler +from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, without_dropped_paths from litellm.llms.anthropic.messages.context_management import ( AnthropicContextManagementError, ) @@ -308,6 +309,7 @@ class _CountTokensBodyOptions(TypedDict): output_format: ReadOnly[JsonValue] context_management: ReadOnly[JsonValue] request_litellm_params: ReadOnly[Mapping[str, JsonValue]] + count_options: ReadOnly[Mapping[str, JsonValue]] def _count_tokens_body_options(data: Mapping[str, JsonValue], auto_drop_params: bool) -> _CountTokensBodyOptions: @@ -318,6 +320,12 @@ def _count_tokens_body_options(data: Mapping[str, JsonValue], auto_drop_params: "request_litellm_params": MappingProxyType( {"drop_params": True, **sent_params} if auto_drop_params else sent_params ), + "count_options": MappingProxyType( + without_dropped_paths( + {name: data[name] for name in COUNT_TOKEN_OPTION_NAMES if name in data}, + data.get("additional_drop_params"), + ) + ), } diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3da65084c7b..6a171c5fe53 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15338,6 +15338,7 @@ async def run_thread( # async def get_available_routes(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)): from litellm.litellm_core_utils.tool_fields import DropPolicy, drop_policy from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCounter +from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, without_dropped_paths from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient @@ -15354,8 +15355,10 @@ from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository from litellm.responses.litellm_completion_transformation.count_request import ( + COUNTED_CHAT_PARAMS, ChatCountBridgeError, chat_count_request, + counted_chat_params, with_provider_file_ids, ) @@ -15468,17 +15471,29 @@ async def _with_resolved_credentials(deployment: Mapping[str, object] | None) -> _DEPLOYMENT_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) +_COUNT_OPTION_DEFAULTS: Final = TypeAdapter(Mapping[str, JsonValue]) +_REQUEST_DROP_SETTINGS: Final = ("drop_params", "additional_drop_params") -def _deployment_with_request_drop_params( +def _count_options_with_defaults( + count_options: Mapping[str, JsonValue], defaults: Mapping[str, object], additional_drop_params: object +) -> Mapping[str, JsonValue]: + default_options: Final = _COUNT_OPTION_DEFAULTS.validate_python( + {name: defaults[name] for name in COUNT_TOKEN_OPTION_NAMES if defaults.get(name) is not None} + ) + return MappingProxyType(without_dropped_paths({**default_options, **count_options}, additional_drop_params)) + + +def _deployment_with_request_drop_settings( deployment: dict[str, object] | None, # mutable-ok: BaseTokenCounter.count_tokens takes a dict deployment - drop_params: object, + sent_params: Mapping[str, object], ) -> dict[str, object] | None: # mutable-ok: BaseTokenCounter.count_tokens takes a dict deployment - if drop_params is None: + overrides: Final = {name: sent_params[name] for name in _REQUEST_DROP_SETTINGS if sent_params.get(name) is not None} + if not overrides: return deployment resolved_deployment: Final = deployment or {} litellm_params: Final = _DEPLOYMENT_LITELLM_PARAMS.validate_python(resolved_deployment.get("litellm_params") or {}) - return {**resolved_deployment, "litellm_params": {**litellm_params, "drop_params": drop_params}} + return {**resolved_deployment, "litellm_params": {**litellm_params, **overrides}} def _counted_request_messages( @@ -15508,6 +15523,7 @@ async def _try_provider_token_count( output_format: JsonValue = None, context_management: CountContextManagement | None = None, chat_params: Mapping[str, JsonValue] = MappingProxyType({}), + count_options: Mapping[str, JsonValue] = MappingProxyType({}), ) -> Union["TokenCountResponse", ProviderCountFailure, None]: """Attempt provider-specific token counting. Returns None to fall through to local counting.""" if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): @@ -15535,6 +15551,7 @@ async def _try_provider_token_count( output_format=output_format, context_management=context_management, chat_params=chat_params, + count_options=count_options, ) except httpx.HTTPStatusError as e: return ProviderCountFailure( @@ -15573,7 +15590,7 @@ async def _provider_token_count( request_model: str, settings: TokenCountSettings, local_estimate_allowed: bool, - forwarded_edits: bool = False, + forwards_request_fields: bool = False, ) -> Union["TokenCountResponse", ProviderCountFailure, None]: """Count with the provider, using the token_count_model in its region when one is set. @@ -15597,7 +15614,7 @@ async def _provider_token_count( pinned_count_model: Final = None if count_model is None or count_model.region is None else count_model region_hint: Final = ( None - if pinned_count_model is None or (forwarded_edits and counted.status_code == 400) + if pinned_count_model is None or (forwards_request_fields and counted.status_code == 400) else f"This model counts tokens with {pinned_count_model.model} in {pinned_count_model.region}. " "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info" ) @@ -16184,6 +16201,7 @@ async def count_request_tokens( responses_request: Mapping[str, object] | None = None, litellm_call_id: str | None = None, client_oauth_token: str | None = None, + count_options: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse: call_id: Final = resolve_litellm_call_id(litellm_call_id) if request.prompt is None and request.messages is None and request.contents is None: @@ -16203,6 +16221,7 @@ async def count_request_tokens( request_litellm_params=request_litellm_params, responses_request=responses_request, client_oauth_token=client_oauth_token, + count_options=count_options, ) return await logged_token_count( _with_anthropic_content_errors(count) if messages_format == "anthropic" else count, @@ -16249,6 +16268,7 @@ async def _count_request_tokens( request_litellm_params: Mapping[str, JsonValue], responses_request: Mapping[str, object] | None, client_oauth_token: str | None, + count_options: Mapping[str, JsonValue], ) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages @@ -16295,7 +16315,11 @@ async def _count_request_tokens( model_to_use = _model sent_params: Final[Mapping[str, object]] = {**_router_default_litellm_params(llm_router), **request_litellm_params} - count_deployment: Final = _deployment_with_request_drop_params(deployment, sent_params.get("drop_params")) + count_deployment: Final = _deployment_with_request_drop_settings(deployment, sent_params) + count_litellm_params: Final = _DEPLOYMENT_LITELLM_PARAMS.validate_python( + (count_deployment or {}).get("litellm_params") or {} + ) + deployment_defaults: Final = {**count_litellm_params, **_router_default_litellm_params(llm_router)} conversation_empty: Final = not messages and not contents if prompt and not conversation_empty and not deployment_drops_params(count_deployment): raise ProxyException( @@ -16363,10 +16387,26 @@ async def _count_request_tokens( _deployment_api_base(deployment_view, routed_request), routed_request, _count_drop_policy(routed_deployment, routed_request), + deployment_params=deployment_defaults, ) ) if isinstance(chat_count, ChatCountBridgeError): raise ProxyException(message=chat_count.message, type="invalid_request_error", param=None, code=400) + chat_params: Final[Mapping[str, JsonValue]] = ( + chat_count.params + if chat_count is not None + else counted_chat_params( + request.model_dump(include=set(COUNTED_CHAT_PARAMS), exclude_none=True), + deployment_defaults, + (), + drop_policy(count_litellm_params), + ) + if messages_format == "openai" + else MappingProxyType({}) + ) + forwarded_count_options: Final = _count_options_with_defaults( + count_options, deployment_defaults, count_litellm_params.get("additional_drop_params") + ) ladder: Final = ( None if provider_counter is None @@ -16397,13 +16437,15 @@ async def _count_request_tokens( messages_format=messages_format, output_format=output_format, context_management=None if edited_request is not None else context_edits, - chat_params=MappingProxyType({}) if chat_count is None else chat_count.params, + chat_params=chat_params, + count_options=forwarded_count_options, ), model_to_use=model_to_use, request_model=request.model, settings=settings, local_estimate_allowed=litellm.disable_token_counter is not True and contents_estimable, - forwarded_edits=context_edits is not None and edited_request is None, + forwards_request_fields=(context_edits is not None and edited_request is None) + or bool(forwarded_count_options), ) ) if isinstance(ladder, TokenCountResponse): diff --git a/litellm/responses/litellm_completion_transformation/count_request.py b/litellm/responses/litellm_completion_transformation/count_request.py index 2141b55850b..a0b25e9604e 100644 --- a/litellm/responses/litellm_completion_transformation/count_request.py +++ b/litellm/responses/litellm_completion_transformation/count_request.py @@ -1,5 +1,6 @@ from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass +from types import MappingProxyType from typing import Final, cast # noqa: TID251 # input items are forwarded as sent, as inference does from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError @@ -11,7 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.file_id_utils import ( update_responses_tools_with_model_file_ids, ) from litellm.litellm_core_utils.tool_fields import DropPolicy -from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body +from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body, responses_request_keys from litellm.responses.additional_tools import with_additional_tools_hoisted from litellm.responses.file_search.emulated_handler import replace_file_search_tools from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary @@ -23,7 +24,8 @@ from litellm.responses.main import previous_response_client_ids from litellm.types.llms.openai import ResponseInputParam from litellm.utils import ProviderConfigManager -_COUNTED_CHAT_PARAMS: Final = ( +COUNTED_CHAT_PARAMS: Final = ( + "thinking", "tool_choice", "parallel_tool_calls", "response_format", @@ -157,19 +159,27 @@ async def _with_previous_turns( ) -def _counted_params( - chat_body: Mapping[str, object], tools: tuple[Mapping[str, object], ...], policy: DropPolicy +def counted_chat_params( + chat_body: Mapping[str, object], + deployment_params: Mapping[str, object], + tools: tuple[Mapping[str, object], ...], + policy: DropPolicy, ) -> Mapping[str, JsonValue]: + inference_params: Final = {**deployment_params, **chat_body} sent: Final = { - name: chat_body[name] - for name in _COUNTED_CHAT_PARAMS - if chat_body.get(name) is not None and name not in policy.additional_drop_params + name: inference_params[name] + for name in COUNTED_CHAT_PARAMS + if inference_params.get(name) is not None and name not in policy.additional_drop_params } return _JSON_PARAMS.validate_python( {**sent, **({"drop_params": True} if policy.drop_params and (sent or tools) else {})} ) +def _fields_outside_responses_api(responses_request: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in responses_request.items() if key not in responses_request_keys()} + + async def chat_count_request( model: str, custom_llm_provider: str | None, @@ -177,6 +187,7 @@ async def chat_count_request( responses_request: Mapping[str, object], policy: DropPolicy, session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None = None, + deployment_params: Mapping[str, object] = MappingProxyType({}), ) -> ChatCountRequest | ChatCountBridgeError | None: """The chat request Responses inference sends a provider with no native Responses API, or None when inference sends this provider the Responses body itself.""" @@ -211,7 +222,9 @@ async def chat_count_request( return ChatCountRequest( messages=tuple(map(_unmarked, _mappings(chat_body.get("messages")))), tools=tools or None, - params=_counted_params(chat_body, tools, policy), + params=counted_chat_params( + {**_fields_outside_responses_api(responses_request), **chat_body}, deployment_params, tools, policy + ), ) except BadRequestError as e: return ChatCountBridgeError(message=str(e.message)) 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 80e8893a5a7..de0ab566dce 100644 --- a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -50870,14 +50870,20 @@ export interface components { }[] | null; /** Model */ model: string; + parallel_tool_calls?: components["schemas"]["JsonValue"]; /** Prompt */ prompt?: string | null; + reasoning_effort?: components["schemas"]["JsonValue"]; + response_format?: components["schemas"]["JsonValue"]; /** System */ system?: unknown | null; + thinking?: components["schemas"]["JsonValue"]; + tool_choice?: components["schemas"]["JsonValue"]; /** Tools */ tools?: { [key: string]: unknown; }[] | null; + web_search_options?: components["schemas"]["JsonValue"]; }; /** TokenCountResponse */ TokenCountResponse: { diff --git a/litellm/utils.py b/litellm/utils.py index 8deda4cd7de..5a12475ab86 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10864,6 +10864,10 @@ class ProviderConfigManager: from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo return AzureFoundryModelInfo(model=model) + elif LlmProviders.DEEPSEEK == provider: + from litellm.llms.deepseek.common_utils import DeepSeekModelInfo + + return DeepSeekModelInfo() elif LlmProviders.OPENROUTER == provider: from litellm.llms.openrouter.common_utils import OpenRouterModelInfo diff --git a/tests/benchmarks/test_count_tokens_benchmarks.py b/tests/benchmarks/test_count_tokens_benchmarks.py new file mode 100644 index 00000000000..fc1a7c615eb --- /dev/null +++ b/tests/benchmarks/test_count_tokens_benchmarks.py @@ -0,0 +1,27 @@ +from typing import Final + +import pytest +from pydantic import JsonValue + +from litellm.llms.anthropic.count_tokens.transformation import without_dropped_paths + +_LARGE_SCHEMA_OPTIONS: Final[dict[str, JsonValue]] = { + "output_config": { + "format": { + "type": "json_schema", + "schema": { + "type": "object", + "properties": {f"field_{index}": {"type": "string", "description": "d" * 8} for index in range(3000)}, + }, + } + }, + "thinking": {"type": "enabled", "budget_tokens": 2048}, +} +_MANY_DROP_PATHS: Final = (*(f"output_config.missing_{index}" for index in range(1000)), "thinking") + + +@pytest.mark.benchmark +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"]} diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_count_tokens_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_count_tokens_wire.py index 1009a954bee..72e1196a29e 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_count_tokens_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_count_tokens_wire.py @@ -173,10 +173,6 @@ def test_anthropic_count_tokens_relays_an_upstream_client_error(gateway: Gateway def test_anthropic_count_tokens_preserves_thinking_tool_choice_and_output_config(gateway: Gateway) -> None: - pytest.skip( - "BUG: /v1/messages/count_tokens drops thinking, tool_choice and output_config before the provider count call" - ) - thinking: Final = {"type": "enabled", "budget_tokens": 1024} tool_choice: Final = {"type": "auto", "disable_parallel_tool_use": True} output_config: Final = {"effort": "medium"} diff --git a/tests/unit/litellm_core_utils/test_dot_notation_indexing.py b/tests/unit/litellm_core_utils/test_dot_notation_indexing.py index 7c670f1dac9..7a64c233205 100644 --- a/tests/unit/litellm_core_utils/test_dot_notation_indexing.py +++ b/tests/unit/litellm_core_utils/test_dot_notation_indexing.py @@ -10,6 +10,7 @@ import pytest from litellm.litellm_core_utils.dot_notation_indexing import ( get_nested_value, delete_nested_value, + delete_nested_values, ) @@ -139,3 +140,19 @@ class TestDeleteNestedValue: def test_delete_accepts_a_read_only_mapping(self): payload: Final = MappingProxyType({"tags": ["a"], "secret": ["s"]}) assert delete_nested_value(payload, "secret") == {"tags": ["a"]} + + +def test_delete_nested_values_removes_every_path_and_leaves_the_input_untouched(): + payload: Final = MappingProxyType( + { + "tools": [{"name": "t1", "secret": "s1"}, {"name": "t2", "secret": "s2"}], + "meta": {"keep": 1, "drop": 2}, + "top": "x", + } + ) + result: Final = delete_nested_values( + payload, ("tools[*].secret", "meta.drop", "top", "top", "missing.path", "tools[9].name") + ) + assert result == {"tools": [{"name": "t1"}, {"name": "t2"}], "meta": {"keep": 1}} + assert payload["tools"] == [{"name": "t1", "secret": "s1"}, {"name": "t2", "secret": "s2"}] + assert payload["meta"] == {"keep": 1, "drop": 2} diff --git a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py index 13c17ebd2bf..11639e80d89 100644 --- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -14,11 +14,16 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCounter +from litellm.exceptions import BadRequestError from litellm.llms.anthropic.count_tokens.transformation import ( + COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig, + without_dropped_paths, ) +from litellm.llms.anthropic.messages.native.transformation import AnthropicMessagesConfig from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) @@ -26,6 +31,8 @@ from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAn from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( AzureAIAnthropicCountTokensConfig, ) +from litellm.llms.azure_ai.anthropic.messages_transformation import AzureAnthropicMessagesConfig +from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES @@ -557,6 +564,8 @@ async def _counted_body( tools: list[dict[str, JsonValue]] | None, messages_format: TokenCountMessagesFormat, chat_params: Mapping[str, JsonValue] = MappingProxyType({}), + count_options: Mapping[str, JsonValue] = MappingProxyType({}), + model: str = "claude-opus-5-5", ) -> tuple[TokenCountResponse | None, tuple[dict[str, JsonValue], ...]]: sent: Final[SimpleQueue[bytes]] = SimpleQueue() @@ -567,14 +576,15 @@ async def _counted_body( with respx.mock: respx.post(count_endpoint).mock(side_effect=count) result: Final = await counter_type().count_tokens( - model_to_use="claude-opus-5-5", + model_to_use=model, messages=messages, contents=None, deployment={"litellm_params": {"api_key": "sk-ant-api03-test-key", "api_base": "https://gateway.example"}}, - request_model="claude-opus-5-5", + request_model=model, tools=tools, messages_format=messages_format, chat_params=chat_params, + count_options=count_options, ) return result, tuple(TypeAdapter(dict[str, JsonValue]).validate_json(sent.get()) for _ in range(sent.qsize())) @@ -919,6 +929,150 @@ async def test_counter_sends_the_structured_output_tool_choice_and_thinking_chat ) +_MESSAGES_CONFIGS: Final = MappingProxyType( + {AnthropicTokenCounter: AnthropicMessagesConfig, AzureAIAnthropicTokenCounter: AzureAnthropicMessagesConfig} +) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +@pytest.mark.parametrize( + ("model", "options"), + ( + ("claude-sonnet-4-5", {"tool_choice": {"type": "any"}, "thinking": {"type": "enabled", "budget_tokens": 2048}}), + ("claude-haiku-5-5", {"thinking": {"type": "enabled", "budget_tokens": 2048}}), + ("claude-opus-5-5", {"thinking": {"type": "adaptive"}, "output_config": {"effort": "low"}}), + ("claude-unlisted-model", {"tool_choice": {"type": "any"}, "thinking": {"type": "enabled", "budget_tokens": 2048}}), + ), +) +async def test_counter_sends_the_thinking_tool_choice_and_effort_messages_inference_sends( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], + count_endpoint: str, + model: str, + options: Mapping[str, JsonValue], +) -> None: + messages: Final[list[dict[str, JsonValue]]] = [{"role": "user", "content": "Weather in SF?"}] + inference: Final = _MESSAGES_CONFIGS[counter_type]().transform_anthropic_messages_request( + model=model, + messages=deepcopy(messages), + anthropic_messages_optional_request_params={**deepcopy(dict(options)), "max_tokens": 4096}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + result, bodies = await _counted_body( + counter_type, count_endpoint, messages, None, "anthropic", count_options=options, model=model + ) + + (body,) = bodies + assert result is not None + assert result.total_tokens == 31 + assert {name: body.get(name) for name in COUNT_TOKEN_OPTION_NAMES} == { + name: inference.get(name) for name in COUNT_TOKEN_OPTION_NAMES + } + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +async def test_counter_leaves_out_an_option_the_deployment_drops( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], count_endpoint: str +) -> None: + sent: Final[SimpleQueue[bytes]] = SimpleQueue() + + def count(request: httpx.Request) -> httpx.Response: + sent.put(request.content) + return httpx.Response(200, json={"input_tokens": 31}) + + with respx.mock: + respx.post(count_endpoint).mock(side_effect=count) + await counter_type().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "Weather in SF?"}], + contents=None, + deployment={ + "litellm_params": { + "api_key": "sk-ant-api03-test-key", + "api_base": "https://gateway.example", + "additional_drop_params": ["tool_choice"], + } + }, + request_model="claude-sonnet-4-5", + messages_format="anthropic", + count_options={"tool_choice": {"type": "any"}, "thinking": {"type": "enabled", "budget_tokens": 2048}}, + ) + + body: Final = TypeAdapter(dict[str, JsonValue]).validate_json(sent.get()) + assert ("tool_choice" in body, body["thinking"]) == (False, {"type": "enabled", "budget_tokens": 2048}) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +@pytest.mark.parametrize("drop_params", (True, False)) +async def test_counter_sends_thinking_display_with_the_beta_messages_inference_adds_for_it( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], + count_endpoint: str, + drop_params: bool, +) -> None: + thinking: Final[dict[str, JsonValue]] = {"type": "adaptive", "display": "updates"} + inference_betas: Final = AnthropicModelInfo().thinking_betas(thinking) + with respx.mock: + route: Final = respx.post(count_endpoint).mock(return_value=httpx.Response(200, json={"input_tokens": 31})) + await counter_type().count_tokens( + model_to_use="claude-opus-5-5", + messages=[{"role": "user", "content": "Weather in SF?"}], + contents=None, + deployment={ + "litellm_params": { + "api_key": "sk-ant-api03-test-key", + "api_base": "https://gateway.example", + "drop_params": drop_params, + } + }, + request_model="claude-opus-5-5", + messages_format="anthropic", + count_options={"thinking": thinking}, + ) + + sent: Final = route.calls.last.request + assert inference_betas + assert frozenset(inference_betas) <= frozenset(sent.headers["anthropic-beta"].split(",")) + assert TypeAdapter(dict[str, JsonValue]).validate_json(sent.content)["thinking"] == thinking + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +async def test_counter_rejects_a_thinking_budget_messages_inference_rejects_without_calling_anthropic( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], count_endpoint: str +) -> None: + options: Final[dict[str, JsonValue]] = {"thinking": {"type": "enabled", "budget_tokens": "lots"}} + with pytest.raises(BadRequestError) as rejected: + _MESSAGES_CONFIGS[counter_type]().transform_anthropic_messages_request( + model="claude-opus-5-5", + messages=[{"role": "user", "content": "Weather in SF?"}], + anthropic_messages_optional_request_params={**deepcopy(options), "max_tokens": 4096}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + result, bodies = await _counted_body( + counter_type, + count_endpoint, + [{"role": "user", "content": "Weather in SF?"}], + None, + "anthropic", + count_options=options, + ) + + assert bodies == () + assert result is not None + assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) + assert rejected.value.message in (result.error_message or "") + + def test_headers_add_the_compaction_beta_only_for_a_compacted_history(): config = AnthropicCountTokensConfig() auth_header = {"x-api-key": "sk-ant-api03-test-key"} @@ -1029,3 +1183,21 @@ async def test_counter_leaves_out_context_management_without_edits(context_manag ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, } + + +@pytest.mark.parametrize( + ("additional_drop_params", "kept"), + ( + (["thinking", 5, None], ("output_config",)), + (("thinking",), ("output_config",)), + ("thinking", ("output_config", "thinking")), + (None, ("output_config", "thinking")), + ), +) +def test_drop_paths_skip_entries_and_values_inference_skips(additional_drop_params: object, kept: tuple[str, ...]) -> None: + options: Final[dict[str, JsonValue]] = { + "output_config": {"effort": "low"}, + "thinking": {"type": "enabled", "budget_tokens": 2048}, + } + + assert without_dropped_paths(options, additional_drop_params) == {name: options[name] for name in kept} diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py index 930a290098f..86d6b3468af 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py @@ -8,6 +8,10 @@ import pytest from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.types.llms.anthropic import ( + ANTHROPIC_THINKING_BINDING_CONTROLS_BETA_HEADER, + ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER, +) from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault MANTLE_COUNT_URL: Final = "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages/count_tokens" @@ -138,6 +142,35 @@ async def test_non_200_answers_raise_bedrock_error_with_mantle_status() -> None: assert "does not exist" in raised.value.message +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("thinking", "expected_beta"), + [ + ({"type": "adaptive", "display": "updates"}, ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER), + ( + {"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}}, + ANTHROPIC_THINKING_BINDING_CONTROLS_BETA_HEADER, + ), + ], + ids=["display_updates", "block_binding"], +) +async def test_a_beta_thinking_field_counts_with_the_beta_mantle_inference_sends( + thinking: dict[str, object], expected_beta: str +) -> None: + mantle: Final = _MantleEndpoint(200, {"input_tokens": 2177}) + + await BedrockMantleCountTokensHandler().handle_count_tokens_request( + request_data={**REQUEST, "thinking": thinking}, + litellm_params=dict(LITELLM_PARAMS), + resolved_model="anthropic.claude-opus-4-8", + client=_client(mantle), + ) + + posted: Final = mantle.only_request() + assert json.loads(posted.content)["thinking"] == thinking + assert expected_beta in posted.headers["anthropic-beta"].split(",") + + @pytest.mark.asyncio async def test_region_and_endpoint_lookups_leave_the_event_loop_free(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py index de0c14558b3..4e231a22919 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -381,6 +381,35 @@ async def test_runtime_count_leaves_out_the_options_only_mantle_counts(chat_para } +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_to_use", "runtime", "posted_hosts"), + ( + ("bedrock_mantle/anthropic.claude-opus-5-5", (200, {"inputTokens": 1}), (MANTLE_HOST,)), + ("us.anthropic.claude-opus-5-5", (400, UNSUPPORTED), (RUNTIME_HOST, MANTLE_HOST)), + ), +) +async def test_messages_count_options_reach_the_mantle_count( + model_to_use: str, runtime: tuple[int, Mapping[str, object]], posted_hosts: tuple[str, ...] +) -> None: + bedrock: Final = _Bedrock(runtime=runtime, mantle=(200, {"input_tokens": 433})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use=model_to_use, + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-5-5", + tools=[{"name": "get_weather", "input_schema": {"type": "object"}}], + count_options={"tool_choice": {"type": "any"}}, + ) + + assert result is not None + assert (result.error, result.total_tokens) == (False, 433) + assert bedrock.posted_hosts() == posted_hosts + assert _JSON_BODY.validate_json(bedrock.requests[-1].content)["tool_choice"] == {"type": "any"} + + @pytest.mark.asyncio async def test_claude_chat_tool_use_is_counted_in_the_shape_claude_inference_sends() -> None: bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) diff --git a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py index 160ed3950b3..25720b631b6 100644 --- a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py +++ b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py @@ -145,7 +145,7 @@ def test_drop_unsupported_tools_keeps_function_tools_only(): "tool_choice": "auto", } - result = DeepSeekChatConfig._drop_unsupported_tools(optional_params) + result = DeepSeekChatConfig.drop_unsupported_tools(optional_params) assert [tool["function"]["name"] for tool in result["tools"]] == [ "shell", @@ -163,7 +163,7 @@ def test_drop_unsupported_tools_drops_dangling_tool_choice_when_none_survive(): "temperature": 0.2, } - result = DeepSeekChatConfig._drop_unsupported_tools(optional_params) + result = DeepSeekChatConfig.drop_unsupported_tools(optional_params) assert "tools" not in result assert "tool_choice" not in result @@ -177,7 +177,7 @@ def test_drop_unsupported_tools_is_noop_for_function_only(): "tool_choice": "auto", } - result = DeepSeekChatConfig._drop_unsupported_tools(optional_params) + result = DeepSeekChatConfig.drop_unsupported_tools(optional_params) assert result is optional_params @@ -185,7 +185,7 @@ def test_drop_unsupported_tools_is_noop_for_function_only(): def test_drop_unsupported_tools_is_noop_without_tools(): optional_params = {"temperature": 0.7} - result = DeepSeekChatConfig._drop_unsupported_tools(optional_params) + result = DeepSeekChatConfig.drop_unsupported_tools(optional_params) assert result is optional_params @@ -733,7 +733,7 @@ class TestDeepSeekThinkingParams: "parallel_tool_calls": True, } - result = self.config._drop_unsupported_tools(optional_params) + result = self.config.drop_unsupported_tools(optional_params) assert result["tools"] == [{"type": "function", "function": {"name": "get_weather"}}] assert "tool_choice" not in result diff --git a/tests/unit/llms/deepseek/count_tokens/__init__.py b/tests/unit/llms/deepseek/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/deepseek/count_tokens/test_token_counter.py b/tests/unit/llms/deepseek/count_tokens/test_token_counter.py new file mode 100644 index 00000000000..2ba99225685 --- /dev/null +++ b/tests/unit/llms/deepseek/count_tokens/test_token_counter.py @@ -0,0 +1,306 @@ +import json +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +import respx +from pydantic import JsonValue + +import litellm +from litellm.llms.anthropic.common_utils import AnthropicError +from litellm.llms.deepseek.count_tokens.token_counter import DEEPSEEK_COUNT_TIMEOUT_SECONDS, DeepSeekTokenCounter +from litellm.llms.deepseek.messages.transformation import DeepSeekAnthropicMessagesConfig +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse + +_DEFAULT_COUNT_ROUTE: Final = "https://api.deepseek.com/anthropic/v1/messages/count_tokens" +_WEATHER_IN_SF: Final[list[dict[str, JsonValue]]] = [{"role": "user", "content": "Weather in SF?"}] +_GET_WEATHER: Final[dict[str, JsonValue]] = { + "name": "get_weather", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}}, +} +_GET_WEATHER_FUNCTION: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": "get_weather", "parameters": _GET_WEATHER["input_schema"]}, +} +_SPAWN_AGENT_NAMESPACE: Final[dict[str, JsonValue]] = { + "type": "namespace", + "name": "multi_agent_v1", + "tools": [{"type": "function", "name": "spawn_agent", "parameters": {"type": "object", "properties": {}}}], +} + + +@pytest.fixture(autouse=True) +def _deepseek_env_unset(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ("DEEPSEEK_API_BASE", "DEEPSEEK_ANTHROPIC_API_BASE", "DEEPSEEK_API_KEY"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + +async def _count( + reply: httpx.Response | Exception, + count_route: str = _DEFAULT_COUNT_ROUTE, + litellm_params: Mapping[str, JsonValue] = MappingProxyType({"api_key": "sk-deepseek-test"}), + messages_format: TokenCountMessagesFormat = "anthropic", + tools: list[dict[str, JsonValue]] | None = None, + chat_params: Mapping[str, JsonValue] = MappingProxyType({}), + count_options: Mapping[str, JsonValue] = MappingProxyType({}), + messages: list[dict[str, JsonValue]] = _WEATHER_IN_SF, +) -> tuple[TokenCountResponse | None, respx.Route]: + with respx.mock: + route: Final = respx.post(count_route).mock( + return_value=reply if isinstance(reply, httpx.Response) else None, + side_effect=None if isinstance(reply, httpx.Response) else reply, + ) + result: Final = await DeepSeekTokenCounter().count_tokens( + model_to_use="deepseek-v4-pro", + messages=messages, + contents=None, + deployment={"litellm_params": dict(litellm_params)}, + request_model="deepseek-pro", + tools=tools, + messages_format=messages_format, + chat_params=chat_params, + count_options=count_options, + ) + return result, route + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "api_base", + ( + None, + "https://api.deepseek.com", + "https://api.deepseek.com/beta", + "https://api.deepseek.com/anthropic/v1/messages", + ), +) +async def test_count_posts_under_the_base_messages_inference_posts_to(api_base: str | None) -> None: + messages_url: Final = DeepSeekAnthropicMessagesConfig().get_complete_url( + api_base=api_base, api_key=None, model="deepseek-v4-pro", optional_params={}, litellm_params={} + ) + + result, route = await _count( + httpx.Response(200, json={"input_tokens": 91}), + count_route=f"{messages_url}/count_tokens", + litellm_params={"api_key": "sk-deepseek-test", **({} if api_base is None else {"api_base": api_base})}, + ) + + assert route.calls.last.request.headers["x-api-key"] == "sk-deepseek-test" + assert result is not None + assert (result.total_tokens, result.error) == (91, False) + + +@pytest.mark.asyncio +async def test_count_without_a_key_sends_no_api_key_header() -> None: + _, route = await _count(httpx.Response(200, json={"input_tokens": 91}), litellm_params={}) + + assert "x-api-key" not in route.calls.last.request.headers + + +@pytest.mark.asyncio +async def test_count_waits_no_longer_than_the_deployment_timeout() -> None: + _, route = await _count( + httpx.Response(200, json={"input_tokens": 91}), + litellm_params={"api_key": "sk-deepseek-test", "timeout": 7}, + ) + + assert route.calls.last.request.extensions["timeout"] == {"connect": 7, "read": 7, "write": 7, "pool": 7} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deployment_timeout", ({}, {"timeout": 600})) +async def test_count_waits_no_longer_than_the_count_cap_when_the_deployment_allows_longer( + deployment_timeout: Mapping[str, JsonValue], +) -> None: + _, route = await _count( + httpx.Response(200, json={"input_tokens": 91}), + litellm_params={"api_key": "sk-deepseek-test", **deployment_timeout}, + ) + + assert set(route.calls.last.request.extensions["timeout"].values()) == {DEEPSEEK_COUNT_TIMEOUT_SECONDS} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("messages_format", ("anthropic", "openai")) +async def test_count_sends_tools_without_the_custom_type_deepseek_refuses( + messages_format: TokenCountMessagesFormat, +) -> None: + tool: Final[dict[str, JsonValue]] = ( + {**_GET_WEATHER, "type": "custom"} + if messages_format == "anthropic" + else { + "type": "function", + "function": {"name": "get_weather", "parameters": _GET_WEATHER["input_schema"]}, + } + ) + + _, route = await _count( + httpx.Response(200, json={"input_tokens": 91}), messages_format=messages_format, tools=[tool] + ) + + assert json.loads(route.calls.last.request.content)["tools"] == [_GET_WEATHER] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tools", "chat_params", "expected_tool_fields"), + ( + ([_GET_WEATHER_FUNCTION, _SPAWN_AGENT_NAMESPACE], {}, {"tools": [_GET_WEATHER]}), + ([_SPAWN_AGENT_NAMESPACE], {"tool_choice": "required"}, {}), + ), +) +async def test_chat_count_leaves_out_the_tools_deepseek_chat_inference_drops( + tools: list[dict[str, JsonValue]], + chat_params: Mapping[str, JsonValue], + expected_tool_fields: Mapping[str, JsonValue], +) -> None: + _, route = await _count( + httpx.Response(200, json={"input_tokens": 91}), messages_format="openai", tools=tools, chat_params=chat_params + ) + + sent: Final = json.loads(route.calls.last.request.content) + assert {name: sent[name] for name in ("tools", "tool_choice") if name in sent} == expected_tool_fields + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("chat_params", "expected_options"), + ( + ({}, {}), + ({"reasoning_effort": "none"}, {"thinking": {"type": "disabled"}}), + ({"reasoning_effort": "max"}, {"thinking": {"type": "enabled"}, "output_config": {"effort": "max"}}), + ({"reasoning_effort": "unlisted"}, {"thinking": {"type": "enabled"}}), + ({"thinking": {"type": "enabled", "budget_tokens": 1024}}, {"thinking": {"type": "enabled"}}), + ), +) +async def test_chat_count_maps_thinking_and_effort_the_way_deepseek_chat_inference_does( + chat_params: Mapping[str, JsonValue], expected_options: Mapping[str, JsonValue] +) -> None: + _, route = await _count( + httpx.Response(200, json={"input_tokens": 91}), messages_format="openai", chat_params=chat_params + ) + + sent: Final = json.loads(route.calls.last.request.content) + assert {name: sent[name] for name in ("thinking", "output_config") if name in sent} == expected_options + + +@pytest.mark.asyncio +async def test_messages_count_sends_the_thinking_messages_inference_sends() -> None: + _, route = await _count( + httpx.Response(200, json={"input_tokens": 91}), + count_options={"thinking": {"type": "disabled"}, "tool_choice": {"type": "any"}}, + ) + + sent: Final = json.loads(route.calls.last.request.content) + assert (sent["thinking"], sent["tool_choice"]) == ({"type": "disabled"}, {"type": "any"}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", (404, 405)) +async def test_a_missing_count_route_leaves_the_count_to_the_estimate(status_code: int) -> None: + result, route = await _count(httpx.Response(status_code, text="no such route")) + + assert route.called + assert result is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("api_base", ("https://inference.baseten.co/v1", "https://gateway.example/deepseek")) +async def test_a_deepseek_model_on_another_host_is_left_to_the_estimate_without_a_count_call(api_base: str) -> None: + messages_url: Final = DeepSeekAnthropicMessagesConfig().get_complete_url( + api_base=api_base, api_key=None, model="deepseek-v4-pro", optional_params={}, litellm_params={} + ) + + result, route = await _count( + httpx.Response(401, json={"error": "Authentication required."}), + count_route=f"{messages_url}/count_tokens", + litellm_params={"api_key": "sk-other-host", "api_base": api_base}, + ) + + assert not route.called + assert result is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "reply", + ( + httpx.Response(408, text="timed out"), + httpx.Response(422, json={"error": {"message": "Failed to deserialize the JSON body"}}), + httpx.Response(429, json={"error": {"message": "Rate limit reached"}}), + httpx.Response(500, text="internal error"), + httpx.Response(503, text="overloaded"), + httpx.ConnectError("connection refused"), + ), +) +async def test_a_count_route_that_cannot_count_leaves_the_count_to_the_estimate( + reply: httpx.Response | Exception, +) -> None: + result, route = await _count(reply) + + assert route.called + assert result is None + + +@pytest.mark.asyncio +async def test_a_rejected_count_answers_with_deepseeks_status() -> None: + result, _ = await _count(httpx.Response(401, json={"error": {"message": "Authentication Fails"}})) + + assert result is not None + assert (result.error, result.status_code, result.total_tokens) == (True, 401, 0) + assert "Authentication Fails" in (result.error_message or "") + + +@pytest.mark.asyncio +async def test_count_without_messages_leaves_the_count_to_the_estimate() -> None: + result, route = await _count(httpx.Response(200, json={"input_tokens": 91}), messages=[]) + + assert (result, route.called) == (None, False) + + +@pytest.mark.asyncio +async def test_messages_count_rejects_an_effort_messages_inference_rejects_without_calling_deepseek() -> None: + options: Final[dict[str, JsonValue]] = {"output_config": {"effort": "unlisted"}} + with pytest.raises(AnthropicError) as rejected: + DeepSeekAnthropicMessagesConfig().transform_anthropic_messages_request( + model="deepseek-v4-pro", + messages=[dict(message) for message in _WEATHER_IN_SF], + anthropic_messages_optional_request_params={**options, "max_tokens": 4096}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + result, route = await _count(httpx.Response(200, json={"input_tokens": 91}), count_options=options) + + assert route.called is False + assert result is not None + assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) + assert (rejected.value.status_code, str(rejected.value.message)) == (400, result.error_message) + + +@pytest.mark.asyncio +async def test_chat_count_rejects_a_tool_call_chat_inference_rejects_without_calling_deepseek() -> None: + messages: Final[list[dict[str, JsonValue]]] = [ + *_WEATHER_IN_SF, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "sunny?"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ] + + result, route = await _count( + httpx.Response(200, json={"input_tokens": 91}), messages_format="openai", messages=messages + ) + + assert route.called is False + assert result is not None + assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) diff --git a/tests/unit/llms/deepseek/test_deepseek_common_utils.py b/tests/unit/llms/deepseek/test_deepseek_common_utils.py new file mode 100644 index 00000000000..8916ff404ee --- /dev/null +++ b/tests/unit/llms/deepseek/test_deepseek_common_utils.py @@ -0,0 +1,15 @@ +import pytest + +import litellm + + +def test_deepseek_model_listing_with_a_provider_check_lists_the_cost_map_models_offline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-deepseek-test") + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + + listed = litellm.get_valid_models(custom_llm_provider="deepseek", check_provider_endpoint=True) + + assert listed + assert sorted(listed) == sorted(litellm.models_by_provider["deepseek"]) diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 32d68249858..eb300d823aa 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1,7 +1,7 @@ import copy import threading import time -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from contextlib import contextmanager from enum import Enum from typing import Final @@ -34,7 +34,7 @@ from litellm.llms.vertex_ai.common_utils import ( ) from litellm.llms.vertex_ai.count_tokens import handler as vertex_count_tokens_handler from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig -from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse @pytest.mark.asyncio @@ -1875,6 +1875,52 @@ async def test_vertex_claude_count_forwards_the_context_edits_and_beta_vertex_in assert headers["anthropic-beta"] == expected_beta +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_drops", "thinking_sent"), + (({}, True), ({"additional_drop_params": ["thinking"]}, False)), +) +async def test_vertex_claude_count_sends_the_thinking_vertex_inference_sends( + deployment_drops: Mapping[str, JsonValue], thinking_sent: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler + + class FakeResponse: + status_code = 200 + + def json(self) -> dict[str, int]: + return {"input_tokens": 21} + + class FakeHttpClient: + posted: tuple[dict[str, object], ...] = () + + async def post( + self, url: str, headers: dict[str, str], json: dict[str, object], timeout: float + ) -> FakeResponse: + self.posted = (*self.posted, json) + return FakeResponse() + + fake_http_client: Final = FakeHttpClient() + application_default_credentials: Final = MagicMock(token="fake-token", valid=True, expiry=None) + monkeypatch.setattr("google.auth.default", lambda *args, **kwargs: (application_default_credentials, "test-project")) + with patch.object(handler, "get_async_httpx_client", return_value=fake_http_client): + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="claude-opus-4-8", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={ + "litellm_params": {"vertex_project": "test-project", "vertex_location": "us-east5", **deployment_drops} + }, + messages_format="anthropic", + count_options={"thinking": {"type": "enabled", "budget_tokens": 2048}}, + ) + + assert result is not None + assert result.total_tokens == 21 + (body,) = fake_http_client.posted + assert ("thinking" in body) is thinking_sent + + @pytest.mark.asyncio async def test_vertex_claude_count_with_an_edit_fetches_the_remote_beta_config_off_the_event_loop( monkeypatch: pytest.MonkeyPatch, @@ -1939,9 +1985,9 @@ _PARTNER_RESPONSES_TOOL: Final[dict[str, object]] = { } -async def _partner_counted_bodies( - model: str, chat_params: dict[str, JsonValue] | None = None -) -> tuple[dict[str, object], ...]: +async def _partner_count( + model: str, messages: list[dict[str, object]], chat_params: dict[str, JsonValue] | None = None +) -> tuple[TokenCountResponse | None, tuple[dict[str, object], ...]]: from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler class FakeResponse: @@ -1970,16 +2016,43 @@ async def _partner_counted_bodies( ): result: Final = await VertexAITokenCounter().count_tokens( model_to_use=model, - messages=copy.deepcopy(_PARTNER_CHAT_TOOL_TURNS), + messages=copy.deepcopy(messages), contents=None, deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-east5"}}, tools=[_PARTNER_RESPONSES_TOOL], chat_params=chat_params or {}, ) + return result, fake_http_client.posted_bodies + +async def _partner_counted_bodies( + model: str, chat_params: dict[str, JsonValue] | None = None +) -> tuple[dict[str, object], ...]: + result, bodies = await _partner_count(model, _PARTNER_CHAT_TOOL_TURNS, chat_params) assert result is not None assert result.total_tokens == 37 - return fake_http_client.posted_bodies + return bodies + + +@pytest.mark.asyncio +async def test_vertex_claude_count_rejects_a_tool_call_chat_inference_rejects_without_calling_vertex() -> None: + unparseable_arguments: Final = [ + _PARTNER_CHAT_TOOL_TURNS[0], + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "sunny?"}} + ], + }, + _PARTNER_CHAT_TOOL_TURNS[2], + ] + + result, bodies = await _partner_count("claude-opus-5-5", unparseable_arguments) + + assert bodies == () + assert result is not None + assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) @pytest.mark.asyncio diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_handler.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_handler.py index c03f08ab610..a0ecc60d6be 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_handler.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_handler.py @@ -3,6 +3,10 @@ from typing import Final import pytest from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_headers +from litellm.types.llms.anthropic import ( + ANTHROPIC_THINKING_BINDING_CONTROLS_BETA_HEADER, + ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER, +) _HISTORY: Final = ({"role": "user", "content": "Hello"},) @@ -58,3 +62,23 @@ def test_count_tokens_on_a_tool_change_request_sends_the_tool_change_beta( "Authorization": "Bearer token", "anthropic-beta": expected_beta, } + + +@pytest.mark.parametrize( + ("thinking", "expected_beta"), + [ + ({"type": "adaptive", "display": "updates"}, ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER), + ( + {"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}}, + ANTHROPIC_THINKING_BINDING_CONTROLS_BETA_HEADER, + ), + ], + ids=["display_updates", "block_binding"], +) +def test_count_tokens_with_a_beta_thinking_field_sends_the_beta_vertex_inference_sends( + thinking: dict[str, object], expected_beta: str +) -> None: + assert count_tokens_headers("token", {"messages": list(_HISTORY), "thinking": thinking}) == { + "Authorization": "Bearer token", + "anthropic-beta": expected_beta, + } diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index d4bdd0ee496..d9ddf0f131a 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -440,6 +440,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): request_litellm_params=None, litellm_call_id=None, client_oauth_token=None, + count_options=None, ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "claude-3-sonnet-20240229" @@ -522,6 +523,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): request_litellm_params=None, litellm_call_id=None, client_oauth_token=None, + count_options=None, ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "gpt-4" @@ -1344,6 +1346,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): request_litellm_params=None, litellm_call_id=None, client_oauth_token=None, + count_options=None, ): raise ProxyException( message='{"detail":{"message":"Input is too long for requested model."}}', @@ -1403,6 +1406,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): request_litellm_params=None, litellm_call_id=None, client_oauth_token=None, + count_options=None, ): raise ProxyException( message='{"Message":"Bearer Token has expired"}', @@ -1462,6 +1466,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): request_litellm_params=None, litellm_call_id=None, client_oauth_token=None, + count_options=None, ): raise ProxyException( message="Rate limit exceeded", @@ -1775,6 +1780,7 @@ class _ScriptedCounter(BaseTokenCounter): output_format: JsonValue = None, context_management: object = None, chat_params: object = None, + count_options: object = None, ) -> TokenCountResponse | None: self.calls.append((model_to_use, messages, tools, system)) self.regions.append(None if deployment is None else deployment["litellm_params"].get("aws_region_name")) @@ -1817,7 +1823,7 @@ async def _ladder( counter: _ScriptedCounter, settings: TokenCountSettings, local_estimate_allowed: bool = True, - forwarded_edits: bool = False, + forwards_request_fields: bool = False, ) -> TokenCountResponse | None: return await _provider_token_count( count_with=partial( @@ -1835,7 +1841,7 @@ async def _ladder( request_model=_COUNT_REQUEST.model, settings=settings, local_estimate_allowed=local_estimate_allowed, - forwarded_edits=forwarded_edits, + forwards_request_fields=forwards_request_fields, ) @@ -2093,7 +2099,7 @@ async def test_a_pinned_token_count_model_rejecting_forwarded_edits_answers_with await _ladder( counter, TokenCountSettings(framing=None, count_model=replace(_OPUS_IN_US_EAST_1, configured=True)), - forwarded_edits=True, + forwards_request_fields=True, ) assert raised.value.code == str(status_code) @@ -4428,13 +4434,14 @@ async def test_messages_count_tokens_on_a_responses_api_bridge_rejects_an_invali class _CountTokensPeer: - def __init__(self, reply: Mapping[str, JsonValue]) -> None: + def __init__(self, reply: Mapping[str, JsonValue], status_code: int = 200) -> None: self._reply: Final = reply + self._status_code: Final = status_code self.requests: tuple[httpx.Request, ...] = () def __call__(self, request: httpx.Request) -> httpx.Response: self.requests = (*self.requests, request) - return httpx.Response(200, json=dict(self._reply), request=request) + return httpx.Response(self._status_code, json=dict(self._reply), request=request) class _BedrockInvokeModel(BaseModel): @@ -4463,12 +4470,13 @@ async def _count_on_a_bridged_provider_count_api( bedrock_model: str = "bedrock/us.amazon.nova-2-lite-v1:0", peer_reply: Mapping[str, JsonValue] = MappingProxyType({"inputTokens": 57}), model_info: Mapping[str, JsonValue] = MappingProxyType({}), + peer_status: int = 200, ) -> tuple[object, tuple[httpx.Request, ...]]: from litellm.llms.bedrock.count_tokens import handler as bedrock_count_tokens_handler from litellm.llms.bedrock.count_tokens import mantle_handler as bedrock_mantle_count_tokens_handler monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) - peer: Final = _CountTokensPeer(peer_reply) + peer: Final = _CountTokensPeer(peer_reply, peer_status) router: Final = Router( model_list=[ { @@ -4582,6 +4590,42 @@ async def test_messages_count_tokens_forwards_bedrock_supported_edits_and_their_ ) +@pytest.mark.parametrize( + ("count_options", "region_hints"), + [({"thinking": {"type": "enabled", "budget_tokens": 1024}}, 0), ({}, 1)], + ids=["count_options_forwarded", "no_count_options"], +) +@pytest.mark.asyncio +async def test_a_mantle_count_rejecting_the_forwarded_count_options_answers_without_the_region_hint( + count_options: Mapping[str, JsonValue], region_hints: int, monkeypatch: pytest.MonkeyPatch +) -> None: + response, sent = await _count_on_a_bridged_provider_count_api( + {"messages": [{"role": "user", "content": "Read the README"}], **count_options}, + monkeypatch, + _BEDROCK_CLAUDE_MODEL, + MappingProxyType( + {"type": "error", "error": {"type": "invalid_request_error", "message": "thinking: invalid value"}} + ), + MappingProxyType( + { + "token_count_model": { + "model": "us.anthropic.claude-sonnet-5-5", + "litellm_provider": "bedrock_mantle", + "region": "us-west-2", + } + } + ), + peer_status=400, + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert len(sent) == 1 + assert ("thinking" in json.loads(sent[0].content)) == bool(count_options) + message: Final = json.loads(bytes(response.body))["error"]["message"] + assert message.count("counts tokens with us.anthropic.claude-sonnet-5-5 in us-west-2") == region_hints + + @pytest.mark.asyncio async def test_messages_count_tokens_sends_no_context_management_to_bedrock_when_no_edit_is_bedrock_supported( monkeypatch: pytest.MonkeyPatch, @@ -5051,7 +5095,9 @@ async def test_framed_stand_in_counts_a_prompt_like_the_same_text_sent_as_a_user @pytest.mark.asyncio @pytest.mark.parametrize( - "model", ["fireworks_ai/accounts/fireworks/models/glm-5", "deepseek/deepseek-v4-flash"], ids=["plain", "framed"] + "model", + ["fireworks_ai/accounts/fireworks/models/glm-5", "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash"], + ids=["plain", "framed"], ) async def test_call_endpoint_counts_a_prompt_on_a_provider_without_a_count_api_as_raw_text( model: str, @@ -5773,3 +5819,375 @@ def test_utils_token_counter_line_stays_one_line_when_the_model_carries_a_newlin assert line["request_model"] == forged_model assert "\n" not in message assert message.startswith(f'token count - litellm_call_id="count log" request_model={json.dumps(forged_model)} ') + + +_COUNTED_OPTION_FIELDS: Final = ("thinking", "tool_choice", "output_config") +_DEEPSEEK_COUNT_ROUTE: Final = "https://api.deepseek.com/anthropic/v1/messages/count_tokens" +_CUSTOM_READ_FILE_TOOL: Final[dict[str, JsonValue]] = {**_READ_FILE_TOOL, "type": "custom"} + + +def _messages_inference_fields( + config: object, model: str, body: Mapping[str, JsonValue], fields: tuple[str, ...] +) -> dict[str, object]: + from litellm.types.router import GenericLiteLLMParams + + inference: Final = config.transform_anthropic_messages_request( + model=model, + messages=json.loads(json.dumps(body["messages"])), + anthropic_messages_optional_request_params={ + **{key: json.loads(json.dumps(value)) for key, value in body.items() if key not in ("model", "messages")}, + "max_tokens": 4096, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + return {field: inference.get(field) for field in fields} + + +async def _count_messages_body( + monkeypatch: pytest.MonkeyPatch, + deployment_params: Mapping[str, JsonValue], + body: Mapping[str, JsonValue], + count_route: str, + provider_reply: httpx.Response, +) -> tuple[object, respx.Route]: + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.delenv("DEEPSEEK_API_BASE", raising=False) + monkeypatch.delenv("DEEPSEEK_ANTHROPIC_API_BASE", raising=False) + monkeypatch.setattr( + proxy_server, + "llm_router", + Router(model_list=[{"model_name": "counted", "litellm_params": dict(deployment_params)}]), + ) + monkeypatch.setattr(proxy_server, "proxy_config", proxy_server.ProxyConfig()) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(count_route).mock(return_value=provider_reply) + response: Final = await anthropic_count_tokens( + _json_request({"model": "counted", **body}, "python-httpx/0.28.1"), Response(), UserAPIKeyAuth() + ) + return response, route + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "options"), + ( + ("claude-sonnet-4-5", {"tool_choice": {"type": "any"}}), + ("claude-sonnet-4-5", {"thinking": {"type": "enabled", "budget_tokens": 2048}}), + ("claude-haiku-5-5", {"thinking": {"type": "enabled", "budget_tokens": 2048}, "tool_choice": {"type": "auto"}}), + ("claude-opus-5-5", {"thinking": {"type": "adaptive"}, "output_config": {"effort": "low"}}), + ), +) +async def test_messages_count_tokens_counts_the_thinking_tool_choice_and_effort_messages_inference_sends( + monkeypatch: pytest.MonkeyPatch, model: str, options: Mapping[str, JsonValue] +) -> None: + from litellm.llms.anthropic.messages.native.transformation import AnthropicMessagesConfig + + body: Final[dict[str, JsonValue]] = { + "messages": [{"role": "user", "content": "Read the README"}], + "tools": [_READ_FILE_TOOL], + **options, + } + + response, route = await _count_messages_body( + monkeypatch, + {"model": f"anthropic/{model}", "api_key": "sk-ant-test"}, + body, + "https://api.anthropic.com/v1/messages/count_tokens", + httpx.Response(200, json={"input_tokens": 120}), + ) + + sent: Final = json.loads(route.calls.last.request.content) + assert response == {"input_tokens": 120, "litellm_count_source": "provider"} + assert {field: sent.get(field) for field in _COUNTED_OPTION_FIELDS} == _messages_inference_fields( + AnthropicMessagesConfig(), model, body, _COUNTED_OPTION_FIELDS + ) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_leaves_out_an_option_the_request_drops(monkeypatch: pytest.MonkeyPatch) -> None: + response, route = await _count_messages_body( + monkeypatch, + {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-ant-test"}, + { + "messages": [{"role": "user", "content": "Read the README"}], + "tools": [_READ_FILE_TOOL], + "tool_choice": {"type": "any"}, + "thinking": {"type": "enabled", "budget_tokens": 2048}, + "additional_drop_params": ["thinking"], + }, + "https://api.anthropic.com/v1/messages/count_tokens", + httpx.Response(200, json={"input_tokens": 120}), + ) + + sent: Final = json.loads(route.calls.last.request.content) + assert response == {"input_tokens": 120, "litellm_count_source": "provider"} + assert (sent["tool_choice"], "thinking" in sent) == ({"type": "any"}, False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("body_options", "sent_thinking_type"), + ( + ({}, "disabled"), + ({"thinking": {"type": "enabled"}}, "enabled"), + ({"additional_drop_params": ["thinking"]}, None), + ), +) +async def test_messages_count_tokens_counts_the_thinking_the_deployment_sets_unless_the_request_overrides_it( + monkeypatch: pytest.MonkeyPatch, body_options: Mapping[str, JsonValue], sent_thinking_type: str | None +) -> None: + _, route = await _count_messages_body( + monkeypatch, + {"model": "deepseek/deepseek-v4-pro", "api_key": "sk-deepseek-test", "thinking": {"type": "disabled"}}, + {"messages": [{"role": "user", "content": "Read the README"}], **body_options}, + _DEEPSEEK_COUNT_ROUTE, + httpx.Response(200, json={"input_tokens": 91}), + ) + + sent: Final = json.loads(route.calls.last.request.content) + assert (sent.get("thinking") or {}).get("type") == sent_thinking_type + + +@pytest.mark.asyncio +async def test_messages_count_tokens_answers_400_for_a_thinking_budget_messages_inference_rejects( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.exceptions import BadRequestError + from litellm.llms.anthropic.messages.native.transformation import AnthropicMessagesConfig + + body: Final[dict[str, JsonValue]] = { + "messages": [{"role": "user", "content": "Read the README"}], + "thinking": {"type": "enabled", "budget_tokens": "lots"}, + } + with pytest.raises(BadRequestError) as rejected: + _messages_inference_fields(AnthropicMessagesConfig(), "claude-opus-5-5", body, _COUNTED_OPTION_FIELDS) + + response, route = await _count_messages_body( + monkeypatch, + {"model": "anthropic/claude-opus-5-5", "api_key": "sk-ant-test"}, + body, + "https://api.anthropic.com/v1/messages/count_tokens", + httpx.Response(200, json={"input_tokens": 120}), + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert rejected.value.message in json.loads(response.body)["error"]["message"] + assert not route.called + + +@pytest.mark.asyncio +async def test_messages_count_tokens_counts_deepseek_on_its_anthropic_count_route( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.llms.deepseek.messages.transformation import DeepSeekAnthropicMessagesConfig + + body: Final[dict[str, JsonValue]] = { + "messages": [{"role": "user", "content": "Read the README"}], + "tools": [_CUSTOM_READ_FILE_TOOL], + "tool_choice": {"type": "any"}, + "thinking": {"type": "enabled"}, + "output_config": {"effort": "max"}, + } + + response, route = await _count_messages_body( + monkeypatch, + {"model": "deepseek/deepseek-v4-pro", "api_key": "sk-deepseek-test"}, + body, + _DEEPSEEK_COUNT_ROUTE, + httpx.Response(200, json={"input_tokens": 91}), + ) + + sent_request: Final = route.calls.last.request + sent: Final = json.loads(sent_request.content) + inference: Final = _messages_inference_fields( + DeepSeekAnthropicMessagesConfig(), "deepseek-v4-pro", body, ("tool_choice", "output_config", "tools", "thinking") + ) + assert response == {"input_tokens": 91, "litellm_count_source": "provider"} + assert sent_request.headers["x-api-key"] == "sk-deepseek-test" + assert (sent["tool_choice"], sent.get("output_config"), sent["tools"], sent["thinking"]["type"]) == ( + inference["tool_choice"], + inference["output_config"], + inference["tools"], + inference["thinking"]["type"], + ) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_on_deepseek_on_another_host_answers_the_estimate_without_a_count_call( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response, route = await _count_messages_body( + monkeypatch, + {"model": "deepseek/deepseek-v4-pro", "api_key": "sk-other-host", "api_base": "https://gateway.example/v1"}, + {"messages": [{"role": "user", "content": "Read the README"}]}, + "https://gateway.example/anthropic/v1/messages/count_tokens", + httpx.Response(401, json={"error": "Authentication required."}), + ) + + assert not route.called + assert isinstance(response, dict) + assert response["litellm_count_source"] == "local_estimate" + assert response["input_tokens"] > 0 + + +async def _count_deepseek_responses( + monkeypatch: pytest.MonkeyPatch, + deployment_options: Mapping[str, JsonValue], + responses_options: Mapping[str, JsonValue] | None, + request_options: Mapping[str, JsonValue] = MappingProxyType({}), +) -> tuple[TokenCountResponse, dict[str, JsonValue]]: + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.delenv("DEEPSEEK_API_BASE", raising=False) + monkeypatch.delenv("DEEPSEEK_ANTHROPIC_API_BASE", raising=False) + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "deepseek-counted", + "litellm_params": { + "model": "deepseek/deepseek-v4-pro", + "api_key": "sk-deepseek-test", + **deployment_options, + }, + } + ] + ), + ) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + with respx.mock: + route: Final = respx.post(_DEEPSEEK_COUNT_ROUTE).mock(return_value=httpx.Response(200, json={"input_tokens": 91})) + result: Final = await count_request_tokens( + request=TokenCountRequest.model_validate( + {"model": "deepseek-counted", "messages": [{"role": "user", "content": "Say hi."}], **request_options} + ), + call_endpoint=True, + responses_request=None + if responses_options is None + else { + "model": "deepseek-counted", + "input": "Say hi.", + "tools": [_WEATHER_TOOL], + **responses_options, + }, + ) + return result, json.loads(route.calls.last.request.content) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("reasoning", "expected_options"), + ( + ({}, {}), + ({"reasoning": {"effort": "none"}}, {"thinking": {"type": "disabled"}}), + ({"reasoning": {"effort": "high"}}, {"thinking": {"type": "enabled"}, "output_config": {"effort": "high"}}), + ({"reasoning": {"effort": "max"}}, {"thinking": {"type": "enabled"}, "output_config": {"effort": "max"}}), + ), +) +async def test_responses_count_on_deepseek_maps_reasoning_effort_the_way_deepseek_chat_inference_does( + monkeypatch: pytest.MonkeyPatch, reasoning: Mapping[str, JsonValue], expected_options: Mapping[str, JsonValue] +) -> None: + result, sent = await _count_deepseek_responses(monkeypatch, {}, reasoning) + + assert (result.total_tokens, result.count_source) == (91, "provider") + assert {field: sent[field] for field in ("thinking", "output_config") if field in sent} == expected_options + assert all("type" not in tool for tool in sent["tools"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_options", "responses_options", "expected_options"), + ( + ({"thinking": {"type": "disabled"}}, {}, {"thinking": {"type": "disabled"}}), + ({"reasoning_effort": "none"}, {}, {"thinking": {"type": "disabled"}}), + ({"thinking": {"type": "disabled"}}, {"additional_drop_params": ["thinking"]}, {}), + ({"thinking": {"type": "disabled"}}, {"thinking": {"type": "enabled"}}, {"thinking": {"type": "enabled"}}), + ), +) +async def test_responses_count_on_deepseek_counts_the_thinking_inference_sends_from_the_request_or_deployment( + monkeypatch: pytest.MonkeyPatch, + deployment_options: Mapping[str, JsonValue], + responses_options: Mapping[str, JsonValue], + expected_options: Mapping[str, JsonValue], +) -> None: + result, sent = await _count_deepseek_responses(monkeypatch, deployment_options, responses_options) + + assert (result.total_tokens, result.count_source) == (91, "provider") + assert {field: sent[field] for field in ("thinking", "output_config") if field in sent} == expected_options + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_options", "expected_options"), + ( + ({"thinking": {"type": "disabled"}}, {"thinking": {"type": "disabled"}}), + ({"reasoning_effort": "high"}, {"thinking": {"type": "enabled"}, "output_config": {"effort": "high"}}), + ({"thinking": {"type": "disabled"}, "additional_drop_params": ["thinking"]}, {}), + ), +) +async def test_chat_count_on_deepseek_counts_the_thinking_the_deployment_sets_unless_it_drops_it( + monkeypatch: pytest.MonkeyPatch, + deployment_options: Mapping[str, JsonValue], + expected_options: Mapping[str, JsonValue], +) -> None: + result, sent = await _count_deepseek_responses(monkeypatch, deployment_options, None) + + assert (result.total_tokens, result.count_source) == (91, "provider") + assert {field: sent[field] for field in ("thinking", "output_config") if field in sent} == expected_options + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_options", "request_options", "expected_options"), + ( + ({}, {"reasoning_effort": "none"}, {"thinking": {"type": "disabled"}}), + ({}, {"reasoning_effort": "max"}, {"thinking": {"type": "enabled"}, "output_config": {"effort": "max"}}), + ({"thinking": {"type": "disabled"}}, {"thinking": {"type": "enabled"}}, {"thinking": {"type": "enabled"}}), + ), +) +async def test_chat_count_on_deepseek_counts_the_thinking_the_request_sends_over_the_deployment_default( + monkeypatch: pytest.MonkeyPatch, + deployment_options: Mapping[str, JsonValue], + request_options: Mapping[str, JsonValue], + expected_options: Mapping[str, JsonValue], +) -> None: + result, sent = await _count_deepseek_responses(monkeypatch, deployment_options, None, request_options) + + assert (result.total_tokens, result.count_source) == (91, "provider") + assert {field: sent[field] for field in ("thinking", "output_config") if field in sent} == expected_options + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_drops", "sent_options"), + ( + ({}, (True, None)), + ({"additional_drop_params": ["tool_choice"]}, (False, "enabled")), + ), +) +async def test_messages_count_tokens_drops_what_the_deployment_drops_unless_the_request_sends_its_own_list( + monkeypatch: pytest.MonkeyPatch, request_drops: Mapping[str, JsonValue], sent_options: tuple[bool, str | None] +) -> None: + _, route = await _count_messages_body( + monkeypatch, + {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-ant-test", "additional_drop_params": ["thinking"]}, + { + "messages": [{"role": "user", "content": "Read the README"}], + "tools": [_READ_FILE_TOOL], + "tool_choice": {"type": "auto"}, + "thinking": {"type": "enabled", "budget_tokens": 2048}, + **request_drops, + }, + "https://api.anthropic.com/v1/messages/count_tokens", + httpx.Response(200, json={"input_tokens": 120}), + ) + + sent: Final = json.loads(route.calls.last.request.content) + assert ("tool_choice" in sent, (sent.get("thinking") or {}).get("type")) == sent_options