145 lines
5.3 KiB
Python
145 lines
5.3 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 convert_message_to_mlflow_chat(message: Union[BaseModel, dict]) -> ChatMessage:
|
|
"""
|
|
Convert Anthropic message object into MLflow's standard format (OpenAI compatible).
|
|
|
|
Ref: https://docs.anthropic.com/en/api/messages#body-messages
|
|
|
|
Args:
|
|
message: Anthropic 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")
|
|
elif isinstance(message, BaseModel):
|
|
content = message.content
|
|
role = message.role
|
|
else:
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Message must be either a dict or a Message object, but got: {type(message)}."
|
|
)
|
|
|
|
if isinstance(content, str):
|
|
return ChatMessage(role=role, content=content)
|
|
|
|
elif isinstance(content, list):
|
|
contents = []
|
|
tool_calls = []
|
|
tool_call_id = None
|
|
for content_block in content:
|
|
if isinstance(content_block, BaseModel):
|
|
if IS_PYDANTIC_V2_OR_NEWER:
|
|
content_block = content_block.model_dump()
|
|
else:
|
|
content_block = content_block.dict()
|
|
|
|
content_type = content_block.get("type")
|
|
if content_type == "tool_use":
|
|
# Anthropic response contains tool calls in the content block
|
|
# Ref: https://docs.anthropic.com/en/docs/build-with-claude/tool-use#example-api-response-with-a-tool-use-content-block
|
|
tool_calls.append(
|
|
ToolCall(
|
|
id=content_block["id"],
|
|
function=Function(
|
|
name=content_block["name"], arguments=json.dumps(content_block["input"])
|
|
),
|
|
type="function",
|
|
)
|
|
)
|
|
elif content_type == "tool_result":
|
|
# In Anthropic, the result of tool execution is returned as a special content type
|
|
# "tool_result" with "user" role, which corresponds to the "tool" role in OpenAI.
|
|
role = "tool"
|
|
tool_call_id = content_block["tool_use_id"]
|
|
if result_content := content_block.get("content"):
|
|
contents.append(_parse_content(result_content))
|
|
else:
|
|
contents.append(TextContentPart(text="", type="text"))
|
|
else:
|
|
contents.append(_parse_content(content_block))
|
|
|
|
message = ChatMessage(role=role, content=contents)
|
|
# Only set tool_calls field when it is present
|
|
if tool_calls:
|
|
message.tool_calls = tool_calls
|
|
if tool_call_id:
|
|
message.tool_call_id = tool_call_id
|
|
return message
|
|
|
|
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":
|
|
source = content["source"]
|
|
return ImageContentPart(
|
|
image_url=ImageUrl(
|
|
url=f"data:{source['media_type']};{source['type']},{source['data']}"
|
|
),
|
|
type="image_url",
|
|
)
|
|
# Claude 3.7 added new "thinking" content block, which is essentially a text block as of now.
|
|
# TODO: We should consider adding a new ContentPart type if more providers support this.
|
|
# https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking
|
|
elif content_type == "thinking":
|
|
return TextContentPart(text=content["thinking"], type="text")
|
|
else:
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Unknown content type: {content_type['type']}. Please make sure the message "
|
|
"is a valid Anthropic 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 Anthropic tool definition into MLflow's standard format (OpenAI compatible).
|
|
|
|
Ref: https://docs.anthropic.com/en/docs/build-with-claude/tool-use
|
|
|
|
Args:
|
|
tool: A dictionary represents a single tool definition in the input request.
|
|
|
|
Returns:
|
|
ChatTool: MLflow's standard tool definition object.
|
|
"""
|
|
return ChatTool(
|
|
type="function",
|
|
function=FunctionToolDefinition(
|
|
name=tool.get("name"),
|
|
description=tool.get("description"),
|
|
parameters=tool.get("input_schema"),
|
|
),
|
|
)
|