Lägga till spårningar i dina agenter

Artikel
08/17/2024

Viktigt!

Den här funktionen finns som allmänt tillgänglig förhandsversion.

Den här artikeln visar hur du lägger till spårningar till dina agenter med hjälp av API:erna Fluent och MLflowClient som gjorts tillgängliga med MLflow Tracing.

Kommentar

Detaljerade API-referens- och kodexempel för MLflow Tracing finns i MLflow-dokumentationen.

Krav

MLflow 2.13.1

Använd automatisk loggning för att lägga till spårningar i LangChain-agenten

Om du använder LangChain använder du MLflow för langchain.autolog() att automatiskt lägga till spårningar till din agent. Detta rekommenderas för Langchain-agenter.

mlflow.langchain.autolog()

MLflow stöder ytterligare bibliotek för automatisk spårningsloggning. En fullständig lista över integrerade bibliotek finns i dokumentationen för MLflow Tracing.

Använd Fluent-API:er för att manuellt lägga till spårningar i din agent

Följande är ett snabbt exempel som använder Fluent API:er: mlflow.trace och mlflow.start_span för att lägga till spårningar i quickstart-agent. Detta rekommenderas för PyFunc-modeller.


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

Utföra slutsatsdragning

När du har instrumenterat koden kan du köra funktionen som vanligt. Följande fortsätter exemplet med predict() funktionen i föregående avsnitt. Spårningarna visas automatiskt när du kör anropsmetoden . 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:er för Fluent

Fluent-API:erna i MLflow konstruerar automatiskt spårningshierarkin baserat på var och när koden körs. I följande avsnitt beskrivs de uppgifter som stöds med hjälp av API:erna för MLflow Tracing Fluent.

Dekorera din funktion

Du kan dekorera din funktion med dekoratören @mlflow.trace för att skapa ett spann för omfånget för den dekorerade funktionen. Intervallet startar när funktionen anropas och slutar när den returneras. MLflow registrerar automatiskt indata och utdata för funktionen, samt eventuella undantag som genereras från funktionen. Om du till exempel kör följande kod skapas ett intervall med namnet "my_function", som samlar in indataargumenten x och y, samt utdata från funktionen.

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

Använda spårningskontexthanteraren

Om du vill skapa ett intervall för ett godtyckligt kodblock, inte bara en funktion, kan du använda mlflow.start_span() som kontexthanterare som omsluter kodblocket. Intervallet börjar när kontexten anges och slutar när kontexten avslutas. Span-indata och utdata ska anges manuellt via settermetoder för span-objektet som returneras från kontexthanteraren.

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

Omsluta en extern funktion

Funktionen mlflow.trace kan användas som omslutning för att spåra valfri funktion. Detta är användbart när du vill spåra funktioner som importerats från externa bibliotek. Det genererar samma spann som du skulle få genom att dekorera den funktionen.


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)

MLflow-klient-API:er

MlflowClient exponerar detaljerade, trådsäkra API:er för att starta och avsluta spårningar, hantera intervall och ange intervallfält. Den ger fullständig kontroll över spårningens livscykel och struktur. Dessa API:er är användbara när Fluent-API:erna inte räcker till för dina krav, till exempel program med flera trådar och återanrop.

Följande är steg för att skapa en fullständig spårning med hjälp av MLflow-klienten.

Skapa en instans av MLflowClient av client = MlflowClient().
Starta en spårning med hjälp av client.start_trace() metoden . Detta initierar spårningskontexten och startar ett absolut rotintervall och returnerar ett rotintervallobjekt. Den här metoden måste köras före API:et start_span() .
1. Ange attribut, indata och utdata för spårningen i client.start_trace().
Kommentar

Det finns inte en motsvarighet till start_trace() metoden i Fluent-API:erna. Det beror på att Fluent-API:erna automatiskt initierar spårningskontexten och avgör om det är rotintervallet baserat på det hanterade tillståndet.
API:et start_trace() returnerar ett spann. Hämta begärande-ID: t, en unik identifierare för spårningen som trace_idäven kallas , och ID för det returnerade intervallet med hjälp av span.request_id och span.span_id.
Starta ett underordnat spann med för client.start_span(request_id, parent_id=span_id) att ange attribut, indata och utdata för intervallet.
1. Den här metoden kräver request_id och parent_id associerar spannet med rätt position i spårningshierarkin. Det returnerar ett annat span-objekt.
Avsluta det underordnade intervallet genom att anropa client.end_span(request_id, span_id).
Upprepa 3–5 för alla underordnade intervall som du vill skapa.
När alla underordnade intervall har avslutats anropar du client.end_trace(request_id) för att stänga hela spårningen och registrera den.

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

Dela via