Трассировка Amazon Bedrock с помощью MLflow

MLflow поддерживает автоматическое отслеживание для Amazon Bedrock, полностью управляемой службы AWS, которая предлагает высокопроизводительные платформы от ведущих поставщиков ИИ, таких как Anthropic, Cohere, Meta, Mistral AI и многих других. Включив автоматическую трассировку для Amazon Bedrock с помощью вызова функции mlflow.bedrock.autolog, MLflow будет записывать трассировки для вызова модели LLM и регистрировать их в активном эксперименте MLflow.

Трассировка MLflow автоматически записывает следующие сведения о вызовах Amazon Bedrock:

• Запросы и результаты выполнения
• Задержки
• Имя модели
• Дополнительные метаданные, такие как температура, max_tokens, если задано.
• Вызов функции, если она возвращается в рамках ответа
• Любое исключение, если оно возникнет

Предпосылки

Чтобы использовать трассировку MLflow с Amazon Bedrock, необходимо установить MLflow и пакет SDK AWS для Python (Boto3).

Развитие

Для сред разработки установите полный пакет MLflow с дополнительными компонентами Databricks и boto3:

pip install --upgrade "mlflow[databricks]>=3.1" boto3

Полный mlflow[databricks] пакет включает все функции для локальной разработки и экспериментирования в Databricks.

Производство

Для рабочих развертываний установите mlflow-tracing и boto3:

pip install --upgrade mlflow-tracing boto3

Пакет mlflow-tracing оптимизирован для использования в рабочей среде.

Прежде чем выполнять приведенные ниже примеры, необходимо настроить среду:

Для пользователей вне ноутбуков Databricks: задайте переменные среды Databricks:

export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="your-personal-access-token"

Для пользователей записных книжек Databricks: эти учетные данные устанавливаются автоматически.

Учетные данные AWS. Убедитесь, что настроены учетные данные AWS для доступа Bedrock. Для использования в производственной среде рекомендуется задействовать роли IAM, AWS Secrets Manager или секреты Databricks вместо переменных окружения (например, используя AWS CLI, роли IAM или переменные окружения).

Поддерживаемые API

MLflow поддерживает автоматическую трассировку для следующих API Amazon Bedrock:

• беседовать
• converse_stream
• invoke_model
• invoke_model_with_response_stream (вызвать_модель_с_потоком_ответов)

Базовый пример

import boto3
import mlflow
import os

# Ensure your AWS credentials are configured in your environment
# (e.g., using AWS CLI `aws configure`, or by setting
#  AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_SESSION_TOKEN, AWS_DEFAULT_REGION)

# Enable auto-tracing for Amazon Bedrock
mlflow.bedrock.autolog()

# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/bedrock-tracing-demo")

# 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
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,
    },
)

Зарегистрированную трассировку, связанную с экспериментом, можно увидеть в пользовательском интерфейсе MLflow.

Необработанные входные и выходные данные

По умолчанию MLflow отображает расширенный пользовательский интерфейс чата для входных и выходных сообщений на вкладке Chat . Чтобы просмотреть необработанные входные и выходные полезные данные, включая параметры конфигурации, щелкните Inputs / Outputs вкладку в пользовательском интерфейсе.

Стриминг

MLflow поддерживает трассировку потоковых вызовов API Amazon Bedrock. Созданная трассировка отображает агрегированное выходное сообщение на 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 автоматически фиксирует метаданные вызовов функций при обращении к API Amazon Bedrock. Определение функции и инструкция в ответе будут выделены на вкладке Chat пользовательского интерфейса трассировки.

Объединяя это с функцией трассировки вручную, можно определить агент вызова функций (ReAct) и отслеживать его выполнение. Вся реализация агента может выглядеть сложно, но трассировочная часть довольно проста: (1) добавьте декоратор @mlflow.trace к функциям, которые нужно трассировать, и (2) включите автоматическую трассировку для Amazon Bedrock, используя mlflow.bedrock.autolog(). MLflow упростит такие задачи, как разрешение цепочек вызовов и запись метаданных выполнения.

import boto3
import mlflow
from mlflow.entities import SpanType
import os

# Ensure your AWS credentials are configured in your environment

# Enable auto-tracing for Amazon Bedrock
mlflow.bedrock.autolog()

# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/bedrock-agent-demo")

# 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 и инструментам.

Отключение автоматической трассировки

Автоматическая трассировка для Amazon Bedrock может быть отключена глобально путем вызова mlflow.bedrock.autolog(disable=True) или mlflow.autolog(disable=True).