首页 / 文章

在Databricks上利用LangGraph构建知识助手智能体

逐步在Databricks上构建LangGraph智能体:从最基础的智能体开始,实现向量搜索RAG功能,利用Lakebase进行短期记忆存储,并设置专门机制用于长期记忆保存。

发布者:REACTAPP.TOP阅读约 7 分钟1720 词

来源: Building Custom Code Agents on Databricks — Getting Started

Databricks上的基础自定义智能体可以从单纯的LLM调用发展成能够搜索文档、记住用户信息的知识助手,每一步进展都很微小。最终的智能体是使用langgraph框架构建的,它具备用于文本检索的向量搜索组件以及记忆模块,因此能够回溯最终用户在当前对话或过往对话中所说的内容。一旦智能体开始运行,接下来就是评估环节:在开发阶段以及上线后都要测量其性能。

本教程通过一系列笔记本进行讲解,这些笔记本被分类放入不同文件夹中,分别对应基础知识、RAG流程以及记忆功能。

准备工作空间

完整的智能体包含向量索引、工具和日志,而这些都需要一个整洁的存储环境。在 Databricks 中,治理层为 Unity Catalog(UC),它通过三层命名空间来管理所有资产。按顺序运行 00_foundations 文件夹中的笔记本,可以创建后续步骤所依赖的所有资产,其中包括智能体将会连接的结构性数据与非结构性数据。

内存数据将存储在 Lakebase 中,这是一种可自动扩展且兼容 Postgres 的数据库。为确保后续能顺利创建表,设置笔记本还会创建一个 Lakebase 项目及相应的分支来存放这些数据。

第一个最简智能体

接下来是 01_rag_pipeline 文件夹下的笔记本,从 0_your_first_agent_on_databricks 开始。

每个智能体都是基于基础模型构建的。Databricks Model Serving会列出可用的基础模型,具体列表可能会因工作区的配置不同而有所差异。这些端点兼容OpenAI API格式,但若需与langgraph深度集成,官方的databricks-langchain库是更好的选择。

简单介绍一下langgraph。实际应用中的智能体很少仅由单一提示词构成,它通常包含根据问题及用户回答情况变化的分支逻辑,还有用于查询外部系统以及根据反馈重新表述查询的循环结构。Langgraph通过将智能体建模为图来组织这些元素:节点负责存储函数和LLM调用指令,边则决定了上下文在节点之间的传递方式。初始的图结构非常简单。

调用该智能体后,会生成一条追踪记录,列出输入问题、模型的输出以及中间的每一步逻辑。这些功能来自 MLflow,它与 langgraph 有原生集成。在笔记本开头调用一次 mlflow.langchain.autolog() 即可开启这些日志与诊断功能。关于 MLflow 的更多内容将在后面详细介绍。

第一个版本仅将用户的问题传递给模型。而真正的智能体还会添加系统提示,用以设定语气、塑造输出内容并提供背景信息,比如助手的用途以及它应该拒绝回答的问题。因此,在第二版中会在流程图中加入系统提示。

添加检索功能

最常见的智能体类型之一是知识检索智能体,也称为RAG智能体。它能够帮助人们在大量文档中搜索并回答相关问题。要赋予智能体这种能力,就需要一个向量搜索索引,本质上就是用于存储非结构化文本的数据库。

笔记01_building_a_doc_store_on_vector_search介绍了如何构建这样的索引。在将索引应用于智能体之前,必须先创建好该索引。这部分内容仅讲解核心步骤,未涉及调优与优化。

构建索引需要执行三个操作:

  • 解析指的是打开文档并读取其内容。
  • 分块是指使用分割工具将文档拆分成多个部分。
  • 嵌入是指将每个分块转换为存储在索引中的语义向量。

每种方法都有其独特之处,早期步骤中的错误会悄悄影响后续的所有处理。本教程全程使用默认设置。在实际应用中,不同的文档可能需要不同的解析器来正确提取文本,不同的使用场景也可能要求对文本、表格和图表采用不同的分块大小,以及不同的嵌入模型。

Databricks提供了相应的工具来简化这一过程。Volumes是一个云原生的、受全面管理的对象存储服务,用于存放原始文件。ai_parse_documents函数几乎无需配置即可提取文本、图表和表格,并且具备良好的扩展性。Spark UDFs可以将提取出的文本按不同粒度分割成多个块,以便被各种嵌入模型使用。

一旦有了嵌入向量,相似度搜索就需要向量搜索索引。新的数据存储似乎会增加工作量,但 Delta Sync 会在数据块被添加、修改或删除时自动对其进行嵌入和索引处理。无需管理新的基础设施,也无需在数据块表与索引之间维护同步任务。

