star

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

@@ -0,0 +1,45 @@
+import logging
+
+from mlflow.utils.annotations import experimental
+from mlflow.utils.autologging_utils import autologging_integration, safe_patch
+
+_logger = logging.getLogger(__name__)
+
+FLAVOR_NAME = "bedrock"
+
+
+@experimental
+@autologging_integration(FLAVOR_NAME)
+def autolog(
+    log_traces: bool = True,
+    disable: bool = False,
+    silent: bool = False,
+):
+    """
+    Enables (or disables) and configures autologging from Amazon Bedrock to MLflow.
+    Only synchronous calls are supported. Asynchnorous APIs and streaming are not recorded.
+
+    Args:
+        log_traces: If ``True``, traces are logged for Bedrock models.
+            If ``False``, no traces are collected during inference. Default to ``True``.
+        disable: If ``True``, disables the Bedrock autologging. Default to ``False``.
+        silent: If ``True``, suppress all event logs and warnings from MLflow during Bedrock
+            autologging. If ``False``, show all events and warnings.
+    """
+    from botocore.client import ClientCreator
+
+    from mlflow.bedrock._autolog import patched_create_client
+
+    # NB: In boto3, the client class for each service is dynamically created at
+    # runtime via the ClientCreator factory class. Therefore, we cannot patch
+    # the service client directly, and instead patch the factory to return
+    # a patched client class.
+    safe_patch(FLAVOR_NAME, ClientCreator, "create_client", patched_create_client)
+
+    # Since we patch the ClientCreator factory, it only takes effect for new client instances.
+    if log_traces:
+        _logger.info(
+            "Enabled auto-tracing for Bedrock. Note that MLflow can only trace boto3 "
+            "service clients that are created after this call. If you have already "
+            "created one, please recreate the client by calling `boto3.client`."
+        )

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

