From 35b9df1f4040ba17af06b936cb65026fb1f90ff6 Mon Sep 17 00:00:00 2001 From: Yogeshwaran Ravichandran <96047771+yogeshwaran10@users.noreply.github.com> Date: Sat, 14 Nov 2026 13:49:57 +0530 Subject: [PATCH] feat(mistral): add native upstream token counting support (#50652) (#50653) * feat(mistral): add native upstream token counting support * test(mistral): standardize on mistral-large-4 and add tool calls tests * fix(mistral): unify credential resolution and guard Azure endpoint boundary * fix(mistral): resolve api base from provider env vars and route responses count natively * fix(test): use groq for bridged proxy tests and increase mistral coverage Switch bridged provider tests in proxy token counter from mistral to groq, with comments noting that PR #50653 added native mistral token counting so groq tests the local estimate fallback. Also add test coverage for mistral error responses, invalid json, and system instruction blocks to satisfy codecov. * fix(lint): fix unit package layout, ruff strict rules, and type discipline in mistral token counter * fix(mistral): suppress basedpyright warnings on untyped httpx handler * fix(mistral): keep chat credential resolution and count with the key chat resolves The token counter now resolves its Mistral key and base through the same helper chat uses, so a count goes to the endpoint chat would call. A deployment that would route to Azure, or an Azure-only env key, falls back to the local estimate instead of sending that key to Mistral. Chat key and base precedence is unchanged from main. * fix(mistral): prefix live /v1/models ids with mistral/ so listed ids route * fix(mistral): count follow-up turns, replayed reasoning, and file_id files the way Mistral chat sends them Mistral has no native Responses API in litellm, so /v1/responses inference runs through the chat bridge, which restores previous_response_id history. The count now takes the same bridge instead of forwarding previous_response_id to Mistral's stateless count endpoint, so the separate Responses count method is gone. The chat-format count adds a reasoning item where Mistral chat replays reasoning as a think chunk, and an input_file part for files sent by file_id, so both reach Mistral's count the way inference sends them. Tests now go through the real handler with respx, and the integration Mistral peer serves the count route. * fix(mistral): leave empty assistant turns out of the count, as Mistral chat does Mistral chat drops an assistant turn with no content and no tool calls before sending, and Mistral's count endpoint rejects one. A follow-up whose stored turn spent its whole output budget on reasoning hit that 400 through previous_response_id, so the count now skips the same turns chat skips. * fix(mistral): count what chat sends and keep the bundled Mistral model list A user turn that mixed file_id files with an image or an inline file lost the image or inline file in the count, since the converted content arrives as a tuple and only lists were read. Both content readers now take either shape A /v1/responses count with a json_schema text.format now sends that format to Mistral, so the count matches what the bridged chat call bills Counts for Mistral models that are not chat models (embeddings, OCR, transcription) now fall back to the local estimate instead of returning Mistral's 400 Count calls now use the deployment's timeout when one is set, and a count response without input_tokens falls back to the local estimate MistralConfig.get_models goes back to the bundled model list, so turning on check_provider_endpoint does not change what /v1/models shows for mistral/* The previous_response_id count test now fakes the spend logs table instead of the history lookup, so the real lookup and its query run * fix(mistral): count tool_choice, the magistral reasoning prompt and system messages like chat, and never send the Azure key to Mistral * fix(mistral): answer 408 on a count timeout and use the shared deployment timeout Main's native counters answer a count timeout with 408 and a connection error with 500. The Mistral counter fell back to a local estimate on both, so a hung Mistral endpoint looked like a successful count. It now raises the same count errors, and it reads the deployment timeout with the shared deployment_timeout helper instead of its own copy * fix(mistral): join count text blocks the way Mistral chat does * fix(mistral): count file ids in order, drop the empty turns chat drops, and join tool result text like Mistral chat * fix(mistral): count the thinking chunks Mistral chat replays in assistant history The OpenAI-format count now takes Mistral chat's own replayed turns, so a Magistral assistant turn whose content carries a native thinking chunk is counted as the reasoning chat sends, and a reasoning_content beside that chunk is skipped the way chat skips it --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../messages/responses_adapters/handler.py | 30 +- .../responses_adapters/transformation.py | 6 + litellm/llms/mistral/chat/transformation.py | 192 ++- litellm/llms/mistral/common_utils.py | 7 + litellm/llms/mistral/count_tokens/__init__.py | 9 + litellm/llms/mistral/count_tokens/handler.py | 108 ++ .../mistral/count_tokens/token_counter.py | 469 +++++++ .../mistral/count_tokens/transformation.py | 42 + .../responses/count_tokens/token_counter.py | 45 +- litellm/utils.py | 4 + .../test_tool_result_text_wire.py | 32 +- .../llms/mistral/count_tokens/__init__.py | 0 .../test_mistral_token_counter.py | 1225 +++++++++++++++++ tests/unit/proxy/test_proxy_token_counter.py | 214 ++- tests/unit/test_count_tokens_public_api.py | 31 +- 15 files changed, 2272 insertions(+), 142 deletions(-) create mode 100644 litellm/llms/mistral/count_tokens/__init__.py create mode 100644 litellm/llms/mistral/count_tokens/handler.py create mode 100644 litellm/llms/mistral/count_tokens/token_counter.py create mode 100644 litellm/llms/mistral/count_tokens/transformation.py create mode 100644 tests/unit/llms/mistral/count_tokens/__init__.py create mode 100644 tests/unit/llms/mistral/count_tokens/test_mistral_token_counter.py 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!"}], + }