教程：生成、评估和部署检索代理

2025-05-09

本教程逐步讲解如何构建一个结合使用检索和工具的 AI 代理。

这是一个中级教程，假定你熟悉在 Databricks 上构建代理的基础知识。如果你不熟悉构建代理，请参阅 AI 代理入门。

示例笔记本包括本教程中使用的所有代码。

本教程介绍了生成生成 AI 应用程序时面临的一些核心挑战：

简化常见任务的开发体验，例如创建工具和调试代理执行。
运营挑战，例如：
- 跟踪代理配置
- 以可预测的方式定义输入和输出
- 管理依赖项的版本
- 版本控制和部署
测量和提高代理的质量和可靠性。

为简单起见，本教程使用内存中方法对包含分块 Databricks 文档的数据集启用关键字搜索。有关使用马赛克 AI 矢量搜索来缩放索引和搜索文档的更现实的示例，请参阅 ChatAgent 示例。

示例笔记本

这款独立笔记本是为您设计的，旨在通过提供示例文档库让您快速开始使用马赛克AI代理。它已准备好运行，无需任何配置或数据。

马赛克 AI 代理演示

获取笔记本

创建代理和工具

马赛克 AI 代理框架支持许多不同的创作框架。此示例使用 LangGraph 来说明概念，但这不是 LangGraph 教程。

有关其他受支持框架的示例，请参阅 ChatAgent 示例。

第一步是创建代理。必须指定 LLM 客户端和工具列表。 Python databricks-langchain 包包括与 LangChain 和 LangGraph 兼容的客户端，这些客户端适用于 Databricks 的 LLM 和 Unity Catalog 中注册的工具。

终结点必须是使用 AI 网关的函数调用基础模型 API 或外部模型。请参阅支持的模型。

from databricks_langchain import ChatDatabricks
llm = ChatDatabricks(endpoint="databricks-meta-llama-3-3-70b-instruct")

以下代码定义了一个函数，该函数从模型和一些工具创建代理，讨论此代理代码的内部内容超出了本指南的范围。有关如何生成 LangGraph 代理的详细信息，请参阅 LangGraph 文档。

from typing import Optional, Sequence, Union

from langchain_core.language_models import LanguageModelLike
from langchain_core.runnables import RunnableConfig, RunnableLambda
from langchain_core.tools import BaseTool
from langgraph.graph import END, StateGraph
from langgraph.graph.graph import CompiledGraph
from langgraph.prebuilt.tool_executor import ToolExecutor
from mlflow.langchain.chat_agent_langgraph import ChatAgentState, ChatAgentToolNode


def create_tool_calling_agent(
  model: LanguageModelLike,
  tools: Union[ToolExecutor, Sequence[BaseTool]],
  agent_prompt: Optional[str] = None,
) -> CompiledGraph:
  model = model.bind_tools(tools)

  def routing_logic(state: ChatAgentState):
    last_message = state["messages"][-1]
    if last_message.get("tool_calls"):
      return "continue"
    else:
      return "end"

  if agent_prompt:
    system_message = {"role": "system", "content": agent_prompt}
    preprocessor = RunnableLambda(
      lambda state: [system_message] + state["messages"]
    )
  else:
    preprocessor = RunnableLambda(lambda state: state["messages"])
  model_runnable = preprocessor | model

  def call_model(
    state: ChatAgentState,
    config: RunnableConfig,
  ):
    response = model_runnable.invoke(state, config)

    return {"messages": [response]}

  workflow = StateGraph(ChatAgentState)

  workflow.add_node("agent", RunnableLambda(call_model))
  workflow.add_node("tools", ChatAgentToolNode(tools))

  workflow.set_entry_point("agent")
  workflow.add_conditional_edges(
    "agent",
    routing_logic,
    {
      "continue": "tools",
      "end": END,
    },
  )
  workflow.add_edge("tools", "agent")

  return workflow.compile()

定义代理工具

