Adicionar MLflow Tracing a agentes AI

MLflow Tracing em Databricks permite que o senhor monitore e analise o comportamento dos agentes generativos de AI. Você pode ativar o rastreamento usando os seguintes métodos:

• Rastreamento automático: rastreie agentes com uma única linha de código.
• Rastreamento manual: instrumente seus agentes de forma mais granular ou trace propriedades que o rastreamento automático não cobre.
• Combine o rastreamento automático e manual para um rastreamento flexível e conveniente.

Rastreamento automático

MLflow O autologging simplifica o início do rastreamento de seus agentes. Basta adicionar a seguinte linha ao seu código:

Python

mlflow.<library>.autolog()

MLflow oferece suporte ao autologging para muitas bibliotecas populares de criação de agentes, incluindo as seguintes:

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()
SDK de agentes da OpenAI	0.0.7 ~ O mais recente	mlflow.openai.autolog()
Índice LlamaM	0.10.44 ~ O mais recente	mlflow.llama_index.autolog()
DSPy	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()
CrewAI	0.80.0 ~ O mais recente	mlflow.crewai.autolog()
LiteLLM	1.52.9 ~ O mais recente	mlflow.litellm.autolog()
Groq	0.13.0 ~ O mais recente	mlflow.groq.autolog()
Mistral	1.0.0 ~ O mais recente	mlflow.mistral.autolog()

Para obter uma lista completa das bibliotecas compatíveis e mais informações, consulte a documentação sobre autologging emMLflow.

Desativar o registro automático

O rastreamento de autologging é ativado por default em Databricks Runtime 15.4 ML e acima para LangChain, Langgraph, OpenAI e LlamaIndex. Para desativar o autologging, execute o seguinte comando em um Notebook:

Python

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

Rastreamento manual

MLflow Tracing APIs permitem que o usuário adicione traços manualmente, caso queira instrumentar o agente de forma mais granular ou adicionar traços adicionais que o autologging não captura.

MLflow Tracing APIs são low-code APIs para adicionar traços sem precisar gerenciar a estrutura de árvore do traço. O MLflow determina automaticamente as relações apropriadas de extensão pai-filho usando a pilha Python.

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 intervalo com o escopo da função decorada. O span registra entradas, saídas, latência e exceções.

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ê pode personalizar o nome e o tipo do intervalo e adicionar atributos personalizados:

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 fornece APIs de baixo nível para controlar explicitamente a estrutura da árvore de rastreamento. Consulte a documentação do MLflow - Instrumentação manual.

Combine registro automático e rastreamento manual

Você pode combinar o rastreamento manual e o registro automático. MLflow mescla os dois tipos de intervalos em um único rastreamento completo.

O exemplo a seguir combina o registro automático do OpenAI e o rastreamento manual:

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)

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 adiciona latência à velocidade de resposta do endpoint, especialmente quando o rastreamento é grande. Teste o site endpoint para entender o impacto antes de implantá-lo na produção.

A tabela a seguir estima o impacto da latência pelo 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.