210 lines
7.7 KiB
Python
210 lines
7.7 KiB
Python
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}")
|