scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

feat(mistral): add native upstream token counting support (#50652) (#50653)

Browse source 
* 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>
This commit is contained in:
Yogeshwaran Ravichandran 2026-11-14 13:49:57 +05:30 • committed by GitHub
parent 9ba216d521
commit 35b9df1f40
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 2272 additions and 142 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
litellm/llms/anthropic/messages/responses_adapters/handler.py
View file

@ -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:

6
litellm/llms/anthropic/messages/responses_adapters/transformation.py
View file

@ -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)

192
litellm/llms/mistral/chat/transformation.py
View file

@ -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

7
litellm/llms/mistral/common_utils.py
View file

@ -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

9
litellm/llms/mistral/count_tokens/__init__.py Normal file
View file

@ -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",
]

108
litellm/llms/mistral/count_tokens/handler.py Normal file
View file

@ -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

469
litellm/llms/mistral/count_tokens/token_counter.py Normal file
View file

@ -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,
}
)

42
litellm/llms/mistral/count_tokens/transformation.py Normal file
View file

@ -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}),
}

45
litellm/llms/openai/responses/count_tokens/token_counter.py
View file

@ -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(

4
litellm/utils.py
View file

@ -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

32
tests/integration/messages_endpoint/test_tool_result_text_wire.py
View file

@ -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

0
tests/unit/llms/mistral/count_tokens/__init__.py Normal file
View file

1225
tests/unit/llms/mistral/count_tokens/test_mistral_token_counter.py Normal file
View file

File diff suppressed because it is too large Load diff

214
tests/unit/proxy/test_proxy_token_counter.py
View file

@ -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"}

31
tests/unit/test_count_tokens_public_api.py
View file

@ -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!"}],
}