From aca1385deffaf24e7dd2d64df7716d333030cc58 Mon Sep 17 00:00:00 2001 From: ElDjeee Date: Thu, 30 Jul 2026 12:31:48 +0200 Subject: [PATCH] refactor(mistral): immutable annotations and a declared premium web search usage field for the ratcheted lint gates --- litellm/llms/mistral/common_utils.py | 42 +++---- .../mistral/conversations/transformation.py | 107 +++++++++--------- litellm/llms/mistral/cost_calculator.py | 13 +-- litellm/types/llms/mistral.py | 9 +- litellm/types/utils.py | 3 + ...st_mistral_conversations_transformation.py | 4 +- .../mistral/test_mistral_cost_calculator.py | 2 +- 7 files changed, 92 insertions(+), 88 deletions(-) diff --git a/litellm/llms/mistral/common_utils.py b/litellm/llms/mistral/common_utils.py index 7c17d1e8592..c5f701df703 100644 --- a/litellm/llms/mistral/common_utils.py +++ b/litellm/llms/mistral/common_utils.py @@ -1,4 +1,4 @@ -from typing import Optional +from typing import Mapping, Sequence import httpx from pydantic import TypeAdapter @@ -10,11 +10,11 @@ from litellm.types.llms.openai import AllMessageValues WEB_SEARCH_TOOL_TYPES: tuple[str, ...] = ("web_search", "web_search_premium") -STR_OBJ_DICT: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object]) -OBJ_LIST: TypeAdapter[list[object]] = TypeAdapter(list[object]) +STR_OBJ_DICT: TypeAdapter[Mapping[str, object]] = TypeAdapter(Mapping[str, object]) +OBJ_LIST: TypeAdapter[Sequence[object]] = TypeAdapter(Sequence[object]) -def is_web_search_request(optional_params: dict) -> bool: +def is_web_search_request(optional_params: Mapping[str, object]) -> bool: """True when a Mistral request should route to the Conversations API for web search.""" params = STR_OBJ_DICT.validate_python(optional_params) if params.get("web_search_options") is not None: @@ -31,33 +31,35 @@ def is_web_search_request(optional_params: dict) -> bool: class MistralModelInfo(BaseLLMModelInfo): def validate_environment( self, - headers: dict, + headers: Mapping[str, str], model: str, - messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - ) -> dict: - if api_key is not None: - headers["Authorization"] = f"Bearer {api_key}" - if "content-type" not in headers and "Content-Type" not in headers: - headers["Content-Type"] = "application/json" - return headers + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: # mutable-ok: BaseLLMModelInfo contract returns the headers dict + auth = {"Authorization": f"Bearer {api_key}"} if api_key is not None else {} + content_type = ( + {} if "content-type" in headers or "Content-Type" in headers else {"Content-Type": "application/json"} + ) + return {**headers, **auth, **content_type} @staticmethod - def get_api_base(api_base: Optional[str] = None) -> Optional[str]: + def get_api_base(api_base: str | None = None) -> str | None: return api_base or get_secret_str("MISTRAL_API_BASE") or "https://api.mistral.ai" @staticmethod - def get_api_key(api_key: Optional[str] = None) -> Optional[str]: + def get_api_key(api_key: str | None = None) -> str | None: return api_key or get_secret_str("MISTRAL_API_KEY") @staticmethod - def get_base_model(model: str) -> Optional[str]: + def get_base_model(model: str) -> str | None: return model.replace("mistral/", "") - def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> list[str]: + def get_models( + self, api_key: str | None = None, api_base: str | None = None + ) -> list[str]: # mutable-ok: BaseLLMModelInfo contract returns List[str] api_base = self.get_api_base(api_base) api_key = self.get_api_key(api_key) if api_base is None or api_key is None: diff --git a/litellm/llms/mistral/conversations/transformation.py b/litellm/llms/mistral/conversations/transformation.py index a497e406567..fa814f40f00 100644 --- a/litellm/llms/mistral/conversations/transformation.py +++ b/litellm/llms/mistral/conversations/transformation.py @@ -12,7 +12,7 @@ annotations for the web sources. Docs - https://docs.mistral.ai/agents/connectors/websearch/ """ -from typing import Optional +from typing import Mapping, Sequence import httpx @@ -66,7 +66,7 @@ def _content_to_str(content: object) -> str: return "" if content is None else str(content) -def _function_call_entry(call: dict[str, object]) -> dict[str, object]: +def _function_call_entry(call: Mapping[str, object]) -> Mapping[str, object]: fn = STR_OBJ_DICT.validate_python(call["function"]) if isinstance(call.get("function"), dict) else {} return { "type": "function.call", @@ -93,26 +93,26 @@ class MistralConversationsConfig(MistralConfig): def should_fake_stream( self, - model: Optional[str], - stream: Optional[bool], - custom_llm_provider: Optional[str] = None, + model: str | None, + stream: bool | None, + custom_llm_provider: str | None = None, ) -> bool: return bool(stream) def get_complete_url( self, - api_base: Optional[str], - api_key: Optional[str], + api_base: str | None, + api_key: str | None, model: str, - optional_params: dict, - litellm_params: dict, - stream: Optional[bool] = None, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + stream: bool | None = None, ) -> str: base = (api_base or "https://api.mistral.ai/v1").rstrip("/") return f"{base}/conversations" @staticmethod - def _build_tools(optional_params: dict[str, object]) -> list[dict[str, object]]: + def _build_tools(optional_params: Mapping[str, object]) -> Sequence[Mapping[str, object]]: raw_tools = optional_params.get("tools") candidate_tools = OBJ_LIST.validate_python(raw_tools) if isinstance(raw_tools, list) else [] dict_tools = [STR_OBJ_DICT.validate_python(tool) for tool in candidate_tools if isinstance(tool, dict)] @@ -127,17 +127,17 @@ class MistralConversationsConfig(MistralConfig): and STR_OBJ_DICT.validate_python(web_search_options).get("premium") is True ) - web_search_tool: tuple[dict[str, object], ...] = ( + web_search_tool: tuple[Mapping[str, object], ...] = ( ({"type": "web_search_premium" if premium else "web_search"},) if wants_web_search else () ) return list(web_search_tool + passthrough) @staticmethod - def _build_completion_args(optional_params: dict[str, object]) -> dict[str, object]: + def _build_completion_args(optional_params: Mapping[str, object]) -> Mapping[str, object]: return {key: optional_params[key] for key in _COMPLETION_ARG_KEYS if optional_params.get(key) is not None} @staticmethod - def _input_entries_for_message(message: AllMessageValues) -> tuple[dict[str, object], ...]: + def _input_entries_for_message(message: AllMessageValues) -> tuple[Mapping[str, object], ...]: """Map one OpenAI message to its Conversations ``inputs`` entries. A ``tool`` message becomes a ``function.result`` entry; an assistant @@ -157,9 +157,9 @@ class MistralConversationsConfig(MistralConfig): }, ) raw_calls = typed.get("tool_calls") if role == "assistant" else None - calls = OBJ_LIST.validate_python(raw_calls) if isinstance(raw_calls, list) else [] + calls = OBJ_LIST.validate_python(raw_calls) if isinstance(raw_calls, list) else () content = _content_to_str(typed.get("content")) - message_entry: tuple[dict[str, object], ...] = ( + message_entry: tuple[Mapping[str, object], ...] = ( ({"role": role, "content": content},) if content or not calls else () ) call_entries = tuple( @@ -170,11 +170,11 @@ class MistralConversationsConfig(MistralConfig): def transform_request( self, model: str, - messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, - headers: dict, - ) -> dict[str, object]: + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + headers: Mapping[str, str], + ) -> dict[str, object]: # mutable-ok: BaseConfig contract; the HTTP handler mutates the returned request body params = STR_OBJ_DICT.validate_python(optional_params) instructions = "\n\n".join( _content_to_str(message.get("content")) @@ -199,24 +199,24 @@ class MistralConversationsConfig(MistralConfig): @staticmethod def _extract_message( - outputs: list[MistralConversationOutput], - ) -> tuple[str, list[ChatCompletionAnnotation]]: + outputs: Sequence[MistralConversationOutput], + ) -> tuple[str, tuple[ChatCompletionAnnotation, ...]]: message_output = next((output for output in outputs if output.type == "message.output"), None) if message_output is None: - return "", [] + return "", () content = message_output.content if isinstance(content, str): - return content, [] - if not isinstance(content, list): - return "", [] + return content, () + if not isinstance(content, tuple): + return "", () text = "".join(chunk.text for chunk in content if chunk.type == "text" and chunk.text) - annotations = [_to_annotation(chunk) for chunk in content if chunk.type == "tool_reference" and chunk.url] + annotations = tuple(_to_annotation(chunk) for chunk in content if chunk.type == "tool_reference" and chunk.url) return text, annotations @staticmethod def _count_web_searches( - usage: Optional[MistralConversationUsage], - outputs: list[MistralConversationOutput], + usage: MistralConversationUsage | None, + outputs: Sequence[MistralConversationOutput], ) -> tuple[int, int]: """Return ``(standard, premium)`` web search call counts. @@ -236,7 +236,7 @@ class MistralConversationsConfig(MistralConfig): @staticmethod def _build_usage( - usage: Optional[MistralConversationUsage], + usage: MistralConversationUsage | None, web_search_requests: int, web_search_premium_requests: int, ) -> Usage: @@ -245,19 +245,23 @@ class MistralConversationsConfig(MistralConfig): total_tokens = ( usage.total_tokens if usage and usage.total_tokens is not None else prompt_tokens + completion_tokens ) - details = PromptTokensDetailsWrapper(web_search_requests=web_search_requests) if web_search_requests else None - usage_obj = Usage( + details = ( + PromptTokensDetailsWrapper( + web_search_requests=web_search_requests, + web_search_premium_requests=web_search_premium_requests or None, + ) + if web_search_requests + else None + ) + return Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens, prompt_tokens_details=details, ) - if web_search_premium_requests: - setattr(usage_obj, "web_search_premium_requests", web_search_premium_requests) - return usage_obj @staticmethod - def _finish_reason(request_data: dict, usage: Optional[MistralConversationUsage]) -> str: + def _finish_reason(request_data: Mapping[str, object], usage: MistralConversationUsage | None) -> str: """The Conversations API returns no finish/stop reason, so infer truncation from the token budget: ``length`` when the completion filled the requested ``max_tokens``, otherwise ``stop``.""" @@ -274,13 +278,13 @@ class MistralConversationsConfig(MistralConfig): raw_response: httpx.Response, model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, - request_data: dict, - messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, + request_data: Mapping[str, object], + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], encoding: object, - api_key: Optional[str] = None, - json_mode: Optional[bool] = None, + api_key: str | None = None, + json_mode: bool | None = None, ) -> ModelResponse: logging_obj.post_call(original_response=raw_response.text) logging_obj.model_call_details["response_headers"] = raw_response.headers @@ -294,15 +298,10 @@ class MistralConversationsConfig(MistralConfig): content=text or None, annotations=annotations or None, ) - model_response.choices = [ - Choices(index=0, message=message, finish_reason=self._finish_reason(request_data, parsed.usage)) - ] - model_response.model = model - if parsed.conversation_id: - model_response.id = parsed.conversation_id - setattr( - model_response, - "usage", - self._build_usage(parsed.usage, standard_searches + premium_searches, premium_searches), + return ModelResponse( + id=parsed.conversation_id or model_response.id, + created=model_response.created, + model=model, + choices=[Choices(index=0, message=message, finish_reason=self._finish_reason(request_data, parsed.usage))], + usage=self._build_usage(parsed.usage, standard_searches + premium_searches, premium_searches), ) - return model_response diff --git a/litellm/llms/mistral/cost_calculator.py b/litellm/llms/mistral/cost_calculator.py index b2e60c0898f..2c897d128ba 100644 --- a/litellm/llms/mistral/cost_calculator.py +++ b/litellm/llms/mistral/cost_calculator.py @@ -6,7 +6,8 @@ model and the response size: web_search at $30 per 1,000 calls ($0.03/call) and web_search_premium at $50 per 1,000 calls ($0.05/call). The Conversations transformation surfaces the total call count as usage.prompt_tokens_details.web_search_requests (the field the shared web search -detection reads) and the premium share as usage.web_search_premium_requests. +detection reads) and the premium share as +usage.prompt_tokens_details.web_search_premium_requests. https://mistral.ai/pricing/ """ @@ -25,12 +26,10 @@ MISTRAL_WEB_SEARCH_PREMIUM_COST_PER_CALL: float = 50.0 / 1000.0 def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> float: """Cost of the web search calls made during a Mistral Conversations request.""" details = getattr(usage, "prompt_tokens_details", None) - total_requests = ( - int(details.web_search_requests) - if isinstance(details, PromptTokensDetailsWrapper) and details.web_search_requests is not None - else 0 - ) - premium_requests = int(getattr(usage, "web_search_premium_requests", 0) or 0) + if not isinstance(details, PromptTokensDetailsWrapper): + return 0.0 + total_requests = int(details.web_search_requests or 0) + premium_requests = int(details.web_search_premium_requests or 0) standard_requests = max(total_requests - premium_requests, 0) return ( standard_requests * MISTRAL_WEB_SEARCH_COST_PER_CALL diff --git a/litellm/types/llms/mistral.py b/litellm/types/llms/mistral.py index fdbb9ffa3b8..f6b7986caf0 100644 --- a/litellm/types/llms/mistral.py +++ b/litellm/types/llms/mistral.py @@ -1,6 +1,7 @@ +from collections.abc import Mapping from typing import Literal -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict from typing_extensions import TypedDict @@ -44,7 +45,7 @@ class MistralConversationOutput(BaseModel): model_config = ConfigDict(extra="allow") type: str | None = None name: str | None = None - content: str | list[MistralConversationContentChunk] | None = None + content: str | tuple[MistralConversationContentChunk, ...] | None = None class MistralConversationUsage(BaseModel): @@ -52,11 +53,11 @@ class MistralConversationUsage(BaseModel): prompt_tokens: int | None = None completion_tokens: int | None = None total_tokens: int | None = None - connectors: dict[str, int] | None = None + connectors: Mapping[str, int] | None = None class MistralConversationsResponse(BaseModel): model_config = ConfigDict(extra="allow") conversation_id: str | None = None - outputs: list[MistralConversationOutput] = Field(default_factory=list) + outputs: tuple[MistralConversationOutput, ...] = () usage: MistralConversationUsage | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 569fce4f7b8..e569294ab87 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1654,6 +1654,9 @@ class PromptTokensDetailsWrapper( google_maps_grounding_requests: int | None = None """Number of Grounding with Google Maps requests made by the tool call. Used for Gemini to calculate Maps cost.""" + web_search_premium_requests: int | None = None + """Number of premium web search requests. Used for Mistral to bill web_search_premium calls.""" + tool_use_tokens: int | None = None """Prompt tokens consumed by server-side tool use (e.g. Gemini grounding via googleSearch).""" diff --git a/tests/test_litellm/llms/mistral/conversations/test_mistral_conversations_transformation.py b/tests/test_litellm/llms/mistral/conversations/test_mistral_conversations_transformation.py index d4136e66391..f3bff71925c 100644 --- a/tests/test_litellm/llms/mistral/conversations/test_mistral_conversations_transformation.py +++ b/tests/test_litellm/llms/mistral/conversations/test_mistral_conversations_transformation.py @@ -342,7 +342,7 @@ def test_transform_response_counts_web_searches_from_connectors(): } ) assert mr.usage.prompt_tokens_details.web_search_requests == 3 - assert mr.usage.web_search_premium_requests == 2 + assert mr.usage.prompt_tokens_details.web_search_premium_requests == 2 def test_transform_response_counts_fall_back_to_tool_execution_by_name(): @@ -359,7 +359,7 @@ def test_transform_response_counts_fall_back_to_tool_execution_by_name(): } ) assert mr.usage.prompt_tokens_details.web_search_requests == 2 - assert mr.usage.web_search_premium_requests == 1 + assert mr.usage.prompt_tokens_details.web_search_premium_requests == 1 def test_transform_response_string_content_no_citations(): diff --git a/tests/test_litellm/llms/mistral/test_mistral_cost_calculator.py b/tests/test_litellm/llms/mistral/test_mistral_cost_calculator.py index a92793dac0d..a96de69821a 100644 --- a/tests/test_litellm/llms/mistral/test_mistral_cost_calculator.py +++ b/tests/test_litellm/llms/mistral/test_mistral_cost_calculator.py @@ -118,7 +118,7 @@ def test_completion_cost_adds_web_search_on_top_of_tokens(): def test_completion_cost_premium_rate(): """web_search_premium calls are billed at $0.05 each.""" response = _conversation_response(connectors={"web_search_premium": 2}) - assert response.usage.web_search_premium_requests == 2 + assert response.usage.prompt_tokens_details.web_search_premium_requests == 2 assert _web_search_surcharge(response) == pytest.approx(0.10)