返回博客
随笔
2026-09-0214 min

记一次 AI 问答 + RAG 检索增强实现

> stardee.cn/ai 中增加 AI 助手「多角色流式对话」+「知识库 RAG 检索增强」的完整实现复盘。 > 主体是通用方法（可复用于同类功能），并在每处标注「本项目选型」作为落地参考。 一次带知识库增强的 AI 功能，由两条链路…

记一次 AI 问答 + RAG 检索增强实现

stardee.cn/ai 中增加 AI 助手「多角色流式对话」+「知识库 RAG 检索增强」的完整实现复盘。

主体是通用方法（可复用于同类功能），并在每处标注「本项目选型」作为落地参考。

一、总体思路

一次带知识库增强的 AI 功能，由两条链路组成：

  • 对话链路：前端发起请求 → 后端组装 system / 历史消息 → 调用 LLM 流式接口 → 把正文与「思考过程」逐段推给前端 → 结束后落库 + 记用量。

  • 检索链路（RAG）：文档入库时拆分为语义连续的片段并向量化存储；用户提问时把问题向量化，拿相似度最高的 Top-K 片段拼进提示词，让模型"带着资料回答"。

RAG 的价值：给没有训练过、或需要最新/私有资料的模型提供上下文；核心是检索质量（分块 + 向量化）与注入方式（如何把上下文交给模型而不污染正常对话）。

二、AI Chat（流式对话）设计

1. Provider 抽象

用一个注册表统一管理多个模型服务商：每个 provider 提供唯一 key、显示名、默认模型、创建 client 的方式、拉取模型列表的方式。模型标识常用复合形式 ${providerKey}::${model}，便于在接口层透传"用哪个 provider 的哪个模型"，前端也可直接复用该字符串作为选项值。

  • 回退原则：指定 provider 未配置 Key → 回退默认 provider；默认也未配置 → 返回失败（调用方给出明确提示，如 503）。

  • 模型列表尽量走服务商官方 /models 接口动态获取，不要写死；带较短 TTL 的内存缓存（如 5 分钟）即可。

本项目选型

  • 注册表内置 DeepSeek（默认）与 OpenAI 两个 provider，统一走 @ai-sdk/openai-compatible 接入 OpenAI 兼容接口。

  • 默认模型 deepseek-chat，由环境变量 AI_DEFAULT_PROVIDER / DEEPSEEK_MODEL 控制。

  • 模型复合标识为 deepseek::deepseek-chat；DeepSeek 模型列表通过 GET /models 动态拉取并缓存 5 分钟。

  • 复合 key 解析：parseProviderKey 以 :: 为分隔，兼容旧调用方。

2. 消息组装与 system 分隔

较新版本的主流 AI SDK（Vercel AI SDK 等）不允许 system 角色出现在 messages 数组中，系统提示词要作为独立的 system 选项传入。因此组装时拆成两部分：

  • system：角色设定 + RAG 上下文（如有）。

  • messages：仅保留 user / assistant 历史 + 最新用户消息；历史中的非 user/assistant 记录要过滤掉，否则可能触发 SDK 校验错误。

通常还会：把角色的专业背景、对话风格、输出偏好等结构化配置拼进 system；首轮对话用用户第一句话截断作会话标题。

本项目选型

  • buildMessages 返回 { system, messages }；有 ragContext 时拼在 system 之后。

  • 角色增强字段（background / conversationStyle / responsePreference）以 【专业背景】… 等小节拼进 system。

  • 首轮（history 为空）用首条用户消息前 20 字作为会话标题。

3. 流式传输

后端用 SSE（text/event-stream） 返回增量，常见事件：reasoning（思考）、text（正文）、done（收尾 + token 用量）、error。

  • 遍历 LLM 流按增量类型分发：text-delta → 累加正文并推送；reasoning-delta → 累加思考并推送；错误抛出。

  • 流结束后取 usage（input/output tokens）。

  • 客户端断开要中止上游：往流里写入失败即视为客户端已离开，AbortController.abort() 终止 LLM 请求避免继续计费，同时停止向下游推送。

  • 持久化、写日志等收尾动作在正文收完后再做，让用户先拿到完整回复。

本项目选型

  • app/api/ai/chat 用 ReadableStream + TextEncoder 手写 SSE，事件为 reasoning / text / done / error。

  • 检测到 clientGone 时 abort 上游；流结束后落库助手回复 + logUsage 记账量 + 首轮更新标题 + send("done") 关流。

