Protokollieren und Registrieren von KI-Agents

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.

Codebasierte Protokollierung im Vergleich zur serialisierungsbasierten Protokollierung

Sie können die codebasierte MLflow-Protokollierung oder die serialisierungsbasierte MLflow-Protokollierung verwenden. Databricks empfiehlt, die codebasierte Protokollierung zu verwenden.

Codebasierte MLflow-Protokollierung: Der Code der Kette wird als Python-Datei erfasst. Die Python-Umgebung wird als Liste der Pakete erfasst. Wenn die Kette bereitgestellt wird, wird die Python-Umgebung wiederhergestellt, und der Code der Kette wird ausgeführt, um die Kette in den Arbeitsspeicher zu laden, damit sie aufgerufen werden kann, wenn der Endpunkt aufgerufen wird.

Serialisierungsbasierte MLflow-Protokollierung: Der Code und der aktuelle Zustand der Kette in der Python-Umgebung wird auf den Datenträger serialisiert, häufig mithilfe von Bibliotheken wie pickle oder joblib. Wenn die Kette bereitgestellt wird, wird die Python-Umgebung wiederhergestellt, und das serialisierte Objekt wird in den Arbeitsspeicher geladen, sodass es aufgerufen werden kann, wenn der Endpunkt aufgerufen wird.

In der Tabelle sind die Vor- und Nachteile der einzelnen Methoden aufgeführt.

Für die codebasierte Protokollierung muss sich der Code, der Ihren Agent oder Ihre Kette protokolliert, in einem separaten Notizbuch von Ihrem Kettencode befinden. Dieses Notizbuch wird als Treibernotizbuch bezeichnet. Ein Beispielnotizbuch finden Sie unter Beispielnotizbücher.

Codebasierte Protokollierung mit LangChain

1. Erstellen Sie ein Notizbuch oder eine Python-Datei mit Ihrem Code. Für dieses Beispiel wird das Notizbuch oder die Datei chain.py genannt. Das Notizbuch oder die Datei muss eine LangChain-Kette enthalten, die hier als lc_chain bezeichnet wird.
2. Fügen Sie mlflow.models.set_model(lc_chain) zum Notizbuch oder zur Datei hinzu.
3. Erstellen Sie ein neues Notizbuch, das als Treibernotizbuch dienen soll (in diesem Beispiel driver.py genannt).
4. Verwenden Sie im Treibernotebook mlflow.lang_chain.log_model(lc_model=”/path/to/chain.py”), um chain.py auszuführen, und protokollieren Sie die Ergebnisse in einem MLflow-Modell.
5. Bereitstellen des Modells. Siehe Bereitstellen eines Agents für generative KI-Anwendung. Die Bereitstellung Ihres Agents hängt möglicherweise von anderen Databricks-Ressourcen ab, z. B. von einem Vektorsuchindex und einem Modellbereitstellungsendpunkten. Für LangChain-Agents:
   • MLflow log_model leitet die von der Kette benötigten Abhängigkeiten ab und protokolliert sie in der MLmodel-Datei im protokollierten Modellartefakt.
   • Während der Bereitstellung erstellt databricks.agents.deploy automatisch die M2M OAuth-Token, die erforderlich sind, um auf diese abgeleiteten Ressourcenabhängigkeiten zuzugreifen und mit ihnen zu kommunizieren.
6. Wenn die Bereitstellungsumgebung geladen wird, wird chain.py ausgeführt.
7. Wenn eine Dienstanforderung erhalten wird, wird lc_chain.invoke(...) aufgerufen.

import mlflow

code_path = "/Workspace/Users/first.last/chain.py"
config_path = "/Workspace/Users/first.last/config.yml"

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-augmented Generation?",
        }
    ]
}

# example using LangChain
with mlflow.start_run():
  logged_chain_info = mlflow.langchain.log_model(
    lc_model=code_path,
    model_config=config_path, # If you specify this parameter, this is the configuration that is used for training the model. The development_config is overwritten.
    artifact_path="chain", # 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 your chain
    example_no_conversion=True, # Required
  )

print(f"MLflow Run: {logged_chain_info.run_id}")
print(f"Model URI: {logged_chain_info.model_uri}")

# To verify that the model has been logged correctly, load the chain and call `invoke`:
model = mlflow.langchain.load_model(logged_chain_info.model_uri)
model.invoke(example)

Codebasierte Protokollierung mit PyFunc

