mistral_common 1.0.0__py3-none-any.whl

mistral_common/base.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class MistralBase(BaseModel):
+    """
+    Base class for all Mistral Pydantic models.
+    """
+
+    model_config = ConfigDict(extra="forbid", strict=True, validate_default=True, use_enum_values=True)

mistral_common/exceptions.py

@@ -0,0 +1,67 @@
+from typing import Optional
+
+
+class MistralCommonException(Exception):
+    message: str = "Internal server error"
+
+    def __init__(
+        self,
+        message: Optional[str] = None,
+    ) -> None:
+        if message:
+            self.message = message
+
+
+class TokenizerException(MistralCommonException):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class UnsupportedTokenizerFeatureException(MistralCommonException):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class InvalidRequestException(MistralCommonException):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class InvalidSystemPromptException(MistralCommonException):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class InvalidMessageStructureException(MistralCommonException):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class InvalidAssistantMessageException(MistralCommonException):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class InvalidToolMessageException(MistralCommonException):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class InvalidToolSchemaException(MistralCommonException):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class InvalidUserMessageException(MistralCommonException):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class InvalidFunctionCallException(MistralCommonException):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class InvalidToolException(MistralCommonException):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)

mistral_common/protocol/base.py

@@ -0,0 +1,9 @@
+from typing import Optional
+
+from mistral_common.base import MistralBase
+
+
+class UsageInfo(MistralBase):
+    prompt_tokens: int = 0
+    total_tokens: int = 0
+    completion_tokens: Optional[int] = 0

mistral_common/protocol/embedding/request.py

@@ -0,0 +1,10 @@
+from typing import List, Optional, Union
+
+from mistral_common.base import MistralBase
+from pydantic import Field
+
+
+class EmbeddingRequest(MistralBase):
+    input: Union[str, List[str]] = Field(description="Text to embed.")
+    model: str = Field(description="ID of the model to use.")
+    encoding_format: Optional[str] = Field(default="float", description="The format to return the embeddings in.")

mistral_common/protocol/embedding/response.py

@@ -0,0 +1,20 @@
+from typing import List
+
+from mistral_common.base import MistralBase
+from mistral_common.protocol.base import UsageInfo
+from mistral_common.protocol.utils import random_uuid
+from pydantic import Field
+
+
+class EmbeddingObject(MistralBase):
+    object: str = Field(default="embedding", description="The type of the object returned.")
+    embedding: List[float] = Field(description="The type of the object returned.")
+    index: int = Field(description="The index of the embedding in the input text.")
+
+
+class EmbeddingResponse(MistralBase):
+    id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
+    object: str = Field(default="list", description="The type of the object returned.")
+    data: List[EmbeddingObject] = Field(description="List of embeddings.")
+    model: str = Field(description="The model used to generate the embeddings.")
+    usage: UsageInfo

mistral_common/protocol/instruct/messages.py

@@ -0,0 +1,56 @@
+from enum import Enum
+from typing import List, Literal, Optional, Union
+
+from pydantic import Field
+from typing_extensions import Annotated  # compatibility with 3.8
+
+from mistral_common.base import MistralBase
+from mistral_common.protocol.instruct.tool_calls import ToolCall
+
+
+class Roles(str, Enum):
+    system = "system"
+    user = "user"
+    assistant = "assistant"
+    tool = "tool"
+
+
+class ChunkTypes(str, Enum):
+    text = "text"
+
+
+class ContentChunk(MistralBase):
+    type: ChunkTypes = ChunkTypes.text
+    text: str
+
+
+class BaseMessage(MistralBase):
+    role: Literal[Roles.system, Roles.user, Roles.assistant, Roles.tool]
+
+
+class UserMessage(BaseMessage):
+    role: Literal[Roles.user] = Roles.user
+    content: Union[str, List[ContentChunk]]
+
+
+class SystemMessage(BaseMessage):
+    role: Literal[Roles.system] = Roles.system
+    content: Union[str, List[ContentChunk]]
+
+
+class AssistantMessage(BaseMessage):
+    role: Literal[Roles.assistant] = Roles.assistant
+    content: Optional[str] = None
+    tool_calls: Optional[List[ToolCall]] = None
+
+
+class ToolMessage(BaseMessage):
+    content: str
+    role: Literal[Roles.tool] = Roles.tool
+    tool_call_id: Optional[str] = None
+
+    # Deprecated in V3 tokenization
+    name: Optional[str] = None
+
+
+ChatMessage = Annotated[Union[SystemMessage, UserMessage, AssistantMessage, ToolMessage], Field(discriminator="role")]

mistral_common/protocol/instruct/request.py

