Pular para o conteúdo principal

Rastreando o Amazon Bedrock com o MLflow

MLflow oferece suporte ao rastreamento automático para Amazon Bedrock, um serviço totalmente gerenciado em AWS que fornece bases de alto desempenho em dos principais provedores de AI, como Anthropic, Cohere, Meta, Mistral AI, entre outros. Ao ativar o rastreamento automático para Amazon Bedrock chamando a função mlflow.bedrock.autolog, MLflow capturará os rastreamentos para a invocação LLM e log para a experiência MLflow ativa.

Rastreamento de agentes Bedrock DIY

MLflow O rastreamento captura automaticamente as seguintes informações sobre as chamadas do site Amazon Bedrock:

  • Solicitações e respostas de conclusão
  • Latências
  • Nome do modelo
  • Metadados adicionais, como temperatura, max_tokens, se especificados.
  • Chamada de função se retornada na resposta
  • Qualquer exceção, se levantada

Pré-requisitos​

Para usar o MLflow Tracing com o Amazon Bedrock, o senhor precisa instalar o MLflow e o AWS SDK for Python (Boto3).

Para ambientes de desenvolvimento, instale o pacote completo do MLflow com os extras do Databricks e boto3:

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

O pacote completo do mlflow[databricks] inclui todos os recursos para desenvolvimento local e experimentação no Databricks.

nota

O MLflow 3 é altamente recomendado para obter a melhor experiência de rastreamento com o Amazon Bedrock.

Antes de executar os exemplos abaixo, o senhor precisará configurar seu ambiente:

Para usuários fora do Databricks Notebook : Defina seu Databricks variável de ambiente:

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

Para usuários do Databricks Notebook : Essas credenciais são definidas automaticamente para o senhor.

Credenciais do AWS : Certifique-se de que suas credenciais do AWS para acesso ao Bedrock estejam configuradas. Para uso em produção, considere usar IAM role, AWS Secrets Manager ou Databricks secrets em vez de variável de ambiente (por exemplo, via AWS CLI, IAM role ou variável de ambiente).

APIs suportadas​

O MLflow oferece suporte ao rastreamento automático para as seguintes APIs do Amazon Bedrock:

Exemplo básico​

Python
import boto3
import mlflow
import os

# Ensure your AWS credentials are configured in your environment
# (e.g., via 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,
    },
)

O rastreamento dos registros, associado ao experimento, pode ser visto na UI MLflow.

Entradas e saídas brutas​

Em default, MLflow renderiza a interface de usuário avançada do tipo chat para mensagens de entrada e saída no Chat tab. Para acessar view a carga útil bruta de entrada e saída, inclusive os parâmetros de configuração, clique em Inputs / Outputs tab na interface do usuário.

nota

O painel Chat é compatível apenas com as APIs converse e converse_stream. Para os outros APIs, MLflow exibe apenas o Inputs / Outputs tab.

transmissão​

MLflow suporta o rastreamento de chamadas de transmissão para Amazon Bedrock APIs. O rastreamento gerado mostra a mensagem de saída agregada no site Chat tab, enquanto os blocos individuais são exibidos no site Events tab.

Python
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)

Rastreamento de transmissão de Bedrock

atenção

MLflow não cria um intervalo imediatamente quando a resposta de transmissão é retornada. Em vez disso, ele cria um intervalo quando os pedaços de transmissão são consumidos , por exemplo, o for-loop no trecho de código acima.

Agente de chamada de funções​

O MLflow Tracing captura automaticamente os metadados de chamada de função ao chamar as APIs do Amazon Bedrock. A definição da função e a instrução na resposta serão destacadas no site Chat tab na UI de rastreamento.

Combinando isso com o recurso de rastreamento manual, o senhor pode definir um agente de chamada de função (ReAct) e rastrear sua execução. Toda a implementação do agente pode parecer complicada, mas a parte do rastreamento é bem simples: (1) adicione o decorador @mlflow.trace às funções a serem rastreadas e (2) ative o rastreamento automático para o Amazon Bedrock com mlflow.bedrock.autolog(). O MLflow cuidará da complexidade, como a resolução de cadeias de chamadas e o registro de metadados de execução.

Python
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)

A execução do código acima criará um único rastreamento que envolve todas as invocações do LLM e as chamadas de ferramentas.

Rastreamento de agentes Bedrock DIY

atenção

Para ambientes de produção, sempre use AI Gateway ou Databricks secrets em vez de valores codificados para um gerenciamento seguro API key .

Desativar o rastreamento automático​

O rastreamento automático para o Amazon Bedrock pode ser desativado globalmente ligando para mlflow.bedrock.autolog(disable=True) ou mlflow.autolog(disable=True).

Última atualização em