1. Erstellen Sie ein Notizbuch oder eine Python-Datei mit Ihrem Code. Für dieses Beispiel wird das Notizbuch oder die Datei chain.py genannt. Das Notizbuch oder die Datei muss eine PyFunc-Klasse enthalten, die hier als PyFuncClass bezeichnet wird.
2. Fügen Sie mlflow.models.set_model(PyFuncClass) zum Notizbuch oder zur Datei hinzu.
3. Erstellen Sie ein neues Notizbuch, das als Treibernotizbuch dienen soll (in diesem Beispiel driver.py genannt).
4. Verwenden Sie im Treibernotebook mlflow.pyfunc.log_model(python_model=”/path/to/chain.py”, resources=”/path/to/resources.yaml”), um chain.py auszuführen, und protokollieren Sie die Ergebnisse in einem MLflow-Modell. Mit dem Parameter resources werden alle Ressourcen deklariert, die für das Modell erforderlich sind, z. B. ein Vektorsuchindex oder ein Endpunkt, der als Basismodell dient. Weitere Informationen finden Sie unter Beispielressourcendatei für PyFunc.
5. Bereitstellen des Modells. Siehe Bereitstellen eines Agents für generative KI-Anwendung.
6. Wenn die Bereitstellungsumgebung geladen wird, wird chain.py ausgeführt.
7. Wenn eine Dienstanforderung erhalten wird, wird PyFuncClass.predict(...) aufgerufen.

import mlflow

code_path = "/Workspace/Users/first.last/chain.py"
config_path = "/Workspace/Users/first.last/config.yml"

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-augmented Generation?",
        }
    ]
}

# example using PyFunc model

resources_path = "/Workspace/Users/first.last/resources.yml"

with mlflow.start_run():
  logged_chain_info = mlflow.pyfunc.log_model(
    python_model=chain_notebook_path,
    artifact_path="chain",
    input_example=input_example,
    resources=resources_path,
    example_no_conversion=True,
  )

print(f"MLflow Run: {logged_chain_info.run_id}")
print(f"Model URI: {logged_chain_info.model_uri}")

# To verify that the model has been logged correctly, load the chain and call `invoke`:
model = mlflow.pyfunc.load_model(logged_chain_info.model_uri)
model.invoke(example)

Angeben von Ressourcen für den PyFunc-Agent

Sie können Ressourcen angeben, die für das Modell erforderlich sind, z. B. einen Vektorsuchindex und einen Bereitstellungsendpunkt. Für LangChain werden Ressourcen automatisch aufgenommen und zusammen mit dem Modell protokolliert.

Beim Bereitstellen eines pyfunc-Agents müssen Sie alle Ressourcenabhängigkeiten des bereitgestellten Agents manuell hinzufügen. Ein M2M OAuth-Token mit Zugriff auf alle angegebenen Ressourcen im resources-Parameter wird für den bereitgestellten Agent erstellt.

Im Folgenden wird gezeigt, wie Sie die Abhängigkeiten für den Bereitstellungsendpunkt und den Vektorsuchindex hinzufügen, indem Sie sie im resources-Parameter angeben.

 with mlflow.start_run():
   logged_chain_info = mlflow.pyfunc.log_model(
     python_model=chain_notebook_path,
     artifact_path="chain",
     input_example=input_example,
     example_no_conversion=True,
     resources=[
            DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
            DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"),
            DatabricksVectorSearchIndex(index_name="rag.studio_bugbash.databricks_docs_index")
        ]
   )

Sie können auch Ressourcen hinzufügen, indem Sie sie in einer resources.yaml-Datei angeben. Sie können im resources-Parameter auf diesen Dateipfad verweisen. Ein M2M OAuth-Token mit Zugriff auf alle angegebenen Ressourcen in resources.yaml wird für den bereitgestellten Agent erstellt.

Im Folgenden finden Sie eine Beispieldatei resources.yaml, die Modellbereitstellungsendpunkte und einen Vektorsuchindex definiert.

api_version: "1"
databricks:
  vector_search_index:
    - name: "catalog.schema.my_vs_index"
  serving_endpoint:
    - name: databricks-dbrx-instruct
    - name: databricks-bge-large-en

Registrieren der Kette im Unity-Katalog

Bevor Sie die Kette bereitstellen, müssen Sie die Kette im Unity-Katalog registrieren. Wenn Sie die Kette registrieren, wird sie als Modell im Unity-Katalog verpackt, und Sie können Unity-Katalogberechtigungen für die Autorisierung für Ressourcen in der Kette verwenden.

import mlflow

mlflow.set_registry_uri("databricks-uc")

catalog_name = "test_catalog"
schema_name = "schema"
model_name = "chain_name"

model_name = catalog_name + "." + schema_name + "." + model_name
uc_model_info = mlflow.register_model(model_uri=logged_chain_info.model_uri, name=model_name)

Nächste Schritte

• Hinzufügen einer Ablaufverfolgung zu einem KI-Agenten.
• Bereitstellen eines KI-Agenten.