記錄和註冊 AI 代理程式

使用 Mosaic AI 代理程式架構記錄 AI 代理程式。 記錄代理程式是開發流程的基礎。 記錄會擷取代理程式程式碼與組態的「時間點」，讓您可以評估組態品質。

需求

在記錄之前建立 AI 代理程式。

代碼型記錄與序列化型記錄

您可以使用代碼型 MLflow 記錄或序列化型 MLflow 記錄。 Databricks 建議您使用代碼型記錄。

代碼型 MLflow 記錄：鏈的代碼將被擷取為 Python 檔案。 Python 環境會被擷取為套件清單。 在部署鏈時，將恢復 Python 環境，並執行鏈代碼來將鏈載入記憶體，以便在呼叫端點時調用鏈。

序列化型 MLflow 記錄：Python 環境鏈的代碼與當前狀態通常使用 pickle 或 joblib 等程式庫序列化至磁碟。 在部署鏈時，會還原 Python 環境，並將序列化物件載入記憶體，以便在呼叫端點時調用。

圖表顯示每種方法的優點與缺點。

針對程式碼型記錄，記錄代理程式或鏈的程式碼必須位於與鏈程式碼不同的筆記本。 此筆記本稱為驅動程式筆記本。 如需範例筆記本，請參閱範例筆記本。

使用 LangChain 的程式碼型記錄

1. 使用您的程式碼建立筆記本或 Python 檔案。 針對此範例目的，筆記本或檔案的名稱為 chain.py。 筆記本或檔案必須包含 LangChain 鏈，這裡稱為 lc_chain。
2. mlflow.models.set_model(lc_chain) 包含在筆記本或檔案。
3. 建立新的筆記本，作為驅動程式筆記本（在此範例稱為 driver.py）。
4. 在驅動程式筆記本，使用 mlflow.lang_chain.log_model(lc_model=”/path/to/chain.py”) 來執行 chain.py 並將結果記錄至 MLflow 模型。
5. 部署模型。 請參閱為生成式 AI 應用程式部署代理程式。 代理程式的部署可能取決於其他 Databricks 資源，例如向量搜尋索引與模型服務端點。 對於 LangChain 代理程式：
   • MLflow log_model 會推斷鏈所需的相依性，並將其記錄到記錄模型成品的 MLmodel 檔案。
   • 在部署期間，databricks.agents.deploy 會自動建立存取這些推斷的資源相依性所需的 M2M OAuth 權杖。
6. 當載入服務環境時，會執行 chain.py。
7. 當有服務請求進入時，會呼叫 lc_chain.invoke(...)。

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)

使用 PyFunc 的程式碼型記錄

1. 使用您的程式碼建立筆記本或 Python 檔案。 針對此範例目的，筆記本或檔案的名稱為 chain.py。 筆記本或檔案必須包含 PyFunc 類別，此處稱為 PyFuncClass。
2. mlflow.models.set_model(PyFuncClass) 包含在筆記本或檔案。
3. 建立新的筆記本，作為驅動程式筆記本（在此範例稱為 driver.py）。
4. 在驅動程式筆記本，使用 mlflow.pyfunc.log_model(python_model=”/path/to/chain.py”, resources=”/path/to/resources.yaml”) 來執行 chain.py 並將結果記錄至 MLflow 模型。 resources 參數會宣告為模型提供服務所需的任何資源，例如為基礎模型提供服務的向量搜尋索引或服務端點。 請參閱 PyFunc 的資源檔案範例。
5. 部署模型。 請參閱為生成式 AI 應用程式部署代理程式。
6. 當載入服務環境時，會執行 chain.py。
7. 當有服務請求進入時，會呼叫 PyFuncClass.predict(...)。

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)

指定 PyFunc 代理程式的資源

您可以指定服務模型所需的資源，例如向量搜尋索引與服務端點。 針對 LangChain，資源會自動拾取並與模型一起記錄。

部署 pyfunc 類型代理程式時，您必須手動新增已部署代理程式的任何資源相依性。 已建立可存取 resources 參數所有指定資源的 M2M OAuth 權杖，並提供給已部署的代理程式。

以下說明如何在 resources 參數指定服務端點與向量搜尋索引相依性。

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

您也可以透過在 resources.yaml 檔案指定資源來新增資源。 您可以在 resources 參數引用該檔案路徑。 已建立可存取 resources.yaml 指定資源的 M2M OAuth 權杖，並提供給已部署的代理程式。

以下是定義模型服務端點與向量搜尋索引的範例 resources.yaml 檔案。

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

將鏈註冊至 Unity 目錄

在部署鏈之前，必須將鏈註冊至 Unity 目錄。 在註冊鏈時，其會被封包為 Unity 目錄的模型，您可以使用 Unity 目錄權限授權鏈的資源。

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)

下一步

• 將追蹤新增至 AI 代理程式。
• 部署 AI 代理程式。