136 lines
4.4 KiB
Python
136 lines
4.4 KiB
Python
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"],
|
|
),
|
|
)
|