Add response_format support to model settings

src/mistral_common/protocol/instruct/request.py

@@ -1,9 +1,10 @@
 from enum import Enum
 from typing import Any, Generic
 
-from pydantic import Field
+from pydantic import ConfigDict, Field, field_validator
 
 from mistral_common.base import MistralBase
+from mistral_common.exceptions import InvalidRequestException
 from mistral_common.protocol.base import BaseCompletionRequest
 from mistral_common.protocol.instruct.converters import (
     convert_openai_messages,
@@ -16,6 +17,7 @@
     ReasoningFieldFormat,
 )
 from mistral_common.protocol.instruct.tool_calls import Tool, ToolChoice, ToolChoiceEnum, ToolType
+from mistral_common.utils.json_utils import validate_json_schema_by_draft7
 
 
 class ResponseFormats(str, Enum):
@@ -24,13 +26,15 @@ class ResponseFormats(str, Enum):
     Attributes:
         text: The response is a plain text.
         json: The response is a JSON object.
+        json_schema: The response follows a custom JSON schema.
 
     Examples:
         >>> response_format = ResponseFormats.text
     """
 
     text = "text"
     json = "json_object"
+    json_schema = "json_schema"
 
 
 class ReasoningEffort(str, Enum):
@@ -55,27 +59,80 @@ class ModelSettings(MistralBase):
     Attributes:
         reasoning_effort: Controls how much reasoning effort the model should apply when
             generating responses. Supported for tokenizer >= v15 and not supported for earlier versions.
+        json_schema: The JSON schema to enforce on the response, derived from the request's
+            response format. Supported for tokenizer >= v15 and not supported for earlier versions.
     """
 
     reasoning_effort: ReasoningEffort | None = None
+    json_schema: dict[str, Any] | None = None
 
     @staticmethod
     def none() -> "ModelSettings":
         r"""Create a ModelSettings instance with default (None) values."""
         return ModelSettings()
 
 
+class JsonSchema(MistralBase):
+    r"""A named JSON schema for structured responses.
+
+    Attributes:
+        name: The schema name.
+        description: An optional description of the schema.
+        custom_schema: The JSON schema (aliased ``schema``).
+        strict: Whether the model must strictly adhere to the schema.
+
+    Examples:
+        >>> schema = JsonSchema(name="obj", schema={"type": "object"})
+    """
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    name: str
+    description: str | None = None
+    custom_schema: dict[str, Any] = Field(..., alias="schema")
+    strict: bool = False
+
+    @field_validator("custom_schema")
+    @classmethod
+    def validate_custom_schema(cls, value: dict[str, Any]) -> dict[str, Any]:
+        r"""Validate the schema against JSON Schema Draft 7."""
+        validate_json_schema_by_draft7(value=value)
+        return value
+
+
 class ResponseFormat(MistralBase):
     r"""The format of the response.
 
     Attributes:
         type: The type of the response.
+        json_schema: The JSON schema when ``type`` is ``json_schema``.
 
     Examples:
         >>> response_format = ResponseFormat(type=ResponseFormats.text)
     """
 
     type: ResponseFormats = ResponseFormats.text
+    json_schema: JsonSchema | None = None
+
+    def get_schema(self) -> dict[str, Any] | None:
+        r"""Return the JSON schema to enforce for this response format.
+
+        Returns:
+            The schema dict, or None when no constraint applies.
+
+        Raises:
+            InvalidRequestException: If ``type`` is ``json_schema`` but no schema is set.
+        """
+        schema: dict[str, Any] | None
+        if self.type == ResponseFormats.json_schema:
+            if self.json_schema is None:
+                raise InvalidRequestException("Response format `json_schema` must define the schema")
+            schema = self.json_schema.custom_schema
+        elif self.type == ResponseFormats.json:
+            schema = {"anyOf": [{"type": "object"}, {"type": "array"}]}
+        else:
+            schema = None
+        return schema
 
 
 class ChatCompletionRequest(BaseCompletionRequest, Generic[ChatMessageType]):
