star

venv/lib/python3.9/site-packages/mlflow/mistral/__init__.py

@@ -0,0 +1,34 @@
+from mlflow.mistral.autolog import patched_class_call
+from mlflow.utils.annotations import experimental
+from mlflow.utils.autologging_utils import autologging_integration, safe_patch
+
+FLAVOR_NAME = "mistral"
+
+
+@experimental
+@autologging_integration(FLAVOR_NAME)
+def autolog(
+    log_traces: bool = True,
+    disable: bool = False,
+    silent: bool = False,
+):
+    """
+    Enables (or disables) and configures autologging from Mistral AI to MLflow.
+    Only synchronous calls to the Text generation API are supported.
+    Asynchronous APIs and streaming are not recorded.
+
+    Args:
+        log_traces: If ``True``, traces are logged for Mistral AI models.
+            If ``False``, no traces are collected during inference. Default to ``True``.
+        disable: If ``True``, disables the Mistral AI autologging. Default to ``False``.
+        silent: If ``True``, suppress all event logs and warnings from MLflow during Mistral AI
+            autologging. If ``False``, show all events and warnings.
+    """
+    from mistralai.chat import Chat
+
+    safe_patch(
+        FLAVOR_NAME,
+        Chat,
+        "complete",
+        patched_class_call,
+    )

venv/lib/python3.9/site-packages/mlflow/mistral/autolog.py

@@ -0,0 +1,61 @@
+import inspect
+import logging
+
+import mlflow
+import mlflow.mistral
+from mlflow.entities import SpanType
+from mlflow.mistral.chat import convert_message_to_mlflow_chat, convert_tool_to_mlflow_chat_tool
+from mlflow.tracing.utils import set_span_chat_messages, set_span_chat_tools
+from mlflow.utils.autologging_utils.config import AutoLoggingConfig
+
+_logger = logging.getLogger(__name__)
+
+
+def _construct_full_inputs(func, *args, **kwargs):
+    signature = inspect.signature(func)
+    # this does not create copy. So values should not be mutated directly
+    arguments = signature.bind_partial(*args, **kwargs).arguments
+
+    if "self" in arguments:
+        arguments.pop("self")
+
+    return arguments
+
+
+def patched_class_call(original, self, *args, **kwargs):
+    config = AutoLoggingConfig.init(flavor_name=mlflow.mistral.FLAVOR_NAME)
+
+    if config.log_traces:
+        with mlflow.start_span(
+            name=f"{self.__class__.__name__}.{original.__name__}",
+            span_type=SpanType.CHAT_MODEL,
+        ) as span:
+            inputs = _construct_full_inputs(original, self, *args, **kwargs)
+            span.set_inputs(inputs)
+
+            if (tools := inputs.get("tools")) is not None:
+                try:
+                    tools = [convert_tool_to_mlflow_chat_tool(tool) for tool in tools if tool]
+                    set_span_chat_tools(span, tools)
+                except Exception as e:
+                    _logger.debug(f"Failed to set tools for {span}. Error: {e}")
+
+            try:
+                messages = [convert_message_to_mlflow_chat(m) for m in inputs.get("messages", [])]
+            except Exception as e:
+                _logger.debug(f"Failed to convert chat messages for {span}. Error: {e}")
+
+            try:
+                outputs = original(self, *args, **kwargs)
+                span.set_outputs(outputs)
+            finally:
+                # Set message attribute once at the end to avoid multiple JSON serialization
+                try:
+                    for choice in getattr(outputs, "choices", []):
+                        choice_message = getattr(choice, "message", {})
+                        messages.append(convert_message_to_mlflow_chat(choice_message))
+                    set_span_chat_messages(span, messages)
+                except Exception as e:
+                    _logger.debug(f"Failed to set chat messages for {span}. Error: {e}")
+
+            return outputs

venv/lib/python3.9/site-packages/mlflow/mistral/chat.py

