Merge pull request #11642 from colesmcintosh/mistral-reasoning

Enhance Mistral model support with reasoning capabilities
This commit is contained in:
Cole McIntosh
2025-06-12 16:42:53 -06:00
committed by GitHub
5 changed files with 569 additions and 19 deletions
+120 -14
View File
@@ -144,20 +144,22 @@ All models listed here https://docs.mistral.ai/platform/endpoints are supported.
:::
| Model Name | Function Call |
|----------------|--------------------------------------------------------------|
| Mistral Small | `completion(model="mistral/mistral-small-latest", messages)` |
| Mistral Medium | `completion(model="mistral/mistral-medium-latest", messages)`|
| Mistral Large 2 | `completion(model="mistral/mistral-large-2407", messages)` |
| Mistral Large Latest | `completion(model="mistral/mistral-large-latest", messages)` |
| Mistral 7B | `completion(model="mistral/open-mistral-7b", messages)` |
| Mixtral 8x7B | `completion(model="mistral/open-mixtral-8x7b", messages)` |
| Mixtral 8x22B | `completion(model="mistral/open-mixtral-8x22b", messages)` |
| Codestral | `completion(model="mistral/codestral-latest", messages)` |
| Mistral NeMo | `completion(model="mistral/open-mistral-nemo", messages)` |
| Mistral NeMo 2407 | `completion(model="mistral/open-mistral-nemo-2407", messages)` |
| Codestral Mamba | `completion(model="mistral/open-codestral-mamba", messages)` |
| Codestral Mamba | `completion(model="mistral/codestral-mamba-latest"", messages)` |
| Model Name | Function Call | Reasoning Support |
|----------------|--------------------------------------------------------------|-------------------|
| Mistral Small | `completion(model="mistral/mistral-small-latest", messages)` | No |
| Mistral Medium | `completion(model="mistral/mistral-medium-latest", messages)`| No |
| Mistral Large 2 | `completion(model="mistral/mistral-large-2407", messages)` | No |
| Mistral Large Latest | `completion(model="mistral/mistral-large-latest", messages)` | No |
| **Magistral Small** | `completion(model="mistral/magistral-small-2506", messages)` | Yes |
| **Magistral Medium** | `completion(model="mistral/magistral-medium-2506", messages)`| Yes |
| Mistral 7B | `completion(model="mistral/open-mistral-7b", messages)` | No |
| Mixtral 8x7B | `completion(model="mistral/open-mixtral-8x7b", messages)` | No |
| Mixtral 8x22B | `completion(model="mistral/open-mixtral-8x22b", messages)` | No |
| Codestral | `completion(model="mistral/codestral-latest", messages)` | No |
| Mistral NeMo | `completion(model="mistral/open-mistral-nemo", messages)` | No |
| Mistral NeMo 2407 | `completion(model="mistral/open-mistral-nemo-2407", messages)` | No |
| Codestral Mamba | `completion(model="mistral/open-codestral-mamba", messages)` | No |
| Codestral Mamba | `completion(model="mistral/codestral-mamba-latest"", messages)` | No |
## Function Calling
@@ -203,6 +205,110 @@ assert isinstance(
)
```
## Reasoning Capabilities (Magistral Models)
Mistral's Magistral models support advanced reasoning capabilities that allow the model to think step-by-step before providing answers. LiteLLM provides seamless integration with these reasoning features through OpenAI-compatible parameters.
### Supported Magistral Models
| Model Name | Function Call |
|----------------|--------------------------------------------------------------|
| Magistral Small | `completion(model="mistral/magistral-small-2506", messages)` |
| Magistral Medium | `completion(model="mistral/magistral-medium-2506", messages)`|
### Using Reasoning Effort
The `reasoning_effort` parameter controls how much effort the model puts into reasoning. When used with magistral models.
```python
from litellm import completion
import os
os.environ['MISTRAL_API_KEY'] = "your-api-key"
response = completion(
model="mistral/magistral-medium-2506",
messages=[
{"role": "user", "content": "What is 15 multiplied by 7?"}
],
reasoning_effort="medium" # Options: "low", "medium", "high"
)
print(response)
```
### Example with System Message
If you already have a system message, LiteLLM will prepend the reasoning instructions:
```python
response = completion(
model="mistral/magistral-medium-2506",
messages=[
{"role": "system", "content": "You are a helpful math tutor."},
{"role": "user", "content": "Explain how to solve quadratic equations."}
],
reasoning_effort="high"
)
# The system message becomes:
# "When solving problems, think step-by-step in <think> tags before providing your final answer...
#
# You are a helpful math tutor."
```
### Usage with LiteLLM Proxy
You can also use reasoning capabilities through the LiteLLM proxy:
<Tabs>
<TabItem value="Curl" label="Curl Request">
```shell
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Content-Type: application/json' \
--data '{
"model": "magistral-medium-2506",
"messages": [
{
"role": "user",
"content": "What is the square root of 144? Show your reasoning."
}
],
"reasoning_effort": "medium"
}'
```
</TabItem>
<TabItem value="openai" label="OpenAI v1.0.0+">
```python
import openai
client = openai.OpenAI(
api_key="anything",
base_url="http://0.0.0.0:4000"
)
response = client.chat.completions.create(
model="magistral-medium-2506",
messages=[
{
"role": "user",
"content": "Calculate the area of a circle with radius 5. Show your work."
}
],
reasoning_effort="high"
)
print(response)
```
</TabItem>
</Tabs>
### Important Notes
- **Model Compatibility**: Reasoning parameters only work with magistral models
- **Backward Compatibility**: Non-magistral models will ignore reasoning parameters and work normally
## Sample Usage - Embedding
```python
from litellm import embedding
+21 -1
View File
@@ -19,6 +19,7 @@ Supported Providers:
- Google AI Studio (`google/`)
- Vertex AI (`vertex_ai/`)
- Perplexity (`perplexity/`)
- Mistral AI (Magistral models) (`mistral/`)
LiteLLM will standardize the `reasoning_content` in the response and `thinking_blocks` in the assistant message.
@@ -39,7 +40,7 @@ LiteLLM will standardize the `reasoning_content` in the response and `thinking_b
## Quick Start
<Tabs>
<TabItem value="sdk" label="SDK">
<TabItem value="anthropic" label="Anthropic">
```python showLineNumbers
from litellm import completion
@@ -57,6 +58,25 @@ response = completion(
print(response.choices[0].message.content)
```
</TabItem>
<TabItem value="mistral" label="Mistral">
```python showLineNumbers
from litellm import completion
import os
os.environ["MISTRAL_API_KEY"] = ""
response = completion(
model="mistral/magistral-medium-2506",
messages=[
{"role": "user", "content": "What is 15 multiplied by 7? Show your reasoning."},
],
reasoning_effort="medium",
)
print(response.choices[0].message.content)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
@@ -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,
@@ -75,7 +75,7 @@ class MistralConfig(OpenAIGPTConfig):
return super().get_config()
def get_supported_openai_params(self, model: str) -> List[str]:
return [
supported_params = [
"stream",
"temperature",
"top_p",
@@ -87,6 +87,12 @@ class MistralConfig(OpenAIGPTConfig):
"stop",
"response_format",
]
# Add reasoning support for magistral models
if "magistral" in model.lower():
supported_params.extend(["thinking", "reasoning_effort"])
return supported_params
def _map_tool_choice(self, tool_choice: str) -> str:
if tool_choice == "auto" or tool_choice == "none":
@@ -96,6 +102,20 @@ class MistralConfig(OpenAIGPTConfig):
else: # openai 'tool_choice' object param not supported by Mistral API
return "any"
@staticmethod
def _get_mistral_reasoning_system_prompt() -> str:
"""
Returns the system prompt for Mistral reasoning models.
Based on Mistral's documentation: https://docs.mistral.ai/capabilities/reasoning/
"""
return """When solving problems, think step-by-step in <think> tags before providing your final answer. Use the following format:
<think>
Your step-by-step reasoning process. Be thorough and work through the problem carefully.
</think>
Then provide a clear, concise answer based on your reasoning."""
def map_openai_params(
self,
non_default_params: dict,
@@ -128,6 +148,12 @@ class MistralConfig(OpenAIGPTConfig):
optional_params["extra_body"] = {"random_seed": value}
if param == "response_format":
optional_params["response_format"] = value
if param == "reasoning_effort" and "magistral" in model.lower():
# Flag that we need to add reasoning system prompt
optional_params["_add_reasoning_prompt"] = True
if param == "thinking" and "magistral" in model.lower():
# Flag that we need to add reasoning system prompt
optional_params["_add_reasoning_prompt"] = True
return optional_params
def _get_openai_compatible_provider_info(
@@ -205,6 +231,57 @@ class MistralConfig(OpenAIGPTConfig):
else:
return super()._transform_messages(new_messages, model, False)
def _add_reasoning_system_prompt_if_needed(
self,
messages: List[AllMessageValues],
optional_params: dict
) -> List[AllMessageValues]:
"""
Add reasoning system prompt for Mistral magistral models when reasoning_effort is specified.
"""
if not optional_params.get("_add_reasoning_prompt", False):
return messages
# Check if there's already a system message
has_system_message = any(msg.get("role") == "system" for msg in messages)
if has_system_message:
# Prepend reasoning instructions to existing system message
for i, msg in enumerate(messages):
if msg.get("role") == "system":
existing_content = msg.get("content", "")
reasoning_prompt = self._get_mistral_reasoning_system_prompt()
# Handle both string and list content, preserving original format
if isinstance(existing_content, str):
# String content - prepend reasoning prompt
new_content: Union[str, list] = f"{reasoning_prompt}\n\n{existing_content}"
elif isinstance(existing_content, list):
# List content - prepend reasoning prompt as text block
new_content = [
{"type": "text", "text": reasoning_prompt + "\n\n"}
] + existing_content
else:
# Fallback for any other type - convert to string
new_content = f"{reasoning_prompt}\n\n{str(existing_content)}"
messages[i] = cast(AllMessageValues, {
**msg,
"content": new_content
})
break
else:
# Add new system message with reasoning instructions
reasoning_message: AllMessageValues = cast(AllMessageValues, {
"role": "system",
"content": self._get_mistral_reasoning_system_prompt()
})
messages = [reasoning_message] + messages
# Remove the internal flag
optional_params.pop("_add_reasoning_prompt", None)
return messages
@classmethod
def _handle_name_in_message(cls, message: AllMessageValues) -> AllMessageValues:
"""
@@ -236,3 +313,31 @@ class MistralConfig(OpenAIGPTConfig):
mistral_tool_calls.append(_tool_call_message)
message["tool_calls"] = mistral_tool_calls # type: ignore
return message
def transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
"""
Transform the overall request to be sent to the API.
For magistral models, adds reasoning system prompt when reasoning_effort is specified.
Returns:
dict: The transformed request. Sent as the body of the API call.
"""
# Add reasoning system prompt if needed (for magistral models)
if "magistral" in model.lower() and optional_params.get("_add_reasoning_prompt", False):
messages = self._add_reasoning_system_prompt_if_needed(messages, optional_params)
# Call parent transform_request which handles _transform_messages
return super().transform_request(
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
+4 -2
View File
@@ -4274,7 +4274,8 @@
"source": "https://mistral.ai/news/magistral",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_reasoning": true
},
"mistral/magistral-small-2506": {
"max_tokens": 40000,
@@ -4287,7 +4288,8 @@
"source": "https://mistral.ai/news/magistral",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_reasoning": true
},
"mistral/mistral-embed": {
"max_tokens": 8192,
@@ -32,3 +32,320 @@ async def test_mistral_chat_transformation():
"is_async": True,
}
)
class TestMistralReasoningSupport:
"""Test suite for Mistral Magistral reasoning functionality."""
def test_get_supported_openai_params_magistral_model(self):
"""Test that magistral models support reasoning parameters."""
mistral_config = MistralConfig()
# Test magistral model supports reasoning parameters
supported_params = mistral_config.get_supported_openai_params("mistral/magistral-medium-2506")
assert "reasoning_effort" in supported_params
assert "thinking" in supported_params
# Test non-magistral model doesn't include reasoning parameters
supported_params_normal = mistral_config.get_supported_openai_params("mistral/mistral-large-latest")
assert "reasoning_effort" not in supported_params_normal
assert "thinking" not in supported_params_normal
def test_map_openai_params_reasoning_effort(self):
"""Test that reasoning_effort parameter is properly mapped for magistral models."""
mistral_config = MistralConfig()
# Test reasoning_effort mapping for magistral model
optional_params = {}
result = mistral_config.map_openai_params(
non_default_params={"reasoning_effort": "low"},
optional_params=optional_params,
model="mistral/magistral-medium-2506",
drop_params=False,
)
assert result.get("_add_reasoning_prompt") is True
# Test reasoning_effort ignored for non-magistral model
optional_params_normal = {}
result_normal = mistral_config.map_openai_params(
non_default_params={"reasoning_effort": "low"},
optional_params=optional_params_normal,
model="mistral/mistral-large-latest",
drop_params=False,
)
assert "_add_reasoning_prompt" not in result_normal
def test_map_openai_params_thinking(self):
"""Test that thinking parameter is properly mapped for magistral models."""
mistral_config = MistralConfig()
# Test thinking mapping for magistral model
optional_params = {}
result = mistral_config.map_openai_params(
non_default_params={"thinking": {"budget": 1000}},
optional_params=optional_params,
model="mistral/magistral-small-2506",
drop_params=False,
)
assert result.get("_add_reasoning_prompt") is True
def test_get_mistral_reasoning_system_prompt(self):
"""Test that the reasoning system prompt is properly formatted."""
prompt = MistralConfig._get_mistral_reasoning_system_prompt()
assert "<think>" in prompt
assert "</think>" in prompt
assert "step-by-step" in prompt
assert isinstance(prompt, str)
assert len(prompt) > 50 # Ensure it's not empty
def test_add_reasoning_system_prompt_no_existing_system_message(self):
"""Test adding reasoning system prompt when no system message exists."""
mistral_config = MistralConfig()
messages = [
{"role": "user", "content": "What is 2+2?"}
]
optional_params = {"_add_reasoning_prompt": True}
result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params)
# Should add a new system message at the beginning
assert len(result) == 2
assert result[0]["role"] == "system"
assert "<think>" in result[0]["content"]
assert result[1]["role"] == "user"
assert result[1]["content"] == "What is 2+2?"
# Should remove the internal flag
assert "_add_reasoning_prompt" not in optional_params
def test_add_reasoning_system_prompt_with_existing_system_message(self):
"""Test adding reasoning system prompt when system message already exists."""
mistral_config = MistralConfig()
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is 2+2?"}
]
optional_params = {"_add_reasoning_prompt": True}
result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params)
# Should modify existing system message
assert len(result) == 2
assert result[0]["role"] == "system"
assert "<think>" in result[0]["content"]
assert "You are a helpful assistant." in result[0]["content"]
assert result[1]["role"] == "user"
# Should remove the internal flag
assert "_add_reasoning_prompt" not in optional_params
def test_add_reasoning_system_prompt_with_existing_list_content(self):
"""Test adding reasoning system prompt when system message has list content."""
mistral_config = MistralConfig()
messages = [
{
"role": "system",
"content": [
{"type": "text", "text": "You are a helpful assistant."},
{"type": "text", "text": "You always provide detailed explanations."}
]
},
{"role": "user", "content": "What is 2+2?"}
]
optional_params = {"_add_reasoning_prompt": True}
result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params)
# Should modify existing system message preserving list format
assert len(result) == 2
assert result[0]["role"] == "system"
assert isinstance(result[0]["content"], list)
# First item should be the reasoning prompt
assert result[0]["content"][0]["type"] == "text"
assert "<think>" in result[0]["content"][0]["text"]
# Original content should be preserved
assert "You are a helpful assistant." in result[0]["content"][1]["text"]
assert "You always provide detailed explanations." in result[0]["content"][2]["text"]
assert result[1]["role"] == "user"
# Should remove the internal flag
assert "_add_reasoning_prompt" not in optional_params
def test_add_reasoning_system_prompt_preserves_content_types(self):
"""Test that reasoning prompt preserves original content types (string vs list)."""
mistral_config = MistralConfig()
# Test with string content
string_messages = [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "Hello"}
]
string_params = {"_add_reasoning_prompt": True}
string_result = mistral_config._add_reasoning_system_prompt_if_needed(string_messages, string_params)
assert isinstance(string_result[0]["content"], str)
assert "<think>" in string_result[0]["content"]
assert "You are helpful." in string_result[0]["content"]
# Test with list content
list_messages = [
{
"role": "system",
"content": [{"type": "text", "text": "You are helpful."}]
},
{"role": "user", "content": "Hello"}
]
list_params = {"_add_reasoning_prompt": True}
list_result = mistral_config._add_reasoning_system_prompt_if_needed(list_messages, list_params)
assert isinstance(list_result[0]["content"], list)
assert list_result[0]["content"][0]["type"] == "text"
assert "<think>" in list_result[0]["content"][0]["text"]
assert "You are helpful." in list_result[0]["content"][1]["text"]
def test_add_reasoning_system_prompt_no_flag(self):
"""Test that no modification happens when _add_reasoning_prompt flag is not set."""
mistral_config = MistralConfig()
messages = [
{"role": "user", "content": "What is 2+2?"}
]
optional_params = {}
result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params)
# Should return messages unchanged
assert result == messages
assert len(result) == 1
def test_transform_request_magistral_with_reasoning(self):
"""Test transform_request method for magistral model with reasoning."""
mistral_config = MistralConfig()
messages = [
{"role": "user", "content": "What is 15 * 7?"}
]
optional_params = {"_add_reasoning_prompt": True}
result = mistral_config.transform_request(
model="mistral/magistral-medium-2506",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
)
# Should have added system message
assert len(result["messages"]) == 2
assert result["messages"][0]["role"] == "system"
assert "<think>" in result["messages"][0]["content"]
assert result["messages"][1]["role"] == "user"
# Should remove internal flag from optional_params
assert "_add_reasoning_prompt" not in result
def test_transform_request_magistral_without_reasoning(self):
"""Test transform_request method for magistral model without reasoning."""
mistral_config = MistralConfig()
messages = [
{"role": "user", "content": "What is 15 * 7?"}
]
optional_params = {}
result = mistral_config.transform_request(
model="mistral/magistral-medium-2506",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
)
# Should not modify messages
assert len(result["messages"]) == 1
assert result["messages"][0]["role"] == "user"
def test_transform_request_non_magistral_with_reasoning_params(self):
"""Test that non-magistral models ignore reasoning parameters."""
mistral_config = MistralConfig()
messages = [
{"role": "user", "content": "What is 15 * 7?"}
]
optional_params = {"_add_reasoning_prompt": True}
result = mistral_config.transform_request(
model="mistral/mistral-large-latest",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
)
# Should not add system message for non-magistral models
assert len(result["messages"]) == 1
assert result["messages"][0]["role"] == "user"
def test_case_insensitive_magistral_detection(self):
"""Test that magistral model detection is case-insensitive."""
mistral_config = MistralConfig()
# Test various case combinations
models_to_test = [
"mistral/Magistral-medium-2506",
"mistral/MAGISTRAL-MEDIUM-2506",
"mistral/magistral-SMALL-2506",
"MaGiStRaL-medium-2506"
]
for model in models_to_test:
supported_params = mistral_config.get_supported_openai_params(model)
assert "reasoning_effort" in supported_params, f"Failed for model: {model}"
def test_end_to_end_reasoning_workflow(self):
"""Test the complete workflow from parameter to system prompt injection."""
mistral_config = MistralConfig()
# Step 1: Map parameters
optional_params = {}
mapped_params = mistral_config.map_openai_params(
non_default_params={"reasoning_effort": "high", "temperature": 0.7},
optional_params=optional_params,
model="mistral/magistral-medium-2506",
drop_params=False,
)
assert mapped_params.get("_add_reasoning_prompt") is True
assert mapped_params.get("temperature") == 0.7
# Step 2: Transform request
messages = [
{"role": "user", "content": "Solve for x: 2x + 5 = 13"}
]
result = mistral_config.transform_request(
model="mistral/magistral-medium-2506",
messages=messages,
optional_params=mapped_params,
litellm_params={},
headers={}
)
# Verify final result
assert len(result["messages"]) == 2
assert result["messages"][0]["role"] == "system"
assert "<think>" in result["messages"][0]["content"]
assert result["messages"][1]["role"] == "user"
assert result["messages"][1]["content"] == "Solve for x: 2x + 5 = 13"
assert result.get("temperature") == 0.7
assert "_add_reasoning_prompt" not in result