工具是构建代理的基本概念。它们提供将 LLM 与人工定义代码集成的功能。当提供提示和工具列表时，调用工具的LLM将生成用于调用该工具的参数。有关工具并将其与马赛克 AI 代理配合使用的详细信息，请参阅 AI 代理工具。

第一步是基于 TF-IDF 创建关键字提取工具。此示例使用 scikit-learn 和 Unity 目录工具。

该 databricks-langchain 包提供了一种使用 Unity 目录工具的便捷方法。以下代码演示如何实现和注册关键字提取程序工具。

注释

Databricks 工作区具有一个内置工具， system.ai.python_exec可用于扩展代理，并能够在沙盒执行环境中执行 Python 脚本。其他有用的内置工具包括外部连接和 AI 函数。

from databricks_langchain.uc_ai import (
  DatabricksFunctionClient,
  UCFunctionToolkit,
  set_uc_function_client,
)

uc_client = DatabricksFunctionClient()
set_uc_function_client(client)

# Change this to your catalog and schema
CATALOG = "main"
SCHEMA = "my_schema"


def tfidf_keywords(text: str) -> list[str]:
  """
  Extracts keywords from the provided text using TF-IDF.

  Args:
    text (string): Input text.
  Returns:
    list[str]: List of extracted keywords in ascending order of importance.
  """
  from sklearn.feature_extraction.text import TfidfVectorizer

  def keywords(text, top_n=5):
    vec = TfidfVectorizer(stop_words="english")
    tfidf = vec.fit_transform([text])  # Convert text to TF-IDF matrix
    indices = tfidf.toarray().argsort()[0, -top_n:]  # Get indices of top N words
    return [vec.get_feature_names_out()[i] for i in indices]

  return keywords(text)


# Create the function in the Unity Catalog catalog and schema specified
# When you use `.create_python_function`, the provided function's metadata
# (docstring, parameters, return type) are used to create a tool in the specified catalog and schema.
function_info = uc_client.create_python_function(
  func=tfidf_keywords,
  catalog=CATALOG,
  schema=SCHEMA,
  replace=True,  # Set to True to overwrite if the function already exists
)

print(function_info)

下面是上述代码的说明：

创建一个客户端，该客户端在 Databricks 工作区中使用 Unity 目录作为“注册表”来创建和发现工具。
定义执行 TF-IDF 关键字提取的 Python 函数。
将 Python 函数注册为 Unity 目录函数。

此工作流解决了几个常见问题。你现在有一个中心注册表，可用于管理与 Unity 目录中的其他对象一样的工具。例如，如果公司有一种计算内部返回率的标准方法，则可以将其定义为 Unity 目录中的函数，并授予对具有该角色的 FinancialAnalyst 所有用户或代理的访问权限。

若要使此工具可由 LangChain 代理使用，请使用 UCFunctionToolkit ，它创建一组工具以提供给 LLM 以供选择：

# Use ".*" here to specify all the tools in the schema, or
# explicitly list functions by name
# uc_tool_names = [f"{CATALOG}.{SCHEMA}.*"]
uc_tool_names = [f"{CATALOG}.{SCHEMA}.tfidf_keywords"]
uc_toolkit = UCFunctionToolkit(function_names=uc_tool_names)

以下代码演示如何测试该工具：

uc_toolkit.tools[0].invoke({ "text": "The quick brown fox jumped over the lazy brown dog." })

以下代码创建使用关键字提取工具的代理。

import mlflow
mlflow.langchain.autolog()

agent = create_tool_calling_agent(llm, tools=[*uc_toolkit.tools])

agent.invoke({"messages": [{"role": "user", "content":"What are the keywords for the sentence: 'the quick brown fox jumped over the lazy brown dog'?"}]})

在生成的跟踪中，可以看到 LLM 选择了该工具。

笔记本中的 MLflow 跟踪输出，其中显示了工具选择。

使用日志调试代理

MLflow 跟踪是用于调试和观察生成 AI 应用程序（包括代理）的强大工具。它通过跨度捕获详细的执行信息，这些范围封装特定的代码段和记录输入、输出和计时数据。

