star

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

@@ -0,0 +1,40 @@
+from mlflow.anthropic.autolog import async_patched_class_call, patched_class_call
+from mlflow.utils.annotations import experimental
+from mlflow.utils.autologging_utils import autologging_integration, safe_patch
+
+FLAVOR_NAME = "anthropic"
+
+
+@experimental
+@autologging_integration(FLAVOR_NAME)
+def autolog(
+    log_traces: bool = True,
+    disable: bool = False,
+    silent: bool = False,
+):
+    """
+    Enables (or disables) and configures autologging from Anthropic to MLflow.
+    Only synchronous calls are supported. Asynchnorous APIs and streaming are not recorded.
+
+    Args:
+        log_traces: If ``True``, traces are logged for Anthropic models.
+            If ``False``, no traces are collected during inference. Default to ``True``.
+        disable: If ``True``, disables the Anthropic autologging. Default to ``False``.
+        silent: If ``True``, suppress all event logs and warnings from MLflow during Anthropic
+            autologging. If ``False``, show all events and warnings.
+    """
+    from anthropic.resources import AsyncMessages, Messages
+
+    safe_patch(
+        FLAVOR_NAME,
+        Messages,
+        "create",
+        patched_class_call,
+    )
+
+    safe_patch(
+        FLAVOR_NAME,
+        AsyncMessages,
+        "create",
+        async_patched_class_call,
+    )

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

@@ -0,0 +1,119 @@
+import logging
+from typing import Any
+
+import mlflow
+import mlflow.anthropic
+from mlflow.anthropic.chat import convert_message_to_mlflow_chat, convert_tool_to_mlflow_chat_tool
+from mlflow.entities import SpanType
+from mlflow.entities.span import LiveSpan
+from mlflow.entities.span_event import SpanEvent
+from mlflow.entities.span_status import SpanStatusCode
+from mlflow.tracing.utils import (
+    construct_full_inputs,
+    end_client_span_or_trace,
+    set_span_chat_messages,
+    set_span_chat_tools,
+    start_client_span_or_trace,
+)
+from mlflow.utils.autologging_utils.config import AutoLoggingConfig
+
+_logger = logging.getLogger(__name__)
+
+
+def patched_class_call(original, self, *args, **kwargs):
+    with TracingSession(original, self, args, kwargs) as manager:
+        output = original(self, *args, **kwargs)
+        manager.output = output
+        return output
+
+
+async def async_patched_class_call(original, self, *args, **kwargs):
+    async with TracingSession(original, self, args, kwargs) as manager:
+        output = await original(self, *args, **kwargs)
+        manager.output = output
+        return output
+
+
+class TracingSession:
+    """Context manager for handling MLflow spans in both sync and async contexts."""
+
+    def __init__(self, original, instance, args, kwargs):
+        self.mlflow_client = mlflow.MlflowClient()
+        self.original = original
+        self.instance = instance
+        self.inputs = construct_full_inputs(original, instance, *args, **kwargs)
+
+        # These attributes are set outside the constructor.
+        self.span = None
+        self.output = None
+
+    def __enter__(self):
+        return self._enter_impl()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._exit_impl(exc_type, exc_val, exc_tb)
+
+    async def __aenter__(self):
+        return self._enter_impl()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self._exit_impl(exc_type, exc_val, exc_tb)
+
+    def _enter_impl(self):
+        config = AutoLoggingConfig.init(flavor_name=mlflow.anthropic.FLAVOR_NAME)
+
+        if config.log_traces:
+            self.span = start_client_span_or_trace(
+                self.mlflow_client,
+                name=f"{self.instance.__class__.__name__}.{self.original.__name__}",
+                span_type=_get_span_type(self.original.__name__),
+                inputs=self.inputs,
+            )
+            _set_tool_attribute(self.span, self.inputs)
+
+        return self
+
+    def _exit_impl(self, exc_type, exc_val, exc_tb) -> None:
+        if self.span:
+            if exc_val:
+                self.span.add_event(SpanEvent.from_exception(exc_val))
+                status = SpanStatusCode.ERROR
+            else:
+                status = SpanStatusCode.OK
+
+            _set_chat_message_attribute(self.span, self.inputs, self.output)
+
+            end_client_span_or_trace(
+                self.mlflow_client,
+                self.span,
+                status=status,
+                outputs=self.output,
+            )
+
+
+def _get_span_type(task_name: str) -> str:
+    # Anthropic has a few APIs in beta, e.g., count_tokens.
+    # Once they are stable, we can add them to the mapping.
+    span_type_mapping = {
+        "create": SpanType.CHAT_MODEL,
+    }
+    return span_type_mapping.get(task_name, SpanType.UNKNOWN)
+
+
+def _set_tool_attribute(span: LiveSpan, inputs: dict[str, Any]):
+    if (tools := inputs.get("tools")) is not None:
+        try:
+            tools = [convert_tool_to_mlflow_chat_tool(tool) for tool in tools]
+            set_span_chat_tools(span, tools)
+        except Exception as e:
+            _logger.debug(f"Failed to set tools for {span}. Error: {e}")
+
+
+def _set_chat_message_attribute(span: LiveSpan, inputs: dict[str, Any], output: Any):
+    try:
+        messages = [convert_message_to_mlflow_chat(msg) for msg in inputs.get("messages", [])]
+        if output is not None:
+            messages.append(convert_message_to_mlflow_chat(output))
+        set_span_chat_messages(span, messages)
+    except Exception as e:
+        _logger.debug(f"Failed to set chat messages for {span}. Error: {e}")

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

@@ -0,0 +1,144 @@
+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"),
+        ),
+    )