From 82334ce52214570944f57c571e054f398ce8bde6 Mon Sep 17 00:00:00 2001 From: Cole McIntosh Date: Wed, 11 Jun 2025 17:22:57 -0600 Subject: [PATCH] [Refactor] Improve type safety in Mistral chat transformation * Added type casting for messages to enhance type safety. * Updated the handling of system messages to ensure correct content formatting. * Cleaned up code for better readability and maintainability. --- litellm/llms/mistral/mistral_chat_transformation.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/llms/mistral/mistral_chat_transformation.py b/litellm/llms/mistral/mistral_chat_transformation.py index 3c90380f12b..b886e3cc964 100644 --- a/litellm/llms/mistral/mistral_chat_transformation.py +++ b/litellm/llms/mistral/mistral_chat_transformation.py @@ -6,7 +6,7 @@ Why separate file? Make it easy to see how transformation works Docs - https://docs.mistral.ai/api/ """ -from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, overload +from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, overload, cast from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_messages_with_content_list_to_str_conversion, @@ -251,17 +251,17 @@ Then provide a clear, concise answer based on your reasoning.""" if msg.get("role") == "system": existing_content = msg.get("content", "") reasoning_prompt = self._get_mistral_reasoning_system_prompt() - messages[i] = { + messages[i] = cast(AllMessageValues, { **msg, "content": f"{reasoning_prompt}\n\n{existing_content}" - } + }) break else: # Add new system message with reasoning instructions - reasoning_message: AllMessageValues = { + reasoning_message: AllMessageValues = cast(AllMessageValues, { "role": "system", "content": self._get_mistral_reasoning_system_prompt() - } + }) messages = [reasoning_message] + messages # Remove the internal flag