@@ -159,7 +216,9 @@ def to_openai(
 
         # Handle messages and tools separately.
         openai_request: dict[str, Any] = self.model_dump(
-            exclude={"messages", "tools", "truncate_for_context_length", "tool_choice"}, exclude_none=True
+            exclude={"messages", "tools", "truncate_for_context_length", "tool_choice"},
+            exclude_none=True,
+            by_alias=True,
         )
 
         # Rename random_seed to seed.

src/mistral_common/tokens/tokenizers/instruct.py

@@ -1401,7 +1401,9 @@ def _encode_settings(
         self._validate_settings(settings)
         if settings == ModelSettings.none():
             return []
-        dumped_settings = json.dumps(settings.model_dump(exclude_none=True), ensure_ascii=False, sort_keys=True)
+        dumped = settings.model_dump(exclude_none=True)
+        ordered = {k: dumped[k] for k in sorted(dumped)}
+        dumped_settings = json.dumps(ordered, ensure_ascii=False)
         setting_json_tokens = self.tokenizer.encode(dumped_settings, bos=False, eos=False)
         settings_tokens = [
             self.BEGIN_MODEL_SETTINGS,

src/mistral_common/tokens/tokenizers/model_settings_builder.py

@@ -1,32 +1,40 @@
 from enum import Enum
-from typing import Any, Generic, TypeVar, final
+from typing import Any, ClassVar, Generic, Literal, TypeAlias, TypeVar, final
 
 from pydantic import model_validator
 
 from mistral_common.base import MistralBase
 from mistral_common.exceptions import InvalidRequestException
-from mistral_common.protocol.instruct.request import ChatCompletionRequest, ModelSettings, ReasoningEffort
+from mistral_common.protocol.instruct.request import (
+    ChatCompletionRequest,
+    ModelSettings,
+    ReasoningEffort,
+    ResponseFormat,
+)
+from mistral_common.utils.json_utils import validate_json_schema_by_draft7
 
 
 class ValidatorType(str, Enum):
     r"""Enumeration of validator types.
 
     Attributes:
         ENUM: Indicates that the validator is for enum values.
+        JSON_SCHEMA: Indicates that the validator is for JSON schema values.
     """
 
     ENUM = "enum"
+    JSON_SCHEMA = "json_schema"
 
 
-T = TypeVar("T")
+InputT = TypeVar("InputT")
+OutputT = TypeVar("OutputT")
+JSONSchemaDict: TypeAlias = dict[str, Any] | None
 
 
-class FieldBuilder(MistralBase, Generic[T]):
+class FieldBuilder(MistralBase, Generic[InputT, OutputT]):
     r"""Base class for field builders.
 
-    This class serves as the base for all field builders in the validation framework.
-    It ensures that all builders have a type attribute that specifies the kind of
-    validation being performed.
+    `InputT` is the request field type, `OutputT` is the converted `ModelSettings` field type.
 
     Attributes:
         type: The type of validator (e.g., ENUM).
@@ -36,7 +44,7 @@ class FieldBuilder(MistralBase, Generic[T]):
 
     type: ValidatorType
     accepts_none: bool
-    default: T | None
+    default: OutputT | None
 
     @model_validator(mode="after")
     def validate_default_accept_none(self) -> "FieldBuilder":
@@ -47,51 +55,42 @@ def validate_default_accept_none(self) -> "FieldBuilder":
             )
         return self
 
-    def _validate_built_value(self, field_name: str, value: Any) -> None:
-        r"""Validate a non-None built value. Must be implemented by subclasses."""
-        raise NotImplementedError(f"{field_name} is not supported")
+    def _convert(self, input_value: InputT) -> OutputT:
+        r"""Convert a request value to the model-settings value."""
+        raise NotImplementedError
 
-    def _build_from_optional(self, field_name: str, value: T | None) -> T | None:
-        r"""Resolve an optional value, substituting the default if value is None.
+    def _build_from_optional(self, field_name: str, input_value: InputT | None) -> OutputT | None:
+        r"""Resolve an optional value, substituting the default if input is None.
 
         Raises:
-            InvalidRequestException: If value is None and the field does not accept None.
+            InvalidRequestException: If input is None and the field does not accept None.
         """
-        if value is None:
+        if input_value is None:
             if not self.accepts_none:
                 raise InvalidRequestException(f"{field_name} should be set for this model.")
             return self.default
-        return value
+        return self._convert(input_value=input_value)
 
-    @final
-    def validate_built_value(self, field_name: str, value: Any) -> None:
-        r"""Validate a fully built value, including None checks.
-
-        Raises:
-            InvalidRequestException: If value is None when not permitted, or fails subclass validation.
-        """
-        if value is None:
-            if not (self.accepts_none and self.default is None):
-                raise InvalidRequestException(f"{field_name} should be set for this model.")
-        else:
-            self._validate_built_value(field_name, value)
+    def validate_built_value(self, field_name: str, built_value: OutputT | None) -> None:
+        r"""Validate a fully built value. Must be implemented by subclasses."""
+        raise NotImplementedError
 
     @final
-    def build_value(self, field_name: str, value: T | None) -> T | None:
+    def build_value(self, field_name: str, input_value: InputT | None) -> OutputT | None:
         r"""Resolve and validate a field value, returning the final built result.
 
         Raises:
             InvalidRequestException: If the value is invalid or missing when required.
         """
-        value = self._build_from_optional(field_name, value)
-        self.validate_built_value(field_name, value)
-        return value
+        built_value = self._build_from_optional(field_name, input_value)
+        self.validate_built_value(field_name, built_value)
+        return built_value
 
 
 E = TypeVar("E", bound=Enum)
 
 
-class EnumBuilder(FieldBuilder[E]):
+class EnumBuilder(FieldBuilder[E, E]):
     r"""Builder for enum fields.
 
     This class validates that enum fields contain only authorized values.
@@ -106,6 +105,10 @@ class EnumBuilder(FieldBuilder[E]):
     type: ValidatorType = ValidatorType.ENUM
     values: list[E]
 
+    def _convert(self, input_value: E) -> E:
+        r"""Enum builders pass values through unchanged."""
+        return input_value
+
     @model_validator(mode="after")
     def validate_unique_values(self) -> "EnumBuilder":
         r"""Ensure no duplicate values are present in the allowed values list."""
@@ -127,16 +130,42 @@ def validate_default(self) -> "EnumBuilder":
             raise ValueError(f"Default value {self.default=} is not in {self.values=}.")
         return self
 
-    def _validate_built_value(self, field_name: str, value: Any) -> None:
-        r"""Check that value is one of the allowed enum values.
+    def validate_built_value(self, field_name: str, built_value: E | None) -> None:
+        r"""Check that the built value is one of the allowed enum values.
 
         Raises:
-            InvalidRequestException: If no values are allowed, or value is not in the allowed list.
+            InvalidRequestException: If unset when required, unsupported, or not allowed.
         """
-        if len(self.values) == 0:
+        if built_value is None:
+            if not (self.accepts_none and self.default is None):
+                raise InvalidRequestException(f"{field_name} should be set for this model.")
+        elif len(self.values) == 0:
             raise InvalidRequestException(f"{field_name} not supported for this model.")
-        if value not in self.values:
-            raise InvalidRequestException(f"{field_name} should be one of {self.values}, got {value}.")
+        elif built_value not in self.values:
+            raise InvalidRequestException(f"{field_name} should be one of {self.values}, got {built_value}.")
+
+
+class JSONSchemaBuilder(FieldBuilder[ResponseFormat, JSONSchemaDict]):
+    r"""Converts a `ResponseFormat` into a JSON-schema dict for model settings.
+
+    Attributes:
+        type: The type of validator (always JSON_SCHEMA for this class).
+        accepts_none: Always False, as a response format is always present on the request.
+        default: Always None, no default schema is supported.
+    """
+
+    type: ValidatorType = ValidatorType.JSON_SCHEMA
+    accepts_none: Literal[False]
+    default: None
+
+    def _convert(self, input_value: ResponseFormat) -> JSONSchemaDict:
+        r"""Render the response format's schema for model-settings encoding."""
+        return input_value.get_schema()
+
+    def validate_built_value(self, field_name: str, built_value: JSONSchemaDict) -> None:
+        r"""Validate the built schema against Draft 7 when present."""
+        if built_value is not None:
+            validate_json_schema_by_draft7(built_value)
 
 
 class ModelSettingsBuilder(MistralBase):
@@ -151,9 +180,16 @@ class ModelSettingsBuilder(MistralBase):
 
     Attributes:
         reasoning_effort: Builder for the allowed ReasoningEffort values, or None if unsupported.
+        json_schema: Builder for the response-format JSON schema, or None if unsupported.
     """
 
+    _SETTINGS_TO_CONV_FIELDS_MAP: ClassVar[dict[str, str]] = {
+        "reasoning_effort": "reasoning_effort",
+        "json_schema": "response_format",
+    }
+
     reasoning_effort: EnumBuilder[ReasoningEffort] | None = None
+    json_schema: JSONSchemaBuilder | None = None
 
     @staticmethod
     def none() -> "ModelSettingsBuilder":
@@ -177,8 +213,8 @@ def build_settings(self, request: ChatCompletionRequest) -> ModelSettings:
         """
         dict_settings = {}
         for field_name in ModelSettingsBuilder.model_fields:
-            # We have a CI test to ensure all fields match between ModelSettings and ModelSettingsEncoder.
-            value = getattr(request, field_name)
+            # We have a CI test to ensure all fields match between ModelSettings and ModelSettingsBuilder.
+            value = getattr(request, self._SETTINGS_TO_CONV_FIELDS_MAP[field_name])
             field_builder: FieldBuilder | None = getattr(self, field_name)
             if field_builder is not None:
                 dict_settings[field_name] = field_builder.build_value(field_name, value)

src/mistral_common/utils/json_utils.py

@@ -0,0 +1,17 @@
+import jsonschema
+import jsonschema.exceptions
+
+
+def validate_json_schema_by_draft7(value: dict) -> None:
+    r"""Validate that a dict is a valid Draft 7 JSON Schema.
+
+    Args:
+        value: The candidate JSON schema.
+
+    Raises:
+        ValueError: If the value is not a valid Draft 7 JSON Schema.
+    """
+    try:
+        jsonschema.Draft7Validator.check_schema(value)
+    except jsonschema.exceptions.SchemaError as e:
+        raise ValueError(f"Invalid JSON Schema: {e.message}") from e