Partilhar via


Adicione rastreamentos aos seus agentes

Importante

Esta funcionalidade está em Pré-visualização Pública.

Este artigo mostra como adicionar rastreamentos aos seus agentes usando as APIs Fluent e MLflowClient disponibilizadas com o MLflow Tracing.

Nota

Para obter referência detalhada da API e exemplos de código para MLflow Tracing, consulte a documentação do MLflow.

Requisitos

  • MLflow 2.13.1

Use o registro automático para adicionar rastreamentos aos seus agentes

Se você estiver usando uma biblioteca GenAI que tenha suporte para rastreamento (como LangChain, LlamaIndex ou OpenAI), poderá habilitar o registro automático MLflow para a integração da biblioteca para habilitar o rastreamento. Por exemplo, use mlflow.langchain.autolog() para adicionar automaticamente rastreamentos ao seu agente baseado em LangChain.

Nota

A partir do Databricks Runtime 15.4 LTS ML, o rastreamento MLflow é habilitado por padrão em notebooks. Para desativar o rastreamento, por exemplo, com LangChain, você pode executar mlflow.langchain.autolog(log_traces=False) em seu bloco de anotações.

mlflow.langchain.autolog()

O MLflow suporta bibliotecas adicionais para registro automático de rastreamento. Consulte a documentação do MLflow Tracing para obter uma lista completa de bibliotecas integradas.

Use APIs Fluent para adicionar rastreamentos manualmente ao seu agente

A seguir está um exemplo rápido que usa as APIs Fluent: mlflow.trace e mlflow.start_span para adicionar rastreamentos ao quickstart-agent. Isto é recomendado para modelos PyFunc.


import mlflow
from mlflow.deployments import get_deploy_client

class QAChain(mlflow.pyfunc.PythonModel):
    def __init__(self):
        self.client = get_deploy_client("databricks")

    @mlflow.trace(name="quickstart-agent")
    def predict(self, model_input, system_prompt, params):
        messages = [
                {
                    "role": "system",
                    "content": system_prompt,
                },
                {
                    "role": "user",
                    "content":  model_input[0]["query"]
                }
          ]

        traced_predict = mlflow.trace(self.client.predict)
        output = traced_predict(
            endpoint=params["model_name"],
            inputs={
                "temperature": params["temperature"],
                "max_tokens": params["max_tokens"],
                "messages": messages,
            },
        )

        with mlflow.start_span(name="_final_answer") as span:
          # Initiate another span generation
            span.set_inputs({"query": model_input[0]["query"]})

            answer = output["choices"][0]["message"]["content"]

            span.set_outputs({"generated_text": answer})
            # Attributes computed at runtime can be set using the set_attributes() method.
            span.set_attributes({
              "model_name": params["model_name"],
                        "prompt_tokens": output["usage"]["prompt_tokens"],
                        "completion_tokens": output["usage"]["completion_tokens"],
                        "total_tokens": output["usage"]["total_tokens"]
                    })
              return answer

Realizar inferência

Depois de instrumentar seu código, você pode executar sua função como faria normalmente. O exemplo a seguir continua com a predict() função na seção anterior. Os rastreamentos são mostrados automaticamente quando você executa o método de invocação, predict().


SYSTEM_PROMPT = """
You are an assistant for Databricks users. You are answering python, coding, SQL, data engineering, spark, data science, DW and platform, API or infrastructure administration question related to Databricks. If the question is not related to one of these topics, kindly decline to answer. If you don't know the answer, just say that you don't know, don't try to make up an answer. Keep the answer as concise as possible. Use the following pieces of context to answer the question at the end:
"""

model = QAChain()

prediction = model.predict(
  [
      {"query": "What is in MLflow 5.0"},
  ],
  SYSTEM_PROMPT,
  {
    # Using Databricks Foundation Model for easier testing, feel free to replace it.
    "model_name": "databricks-dbrx-instruct",
    "temperature": 0.1,
    "max_tokens": 1000,
  }
)

APIs fluentes

As APIs Fluent no MLflow constroem automaticamente a hierarquia de rastreamento com base em onde e quando o código é executado. As seções a seguir descrevem as tarefas suportadas usando as APIs MLflow Tracing Fluent.

Decore a sua função