@@ -0,0 +1,29 @@
+from enum import Enum
+from typing import List, Optional
+
+from pydantic import Field
+
+from mistral_common.base import MistralBase
+from mistral_common.protocol.instruct.messages import ChatMessage
+from mistral_common.protocol.instruct.tool_calls import Tool, ToolChoice
+
+
+class ResponseFormats(str, Enum):
+    text: str = "text"
+    json: str = "json_object"
+
+
+class ResponseFormat(MistralBase):
+    type: ResponseFormats = ResponseFormats.text
+
+
+class ChatCompletionRequest(MistralBase):
+    model: Optional[str] = None
+    messages: List[ChatMessage]
+    response_format: ResponseFormat = Field(default_factory=ResponseFormat)
+    tools: List[Tool] = Field(default_factory=list)
+    tool_choice: ToolChoice = ToolChoice.auto
+    temperature: float = Field(default=0.7, ge=0.0, le=1.0)
+    top_p: float = Field(default=1.0, ge=0.0, le=1.0)
+    max_tokens: Optional[int] = Field(default=None, ge=0)
+    random_seed: Optional[int] = Field(default=None, ge=0)

mistral_common/protocol/instruct/response.py

@@ -0,0 +1,66 @@
+import time
+from enum import Enum
+from typing import List, Optional
+
+from pydantic import Field
+
+from mistral_common.base import MistralBase
+from mistral_common.protocol.base import UsageInfo
+from mistral_common.protocol.instruct.tool_calls import ToolCall
+from mistral_common.protocol.utils import random_uuid
+
+
+class FinishReason(str, Enum):
+    stop: str = "stop"
+    length: str = "length"
+    model_length: str = "model_length"
+    error: str = "error"
+    tool_call: str = "tool_calls"
+
+
+class ChatCompletionTokenLogprobs(MistralBase):
+    token: str
+    logprob: float
+    bytes: List[int]
+
+
+class ChatCompletionResponseChoiceLogprobs(MistralBase):
+    content: List[ChatCompletionTokenLogprobs]
+
+
+class DeltaMessage(MistralBase):
+    role: Optional[str] = None
+    content: Optional[str] = None
+    tool_calls: Optional[List[ToolCall]] = None
+
+
+class ChatCompletionResponseChoice(MistralBase):
+    index: int
+    message: DeltaMessage
+    finish_reason: Optional[FinishReason] = None
+    logprobs: Optional[ChatCompletionResponseChoiceLogprobs] = None
+
+
+class ChatCompletionResponse(MistralBase):
+    id: str = Field(default_factory=lambda: f"chatcmpl-{random_uuid()}")
+    object: str = "chat.completion"
+    created: int = Field(default_factory=lambda: int(time.time()))
+    model: str
+    choices: List[ChatCompletionResponseChoice]
+    usage: UsageInfo
+
+
+class ChatCompletionResponseStreamChoice(MistralBase):
+    index: int
+    delta: DeltaMessage
+    finish_reason: Optional[FinishReason] = None
+    logprobs: Optional[ChatCompletionResponseChoiceLogprobs] = None
+
+
+class ChatCompletionStreamResponse(MistralBase):
+    id: str = Field(default_factory=lambda: f"chatcmpl-{random_uuid()}")
+    object: str = "chat.completion.chunk"
+    created: int = Field(default_factory=lambda: int(time.time()))
+    model: str
+    choices: List[ChatCompletionResponseStreamChoice]
+    usage: Optional[UsageInfo] = None

mistral_common/protocol/instruct/tool_calls.py

@@ -0,0 +1,36 @@
+from enum import Enum
+from typing import Any, Dict
+
+from mistral_common.base import MistralBase
+
+
+class Function(MistralBase):
+    name: str
+    description: str = ""
+    parameters: Dict[str, Any]
+
+
+class ToolType(str, Enum):
+    function = "function"
+
+
+class ToolChoice(str, Enum):
+    auto: str = "auto"
+    none: str = "none"
+    any: str = "any"
+
+
+class Tool(MistralBase):
+    type: ToolType = ToolType.function
+    function: Function
+
+
+class FunctionCall(MistralBase):
+    name: str
+    arguments: str
+
+
+class ToolCall(MistralBase):
+    id: str = "null"  # required for V3 tokenization
+    type: ToolType = ToolType.function
+    function: FunctionCall

mistral_common/protocol/instruct/validator.py

