記錄和註冊 AI 代理程式

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

Requirements

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

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

% pip install databricks-sdk

以程式碼為基礎的記錄

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

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

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

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

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

如果未使用 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 管理資源，例如向量搜尋索引或用來提供基礎模型服務的端點。 如需詳細資訊，請參閱 實作自動驗證傳遞。

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 管理資源，例如向量搜尋索引或用來提供基礎模型服務的端點。 如需詳細資訊，請參閱 實作自動驗證傳遞。

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()。

後續步驟

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