記錄並註冊代理程式(模型服務)

這很重要

對於新代理,Databricks 建議使用 Agent Bricks CLI。 CLI 會用 Agent Bricks 代理伺服器 在 Agent Runtime 部署你的代理程式。 請參閱 在 Azure Databricks 上部署代理程式。 若要將現有代理從模型服務遷移到 Databricks Apps,請參見「將代理從模型服務遷移到 Databricks Apps」。

使用自訂代理程式記錄代理程式。 記錄代理是開發流程的基礎。 日誌記錄會擷取代理程式的程式碼和設定在某一時間點的狀態,讓你能評估其品質。

Requirements

在登錄前先建立一個代理。

Databricks 建議安裝最新版的 databricks-sdk。

% pip install databricks-sdk

以程式碼為基礎的記錄

Databricks 建議在記錄代理時使用 MLflow 的從程式代碼功能的模型。

在此方法中,代理程式的程式代碼會擷取為 Python 檔案,而 Python 環境會擷取為套件清單。 部署代理程式時,會還原 Python 環境,並執行代理程式的程式代碼,將代理程式載入記憶體,以便在呼叫端點時叫用它。

您可以結合此方法與使用預先部署驗證 API,例如 mlflow.models.predict(),以確保代理程式在部署以供服務時可靠地執行。

若要查看程式代碼型記錄的範例,請參閱 ResponsesAgent 撰寫範例筆記本。

在記錄的過程中推斷模型特徵

Note

Databricks 建議使用 ResponsesAgent 介面撰寫代理程式。 如果使用 ResponsesAgent,您可以略過本節;MLflow 會自動推斷代理程式的有效簽章。

如果未使用 ResponsesAgent 介面,您必須在記錄時間使用下列其中一種方法來指定代理程式的 MLflow 模型簽章:

  1. 手動定義簽章
  2. 使用 MLflow 的模型簽章推斷功能,根據您提供的輸入範例自動產生代理程式的簽章。 這種方法比手動定義簽章更方便。

MLflow 模型簽章會驗證輸入和輸出,以確保代理程式與 AI 遊樂場和檢閱應用程式等下游工具正確互動。 它也會引導其他應用程式有效地使用代理程式。

下列 LangChain 和 PyFunc 範例使用模型簽章推斷。

如果您想要在記錄時自行明確定義模型簽章,請參閱 MLflow 檔 - 如何使用簽章記錄模型。

基於程式碼的日誌記錄使用 LangChain

下列指示和程式代碼範例示範如何使用 LangChain 來記錄代理程式。

  1. 使用您的程式碼建立筆記本或 Python 檔案。 在這裡範例中,筆記本或檔案會命名為 agent.py。 筆記本或檔案必須包含 LangChain 代理程式,這裡稱為 lc_agent。

  2. 在筆記本或檔案中包含 mlflow.models.set_model(lc_agent)。

  3. 建立新的筆記本,作為驅動程式筆記本(在此範例稱為 driver.py)。

  4. 在驅動程序筆記本中,使用下列程式代碼來執行 agent.py 並將結果記錄至 MLflow 模型:

    mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)
    

    resources 參數用於宣告為該代理提供服務所需的、由 Databricks 管理的資源,例如 AI 搜尋索引或用於提供基礎模型服務的服務端點。 如需詳細資訊,請參閱 實作自動驗證傳遞。

  5. 部署模型。 參見「 部署代理以支援 AI 應用(模型服務)」。

  6. 載入服務環境時, agent.py 會執行 。

  7. 當有服務請求進入時,會呼叫 lc_agent.invoke(...)。


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)

使用 PyFunc 的基於程式代碼的日誌記錄

下列指示和程式代碼範例示範如何使用 PyFunc 記錄代理程式。

  1. 使用您的程式碼建立筆記本或 Python 檔案。 在這裡範例中,筆記本或檔案會命名為 agent.py。 筆記本或檔案必須包含名為 PyFuncClass的 PyFunc 類別。

  2. 將 mlflow.models.set_model(PyFuncClass) 包含在筆記本或檔案中。

  3. 建立新的筆記本,作為驅動程式筆記本(在此範例稱為 driver.py)。

  4. 在驅動程式筆記本中,使用下列程式代碼來執行 agent.py ,並使用 log_model() 將結果記錄至 MLflow 模型:

    mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)
    

    resources 參數用於宣告為該代理提供服務所需的、由 Databricks 管理的資源,例如 AI 搜尋索引或用於提供基礎模型服務的服務端點。 如需詳細資訊,請參閱 實作自動驗證傳遞。

  5. 部署模型。 參見「 部署代理以支援 AI 應用(模型服務)」。

  6. 載入服務環境時, agent.py 會執行 。

  7. 當有服務請求進入時,會呼叫 PyFuncClass.predict(...)。

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-meta-llama-3-3-70b-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)

Databricks 資源的驗證

代理通常需要與其他資源進行認證才能完成任務。 例如,部署中的代理可能需要存取 AI 搜尋索引來查詢非結構化資料,或存取提示登錄表來載入動態提示。

自動驗證傳遞 和 代表驗證 需要在代理程式日誌記錄期間進行設定。

將代理程序註冊至 Unity 目錄

在部署代理程式之前,您必須將代理程式註冊到 Unity 目錄。 在 Unity Catalog 中註冊代理程式,將其封裝為模型。 因此,您可以使用 Unity 目錄權限來授權代理程式中的資源。

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)

參見 mlflow.register_model()。

其他資源