对 LangChain 等热门库启用自动跟踪，请使用 mlflow.langchain.autolog()。还可以用于 mlflow.start_span() 自定义跟踪。例如，可以添加自定义数据值字段或标记，以便观察。在该范围上下文中运行的代码与定义的字段相关联。在这个内存中的 TF-IDF 示例中，为其命名并指定一个范围类型。

若要了解有关跟踪的详细信息，请参阅 MLflow 跟踪的代理可观测性。

以下示例使用简单的内存中 TF-IDF 索引创建检索器工具。它展示了工具执行的自动记录，以及自定义跨度跟踪以增强可观测性。

from sklearn.feature_extraction.text import TfidfVectorizer
import mlflow
from langchain_core.tools import tool


documents = parsed_docs_df
doc_vectorizer = TfidfVectorizer(stop_words="english")
tfidf_matrix = doc_vectorizer.fit_transform(documents["content"])


@tool
def find_relevant_documents(query, top_n=5):
  """gets relevant documents for the query"""
  with mlflow.start_span(name="LittleIndex", span_type="RETRIEVER") as retriever_span:
    retriever_span.set_inputs({"query": query})
    retriever_span.set_attributes({"top_n": top_n})

    query_tfidf = doc_vectorizer.transform([query])
    similarities = (tfidf_matrix @ query_tfidf.T).toarray().flatten()
    ranked_docs = sorted(enumerate(similarities), key=lambda x: x[1], reverse=True)

    result = []
    for idx, score in ranked_docs[:top_n]:
      row = documents.iloc[idx]
      content = row["content"]
      doc_entry = {
        "page_content": content,
        "metadata": {
          "doc_uri": row["doc_uri"],
          "score": score,
        },
      }
      result.append(doc_entry)

    retriever_span.set_outputs(result)
    return result

此代码使用为检索器工具保留的特殊范围类型 RETRIEVER。其他马赛克 AI 代理功能（如 AI Playground、审阅 UI 和评估）使用 RETRIEVER 范围类型显示检索结果。

检索器工具要求你指定其架构，以确保与下游 Databricks 功能兼容。有关详细信息 mlflow.models.set_retriever_schema，请参阅指定自定义检索器架构。

import mlflow
from mlflow.models import set_retriever_schema

uc_toolkit = UCFunctionToolkit(function_names=[f"{CATALOG}.{SCHEMA}.*"])

graph = create_tool_calling_agent(llm, tools=[*uc_toolkit.tools, find_relevant_documents])

mlflow.langchain.autolog()
set_retriever_schema(
  primary_key="chunk_id",
  text_column="chunk_text",
  doc_uri="doc_uri",
  other_columns=["title"],
)

graph.invoke(input = {"messages": [("user", "How do the docs say I use llm judges on databricks?")]})

显示元数据的检索结果。

定义代理

下一步是评估代理并将其准备用于部署。概括而言，这涉及以下内容：

使用签名定义代理的可预测 API。
添加模型配置，以便于配置参数。
记录模型及其依赖项以创建可重现环境，并允许您配置身份验证以访问其他服务。

MLflow ChatAgent 接口简化了代理输入和输出的定义。若要使用它，请将您的代理定义为ChatAgent的子类，使用predict函数实现非流式推理，并使用predict_stream函数进行流式推理。

ChatAgent 与你选择的代理创作框架无关，允许你轻松测试和使用不同的框架和代理实现 - 唯一的要求是实现 predict 和 predict_stream 接口。

使用 ChatAgent 创建代理可提供许多好处，包括：

输出流支持
全面的工具呼叫消息历史记录：返回多个消息，包括中间工具呼叫消息，以提高质量和聊天管理。
多代理系统支持
Databricks 功能集成： 与 AI Playground、代理评估和代理监视的现式兼容性。
类型化创作接口：使用类型化的 Python 类编写代理代码，受益于 IDE 和笔记本自动完成。

