Menambahkan jejak ke agen Anda

Artikel
07/27/2024

Penting

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 Langchain Anda

Jika Anda menggunakan Langchain, gunakan MLflow langchain.autolog() untuk menambahkan jejak secara otomatis ke agen Anda. Ini direkomendasikan untuk agen Langchain.

mlflow.langchain.autolog()

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 sebelum start_span() API.
1. 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.
API start_trace() mengembalikan rentang. Dapatkan ID permintaan, pengidentifikasi unik dari jejak yang juga disebut sebagai trace_id, dan ID rentang yang dikembalikan menggunakan span.request_id dan span.span_id.
Mulai rentang anak menggunakan client.start_span(request_id, parent_id=span_id) untuk mengatur atribut, input, dan output Anda untuk rentang.
1. Metode ini memerlukan request_id dan parent_id untuk mengaitkan rentang dengan posisi yang benar dalam hierarki jejak. Ini mengembalikan objek rentang lain.
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})

Bagikan melalui