Observabilidade do agente com MLflow Tracing

Este artigo descreve como adicionar observabilidade aos seus aplicativos de IA generativa com o MLflow Tracing no Databricks.

O que é MLflow Tracing?

MLflow Tracing fornece observabilidade de ponta a ponta para aplicativos de IA generativa desde o desenvolvimento até a implantação. O rastreamento é totalmente integrado com o conjunto de ferramentas de IA de geração da Databricks, capturando insights detalhados em todo o ciclo de vida de desenvolvimento e produção.

A seguir estão os principais casos de uso para rastreamento em aplicativos de IA de geração:

• Depuração simplificada: O rastreamento fornece visibilidade em cada etapa da sua aplicação de IA generativa, tornando mais fácil o diagnóstico e a resolução de problemas.

• Avaliação offline: O rastreamento gera dados valiosos para a avaliação dos agentes, permitindo medir e melhorar a qualidade dos agentes ao longo do tempo.

• Monitorização da Produção: O rastreamento fornece visibilidade sobre o comportamento de agentes e etapas detalhadas de execução, permitindo monitorizar e otimizar o desempenho dos agentes em produção.

• Logs de auditoria: MLflow Tracing gera logs de auditoria abrangentes de ações e decisões do agente. Isso é vital para garantir a conformidade e dar suporte à depuração quando surgem problemas inesperados.

Requerimentos

O MLflow Tracing está disponível nas versões 2.13.0 e superiores do MLflow. O Databricks recomenda a instalação da versão mais recente do MLflow para acessar os recursos e melhorias mais recentes.

%pip install mlflow>=2.13.0 -qqqU
%restart_python

Rastreio automático

de registro automático do MLflow permite que você instrumente rapidamente seu agente adicionando uma única linha ao seu código, mlflow.<library>.autolog().

O MLflow suporta registro automático para as bibliotecas de criação de agentes mais populares. Para obter mais informações sobre cada biblioteca de criação, consulte documentação de registro automático do MLflow:

Desativar o registo automático

O rastreamento automático de logs é ativado por padrão no Databricks Runtime 15.4 ML e posteriores para as seguintes bibliotecas.

• LangChain
• Langgraph
• OpenAI
• LlamaIndex

Para desativar o rastreamento de registro automático para essas bibliotecas, execute o seguinte comando em um bloco de anotações:

`mlflow.<library>.autolog(log_traces=False)`

Adicionar rastreamentos manualmente

Embora o registro automático forneça uma maneira conveniente de instrumentar agentes, você pode querer instrumentar seu agente de forma mais granular ou adicionar rastreamentos adicionais que o registro automático não captura. Nesses casos, use APIs de rastreamento MLflow para adicionar rastreamentos manualmente.

As APIs de rastreamento do MLflow são APIs de baixo código para adicionar rastreabilidades sem se preocupar em gerenciar a estrutura de árvore das rastreabilidades. MLflow determina automaticamente as relações apropriadas de abrangência pai-filho usando a stack Python.

Combine registro automático e rastreamento manual

As APIs de rastreamento manual podem ser usadas com o registro automático. O MLflow combina os spans criados pelo registo automático e pela rastreação manual para criar uma trilha completa da execução do agente. Para obter um exemplo de combinação de registo automático e rastreamento manual, veja Instrumentando um agente de chamada de ferramenta com Rastreamento MLflow.

Rastreie funções usando o decorador @mlflow.trace

A maneira mais simples de instrumentar manualmente o seu código é decorar uma função com o decorador de @mlflow.trace. O decorador de rastreamento MLflow cria uma "extensão" com o escopo da função decorada, que representa uma unidade de execução em um rastreamento e é exibida como uma única linha na visualização de rastreamento. O span captura a entrada e a saída da função, a latência e quaisquer exceções levantadas pela função.

Por exemplo, o código a seguir cria uma extensão chamada my_function que captura argumentos de entrada x e y e a saída.

import mlflow