@@ -0,0 +1,209 @@
+import io
+import json
+import logging
+from typing import Any, Optional, Union
+
+from botocore.client import BaseClient
+from botocore.response import StreamingBody
+
+import mlflow
+from mlflow.bedrock import FLAVOR_NAME
+from mlflow.bedrock.chat import convert_message_to_mlflow_chat, convert_tool_to_mlflow_chat_tool
+from mlflow.bedrock.stream import ConverseStreamWrapper, InvokeModelStreamWrapper
+from mlflow.bedrock.utils import skip_if_trace_disabled
+from mlflow.entities import SpanType
+from mlflow.tracing.utils import (
+    set_span_chat_messages,
+    set_span_chat_tools,
+    start_client_span_or_trace,
+)
+from mlflow.utils.autologging_utils import safe_patch
+
+_BEDROCK_RUNTIME_SERVICE_NAME = "bedrock-runtime"
+_BEDROCK_SPAN_PREFIX = "BedrockRuntime."
+
+_logger = logging.getLogger(__name__)
+
+
+def patched_create_client(original, self, *args, **kwargs):
+    """
+    Patched version of the boto3 ClientCreator.create_client method that returns
+    a patched client class.
+    """
+    if kwargs.get("service_name") != _BEDROCK_RUNTIME_SERVICE_NAME:
+        return original(self, *args, **kwargs)
+
+    client = original(self, *args, **kwargs)
+    patch_bedrock_runtime_client(client.__class__)
+
+    return client
+
+
+def patch_bedrock_runtime_client(client_class: type[BaseClient]):
+    """
+    Patch the BedrockRuntime client to log traces and models.
+    """
+    # The most basic model invocation API
+    safe_patch(FLAVOR_NAME, client_class, "invoke_model", _patched_invoke_model)
+    safe_patch(
+        FLAVOR_NAME,
+        client_class,
+        "invoke_model_with_response_stream",
+        _patched_invoke_model_with_response_stream,
+    )
+
+    if hasattr(client_class, "converse"):
+        # The new "converse" API was introduced in boto3 1.35 to access all models
+        # with the consistent chat format.
+        # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/converse.html
+        safe_patch(FLAVOR_NAME, client_class, "converse", _patched_converse)
+
+    if hasattr(client_class, "converse_stream"):
+        safe_patch(FLAVOR_NAME, client_class, "converse_stream", _patched_converse_stream)
+
+
+@skip_if_trace_disabled
+def _patched_invoke_model(original, self, *args, **kwargs):
+    with mlflow.start_span(name=f"{_BEDROCK_SPAN_PREFIX}{original.__name__}") as span:
+        # NB: Bedrock client doesn't accept any positional arguments
+        span.set_inputs(kwargs)
+
+        result = original(self, *args, **kwargs)
+
+        result["body"] = _buffer_stream(result["body"])
+        parsed_response_body = _parse_invoke_model_response_body(result["body"])
+
+        # Determine the span type based on the key in the response body.
+        # As of 2024 Dec 9th, all supported embedding models in Bedrock returns the response body
+        # with the key "embedding". This might change in the future.
+        span_type = SpanType.EMBEDDING if "embedding" in parsed_response_body else SpanType.LLM
+        span.set_span_type(span_type)
+        span.set_outputs({**result, "body": parsed_response_body})
+
+        return result
+
+
+@skip_if_trace_disabled
+def _patched_invoke_model_with_response_stream(original, self, *args, **kwargs):
+    client = mlflow.MlflowClient()
+
+    span = start_client_span_or_trace(
+        client=client,
+        name=f"{_BEDROCK_SPAN_PREFIX}{original.__name__}",
+        # NB: Since we don't inspect the response body for this method, the span type is unknown.
+        # We assume it is LLM as using streaming for embedding is not common.
+        span_type=SpanType.LLM,
+        inputs=kwargs,
+    )
+
+    result = original(self, *args, **kwargs)
+
+    # To avoid consuming the stream during serialization, set dummy outputs for the span.
+    span.set_outputs({**result, "body": "EventStream"})
+
+    result["body"] = InvokeModelStreamWrapper(stream=result["body"], client=client, span=span)
+    return result
+
+
+def _buffer_stream(raw_stream: StreamingBody) -> StreamingBody:
+    """
+    Create a buffered stream from the raw byte stream.
+
+    The boto3's invoke_model() API returns the LLM response as a byte stream.
+    We need to read the stream data to set the span outputs, however, the stream
+    can only be read once and not seekable (https://github.com/boto/boto3/issues/564).
+    To work around this, we create a buffered stream that can be read multiple times.
+    """
+    buffered_response = io.BytesIO(raw_stream.read())
+    buffered_response.seek(0)
+    return StreamingBody(buffered_response, raw_stream._content_length)
+
+
+def _parse_invoke_model_response_body(response_body: StreamingBody) -> Union[dict[str, Any], str]:
+    content = response_body.read()
+    try:
+        return json.loads(content)
+    except Exception:
+        # When failed to parse the response body as JSON, return the raw response
+        return content
+    finally:
+        # Reset the stream position to the beginning
+        response_body._raw_stream.seek(0)
+        # Boto3 uses this attribute to validate the amount of data read from the stream matches
+        # the content length, so we need to reset it as well.
+        # https://github.com/boto/botocore/blob/f88e981cb1a6cd0c64bc89da262ab76f9bfa9b7d/botocore/response.py#L164C17-L164C32
+        response_body._amount_read = 0
+
+
+@skip_if_trace_disabled
+def _patched_converse(original, self, *args, **kwargs):
+    with mlflow.start_span(
+        name=f"{_BEDROCK_SPAN_PREFIX}{original.__name__}",
+        span_type=SpanType.CHAT_MODEL,
+    ) as span:
+        # NB: Bedrock client doesn't accept any positional arguments
+        span.set_inputs(kwargs)
+        _set_tool_attributes(span, kwargs)
+
+        result = None
+        try:
+            result = original(self, *args, **kwargs)
+            span.set_outputs(result)
+        finally:
+            _set_chat_messages_attributes(span, kwargs.get("messages", []), result)
+        return result
+
+
+@skip_if_trace_disabled
+def _patched_converse_stream(original, self, *args, **kwargs):
+    # NB: Do not use fluent API to create a span for streaming response. If we do so,
+    # the span context will remain active until the stream is fully exhausted, which
+    # can lead to super hard-to-debug issues.
+    client = mlflow.MlflowClient()
+    span = start_client_span_or_trace(
+        client=client,
+        name=f"{_BEDROCK_SPAN_PREFIX}{original.__name__}",
+        span_type=SpanType.CHAT_MODEL,
+        inputs=kwargs,
+    )
+    _set_tool_attributes(span, kwargs)
+
+    result = original(self, *args, **kwargs)
+
+    if span:
+        result["stream"] = ConverseStreamWrapper(
+            stream=result["stream"],
+            span=span,
+            client=client,
+            inputs=kwargs,
+        )
+
+    return result
+
+
+def _set_chat_messages_attributes(span, messages: list[dict], response: Optional[dict]):
+    """
+    Extract standard chat span attributes for the Bedrock Converse API call.
+
+    NB: We only support standard attribute extraction for the Converse API, because
+    the InvokeModel API exposes the raw API spec from each LLM provider, hence
+    maintaining the compatibility for all providers is significantly cumbersome.
+    """
+    try:
+        messages = [*messages]  # shallow copy to avoid appending to the original list
+        if response:
+            messages.append(response["output"]["message"])
+        messages = [convert_message_to_mlflow_chat(msg) for msg in messages]
+        set_span_chat_messages(span, messages)
+    except Exception as e:
+        _logger.debug(f"Failed to set messages for {span}. Error: {e}")
+
+
+def _set_tool_attributes(span, kwargs):
+    """Extract tool attributes for the Bedrock Converse API call."""
+    if tool_config := kwargs.get("toolConfig"):
+        try:
+            tools = [convert_tool_to_mlflow_chat_tool(tool) for tool in tool_config["tools"]]
+            set_span_chat_tools(span, tools)
+        except Exception as e:
+            _logger.debug(f"Failed to set tools for {span}. Error: {e}")

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"),
+        ),
+    )

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