有关创作ChatAgent的信息，请参阅使用ChatAgent创作代理。

from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import (
  ChatAgentChunk,
  ChatAgentMessage,
  ChatAgentResponse,
  ChatContext,
)
from typing import Any, Optional


class DocsAgent(ChatAgent):
  def __init__(self, agent):
    self.agent = agent
    set_retriever_schema(
      primary_key="chunk_id",
      text_column="chunk_text",
      doc_uri="doc_uri",
      other_columns=["title"],
    )

  def predict(
    self,
    messages: list[ChatAgentMessage],
    context: Optional[ChatContext] = None,
    custom_inputs: Optional[dict[str, Any]] = None,
  ) -> ChatAgentResponse:
    # ChatAgent has a built-in helper method to help convert framework-specific messages, like langchain BaseMessage to a python dictionary
    request = {"messages": self._convert_messages_to_dict(messages)}

    output = agent.invoke(request)
    # Here 'output' is already a ChatAgentResponse, but to make the ChatAgent signature explicit for this demonstration, the code returns a new instance
    return ChatAgentResponse(**output)

以下代码演示如何使用 ChatAgent.

AGENT = DocsAgent(agent=agent)
AGENT.predict(
  {
    "messages": [
      {"role": "user", "content": "What is DLT in Databricks?"},
    ]
  }
)

使用参数配置代理

使用 Agent Framework，可以使用参数控制代理执行。这意味着你可以快速测试不同的代理配置，例如切换 LLM 终结点或尝试不同的工具，而无需更改基础代码。

以下代码创建一个配置字典，用于在初始化模型时设置代理参数。

有关参数化代理的更多详细信息，请参阅参数化代理代码，以便跨环境进行部署。

)

from mlflow.models import ModelConfig

baseline_config = {
  "endpoint_name": "databricks-meta-llama-3-3-70b-instruct",
  "temperature": 0.01,
  "max_tokens": 1000,
  "system_prompt": """You are a helpful assistant that answers questions about Databricks. Questions unrelated to Databricks are irrelevant.


  You answer questions using a set of tools. If needed, you ask the user follow-up questions to clarify their request.
  """,
  "tool_list": ["catalog.schema.*"],
}


class DocsAgent(ChatAgent):
  def __init__(self):
    self.config = ModelConfig(development_config=baseline_config)
    self.agent = self._build_agent_from_config()


def _build_agent_from_config(self):
  temperature = config.get("temperature", 0.01)
  max_tokens = config.get("max_tokens", 1000)
  system_prompt = config.get("system_prompt", """You are a helpful assistant.
    You answer questions using a set of tools. If needed you ask the user follow-up questions to clarify their request.""")
  llm_endpoint_name = config.get("endpoint_name", "databricks-meta-llama-3-3-70b-instruct")
  tool_list = config.get("tool_list", [])

  llm = ChatDatabricks(endpoint=llm_endpoint_name, temperature=temperature, max_tokens=max_tokens)
  toolkit = UCFunctionToolkit(function_names=tool_list)
  agent = create_tool_calling_agent(llm, tools=[*toolkit.tools, find_relevant_documents], prompt=system_prompt)

  return agent

记录代理信息

定义代理后，现在可以日志化该代理了。在 MLflow 中，日志记录代理意味着保存代理的配置（包括依赖项），以便可用于评估和部署。

注释

在笔记本中开发代理时，MLflow 会从笔记本环境中推断代理的依赖项。

若要从笔记本中记录代理，可以在单个单元格中编写定义模型的所有代码，然后使用 %%writefile magic 命令将代理的定义保存到文件中：

%%writefile agent.py
...
<Code that defines the agent>

如果代理需要访问外部资源（例如 Unity 目录）才能执行关键字提取工具，则必须为代理配置身份验证，以便它可以在部署资源时访问资源。

若要简化 Databricks 资源的身份验证，请启用自动身份验证直通：

