AI 代理入门

2025-05-14

使用马赛克 AI 代理框架生成第一个 AI 代理。在本教程中，你将：

使用代理框架创作代理。
将工具添加到代理。
将代理部署到 Databricks 模型服务端点。

有关代理和其他第一代 AI 应用的概念性简介，请参阅什么是 Gen AI 应用？

要求

工作区必须启用以下功能：

Unity Catalog
Mosaic AI 代理框架
基础模型（按令牌付费、预配的吞吐量或外部模型）。请参阅具有有限区域可用性的功能

示例笔记本

此笔记本包含创作和部署第一个 AI 代理所需的所有代码。将笔记本导入 Azure Databricks 工作区以运行它。

马赛克 AI 代理演示

获取笔记本

定义代理

AI 代理人包括以下组成部分：

可以推理和做出决策的大型语言模型（LLM）
大语言模型可以使用的工具，不仅可以用于生成文本，还可以执行更多操作，例如运行 Python 代码或获取数据。

在 Databricks 笔记本中运行以下代码以定义简单的工具调用代理：

安装所需的 Python 包：
```
%pip install -U -qqqq mlflow databricks-openai databricks-agents
dbutils.library.restartPython()
```
- mlflow：用于代理开发和代理跟踪。
- databricks-openai：用于连接到由 Databricks 托管的 LLM。
- databricks-agent：用于打包和部署代理。

定义代理。此代码片段执行以下作：

使用 OpenAI 客户端连接到 Databricks 模型服务终结点。
使用 autolog(). 启用 MLflow 跟踪。这会添加检测工具，以便在提交查询时查看代理的工作。
将 system.ai.python_exec 工具添加到代理。此内置 Unity 目录函数允许代理运行 Python 代码。
定义一个函数，该函数使用提示查询 LLM 并处理响应。

import mlflow
import json
from databricks.sdk import WorkspaceClient
from databricks_openai import UCFunctionToolkit, DatabricksFunctionClient

# Get an OpenAI client configured to connect to Databricks model serving endpoints
# Use this client to query the LLM
openai_client = WorkspaceClient().serving_endpoints.get_open_ai_client()

# Enable automatic tracing for easier debugging
mlflow.openai.autolog()

# Load Databricks built-in tools (Python code interpreter)
client = DatabricksFunctionClient()
builtin_tools = UCFunctionToolkit(function_names=["system.ai.python_exec"], client=client).tools
for tool in builtin_tools:
  del tool["function"]["strict"]


def call_tool(tool_name, parameters):
  if tool_name == "system__ai__python_exec":
    return DatabricksFunctionClient().execute_function("system.ai.python_exec", parameters=parameters)
  raise ValueError(f"Unknown tool: {tool_name}")

def run_agent(prompt):
  """
  Send a user prompt to the LLM and return a list of LLM response messages
  The LLM is allowed to call the code interpreter tool, if needed, to respond to the user
  """
  result_msgs = []
  response = openai_client.chat.completions.create(
    model="databricks-claude-3-7-sonnet",
    messages=[{"role": "user", "content": prompt}],
    tools=builtin_tools,
  )
  msg = response.choices[0].message
  result_msgs.append(msg.to_dict())

  # If the model executed a tool, call it
  if msg.tool_calls:
    call = msg.tool_calls[0]
    tool_result = call_tool(call.function.name, json.loads(call.function.arguments))
    result_msgs.append({"role": "tool", "content": tool_result.value, "name": call.function.name, "tool_call_id": call.id})
  return result_msgs

测试代理

使用需要运行 Python 代码的提示对其进行查询来测试代理：

answer = run_agent("What is the 100th fibonacci number?")
for message in answer:
  print(f'{message["role"]}: {message["content"]}')

除了 LLM 的输出，你还将直接在笔记本中看到详细的跟踪信息。这些跟踪有助于调试慢速或失败的代理调用。这些跟踪是使用自动添加的 mlflow.openai.autolog() 。

部署代理

拥有代理后，可以打包并将其部署到 Databricks 服务终结点。开始收集已部署代理的反馈，方法是与他人共享该代理，并使用内置聊天 UI 与之聊天。

准备代理代码进行部署

若要为部署准备代理代码，请使用 MLflow 的 ChatAgent 接口包装它。此 ChatAgent 接口是打包代理以在 Azure Databricks 上进行部署的建议方法。

若要实现 ChatAgent 接口，必须定义一个 predict() 函数，该函数将用户的消息发送到代理，收集代理的响应，然后以格式返回它 ChatAgentResponses 。

import uuid
from typing import Any, Optional

from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import ChatAgentMessage, ChatAgentResponse, ChatContext

class QuickstartAgent(ChatAgent):
  def predict(
    self,
    messages: list[ChatAgentMessage],
    context: Optional[ChatContext] = None,
    custom_inputs: Optional[dict[str, Any]] = None,
  ) -> ChatAgentResponse:
    # 1. Extract the last user prompt from the input messages
    prompt = messages[-1].content

    # 2. Call run_agent to get back a list of response messages
    raw_msgs = run_agent(prompt)

    # 3. Map each response message into a ChatAgentMessage and return
    # the response
    out = []
    for m in raw_msgs:
      out.append(ChatAgentMessage(
        id=uuid.uuid4().hex,
        **m
      ))

    return ChatAgentResponse(messages=out)