@@ -0,0 +1,167 @@
+import json
+import logging
+from typing import Any, Optional
+
+from botocore.eventstream import EventStream
+
+from mlflow.bedrock.chat import convert_message_to_mlflow_chat
+from mlflow.bedrock.utils import capture_exception
+from mlflow.entities.span import LiveSpan
+from mlflow.entities.span_event import SpanEvent
+from mlflow.tracing.utils import set_span_chat_messages
+from mlflow.tracking.client import MlflowClient
+
+_logger = logging.getLogger(__name__)
+
+
+class BaseEventStreamWrapper:
+    """
+    A wrapper class for a event stream to record events and accumulated response
+    in an MLflow span if possible.
+
+    A span should be ended when the stream is exhausted rather than when it is created.
+
+    Args:
+        stream: The original event stream to wrap.
+        client: The MLflow client to end the span.
+        span: The span to record events and response in.
+        inputs: The inputs to the converse API.
+    """
+
+    def __init__(
+        self,
+        stream: EventStream,
+        client: MlflowClient,
+        span: LiveSpan,
+        inputs: Optional[dict[str, Any]] = None,
+    ):
+        self._stream = stream
+        self._span = span
+        self._client = client
+        self._inputs = inputs
+
+    def __iter__(self):
+        for event in self._stream:
+            self._handle_event(self._span, event)
+            yield event
+
+        # End the span when the stream is exhausted
+        self._close()
+
+    def __getattr__(self, attr):
+        """Delegate all other attributes to the original stream."""
+        return getattr(self._stream, attr)
+
+    def _handle_event(self, span, event):
+        """Process a single event from the stream."""
+        raise NotImplementedError
+
+    def _close(self):
+        """End the span and run any finalization logic."""
+        raise NotImplementedError
+
+    @capture_exception("Failed to handle event for the stream")
+    def _end_span(self):
+        """End the span."""
+        if self._span.parent_id:
+            self._client.end_span(self._span.request_id, self._span.span_id)
+        else:
+            self._client.end_trace(self._span.request_id)
+
+
+class InvokeModelStreamWrapper(BaseEventStreamWrapper):
+    """A wrapper class for a event stream returned by the InvokeModelWithResponseStream API."""
+
+    @capture_exception("Failed to handle event for the stream")
+    def _handle_event(self, span, event):
+        chunk = json.loads(event["chunk"]["bytes"])
+        self._span.add_event(SpanEvent(name=chunk["type"], attributes={"json": json.dumps(chunk)}))
+
+    def _close(self):
+        self._end_span()
+
+
+class ConverseStreamWrapper(BaseEventStreamWrapper):
+    """A wrapper class for a event stream returned by the ConverseStream API."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._response_builder = _ConverseMessageBuilder()
+
+    def __getattr__(self, attr):
+        """Delegate all other attributes to the original stream."""
+        return getattr(self._stream, attr)
+
+    @capture_exception("Failed to handle event for the stream")
+    def _handle_event(self, span, event):
+        """
+        Process a single event from the stream.
+
+        Refer to the following documentation for the event format:
+        https://boto3.amazonaws.com/v1/documentation/api/1.35.8/reference/services/bedrock-runtime/client/converse_stream.html
+        """
+        event_name = list(event.keys())[0]
+        self._response_builder.process_event(event_name, event[event_name])
+        # Record raw event as a span event
+        self._span.add_event(
+            SpanEvent(name=event_name, attributes={"json": json.dumps(event[event_name])})
+        )
+
+    @capture_exception("Failed to record the accumulated response in the span")
+    def _close(self):
+        # Record the accumulated response as the output of the span
+        converse_response = self._response_builder.build()
+        self._span.set_outputs(converse_response)
+
+        # Record the chat message attributes in the MLflow's standard format
+        messages = self._inputs.get("messages", []) + [converse_response["output"]["message"]]
+        mlflow_messages = [convert_message_to_mlflow_chat(m) for m in messages]
+        set_span_chat_messages(self._span, mlflow_messages)
+
+        self._end_span()
+
+
+class _ConverseMessageBuilder:
+    """A helper class to accumulate the chunks of a streaming Converse API response."""
+
+    def __init__(self):
+        self._role = "assistant"
+        self._text_content_buffer = ""
+        self._tool_use = {}
+        self._response = {}
+
+    def process_event(self, event_name: str, event_attr: dict):
+        if event_name == "messageStart":
+            self._role = event_attr["role"]
+        elif event_name == "contentBlockStart":
+            # ContentBlockStart event is only used for tool usage. It carries the tool id
+            # and the name, but not the input arguments.
+            self._tool_use = {
+                # In streaming, input is always string
+                "input": "",
+                **event_attr["start"]["toolUse"],
+            }
+        elif event_name == "contentBlockDelta":
+            delta = event_attr["delta"]
+            if text := delta.get("text"):
+                self._text_content_buffer += text
+            if tool_use := delta.get("toolUse"):
+                self._tool_use["input"] += tool_use["input"]
+        elif event_name == "contentBlockStop":
+            pass
+        elif event_name == "messageStop" or event_name == "metadata":
+            self._response.update(event_attr)
+        else:
+            _logger.debug(f"Unknown event, skipping: {event_name}")
+
+    def build(self) -> dict[str, Any]:
+        message = {
+            "role": self._role,
+            "content": [{"text": self._text_content_buffer}],
+        }
+        if self._tool_use:
+            message["content"].append({"toolUse": self._tool_use})
+
+        self._response.update({"output": {"message": message}})
+
+        return self._response

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

@@ -0,0 +1,43 @@
+import logging
+from typing import Any, Callable
+
+from mlflow.bedrock import FLAVOR_NAME
+from mlflow.environment_variables import _MLFLOW_TESTING
+from mlflow.utils.autologging_utils.config import AutoLoggingConfig
+
+_logger = logging.getLogger(__name__)
+
+
+def capture_exception(logging_message: str):
+    """
+    A decorator to capture exceptions during a function execution.
+    """
+
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except Exception:
+                _logger.debug(logging_message)
+                if _MLFLOW_TESTING:
+                    raise
+
+        return wrapper
+
+    return decorator
+
+
+def skip_if_trace_disabled(func: Callable[..., Any]) -> Callable[..., Any]:
+    """
+    A decorator to apply the function only if trace autologging is enabled.
+    This decorator is used to skip the test if the trace autologging is disabled.
+    """
+
+    def wrapper(original, self, *args, **kwargs):
+        config = AutoLoggingConfig.init(flavor_name=FLAVOR_NAME)
+        if not config.log_traces:
+            return original(self, *args, **kwargs)
+
+        return func(original, self, *args, **kwargs)
+
+    return wrapper