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.
Crie uma instância de MLflowClient por
client = MlflowClient()
.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 dastart_span()
API.- 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.- Defina seus atributos, entradas e saídas para o rastreamento no
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 usandospan.request_id
espan.span_id
.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.- Este método requer
request_id
eparent_id
para associar a extensão com a posição correta na hierarquia de rastreamento. Ele retorna outro objeto span.
- Este método requer
Termine a extensão da criança chamando
client.end_span(request_id, span_id)
.Repita de 3 a 5 para todas as extensões de crianças que você deseja criar.
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})