/
OS-Worldb968155# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0
#
# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
# SPDX-License-Identifier: MIT
"""Create an OpenAI-compatible client using Mistral.AI's API.
Example:
```python
llm_config = {
"config_list": [
{"api_type": "mistral", "model": "open-mixtral-8x22b", "api_key": os.environ.get("MISTRAL_API_KEY")}
]
}
agent = autogen.AssistantAgent("my_agent", llm_config=llm_config)
```
Install Mistral.AI python library using: pip install --upgrade mistralai
Resources:
- https://docs.mistral.ai/getting-started/quickstart/
NOTE: Requires mistralai package version >= 1.0.1
"""
import json
import os
import time
import warnings
from typing import Any, Literal, Optional, Union
from pydantic import Field
from ..import_utils import optional_import_block, require_optional_import
from ..llm_config import LLMConfigEntry, register_llm_config
from .client_utils import should_hide_tools, validate_parameter
from .oai_models import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageToolCall, Choice, CompletionUsage
with optional_import_block():
# Mistral libraries
# pip install mistralai
from mistralai import (
AssistantMessage,
Function,
FunctionCall,
Mistral,
SystemMessage,
ToolCall,
ToolMessage,
UserMessage,
)
@register_llm_config
class MistralLLMConfigEntry(LLMConfigEntry):
api_type: Literal["mistral"] = "mistral"
temperature: float = Field(default=0.7)
top_p: Optional[float] = None
max_tokens: Optional[int] = Field(default=None, ge=0)
safe_prompt: bool = False
random_seed: Optional[int] = None
stream: bool = False
hide_tools: Literal["if_all_run", "if_any_run", "never"] = "never"
tool_choice: Optional[Literal["none", "auto", "any"]] = None
def create_client(self):
raise NotImplementedError("MistralLLMConfigEntry.create_client is not implemented.")
@require_optional_import("mistralai", "mistral")
class MistralAIClient:
"""Client for Mistral.AI's API."""
def __init__(self, **kwargs):
"""Requires api_key or environment variable to be set
Args:
**kwargs: Additional keyword arguments to pass to the Mistral client.
"""
# Ensure we have the api_key upon instantiation
self.api_key = kwargs.get("api_key")
if not self.api_key:
self.api_key = os.getenv("MISTRAL_API_KEY", None)
assert self.api_key, (
"Please specify the 'api_key' in your config list entry for Mistral or set the MISTRAL_API_KEY env variable."
)
if "response_format" in kwargs and kwargs["response_format"] is not None:
warnings.warn("response_format is not supported for Mistral.AI, it will be ignored.", UserWarning)
self._client = Mistral(api_key=self.api_key)
def message_retrieval(self, response: ChatCompletion) -> Union[list[str], list[ChatCompletionMessage]]:
"""Retrieve the messages from the response."""
return [choice.message for choice in response.choices]
def cost(self, response) -> float:
return response.cost
@require_optional_import("mistralai", "mistral")
def parse_params(self, params: dict[str, Any]) -> dict[str, Any]:
"""Loads the parameters for Mistral.AI API from the passed in parameters and returns a validated set. Checks types, ranges, and sets defaults"""
mistral_params = {}
# 1. Validate models
mistral_params["model"] = params.get("model")
assert mistral_params["model"], (
"Please specify the 'model' in your config list entry to nominate the Mistral.ai model to use."
)
# 2. Validate allowed Mistral.AI parameters
mistral_params["temperature"] = validate_parameter(params, "temperature", (int, float), True, 0.7, None, None)
mistral_params["top_p"] = validate_parameter(params, "top_p", (int, float), True, None, None, None)
mistral_params["max_tokens"] = validate_parameter(params, "max_tokens", int, True, None, (0, None), None)
mistral_params["safe_prompt"] = validate_parameter(
params, "safe_prompt", bool, False, False, None, [True, False]
)
mistral_params["random_seed"] = validate_parameter(params, "random_seed", int, True, None, False, None)
mistral_params["tool_choice"] = validate_parameter(
params, "tool_choice", str, False, None, None, ["none", "auto", "any"]
)
# TODO
if params.get("stream", False):
warnings.warn(
"Streaming is not currently supported, streaming will be disabled.",
UserWarning,
)
# 3. Convert messages to Mistral format
mistral_messages = []
tool_call_ids = {} # tool call ids to function name mapping
for message in params["messages"]:
if message["role"] == "assistant" and "tool_calls" in message and message["tool_calls"] is not None:
# Convert OAI ToolCall to Mistral ToolCall
mistral_messages_tools = []
for toolcall in message["tool_calls"]:
mistral_messages_tools.append(
ToolCall(
id=toolcall["id"],
function=FunctionCall(
name=toolcall["function"]["name"],
arguments=json.loads(toolcall["function"]["arguments"]),
),
)
)
mistral_messages.append(AssistantMessage(content="", tool_calls=mistral_messages_tools))
# Map tool call id to the function name
for tool_call in message["tool_calls"]:
tool_call_ids[tool_call["id"]] = tool_call["function"]["name"]
elif message["role"] == "system":
if len(mistral_messages) > 0 and mistral_messages[-1].role == "assistant":
# System messages can't appear after an Assistant message, so use a UserMessage
mistral_messages.append(UserMessage(content=message["content"]))
else:
mistral_messages.append(SystemMessage(content=message["content"]))
elif message["role"] == "assistant":
mistral_messages.append(AssistantMessage(content=message["content"]))
elif message["role"] == "user":
mistral_messages.append(UserMessage(content=message["content"]))
elif message["role"] == "tool":
# Indicates the result of a tool call, the name is the function name called
mistral_messages.append(
ToolMessage(
name=tool_call_ids[message["tool_call_id"]],
content=message["content"],
tool_call_id=message["tool_call_id"],
)
)
else:
warnings.warn(f"Unknown message role {message['role']}", UserWarning)
# 4. Last message needs to be user or tool, if not, add a "please continue" message
if not isinstance(mistral_messages[-1], UserMessage) and not isinstance(mistral_messages[-1], ToolMessage):
mistral_messages.append(UserMessage(content="Please continue."))
mistral_params["messages"] = mistral_messages
# 5. Add tools to the call if we have them and aren't hiding them
if "tools" in params:
hide_tools = validate_parameter(
params, "hide_tools", str, False, "never", None, ["if_all_run", "if_any_run", "never"]
)
if not should_hide_tools(params["messages"], params["tools"], hide_tools):
mistral_params["tools"] = tool_def_to_mistral(params["tools"])
return mistral_params
@require_optional_import("mistralai", "mistral")
def create(self, params: dict[str, Any]) -> ChatCompletion:
# 1. Parse parameters to Mistral.AI API's parameters
mistral_params = self.parse_params(params)
# 2. Call Mistral.AI API
mistral_response = self._client.chat.complete(**mistral_params)
# TODO: Handle streaming
# 3. Convert Mistral response to OAI compatible format
if mistral_response.choices[0].finish_reason == "tool_calls":
mistral_finish = "tool_calls"
tool_calls = []
for tool_call in mistral_response.choices[0].message.tool_calls:
tool_calls.append(
ChatCompletionMessageToolCall(
id=tool_call.id,
function={"name": tool_call.function.name, "arguments": tool_call.function.arguments},
type="function",
)
)
else:
mistral_finish = "stop"
tool_calls = None
message = ChatCompletionMessage(
role="assistant",
content=mistral_response.choices[0].message.content,
function_call=None,
tool_calls=tool_calls,
)
choices = [Choice(finish_reason=mistral_finish, index=0, message=message)]
response_oai = ChatCompletion(
id=mistral_response.id,
model=mistral_response.model,
created=int(time.time()),
object="chat.completion",
choices=choices,
usage=CompletionUsage(
prompt_tokens=mistral_response.usage.prompt_tokens,
completion_tokens=mistral_response.usage.completion_tokens,
total_tokens=mistral_response.usage.prompt_tokens + mistral_response.usage.completion_tokens,
),
cost=calculate_mistral_cost(
mistral_response.usage.prompt_tokens, mistral_response.usage.completion_tokens, mistral_response.model
),
)
return response_oai
@staticmethod
def get_usage(response: ChatCompletion) -> dict:
return {
"prompt_tokens": response.usage.prompt_tokens if response.usage is not None else 0,
"completion_tokens": response.usage.completion_tokens if response.usage is not None else 0,
"total_tokens": (
response.usage.prompt_tokens + response.usage.completion_tokens if response.usage is not None else 0
),
"cost": response.cost if hasattr(response, "cost") else 0,
"model": response.model,
}
@require_optional_import("mistralai", "mistral")
def tool_def_to_mistral(tool_definitions: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Converts AG2 tool definition to a mistral tool format"""
mistral_tools = []
for autogen_tool in tool_definitions:
mistral_tool = {
"type": "function",
"function": Function(
name=autogen_tool["function"]["name"],
description=autogen_tool["function"]["description"],
parameters=autogen_tool["function"]["parameters"],
),
}
mistral_tools.append(mistral_tool)
return mistral_tools
def calculate_mistral_cost(input_tokens: int, output_tokens: int, model_name: str) -> float:
"""Calculate the cost of the mistral response."""
# Prices per 1 thousand tokens
# https://mistral.ai/technology/
model_cost_map = {
"open-mistral-7b": {"input": 0.00025, "output": 0.00025},
"open-mixtral-8x7b": {"input": 0.0007, "output": 0.0007},
"open-mixtral-8x22b": {"input": 0.002, "output": 0.006},
"mistral-small-latest": {"input": 0.001, "output": 0.003},
"mistral-medium-latest": {"input": 0.00275, "output": 0.0081},
"mistral-large-latest": {"input": 0.0003, "output": 0.0003},
"mistral-large-2407": {"input": 0.0003, "output": 0.0003},
"open-mistral-nemo-2407": {"input": 0.0003, "output": 0.0003},
"codestral-2405": {"input": 0.001, "output": 0.003},
}
# Ensure we have the model they are using and return the total cost
if model_name in model_cost_map:
costs = model_cost_map[model_name]
return (input_tokens * costs["input"] / 1000) + (output_tokens * costs["output"] / 1000)
else:
warnings.warn(f"Cost calculation is not implemented for model {model_name}, will return $0.", UserWarning)
return 0