Protokollieren und Registrieren von KI-Agents
Wichtig
Dieses Feature befindet sich in der Public Preview.
Protokollieren von KI-Agenten mithilfe von Mosaik AI Agent Framework. Das Protokollieren eines Agents ist die Grundlage des Entwicklungsprozesses. Die Protokollierung erfasst einen "Zeitpunkt" des Codes und der Konfiguration des Agents, damit Sie die Qualität der Konfiguration auswerten können.
Anforderungen
Erstellen Sie einen KI-Agent, bevor Sie ihn protokollieren.
Databricks empfiehlt, die neueste Version des databricks-sdkzu installieren.
% pip install databricks-sdk
Codebasierte Protokollierung
Databricks empfiehlt die Verwendung der Modelle von MLflow aus Codefunktionalität beim Protokollieren von Agents.
Bei diesem Ansatz wird der Code des Agents als Python-Datei erfasst, und die Python-Umgebung wird als Liste der Pakete erfasst. Wenn der Agent bereitgestellt wird, wird die Python-Umgebung wiederhergestellt, und der Code des Agents wird ausgeführt, um den Agent in den Arbeitsspeicher zu laden, damit er aufgerufen werden kann, wenn der Endpunkt aufgerufen wird.
Sie können diesen Ansatz mit der Verwendung von Validierungs-APIs vor der Bereitstellung wie mlflow.models.predict() koppeln, um sicherzustellen, dass der Agent zuverlässig ausgeführt wird, wenn er im Einsatz ist.
Der Code, der den Agenten oder Agenten aufzeichnet, muss sich in einem separatem Notizbuch vom Agent-Code befinden. Dieses Notizbuch wird als Treibernotizbuch bezeichnet. Ein Beispielnotizbuch finden Sie unter Beispielnotizbücher.
Modellsignatur bei der Protokollierung ermitteln
Während der Protokollierung müssen Sie einen MLflow-Modellsignatur-definieren, der das Eingabe- und Ausgabeschema des Agents angibt. Die Signatur überprüft Eingaben und Ausgaben, um sicherzustellen, dass der Agent ordnungsgemäß mit nachgelagerten Tools wie AI Playground und der Überprüfungs-App interagiert. Es führt auch andere Anwendungen dazu, wie der Agent effektiv verwendet werden kann.
Databricks empfiehlt die Verwendung der Modellsignatur-Inferencing-Funktionen von MLflow, um die Signatur des Agents basierend auf einem von Ihnen bereitgestellten Eingabebeispiel automatisch zu generieren. Dieser Ansatz ist praktischer als die manuelle Definition der Signatur.
Die LangChain und PyFunc Beispiele unten verwenden Modellsignatur-Ableitungen.
Wenn Sie zum Zeitpunkt der Protokollierung lieber selbst eine Modell-Signatur explizit definieren möchten, finden Sie Informationen dazu in MLflow-Dokumentation – So geht's: Modelle mit Signaturen protokollieren.
Codebasierte Protokollierung mit LangChain
Die folgenden Anweisungen und codebeispiele zeigen Ihnen, wie Sie einen Agent mit LangChain protokollieren.
Erstellen Sie ein Notizbuch oder eine Python-Datei mit Ihrem Code. In diesem Beispiel wird das Notizbuch oder die Datei
agent.pybenannt. Das Notizbuch oder die Datei muss einen LangChain-Agent enthalten, der hier alslc_agentbezeichnet wird.Schließen Sie mlflow.models.set_model(lc_agent) in das Notizbuch oder die Datei ein.
Erstellen Sie ein neues Notizbuch, das als Treibernotizbuch dienen soll (in diesem Beispiel
driver.pygenannt).Verwenden Sie im Treibernotizbuch den folgenden Code, um
agent.pyauszuführen und die Ergebnisse in einem MLflow-Modell zu protokollieren:mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)Der
resources-Parameter deklariert databricks-verwaltete Ressourcen, die erforderlich sind, um den Agent zu bedienen, z. B. einen Vektorsuchindex oder einen Endpunkt, der einem Foundation-Modell dient. Weitere Informationen finden Sie unter Ressourcen für die automatische Authentifizierungsweiterleitung festlegen.Bereitstellen des Modells. Siehe Bereitstellen eines Agents für generative KI-Anwendung.
Wenn die Bereitstellungsumgebung geladen wird, wird
agent.pyausgeführt.Wenn eine Dienstanforderung erhalten wird, wird
lc_agent.invoke(...)aufgerufen.
import mlflow
code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"
# Input example used by MLflow to infer Model Signature
input_example = {
"messages": [
{
"role": "user",
"content": "What is Retrieval-augmented Generation?",
}
]
}
# example using langchain
with mlflow.start_run():
logged_agent_info = mlflow.langchain.log_model(
lc_model=code_path,
model_config=config_path, # If you specify this parameter, this configuration is used by agent code. The development_config is overwritten.
artifact_path="agent", # This string is used as the path inside the MLflow model where artifacts are stored
input_example=input_example, # Must be a valid input to the agent
example_no_conversion=True, # Required
)
print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")
# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.langchain.load_model(logged_agent_info.model_uri)
model.invoke(example)
Codebasierte Protokollierung mit PyFunc
Die folgenden Anweisungen und Codebeispiele zeigen Ihnen, wie Sie einen Agent mit PyFunc protokollieren.
Erstellen Sie ein Notizbuch oder eine Python-Datei mit Ihrem Code. In diesem Beispiel wird das Notizbuch oder die Datei
agent.pybenannt. Das Notizbuch oder die Datei muss eine PyFunc-Klasse mit dem NamenPyFuncClassenthalten.Fügen Sie
mlflow.models.set_model(PyFuncClass)zum Notizbuch oder zur Datei hinzu.Erstellen Sie ein neues Notizbuch, das als Treibernotizbuch dienen soll (in diesem Beispiel
driver.pygenannt).Verwenden Sie im Treibernotizbuch den folgenden Code, um
agent.pyauszuführen und die Ergebnisse in einem MLflow-Modell zu protokollieren:mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)Der
resources-Parameter deklariert databricks-verwaltete Ressourcen, die erforderlich sind, um den Agent zu bedienen, z. B. einen Vektorsuchindex oder einen Endpunkt, der einem Foundation-Modell dient. Weitere Informationen finden Sie unter Angabe von Ressourcen für die automatische Authentifizierungsweitergabe.Bereitstellen des Modells. Siehe Bereitstellen eines Agents für generative KI-Anwendung.
Wenn die Bereitstellungsumgebung geladen wird, wird
agent.pyausgeführt.Wenn eine Dienstanforderung erhalten wird, wird
PyFuncClass.predict(...)aufgerufen.
import mlflow
from mlflow.models.resources import (
DatabricksServingEndpoint,
DatabricksVectorSearchIndex,
)
code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"
# Input example used by MLflow to infer Model Signature
input_example = {
"messages": [
{
"role": "user",
"content": "What is Retrieval-augmented Generation?",
}
]
}
with mlflow.start_run():
logged_agent_info = mlflow.pyfunc.log_model(
python_model=agent_notebook_path,
artifact_path="agent",
input_example=input_example,
resources=resources_path,
example_no_conversion=True,
resources=[
DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
]
)
print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")
# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.pyfunc.load_model(logged_agent_info.model_uri)
model.invoke(example)
Angeben von Ressourcen für die automatische Authentifizierungspassthrough
KI-Agents müssen sich häufig bei anderen Ressourcen authentifizieren, um Aufgaben auszuführen. Beispielsweise muss ein Agent möglicherweise auf einen Vektorsuchindex zugreifen, um unstrukturierte Daten abzufragen.
Wie in Authentifizierung für abhängige Ressourcenbeschrieben, unterstützt Model Serving die Authentifizierung sowohl für vom Databricks verwaltete als auch für externe Ressourcen, wenn Sie den Agent bereitstellen.
Für die am häufigsten verwendeten Databricks-Ressourcentypen unterstützt Databricks und empfiehlt, Während der Protokollierung Ressourcenabhängigkeiten für den Agent vorab zu deklarieren. Dies ermöglicht den automatischen Authentifizierungsdurchlauf, wenn Sie den Agenten bereitstellen – Databricks stellt, rotiert und verwaltet automatisch kurzfristige Berechtigungsnachweise, um innerhalb des Agent-Endpunkts sicher auf diese Ressourcenabhängigkeiten zuzugreifen.
Zum Aktivieren der automatischen Authentifizierungsdurchlauf geben Sie abhängige Ressourcen mithilfe des resources-Parameters der log_model()-API an, wie im folgenden Code gezeigt.
import mlflow
from mlflow.models.resources import (
DatabricksVectorSearchIndex,
DatabricksServingEndpoint,
DatabricksSQLWarehouse,
DatabricksFunction,
DatabricksGenieSpace,
DatabricksTable,
)
with mlflow.start_run():
logged_agent_info = mlflow.pyfunc.log_model(
python_model=agent_notebook_path,
artifact_path="agent",
input_example=input_example,
example_no_conversion=True,
# Specify resources for automatic authentication passthrough
resources=[
DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"),
DatabricksSQLWarehouse(warehouse_id="your_warehouse_id"),
DatabricksFunction(function_name="ml.tools.python_exec"),
DatabricksGenieSpace(genie_space_id="your_genie_space_id"),
DatabricksTable(table_name="your_table_name"),
]
)
Databricks empfiehlt, resources für alle Agent-Aromen manuell anzugeben.
Hinweis
Wenn Sie beim Protokollieren von LangChain-Agents mit mlflow.langchain.log_model(...)keine Ressourcen angeben, führt MLflow eine automatische Schätzung der Ressourcen nach bestem Bemühen durch. Dies kann jedoch nicht alle Abhängigkeiten erfassen, was zu Autorisierungsfehlern bei der Bereitstellung oder Abfrage des Agents führt.
In der folgenden Tabelle sind die Databricks-Ressourcen aufgeführt, die die automatische Authentifizierungsweiterleitung unterstützen, sowie die mindestens erforderliche mlflow-Version, um die Ressource zu protokollieren.
| Ressourcentyp | Mindestens mlflow Version erforderlich, um die Ressource zu protokollieren |
|---|---|
| Vektorsuchindex | Erfordert mlflow 2.13.1 oder höher |
| Modell, das Endpunkt bedient | Erfordert mlflow 2.13.1 oder höher |
| SQL-Warehouse | Erfordert mlflow 2.16.1 oder höher |
| Unity-Katalogfunktion | Erfordert mlflow 2.16.1 oder höher |
| Genie Space | Erfordert mlflow 2.17.1 oder höher |
| Unity-Katalogtabelle | Erfordert mlflow 2.18.0 oder höher |
Automatische Authentifizierung für OpenAI-Clients
Wenn Ihr Agent den OpenAI-Client verwendet, verwenden Sie das Databricks SDK, um sich während der Bereitstellung automatisch zu authentifizieren. Databricks SDK bietet einen Wrapper zum Erstellen des OpenAI-Clients mit automatisch konfigurierter Autorisierung. Führen Sie folgendes in Ihrem Notizbuch aus:
% pip install databricks-sdk[openai]
from databricks.sdk import WorkspaceClient
def openai_client(self):
w = WorkspaceClient()
return w.serving_endpoints.get_open_ai_client()
Geben Sie dann den Model Serving-Endpunkt als Teil von resources an, um sich automatisch bei der Bereitstellung zu authentifizieren.
Registrieren des Agents im Unity-Katalog
Bevor Sie den Agent bereitstellen, müssen Sie den Agent im Unity-Katalog registrieren. Das Registrieren der Agent-Pakete paketiert es als Modell im Unity-Katalog. Daher können Sie Unity-Katalogberechtigungen für die Autorisierung für Ressourcen im Agent verwenden.
import mlflow
mlflow.set_registry_uri("databricks-uc")
catalog_name = "test_catalog"
schema_name = "schema"
model_name = "agent_name"
model_name = catalog_name + "." + schema_name + "." + model_name
uc_model_info = mlflow.register_model(model_uri=logged_agent_info.model_uri, name=model_name)