refactor(mistral): immutable annotations and a declared premium web search usage field for the ratcheted lint gates

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:

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

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

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

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)."""
 

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():

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)