记录和注册 AI 代理

项目
02/08/2025

重要

此功能目前以公共预览版提供。

使用 Mosaic AI 代理框架记录 AI 代理。记录代理是开发过程的基础。日志记录将捕获代理代码和配置的“时间点”，以便你可以评估配置的质量。

要求

在记录 AI 代理之前创建 AI 代理。

基于代码的日志记录

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

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

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

记录代理或代理的代码必须位于独立于代理代码的笔记本中。此笔记本称为驱动程序笔记本。有关示例笔记本，请参阅示例笔记本。

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

在日志记录期间，必须定义一个 MLflow 模型签名，该签名指定代理的输入和输出架构。签名验证输入和输出，以确保代理与 AI Playground 和审阅应用等下游工具正确交互。它还指导其他应用程序如何有效地使用代理。

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

以下 LangChain 和 PyFunc 示例使用模型签名推理。

如果想要在日志记录时自行显式定义模型签名，请参阅 MLflow 文档 - 如何使用签名记录模型。

使用 LangChain 的基于代码的日志记录

以下说明和代码示例演示如何使用 LangChain 记录代理。

使用代码创建一个笔记本或 Python 文件。在本示例中，笔记本或文件命名为 agent.py。笔记本或文件必须包含一个 LangChain 代理，此处称为 lc_agent。
在笔记本或文件中包括 mlflow.models.set_model（lc_agent）。
创建新的笔记本作为驱动程序笔记本（在此示例中称为 driver.py）。
在驱动程序笔记本中，使用以下代码运行 agent.py 并将结果记录到 MLflow 模型：
```
mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)
```
resources 参数声明了为代理提供服务所需的 Databricks 托管资源，例如向量搜索索引或用于提供基础模型的服务端点。有关详细信息，请参阅为自动身份验证直通指定资源。
部署模型。请参阅为生成式 AI 应用程序部署代理。
加载服务环境时，将执行 agent.py。
传入服务请求时，将调用 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 记录代理。

使用代码创建一个笔记本或 Python 文件。在本示例中，笔记本或文件命名为 agent.py。笔记本或文件必须包含名为 PyFuncClass的 PyFunc 类。
在笔记本或文件中包含 mlflow.models.set_model(PyFuncClass)。
创建新的笔记本作为驱动程序笔记本（在此示例中称为 driver.py）。
在驱动程序笔记本中，使用以下代码运行 agent.py 并将结果记录到 MLflow 模型：
```
mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)
```
resources 参数声明了为代理提供支持所需的 Databricks 托管资源，例如用于基础模型的向量搜索索引或服务终结点。有关详细信息，请参阅为自动身份验证直通指定资源。
部署模型。请参阅为生成式 AI 应用程序部署代理。
加载服务环境时，将执行 agent.py。
传入服务请求时，将调用 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-mixtral-8x7b-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)

为自动身份验证直通指定资源

AI 代理通常需要向其他资源进行身份验证才能完成任务。例如，代理可能需要访问矢量搜索索引来查询非结构化数据。

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

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

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

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

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-mixtral-8x7b-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"),
    ]
  )

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

注意

如果在使用 mlflow.langchain.log_model(...)记录 LangChain 代理时未指定资源，则 MLflow 会尽力自动推断资源。但是，这可能不会捕获所有依赖项，导致在服务或查询代理时出现授权错误。

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

资源类型	记录资源所需的最低 `mlflow` 版本
矢量搜索索引	需要 `mlflow` 2.13.1 或更高版本
模型服务终结点	需要 `mlflow` 2.13.1 或更高版本
SQL 仓库	需要 `mlflow` 2.16.1 或更高版本
Unity Catalog 函数	需要 `mlflow` 2.16.1 或更高版本
精灵空间	需要 `mlflow` 2.17.1 或更高版本
Unity Catalog 表	需要 `mlflow` 2.18.0 或更高版本

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

后续步骤

向 AI 代理添加跟踪。
部署 AI 代理。

通过