记录和注册 AI 代理

使用马赛克 AI 代理框架记录 AI 代理。 记录代理活动是开发过程的基础。 日志记录用来捕捉代理代码和配置在特定时刻的状态，以便评估配置的质量。

要求

创建 AI 代理后再记录它。

Databricks 建议安装最新版本的 databricks-sdk。

% pip install databricks-sdk

基于代码的日志记录

Databricks 建议在记录代理时使用 MLflow 的 代码功能中的模型功能。

在此方法中，代理的代码作为 Python 文件捕获，Python 环境将捕获为包列表。 部署代理后，将还原 Python 环境，并执行代理的代码将代理加载到内存中，以便在调用终结点时调用代理。

可以将此方法与使用预部署验证 API（如 mlflow.models.predict（） 结合使用，以确保代理在部署服务时可靠运行。

若要查看基于代码的日志记录示例，请参阅 ChatAgent 创作示例笔记本。

在日志记录期间推断模型的签名

如果未使用 ChatAgent 接口，则必须使用以下方法之一在日志记录时指定代理的 MLflow 模型签名：

1. 手动定义签名
2. 使用 MLflow 的模型签名推理功能根据提供的输入示例自动生成代理的签名。 此方法比手动定义签名更方便。

MLflow 模型签名验证输入和输出，以确保代理与 AI Playground 和审阅应用等下游工具正确交互。 它还指导其他应用程序如何有效地使用代理。

以下 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 托管资源，例如向量搜索索引或用于提供基础模型的服务端点。 有关详细信息，请参阅 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 托管资源，例如向量搜索索引或用于提供基础模型的服务端点。 有关详细信息，请参阅 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 代理通常必须向其他资源进行身份验证才能完成任务。 例如，代理可能需要访问矢量搜索索引来查询非结构化数据。

如 依赖资源的身份验证中所述，模型服务支持在部署代理时对 Databricks 托管的资源和外部资源进行身份验证。

模型服务支持 Databricks 托管资源的两种不同类型的身份验证：

1. 系统身份验证：允许代理服务主体访问在代理日志记录时指定的任何依赖资源。 这可用于访问共享或非敏感资源，例如包含公共文档的矢量搜索索引
2. [Beta] 代理用户身份验证：允许代理使用终端用户凭据访问 Databricks 资源。 这在代理需要访问敏感数据或查询远程 API 以根据每个用户采取行动的情境中非常有用。

为自动身份验证直通指定资源（系统身份验证）

对于最常见的 Databricks 资源类型，Databricks 支持并建议在日志记录期间为代理声明资源依赖项。 这将在部署代理时启用自动身份验证直通 - Databricks 会自动预配、轮换和管理生存期较短的凭据，以便从代理终结点内安全地访问这些资源依赖项。

若要启用自动身份验证传递，请使用 resources API 的参数 log_model() 指定依赖资源，如以下代码所示。

import mlflow
from mlflow.models.resources import (
  DatabricksVectorSearchIndex,
  DatabricksServingEndpoint,
  DatabricksSQLWarehouse,
  DatabricksFunction,
  DatabricksGenieSpace,
  DatabricksTable,
  DatabricksUCConnection
)

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-meta-llama-3-3-70b-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"),
      DatabricksUCConnection(connection_name="your_connection_name"),
    ]
  )

Databricks 建议手动为所有代理版本指定 resources。

下表列出了支持自动身份验证直通的 Databricks 资源，以及记录资源所需的最低 mlflow 版本。

代表用户身份验证

使用代表用户身份验证记录代码时，请指定在代理代码中以最终用户身份执行操作时所需的最小 Databricks API 范围集。 这可确保代理在部署时具有代表最终用户执行操作的最低特权访问权限，从而通过防止未经授权的操作来提高安全性，并最大程度地降低令牌滥用的风险。

下面是访问多个常见 Databricks 资源类型所需的范围列表。

若要启用代表用户身份验证，请将 MLflow AuthPolicy 传递给 log_model()，如以下示例所示。 MLflow AuthPolicy 有两个组件：

• system_auth_policy：指定用于系统身份验证的资源。 通常，代理将为共享资源使用系统身份验证（例如，查询服务终结点的模型），同时结合代表用户身份验证来访问敏感资源或 API
• user_auth_policy：为代表用户身份验证指定代理所需的 API 范围

from mlflow.models.resources import DatabricksServingEndpoint
from mlflow.models.auth_policy import SystemAuthPolicy, UserAuthPolicy, AuthPolicy

resources = [
    DatabricksServingEndpoint(endpoint_name="databricks-meta-llama-3-3-70b-instruct")
]
# Specify resources here for system authentication
system_auth_policy = SystemAuthPolicy(resources=resources)

# Specify the minimal set of API scopes needed for on-behalf-of-user authentication
	# When deployed, the agent can access Databricks resources and APIs
	# on behalf of the end user, but only via REST APIs that are covered by the list of
	# scopes below

user_auth_policy = UserAuthPolicy(
    api_scopes=[
        "serving.serving-endpoints",
        "vectorsearch.vector-search-endpoints",
        "vectorsearch.vector-search-indexes",
    ]
)

with mlflow.start_run():
    logged_agent_info = mlflow.pyfunc.log_model(
        ...
        # Instead of passing `resources` (which only supports system authentication),
	        # pass an auth_policy to log_model to enable both system authentication and
	        # on-behalf-of-user authentication
        auth_policy=AuthPolicy(
            system_auth_policy=system_auth_policy,
            user_auth_policy=user_auth_policy
        )
    )

OpenAI 客户端的自动身份验证

如果代理使用 OpenAI 客户端，请使用 Databricks SDK 在部署期间自动进行身份验证。 Databricks SDK 提供了一个封装器，用于构造具有授权自动配置的 OpenAI 客户端。 在笔记本中运行以下命令：

% pip install databricks-sdk[openai]

from databricks.sdk import WorkspaceClient
def openai_client(self):
  w = WorkspaceClient()
  return w.serving_endpoints.get_open_ai_client()

然后，将模型服务终结点指定为 resources 的一部分，以便在部署时自动进行身份验证。

将代理注册到 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 代理。