在笔记本的前两个步骤中,设置过程中生成的 Markdown 文件会被读取并分割成长度为 800 字符的重叠数据块,随后加载到 Databricks Vector Search 中。databricks_langchain 库提供了一个现成的 VectorSearchRetrieverTool,用于将索引与 langgraph 代理连接起来。

在索引已构建并经过测试后,第三步是对智能体进行重新定义。它与之前的版本非常相似,只是多了一个元素:即ToolNode,该节点会接收一个工具列表。在此示例中,该列表仅包含VectorSearchRetrieverTool。

在最终的流程中,用户提出的问题会先进入提示模板。模型会看到提示内容以及各类工具及其描述,然后决定是否需要查询索引。它可以重复执行这一过程,只有找到所需信息后才会向用户回复。

这就是一个可运行的RAG智能体。其结构依然十分简洁,仅增加了工具节点,这体现了检索循环其实并不需要复杂的结构。不过,知识助手还需要另一个组件:内存。

为智能体添加内存

目前构建的版本中,该智能体完全没有记忆功能。任何试图引用之前回答的后续问题都会失败,因为模型根本记不住刚才说了什么。

为了让智能体能够记住之前的对话内容,langgraph需要一个用于存储聊天历史的检查点文件。这类检查点文件可以是纯文本文件,也可以是完整的数据库。从可扩展性、延迟和安全性考虑,数据库是更优的选择,在Databricks平台上则对应Lakebase。

记忆分为两种类型:短期记忆仅涵盖当前正在进行的会话;长期记忆则能在不同会话之间保留信息。几乎所有应用都需要短期记忆,而是否需要长期记忆则取决于具体应用场景。

短期记忆

02_memory文件夹中包含01_short_term_memory,后者会添加一个检查点。第1至3步会构建与之前相同的StateGraph,而第4步则展示了没有记忆机制时对话会变得多么混乱。

通过创建与Lakebase关联的检查点可以恢复对话的连贯性,databricks_langchain提供了内置的CheckpointSaver类,可用于实例化该检查点并将其附加到图结构中。

智能体是供多个用户共享的服务,因此保存的对话需要有一个所有者。thread_id便承担了这一功能。在实际应用中,前端应为每个活跃会话生成唯一的thread_id,这样每位用户只能看到自己的聊天记录。

到第6步时,检查点已设置完成,对话也能自然进行。

长期记忆

会话结束时,短期记忆就会消失。为了记住用户及其偏好,智能体需要长期记忆,而是否需要这种功能则取决于具体产品。一个长期运行的个人助手应当记住之前的对话、用户的姓名等事实信息,以及用户更倾向于简短直接的回复还是详细的解释等偏好。而客户门户上的客服服务则未必需要保存这些细节,因为那样就需要长期承担保护、加密和管理个人数据的重任。

短期记忆只需CheckpointSaver即可实现。而长期记忆则需要一定的设计考量。通常并不需要将每次对话都记录下来以供日后查看。值得保存的是姓名、家乡等地理信息以及用户偏好等事实内容。

02_long_term_memory 这个笔记本详细介绍了相关实现机制。第1至4步用于构建带有短期记忆的RAG智能体,而第5步则展示了当会话及thread_id发生变化时,用户偏好信息会如何消失。

第6步引入了DatabricksStore对象,用于将数据持久化到Lakebase中。这部分实现相对简单,真正的设计难点在于决定需要持久化的内容是什么。第7步采用了一种较为基础的方法:将少量硬编码的事实直接写入存储中。

为使用这些数据,第8步会重新创建该智能体。此前,智能体节点只是一个简单的函数,仅负责调用基础模型中的invoke方法;而现在它则变成了一个更为复杂的函数——该函数会接收用户ID(演示中为硬编码值),从DatabricksStore中获取该用户的相关事实,并将其整合到发送给模型的上下文中。显然,要使这种设计能够稳定运行,还需要进行精心设计。

采用更灵活的方式可以让模型自行决定何时保存、删除和获取记忆。第12步实现了基础版本:三个新工具与内存模块交互,系统提示也会相应更改,向模型说明这些工具的存在并鼓励其使用。

这样就能打造出一个功能完备的助手,而且由于各项功能是通过独立的笔记本文件以及图结构的独立修改实现的,因此可以单独对每项功能进行测试。它能通过向量索引搜索文档,利用短期记忆跟踪当前对话,并借助长期记忆记住不同会话中用户的详细信息。

未来的发展方向

该路径从将问题直接传递给大语言模型的机器人,发展为具备提示词模板、短期与长期记忆功能以及用于上下文检索的向量搜索功能的版本。后续工作则是将其投入实际应用:评估该智能体的性能,添加更多功能,并作为服务提供给最终用户。