Você pode decorar sua função com o @mlflow.trace decorador para criar um espaço para o escopo da função decorada. A extensão começa quando a função é invocada e termina quando ela retorna. MLflow registra automaticamente a entrada e saída da função, bem como quaisquer exceções geradas pela função. Por exemplo, executar o código a seguir criará uma extensão com o nome "my_function", capturando os argumentos de entrada x e y, bem como a saída da função.

@mlflow.trace(name="agent", span_type="TYPE", attributes={"key": "value"})
def my_function(x, y):
    return x + y

Usar o gerenciador de contexto de rastreamento

Se você quiser criar uma extensão para um bloco arbitrário de código, não apenas uma função, você pode usar mlflow.start_span() como um gerenciador de contexto que encapsula o bloco de código. A extensão começa quando o contexto é inserido e termina quando o contexto é encerrado. As entradas e saídas span devem ser fornecidas manualmente por meio de métodos setter do objeto span que é produzido a partir do gerenciador de contexto.

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

Envolver uma função externa

A mlflow.trace função pode ser usada como um invólucro para rastrear uma função de sua escolha. Isso é útil quando você deseja rastrear funções importadas de bibliotecas externas. Ele gera o mesmo espaço que você obteria decorando essa função.


from sklearn.metrics import accuracy_score

y_pred = [0, 2, 1, 3]
y_true = [0, 1, 2, 3]

traced_accuracy_score = mlflow.trace(accuracy_score)
traced_accuracy_score(y_true, y_pred)

APIs do cliente MLflow

MlflowClient expõe APIs granulares e seguras para iniciar e terminar rastreamentos, gerenciar extensões e definir campos de extensão. Ele fornece controle total do ciclo de vida e da estrutura do rastreamento. Essas APIs são úteis quando as APIs Fluent não são suficientes para suas necessidades, como aplicativos multi-threaded e retornos de chamada.

A seguir estão as etapas para criar um rastreamento completo usando o MLflow Client.

  1. Crie uma instância de MLflowClient por client = MlflowClient().

  2. Inicie um rastreamento usando o client.start_trace() método. Isso inicia o contexto de rastreamento e inicia uma extensão raiz absoluta e retorna um objeto de extensão raiz. Esse método deve ser executado antes da start_span() API.

    1. Defina seus atributos, entradas e saídas para o rastreamento no client.start_trace().

    Nota

    Não há um equivalente ao start_trace() método nas APIs Fluent. Isso ocorre porque as APIs Fluent inicializam automaticamente o contexto de rastreamento e determinam se é a extensão raiz com base no estado gerenciado.

  3. A API start_trace() retorna uma extensão. Obtenha a ID da solicitação, um identificador exclusivo do rastreamento também conhecido como trace_id, e a ID da extensão retornada usando span.request_id e span.span_id.

  4. Inicie uma extensão filho usando client.start_span(request_id, parent_id=span_id) para definir seus atributos, entradas e saídas para a extensão.

    1. Este método requer request_id e parent_id para associar a extensão com a posição correta na hierarquia de rastreamento. Ele retorna outro objeto span.
  5. Termine a extensão da criança chamando client.end_span(request_id, span_id).

  6. Repita de 3 a 5 para todas as extensões de crianças que você deseja criar.

  7. Depois que todos os intervalos das crianças terminarem, ligue client.end_trace(request_id) para fechar todo o rastreamento e gravá-lo.

from mlflow.client import MlflowClient

mlflow_client = MlflowClient()

root_span = mlflow_client.start_trace(
  name="simple-rag-agent",
  inputs={
          "query": "Demo",
          "model_name": "DBRX",
          "temperature": 0,
          "max_tokens": 200
         }
  )

request_id = root_span.request_id

# Retrieve documents that are similar to the query
similarity_search_input = dict(query_text="demo", num_results=3)

span_ss = mlflow_client.start_span(
      "search",
      # Specify request_id and parent_id to create the span at the right position in the trace
        request_id=request_id,
        parent_id=root_span.span_id,
        inputs=similarity_search_input
  )
retrieved = ["Test Result"]

# Span has to be ended explicitly
mlflow_client.end_span(request_id, span_id=span_ss.span_id, outputs=retrieved)

root_span.end_trace(request_id, outputs={"output": retrieved})