Přidání trasování do agentů

Tento článek ukazuje, jak přidat trasování do vašich agentů pomocí rozhraní Fluent a MLflowClient API zpřístupněných pomocí trasování MLflow.

Požadavky

• MLflow 2.13.1

Přidání trasování do agentů pomocí automatického přihlašování

Pokud používáte knihovnu GenAI, která podporuje trasování (například LangChain, LlamaIndex nebo OpenAI), můžete povolit automatickélogování MLflow pro integraci knihovny a povolit trasování. Můžete například použít mlflow.langchain.autolog() k automatickému přidání trasování do agenta založeného na jazyce LangChain.

mlflow.langchain.autolog()

MLflow podporuje další knihovny pro automatické protokolování trasování. Úplný seznam integrovaných knihoven najdete v dokumentaci k trasování MLflow.

Použití rozhraní Fluent API k ručnímu přidání trasování do agenta

Následuje rychlý příklad, který používá rozhraní Fluent API: mlflow.trace a mlflow.start_span k přidání trasování do quickstart-agent. To se doporučuje pro modely 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

Provést odvozování

Po instrumentaci kódu můžete funkci spustit stejně jako obvykle. Následující příklad pokračuje s predict() funkcí v předchozí části. Trasování se automaticky zobrazí při spuštění metody vyvolá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,
  }
)

Rozhraní API fluentu

Rozhraní Fluent API v MLflow automaticky sestavují hierarchii trasování na základě toho, kde a kdy se kód spustí. Následující části popisují podporované úlohy pomocí rozhraní FLUENT API trasování MLflow.

Vyzdobit funkci

Funkci můžete ozdobit dekorátorem @mlflow.trace a vytvořit rozpětí pro rozsah dekorované funkce. Rozsah se spustí při vyvolání funkce a končí, když se vrátí. MLflow automaticky zaznamenává vstup a výstup funkce a také všechny výjimky vyvolané funkcí. Spuštěním následujícího kódu například vytvoříte rozsah s názvem "my_function", zachytíte vstupní argumenty x a y a také výstup funkce.

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

Použití správce kontextu trasování

Pokud chcete vytvořit rozsah pro libovolný blok kódu, nejen funkci, můžete použít mlflow.start_span() jako správce kontextu, který zabalí blok kódu. Rozsah začíná při zadání kontextu a končí při ukončení kontextu. Vstupy a výstupy rozsahu by měly být poskytovány ručně prostřednictvím metod setter objektu span, který je poskytován z správce kontextu.

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

Zabalení externí funkce

Funkci mlflow.trace lze použít jako obálku ke sledování funkce podle vašeho výběru. To je užitečné, když chcete trasovat funkce importované z externích knihoven. Vygeneruje stejné rozpětí, jaké byste získali dekorací této funkce.

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)

Klientská rozhraní API MLflow

MlflowClient zpřístupňuje podrobná rozhraní API bezpečná pro vlákna pro počáteční a koncové trasování, správu rozsahů a nastavení polí rozsahu. Poskytuje úplnou kontrolu nad životním cyklem a strukturou trasování. Tato rozhraní API jsou užitečná, pokud rozhraní Fluent API nestačí vašim požadavkům, jako jsou aplikace s více vlákny a zpětná volání.

Následující postup slouží k vytvoření kompletního trasování pomocí klienta MLflow.

1. Vytvořte instanci MLflowClient pomocí client = MlflowClient().

2. Spusťte trasování pomocí client.start_trace() metody. Tím se zahájí kontext trasování a spustí se absolutní kořenová délka a vrátí objekt kořenového rozsahu. Tato metoda musí být spuštěna před rozhraním start_span() API.

   1. Nastavte atributy, vstupy a výstupy pro trasování v client.start_trace().

   Poznámka:

   V rozhraních Fluent API neexistuje ekvivalent start_trace() metody. Důvodem je to, že rozhraní Fluent API automaticky inicializují kontext trasování a určují, jestli se jedná o kořenové rozpětí na základě spravovaného stavu.

3. Rozhraní API start_trace() vrátí rozsah. Získejte ID požadavku, jedinečný identifikátor trasování, který se označuje také jako trace_id, a ID vráceného rozsahu pomocí span.request_id a span.span_id.

4. Spuštění podřízeného rozsahu pomocí client.start_span(request_id, parent_id=span_id) nastavení atributů, vstupů a výstupů pro rozsah

   1. Tato metoda vyžaduje request_id a parent_id přidružuje rozsah ke správné pozici v hierarchii trasování. Vrátí další objekt spanu.

5. Ukončete rozsah podřízenosti voláním client.end_span(request_id, span_id).

6. Opakujte 3 až 5 pro všechny podřízené rozsahy, které chcete vytvořit.

7. Po ukončení všech podřízených rozsahů volání client.end_trace(request_id) zavřete celé trasování a zaznamenejte ho.

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