使用 MLflow 追踪 Amazon Bedrock

MLflow 支持对 Amazon Bedrock 进行自动追踪，Amazon Bedrock 是 AWS 上的一项完全托管服务，提供来自领先 AI 提供商（如 Anthropic、Cohere、Meta、Mistral AI 等）的高性能基础模型。通过调用 mlflow.bedrock.autolog() 函数启用 Amazon Bedrock 的自动追踪后，MLflow 将捕获 LLM 调用追踪并将其记录到活动 MLflow 实验中。

import mlflow

mlflow.bedrock.autolog()

MLflow 追踪会自动捕获有关 Amazon Bedrock 调用的以下信息：

• 提示词和完成响应
• 延迟
• 模型名称
• 其他元数据，如温度、max_tokens（如果指定）。
• 如果在响应中返回函数调用信息
• 如果引发任何异常

支持的 API

MLflow 支持对以下 Amazon Bedrock API 进行自动追踪：

• converse
• converse_stream
• invoke_model
• invoke_model_with_response_stream

基本示例

import boto3
import mlflow

# Enable auto-tracing for Amazon Bedrock
mlflow.bedrock.autolog()
mlflow.set_experiment("Bedrock")
# Create a boto3 client for invoking the Bedrock API
bedrock = boto3.client(
    service_name="bedrock-runtime",
    region_name="<REPLACE_WITH_YOUR_AWS_REGION>",
)
# MLflow will log a trace for Bedrock API call in the "Bedrock" experiment created above
response = bedrock.converse(
    modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
    messages=[
        {
            "role": "user",
            "content": "Describe the purpose of a 'hello world' program in one line.",
        }
    ],
    inferenceConfig={
        "maxTokens": 512,
        "temperature": 0.1,
        "topP": 0.9,
    },
)

关联到 Bedrock 实验的记录追踪可以在 MLflow UI 中查看。

原始输入和输出

默认情况下，MLflow 在 Chat 标签页中为输入和输出消息渲染丰富的聊天式 UI。要查看原始输入和输出载荷，包括配置参数，请点击 UI 中的 Inputs / Outputs 标签页。

注意

Chat 面板仅支持 converse 和 converse_stream API。对于其他 API，MLflow 仅显示 Inputs / Outputs 标签页。

流式传输

MLflow 支持追踪对 Amazon Bedrock API 的流式调用。生成的追踪在 Chat 标签页中显示聚合的输出消息，而单个块则在 Events 标签页中显示。

response = bedrock.converse_stream(
    modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
    messages=[
        {
            "role": "user",
            "content": [
                {"text": "Describe the purpose of a 'hello world' program in one line."}
            ],
        }
    ],
    inferenceConfig={
        "maxTokens": 300,
        "temperature": 0.1,
        "topP": 0.9,
    },
)

for chunk in response["stream"]:
    print(chunk)

警告

当流式响应返回时，MLflow 不会立即创建 span。相反，它会在流式块被消费时创建 span，例如上面代码片段中的 for 循环。

函数调用 Agent

调用 Amazon Bedrock API 时，MLflow Tracing 会自动捕获函数调用元数据。响应中的函数定义和指令将在追踪 UI 的 Chat 标签页中高亮显示。

结合手动追踪功能，您可以定义一个函数调用 agent (ReAct) 并追踪其执行。整个 agent 实现可能看起来很复杂，但追踪部分非常简单：(1) 将 @mlflow.trace 装饰器添加到要追踪的函数上，以及 (2) 使用 mlflow.bedrock.autolog() 为 Amazon Bedrock 启用自动追踪。MLflow 将处理诸如解析调用链和记录执行元数据等复杂工作。

import boto3
import mlflow
from mlflow.entities import SpanType

# Enable auto-tracing for Amazon Bedrock
mlflow.bedrock.autolog()
mlflow.set_experiment("Bedrock")
# Create a boto3 client for invoking the Bedrock API
bedrock = boto3.client(
    service_name="bedrock-runtime",
    region_name="<REPLACE_WITH_YOUR_AWS_REGION>",
)
model_id = "anthropic.claude-3-5-sonnet-20241022-v2:0"


# Define the tool function. Decorate it with `@mlflow.trace` to create a span for its execution.
@mlflow.trace(span_type=SpanType.TOOL)
def get_weather(city: str) -> str:
    """ "Get the current weather in a given location"""
    return "sunny" if city == "San Francisco, CA" else "unknown"


# Define the tool configuration passed to Bedrock
tools = [
    {
        "toolSpec": {
            "name": "get_weather",
            "description": "Get the current weather in a given location",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "city": {
                            "type": "string",
                            "description": "The city and state, e.g., San Francisco, CA",
                        },
                    },
                    "required": ["city"],
                }
            },
        }
    }
]
tool_functions = {"get_weather": get_weather}


# Define a simple tool calling agent
@mlflow.trace(span_type=SpanType.AGENT)
def run_tool_agent(question: str) -> str:
    messages = [{"role": "user", "content": [{"text": question}]}]
    # Invoke the model with the given question and available tools
    response = bedrock.converse(
        modelId=model_id,
        messages=messages,
        toolConfig={"tools": tools},
    )
    assistant_message = response["output"]["message"]
    messages.append(assistant_message)
    # If the model requests tool call(s), invoke the function with the specified arguments
    tool_use = next(
        (c["toolUse"] for c in assistant_message["content"] if "toolUse" in c), None
    )
    if tool_use:
        tool_func = tool_functions[tool_use["name"]]
        tool_result = tool_func(**tool_use["input"])
        messages.append(
            {
                "role": "user",
                "content": [
                    {
                        "toolResult": {
                            "toolUseId": tool_use["toolUseId"],
                            "content": [{"text": tool_result}],
                        }
                    }
                ],
            }
        )
        # Send the tool results to the model and get a new response
        response = bedrock.converse(
            modelId=model_id,
            messages=messages,
            toolConfig={"tools": tools},
        )
    return response["output"]["message"]["content"][0]["text"]


# Run the tool calling agent
question = "What's the weather like in San Francisco today?"
answer = run_tool_agent(question)

执行上面的代码将创建一个包含所有 LLM 调用和工具调用的单一追踪。

禁用自动追踪

可以通过调用 mlflow.bedrock.autolog(disable=True) 或 mlflow.autolog(disable=True) 来全局禁用 Amazon Bedrock 的自动追踪。