star

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

@@ -0,0 +1,122 @@
+import base64
+import json
+import logging
+from typing import Optional, Union
+
+from mlflow.types.chat import (
+    ChatMessage,
+    ChatTool,
+    Function,
+    FunctionToolDefinition,
+    ImageContentPart,
+    ImageUrl,
+    TextContentPart,
+    ToolCall,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+def convert_message_to_mlflow_chat(message: dict) -> ChatMessage:
+    """
+    Convert Bedrock Converse API's message object into MLflow's standard format (OpenAI compatible).
+
+    Ref: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Message.html
+
+    Args:
+        message: Bedrock Converse API's message object.
+
+    Returns:
+        ChatMessage: MLflow's standard chat message object.
+    """
+    role = message["role"]
+    contents = []
+    tool_calls = []
+    tool_call_id = None
+    for content in message["content"]:
+        if tool_call := content.get("toolUse"):
+            input = tool_call.get("input")
+            tool_calls.append(
+                ToolCall(
+                    id=tool_call["toolUseId"],
+                    function=Function(
+                        name=tool_call["name"],
+                        arguments=input if isinstance(input, str) else json.dumps(input),
+                    ),
+                    type="function",
+                )
+            )
+        elif tool_result := content.get("toolResult"):
+            tool_call_id = tool_result["toolUseId"]
+            # "tool_result" content corresponds to the "tool" message in OpenAI.
+            # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolResultContentBlock.html
+            role = "tool"
+            for content in tool_result["content"]:
+                parsed_content = _parse_content(content)
+                if parsed_content:
+                    contents.append(parsed_content)
+        else:
+            parsed_content = _parse_content(content)
+            if parsed_content:
+                contents.append(parsed_content)
+
+    message = ChatMessage(role=role, content=contents)
+    if tool_calls:
+        message.tool_calls = tool_calls
+    if tool_call_id:
+        message.tool_call_id = tool_call_id
+    return message
+
+
+def _parse_content(content: dict) -> Optional[Union[TextContentPart, ImageContentPart]]:
+    """
+    Parse a single content block in the Bedrock message object.
+
+    Some content types like video and document are not supported by OpenAI's spec. This
+    function returns None for those content types.
+
+    Ref: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ContentBlock.html
+    """
+    if text := content.get("text"):
+        return TextContentPart(text=text, type="text")
+    elif json_content := content.get("json"):
+        return TextContentPart(text=json.dumps(json_content), type="text")
+    elif image := content.get("image"):
+        # Bedrock support passing images in both raw bytes and base64 encoded strings.
+        # OpenAI spec only supports base64 encoded images, so we encode the raw bytes to base64.
+        # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ImageBlock.html
+        image_bytes = image["source"]["bytes"]
+        if isinstance(image_bytes, bytes):
+            data = base64.b64encode(image_bytes).decode("utf-8")
+        else:
+            data = image_bytes
+        format = "image/" + image["format"]
+        image_url = ImageUrl(url=f"data:{format};base64,{data}", detail="auto")
+        return ImageContentPart(type="image_url", image_url=image_url)
+    # NB: Video and Document content type are not supported by OpenAI's spec, so recording as text.
+    else:
+        _logger.debug(f"Received an unsupported content type: {list(content.keys())[0]}")
+        return None
+
+
+def convert_tool_to_mlflow_chat_tool(tool: dict) -> ChatTool:
+    """
+    Convert Bedrock tool definition into MLflow's standard format (OpenAI compatible).
+
+    Ref: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Tool.html
+
+    Args:
+        tool: A dictionary represents a single tool definition in the input request.
+
+    Returns:
+        ChatTool: MLflow's standard tool definition object.
+    """
+    tool_spec = tool["toolSpec"]
+    return ChatTool(
+        type="function",
+        function=FunctionToolDefinition(
+            name=tool_spec["name"],
+            description=tool_spec.get("description"),
+            parameters=tool_spec["inputSchema"].get("json"),
+        ),
+    )