Traceringen toevoegen aan uw agents
Belangrijk
Deze functie is beschikbaar als openbare preview.
In dit artikel wordt beschreven hoe u traceringen toevoegt aan uw agents met behulp van de Fluent- en MLflowClient-API's die beschikbaar zijn gemaakt met MLflow Tracing.
Notitie
Vereisten
- MLflow 2.13.1
Autologging gebruiken om traceringen toe te voegen aan uw agents
Als u een GenAI-bibliotheek gebruikt die ondersteuning biedt voor tracering (zoals LangChain, LlamaIndex of OpenAI), kunt u de automatische registratie van MLflow inschakelen voor de bibliotheekintegratie om tracering in te schakelen.
Gebruik bijvoorbeeld mlflow.langchain.autolog() om automatisch traceringen toe te voegen aan uw LangChain-agent.
Notitie
Vanaf Databricks Runtime 15.4 LTS ML is MLflow-tracering standaard ingeschakeld in notebooks. Als u tracering wilt uitschakelen, bijvoorbeeld met LangChain, kunt u uitvoeren mlflow.langchain.autolog(log_traces=False) in uw notebook.
mlflow.langchain.autolog()
MLflow biedt ondersteuning voor aanvullende bibliotheken voor automatisch logboeken traceren. Zie de documentatie voor MLflow Tracing voor een volledige lijst met geïntegreerde bibliotheken.
Fluent-API's gebruiken om handmatig traceringen aan uw agent toe te voegen
Hier volgt een snel voorbeeld waarin de Fluent-API's worden gebruikt: mlflow.trace en mlflow.start_span traceringen toevoegen aan de quickstart-agent. Dit wordt aanbevolen voor PyFunc-modellen.
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
Deductie uitvoeren
Nadat u uw code hebt geïnstrueerd, kunt u uw functie op de gebruikelijke manier uitvoeren. Hieronder wordt het voorbeeld voortgezet met de predict() functie in de vorige sectie. De traceringen worden automatisch weergegeven wanneer u de aanroepmethode uitvoert. 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,
}
)
Fluent-API's
De Fluent-API's in MLflow maken automatisch de traceringshiërarchie op basis van waar en wanneer de code wordt uitgevoerd. In de volgende secties worden de ondersteunde taken beschreven met behulp van de Fluent-API's voor MLflow Tracing.
Uw functie versieren
U kunt uw functie versieren met de @mlflow.trace decorator om een spanwijdte te creëren voor het bereik van de versierde functie. De periode begint wanneer de functie wordt aangeroepen en eindigt wanneer deze wordt geretourneerd. MLflow registreert automatisch de invoer en uitvoer van de functie, evenals eventuele uitzonderingen die van de functie zijn gegenereerd. Als u bijvoorbeeld de volgende code uitvoert, wordt een span gemaakt met de naam 'my_function', waarbij de invoerargumenten x en y worden vastgelegd, evenals de uitvoer van de functie.
@mlflow.trace(name="agent", span_type="TYPE", attributes={"key": "value"})
def my_function(x, y):
return x + y
Contextbeheer voor tracering gebruiken
Als u een bereik wilt maken voor een willekeurig codeblok, niet alleen een functie, kunt u gebruiken mlflow.start_span() als contextbeheerder die het codeblok verpakt. De periode begint wanneer de context wordt ingevoerd en eindigt wanneer de context wordt afgesloten. De spaninvoer en uitvoer moeten handmatig worden opgegeven via settermethoden van het spanobject dat wordt geretourneerd door de contextbeheerder.
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")
Een externe functie verpakken
De mlflow.trace functie kan worden gebruikt als wrapper om een functie van uw keuze te traceren. Dit is handig als u functies wilt traceren die zijn geïmporteerd uit externe bibliotheken. Het genereert dezelfde periode als u zou krijgen door die functie te decoreren.
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-client-API's
MlflowClient biedt gedetailleerde, threadveilige API's voor het starten en beëindigen van traceringen, het beheren van spanvakken en het instellen van spanvelden. Het biedt volledige controle over de levenscyclus en structuur van de tracering. Deze API's zijn handig wanneer de Fluent-API's niet voldoende zijn voor uw vereisten, zoals toepassingen met meerdere threads en callbacks.
Hier volgen de stappen voor het maken van een volledige tracering met behulp van de MLflow-client.
Maak een exemplaar van MLflowClient door
client = MlflowClient().Start een tracering met behulp van de
client.start_trace()methode. Hiermee wordt de traceringscontext gestart en wordt een absolute hoofdspanne gestart en wordt een hoofdspanneobject geretourneerd. Deze methode moet worden uitgevoerd vóór destart_span()API.- Stel uw kenmerken, invoer en uitvoer voor de tracering in
client.start_trace().
Notitie
Er is geen equivalent aan de
start_trace()methode in de Fluent-API's. Dit komt doordat de Fluent-API's de traceringscontext automatisch initialiseren en bepalen of het de hoofdspanne is op basis van de beheerde status.- Stel uw kenmerken, invoer en uitvoer voor de tracering in
De START_TRACE() API retourneert een periode. Haal de aanvraag-id op, een unieke id van de tracering, ook wel aangeduid als de
trace_id, en de id van de geretourneerde spanwijdte met behulpspan.request_idvan enspan.span_id.Start een onderliggend bereik met behulp van
client.start_span(request_id, parent_id=span_id)het instellen van uw kenmerken, invoer en uitvoer voor de periode.- Deze methode vereist
request_idenparent_idkoppelt de spaniteit aan de juiste positie in de traceringshiërarchie. Het retourneert een ander spanobject.
- Deze methode vereist
Beëindig het bereik van het kind door aan te roepen
client.end_span(request_id, span_id).Herhaal 3 - 5 voor alle onderliggende groepen die u wilt maken.
Nadat alle onderliggende items zijn beëindigd,
client.end_trace(request_id)roept u aan om de volledige tracering te sluiten en vast te leggen.
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})