@mlflow.trace
def add(x: int, y: int) -> int:
  return x + y

Você também pode personalizar o nome da extensão, o tipo de extensão e adicionar atributos personalizados à extensão:

from mlflow.entities import SpanType

@mlflow.trace(
  # By default, the function name is used as the span name. You can override it with the `name` parameter.
  name="my_add_function",
  # Specify the span type using the `span_type` parameter.
  span_type=SpanType.TOOL,
  # Add custom attributes to the span using the `attributes` parameter. By default, MLflow only captures input and output.
  attributes={"key": "value"}
)
def add(x: int, y: int) -> int:
  return x + y

Rastrear blocos de código arbitrário usando o gerenciador de contexto

Para criar uma extensão para um bloco arbitrário de código, não apenas uma função, use mlflow.start_span() como um gerenciador de contexto que encapsula o bloco de código. O período começa quando o contexto é inserido e termina quando o contexto é encerrado. As entradas e saídas de span devem ser fornecidas manualmente, usando métodos setter do objeto span produzido pelo gestor de contexto. Para obter mais informações, consulte documentação do MLflow - manipulador de contexto.

with mlflow.start_span(name="my_span") as span:
  span.set_inputs({"x": x, "y": y})
  result = x + y
  span.set_outputs(result)
  span.set_attribute("key", "value")

Bibliotecas de rastreamento de nível inferior

O MLflow também fornece APIs de baixo nível para controlar explicitamente a estrutura da árvore de rastreamento. Consulte documentação do MLflow - Manual Instrumentation.

Exemplo de rastreamento: combinar registro automático e rastreamentos manuais

O exemplo seguinte combina o registo automático do OpenAI com o rastreamento manual para configurar completamente um agente de execução de ferramentas.

import json
from openai import OpenAI
import mlflow
from mlflow.entities import SpanType

client = OpenAI()

# Enable OpenAI autologging to capture LLM API calls
# (*Not necessary if you are using the Databricks Runtime 15.4 ML and above, where OpenAI autologging is enabled by default)
mlflow.openai.autolog()

# 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:
  if city == "Tokyo":
    return "sunny"
  elif city == "Paris":
    return "rainy"
  return "unknown"


tools = [
  {
    "type": "function",
    "function": {
      "name": "get_weather",
      "parameters": {
        "type": "object",
        "properties": {"city": {"type": "string"}},
      },
    },
  }
]

_tool_functions = {"get_weather": get_weather}

# Define a simple tool-calling agent
@mlflow.trace(span_type=SpanType.AGENT)
def run_tool_agent(question: str):
  messages = [{"role": "user", "content": question}]

  # Invoke the model with the given question and available tools
  response = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    tools=tools,
  )
  ai_msg = response.choices[0].message
  messages.append(ai_msg)

  # If the model requests tool calls, invoke the function(s) with the specified arguments
  if tool_calls := ai_msg.tool_calls:
    for tool_call in tool_calls:
      function_name = tool_call.function.name
      if tool_func := _tool_functions.get(function_name):
        args = json.loads(tool_call.function.arguments)
        tool_result = tool_func(**args)
      else:
        raise RuntimeError("An invalid tool is returned from the assistant!")

      messages.append(
        {
          "role": "tool",
          "tool_call_id": tool_call.id,
          "content": tool_result,
        }
      )

    # Send the tool results to the model and get a new response
    response = client.chat.completions.create(
      model="gpt-4o-mini", messages=messages
    )

  return response.choices[0].message.content

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

Anotar traços com tags

As tags de rastreamento do MLflow são pares chave-valor que permitem adicionar metadados personalizados aos rastreamentos, como um ID de conversa, um ID de usuário, hash de confirmação do Git, etc. As tags são exibidas na UI do MLflow para filtrar e pesquisar rastreamentos.

As tags podem ser definidas como um rastreamento contínuo ou concluído usando APIs MLflow ou a interface do usuário MLflow. O exemplo a seguir demonstra a adição de uma tag a um rastreamento contínuo usando a API mlflow.update_current_trace().