from mlflow.models.resources import DatabricksFunction, DatabricksServingEndpoint


resources = [
  DatabricksServingEndpoint(endpoint_name=LLM_ENDPOINT_NAME),
  DatabricksFunction(function_name=tool.uc_function_name),
]


with mlflow.start_run():
  logged_agent_info = mlflow.pyfunc.log_model(
    artifact_path="agent",
    python_model="agent.py",
    pip_requirements=[
      "mlflow",
      "langchain",
      "langgraph",
      "databricks-langchain",
      "unitycatalog-langchain[databricks]",
      "pydantic",
    ],
    resources=resources,
  )

若要了解有关日志记录代理的详细信息，请参阅基于代码的日志记录。

评估代理

下一步是评估代理以查看其执行方式。代理评估具有挑战性，并提出了许多问题，例如：

评估质量的正确指标是什么？如何信任这些指标的输出？
我需要评估许多想法，我怎么做？
- 可以快速运行评估，以便我的大部分时间不花在等待上吗？
- 快速比较我的代理的不同版本在质量、成本和延迟上的表现。
如何快速确定任何质量问题的根本原因？

作为数据科学家或开发人员，你可能不是实际的主题专家。本部分的其余部分介绍有助于定义良好输出的代理评估工具。

创建评估集

若要定义代理的质量含义，请使用指标来衡量代理对评估集的性能。请参阅 “定义质量”：评估集。

使用代理评估功能，您可以通过运行评估来创建综合评估集，并在评估过程中度量质量。这个想法是从事实开始，比如一组文档，然后通过这些事实生成一组问题，以实现“反向推理”。可以通过提供一些准则来条件生成的问题：

from databricks.agents.evals import generate_evals_df
import pandas as pd


databricks_docs_url = "https://raw.githubusercontent.com/databricks/genai-cookbook/refs/heads/main/quick_start_demo/chunked_databricks_docs_filtered.jsonl"
parsed_docs_df = pd.read_json(databricks_docs_url, lines=True)


agent_description = f"""
The agent is a RAG chatbot that answers questions about Databricks. Questions unrelated to Databricks are irrelevant.
"""
question_guidelines = f"""
# User personas
- A developer who is new to the Databricks platform
- An experienced, highly technical Data Scientist or Data Engineer


# Example questions
- what API lets me parallelize operations over rows of a delta table?
- Which cluster settings will give me the best performance when using Spark?


# Additional Guidelines
- Questions should be succinct, and human-like
"""


num_evals = 25
evals = generate_evals_df(
  docs=parsed_docs_df[
    :500
  ],  # Pass your docs. They should be in a Pandas or Spark DataFrame with columns `content STRING` and `doc_uri STRING`.
  num_evals=num_evals,  # How many synthetic evaluations to generate
  agent_description=agent_description,
  question_guidelines=question_guidelines,
)

生成的评估包括：

类似于前面提到的请求字段 ChatAgentRequest ：

{"messages":[{"content":"What command must be run at the start of your workload to explicitly target the Workspace Model Registry if your workspace default catalog is in Unity Catalog and you use Databricks Runtime 13.3 LTS or above?","role":"user"}]}

“预期检索的内容”的列表。检索器架构是使用 content 和 doc_uri 字段定义的。

[{"content":"If your workspace's [default catalog](https://docs.databricks.com/data-governance/unity-catalog/create-catalogs.html#view-the-current-default-catalog) is in Unity Catalog (rather than `hive_metastore`) and you are running a cluster using Databricks Runtime 13.3 LTS or above, models are automatically created in and loaded from the workspace default catalog, with no configuration required. To use the Workspace Model Registry in this case, you must explicitly target it by running `import mlflow; mlflow.set_registry_uri(\"databricks\")` at the start of your workload.","doc_uri":"https://docs.databricks.com/machine-learning/manage-model-lifecycle/workspace-model-registry.html"}]

