Menambahkan jejak ke agen Anda
Penting
Fitur ini ada di Pratinjau Publik.
Artikel ini memperlihatkan cara menambahkan jejak ke agen Anda menggunakan API Fluent dan MLflowClient yang tersedia dengan Pelacakan MLflow.
Catatan
Untuk referensi API terperinci dan contoh kode untuk Pelacakan MLflow, lihat dokumentasi MLflow.
Persyaratan
- MLflow 2.13.1
Menggunakan autologging untuk menambahkan jejak ke agen Anda
Jika Anda menggunakan pustaka GenAI yang memiliki dukungan untuk pelacakan (seperti LangChain, LlamaIndex, atau OpenAI), Anda dapat mengaktifkan autologging MLflow untuk integrasi pustaka untuk mengaktifkan pelacakan.
Misalnya, gunakan mlflow.langchain.autolog() untuk menambahkan jejak secara otomatis ke agen berbasis LangChain Anda.
Catatan
Pada Databricks Runtime 15.4 LTS ML, pelacakan MLflow diaktifkan secara default dalam notebook. Untuk menonaktifkan pelacakan, misalnya dengan LangChain, Anda dapat menjalankan mlflow.langchain.autolog(log_traces=False) di buku catatan Anda.
mlflow.langchain.autolog()
MLflow mendukung pustaka tambahan untuk penelusuran otomatis pelacakan. Lihat dokumentasi Pelacakan MLflow untuk daftar lengkap pustaka terintegrasi.
Gunakan API Fasih untuk menambahkan jejak secara manual ke agen Anda
Berikut ini adalah contoh cepat yang menggunakan API Fasih: mlflow.trace dan mlflow.start_span untuk menambahkan jejak ke quickstart-agent. Ini direkomendasikan untuk model 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
Melakukan inferensi
Setelah melengkapi kode, Anda dapat menjalankan fungsi seperti biasa. Berikut ini melanjutkan contoh dengan predict() fungsi di bagian sebelumnya. Jejak secara otomatis ditampilkan saat Anda menjalankan metode pemanggilan, 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 Fasih
API Fasih di MLflow secara otomatis membuat hierarki jejak berdasarkan di mana dan kapan kode dijalankan. Bagian berikut menjelaskan tugas yang didukung menggunakan MLflow Tracing Fluent API.
Menghias fungsi Anda
Anda dapat menghias fungsi Anda dengan @mlflow.trace dekorator untuk membuat rentang untuk cakupan fungsi yang dihiasi. Rentang dimulai ketika fungsi dipanggil dan berakhir ketika kembali. MLflow secara otomatis merekam input dan output fungsi, serta pengecualian apa pun yang dinaikkan dari fungsi. Misalnya, menjalankan kode berikut akan membuat rentang dengan nama "my_function", menangkap argumen input x dan y, serta output fungsi.
@mlflow.trace(name="agent", span_type="TYPE", attributes={"key": "value"})
def my_function(x, y):
return x + y
Menggunakan pengelola konteks pelacakan
Jika Anda ingin membuat rentang untuk blok kode arbitrer, bukan hanya fungsi, Anda dapat menggunakan mlflow.start_span() sebagai manajer konteks yang membungkus blok kode. Rentang dimulai ketika konteks dimasukkan dan berakhir saat konteks keluar. Rentang input dan output harus disediakan secara manual melalui metode setter objek rentang yang dihasilkan dari manajer konteks.
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")
Membungkus fungsi eksternal
Fungsi ini mlflow.trace dapat digunakan sebagai pembungkus untuk melacak fungsi pilihan Anda. Ini berguna ketika Anda ingin melacak fungsi yang diimpor dari pustaka eksternal. Ini menghasilkan rentang yang sama seperti yang Akan Anda dapatkan dengan menghias fungsi tersebut.
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 Klien MLflow
MlflowClient mengekspos API terperinci dan aman utas untuk memulai dan mengakhiri jejak, mengelola rentang, dan mengatur bidang rentang. Ini memberikan kontrol penuh atas siklus hidup dan struktur jejak. API ini berguna ketika API Fasih tidak cukup untuk kebutuhan Anda, seperti aplikasi multi-utas dan panggilan balik.
Berikut ini adalah langkah-langkah untuk membuat jejak lengkap menggunakan Klien MLflow.
Buat instans MLflowClient dengan
client = MlflowClient().Mulai jejak menggunakan
client.start_trace()metode . Ini memulai konteks pelacakan dan memulai rentang akar absolut dan mengembalikan objek rentang akar. Metode ini harus dijalankan sebelumstart_span()API.- Atur atribut, input, dan output Anda untuk pelacakan di
client.start_trace().
Catatan
Tidak ada yang setara
start_trace()dengan metode dalam API Fasih. Ini karena API Fasih secara otomatis menginisialisasi konteks pelacakan dan menentukan apakah itu rentang akar berdasarkan status terkelola.- Atur atribut, input, dan output Anda untuk pelacakan di
API start_trace() mengembalikan rentang. Dapatkan ID permintaan, pengidentifikasi unik dari jejak yang juga disebut sebagai
trace_id, dan ID rentang yang dikembalikan menggunakanspan.request_iddanspan.span_id.Mulai rentang anak menggunakan
client.start_span(request_id, parent_id=span_id)untuk mengatur atribut, input, dan output Anda untuk rentang.- Metode ini memerlukan
request_iddanparent_iduntuk mengaitkan rentang dengan posisi yang benar dalam hierarki jejak. Ini mengembalikan objek rentang lain.
- Metode ini memerlukan
Akhiri rentang anak dengan memanggil
client.end_span(request_id, span_id).Ulangi 3 - 5 untuk setiap rentang anak yang ingin Anda buat.
Setelah semua rentang anak-anak berakhir, panggil
client.end_trace(request_id)untuk menutup seluruh jejak dan merekamnya.
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})