@@ -0,0 +1,282 @@
+import re
+from enum import Enum
+from typing import List
+
+from jsonschema import Draft7Validator, SchemaError
+
+from mistral_common.exceptions import (
+    InvalidAssistantMessageException,
+    InvalidFunctionCallException,
+    InvalidMessageStructureException,
+    InvalidRequestException,
+    InvalidSystemPromptException,
+    InvalidToolException,
+    InvalidToolMessageException,
+    InvalidToolSchemaException,
+)
+from mistral_common.protocol.instruct.messages import (
+    AssistantMessage,
+    ChatMessage,
+    Roles,
+    SystemMessage,
+    ToolMessage,
+    UserMessage,
+)
+from mistral_common.protocol.instruct.request import ChatCompletionRequest
+from mistral_common.protocol.instruct.tool_calls import (
+    Function,
+    FunctionCall,
+    Tool,
+    ToolCall,
+)
+
+
+class ValidationMode(Enum):
+    serving = "serving"
+    test = "test"
+
+
+class MistralRequestValidator:
+    def __init__(self, mode: ValidationMode = ValidationMode.test):
+        self._mode = mode
+
+    def validate_messages(self, messages: List[ChatMessage]) -> None:
+        """
+        Validates the list of messages
+        """
+        self._validate_message_list_structure(messages)
+        self._validate_message_list_content(messages)
+
+    def validate_request(self, request: ChatCompletionRequest) -> ChatCompletionRequest:
+        """
+        Validates the request
+        """
+
+        if self._mode == ValidationMode.serving:
+            if request.model is None:
+                raise InvalidRequestException("Model name parameter is required for serving mode")
+
+        # Validate the messages
+        self.validate_messages(request.messages)
+
+        # Validate the tools
+        self._validate_tools(request.tools)
+
+        return request
+
+    def _validate_function(self, function: Function) -> None:
+        """
+        Checks:
+        - That the function schema is valid
+        """
+        try:
+            Draft7Validator.check_schema(function.parameters)
+        except SchemaError as e:
+            raise InvalidToolSchemaException(f"Invalid tool schema: {e.message}")
+
+        if not re.match(r"^[a-zA-Z0-9_-]{1,64}$", function.name):
+            raise InvalidToolException(
+                f"Function name was {function.name} but must be a-z, A-Z, 0-9, "
+                "or contain underscores and dashes, with a maximum length of 64."
+            )
+
+    def _validate_tools(self, tools: List[Tool]) -> None:
+        """
+        Checks:
+        - That the tool schemas are valid
+        """
+
+        for tool in tools:
+            self._validate_function(tool.function)
+
+    def _validate_user_message(self, message: UserMessage) -> None:
+        pass
+
+    def _validate_tool_message(self, message: ToolMessage) -> None:
+        """
+        Checks:
+        - The tool name is valid
+        """
+        if message.name is not None:
+            if not re.match(r"^[a-zA-Z0-9_-]{1,64}$", message.name):
+                raise InvalidToolMessageException(
+                    f"Function name was {message.name} but must be a-z, A-Z, 0-9, "
+                    "or contain underscores and dashes, with a maximum length of 64."
+                )
+
+    def _validate_system_message(self, message: SystemMessage) -> None:
+        """
+        Checks:
+        - That the system prompt has content
+        """
+        if message.content is None:
+            raise InvalidSystemPromptException("System prompt must have content")
+
+    def _validate_function_call(self, function_call: FunctionCall) -> None:
+        """
+        Checks:
+        - That the function call has a valid name
+        """
+        if not re.match(r"^[a-zA-Z0-9_-]{1,64}$", function_call.name):
+            raise InvalidFunctionCallException(
+                f"Function name was {function_call.name} but must be a-z, A-Z, 0-9, "
+                "or contain underscores and dashes, with a maximum length of 64."
+            )
+
+    def _validate_tool_call(self, tool_call: ToolCall) -> None:
+        """
+        Checks:
+        - That the tool call has a valid function
+        """
+
+        self._validate_function_call(tool_call.function)
+
+    def _validate_assistant_message(self, message: AssistantMessage) -> None:
+        """
+        Checks:
+        - That the assistant message has either text or tool_calls, but not both
+        - That the tool calls are valid
+        """
+
+        # Validate that the message has either text or tool_calls, but not both
+        if (message.content is None) == (message.tool_calls is None):
+            raise InvalidAssistantMessageException("Assistant message must have either content or tool_calls")
+
+        # If we have tool calls, validate them
+        if message.tool_calls is not None:
+            # Validate that the tool calls are valid
+            for tool_call in message.tool_calls:
+                self._validate_tool_call(tool_call)
+
+    def _validate_tool_calls_followed_by_tool_messages(self, messages: List[ChatMessage]) -> None:
+        """
+        Checks:
+        - That the number of tool calls and tool messages are the same
+        - That the tool calls are followed by tool messages
+        """
+        prev_role = None
+        expected_tool_messages = 0
+        for message in messages:
+            if prev_role is None:
+                prev_role = message.role
+                continue
+
+            if message.role == Roles.tool:
+                expected_tool_messages -= 1
+            elif message.role == Roles.assistant:
+                # if we have an assistant message and we have not recieved all the function calls
+                # we need to raise an exception
+                if expected_tool_messages != 0:
+                    raise InvalidMessageStructureException("Not the same number of function calls and responses")
+
+                if message.tool_calls is not None:
+                    # Validate that the number of function calls and responses are the same
+                    expected_tool_messages = len(message.tool_calls)
+
+            prev_role = message.role
+
+        if expected_tool_messages != 0:
+            raise InvalidMessageStructureException("Not the same number of function calls and responses")
+
+    def _validate_message_order(self, messages: List[ChatMessage]) -> None:
+        """
+        Validates the order of the messages, for example user -> assistant -> user -> assistant -> ...
+        """
+        previous_role = None
+        for message in messages:
+            current_role = message.role
+
+            if previous_role is not None:
+                if previous_role == Roles.system:
+                    expected_roles = {Roles.user, Roles.assistant, Roles.system}
+                elif previous_role == Roles.user:
+                    expected_roles = {Roles.assistant, Roles.system, Roles.user}
+                elif previous_role == Roles.assistant:
+                    expected_roles = {Roles.user, Roles.tool}
+                elif previous_role == Roles.tool:
+                    expected_roles = {Roles.assistant, Roles.tool}
+
+                if current_role not in expected_roles:
+                    raise InvalidMessageStructureException(
+                        f"Unexpected role '{current_role.value}' after role '{previous_role.value}'"
+                    )
+
+            previous_role = current_role
+
+    def _validate_message_list_structure(self, messages: List[ChatMessage]) -> None:
+        """
+        Validates the structure of the list of messages
+
+        For example the messages must be in the correct order of user/assistant/tool
+        """
+
+        if len(messages) == 0:
+            raise InvalidMessageStructureException("Conversation must have at least one message")
+
+        # If we have one message it must be a user or a system message
+        if len(messages) == 1:
+            if messages[0].role not in {Roles.user, Roles.system}:
+                raise InvalidMessageStructureException("Conversation must start with a user message or system message")
+
+        else:
+            # The last message must be a user or tool message
+            last_message_role = messages[-1].role
+            expected_roles = {Roles.user, Roles.tool}
+
+            if last_message_role not in expected_roles:
+                expected_roles_str = ", ".join(role.value for role in expected_roles)
+                raise InvalidMessageStructureException(
+                    f"Expected last role to be one of: [{expected_roles_str}] but got {last_message_role.value}"
+                )
+
+        self._validate_message_order(messages)
+        self._validate_tool_calls_followed_by_tool_messages(messages)
+
+    def _validate_message_list_content(self, messages: List[ChatMessage]) -> None:
+        """
+        Validates the content of the messages
+        """
+
+        for message in messages:
+            if message.role == Roles.user:
+                self._validate_user_message(message)
+            elif message.role == Roles.assistant:
+                self._validate_assistant_message(message)
+            elif message.role == Roles.tool:
+                self._validate_tool_message(message)
+            elif message.role == Roles.system:
+                self._validate_system_message(message)
+            else:
+                raise InvalidRequestException(f"Unsupported message type {type(message)}")
+
+
+class MistralRequestValidatorV3(MistralRequestValidator):
+    def _validate_tool_message(self, message: ToolMessage) -> None:
+        """
+        Checks:
+        - The tool name is valid
+        """
+        if message.name is not None:
+            if not re.match(r"^[a-zA-Z0-9_-]{1,64}$", message.name):
+                raise InvalidToolMessageException(
+                    f"Function name was {message.name} but must be a-z, A-Z, 0-9, "
+                    "or contain underscores and dashes, with a maximum length of 64."
+                )
+
+        if message.tool_call_id is not None:
+            if not re.match(r"^[a-zA-Z0-9]{9}$", message.tool_call_id):
+                raise InvalidToolMessageException(
+                    f"Tool call id was {message.tool_call_id} but must be a-z, A-Z, 0-9, with a length of 9."
+                )
+
+    def _validate_tool_call(self, tool_call: ToolCall) -> None:
+        """
+        Validate that the tool call has a valid ID
+        """
+
+        if not re.match(r"^[a-zA-Z0-9]{9}$", tool_call.id):
+            raise InvalidFunctionCallException(
+                f"Tool call id was {tool_call.id} but must be a-z, A-Z, 0-9, with a length of 9."
+            )
+
+        self._validate_function_call(tool_call.function)

mistral_common/protocol/utils.py

@@ -0,0 +1,5 @@
+import uuid
+
+
+def random_uuid() -> str:
+    return str(uuid.uuid4().hex)