@@ -0,0 +1,135 @@
+import json
+from typing import Union
+
+from pydantic import BaseModel
+
+from mlflow.exceptions import MlflowException
+from mlflow.types.chat import (
+    ChatMessage,
+    ChatTool,
+    Function,
+    FunctionToolDefinition,
+    ImageContentPart,
+    ImageUrl,
+    TextContentPart,
+    ToolCall,
+)
+from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
+
+
+def _to_dict(obj: BaseModel):
+    if IS_PYDANTIC_V2_OR_NEWER:
+        return obj.model_dump()
+    return obj.dict()
+
+
+def convert_message_to_mlflow_chat(message: Union[BaseModel, dict]) -> ChatMessage:
+    """
+    Convert Mistral AI message object into MLflow's standard format (OpenAI compatible).
+
+    Ref: https://docs.mistral.ai/api/#tag/chat/operation/chat_completion_v1_chat_completions_post
+
+    Args:
+        message: Mistral AI message object or a dictionary representing the message.
+
+    Returns:
+        ChatMessage: MLflow's standard chat message object.
+    """
+    if isinstance(message, dict):
+        content = message.get("content")
+        role = message.get("role")
+        tool_calls = message.get("tool_calls")
+        tool_call_id = message.get("tool_call_id")
+    elif isinstance(message, BaseModel):
+        content = message.content
+        role = message.role
+        # tool_calls is available if message is an AssistantMessage object
+        tool_calls = getattr(message, "tool_calls", None)
+        if tool_calls:
+            tool_calls = [_to_dict(tool_call) for tool_call in tool_calls]
+        # tool_call_id is available if message is a ToolMessage object
+        tool_call_id = getattr(message, "tool_call_id", None)
+    else:
+        raise MlflowException.invalid_parameter_value(
+            f"Message must be either a dict or a Message object, but got: {type(message)}."
+        )
+
+    if tool_calls:
+        tool_calls = [
+            ToolCall(
+                id=tool_call["id"],
+                function=Function(
+                    name=tool_call["function"]["name"],
+                    arguments=json.dumps(tool_call["function"]["arguments"]),
+                ),
+                type="function",
+            )
+            for tool_call in tool_calls
+        ]
+
+    if isinstance(content, str):
+        return ChatMessage(
+            role=role, content=content, tool_calls=tool_calls, tool_call_id=tool_call_id
+        )
+
+    elif isinstance(content, list):
+        contents = []
+        tool_calls = []
+        tool_call_id = None
+        for content_chunk in content:
+            if isinstance(content_chunk, BaseModel):
+                content_chunk = _to_dict(content_chunk)
+            contents.append(_parse_content(content_chunk))
+
+        return ChatMessage(
+            role=role, content=contents, tool_calls=tool_calls, tool_call_id=tool_call_id
+        )
+
+    else:
+        raise MlflowException.invalid_parameter_value(
+            f"Invalid content type. Must be either a string or a list, but got: {type(content)}."
+        )
+
+
+def _parse_content(content: Union[str, dict]) -> Union[TextContentPart, ImageContentPart]:
+    if isinstance(content, str):
+        return TextContentPart(text=content, type="text")
+
+    content_type = content.get("type")
+    if content_type == "text":
+        return TextContentPart(text=content["text"], type="text")
+    elif content_type == "image_url":
+        return ImageContentPart(
+            image_url=ImageUrl(url=content["image_url"], detail="auto"),
+            type="image_url",
+        )
+    else:
+        raise MlflowException.invalid_parameter_value(
+            f"Unknown content type: {content_type['type']}. Please make sure the message "
+            "is a valid Mistral AI message object. If it is a valid type, contact to the "
+            "MLflow maintainer via https://github.com/mlflow/mlflow/issues/new/choose for "
+            "requesting support for a new message type."
+        )
+
+
+def convert_tool_to_mlflow_chat_tool(tool: dict) -> ChatTool:
+    """
+    Convert Mistral AI tool definition into MLflow's standard format (OpenAI compatible).
+
+    Ref: https://docs.mistral.ai/capabilities/function_calling/#tools
+
+    Args:
+        tool: A dictionary represents a single tool definition in the input request.
+
+    Returns:
+        ChatTool: MLflow's standard tool definition object.
+    """
+    function = tool["function"]
+    return ChatTool(
+        type="function",
+        function=FunctionToolDefinition(
+            name=function["name"],
+            description=function.get("description"),
+            parameters=function["parameters"],
+        ),
+    )