Verfolgen von Amazon Bedrock mit MLflow

MLflow unterstützt das automatische Tracking für Amazon Bedrock, einen vollständig verwalteten Service auf AWS, der leistungsstarke Plattformen von führenden KI-Anbietern wie Anthropic, Cohere, Meta, Mistral AI und vielen mehr bietet. Wenn Sie die automatische Ablaufverfolgung für Amazon Bedrock durch die Anfrage der Funktion mlflow.bedrock.autolog aktivieren, erfasst MLflow Ablaufverfolgungen für LLM-Aufrufe und protokolliert sie im aktiven MLflow-Experiment.

Die MLflow-Ablaufverfolgung erfasst automatisch die folgenden Informationen zu Amazon Bedrock-Anrufen:

• Prompts und Abschlussantworten
• Wartezeiten
• Modellname
• Zusätzliche Metadaten wie Temperatur, max_tokens, falls angegeben.
• Funktionsaufrufe bei Rückgabe in der Antwort
• Jede ausgelöste Ausnahme

Voraussetzungen

Um die MLflow-Ablaufverfolgung mit Amazon Bedrock zu verwenden, müssen Sie MLflow und das AWS SDK für Python (Boto3) installieren.

Entwicklung

Für Entwicklungsumgebungen installieren Sie das vollständige MLflow-Paket mit Databricks-Zusatzkomponenten und boto3:

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

Das vollständige mlflow[databricks] Paket enthält alle Features für die lokale Entwicklung und Das Experimentieren mit Databricks.

Produktion

Für die Produktionsimplementierung installieren Sie mlflow-tracing und boto3.

pip install --upgrade mlflow-tracing boto3

Das mlflow-tracing-Paket ist für den Produktionseinsatz optimiert.

Bevor Sie die folgenden Beispiele ausführen, müssen Sie Ihre Umgebung konfigurieren:

Für Benutzer außerhalb von Databricks-Notizbüchern: Legen Sie Ihre Databricks-Umgebungsvariablen fest:

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

Für Benutzer innerhalb von Databricks-Notizbüchern: Diese Anmeldeinformationen werden automatisch für Sie festgelegt.

AWS-Anmeldeinformationen: Stellen Sie sicher, dass Ihre AWS-Anmeldeinformationen für den Bedrock-Zugriff konfiguriert sind. Bei der Produktionsverwendung sollten Sie die Verwendung von IAM-Rollen, AWS Secrets Manager oder Databricks-Geheimnissen anstelle von Umgebungsvariablen in Betracht ziehen (z. B. die Verwendung von AWS CLI-, IAM-Rollen oder Umgebungsvariablen).

Unterstützte APIs

MLflow unterstützt die automatische Verfolgung für die folgenden Amazon Bedrock-APIs:

• Unterhaltung
• converse_stream
• invoke_model
• Modell_mit_Response_Stream_aufrufen

Elementares Beispiel

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

Die protokollierte Ablaufverfolgung, die mit dem Test verbunden ist, kann in der Benutzeroberfläche von MLflow eingesehen werden.

Roheingaben und -ausgaben

Standardmäßig rendert MLflow die umfangreiche Chat-ähnliche Benutzeroberfläche für Eingabe- und Ausgabenachrichten auf der Chat Registerkarte. Zum Anzeigen der unformatierten Eingabe- und Ausgabenutzlast, einschließlich Konfigurationsparametern, klicken Sie auf die Inputs / Outputs Registerkarte in der Benutzeroberfläche.

Streamen

MLflow unterstützt die Ablaufverfolgung von Streaming-Aufrufen an Amazon Bedrock-APIs. Die generierte Ablaufverfolgung zeigt die aggregierte Ausgabemeldung auf der Registerkarte Chat, während die einzelnen Chunks auf der Registerkarte Events angezeigt werden.

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)

Funktionsaufruf-Agent

Die MLflow-Ablaufverfolgung erfasst automatisch Funktionsaufrufmetadaten beim Aufrufen von Amazon Bedrock-APIs. Die Funktionsdefinition und die Anweisung in der Antwort werden auf der Registerkarte Chat der Ablaufverfolgung in der Benutzeroberfläche hervorgehoben.

Wenn Sie dies mit dem feature für die manuelle Ablaufverfolgung kombinieren, können Sie einen Funktionsaufruf-Agent (ReAct) definieren und dessen Ausführung nachverfolgen. Die gesamte Agentenimplementierung mag kompliziert aussehen, aber der Teil der Ablaufverfolgung ist ziemlich einfach: (1) fügen Sie den @mlflow.trace-Dekorator zu den Funktionen für die Ablaufverfolgung hinzu und (2) aktivieren Sie die automatische Ablaufverfolgung für Amazon Bedrock mit mlflow.bedrock.autolog(). MLflow kümmert sich um die Komplexität, z. B. das Auflösen von Anrufketten und die Aufzeichnung von Ausführungsmetadaten.

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)

Wenn Sie den obigen Code ausführen, wird eine einzige Ablaufverfolgung erstellt, die alle LLM-Aufrufe und die Toolaufrufe umfasst.

Automatische Ablaufverfolgung deaktivieren

Die automatische Ablaufverfolgung für Amazon Bedrock kann durch eine Anfrage an mlflow.bedrock.autolog(disable=True) oder mlflow.autolog(disable=True) global deaktiviert werden.