mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-11 00:25:15 +00:00
Merge pull request #11642 from colesmcintosh/mistral-reasoning
Enhance Mistral model support with reasoning capabilities
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user