diff --git a/litellm/llms/anthropic/messages/responses_adapters/handler.py b/litellm/llms/anthropic/messages/responses_adapters/handler.py index a4114d03ab4..0c5111eefca 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/handler.py +++ b/litellm/llms/anthropic/messages/responses_adapters/handler.py @@ -8,6 +8,7 @@ from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Mappin from typing import Final, TypeAlias from pydantic import TypeAdapter, ValidationError +from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger @@ -41,7 +42,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) from litellm.types.llms.openai import ResponsesAPIResponse -from litellm.types.utils import CountContextManagement +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat from litellm.utils import ProviderConfigManager from ..adapters.handler import ( @@ -158,6 +159,33 @@ async def count_view_messages( return messages if result is None else result.messages +async def count_messages_as_sent( + *, + model: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + system: object, + messages_format: TokenCountMessagesFormat, + context_management: CountContextManagement | None, +) -> Sequence[Mapping[str, object]] | AnthropicContextManagementError: + match messages_format: + case "openai": + return messages + case "anthropic": + try: + return await count_view_messages( + model=model, + messages=messages, + tools=tools, + system=system, + context_management=context_management, + ) + except AnthropicContextManagementError as error: + return error + case _: + assert_never(messages_format) + + def _with_applied_edits( response: AnthropicMessagesResponse, polyfill_result: PolyfillResult | None ) -> AnthropicMessagesResponse: diff --git a/litellm/llms/anthropic/messages/responses_adapters/transformation.py b/litellm/llms/anthropic/messages/responses_adapters/transformation.py index 1d186c8aee4..a63884f9afa 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/messages/responses_adapters/transformation.py @@ -589,6 +589,12 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if block.get("type") == "text" and isinstance(text := block.get("text"), str) and text ) + @classmethod + def system_instructions(cls, system: object) -> str | None: + if system is None or isinstance(system, str): + return system + return cls.system_blocks_to_instructions(system) or None + @staticmethod def _resolve_output_format_strict(strict: object) -> bool: return True if strict is None else bool(strict) diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 76816700e20..c454c37c01b 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -19,7 +19,7 @@ from typing_extensions import TypeIs, assert_never # noqa: TID251 # these are import litellm from litellm._logging import verbose_logger from litellm.exceptions import UnsupportedParamsError -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import is_str_keyed_dict from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -33,12 +33,11 @@ from litellm.litellm_core_utils.reasoning_effort_capability import ( nearest_declared_reasoning_effort, ) from litellm.litellm_core_utils.reasoning_effort_utils import reasoning_effort_from_thinking, warn_dropped_reasoning -from litellm.litellm_core_utils.url_utils import ensure_api_route +from litellm.llms.mistral.common_utils import get_mistral_chat_credentials from litellm.llms.openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig, ) -from litellm.secret_managers.main import get_secret_str from litellm.types.llms.mistral import ( MistralCacheKeyParams, MistralFilePart, @@ -60,6 +59,7 @@ from litellm.utils import convert_to_model_response_object, supports_reasoning if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer + from litellm.llms.base_llm.base_utils import BaseTokenCounter def _accepted_reasoning_effort(model: str, requested: str, custom_llm_provider: str) -> str: @@ -185,14 +185,19 @@ class MistralConfig(OpenAIGPTConfig): def custom_llm_provider(self) -> str: return "mistral" - def _accepts_reasoning(self, model: str) -> bool: + def get_token_counter(self) -> "BaseTokenCounter | None": + from litellm.llms.mistral.count_tokens.token_counter import MistralTokenCounter + + return MistralTokenCounter() + + def accepts_reasoning(self, model: str) -> bool: return "magistral" in model.lower() or supports_reasoning( model=model, custom_llm_provider=self.custom_llm_provider ) def get_supported_openai_params(self, model: str) -> list[str]: is_magistral: Final = "magistral" in model.lower() - accepts_reasoning_effort: Final = self._accepts_reasoning(model) + accepts_reasoning_effort: Final = self.accepts_reasoning(model) accepts_prompt_cache_key: Final = self.custom_llm_provider == "mistral" return [ "stream", @@ -212,7 +217,7 @@ class MistralConfig(OpenAIGPTConfig): ] @staticmethod - def _map_tool_choice(tool_choice: object) -> str | ChatCompletionToolChoiceObjectParam | None: + def map_tool_choice(tool_choice: object) -> str | ChatCompletionToolChoiceObjectParam | None: match tool_choice: case str(choice) if choice in ("auto", "none"): return choice @@ -228,7 +233,7 @@ class MistralConfig(OpenAIGPTConfig): def _tool_choice_to_send( self, model: str, tool_choice: object, drop_params: bool ) -> str | ChatCompletionToolChoiceObjectParam | UnsupportedParamsError | None: - mapped: Final = self._map_tool_choice(tool_choice) + mapped: Final = self.map_tool_choice(tool_choice) if mapped is not None: return mapped if drop_params or litellm.drop_params is True: @@ -370,22 +375,23 @@ class MistralConfig(OpenAIGPTConfig): model=model, ) - def _get_openai_compatible_provider_info(self, api_base: str | None, api_key: str | None) -> tuple[str, str | None]: - # mistral is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.mistral.ai - api_base = ( - api_base - or get_secret_str("MISTRAL_AZURE_API_BASE") # for Azure AI Mistral - or "https://api.mistral.ai/v1" - ) + @staticmethod + def get_api_key(api_key: str | None = None) -> str | None: + return get_mistral_chat_credentials(api_base=None, api_key=api_key)[1] - # if api_base does not end with /v1 we add it - api_base = ensure_api_route(api_base, "v1") - dynamic_api_key: Final = ( - api_key - or get_secret_str("MISTRAL_AZURE_API_KEY") # for Azure AI Mistral - or get_secret_str("MISTRAL_API_KEY") - ) - return api_base, dynamic_api_key + @staticmethod + def get_api_base(api_base: str | None = None) -> str | None: + return get_mistral_chat_credentials(api_base=api_base, api_key=None)[0] + + def get_models( + self, api_key: str | None = None, api_base: str | None = None + ) -> list[str]: # mutable-ok: matches BaseConfig contract + from litellm.utils import bundled_provider_models + + return list(bundled_provider_models(self.custom_llm_provider)) + + def _get_openai_compatible_provider_info(self, api_base: str | None, api_key: str | None) -> tuple[str, str | None]: + return get_mistral_chat_credentials(api_base, api_key) # fmt: off @@ -418,53 +424,26 @@ class MistralConfig(OpenAIGPTConfig): The above statement is not valid now. Need to plan to remove all the #1,2,3 Mistral API supports content as a list. """ - keeps_reasoning: Final = self._accepts_reasoning(model) - final_position: Final = len(messages) - 1 - replayed_messages: Final = [ - MistralConfig._with_replayed_reasoning(m, keeps_reasoning, position == final_position) - for position, m in enumerate(messages) - ] + keeps_reasoning: Final = self.accepts_reasoning(model) + replayed_messages: Final = MistralConfig._replayed_messages(messages, keeps_reasoning) ## 1. If 'image_url' or 'file' in content, then transform with base class and mistral-specific handling - for m in replayed_messages: - _content_block = m.get("content") - if _content_block and isinstance(_content_block, list): - if any(c.get("type") in ["image_url", "file"] for c in _content_block): - if is_async: - return self._transform_messages_async(replayed_messages, model) - else: - return self._transform_messages_sync(replayed_messages, model) + if MistralConfig._sends_content_as_is(replayed_messages): + if is_async: + return self._transform_messages_async(list(replayed_messages), model) + else: + return self._transform_messages_sync(list(replayed_messages), model) ## 2. If content is list, then convert to string - flattened_messages: Final = tuple( - MistralConfig._with_flattened_content(m, keeps_reasoning) for m in replayed_messages - ) - ## 3. Handle name in message prepared_messages: Final = tuple( - MistralConfig._handle_tool_call_message(MistralConfig._handle_name_in_message(m)) - for m in flattened_messages - ) - empty_flags: Final = tuple(map(MistralConfig._is_empty_assistant_message, prepared_messages)) - tool_call_turns: Final = accumulate( - zip(count(), prepared_messages, empty_flags, strict=False), - MistralConfig._tool_call_turn, - initial=_NO_TOOL_CALL_TURN, - ) - tool_reply_keys: Final = tuple(map(MistralConfig._tool_reply_key, prepared_messages, tool_call_turns)) - first_reply_positions: Final = { - key: position for position, key in reversed(tuple(enumerate(tool_reply_keys))) if key is not None - } - keys_answered_with_content: Final = frozenset( - key for key, is_empty in zip(tool_reply_keys, empty_flags, strict=True) if not is_empty - ) - answering_positions: Final = frozenset( - position for key, position in first_reply_positions.items() if key not in keys_answered_with_content + MistralConfig._handle_tool_call_message(m) + for m in MistralConfig._flattened_messages(replayed_messages, keeps_reasoning) ) new_messages: Final[list[AllMessageValues]] = [ strip_none_values_from_message(MistralConfig._with_tool_reply_content(m)) - for position, (m, is_empty) in enumerate(zip(prepared_messages, empty_flags, strict=True)) - if not is_empty or position in answering_positions + for m, is_sent in zip(prepared_messages, MistralConfig._sent_flags(prepared_messages), strict=True) + if is_sent ] if is_async: @@ -583,6 +562,14 @@ class MistralConfig(OpenAIGPTConfig): case _: return None + @staticmethod + def think_chunk_texts(part: object) -> tuple[str, ...] | None: + match part: + case {"type": "thinking", "thinking": object() as items} if _is_json_array(items): + return tuple(text for text in map(MistralConfig._text_item_text, items) if text is not None) + case _: + return None + @staticmethod def _sendable_think_chunk( chunk: object, items: Sequence[object], may_stay_open: bool @@ -633,16 +620,12 @@ class MistralConfig(OpenAIGPTConfig): return "" @staticmethod - def _replayed_reasoning_text(message: ChatCompletionAssistantMessage) -> str | None: + def replayed_reasoning_text(message: Mapping[str, object]) -> str | None: reasoning_content: Final = message.get("reasoning_content") if isinstance(reasoning_content, str) and reasoning_content.strip(): return reasoning_content - thinking_blocks: Final = message.get("thinking_blocks") - thinking_text: Final = ( - "".join(map(MistralConfig._thinking_block_text, thinking_blocks)) - if isinstance(thinking_blocks, list) - else "" - ) + thinking_blocks: Final = as_object_list(message.get("thinking_blocks")) or () + thinking_text: Final = "".join(map(MistralConfig._thinking_block_text, thinking_blocks)) return thinking_text if thinking_text.strip() else None @staticmethod @@ -665,6 +648,79 @@ class MistralConfig(OpenAIGPTConfig): case _: return False + def replayed_turns(self, messages: Sequence[Mapping[str, object]], model: str) -> tuple[AllMessageValues, ...]: + chat_messages: Final = cast( # cast-ok: a count holds the chat turns as plain mappings + list[AllMessageValues], [dict(message) for message in messages] + ) + return MistralConfig._replayed_messages(chat_messages, self.accepts_reasoning(model)) + + def sent_turns(self, messages: Sequence[Mapping[str, object]], model: str) -> tuple[bool, ...]: + replayed_messages: Final = self.replayed_turns(messages, model) + if MistralConfig._sends_content_as_is(replayed_messages): + return tuple(True for _ in replayed_messages) + return MistralConfig._sent_flags( + tuple( + map( + MistralConfig._handle_tool_call_message, + MistralConfig._flattened_messages(replayed_messages, self.accepts_reasoning(model)), + ) + ) + ) + + @staticmethod + def _replayed_messages(messages: Sequence[AllMessageValues], keeps_reasoning: bool) -> tuple[AllMessageValues, ...]: + final_position: Final = len(messages) - 1 + return tuple( + MistralConfig._with_replayed_reasoning(m, keeps_reasoning, position == final_position) + for position, m in enumerate(messages) + ) + + @staticmethod + def _sends_content_as_is(messages: Sequence[Mapping[str, object]]) -> bool: + return any(map(MistralConfig._has_image_or_file_part, messages)) + + @staticmethod + def _has_image_or_file_part(message: Mapping[str, object]) -> bool: + content: Final = message.get("content") + return _is_json_array(content) and any(map(MistralConfig._is_image_or_file_part, content)) + + @staticmethod + def _is_image_or_file_part(part: object) -> bool: + match part: + case {"type": "image_url" | "file"}: + return True + case _: + return False + + @staticmethod + def _flattened_messages( + messages: Sequence[AllMessageValues], keeps_reasoning: bool + ) -> tuple[AllMessageValues, ...]: + return tuple( + MistralConfig._handle_name_in_message(MistralConfig._with_flattened_content(m, keeps_reasoning)) + for m in messages + ) + + @staticmethod + def _sent_flags(messages: Sequence[AllMessageValues]) -> tuple[bool, ...]: + empty_flags: Final = tuple(map(MistralConfig._is_empty_assistant_message, messages)) + tool_call_turns: Final = accumulate( + zip(count(), messages, empty_flags, strict=False), + MistralConfig._tool_call_turn, + initial=_NO_TOOL_CALL_TURN, + ) + tool_reply_keys: Final = tuple(map(MistralConfig._tool_reply_key, messages, tool_call_turns)) + first_reply_positions: Final = { + key: position for position, key in reversed(tuple(enumerate(tool_reply_keys))) if key is not None + } + keys_answered_with_content: Final = frozenset( + key for key, is_empty in zip(tool_reply_keys, empty_flags, strict=True) if not is_empty + ) + answering_positions: Final = frozenset( + position for key, position in first_reply_positions.items() if key not in keys_answered_with_content + ) + return tuple(not is_empty or position in answering_positions for position, is_empty in enumerate(empty_flags)) + @staticmethod def _with_replayed_reasoning( message: AllMessageValues, keeps_reasoning: bool, is_final_turn: bool @@ -679,7 +735,7 @@ class MistralConfig(OpenAIGPTConfig): return MistralConfig._with_sendable_think_chunks( stripped_message, content_parts, is_final_turn and MistralConfig._is_prefix_turn(message) ) - reasoning_text: Final = MistralConfig._replayed_reasoning_text(message) + reasoning_text: Final = MistralConfig.replayed_reasoning_text(message) has_turn_output: Final = bool(content_parts) or bool(message.get("tool_calls")) if reasoning_text is None or content_parts is None or not has_turn_output: return stripped_message diff --git a/litellm/llms/mistral/common_utils.py b/litellm/llms/mistral/common_utils.py index ef354047b06..26a05324049 100644 --- a/litellm/llms/mistral/common_utils.py +++ b/litellm/llms/mistral/common_utils.py @@ -3,6 +3,7 @@ from typing import Final import httpx +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str @@ -20,6 +21,12 @@ def get_mistral_api_base(api_base: str | None) -> str: return resolved.removesuffix("/v1") +def get_mistral_chat_credentials(api_base: str | None, api_key: str | None) -> tuple[str, str | None]: + resolved_base: Final = api_base or get_secret_str("MISTRAL_AZURE_API_BASE") or f"{MISTRAL_API_BASE}/v1" + resolved_key: Final = api_key or get_secret_str("MISTRAL_AZURE_API_KEY") or get_secret_str(MISTRAL_API_KEY_ENV_VAR) + return ensure_api_route(resolved_base, "v1"), resolved_key + + def get_mistral_auth_headers( headers: Mapping[str, str], api_key: str | None ) -> dict[str, str]: # mutable-ok: BaseConfig.validate_environment contract returns dict diff --git a/litellm/llms/mistral/count_tokens/__init__.py b/litellm/llms/mistral/count_tokens/__init__.py new file mode 100644 index 00000000000..8895acbe45f --- /dev/null +++ b/litellm/llms/mistral/count_tokens/__init__.py @@ -0,0 +1,9 @@ +from litellm.llms.mistral.count_tokens.handler import MistralCountTokensHandler +from litellm.llms.mistral.count_tokens.token_counter import MistralTokenCounter +from litellm.llms.mistral.count_tokens.transformation import MistralCountTokensConfig + +__all__ = [ + "MistralCountTokensConfig", + "MistralCountTokensHandler", + "MistralTokenCounter", +] diff --git a/litellm/llms/mistral/count_tokens/handler.py b/litellm/llms/mistral/count_tokens/handler.py new file mode 100644 index 00000000000..664970998cc --- /dev/null +++ b/litellm/llms/mistral/count_tokens/handler.py @@ -0,0 +1,108 @@ +from collections.abc import Mapping, Sequence +from typing import Final + +import httpx +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # factory stub leaves params untyped +) +from litellm.llms.mistral.common_utils import MistralError +from litellm.llms.mistral.count_tokens.transformation import ( + MistralCountTokensConfig, +) + +_DECODED_JSON: Final = TypeAdapter(object) + + +class MistralCountTokensHandler(MistralCountTokensConfig): + async def handle_count_tokens_request( + self, + model: str, + input: str | Sequence[object], + api_key: str, + api_base: str | None = None, + timeout: float | httpx.Timeout | None = None, + tools: Sequence[Mapping[str, object]] | None = None, + instructions: str | None = None, + text: Mapping[str, object] | None = None, + tool_choice: JsonValue = None, + ) -> object: + try: + self.validate_request(model, input) + verbose_logger.debug("Processing Mistral CountTokens request for model: %s", model) + request_body: Final = self.transform_request_to_count_tokens( + model=model, + input=input, + tools=tools, + instructions=instructions, + text=text, + tool_choice=tool_choice, + ) + except ValueError as e: + verbose_logger.error("Error in Mistral CountTokens handler: %s", e) + raise MistralError( + status_code=500, message=f"CountTokens processing error: {e}", headers=httpx.Headers() + ) from e + return await self.post_count_request(request_body, api_key, api_base, timeout) + + async def post_count_request( + self, + request_body: Mapping[str, object], + api_key: str, + api_base: str | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> object: + try: + endpoint_url: Final = self.get_mistral_count_tokens_endpoint(api_base) + verbose_logger.debug("Making Mistral count request to: %s", endpoint_url) + headers: Final = self.get_required_headers(api_key) + async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.MISTRAL) + request_timeout: Final = timeout if timeout is not None else litellm.request_timeout + + response: Final = await async_client.post( # pyright: ignore[reportUnknownMemberType] # HTTPHandler.post stub leaves params/headers untyped + endpoint_url, + headers=headers, + json=dict(request_body), + timeout=request_timeout, + ) + verbose_logger.debug("Mistral count response status: %s", response.status_code) + + if response.status_code != 200: + error_text: Final = response.text + verbose_logger.error("Mistral API error: %s", error_text) + raise MistralError( + status_code=response.status_code, + message=error_text, + headers=response.headers, + ) + + mistral_response: Final = _DECODED_JSON.validate_python(response.json()) + verbose_logger.debug("Mistral count response: %s", mistral_response) + except litellm.Timeout as e: + raise MistralError(status_code=e.status_code, message=str(e), headers=httpx.Headers()) from e + except (httpx.RequestError, litellm.APIConnectionError, OSError) as e: + verbose_logger.error("Connection error in Mistral CountTokens handler: %s", e) + raise MistralError( + status_code=500, message=f"CountTokens connection error: {e}", headers=httpx.Headers() + ) from e + except MistralError: + raise + except httpx.HTTPStatusError as e: + verbose_logger.error("HTTP error in Mistral CountTokens handler: %s", e) + raise MistralError( + status_code=e.response.status_code, + message=e.response.text, + headers=e.response.headers, + ) from e + except ValueError as e: + verbose_logger.error("Error in Mistral CountTokens handler: %s", e) + raise MistralError( + status_code=500, + message=f"CountTokens processing error: {e}", + headers=httpx.Headers(), + ) from e + else: + return mistral_response diff --git a/litellm/llms/mistral/count_tokens/token_counter.py b/litellm/llms/mistral/count_tokens/token_counter.py new file mode 100644 index 00000000000..ee8616103bb --- /dev/null +++ b/litellm/llms/mistral/count_tokens/token_counter.py @@ -0,0 +1,469 @@ +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain, groupby +from types import MappingProxyType +from typing import Final +from urllib.parse import urlparse + +import httpx +from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never + +import litellm +from litellm._logging import verbose_logger +from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, +) +from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object +from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import parts_of +from litellm.llms.anthropic.count_tokens.transformation import ChatCountTranslationError +from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter +from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError +from litellm.llms.anthropic.messages.convertible_blocks import deployment_drops_params, unconvertible_blocks_error +from litellm.llms.anthropic.messages.responses_adapters.handler import count_messages_as_sent +from litellm.llms.anthropic.messages.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) +from litellm.llms.base_llm.base_utils import BaseTokenCounter, deployment_timeout +from litellm.llms.mistral.chat.transformation import MistralConfig +from litellm.llms.mistral.common_utils import MistralError, get_mistral_chat_credentials +from litellm.llms.mistral.count_tokens.handler import ( + MistralCountTokensHandler, +) +from litellm.llms.openai.responses.count_tokens.transformation import ( + OpenAICountTokensConfig, + ResponsesCountRequestParts, + chat_content_blocks_to_responses_content, +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.anthropic import AnthropicMessagesRequest, AnthropicMessagesToolChoice +from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse +from litellm.utils import catalog_with_runtime_registrations + +mistral_count_tokens_handler: Final = MistralCountTokensHandler() +_MISTRAL_CHAT_CONFIG: Final = MistralConfig() +_CHAT_TO_RESPONSES: Final = LiteLLMResponsesTransformationHandler() +_MODES_MISTRAL_COUNTS: Final = frozenset({"chat", "responses"}) +_ANTHROPIC_TOOL_CHOICE: Final[TypeAdapter[AnthropicMessagesToolChoice | None]] = TypeAdapter( + AnthropicMessagesToolChoice | None +) +_ANTHROPIC_OPTIONS_AS_CHAT: Final = TypeAdapter(AnthropicMessagesRequest) +_CHAT_PARAMS: Final = TypeAdapter(Mapping[str, JsonValue]) +_CHAT_REQUEST: Final = TypeAdapter(Mapping[str, object]) +_CHAT_PARAMS_FROM_ANTHROPIC_OPTIONS: Final = ("tool_choice",) + + +@dataclass(frozen=True, slots=True) +class _MistralCountCredential: + api_key: str + api_base: str | None + + +def _is_azure_endpoint(api_base: str) -> bool: + return "azure" in (urlparse(api_base).hostname or "").lower() + + +def _count_credential( + deployment: Mapping[str, object] | None, model_label: str +) -> _MistralCountCredential | MistralError | None: + litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) or {} + if litellm_params.get("custom_llm_provider") in ("azure", "azure_ai"): + return None + raw_api_base: Final = litellm_params.get("api_base") + raw_api_key: Final = litellm_params.get("api_key") + deployment_key: Final = raw_api_key if isinstance(raw_api_key, str) and raw_api_key else None + api_base, api_key = get_mistral_chat_credentials( + api_base=raw_api_base if isinstance(raw_api_base, str) else None, api_key=deployment_key + ) + azure_key: Final = get_secret_str("MISTRAL_AZURE_API_KEY") + if _is_azure_endpoint(api_base) or (bool(azure_key) and api_key == azure_key): + verbose_logger.debug( + "Mistral native count tokens skipped for Azure deployment/credential, falling back to local: base=%s", + api_base, + ) + return None + + resolved_key: Final = api_key or litellm.api_key + if not resolved_key: + return MistralError( + status_code=401, + message=( + f"No Mistral API key is set for model {model_label}, so tokens can't be counted. Set api_key on " + "the deployment or MISTRAL_API_KEY in the environment." + ), + headers=httpx.Headers(), + ) + return _MistralCountCredential(api_key=resolved_key, api_base=api_base) + + +def _counts_on_mistral(model: str) -> bool: + cost_entry: Final = catalog_with_runtime_registrations().get( + model if model.startswith("mistral/") else f"mistral/{model}" + ) + mode: Final = None if cost_entry is None else cost_entry.get("mode") + return mode is None or mode in _MODES_MISTRAL_COUNTS + + +def _count_error(model_to_use: str, request_model: str, message: str, status_code: int) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="mistral_api", + error=True, + error_message=message, + status_code=status_code, + ) + + +def _file_id(block: object) -> str | None: + match block: + case {"type": "file", "file": {"file_id": str(file_id)}} if file_id: + return file_id + case {"type": "file", "file_id": str(file_id)} if file_id: + return file_id + case _: + return None + + +def _is_file_id_block(block: object) -> bool: + return _file_id(block) is not None + + +def _sends_file_ids(message: Mapping[str, object]) -> bool: + return message.get("role") == "user" and any(map(_is_file_id_block, parts_of(message.get("content")))) + + +def _run_parts(is_file_id: bool, run: Iterable[object]) -> tuple[object, ...]: + blocks: Final = tuple(run) + if is_file_id: + return tuple({"type": "input_file", "file_id": _file_id(block)} for block in blocks) + return _input_content_parts(chat_content_blocks_to_responses_content(blocks, "user")) + + +def _user_parts_in_order(content: object) -> tuple[object, ...]: + runs: Final = groupby(parts_of(content), key=_is_file_id_block) + return tuple(chain.from_iterable(_run_parts(is_file_id, run) for is_file_id, run in runs)) + + +def _is_think_chunk(part: object) -> bool: + return MistralConfig.think_chunk_texts(part) is not None + + +def _reasoning_item(texts: tuple[str, ...]) -> Mapping[str, object]: + return {"type": "reasoning", "summary": [{"type": "summary_text", "text": text} for text in texts]} + + +def _replayed_reasoning_items(message: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + if message.get("role") != "assistant": + return () + chunk_texts: Final = map(MistralConfig.think_chunk_texts, parts_of(message.get("content"))) + return tuple(_reasoning_item(texts) for texts in chunk_texts if texts) + + +def _without_think_chunks(message: Mapping[str, object]) -> Mapping[str, object]: + parts: Final = parts_of(message.get("content")) + if message.get("role") != "assistant" or not any(map(_is_think_chunk, parts)): + return message + return {**message, "content": [part for part in parts if not _is_think_chunk(part)] or None} + + +def _input_content_parts(content: object) -> tuple[object, ...]: + match content: + case str(text) if text: + return ({"type": "input_text", "text": text},) + case _: + return parts_of(content) + + +def _part_text(part: object) -> str: + match part: + case {"text": str(text)}: + return text + case _: + return "" + + +def _system_text(content: object) -> str: + return content if isinstance(content, str) else "".join(map(_part_text, parts_of(content))) + + +def _is_system_turn(message: Mapping[str, object]) -> bool: + return message.get("role") in ("system", "developer") + + +def _is_text_block(block: object) -> bool: + match block: + case {"type": "text"}: + return True + case _: + return False + + +def _joined_run(is_text: bool, run: Iterable[object]) -> tuple[object, ...]: + blocks: Final = tuple(run) + return ({"type": "text", "text": "".join(map(_part_text, blocks))},) if is_text else blocks + + +def _joined_text_runs(content: object) -> tuple[object, ...]: + runs: Final = groupby(parts_of(content), key=_is_text_block) + return tuple(chain.from_iterable(_joined_run(is_text, run) for is_text, run in runs)) + + +def _with_joined_tool_result_text(block: object) -> object: + tool_result: Final = as_str_mapping(block) + if tool_result is None or tool_result.get("type") != "tool_result": + return block + if not isinstance(tool_result.get("content"), list): + return block + return {**tool_result, "content": list(_joined_text_runs(tool_result.get("content")))} + + +def _with_tool_result_text_joins(message: Mapping[str, object]) -> Mapping[str, object]: + if not isinstance(message.get("content"), list): + return message + return {**message, "content": list(map(_with_joined_tool_result_text, parts_of(message.get("content"))))} + + +def _with_chat_text_joins(message: Mapping[str, object]) -> Mapping[str, object]: + if not isinstance(message.get("content"), list): + return message + blocks: Final = _joined_text_runs(message.get("content")) + match blocks: + case ({"type": "text", "text": str(text)},): + return {**message, "content": text} + case _: + return {**message, "content": list(blocks)} + + +def _chat_message_items(message: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + if _is_system_turn(message): + return ({"role": "system", "content": _system_text(message.get("content"))},) + joined_message: Final = _with_chat_text_joins(_without_think_chunks(message)) + items, _ = OpenAICountTokensConfig.messages_to_responses_input((joined_message,)) + sent_items: Final = ( + tuple({**item, "content": _user_parts_in_order(joined_message.get("content"))} for item in items) + if _sends_file_ids(message) + else tuple(items) + ) + return (*_replayed_reasoning_items(message), *sent_items) + + +def _counted_tool_choice(tool_choice: object) -> JsonValue: + match MistralConfig.map_tool_choice(tool_choice): + case "any": + return "required" + case {"function": {"name": str(name)}}: + return {"type": "function", "name": name} + case str(choice): + return choice + case _: + return None + + +def _anthropic_options_as_chat_params( + model: str, count_options: Mapping[str, JsonValue], drop_params: bool +) -> Mapping[str, JsonValue] | ChatCountTranslationError: + try: + _ = _ANTHROPIC_TOOL_CHOICE.validate_python(count_options.get("tool_choice")) + chat_request, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + _ANTHROPIC_OPTIONS_AS_CHAT.validate_python({**count_options, "model": model, "messages": []}), + custom_llm_provider=LlmProviders.MISTRAL.value, + drop_params=drop_params, + ) + except (ValidationError, UnsupportedParamsError) as e: + return ChatCountTranslationError(message=str(e)) + sent: Final = _CHAT_REQUEST.validate_python(chat_request) + return _CHAT_PARAMS.validate_python( + {name: sent[name] for name in _CHAT_PARAMS_FROM_ANTHROPIC_OPTIONS if sent.get(name) is not None} + ) + + +def _anthropic_system_as_chat_sends_it(system: object) -> object: + return _system_text(system) if system else system + + +class MistralTokenCounter(BaseTokenCounter): + def should_use_token_counting_api( + self, + custom_llm_provider: str | None = None, + ) -> bool: + return custom_llm_provider == LlmProviders.MISTRAL.value + + @staticmethod + def _responses_request_parts( + model_to_use: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + system: object, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue, + ) -> ResponsesCountRequestParts: + match messages_format: + case "anthropic": + return OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model=model_to_use, + messages=tuple(map(_with_tool_result_text_joins, messages)), + system=system, + tools=tools, + output_format=output_format, + ) + case "openai": + replayed_turns: Final = _MISTRAL_CHAT_CONFIG.replayed_turns(messages, model_to_use) + sent_turns: Final = _MISTRAL_CHAT_CONFIG.sent_turns(messages, model_to_use) + input_items: Final = tuple( + chain.from_iterable( + _chat_message_items(message) + for message, is_sent in zip(replayed_turns, sent_turns, strict=True) + if is_sent + ) + ) + return OpenAICountTokensConfig.count_request_parts( + input_items, + None + if any(map(_is_system_turn, messages)) + else LiteLLMAnthropicToResponsesAPIAdapter.system_instructions(system), + tools, + ) + case _: + assert_never(messages_format) + + async def count_tokens( + self, + model_to_use: str, + messages: list[dict[str, object]] | None, # mutable-ok: matches BaseTokenCounter signature + contents: list[dict[str, object]] | None, # mutable-ok: matches BaseTokenCounter signature + deployment: dict[str, object] | None = None, # mutable-ok: matches BaseTokenCounter signature + request_model: str = "", + tools: list[dict[str, object]] | None = None, # mutable-ok: matches BaseTokenCounter signature + system: object = None, + messages_format: TokenCountMessagesFormat = "openai", + 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 or not _counts_on_mistral(model_to_use): + return None + + malformed_error: Final = ( + OpenAICountTokensConfig.malformed_anthropic_messages_error( + messages, drop_params=deployment_drops_params(deployment) + ) + if messages_format == "anthropic" + else None + ) + request_error: Final = malformed_error or unconvertible_blocks_error(messages, messages_format) + if request_error is not None: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="mistral_api", + error=True, + error_message=request_error, + status_code=400, + rejected_before_provider_call=True, + ) + sent_chat_params: Final = ( + _anthropic_options_as_chat_params(model_to_use, count_options, deployment_drops_params(deployment)) + if messages_format == "anthropic" + else chat_params + ) + if isinstance(sent_chat_params, ChatCountTranslationError): + return sent_chat_params.response(request_model, model_to_use, "mistral_api") + + sent_messages: Final = await count_messages_as_sent( + model=model_to_use, + messages=messages, + tools=tools, + system=system, + messages_format=messages_format, + context_management=context_management, + ) + if isinstance(sent_messages, AnthropicContextManagementError): + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="mistral_api", + error=True, + error_message=sent_messages.message, + status_code=sent_messages.status_code, + rejected_before_provider_call=True, + ) + + resolved: Final = _count_credential(deployment, request_model or model_to_use) + if resolved is None: + return None + if isinstance(resolved, MistralError): + return _count_error(model_to_use, request_model, resolved.message, resolved.status_code) + + input_items, instructions, counted_tools, text = self._responses_request_parts( + model_to_use=model_to_use, + messages=sent_messages, + tools=tools, + system=_anthropic_system_as_chat_sends_it(system) if messages_format == "anthropic" else system, + messages_format=messages_format, + output_format=output_format, + ) + + if not input_items: + return None + counted_text: Final = ( + text + if text is not None + else _CHAT_TO_RESPONSES.transform_response_format_to_text_format(chat_params.get("response_format")) + ) + timeout: Final = deployment_timeout( + as_str_mapping(None if deployment is None else deployment.get("litellm_params")) or {} + ) + tool_choice: Final = _counted_tool_choice(sent_chat_params.get("tool_choice")) if counted_tools else None + + return await self._counted( + resolved, + lambda api_key, api_base: mistral_count_tokens_handler.handle_count_tokens_request( + model=model_to_use, + input=input_items, + api_key=api_key, + api_base=api_base, + timeout=timeout, + tools=counted_tools, + instructions=instructions, + text=counted_text, + tool_choice=tool_choice, + ), + model_to_use, + request_model, + ) + + @staticmethod + async def _counted( + resolved: _MistralCountCredential, + send: Callable[[str, str | None], Awaitable[object]], + model_to_use: str, + request_model: str, + ) -> TokenCountResponse | None: + try: + result: Final = await send(resolved.api_key, resolved.api_base) + except MistralError as e: + verbose_logger.warning("Mistral CountTokens API error: status=%s, message=%s", e.status_code, e.message) + return _count_error(model_to_use, request_model, e.message, e.status_code) + except (httpx.HTTPError, ValueError) as e: + verbose_logger.warning("Error calling Mistral CountTokens API: %s", e) + return _count_error(model_to_use, request_model, str(e), 500) + else: + input_tokens: Final = None if result is None else decoded_json_object(result).get("input_tokens") + if not isinstance(input_tokens, int) or isinstance(input_tokens, bool): + verbose_logger.warning("Mistral count response had no input_tokens, falling back to local: %s", result) + return None + return TokenCountResponse.model_validate( + { + "total_tokens": input_tokens, + "request_model": request_model, + "model_used": model_to_use, + "tokenizer_type": "mistral_api", + "original_response": result, + } + ) diff --git a/litellm/llms/mistral/count_tokens/transformation.py b/litellm/llms/mistral/count_tokens/transformation.py new file mode 100644 index 00000000000..043d3722324 --- /dev/null +++ b/litellm/llms/mistral/count_tokens/transformation.py @@ -0,0 +1,42 @@ +from collections.abc import Mapping, Sequence + +from pydantic import JsonValue + +from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.mistral.common_utils import get_mistral_api_base +from litellm.llms.openai.responses.count_tokens.transformation import OpenAICountTokensConfig + + +class MistralCountTokensConfig(OpenAICountTokensConfig): + def get_mistral_count_tokens_endpoint(self, api_base: str | None = None) -> str: + return join_api_base(get_mistral_api_base(api_base), "v1/responses/input_tokens") + + def get_required_headers( + self, api_key: str + ) -> dict[str, str]: # mutable-ok: matches OpenAICountTokensConfig contract + return { + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}", + } + + def transform_request_to_count_tokens( + self, + model: str, + input: str | Sequence[object], + tools: Sequence[Mapping[str, object]] | None = None, + instructions: str | None = None, + text: Mapping[str, object] | None = None, + tool_choice: JsonValue = None, + ) -> dict[str, object]: # mutable-ok: matches OpenAICountTokensConfig contract + return { + "model": model, + "input": input, + **({} if instructions is None else {"instructions": instructions}), + **( + {} + if tools is None + else {"tools": self._transform_tools_for_responses_api([dict(tool) for tool in tools])} + ), + **({} if text is None else {"text": dict(text)}), + **({} if tool_choice is None else {"tool_choice": tool_choice}), + } diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index 787d1b840a2..91966c83ad2 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -18,7 +18,7 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.tool_fields import DropPolicy from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError from litellm.llms.anthropic.messages.convertible_blocks import deployment_drops_params, unconvertible_blocks_error -from litellm.llms.anthropic.messages.responses_adapters.handler import count_view_messages +from litellm.llms.anthropic.messages.responses_adapters.handler import count_messages_as_sent from litellm.llms.anthropic.messages.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, ) @@ -101,32 +101,6 @@ def _count_error(model_to_use: str, request_model: str, message: str, status_cod ) -async def _sent_messages( - model: str, - messages: Sequence[Mapping[str, object]], - tools: Sequence[Mapping[str, object]] | None, - system: object, - messages_format: TokenCountMessagesFormat, - context_management: CountContextManagement | None, -) -> Sequence[Mapping[str, object]] | AnthropicContextManagementError: - match messages_format: - case "openai": - return messages - case "anthropic": - try: - return await count_view_messages( - model=model, - messages=messages, - tools=tools, - system=system, - context_management=context_management, - ) - except AnthropicContextManagementError as error: - return error - case _: - assert_never(messages_format) - - class OpenAITokenCounter(BaseTokenCounter): """Token counter implementation for OpenAI provider using the Responses API.""" @@ -136,12 +110,6 @@ class OpenAITokenCounter(BaseTokenCounter): ) -> bool: return custom_llm_provider == LlmProviders.OPENAI.value - @staticmethod - def _system_instructions(system: object) -> str | None: - if system is None or isinstance(system, str): - return system - return LiteLLMAnthropicToResponsesAPIAdapter.system_blocks_to_instructions(system) or None - @staticmethod def _responses_request_parts( model_to_use: str, @@ -160,7 +128,7 @@ class OpenAITokenCounter(BaseTokenCounter): input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( [dict(message) for message in messages] ) - fallback_instructions: Final = OpenAITokenCounter._system_instructions(system) + fallback_instructions: Final = LiteLLMAnthropicToResponsesAPIAdapter.system_instructions(system) return OpenAICountTokensConfig.count_request_parts( input_items, instructions if instructions is not None else fallback_instructions, tools ) @@ -208,8 +176,13 @@ class OpenAITokenCounter(BaseTokenCounter): rejected_before_provider_call=True, ) - sent_messages: Final = await _sent_messages( - model_to_use, messages, tools, system, messages_format, context_management + sent_messages: Final = await count_messages_as_sent( + model=model_to_use, + messages=messages, + tools=tools, + system=system, + messages_format=messages_format, + context_management=context_management, ) if isinstance(sent_messages, AnthropicContextManagementError): return TokenCountResponse( diff --git a/litellm/utils.py b/litellm/utils.py index 205dec2f9f9..3f4ebf7a4fb 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10856,6 +10856,10 @@ class ProviderConfigManager: return litellm.AnthropicModelInfo() elif LlmProviders.XAI == provider: return litellm.XAIModelInfo() + elif LlmProviders.MISTRAL == provider: + from litellm.llms.mistral.chat.transformation import MistralConfig + + return MistralConfig() elif LlmProviders.OLLAMA == provider or LlmProviders.OLLAMA_CHAT == provider: # Dynamic model listing for Ollama server from litellm.llms.ollama.common_utils import OllamaModelInfo diff --git a/tests/integration/messages_endpoint/test_tool_result_text_wire.py b/tests/integration/messages_endpoint/test_tool_result_text_wire.py index f8c6eb6b0ca..f652eea5b49 100644 --- a/tests/integration/messages_endpoint/test_tool_result_text_wire.py +++ b/tests/integration/messages_endpoint/test_tool_result_text_wire.py @@ -24,6 +24,7 @@ _CLAUDE_CODE: Final = MappingProxyType({"user-agent": "claude-cli/2.1.0 (externa _OPENAI_COUNT: Final = 4242 _GOOGLE_COUNT: Final = 777 _ANTHROPIC_COUNT: Final = 58 +_MISTRAL_COUNT: Final = 2929 _BEDROCK_COUNT: Final = 1345 _GEMINI: Final = "gemini-4-argon" _VERTEX_PROJECT: Final = "scripted-project" @@ -119,7 +120,9 @@ def _openai_peer(request: Request) -> Reply: return Reply(status=404, body=json.dumps({"error": request.target}).encode()) -def _chat_peer(request: Request) -> Reply: +def _mistral_peer(request: Request) -> Reply: + if request.target.endswith("/responses/input_tokens"): + return Reply(body=json.dumps({"object": "response.input_tokens", "input_tokens": _MISTRAL_COUNT}).encode()) if request.target.endswith("/chat/completions"): return Reply(body=_chat_reply()) return Reply(status=404, body=json.dumps({"error": request.target}).encode()) @@ -286,8 +289,8 @@ def test_explicit_drop_params_false_wins_on_both_routes(gateway: Gateway) -> Non assert wire.drain() == () -def test_chat_bridge_and_local_count_reject_a_non_string_tool_result_text(gateway: Gateway) -> None: - with wire_server(_chat_peer) as wire, gateway.scenario() as scenario: +def test_chat_bridge_and_count_reject_a_non_string_tool_result_text(gateway: Gateway) -> None: + with wire_server(_mistral_peer) as wire, gateway.scenario() as scenario: model: Final = _mistral_model(scenario, wire) body: Final = _body(model, _NULL_TEXT) _assert_field_rejection(gateway, "/v1/messages", {**body, "max_tokens": 64}) @@ -299,14 +302,17 @@ def test_chat_bridge_and_local_count_reject_a_non_string_tool_result_text(gatewa @pytest.mark.parametrize("source", ("deployment", "body")) def test_chat_bridge_drop_params_drops_the_malformed_text(gateway: Gateway, source: str) -> None: deployment, extra, _ = _DROP_SOURCES[source] - with wire_server(_chat_peer) as wire, gateway.scenario() as scenario: + with wire_server(_mistral_peer) as wire, gateway.scenario() as scenario: model: Final = _mistral_model(scenario, wire, **deployment) body: Final = _body(model, _NUMBER_THEN_WEATHER, **extra) created: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 64}) counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) assert created.status_code == 200, created.text assert counted.status_code == 200, counted.text - assert _tool_message_text(_only(wire.drain(), "/chat/completions")) == _WEATHER + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _MISTRAL_COUNT, counted.text + requests: Final = wire.drain() + assert _tool_message_text(_only(requests, "/chat/completions")) == _WEATHER + assert _function_call_output(_only(requests, "/responses/input_tokens")) == _WEATHER def test_count_without_a_deployment_rejects_and_honors_body_drop_params(gateway: Gateway) -> None: @@ -386,7 +392,7 @@ def test_bedrock_count_with_a_claude_code_user_agent_still_counts(gateway: Gatew @pytest.mark.parametrize("provider", ("openai", "mistral")) def test_valid_tool_result_text_reaches_the_upstream(gateway: Gateway, provider: str) -> None: - peer: Final = _openai_peer if provider == "openai" else _chat_peer + peer: Final = _openai_peer if provider == "openai" else _mistral_peer with wire_server(peer) as wire, gateway.scenario() as scenario: model: Final = (_openai_model if provider == "openai" else _mistral_model)(scenario, wire) body: Final = _body(model, _WEATHER_THEN_MISSING) @@ -399,7 +405,11 @@ def test_valid_tool_result_text_reaches_the_upstream(gateway: Gateway, provider: assert _function_call_output(_only(requests, "/responses")) == f"{_WEATHER}\n" assert _function_call_output(_only(requests, "/responses/input_tokens")) == f"{_WEATHER}\n" else: - assert _tool_message_text(_only(requests, "/chat/completions")) == _WEATHER + assert ( + _function_call_output(_only(requests, "/responses/input_tokens")) + == _tool_message_text(_only(requests, "/chat/completions")) + == _WEATHER + ) def _rewritten(content: JsonValue, rewrite: Callable[[str], str]) -> JsonValue: @@ -626,7 +636,7 @@ def test_rewriting_guardrail_leaves_the_rejection_to_the_chat_bridge_for_any_too gateway: Gateway, tool_use_id: str ) -> None: with ( - wire_server(_chat_peer) as provider, + wire_server(_mistral_peer) as provider, wire_server(_guardrail_peer) as guard, gateway.scenario() as scenario, ): @@ -734,7 +744,7 @@ def _forwarded_chat_turns(wire: Wire) -> tuple[dict[str, JsonValue], ...]: def test_mistral_receives_the_tool_reply_whose_only_text_drop_params_dropped(gateway: Gateway) -> None: - with wire_server(_chat_peer) as wire, gateway.scenario() as scenario: + with wire_server(_mistral_peer) as wire, gateway.scenario() as scenario: model: Final = _mistral_model(scenario, wire) created: Final = gateway.request( "POST", "/v1/messages", _body(model, _NULL_TEXT, max_tokens=64, drop_params=True) @@ -764,7 +774,7 @@ def test_mistral_chat_completions_keeps_an_empty_tool_reply(gateway: Gateway) -> {"role": "tool", "tool_call_id": "call_weather", "content": ""}, {"role": "user", "content": ""}, ] - with wire_server(_chat_peer) as wire, gateway.scenario() as scenario: + with wire_server(_mistral_peer) as wire, gateway.scenario() as scenario: model: Final = _mistral_model(scenario, wire) response: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": messages}) assert response.status_code == 200, response.text @@ -787,7 +797,7 @@ def test_mistral_chat_completions_drops_an_orphan_tool_message_and_sends_a_null_ {"role": "assistant", "content": None, "tool_calls": [_WEATHER_CALL]}, {"role": "tool", "tool_call_id": "call_weather", "content": None}, ] - with wire_server(_chat_peer) as wire, gateway.scenario() as scenario: + with wire_server(_mistral_peer) as wire, gateway.scenario() as scenario: model: Final = _mistral_model(scenario, wire) response: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": messages}) assert response.status_code == 200, response.text diff --git a/tests/unit/llms/mistral/count_tokens/__init__.py b/tests/unit/llms/mistral/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/mistral/count_tokens/test_mistral_token_counter.py b/tests/unit/llms/mistral/count_tokens/test_mistral_token_counter.py new file mode 100644 index 00000000000..e6095fd1ea8 --- /dev/null +++ b/tests/unit/llms/mistral/count_tokens/test_mistral_token_counter.py @@ -0,0 +1,1225 @@ +import json +from collections.abc import Mapping +from typing import Final + +import httpx +import pytest +import respx +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm.litellm_core_utils.core_helpers import decoded_json_object +from litellm.llms.anthropic.messages.adapters.transformation import AnthropicAdapter +from litellm.llms.mistral.chat.transformation import MistralConfig +from litellm.llms.mistral.count_tokens.token_counter import MistralTokenCounter +from litellm.llms.mistral.count_tokens.transformation import MistralCountTokensConfig +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse + +COUNT_URL: Final = "https://api.mistral.ai/v1/responses/input_tokens" + + +def test_get_mistral_count_tokens_endpoint(): + config: Final = MistralCountTokensConfig() + assert config.get_mistral_count_tokens_endpoint() == COUNT_URL + assert ( + config.get_mistral_count_tokens_endpoint("https://custom.mistral.ai") + == "https://custom.mistral.ai/v1/responses/input_tokens" + ) + assert ( + config.get_mistral_count_tokens_endpoint("https://custom.mistral.ai/v1") + == "https://custom.mistral.ai/v1/responses/input_tokens" + ) + + +def test_get_required_headers(): + config: Final = MistralCountTokensConfig() + headers: Final = config.get_required_headers("test-key") + assert headers["Authorization"] == "Bearer test-key" + assert headers["Content-Type"] == "application/json" + + +def test_transform_request_to_count_tokens(): + config: Final = MistralCountTokensConfig() + transformed: Final = config.transform_request_to_count_tokens( + model="mistral-large-4", + input="Hello world", + instructions="Be concise", + tools=[ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Lookup data", + "parameters": {"type": "object"}, + }, + } + ], + ) + assert transformed["model"] == "mistral-large-4" + assert transformed["input"] == "Hello world" + assert transformed["instructions"] == "Be concise" + assert transformed["tools"] == [ + { + "type": "function", + "name": "lookup", + "description": "Lookup data", + "parameters": {"type": "object"}, + } + ] + + +def test_should_use_token_counting_api(): + counter: Final = MistralTokenCounter() + assert counter.should_use_token_counting_api("mistral") is True + assert counter.should_use_token_counting_api("openai") is False + assert counter.should_use_token_counting_api(None) is False + + +@pytest.mark.asyncio +async def test_count_tokens_success(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + counter: Final = MistralTokenCounter() + with respx.mock: + respx.post(COUNT_URL).mock( + return_value=httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 42}) + ) + result: Final = await counter.count_tokens( + model_to_use="mistral-large-4", + messages=[{"role": "user", "content": "Tell me a joke"}], + contents=None, + deployment={"litellm_params": {"api_key": "test-key"}}, + request_model="mistral/mistral-large-4", + ) + + assert result is not None + assert result.total_tokens == 42 + assert result.tokenizer_type == "mistral_api" + assert result.model_used == "mistral-large-4" + assert result.request_model == "mistral/mistral-large-4" + assert result.error is False + + +@pytest.mark.asyncio +async def test_count_tokens_missing_api_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("MISTRAL_API_KEY", raising=False) + monkeypatch.delenv("MISTRAL_AZURE_API_KEY", raising=False) + monkeypatch.setattr(litellm, "api_key", None) + counter: Final = MistralTokenCounter() + result: Final = await counter.count_tokens( + model_to_use="mistral-large-4", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="mistral/mistral-large-4", + ) + assert result is not None + assert result.error is True + assert result.status_code == 401 + assert "No Mistral API key is set" in (result.error_message or "") + + +@pytest.mark.parametrize( + ("upstream_error", "status_code"), + [(httpx.ConnectTimeout("Connection timed out"), 408), (httpx.ConnectError("Connection refused"), 500)], + ids=["timeout", "connect_error"], +) +@pytest.mark.asyncio +async def test_count_tokens_reports_a_failed_connection_as_a_count_error( + monkeypatch: pytest.MonkeyPatch, upstream_error: httpx.RequestError, status_code: int +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + counter: Final = MistralTokenCounter() + with respx.mock: + respx.post(COUNT_URL).mock(side_effect=upstream_error) + result: Final = await counter.count_tokens( + model_to_use="mistral-large-4", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={"litellm_params": {"api_key": "test-key"}}, + request_model="mistral/mistral-large-4", + ) + assert result is not None + assert (result.error, result.status_code, result.tokenizer_type) == (True, status_code, "mistral_api") + + +@pytest.mark.asyncio +async def test_count_tokens_upstream_error_reported(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + counter: Final = MistralTokenCounter() + with respx.mock: + respx.post(COUNT_URL).mock(return_value=httpx.Response(400, text="Bad Request")) + result: Final = await counter.count_tokens( + model_to_use="mistral-large-4", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={"litellm_params": {"api_key": "test-key"}}, + request_model="mistral/mistral-large-4", + ) + assert result is not None + assert result.error is True + assert result.status_code == 400 + + +@pytest.mark.asyncio +async def test_count_tokens_empty_messages(): + counter: Final = MistralTokenCounter() + result: Final = await counter.count_tokens( + model_to_use="mistral-large-4", + messages=[], + contents=None, + deployment={"litellm_params": {"api_key": "test-key"}}, + request_model="mistral/mistral-large-4", + ) + assert result is None + + +@pytest.mark.asyncio +async def test_count_tokens_openai_format_with_tool_calls_and_tool_results(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + counter: Final = MistralTokenCounter() + messages: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_123", "content": '{"temp": 20}'}, + ] + tools: Final[list[dict[str, object]]] = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, + } + ] + with respx.mock: + route: Final = respx.post(COUNT_URL).mock( + return_value=httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 58}) + ) + result: Final = await counter.count_tokens( + model_to_use="mistral-large-4", + messages=messages, + contents=None, + tools=tools, + system="Answer concisely.", + deployment={"litellm_params": {"api_key": "test-key"}}, + request_model="mistral/mistral-large-4", + messages_format="openai", + ) + + assert result is not None + assert result.total_tokens == 58 + assert result.error is False + sent: Final = decoded_json_object(json.loads(route.calls.last.request.content)) + assert sent["instructions"] == "Answer concisely." + assert sent["tools"] == [ + { + "type": "function", + "name": "get_weather", + "description": "Get weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + } + ] + assert sent["input"] == [ + {"role": "user", "content": "What's the weather?"}, + { + "type": "function_call", + "call_id": "call_123", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + }, + { + "type": "function_call_output", + "call_id": "call_123", + "output": '{"temp": 20}', + }, + ] + + +@pytest.mark.asyncio +async def test_count_tokens_anthropic_format_with_tool_use_and_tool_result(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + counter: Final = MistralTokenCounter() + messages: Final[list[dict[str, object]]] = [ + {"role": "user", "content": [{"type": "text", "text": "What's the weather?"}]}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_123", + "name": "get_weather", + "input": {"city": "Paris"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_123", + "content": [{"type": "text", "text": '{"temp": 20}'}], + } + ], + }, + ] + tools: Final[list[dict[str, object]]] = [ + { + "name": "get_weather", + "description": "Get weather", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}}, + } + ] + with respx.mock: + route: Final = respx.post(COUNT_URL).mock( + return_value=httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 64}) + ) + result: Final = await counter.count_tokens( + model_to_use="mistral-large-4", + messages=messages, + contents=None, + tools=tools, + system="Be precise.", + deployment={"litellm_params": {"api_key": "test-key"}}, + request_model="mistral/mistral-large-4", + messages_format="anthropic", + ) + + assert result is not None + assert result.total_tokens == 64 + sent: Final = decoded_json_object(json.loads(route.calls.last.request.content)) + assert sent["instructions"] == "Be precise." + assert sent["tools"] == [ + { + "type": "function", + "name": "get_weather", + "description": "Get weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + "strict": False, + } + ] + assert { + "type": "function_call", + "call_id": "toolu_123", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + } in sent["input"] + assert { + "type": "function_call_output", + "call_id": "toolu_123", + "output": '{"temp": 20}', + } in sent["input"] + + +@pytest.mark.asyncio +async def test_count_tokens_malformed_tool_result_rejected_before_provider_call(): + counter: Final = MistralTokenCounter() + messages: Final[list[dict[str, object]]] = [ + {"role": "user", "content": [{"type": "text", "text": "Hi"}]}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "t1", "name": "fn", "input": {}}]}, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t1", "content": [{"type": "text", "text": 123}]}], + }, + ] + result: Final = await counter.count_tokens( + model_to_use="mistral-large-4", + messages=messages, + contents=None, + deployment={"litellm_params": {"api_key": "test-key", "drop_params": False}}, + request_model="mistral/mistral-large-4", + messages_format="anthropic", + ) + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert result.rejected_before_provider_call is True + + +@pytest.mark.asyncio +async def test_count_tokens_azure_endpoint_falls_back_to_none(): + counter: Final = MistralTokenCounter() + result: Final = await counter.count_tokens( + model_to_use="mistral-large-4", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={ + "litellm_params": { + "api_key": "test-key", + "api_base": "https://my-mistral.inference.ai.azure.com", + } + }, + request_model="mistral/mistral-large-4", + ) + assert result is None + + +@pytest.mark.parametrize("mistral_api_key", [None, "env-mistral-key"]) +@pytest.mark.asyncio +async def test_count_tokens_never_sends_the_azure_env_key_to_mistral( + monkeypatch: pytest.MonkeyPatch, mistral_api_key: str | None +): + monkeypatch.setenv("MISTRAL_AZURE_API_KEY", "azure-test-key") + if mistral_api_key is None: + monkeypatch.delenv("MISTRAL_API_KEY", raising=False) + else: + monkeypatch.setenv("MISTRAL_API_KEY", mistral_api_key) + monkeypatch.delenv("MISTRAL_AZURE_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock(assert_all_called=False) as router: + route: Final = router.post(COUNT_URL).mock( + return_value=httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 9}) + ) + result: Final = await MistralTokenCounter().count_tokens( + model_to_use="mistral-large-4", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="mistral/mistral-large-4", + ) + assert result is None + assert not route.called + + +@pytest.mark.asyncio +async def test_acount_tokens_never_sends_the_azure_env_key_to_mistral(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ("MISTRAL_API_KEY", "MISTRAL_API_BASE", "MISTRAL_AZURE_API_BASE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("MISTRAL_AZURE_API_KEY", "azure-test-key") + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock(assert_all_called=False) as router: + route: Final = router.post(COUNT_URL).mock( + return_value=httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 9}) + ) + result: Final = await litellm.acount_tokens( + model="mistral/mistral-large-4", messages=[{"role": "user", "content": "Hello"}] + ) + assert (route.called, result.tokenizer_type) == (False, "local_tokenizer") + + +@pytest.mark.asyncio +async def test_count_tokens_mistral_env_key_takes_precedence_over_generic_litellm_api_key( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("MISTRAL_API_KEY", "real-mistral-key") + monkeypatch.delenv("MISTRAL_AZURE_API_KEY", raising=False) + monkeypatch.setattr(litellm, "api_key", "generic-openai-key") + counter: Final = MistralTokenCounter() + with respx.mock: + route: Final = respx.post(COUNT_URL).mock( + return_value=httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 15}) + ) + result: Final = await counter.count_tokens( + model_to_use="mistral-large-4", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="mistral/mistral-large-4", + ) + assert result is not None + assert result.total_tokens == 15 + assert route.calls.last.request.headers["authorization"] == "Bearer real-mistral-key" + + +_MISTRAL_ENV_VARS: Final = ("MISTRAL_API_KEY", "MISTRAL_API_BASE", "MISTRAL_AZURE_API_KEY", "MISTRAL_AZURE_API_BASE") + + +def _set_mistral_env(monkeypatch: pytest.MonkeyPatch, env: dict[str, str]) -> None: + for name in _MISTRAL_ENV_VARS: + monkeypatch.delenv(name, raising=False) + for name, value in env.items(): + monkeypatch.setenv(name, value) + + +@pytest.mark.parametrize( + ("env", "litellm_params"), + [ + ({"MISTRAL_API_KEY": "env-mistral-key"}, {}), + ({"MISTRAL_API_KEY": "env-mistral-key", "MISTRAL_API_BASE": "https://egress.example.com"}, {}), + ({"MISTRAL_API_KEY": "env-mistral-key"}, {"api_key": "deployment-key"}), + ({}, {"api_key": "deployment-key", "api_base": "https://mistral-gateway.example.com"}), + ({}, {"api_key": "deployment-key", "api_base": "https://mistral-gateway.example.com/v1"}), + ], +) +@pytest.mark.asyncio +async def test_count_tokens_sends_the_key_and_base_chat_resolves( + monkeypatch: pytest.MonkeyPatch, env: dict[str, str], litellm_params: dict[str, str] +): + _set_mistral_env(monkeypatch, env) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + _, _, chat_key, chat_base = litellm.get_llm_provider( + model="mistral/mistral-large-4", + api_base=litellm_params.get("api_base"), + api_key=litellm_params.get("api_key"), + ) + with respx.mock: + route: Final = respx.post(f"{chat_base}/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 7}) + ) + result: Final = await MistralTokenCounter().count_tokens( + model_to_use="mistral-large-4", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={"litellm_params": litellm_params}, + request_model="mistral/mistral-large-4", + ) + assert result is not None + assert result.total_tokens == 7 + assert route.calls.last.request.headers["authorization"] == f"Bearer {chat_key}" + + +@pytest.mark.parametrize( + "env", + [ + {"MISTRAL_API_KEY": "mistral-key", "MISTRAL_API_BASE": "https://egress.example.com"}, + { + "MISTRAL_API_KEY": "mistral-key", + "MISTRAL_AZURE_API_KEY": "azure-key", + "MISTRAL_AZURE_API_BASE": "https://mistral-resource.services.ai.azure.com/models", + }, + ], +) +def test_mistral_model_info_key_and_base_match_chat(monkeypatch: pytest.MonkeyPatch, env: dict[str, str]): + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + + _set_mistral_env(monkeypatch, env) + _, _, chat_key, chat_base = litellm.get_llm_provider(model="mistral/mistral-large-4") + model_info: Final = ProviderConfigManager.get_provider_model_info(model=None, provider=LlmProviders.MISTRAL) + assert isinstance(model_info, MistralConfig) + assert (model_info.get_api_base(), model_info.get_api_key()) == (chat_base, chat_key) + + +def test_check_provider_endpoint_lists_the_bundled_mistral_models_without_calling_mistral( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.utils import get_valid_models + + monkeypatch.setenv("MISTRAL_API_KEY", "test-key") + with respx.mock: + listed: Final = get_valid_models(check_provider_endpoint=True, custom_llm_provider="mistral") + + assert listed + assert listed == get_valid_models(check_provider_endpoint=False, custom_llm_provider="mistral") + + +@pytest.mark.asyncio +async def test_mistral_count_tokens_handler_raises_mistral_error_on_non_200(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.mistral.common_utils import MistralError + from litellm.llms.mistral.count_tokens.handler import MistralCountTokensHandler + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + handler: Final = MistralCountTokensHandler() + with respx.mock: + respx.post(COUNT_URL).mock(return_value=httpx.Response(400, text="Bad Request")) + with pytest.raises(MistralError) as exc_info: + await handler.post_count_request({"model": "mistral-large-4", "input": "hi"}, api_key="test-key") + + assert exc_info.value.status_code == 400 + assert "Bad Request" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_mistral_count_tokens_handler_raises_mistral_error_on_invalid_json( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.llms.mistral.common_utils import MistralError + from litellm.llms.mistral.count_tokens.handler import MistralCountTokensHandler + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + handler: Final = MistralCountTokensHandler() + with respx.mock: + respx.post(COUNT_URL).mock(return_value=httpx.Response(200, text="not-valid-json")) + with pytest.raises(MistralError) as exc_info: + await handler.post_count_request({"model": "mistral-large-4", "input": "hi"}, api_key="test-key") + + assert exc_info.value.status_code == 500 + + +@pytest.mark.asyncio +async def test_mistral_count_tokens_handler_rejects_an_empty_input_without_calling_mistral( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.llms.mistral.common_utils import MistralError + from litellm.llms.mistral.count_tokens.handler import MistralCountTokensHandler + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock(assert_all_called=False) as router: + route: Final = router.post(COUNT_URL).mock( + return_value=httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 1}) + ) + with pytest.raises(MistralError) as exc_info: + await MistralCountTokensHandler().handle_count_tokens_request( + model="mistral-large-4", input="", api_key="test-key" + ) + + assert exc_info.value.status_code == 500 + assert not route.called + + +@pytest.mark.asyncio +async def test_count_tokens_system_blocks_instruction(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + counter: Final = MistralTokenCounter() + with respx.mock: + route: Final = respx.post(COUNT_URL).mock( + return_value=httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 12}) + ) + result: Final = await counter.count_tokens( + model_to_use="mistral-large-4", + messages=[{"role": "user", "content": "hi"}], + contents=None, + system=[{"type": "text", "text": "system block"}], + deployment={"litellm_params": {"api_key": "test-key"}}, + request_model="mistral/mistral-large-4", + ) + + assert result is not None + assert result.total_tokens == 12 + assert route.called + sent_json: Final = decoded_json_object(json.loads(route.calls.last.request.content)) + assert sent_json["instructions"] == "system block" + + +@pytest.mark.asyncio +async def test_count_tokens_returns_the_mistral_error_status_instead_of_a_count( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + respx.post(COUNT_URL).mock( + return_value=httpx.Response(401, json={"detail": "Unauthorized"}), + ) + result: Final = await MistralTokenCounter().count_tokens( + model_to_use="mistral-large-4", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "bad-key"}}, + request_model="mistral/mistral-large-4", + ) + + assert result is not None + assert (result.error, result.status_code, result.total_tokens) == (True, 401, 0) + assert "Unauthorized" in (result.error_message or "") + + +async def _sent_count_body( + monkeypatch: pytest.MonkeyPatch, + model: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None = None, + chat_params: dict[str, JsonValue] | None = None, + messages_format: TokenCountMessagesFormat = "openai", + system: object = None, + count_options: dict[str, JsonValue] | None = None, +) -> Mapping[str, object]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock( + return_value=httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 3}) + ) + result: Final = await MistralTokenCounter().count_tokens( + model_to_use=model, + messages=messages, + contents=None, + deployment={"litellm_params": {"api_key": "test-key"}}, + request_model=f"mistral/{model}", + tools=tools, + chat_params=chat_params or {}, + messages_format=messages_format, + system=system, + count_options=count_options or {}, + ) + assert result is not None and result.total_tokens == 3 + return decoded_json_object(json.loads(route.calls.last.request.content)) + + +async def _sent_count_input( + monkeypatch: pytest.MonkeyPatch, model: str, messages: list[dict[str, object]] +) -> list[object]: + sent_input: Final = (await _sent_count_body(monkeypatch, model, messages))["input"] + assert isinstance(sent_input, list) + return sent_input + + +_REASONED_TURN: Final = {"role": "assistant", "content": "Paris.", "reasoning_content": "France's capital is Paris."} +_REASONING_ITEM: Final = { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "France's capital is Paris."}], +} +_ANSWER_ITEM: Final = {"role": "assistant", "content": "Paris."} +_THINK_CHUNK_TURN: Final = { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": [{"type": "text", "text": "France's capital is Paris."}]}, + {"type": "text", "text": "Paris."}, + ], +} + + +@pytest.mark.parametrize( + ("model", "assistant_turn", "counted_turn"), + [ + ("magistral-medium-latest", _REASONED_TURN, [_REASONING_ITEM, _ANSWER_ITEM]), + ( + "magistral-medium-latest", + { + "role": "assistant", + "content": "Paris.", + "thinking_blocks": [{"type": "thinking", "thinking": "France's capital is Paris."}], + }, + [_REASONING_ITEM, _ANSWER_ITEM], + ), + ("magistral-medium-latest", {**_REASONED_TURN, "content": None}, []), + ("magistral-medium-latest", {"role": "assistant", "content": ""}, []), + ("mistral-no-reasoning-test-model", _REASONED_TURN, [_ANSWER_ITEM]), + ("magistral-medium-latest", _THINK_CHUNK_TURN, [_REASONING_ITEM, _ANSWER_ITEM]), + ( + "magistral-medium-latest", + {**_THINK_CHUNK_TURN, "reasoning_content": "A reasoning text chat skips beside a thinking chunk."}, + [_REASONING_ITEM, _ANSWER_ITEM], + ), + ("mistral-no-reasoning-test-model", _THINK_CHUNK_TURN, [_ANSWER_ITEM]), + ], + ids=[ + "reasoning_content", + "thinking_blocks", + "no_turn_output", + "empty_text", + "model_without_reasoning", + "thinking_chunk", + "thinking_chunk_beside_reasoning_content", + "thinking_chunk_model_without_reasoning", + ], +) +@pytest.mark.asyncio +async def test_count_tokens_counts_reasoning_only_where_mistral_chat_replays_it( + monkeypatch: pytest.MonkeyPatch, + model: str, + assistant_turn: dict[str, object], + counted_turn: list[dict[str, object]], +) -> None: + sent_input: Final = await _sent_count_input( + monkeypatch, + model, + [ + {"role": "user", "content": "Capital of France?"}, + assistant_turn, + {"role": "user", "content": "And of Spain?"}, + ], + ) + + assert sent_input == [ + {"role": "user", "content": "Capital of France?"}, + *counted_turn, + {"role": "user", "content": "And of Spain?"}, + ] + + +_IMAGE_URL: Final = "https://example.com/cat.png" + + +@pytest.mark.parametrize( + ("sibling_blocks", "sibling_parts"), + [ + ([], []), + ( + [{"type": "image_url", "image_url": {"url": _IMAGE_URL}}], + [{"type": "input_image", "image_url": _IMAGE_URL, "detail": "auto"}], + ), + ( + [{"type": "file", "file": {"filename": "notes.txt", "file_data": "data:text/plain;base64,aGk="}}], + [{"type": "input_file", "filename": "notes.txt", "file_data": "data:text/plain;base64,aGk="}], + ), + ], + ids=["text_only", "with_image", "with_inline_file"], +) +@pytest.mark.asyncio +async def test_count_tokens_counts_files_sent_by_file_id( + monkeypatch: pytest.MonkeyPatch, + sibling_blocks: list[dict[str, object]], + sibling_parts: list[dict[str, object]], +) -> None: + sent_input: Final = await _sent_count_input( + monkeypatch, + "mistral-large-4", + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize both"}, + *sibling_blocks, + {"type": "file", "file": {"file_id": "file-openai-shape"}}, + {"type": "file", "file_id": "file-mistral-shape"}, + ], + } + ], + ) + + assert sent_input == [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize both"}, + *sibling_parts, + {"type": "input_file", "file_id": "file-openai-shape"}, + {"type": "input_file", "file_id": "file-mistral-shape"}, + ], + } + ] + + +def test_get_mistral_count_tokens_endpoint_other_path() -> None: + config: Final = MistralCountTokensConfig() + endpoint: Final = config.get_mistral_count_tokens_endpoint("https://custom.mistral.ai/api") + assert endpoint == "https://custom.mistral.ai/api/v1/responses/input_tokens" + + +async def _count_with( + monkeypatch: pytest.MonkeyPatch, + model: str, + deployment_params: dict[str, object], + count_response: dict[str, object], + chat_params: dict[str, JsonValue] | None = None, +) -> tuple[TokenCountResponse | None, list[httpx.Request]]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json=count_response)) + result: Final = await MistralTokenCounter().count_tokens( + model_to_use=model, + messages=[{"role": "user", "content": "Give me a city as JSON"}], + contents=None, + deployment={"litellm_params": {"api_key": "test-key", **deployment_params}}, + request_model=f"mistral/{model}", + chat_params=chat_params or {}, + ) + return result, [call.request for call in route.calls] + + +_COUNTED: Final = {"object": "response.input_tokens", "input_tokens": 166} + + +@pytest.mark.asyncio +async def test_count_tokens_sends_the_bridged_response_format_as_text_format( + monkeypatch: pytest.MonkeyPatch, +) -> None: + schema: Final = {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]} + result, requests = await _count_with( + monkeypatch, + "mistral-large-4", + {}, + _COUNTED, + {"response_format": {"type": "json_schema", "json_schema": {"name": "city", "schema": schema, "strict": True}}}, + ) + + assert result is not None and result.total_tokens == 166 + (request,) = requests + assert json.loads(request.content)["text"] == { + "format": {"type": "json_schema", "name": "city", "schema": schema, "strict": True} + } + + +@pytest.mark.parametrize("model", ["mistral-embed", "mistral-ocr-latest"]) +@pytest.mark.asyncio +async def test_count_tokens_leaves_models_mistral_cannot_count_to_the_local_estimate( + monkeypatch: pytest.MonkeyPatch, model: str +) -> None: + result, requests = await _count_with(monkeypatch, model, {}, _COUNTED) + + assert (result, requests) == (None, []) + + +@pytest.mark.asyncio +async def test_count_tokens_uses_the_deployment_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + _, requests = await _count_with(monkeypatch, "mistral-large-4", {"timeout": 7}, _COUNTED) + + (request,) = requests + assert request.extensions["timeout"]["read"] == 7.0 + + +@pytest.mark.asyncio +async def test_count_tokens_falls_back_when_mistral_answers_without_a_count(monkeypatch: pytest.MonkeyPatch) -> None: + result, requests = await _count_with(monkeypatch, "mistral-large-4", {}, {"object": "response.input_tokens"}) + + assert (result, len(requests)) == (None, 1) + + +_LOOKUP_TOOL: Final[dict[str, object]] = { + "type": "function", + "function": {"name": "lookup", "description": "Look up a city", "parameters": {"type": "object"}}, +} +_LOOKUP_ASK: Final[list[dict[str, object]]] = [{"role": "user", "content": "Look up Paris"}] + + +@pytest.mark.parametrize( + ("tool_choice", "counted_tool_choice"), + [ + ("none", "none"), + ("auto", "auto"), + ("required", "required"), + ("any", "required"), + ({"type": "function", "function": {"name": "lookup"}}, {"type": "function", "name": "lookup"}), + ], + ids=["none", "auto", "required", "mistral_any", "named_function"], +) +@pytest.mark.asyncio +async def test_count_tokens_sends_the_tool_choice_chat_sends_in_the_responses_shape( + monkeypatch: pytest.MonkeyPatch, tool_choice: JsonValue, counted_tool_choice: JsonValue +) -> None: + sent: Final = await _sent_count_body( + monkeypatch, "mistral-large-4", _LOOKUP_ASK, tools=[_LOOKUP_TOOL], chat_params={"tool_choice": tool_choice} + ) + + assert sent["tool_choice"] == counted_tool_choice + + +@pytest.mark.asyncio +async def test_count_tokens_sends_no_tool_choice_without_tools_like_chat(monkeypatch: pytest.MonkeyPatch) -> None: + sent: Final = await _sent_count_body(monkeypatch, "mistral-large-4", _LOOKUP_ASK, chat_params={"tool_choice": "none"}) + + assert "tool_choice" not in sent + + +def _chat_messages(model: str, messages: list[dict[str, object]], params: dict[str, object]) -> list[object]: + config: Final = MistralConfig() + optional_params: Final = config.map_openai_params( + non_default_params=params, optional_params={}, model=model, drop_params=False + ) + sent_messages: Final = config.transform_request( + model=model, + messages=json.loads(json.dumps(messages)), + optional_params=optional_params, + litellm_params={}, + headers={}, + )["messages"] + assert isinstance(sent_messages, list) + return sent_messages + + +@pytest.mark.parametrize( + "messages", + [ + [{"role": "user", "content": "Capital of France?"}], + [{"role": "system", "content": "Answer in one word."}, {"role": "user", "content": "Capital of France?"}], + [ + {"role": "system", "content": [{"type": "text", "text": "Answer in one word."}]}, + {"role": "user", "content": "Capital of France?"}, + ], + ], + ids=["no_system_message", "system_text", "system_blocks"], +) +@pytest.mark.asyncio +async def test_count_tokens_counts_the_messages_mistral_chat_sends_magistral_with_reasoning( + monkeypatch: pytest.MonkeyPatch, messages: list[dict[str, object]] +) -> None: + model: Final = "magistral-medium-latest" + chat_sent: Final = _chat_messages(model, messages, {"reasoning_effort": "high"}) + counted: Final = await _sent_count_body(monkeypatch, model, messages, chat_params={"reasoning_effort": "high"}) + without_reasoning: Final = await _sent_count_body(monkeypatch, model, messages) + + counted_input: Final = counted["input"] + assert isinstance(counted_input, list) + assert len(counted_input) == len(chat_sent) + assert counted == without_reasoning + + +@pytest.mark.asyncio +async def test_count_tokens_sends_each_system_message_where_chat_sends_it(monkeypatch: pytest.MonkeyPatch) -> None: + sent: Final = await _sent_count_body( + monkeypatch, + "mistral-large-4", + [ + {"role": "system", "content": "You are a terse assistant."}, + {"role": "user", "content": "Say hello."}, + {"role": "assistant", "content": "Hello."}, + {"role": "system", "content": [{"type": "text", "text": "Now answer "}, {"type": "text", "text": "in French."}]}, + {"role": "user", "content": "Say goodbye."}, + ], + ) + + assert (sent["input"], "instructions" in sent) == ( + [ + {"role": "system", "content": "You are a terse assistant."}, + {"role": "user", "content": "Say hello."}, + {"role": "assistant", "content": "Hello."}, + {"role": "system", "content": "Now answer in French."}, + {"role": "user", "content": "Say goodbye."}, + ], + False, + ) + + +def _text_of(item: object, keys: tuple[str, ...]) -> object: + fields: Final = decoded_json_object(item) + return next((fields[key] for key in keys if fields.get(key)), None) + + +def _texts(items: list[object], *keys: str) -> list[object]: + return list(filter(None, (_text_of(item, keys) for item in items))) + + +@pytest.mark.asyncio +async def test_count_tokens_joins_text_blocks_the_way_mistral_chat_does(monkeypatch: pytest.MonkeyPatch) -> None: + messages: Final[list[dict[str, object]]] = [ + { + "role": "user", + "content": [{"type": "text", "text": "Summarize the plan"}, {"type": "text", "text": " for the release."}], + }, + {"role": "assistant", "content": [{"type": "text", "text": "Ship the "}, {"type": "text", "text": "fix."}]}, + {"role": "user", "content": [{"type": "text", "text": "Look it "}, {"type": "text", "text": "up."}]}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": '{"q": "plan"}'}} + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [{"type": "text", "text": "Release "}, {"type": "text", "text": "on Monday."}], + }, + ] + chat_sent: Final = _chat_messages("mistral-large-4", messages, {}) + counted: Final = await _sent_count_body(monkeypatch, "mistral-large-4", messages) + + counted_input: Final = counted["input"] + assert isinstance(counted_input, list) + assert _texts(counted_input, "content", "output") == _texts(chat_sent, "content") == [ + "Summarize the plan for the release.", + "Ship the fix.", + "Look it up.", + "Release on Monday.", + ] + + +@pytest.mark.asyncio +async def test_count_tokens_keeps_file_ids_where_mistral_chat_sends_them(monkeypatch: pytest.MonkeyPatch) -> None: + messages: Final[list[dict[str, object]]] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Compare"}, + {"type": "file", "file": {"file_id": "file-release-notes"}}, + {"type": "text", "text": "with the plan"}, + ], + } + ] + chat_sent: Final = _chat_messages("mistral-large-4", messages, {}) + sent_input: Final = await _sent_count_input(monkeypatch, "mistral-large-4", messages) + + assert chat_sent == [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Compare"}, + {"type": "file", "file_id": "file-release-notes"}, + {"type": "text", "text": "with the plan"}, + ], + } + ] + assert sent_input == [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Compare"}, + {"type": "input_file", "file_id": "file-release-notes"}, + {"type": "input_text", "text": "with the plan"}, + ], + } + ] + + +_CAPITALS_ASKED: Final = ( + {"role": "user", "content": "Capital of France?"}, + {"role": "assistant", "content": "Paris."}, + {"role": "user", "content": "And of Spain?"}, +) + + +@pytest.mark.parametrize( + "messages", + [ + [ + {"role": "system", "content": ""}, + _CAPITALS_ASKED[0], + {"role": "user", "content": ""}, + _CAPITALS_ASKED[1], + {"role": "user", "content": "", "name": "reviewer"}, + _CAPITALS_ASKED[2], + ], + [ + _CAPITALS_ASKED[0], + {"role": "tool", "tool_call_id": "call_never_made", "content": ""}, + _CAPITALS_ASKED[1], + _CAPITALS_ASKED[2], + ], + ], + ids=["empty_system_and_user_turns", "empty_reply_to_no_tool_call"], +) +@pytest.mark.asyncio +async def test_count_tokens_drops_the_empty_turns_mistral_chat_drops( + monkeypatch: pytest.MonkeyPatch, messages: list[dict[str, object]] +) -> None: + chat_sent: Final = _chat_messages("mistral-large-4", messages, {}) + sent_input: Final = await _sent_count_input(monkeypatch, "mistral-large-4", messages) + + assert sent_input == chat_sent == list(_CAPITALS_ASKED) + + +@pytest.mark.asyncio +async def test_count_tokens_joins_anthropic_tool_result_text_the_way_mistral_chat_does( + monkeypatch: pytest.MonkeyPatch, +) -> None: + messages: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "What's the weather in Paris?"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Paris"}}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [{"type": "text", "text": "sunny, "}, {"type": "text", "text": "21C"}], + } + ], + }, + ] + sent: Final = await _sent_count_body( + monkeypatch, + "mistral-large-4", + messages, + tools=[{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}], + messages_format="anthropic", + ) + + sent_input: Final = sent["input"] + assert isinstance(sent_input, list) + assert {"type": "function_call_output", "call_id": "toolu_1", "output": "sunny, 21C"} in sent_input + + +_ANTHROPIC_LOOKUP_TOOL: Final[dict[str, object]] = { + "name": "lookup", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}}, +} + + +_CHAT_MESSAGES: Final = TypeAdapter(list[dict[str, object]]) + + +def _chat_request_for_anthropic(model: str, body: dict[str, object]) -> Mapping[str, object]: + chat_request: Final = AnthropicAdapter().translate_completion_input_params({**body, "model": model}) + assert chat_request is not None + return chat_request + + +@pytest.mark.parametrize( + ("thinking", "system"), + [ + ({"type": "enabled", "budget_tokens": 2048}, None), + ({"type": "enabled", "budget_tokens": 2048}, "Answer in one word."), + ({"type": "adaptive"}, [{"type": "text", "text": "Answer in one word."}]), + ], + ids=["enabled_no_system", "enabled_system_text", "adaptive_system_blocks"], +) +@pytest.mark.asyncio +async def test_anthropic_count_with_thinking_counts_the_system_mistral_chat_sends_magistral( + monkeypatch: pytest.MonkeyPatch, thinking: dict[str, JsonValue], system: object +) -> None: + model: Final = "magistral-medium-latest" + messages: Final[list[dict[str, object]]] = [{"role": "user", "content": "Capital of France?"}] + chat_request: Final = _chat_request_for_anthropic( + model, {"messages": messages, "thinking": thinking, **({} if system is None else {"system": system})} + ) + chat_messages: Final = _CHAT_MESSAGES.validate_python(chat_request["messages"]) + chat_params: Final = {name: chat_request[name] for name in ("reasoning_effort", "thinking") if name in chat_request} + chat_sent: Final = _chat_messages(model, chat_messages, chat_params) + chat_systems: Final = [m["content"] for m in chat_sent if isinstance(m, dict) and m["role"] == "system"] + + counted: Final = await _sent_count_body( + monkeypatch, model, messages, messages_format="anthropic", system=system, count_options={"thinking": thinking} + ) + without_thinking: Final = await _sent_count_body( + monkeypatch, model, messages, messages_format="anthropic", system=system + ) + + assert counted.get("instructions") == (chat_systems[0] if chat_systems else None) + assert counted == without_thinking + + +@pytest.mark.asyncio +async def test_anthropic_count_joins_system_blocks_the_way_mistral_chat_does(monkeypatch: pytest.MonkeyPatch) -> None: + model: Final = "mistral-large-4" + messages: Final[list[dict[str, object]]] = [{"role": "user", "content": "Capital of France?"}] + system: Final = [{"type": "text", "text": "Answer in one word."}, {"type": "text", "text": "Be terse."}] + chat_request: Final = _chat_request_for_anthropic(model, {"messages": messages, "system": system}) + chat_messages: Final = _CHAT_MESSAGES.validate_python(chat_request["messages"]) + chat_system: Final = _chat_messages(model, chat_messages, {})[0] + assert isinstance(chat_system, dict) + + counted: Final = await _sent_count_body(monkeypatch, model, messages, messages_format="anthropic", system=system) + + assert counted["instructions"] == chat_system["content"] + + +@pytest.mark.parametrize( + ("tool_choice", "counted_tool_choice"), + [ + ({"type": "none"}, "none"), + ({"type": "auto"}, "auto"), + ({"type": "any"}, "required"), + ({"type": "tool", "name": "lookup"}, {"type": "function", "name": "lookup"}), + ], + ids=["none", "auto", "any", "named_tool"], +) +@pytest.mark.asyncio +async def test_anthropic_count_sends_the_tool_choice_mistral_chat_sends( + monkeypatch: pytest.MonkeyPatch, tool_choice: dict[str, JsonValue], counted_tool_choice: JsonValue +) -> None: + sent: Final = await _sent_count_body( + monkeypatch, + "mistral-large-4", + _LOOKUP_ASK, + tools=[_ANTHROPIC_LOOKUP_TOOL], + messages_format="anthropic", + count_options={"tool_choice": tool_choice}, + ) + + assert sent["tool_choice"] == counted_tool_choice + + +@pytest.mark.asyncio +async def test_anthropic_count_rejects_an_unknown_tool_choice_before_calling_mistral() -> None: + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json=_COUNTED)) + result: Final = await MistralTokenCounter().count_tokens( + model_to_use="mistral-large-4", + messages=_LOOKUP_ASK, + contents=None, + deployment={"litellm_params": {"api_key": "test-key"}}, + request_model="mistral/mistral-large-4", + tools=[_ANTHROPIC_LOOKUP_TOOL], + messages_format="anthropic", + count_options={"tool_choice": {"type": "sometimes"}}, + ) + + assert result is not None + assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) + assert not route.called diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index d9ddf0f131a..d9f523b951b 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -17,7 +17,7 @@ from contextlib import contextmanager from dataclasses import replace from functools import partial from pathlib import Path -from types import MappingProxyType +from types import MappingProxyType, SimpleNamespace from typing import Final, NoReturn from unittest.mock import AsyncMock, MagicMock, patch @@ -63,10 +63,10 @@ from litellm.llms.bedrock.count_tokens.kimi_k3 import ( ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler +from litellm.llms.mistral.count_tokens import handler as mistral_count_tokens_http from litellm.llms.openai.responses.count_tokens import handler as openai_count_tokens_http from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import VertexAIPartnerModelsTokenCounter from litellm.proxy._types import ProxyException, TokenCountRequest, UserAPIKeyAuth -from litellm.responses.litellm_completion_transformation.session_handler import ResponsesSessionHandler from litellm.proxy.anthropic_endpoints.endpoints import ( count_tokens as anthropic_count_tokens, ) @@ -84,6 +84,7 @@ from litellm.proxy.proxy_server import ( _try_provider_token_count, count_request_tokens, ) +from litellm.responses.litellm_completion_transformation.session_handler import ResponsesSessionHandler from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES @@ -95,7 +96,6 @@ from litellm.types.utils import ( TokenCountMessagesFormat, TokenCountResponse, ) - from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs @@ -1590,7 +1590,7 @@ async def _count_on_a_model_without_a_count_api( model_list=[ { "model_name": "mistral-tool-result", - "litellm_params": {"model": "mistral/mistral-small-latest", "api_key": "sk-test"}, + "litellm_params": {"model": "mistral/mistral-large-4", "api_key": "sk-test"}, } ], default_litellm_params=router_defaults, @@ -1598,7 +1598,11 @@ async def _count_on_a_model_without_a_count_api( request: Final = _json_request(body, "python-httpx/0.28.1") monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setattr(proxy_server, "proxy_config", proxy_server.ProxyConfig()) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) with respx.mock: + respx.post("https://api.mistral.ai/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 7}) + ) return await anthropic_count_tokens(request, Response(), UserAPIKeyAuth()) @@ -4169,18 +4173,18 @@ async def test_gemini_messages_count_tokens_rejects_an_invalid_edit_without_drop assert sent_bodies == () -_MISTRAL_MODEL: Final = "mistral/mistral-large-latest" +_GROQ_MODEL: Final = "groq/llama-3.3-70b-versatile" -def _mistral_router( +def _groq_router( litellm_params: Mapping[str, JsonValue] = MappingProxyType({}), router_defaults: Mapping[str, JsonValue] = MappingProxyType({}), ) -> Router: return Router( model_list=[ { - "model_name": "mistral-agent", - "litellm_params": {"model": _MISTRAL_MODEL, "api_key": "test-key", **litellm_params}, + "model_name": "groq-agent", + "litellm_params": {"model": _GROQ_MODEL, "api_key": "test-key", **litellm_params}, } ], default_litellm_params=dict(router_defaults), @@ -4228,9 +4232,9 @@ async def _messages_the_inference_polyfill_sends(model: str) -> list[dict[str, J @pytest.mark.asyncio async def test_messages_count_tokens_counts_the_body_the_inference_polyfill_sends_on_a_bridged_provider() -> None: - router: Final = _mistral_router() + router: Final = _groq_router() unedited_body: Final[dict[str, JsonValue]] = { - "model": "mistral-agent", + "model": "groq-agent", "messages": _agent_turns_with_tool_results(3), "tools": [_READ_FILE_TOOL], } @@ -4238,13 +4242,57 @@ async def test_messages_count_tokens_counts_the_body_the_inference_polyfill_send edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES}) unedited: Final = await _input_tokens(router, unedited_body) polyfilled: Final = await _input_tokens( - router, {**unedited_body, "messages": await _messages_the_inference_polyfill_sends(_MISTRAL_MODEL)} + router, {**unedited_body, "messages": await _messages_the_inference_polyfill_sends(_GROQ_MODEL)} ) assert edited < unedited assert edited == polyfilled +class _MistralCountTokensPeer: + def __init__(self) -> None: + self.request_bodies: tuple[bytes, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.request_bodies = (*self.request_bodies, request.content) + return httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 57}, request=request) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_sends_mistral_the_messages_the_inference_polyfill_edits() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "mistral-agent", + "litellm_params": {"model": "mistral/mistral-large-4", "api_key": "test-mistral-key"}, + } + ] + ) + unedited_body: Final[dict[str, JsonValue]] = { + "model": "mistral-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + } + polyfilled_messages: Final = await _messages_the_inference_polyfill_sends("mistral/mistral-large-4") + peer: Final = _MistralCountTokensPeer() + + with patch.object( + mistral_count_tokens_http, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ): + for request_body in ( + {**unedited_body, "context_management": _CLEAR_TOOL_USES}, + {**unedited_body, "messages": polyfilled_messages}, + unedited_body, + ): + assert await _input_tokens(router, request_body) == 57 + + edited, polyfilled, unedited = (json.loads(body)["input"] for body in peer.request_bodies) + assert edited == polyfilled + assert edited != unedited + + @pytest.mark.asyncio @pytest.mark.parametrize( ("deployment_params", "router_defaults", "request_options"), @@ -4260,9 +4308,9 @@ async def test_messages_count_tokens_on_a_bridged_provider_keeps_the_unedited_co router_defaults: dict[str, JsonValue], request_options: dict[str, JsonValue], ) -> None: - router: Final = _mistral_router(deployment_params, router_defaults) + router: Final = _groq_router(deployment_params, router_defaults) unedited_body: Final[dict[str, JsonValue]] = { - "model": "mistral-agent", + "model": "groq-agent", "messages": _agent_turns_with_tool_results(3), "tools": [_READ_FILE_TOOL], **request_options, @@ -4277,9 +4325,9 @@ async def test_messages_count_tokens_on_a_bridged_provider_keeps_the_unedited_co async def test_messages_count_tokens_on_a_bridged_provider_rejects_an_invalid_edit_without_drop_params() -> None: with patch.object(litellm, "drop_params", False): response: Final = await _messages_endpoint_count( - _mistral_router(), + _groq_router(), { - "model": "mistral-agent", + "model": "groq-agent", "messages": _agent_turns_with_tool_results(3), "tools": [_READ_FILE_TOOL], "context_management": _INVALID_COMPACT_EDIT, @@ -4292,9 +4340,9 @@ async def test_messages_count_tokens_on_a_bridged_provider_rejects_an_invalid_ed @pytest.mark.asyncio async def test_messages_count_tokens_on_a_bridged_provider_counts_without_an_invalid_edit_under_drop_params() -> None: - router: Final = _mistral_router() + router: Final = _groq_router() unedited_body: Final[dict[str, JsonValue]] = { - "model": "mistral-agent", + "model": "groq-agent", "messages": _agent_turns_with_tool_results(3), "tools": [_READ_FILE_TOOL], "drop_params": True, @@ -4319,9 +4367,9 @@ def _turns_compacted_by_the_client() -> list[dict[str, JsonValue]]: async def test_messages_count_tokens_on_a_bridged_provider_counts_only_what_follows_a_client_compaction_block() -> None: from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request - router: Final = _mistral_router() + router: Final = _groq_router() sent: Final = await prepare_context_managed_request( - model=_MISTRAL_MODEL, + model=_GROQ_MODEL, messages=[dict(message) for message in _turns_compacted_by_the_client()], tools=[dict(_READ_FILE_TOOL)], system=None, @@ -4333,15 +4381,15 @@ async def test_messages_count_tokens_on_a_bridged_provider_counts_only_what_foll assert sent is not None compacted: Final = await _input_tokens( - router, {"model": "mistral-agent", "messages": _turns_compacted_by_the_client(), "tools": [_READ_FILE_TOOL]} + router, {"model": "groq-agent", "messages": _turns_compacted_by_the_client(), "tools": [_READ_FILE_TOOL]} ) uncompacted: Final = await _input_tokens( - router, {"model": "mistral-agent", "messages": _agent_turns_with_tool_results(3), "tools": [_READ_FILE_TOOL]} + router, {"model": "groq-agent", "messages": _agent_turns_with_tool_results(3), "tools": [_READ_FILE_TOOL]} ) sent_body_count: Final = await _input_tokens( router, { - "model": "mistral-agent", + "model": "groq-agent", "messages": TypeAdapter(list[dict[str, JsonValue]]).validate_python(sent.messages), "system": TypeAdapter(JsonValue).validate_python(sent.system), "tools": [_READ_FILE_TOOL], @@ -5821,6 +5869,128 @@ def test_utils_token_counter_line_stays_one_line_when_the_model_carries_a_newlin assert message.startswith(f'token count - litellm_call_id="count log" request_model={json.dumps(forged_model)} ') +@pytest.mark.asyncio +async def test_mistral_proxy_token_counter_chat_and_responses(monkeypatch: pytest.MonkeyPatch) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "mistral-proxy-model", + "litellm_params": { + "model": "mistral/mistral-large-4", + "api_key": "test-mistral-key", + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + + def count_peer(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 33}, request=request) + + with patch.object( + mistral_count_tokens_http, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(count_peer)), + ): + utils_result: Final = await count_request_tokens( + request=TokenCountRequest( + model="mistral-proxy-model", messages=[{"role": "user", "content": "hello mistral"}] + ), + call_endpoint=True, + ) + responses_result: Final = await count_request_tokens( + request=TokenCountRequest( + model="mistral-proxy-model", messages=[{"role": "user", "content": "hello mistral"}] + ), + call_endpoint=True, + responses_request={"model": "mistral-proxy-model", "input": "hello mistral"}, + ) + + assert utils_result.total_tokens == 33 + assert utils_result.tokenizer_type == "mistral_api" + assert responses_result.total_tokens == 33 + assert responses_result.tokenizer_type == "mistral_api" + + +def _input_texts(item: JsonValue) -> list[str]: + assert isinstance(item, dict) + content: Final = item.get("content") + if isinstance(content, str): + return [content] + assert isinstance(content, list) + return [str(part["text"]) for part in content if isinstance(part, dict) and "text" in part] + + +class _SpendLogsTable: + def __init__(self, rows: tuple[Mapping[str, object], ...]) -> None: + self.rows: Final = rows + + async def query_raw(self, query: str, request_ids: list[str]) -> list[Mapping[str, object]]: + sessions: Final = {row["session_id"] for row in self.rows if row["request_id"] in request_ids} + return [row for row in self.rows if row["session_id"] in sessions] + + +@pytest.mark.asyncio +async def test_responses_count_on_mistral_counts_the_turns_previous_response_id_replays( + monkeypatch: pytest.MonkeyPatch, +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "mistral-counted", + "litellm_params": {"model": "mistral/mistral-large-4", "api_key": "test-mistral-key"}, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + stored_turn: Final = { + "request_id": "chatcmpl-stored", + "call_type": "aresponses", + "session_id": "session-1", + "proxy_server_request": {"input": "Who won in 1998?", "model": "mistral-counted"}, + "response": { + "id": "chatcmpl-stored", + "object": "chat.completion", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "The Bulls."}, "finish_reason": "stop"} + ], + }, + "status": "success", + } + bodies: Final[list[dict[str, JsonValue]]] = [] # mutable-ok: collects what the mock peer receives + + def count_peer(request: httpx.Request) -> httpx.Response: + bodies.append(json.loads(request.content)) + return httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 41}, request=request) + + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", SimpleNamespace(db=_SpendLogsTable((stored_turn,)))) + with patch.object( + mistral_count_tokens_http, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(count_peer)), + ): + result: Final = await count_request_tokens( + request=TokenCountRequest(model="mistral-counted", messages=[{"role": "user", "content": "And in 1999?"}]), + call_endpoint=True, + responses_request={ + "model": "mistral-counted", + "input": "And in 1999?", + "previous_response_id": "chatcmpl-stored", + }, + ) + + assert (result.error, result.total_tokens, result.tokenizer_type) == (False, 41, "mistral_api") + (body,) = bodies + assert "previous_response_id" not in body + sent_input: Final = body["input"] + assert isinstance(sent_input, list) + assert list(itertools.chain.from_iterable(map(_input_texts, sent_input))) == [ + "Who won in 1998?", + "The Bulls.", + "And in 1999?", + ] + + _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"} diff --git a/tests/unit/test_count_tokens_public_api.py b/tests/unit/test_count_tokens_public_api.py index 581682be474..d6a0e070d72 100644 --- a/tests/unit/test_count_tokens_public_api.py +++ b/tests/unit/test_count_tokens_public_api.py @@ -4,7 +4,6 @@ Tests for litellm.acount_tokens() public API. import asyncio import json -import os from unittest.mock import AsyncMock, patch import httpx @@ -12,7 +11,6 @@ import pytest import respx import litellm -from litellm.types.utils import TokenCountResponse def test_acount_tokens_routes_to_openai(): @@ -161,8 +159,6 @@ def test_acount_tokens_no_api_key_falls_back(monkeypatch): assert result.tokenizer_type == "local_tokenizer" - - @pytest.fixture def httpx_transport_clients(monkeypatch): monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) @@ -203,3 +199,30 @@ def test_acount_tokens_sends_chat_tool_use_to_anthropic_in_the_shape_claude_infe assert sent["tools"] == [{"name": "get_weather", "input_schema": schema, "type": "custom"}] assert [message["role"] for message in sent["messages"]] == ["user", "assistant", "user"] assert sent["messages"][2]["content"] == [{"type": "tool_result", "tool_use_id": "call_1", "content": "sunny"}] + + +def test_acount_tokens_counts_mistral_on_the_mistral_count_endpoint(httpx_transport_clients, monkeypatch): + for name in ("MISTRAL_API_BASE", "MISTRAL_AZURE_API_BASE", "MISTRAL_AZURE_API_KEY"): + monkeypatch.delenv(name, raising=False) + with respx.mock(assert_all_called=True) as router: + route = router.post("https://api.mistral.ai/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 25}) + ) + result = asyncio.run( + litellm.acount_tokens( + model="mistral/mistral-large-4", + messages=[{"role": "user", "content": "Hello Mistral!"}], + api_key="test-mistral-key", + ) + ) + + assert (result.total_tokens, result.tokenizer_type, result.request_model) == ( + 25, + "mistral_api", + "mistral/mistral-large-4", + ) + assert route.calls.last.request.headers["authorization"] == "Bearer test-mistral-key" + assert json.loads(route.calls.last.request.content) == { + "model": "mistral-large-4", + "input": [{"role": "user", "content": "Hello Mistral!"}], + }