Adición de seguimientos a los agentes

Artículo
11/08/2024

Importante

Esta característica está en versión preliminar pública.

En este artículo se muestra cómo agregar seguimientos a los agentes mediante las API Fluent y MLflowClient disponibles con seguimiento de MLflow.

Nota:

Para obtener ejemplos detallados de código y referencia de API para el seguimiento de MLflow, vea la documentación de MLflow.

Requisitos

MLflow 2.13.1

Uso del registro automático para agregar seguimientos a los agentes

Si usa una biblioteca de GenAI que tiene compatibilidad con el seguimiento (como LangChain, LlamaIndex o OpenAI), puede habilitar el registro automático de MLflow para que la integración de la biblioteca habilite el seguimiento. Por ejemplo, use mlflow.langchain.autolog() para agregar automáticamente seguimientos al agente basado en LangChain.

Nota:

A partir de Databricks Runtime 15.4 LTS ML, el seguimiento de MLflow está habilitado de forma predeterminada en cuadernos. Para deshabilitar el seguimiento, por ejemplo con LangChain, puede ejecutar mlflow.langchain.autolog(log_traces=False) en el cuaderno.

mlflow.langchain.autolog()

MLflow admite bibliotecas adicionales para el registro automático de seguimiento. Consulte la documentación de seguimiento de MLflow para obtener una lista completa de las bibliotecas integradas.

Uso de las API de Fluent para agregar seguimientos manualmente al agente

A continuación se muestra un ejemplo rápido que usa las API de Fluent: mlflow.trace y mlflow.start_span para agregar seguimientos a la quickstart-agent. Esto se recomienda para los 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 la inferencia

Después de instrumentar el código, puede ejecutar la función como lo haría normalmente. A continuación se continúa el ejemplo con la predict() función de la sección anterior. Los seguimientos se muestran automáticamente al ejecutar el método de invocación, 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,
  }
)

API de Fluent

Las API de Fluent de MLflow construyen automáticamente la jerarquía de seguimiento en función de dónde y cuándo se ejecuta el código. En las secciones siguientes se describen las tareas admitidas mediante las API de Fluent de seguimiento de MLflow.

Decorar la función

Puede decorar la función con el @mlflow.trace decorador para crear un intervalo para el ámbito de la función decorada. El intervalo se inicia cuando se invoca la función y finaliza cuando devuelve. MLflow registra automáticamente la entrada y salida de la función, así como las excepciones generadas a partir de la función. Por ejemplo, al ejecutar el código siguiente se creará un intervalo con el nombre “my_function”, se capturarán los argumentos de entrada x e y, así como la salida de la función.

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

Uso del administrador de contextos de seguimiento

Si desea crear un intervalo para un bloque arbitrario de código, no solo una función, puede usar mlflow.start_span() como administrador de contexto que encapsula el bloque de código. El intervalo se inicia cuando se escribe el contexto y finaliza cuando se cierra el contexto. Las entradas y salidas span deben proporcionarse manualmente a través de los métodos de establecimiento del objeto span que se obtiene del administrador de contextos.

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

Ajuste de una función externa

La función mlflow.trace se puede usar como contenedor para realizar un seguimiento de una función de su elección. Esto resulta útil cuando desea realizar un seguimiento de las funciones importadas desde bibliotecas externas. Genera el mismo intervalo que obtendría al decorar esa función.


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)

API de cliente de MLflow

MlflowClient expone API granulares seguras para subprocesos para iniciar y finalizar seguimientos, administrar intervalos y establecer campos de intervalo. Proporciona control total del ciclo de vida y la estructura del seguimiento. Estas API son útiles cuando las API de Fluent no son suficientes para sus requisitos, como aplicaciones multiproceso y devoluciones de llamada.

A continuación se indican los pasos para crear un seguimiento completo mediante el cliente de MLflow.

Cree una instancia de MLflowClient mediante client = MlflowClient().
Inicie un seguimiento mediante el método client.start_trace(). Esto inicia el contexto de seguimiento e inicia un intervalo raíz absoluto y devuelve un objeto de intervalo raíz. Este método debe ejecutarse antes de la API de start_span().
1. Establezca los atributos, las entradas y las salidas del seguimiento en client.start_trace().
Nota:

No hay un equivalente al método start_trace() en las API de Fluent. Esto se debe a que las API de Fluent inicializan automáticamente el contexto de seguimiento y determinan si es el intervalo raíz basado en el estado administrado.
La API start_trace() devuelve un intervalo. Obtenga el identificador de solicitud, un identificador único del seguimiento también denominado trace_idy el identificador del intervalo devuelto mediante span.request_id y span.span_id.
Inicie un intervalo secundario mediante client.start_span(request_id, parent_id=span_id) para establecer los atributos, las entradas y las salidas del intervalo.
1. Este método requiere request_id y parent_id asociar el intervalo con la posición correcta en la jerarquía de seguimiento. Devuelve otro objeto span.
Finalice el intervalo secundario llamando a client.end_span(request_id, span_id).
Repita 3 - 5 para los intervalos secundarios que quiera crear.
Una vez finalizados todos los intervalos secundarios, llame a client.end_trace(request_id) para cerrar todo el seguimiento y registrarlo.

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

Compartir a través de