将以下代码添加到笔记本来测试 ChatAgent 类：

AGENT = QuickstartAgent()
for response_message in AGENT.predict({"messages": [{"role": "user", "content": "What's the 100th fibonacci number?"}]}).messages:
  print(f"role: {response_message.role}, content: {response_message.content}")

将所有代理代码合并到单个文件中，以便可以记录和部署它。

将所有代理代码合并到一个笔记本单元中。
在单元格顶部，添加 %%writefile quickstart_agent.py magic 命令，将代理保存到文件中。
在单元格底部，使用代理对象进行调用 mlflow.models.set_model() 。这会告知 MLflow 在提供预测时要使用的代理对象。此步骤正确地配置我们的代理程序代码的入口。

笔记本单元格应如下所示：

%%writefile quickstart_agent.py

import json
import uuid
from databricks.sdk import WorkspaceClient
from databricks_openai import UCFunctionToolkit, DatabricksFunctionClient
from typing import Any, Optional

import mlflow
from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import ChatAgentMessage, ChatAgentResponse, ChatContext

# Add an mlflow.openai.autolog() call to capture traces in the serving endpoint

# Get an OpenAI client configured to talk to Databricks model serving endpoints
# We'll use this to query an LLM in our agent
openai_client = WorkspaceClient().serving_endpoints.get_open_ai_client()

# Load Databricks built-in tools (a stateless Python code interpreter tool)
client = DatabricksFunctionClient()
builtin_tools = UCFunctionToolkit(function_names=["system.ai.python_exec"], client=client).tools
for tool in builtin_tools:
  del tool["function"]["strict"]


def call_tool(tool_name, parameters):
  if tool_name == "system__ai__python_exec":
    return DatabricksFunctionClient().execute_function("system.ai.python_exec", parameters=parameters)
  raise ValueError(f"Unknown tool: {tool_name}")

def run_agent(prompt):
  """
  Send a user prompt to the LLM, and return a list of LLM response messages
  The LLM is allowed to call the code interpreter tool if needed, to respond to the user
  """
  result_msgs = []
  response = openai_client.chat.completions.create(
    model="databricks-claude-3-7-sonnet",
    messages=[{"role": "user", "content": prompt}],
    tools=builtin_tools,
  )
  msg = response.choices[0].message
  result_msgs.append(msg.to_dict())

  # If the model executed a tool, call it
  if msg.tool_calls:
    call = msg.tool_calls[0]
    tool_result = call_tool(call.function.name, json.loads(call.function.arguments))
    result_msgs.append({"role": "tool", "content": tool_result.value, "name": call.function.name, "tool_call_id": call.id})
  return result_msgs

class QuickstartAgent(ChatAgent):
  def predict(
    self,
    messages: list[ChatAgentMessage],
    context: Optional[ChatContext] = None,
    custom_inputs: Optional[dict[str, Any]] = None,
  ) -> ChatAgentResponse:
    prompt = messages[-1].content
    raw_msgs = run_agent(prompt)
    out = []
    for m in raw_msgs:
      out.append(ChatAgentMessage(
        id=uuid.uuid4().hex,
        **m
      ))

    return ChatAgentResponse(messages=out)

AGENT = QuickstartAgent()
mlflow.models.set_model(AGENT)

记录代理信息

登陆你的代理并将其注册到 Unity 目录服务。这会将代理及其依赖项打包到单个项目进行部署。

import mlflow
from mlflow.models.resources import DatabricksFunction, DatabricksServingEndpoint
from pkg_resources import get_distribution

# Change the catalog name ("main") and schema name ("default") to register the agent to a different ___location
registered_model_name = "main.default.quickstart_agent"

# Specify Databricks resources that the agent needs to access.
# This step lets Databricks automatically configure authentication
# so the agent can access these resources when it's deployed.
resources = [
  DatabricksServingEndpoint(endpoint_name="databricks-claude-3-7-sonnet"),
  DatabricksFunction(function_name="system.ai.python_exec"),
]

mlflow.set_registry_uri("databricks-uc")
logged_agent_info = mlflow.pyfunc.log_model(
  artifact_path="agent",
  python_model="quickstart_agent.py",
  extra_pip_requirements=[f"databricks-connect=={get_distribution('databricks-connect').version}"],
  resources=resources,
  registered_model_name=registered_model_name
)

部署代理

将已注册的代理部署到服务端点：

from databricks import agents

deployment_info = agents.deploy(
  model_name=registered_model_name, model_version=logged_agent_info.registered_model_version
)

代理终结点启动后，可以使用 AI Playground 与其聊天，或与利益干系人共享以获取反馈。

后续步骤

根据目标选择下一步的位置：

衡量和改进代理的质量：请参阅代理评估快速入门。

生成更高级代理：创建使用非结构化数据执行 RAG 的代理，处理多轮次对话，并使用代理评估来衡量质量。请参阅教程：生成、评估和部署检索代理。

了解如何使用其他框架生成代理：了解如何使用 LangGraph、纯 Python 和 OpenAI 等常用库生成代理。请参阅使用ChatAgent用于创作代理

通过

AI 代理入门

要求

示例笔记本

马赛克 AI 代理演示

定义代理

测试代理

部署代理

准备代理代码进行部署

记录代理信息

部署代理

后续步骤

反馈

其他资源