预期事实列表。比较两个响应时，很难找到它们之间的小差异。期望的事实将正确答案与部分正确答案及错误答案区分开，并改善 AI 评审的质量以及与代理相关的人员的经验。
```
["The command must import the MLflow module.","The command must set the registry URI to \"databricks\"."]
```
source_id字段在此处为SYNTHETIC_FROM_DOC。生成更完整的评估集时，示例将来自不同的源，因此此字段区分它们。

若要详细了解如何创建评估集，请参阅合成评估集。

使用 LLM 法官评估代理

手动评估代理在众多生成示例上的性能难以实现可扩展性。大规模地，使用 LLM 作为法官是一个更合理的解决方案。若要使用使用代理评估时可用的内置法官，请使用以下代码：

with mlflow.start_run(run_name="my_agent"):
  eval_results = mlflow.evaluate(
    data=evals,  # Your evaluation set
    model=model_info.model_uri,  # Logged agent from above
    model_type="databricks-agent",  # activate Mosaic AI Agent Evaluation
)

MLflow 试验 - 评估结果。

简单智能体整体得分为 68%。根据所使用的配置，结果可能在此处有所不同。运行一个试验来比较三个不同的 LLM 的成本和质量与更改配置和重新评估一样简单。

请考虑更改模型配置以使用不同的 LLM、系统提示或温度设置。

这些评估员可以被定制，以遵循与人类专家评价响应时相同的准则。有关 LLM 法官的详细信息，请参阅内置 AI 法官。

您可以通过使用代理评估工具，来利用自定义指标定制衡量特定代理的质量的方法。可以将评估视为集成测试，将单个指标视为单元测试。以下示例使用布尔指标来检查代理是否对给定请求使用了关键字提取和检索器：

from databricks.agents.evals import metric

@metric
def uses_keywords_and_retriever(request, trace):
  retriever_spans = trace.search_spans(span_type="RETRIEVER")
  keyword_tool_spans = trace.search_spans(name=f"{CATALOG}__{SCHEMA}__tfidf_keywords")
  return len(keyword_tool_spans) > 0 and len(retriever_spans) > 0


# same evaluate as above, with the addition of 'extra_metrics'
with mlflow.start_run(run_name="my_agent"):
  eval_results = mlflow.evaluate(
    data=evals,  # Your evaluation set
    model=model_info.model_uri,  # Logged agent from above
    model_type="databricks-agent",  # activate Mosaic AI Agent Evaluation,
    extra_metrics=[uses_keywords_and_retriever],
  )

请注意，代理永远不会使用关键字提取。如何解决此问题？

显示自定义指标输出的评估结果。

部署和监控代理

准备好开始与真实用户测试代理时，Agent Framework 提供了一个生产就绪的解决方案，用于在马赛克 AI 模型服务上为代理提供服务。

将代理部署到 Model Services 具有以下优势：

模型服务管理自动缩放、日志记录、版本控制和访问控制，使你能够专注于开发质量代理。
主题专家可以使用评审应用与代理交互，并提供反馈，这些反馈可以合并到监视和评估中。
您可以通过对实时流量进行评估来监视代理程序。尽管用户流量不包括真实数据，但 LLM 评委（以及你创建的自定义指标）进行无监督评估。

以下代码将代理部署到服务终结点。有关详细信息，请参阅为生成式 AI 应用程序部署代理。

from databricks import agents
import mlflow

# Connect to the Unity Catalog model registry
mlflow.set_registry_uri("databricks-uc")

# Configure UC model ___location
UC_MODEL_NAME = f"{CATALOG}.{SCHEMA}.getting_started_agent"
# REPLACE WITH UC CATALOG/SCHEMA THAT YOU HAVE `CREATE MODEL` permissions in

# Register to Unity Catalog
uc_registered_model_info = mlflow.register_model(
  model_uri=model_info.model_uri, name=UC_MODEL_NAME
)
# Deploy to enable the review app and create an API endpoint
deployment_info = agents.deploy(
  model_name=UC_MODEL_NAME, model_version=uc_registered_model_info.version
)

通过