4. 思考深度 / 档位映射

若模型支持"思考式"推理，可用档位（如 off/low/high/max）控制强度。档位名用稳定英文枚举；校验并规范化（请求显式值 > 角色默认值 > 兜底值），非法值回退；经 providerOptions 透传，不支持的模型静默忽略。

本项目选型

  • 档位枚举 off/low/high/max，normalizeThinkingDepth 规范化。

  • DeepSeek 映射：off → { thinking: { type: "disabled" } }（v4 思考模式默认开启需显式禁用）；其余 → { thinking: { type: "enabled" }, reasoningEffort: depth }。

  • OpenAI 等 provider 不传 thinking 参数（静默忽略）。

5. 配额与用量

  • 双重限制：次数（限流器、滚动窗口）+ token 累计（按日志聚合），管理员/内部账号豁免。

  • 预检前置：在调用 LLM、持久化任何数据之前校验，被拒请求零写入。

  • 估算 prompt token：发送前按字符数粗略估算（如 字符数 / 4）预判超额度。

  • 剩余额度可用来限制本次输出 maxOutputTokens。

  • 用量日志统一落一张表（类型 chat/embedding/rag_search、模型、token、耗时、成功与否），既是成本核算也是统计基础。

  • DB 查询失败用 fail-open（降级放行），避免底层抖动阻塞核心链路。

本项目选型

  • 非管理员默认 2 次/天 + 10000 token/天，滚动窗口 86400s，均由环境变量覆盖。

  • estimatePromptTokens 用 字符数 / 4；provider 输出上限 deepseek 8192 / openai 16384，maxTokens 取 min(剩余额度, 上限)。

三、RAG（检索增强）设计

1. 存储选型

常用 PostgreSQL + pgvector：在分块表上加 vector(dim) 列，长度与 embedding 模型输出维度一致。向量列多数 ORM "不认识"，写入与检索走原生 SQL（::vector 转换、<=> 操作符）。embedding provider 通常独立于对话 LLM，可单独配置。

本项目选型

  • embedding模型使用BAAI/bge-m3(适合中文)，硅基流动有可以免费使用的版本适合个人测试和验证使用

  • KnowledgeChunk.embedding 用 Unsupported("vector(1024)")；Prisma 无法对它做普通 CRUD，故读写都用 $executeRaw / $queryRaw。

  • 维度 1024 与所选 embedding 模型输出一致。

2. 文档入库链路

流程：建文档记录(处理中) → 文本分块 → 批量向量化 → 批量写分块表(带向量) → 文档置就绪 + 记录分块数。

  • 分块：纯函数便于单测。

  • 向量化：批量调用（embedMany）；生成失败（provider 未配 / 报错）时把文档标记为 failed，明确展示，而非静默保留半成品。

  • 写入：用事务包装保证原子性；主键可由数据库生成；向量以 [...]::vector 字符串转换写入。

  • 文档的 chunkCount、status（processing/ready/failed）是后续筛选（只检索 ready）与统计的基础。

本项目选型

  • addDocument 分块 → generateEmbeddings 批量向量化 → prisma.$transaction(...$executeRaw) 写入，主键 gen_random_uuid()::text 由 PG 生成。

  • 同步 logUsage 记录 embedding 用量（模型、token 估算、耗时）。

3. 文本分块（chunking）策略

分块质量直接决定检索精度。目标是：每块语义相对完整、长度适中、块间边缘尽量连续。常用策略（按优先级组合）：

  1. 按段落切分：先以空行（\n\n）分出语义块，段落是天然话题边界。

  2. 长段落再按句子细化：超过上限 maxChunkSize 的段落，按句末标点（中英文 。！？.!?）切句、逐句累积拼接，保证单块不超上限且不切断句子。

  3. 合并过短片段：把过短小段并入前一片，让每块贴近目标长度，避免大量碎片向量稀释检索质量。

  4. 加重叠（overlap）：相邻块共享上一块末尾一段字符作下一块开头，保住跨块边界的上下文。

参数经验：maxChunkSize ≈ 512 token（约 15002000 字符，按 3 字符/token 折算）；overlap ≈ 50 token（约 150200 字符）；具体随 embedding 模型 / 检索方式实测调整。

