Observabilidade do agente com o MLflow Tracing

Este artigo descreve como adicionar observabilidade aos seus aplicativos generativos AI com MLflow Tracing em Databricks.

O que é MLflow Tracing?

MLflow O rastreamento oferece observabilidade de ponta a ponta para aplicativos generativos do AI, desde o desenvolvimento até a implementação. O rastreamento é totalmente integrado ao Databricks conjunto AI de ferramentas gen da, capturando percepções detalhadas em todo o ciclo de vida de desenvolvimento e produção.

Veja a seguir os casos de uso do key para rastreamento em aplicativos do AI:

• Depuração simplificada : O rastreamento oferece visibilidade de cada etapa do seu aplicativo gen AI, facilitando o diagnóstico e a resolução de problemas.

• Avaliação off-line : o rastreamento gera dados valiosos para a avaliação dos agentes, permitindo que você meça e melhore a qualidade dos agentes ao longo do tempo.

• Monitoramento da produção : O rastreamento oferece visibilidade do comportamento do agente e das etapas de execução detalhadas, permitindo que o senhor monitore e otimize o desempenho do agente na produção.

• Auditoria logs : MLflow O rastreamento gera uma auditoria abrangente logs das ações e decisões do agente. Isso é fundamental para garantir compliance e apoiar a depuração quando surgirem problemas inesperados.

Requisitos

MLflow O rastreamento está disponível em MLflow nas versões 2.13.0 e acima. Databricks recomenda a instalação da versão mais recente do site MLflow para acessar os recursos e aprimoramentos mais recentes.

Bash

%pip install mlflow>=2.13.0 -qqqU
%restart_python

Rastreamento automático

O autologging do MLflow permite que o senhor instrumentalize rapidamente seu agente adicionando uma única linha ao seu código, mlflow.<library>.autolog().

MLflow oferece suporte a autologging para a biblioteca de criação de agentes mais popular. Para obter mais informações sobre cada biblioteca de autoria, consulte a documentação de autoriaMLflow:

Biblioteca	Suporte à versão de registro automático	Comando de autologação
LangChain	0.1.0 ~ Mais recente	mlflow.langchain.autolog()
Idioma	0.1.1 ~ O mais recente	mlflow.langgraph.autolog()
Inteligência artificial aberta	1.0.0 ~ O mais recente	mlflow.openai.autolog()
Índice LlamaM	0.10.44 ~ O mais recente	mlflow.llamaindex.autolog()
Espião	2.5.17 ~ O mais recente	mlflow.dspy.autolog()
Amazon Bedrock	1.33.0 ~ O mais recente (boto3)	mlflow.bedrock.autolog()
Anthropic	0.30.0 ~ O mais recente	mlflow.anthropic.autolog()
AutoGen	0,2,36 ~ 0,2,40	mlflow.autogen.autolog()
Google Gêmeos	1.0.0 ~ O mais recente	mlflow.gemini.autolog()
Tripulação AI	0.80.0 ~ O mais recente	mlflow.crewai.autolog()
Lite LM	1.52.9 ~ O mais recente	mlflow.litellm.autolog()
Cresça	0.13.0 ~ O mais recente	mlflow.groq.autolog()
Mistral	1.0.0 ~ O mais recente	mlflow.mistral.autolog()

Desativar o registro automático

O rastreamento de autologging é ativado por default em Databricks Runtime 15.4 ML e acima para a seguinte biblioteca:

• LangChain
• Idioma
• Inteligência artificial aberta
• Índice LlamaM

Para desativar o rastreamento de autologging para essas bibliotecas, execute o seguinte comando em um Notebook:

Python

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

Adicione traços manualmente

Embora o registro automático ofereça uma maneira conveniente de instrumentar agentes, talvez você queira instrumentar seu agente de forma mais granular ou adicionar traços adicionais que o registro automático não captura. Nesses casos, use as APIs de rastreamento do MLflow para adicionar rastreamentos manualmente.

As APIs de rastreamento do MLflow são APIs de baixo código para adicionar rastreamentos sem se preocupar com o gerenciamento da estrutura de árvore do rastreamento. O MLflow determina automaticamente as relações apropriadas de extensão pai-filho usando a pilha Python.

Combine registro automático e rastreamento manual