@mlflow.trace
def my_func(x):
    mlflow.update_current_trace(tags={"fruit": "apple"})
    return x + 1

Para saber mais sobre como marcar rastreamentos e como usá-los para filtrar e pesquisar rastreamentos, consulte documentação do MLflow - Definindo marcas de rastreamento.

Rever rastreamentos

Para revisar os rastreamentos depois de executar o agente, use uma das seguintes opções:

• Visualização em linha: Em blocos de anotações Databricks, os traços são renderizados em linha na saída da célula.
• experimento MLflow: No Databricks, vá para Experimentos> Selecione um experimento >Rastreamentos para visualizar e pesquisar todos os rastreamentos de um experimento.
• Execução de MLflow: Quando o agente é executado sob uma execução ativa do MLflow, os rastreamentos aparecem na página de execução da interface do MLflow.
• Agent Evaluation UI: No Mosaic AI Agent Evaluation, você pode revisar os registos para cada execução do agente clicando em Ver exibição detalhada do registo no resultado da avaliação.
• API de Pesquisa de Rastreamento: Para recuperar rastreamentos programaticamente, use a API de Pesquisa de Rastreamento .

Avaliar agentes usando rastreamentos

Os dados de rastreamento servem como um recurso valioso para avaliar seus agentes. Ao capturar informações detalhadas sobre a execução de seus modelos, o MLflow Tracing é fundamental na avaliação offline. Você pode usar os dados de rastreamento para avaliar o desempenho do agente em relação a um conjunto de dados dourado, identificar problemas e melhorar o desempenho do agente.

%pip install -U mlflow databricks-agents
%restart_python

import mlflow

# Get the recent 50 successful traces from the experiment
traces = mlflow.search_traces(
    max_results=50,
    filter_string="status = 'OK'",
)

traces.drop_duplicates("request", inplace=True) # Drop duplicate requests.
traces["trace"] = traces["trace"].apply(lambda x: x.to_json()) # Convert the trace to JSON format.

# Evaluate the agent with the trace data
mlflow.evaluate(data=traces, model_type="databricks-agent")

Para saber mais sobre a avaliação de agentes, consulte Realizar uma avaliação e ver os resultados.

Monitorar agentes implantados com tabelas de inferência

Depois que um agente é implantado no Mosaic AI Model Serving, você pode usar tabelas de inferência para monitorar o agente. As tabelas de inferência contêm logs detalhados de solicitações, respostas, rastreamentos de agentes e comentários de agentes do aplicativo de revisão. Essas informações permitem depurar problemas, monitorar o desempenho e criar um conjunto de dados dourado para avaliação offline.

Para habilitar tabelas de inferência para implantações de agentes, consulte Habilitar tabelas de inferência para agentes de IA.

Consultar rastreamentos online

Use um bloco de anotações para consultar a tabela de inferência e analisar os resultados.

Para visualizar rastreamentos, execute display(<the request logs table>) e selecione linhas para inspecionar:

# Query the inference table
df = spark.sql("SELECT * FROM <catalog.schema.my-inference-table-name>")
display(df)

Agentes de monitoramento

Veja Como monitorizar a sua aplicação de IA generativa.

Rastrear latência de sobrecarga

Os rastros são gravados de forma assíncrona para minimizar o impacto no desempenho. No entanto, o rastreamento ainda adiciona latência à velocidade de resposta do ponto final, especialmente quando o tamanho do rastreamento para cada solicitação de inferência é grande. A Databricks recomenda testar seu endpoint para entender os impactos da latência de rastreamento antes de implantar na produção.

A tabela a seguir fornece estimativas aproximadas do impacto da latência por tamanho de rastreamento:

Solução de problemas

Para solução de problemas e perguntas comuns, consulte a documentação do MLflow: Guia de instruções de rastreamento e a documentação do MLflow: Perguntas Frequentes .