进阶方案：按 Markdown 标题 / 代码块 / 语义（小模型打分）切分。多数场景规则式分块已足够好、成本低、可单测。

本项目选型

  • chunkText 完全按上述规则：段落 → 超 maxChunkSize 按句切分 → 合并 + overlap。

  • 参数：maxChunkSize = 2000 字符、overlap = 200 字符；token 估算用 字符数 / 3。

4. 查询检索

流程：对用户问题生成 embedding → COS 距离排序取 Top-K → 返回片段文本。

sql

SELECT chunkId, content, similarity

FROM chunk ch JOIN doc ON ch.document_id = doc.id

WHERE doc.kb_id = :kbId AND doc.status = 'ready'

ORDER BY ch.embedding <=> :queryVector

LIMIT :k;

要点：

  • <=> 是 pgvector 余弦距离，相似度 = 1 - 距离。

  • 只检索 status = 'ready' 的文档。

  • Top-K 优先级：显式参数 > 环境变量 > 默认值。

  • 原生查询数值可能以字符串返回，统一转 number。

  • 失败降级：检索异常/无命中返回空，不拖垮对话（记日志即可）。

  • 可选：检索后过滤低于相似度阈值的片段，避免注入无关内容。

本项目选型

  • searchChunks 对查询 generateEmbedding 后，用 1 - (embedding <=> :q::vector) AS similarity 排序，取前 k。

  • Top-K 优先级：参数 > AI_RAG_TOP_K > 默认 4；结果统一 Number() 规范化。

  • 检索失败仅 logger.warn，聊天照常进行。

5. 上下文注入

  • 把 Top-K 片段的文本用分隔符（空行 + ---）连接成"知识库上下文"，拼进 system 提示词，并结构化提示模型据此作答。

  • 注入 system 而非混入对话数组：不污染正常对话历史、与"system 不进 messages"约束一致。

  • 可选：附带来源文档/序号，便于模型说明出处、前端展示引用定位。

本项目选型

  • ragContext = results.map(r => r.content).join("\n\n---\n\n")，拼在 system 后。

四、本项目数据模型

  • 角色 AIRole：systemPrompt + background/conversationStyle/responsePreference + thinkingDepth + ownerId（内置/自定义）。

  • 会话 / 消息：ChatSession（按 userId 隔离）+ ChatMessage（role user/assistant/system、content、tokens）。

  • 知识库 / 文档 / 分块：KnowledgeBase（ownerId、category）→ KnowledgeDocument（status、chunkCount）→ KnowledgeChunk（content、embedding vector(1024)、seqIndex 保留顺序）。

  • 用量日志 AIUsageLog：type(chat/embedding/rag_search)、model、tokens、duration、success。

  • 级联删除：删库 → 删文档 → 删分块，交给 Prisma onDelete: Cascade。

五、通用注意点（坑）

  1. 权限隔离在服务端再做一遍：前端/中间件不是安全边界；会话、知识库归属要在路由处理器 / 数据层用 userId/ownerId 复核。

  2. 系统提示词一个口子：新版 AI SDK 禁 system 进消息数组，走独立 system 选项并过滤历史 system 记录。

  3. 配额预检写在副作用之前：被拒请求零写入。

  4. 检索失败降级、用量查询 fail-open：底层组件不稳不能拖垮核心链路。

  5. 原生 SQL 结果/入参规范化：向量列走原生 SQL 时注意入参类型转换与数值返回类型。

  6. 分块与 embedding 解耦：分块纯函数方便单测；向量化独立 provider，各自可替换。

  7. 流式中断处理：客户端断开立即 abort 上游，避免无效计费。

六、可直接测试的核心单元

  • 分块：输入文本 → 断言块数、每块不超上限、跨边界未被生切、有重叠。

  • 消息组装：system 与 messages 分离、history 过滤非 user/assistant、RAG 上下文对 system 的影响。

  • 思考档位映射：off → 显式禁用；其他 → enabled + effort；非法值回退。

  • 配额：次数/token 边界、管理员豁免、被拒不消耗额度。

  • 流式：模拟文本流 + 推理流，断言正文/思考增量按序分发、汇总、用量返回。

本项目现状：以上各点均有对应单测（lib/ai/*.test.ts），覆盖 buildMessages、streamChat、chunkText、thinking、quota 等，可作回归保障。

AI