As APIs de rastreamento manual podem ser usadas com autologging. O MLflow combina os intervalos criados pelo autologging e pelo rastreamento manual para criar um rastreamento completo da execução do agente. Para obter um exemplo de combinação de autologging e rastreamento manual, consulte Instrumentação de um agente de chamada de ferramenta com o MLflow Tracing.

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

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

Por exemplo, o código a seguir cria um intervalo chamado my_function que captura os argumentos de entrada x e y e a saída.

Python

import mlflow

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

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

Python

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

Rastreie blocos de código arbitrários usando o gerenciador de contexto

Para criar uma extensão para um bloco arbitrário de código, não apenas para uma função, use mlflow.start_span() como um gerenciador de contexto que envolve o bloco de código. O intervalo começa quando o contexto é inserido e termina quando o contexto é encerrado. A entrada e as saídas de span devem ser fornecidas manualmente usando os métodos setter do objeto span gerado pelo gerenciador de contexto. Para obter mais informações, consulte a documentaçãoMLflow - context handler.

Python

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

Biblioteca de rastreamento de nível inferior

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

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

O exemplo a seguir combina o registro automático do OpenAI e o rastreamento manual para instrumentar totalmente um agente de chamada de ferramentas.

Python

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 etiquetas

MLflow As tags de rastreamento são par key-value que permitem adicionar metadados personalizados aos rastreamentos, como ID de conversa, ID de usuário, Git commit hash etc. As tags são exibidas na interface do usuário MLflow para filtrar e pesquisar rastreamentos.

As tags podem ser definidas para um rastreamento em andamento ou concluído usando as APIs do MLflow ou a interface do usuário do MLflow. O exemplo a seguir demonstra como adicionar uma tag a um rastreamento em andamento usando a API mlflow.update_current_trace().

Python

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

Para saber mais sobre a marcação de traços e como usá-los para filtrar e pesquisar traços, consulte a documentação do MLflow - Definição de tags de traço.

Analisar rastros

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

• Visualização em linha : No Databricks Notebook, os traços são renderizados em linha na saída da célula.
• MLflow experimento : Em Databricks, vá para Experiments > Select an experiment > Traces to view e pesquise todos os traços de um experimento.
• MLflow execução : Quando o agente executa sob uma execução ativa do MLflow, os rastros aparecem na página de execução da UI do MLflow.
• UI de avaliação do agente : Em Mosaic AI Agent Evaluation, o senhor pode revisar os traços de cada execução do agente clicando em See detailed trace view no resultado da avaliação.
• API de pesquisa de traços : Para recuperar traços de forma programática, use a API de pesquisa de traços.

Avalie 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 para a avaliação off-line. O senhor pode usar os dados de rastreamento para avaliar o desempenho do seu agente em relação a um golden dataset, identificar problemas e melhorar o desempenho do agente.

Python

%pip install -U mlflow databricks-agents
%restart_python

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 executar uma avaliação e view os resultados.

Monitorar agentes implantados com tabelas de inferência

Depois que um agente é implantado no Mosaic AI Model Serving, o senhor pode usar tabelas de inferência para monitorar o agente. As tabelas de inferência contêm detalhes logs de solicitações, respostas, rastreamentos de agentes e feedback de agentes do aplicativo de análise. Essas informações permitem que o senhor depure problemas, monitore o desempenho e crie um golden dataset para avaliação off-line.

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

Consulte rastros on-line

Use um Notebook para consultar a tabela de inferência e analisar os resultados.

Para visualizar os traços, execute display(<the request logs table>) e selecione as linhas a serem inspecionadas:

Python

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

Monitore agentes

Consulte Como monitorar seu aplicativo gen AI.

Rastreie a latência da sobrecarga

Os rastreamentos são gravados de forma assíncrona para minimizar o impacto no desempenho. No entanto, o rastreamento ainda acrescenta latência à velocidade de resposta do endpoint, principalmente quando o tamanho do rastreamento para cada solicitação de inferência é grande. Databricks recomenda testar o site endpoint para entender os impactos da latência do rastreamento antes de implantá-lo na produção.

A tabela a seguir fornece estimativas aproximadas do impacto da latência por tamanho do traço:

Tamanho do traço por solicitação	Impacto na latência da velocidade de resposta (ms)
~10 KB	~ 1 ms
~ 1 MB	50 ~ 100 ms
10 MB	150 ms ~

Solução de problemas

Para solução de problemas e perguntas comuns, consulte a documentaçãoMLflow: Tracing How-to guide (Guia prático de rastreamento ) e a documentaçãoMLflow: FAQ