数据截至 (上游 commit 5215e9791133)
对话层:AutoQuery、多 KB 检索与带引用回答
30 秒导读: 前面四章把一个
KnowledgeBase讲透了——文档进去(摄取、补上下文),一句查询进去、相关长段(RSE)出来。但kb.query()只吃一句独立查询、只认一个库、返回的是裸文本段。本章讲 dsRAG 怎么在它之上再包一层,让它变成一个真正的对话助手:能记住多轮历史、能同时查好几个知识库、能在回答里标出"这句话来自第几号来源第几页",并且流式或一次性都行。
本章假设你已经理解 kb.query() 会跑 RSE 返回相关段落(见第 03 章)和各可插拔组件的职责(见第 04 章)。这里不重复 RSE 内部,只讲对话层怎么调它、怎么把结果拼成能引用的回答。
1. 这是什么(零基础也能懂)
一句话定义: 对话层是一组无状态的函数(不是一 个类),把一个或多个 KnowledgeBase 组装成"可多轮对话、带引用回答"的 RAG 助手。
它解决的问题: 直接用 kb.query() 有三个缺口——
| 缺口 | 裸 kb.query() | 对话层补上 |
|---|---|---|
| 多轮 | 只吃一句查询,不知道上文 | 存历史、把历史一起喂给"查询生成器" |
| 多库 | 一次只查一个 KB | 让模型自己决定每条查询打到哪个 KB |
| 可信 | 返回裸文本,不知出处 | 回答里每句都带 来源编号 + 页码 + 原文 |
用起来什么样: 三步——建库、建线程、问问题。下面是一段贴近真实 API 的最小示例:
# 示意,非源码:展示对话层的三步用法
from dsrag.chat.chat import create_new_chat_thread, get_chat_thread_response
from dsrag.database.chat_thread.sqlite_db import SQLiteChatThreadDB
from dsrag.chat.chat_types import ChatResponseInput
chat_db = SQLiteChatThreadDB() # 线程持久化(sqlite)
thread_id = create_new_chat_thread( # 建一个对话线程
{"kb_ids": ["finance_kb", "legal_kb"]}, chat_db # 一个线程可绑多个 KB
)
resp = get_chat_thread_response( # 问一句
thread_id,
ChatResponseInput(user_input="2023 年营收和相关合同条款?"),
chat_db,
knowledge_bases={"finance_kb": kb1, "legal_kb": kb2},
)
print(resp["model_response"]["content"]) # 回答正文
print(resp["model_response"]["citations"]) # [{source_index, page_number, cited_text, doc_id, kb_id}, ...]
一句话直觉: 把 kb.query() 当成一个"只会答单题、不记事、不署名"的图书管理员;对话层给他配了个秘书:秘书听你连续说话、替你把问题拆成几条精确检索、分派到对的书库、最后把答案连同"翻到哪本书哪页"一起整理给你。
本节不出现底层代码细节。记住一 件事:对话层是函数 + 一个线程数据库,不是有状态的对象。
2. 顶层全景(它大概怎么转)
2.1 部件一句话职责
| 部件 | 干什么 | 文件 |
|---|---|---|
create_new_chat_thread | 建线程、灌默认参数、写库 | dsrag/chat/chat.py:76 |
ChatThreadParams | 线程的全部配置(模型/温度/KB 列表/历史上限…) | dsrag/chat/chat_types.py:5 |
ChatThreadDB | 线程 + 交互记录的持久化抽象 | dsrag/database/chat_thread/db.py:3 |
get_search_queries(AutoQuery) | 把"历史+新输入"变成多条绑定到具体 KB 的查询 | dsrag/chat/auto_query.py:60 |
_prepare_chat_context | 组装上下文:跑查询、调 kb.query、编来源号、拼 system 消息 | dsrag/chat/chat.py:243 |
format_sources_for_context | 把检索结果包成带 source_index/页码标签的文本 | dsrag/chat/citations.py:42 |
get_response(instructor) | 统一多家 LLM,产出结构化的 ResponseWithCitations | dsrag/chat/instructor_get_response.py:14 |
get_chat_thread_response | 对外总入口,按 stream 分流 | dsrag/chat/chat.py:848 |
2.2 一次问答的主线(从输入到带引用的回答)
下面这张图是本章的骨架。怎么读:从上往下是一次 get_chat_thread_response 的完整生命周期;左侧是数据,右侧标出负责的函数。
用户输入 + thread_id
│
▼
读线程(历史 + 参数) get_chat_thread_response chat.py:848
│ → 按 stream 分流
▼
┌───────────────────── ─────────────────────────┐
│ _prepare_chat_context chat.py:243 │ ← 核心组装,流式/非流式共用
│ │
│ ① 拼历史 + 新输入 → 截断 limit_chat_messages chat.py:146
│ ② AutoQuery:生成多条查询 get_search_queries auto_query.py:60
│ ③ 按 KB 分组查询 chat.py:326-332
│ ④ 逐 KB 调 kb.query 跑 RSE chat.py:340-342
│ ⑤ 给每条结果编 source_index chat.py:344-350
│ ⑥ 包成带页码标签的来源文本 format_sources_for_context citations.py:42
│ ⑦ 填进 MAIN_SYSTEM_MESSAGE chat.py:398-404
└──────────────────────────────────────────────┘
│
▼
LLM 结构化生成 get_response instructor_get_response.py:14
→ ResponseWithCitations citations.py:10
│
▼
引用回填:source_index→doc_id→kb_id chat.py:619-636 / 515-540
│
▼
写库(add/update_interaction) ChatThreadDB
│
▼
返回 interaction(正文 + citations + 检索段)
一句话抓住主线:AutoQuery 拆问题 → 分派多库检索 → 编号 → LLM 带号回答 → 把号翻译回真实文档。后面各节逐个拆。
3. 会话线程与参数(状态放哪、怎么配)
这节讲:对话层"无状态函数"里那点必须持久化的状态,存在哪、长什么样。
3.1 一个线程 = 一份参数 + 一串交互
create_new_chat_thread 做的事很轻:生成 UUID、补 supp_id、灌默认值、写库、返回 thread_id。
# 真实源码节选 dsrag/chat/chat.py:116-123 create_new_chat_thread
thread_id = str(uuid.uuid4())
chat_thread_params["thread_id"] = thread_id
if "supp_id" not in chat_thread_params:
chat_thread_params["supp_id"] = ""
chat_thread_params = _set_chat_thread_params(chat_thread_params) # 灌默认
chat_thread_db.create_chat_thread(chat_thread_params=chat_thread_params)
return thread_id
它把"填默认值"这件事全权交给 _set_chat_thread_params(chat.py:164)。这个函数的模式是**"传了就用,没传就填默认"**,逐个字段兜底:
| 参数 | 默认值 | 作用 |
|---|---|---|
kb_ids | [] | 这个线程能查哪些知识库 |
model | "gpt-4o-mini" | 生成回答的 LLM |
temperature | 0.2 | 采样温度 |
system_message | "" | 用户自定义的角色/指令,注入进大 system prompt |
auto_query_model | "gpt-4o-mini" | 生成检索查询用的模型 |
auto_query_guidance | "" | 给 AutoQuery 的额外提示 |
target_output_length | "medium" | 回答长短(short/medium/long) |
max_chat_history_tokens | 8000 | 历史最多带多少 token |
rse_params | {} | 透传给 kb.query 的 RSE 参数(见第 03 章) |
字段的权威定义是 ChatThreadParams(chat_types.py:5),一个 TypedDict——注意它是结构约定而非运行时校验,真正的兜底靠 _set_chat_thread_params。
3.2 参数怎么被"两次"确定
一个容易忽略的细节:参数会被兜底两次。建线程时 _set_chat_thread_params 兜一次并写库;每次回答时 _prepare_chat_context 又从库里读出来、再兜一次(chat.py:276-287)。这让旧线程也能用上新加的默认字段,不至于因缺字段崩掉。
另外,单次请求可以临时覆盖整份参数:ChatResponseInput.chat_thread_params 非空时,直接顶掉线程里存的那份(chat.py:712-716),用于"这一问我想换个模型/换套 RSE"这种场景。
3.3 两种线程数据库
ChatThreadDB(db.py:3)是抽象基类,定义了线程的增删改查 + 两个交互操作 add_interaction / update_interaction。两个实现:
| 实现 | 存哪 | 适合 |
|---|---|---|
BasicChatThreadDB | 单个 chat_thread_db.json,全内存 + 落盘 | 原型、单机、量小 |
SQLiteChatThreadDB | ~/dsRAG/chat_thread.db 两张表 | 稍正式的持久化 |
SQLiteChatThreadDB 值得点两处工程细节:
- 列打平存储。
kb_ids列表被",".join成字符串存、读时再split(",");rse_params字典走json.dumps/loads(sqlite_db.py:41-44、77-83)。SQLite 没有原生列表/字典列,这是常规打平手法。 - 自动迁移。 老库缺
citations/model_response_status/rse_params列时,_check_and_migrate_db(sqlite_db.py:258)用ALTER TABLE ADD COLUMN补上,老用户升级不用手动改表。
读线程时 get_chat_thread(sqlite_db.py:88)把两张表重组成 {id, params, interactions} 的嵌套结构,其中每条 interaction 被还原成 user_input / model_response / relevant_segments / search_queries 的嵌套形状(sqlite_db.py:124-140)——这正是 _prepare_chat_context 拼历史时期待的形状。
4. AutoQuery:把一段话拆成"打到哪个库"的多条查询
这节讲本章第一支核心机制:用户说的是自然语言、还带着上文,怎么变成
kb.query()能吃的、还标明去哪个库的检索查询。
4.1 要解决的小问题
kb.query() 只吃 list[str],而且不知道该查哪个库。但真实提问往往:(a) 带指代("那它的营收呢"里的"它"要靠历史);(b) 一句话里塞了好几个信息点;(c) 不同信息点属于不同库。AutoQuery 就是把这团东西结构化成清晰的检索计划。
4.2 思路:让模型输出一个带 kb_id 的查询列表
核心是两个 Pydantic 模型 + 一次结构化 LLM 调用。模型被要求为每条查询指定 knowledge_base_id:
# 真实源码 dsrag/chat/auto_query.py:17-22
class Query(BaseModel):
query: str
knowledge_base_id: str
class Queries(BaseModel):
queries: List[Query]
get_search_queries(auto_query.py:60)把各 KB 的标题+描述拼进 system prompt,连同完整对话历史(含新输入)一起交给 instructor,拿回一个 Queries。因为历史在场,模型能自己消解"它/那个"这类指代——多轮能力 就落在这里。
4.3 关键细节一:双模型 fallback
结构化输出偶尔会失败(小模型 JSON 生成不稳)。AutoQuery 的对策是弱模型先试、失败换强模型:
# 真实源码节选 dsrag/chat/auto_query.py:83-100
try:
queries = get_response(messages=..., model_name=auto_query_model, response_model=Queries, ...)
except:
# gpt 系列就退到 gpt-4o,否则退到 claude-3-5-sonnet
fallback_model = "claude-3-5-sonnet-20241022" if "gpt" in auto_query_model else "gpt-4o"
queries = get_response(messages=..., model_name=fallback_model, response_model=Queries, ...)
上一层 _prepare_chat_context 还套了一层"整体失败就当空查询"的兜底(chat.py:319-323)——AutoQuery 崩了不至于让整次对话崩,只是这轮没检索。
4.4 关键细节二:validate_queries 的自愈
模型可能吐出一个不存在的 knowledge_base_id(幻觉)。validate_queries(auto_query.py:31)不直接丢弃,而是分情况自愈:
模型给的 kb_id 合法吗?
│
┌──┴───────────────┐
是 否
│ │
直接采用 只有 1 个 KB? ──是──▶ 强制用那唯一的 KB
│
否(多个 KB)
│
把这条查询 fan-out 到每一个 KB
也就是说:分不清该去哪时,宁可多查也不漏查。最后再对结果和数量都做 [:max_queries] 截断(auto_query.py:103)。
4.5 一个诚实的说明:两个 auto_query 文件
仓库里有两个同名 get_search_queries:
dsrag/chat/auto_query.py—— 本章讲的、在用的这个,支持多 KB + 历史。dsrag/auto_query.py—— 顶部自注NOTE: this is a legacy file and is not used(dsrag/auto_query.py:1),只吃单串输入、不认 KB。别读错文件。
5. 上下文组装 _prepare_chat_context(把检索结果编号并拼进 prompt)
这节讲流式/非流式共用的组装心脏:从"有哪些查询"到"一条能喂给 LLM 的完整 messages"。
5.1 历史拼接与截断
先把库里的历次交互摊平成 user/assistant 轮次,末尾追加本轮输入(chat.py:305-309),再按 token 预算截断:
# 真实源码 dsrag/chat/chat.py:146-162 limit_chat_messages(节选)
for message in reversed(chat_messages): # 从最新往回数
message_tokens = count_tokens(message['content'])
if total_tokens + message_tokens <= max_tokens:
limited_messages.insert(0, message) # 塞回开头,保持时序
total_tokens += message_tokens
else:
break # 预算用完,更旧的丢掉
策略很直白:保新弃旧,整条消息为单位(不切半条),token 用 tiktoken 的 gpt-4o 编码估算(chat.py:141-144)。
5.2 分组、检索、编号——引用系统的地基
这是全章最关键的一小段。查询按 KB 分组后逐库调 kb.query(把 rse_params、metadata_filter 透传下去):
# 真实源码节选 dsrag/chat/chat.py:340-350
for kb_id, queries in search_queries_by_kb.items():
kb = kbs.get(kb_id)
search_results[kb_id] = kb.query(search_queries=queries, rse_params=rse_params, metadata_filter=metadata_filter)
# 关键:给跨库的每条结果编一个全局连续号 source_index
i = 0
for kb_id, results in search_results.items():
for result in results:
result["source_index"] = i
source_index_to_doc_id[i] = result["doc_id"] # 记下:号 → 真实文档
i += 1
为什么要这个 source_index? 因为 LLM 不该看到又长又乱的 doc_id/kb_id,更不该被要求原样吐回来。dsRAG 的做法是给每个来源发一个**简单的整数号(0,1,2…)**给模型看;同时在 source_index_to_doc_id 里悄悄记下"几号=哪个真实文档"。等模型带号引用完,再翻译回真身(见第 6 节)。这是整套引用能可核查的地基。
5.3 包装成带页码标签的来源文本
编好号的结果交给 format_sources_for_context(citations.py:42)包成 LLM 友好的文本。有页码就带页码标签,没有就退回裸内容:
<source_index: 3>
<page_12>
……这一页的正文……
</page_12>
</source_index: 3>
页码文本由 get_source_text(citations.py:21)从 file_system.load_page_content_range 取——只有摄取时存过页内容才有(见 5.5)。整段来源文本连同各 KB 描述、用户自定义 system 消息、长度指引,一起填进 MAIN_SYSTEM_MESSAGE(chat.py:398-404),它内含完整的引用格式契约(chat.py:39-60:告诉模型必须回一个含 source_index/page_number/cited_text 的对象)。
5.4 一个诚实提醒:被架空的 format_relevant_knowledge_str
chat.py:135 定义了 format_relevant_knowledge_str(简单把各段 text 拼起来),但当前的 _prepare_chat_context 并不调用它——实际拼来源文本走的是 format_sources_for_context(带 source_index/页码,才能支撑引用)。前者更像早期遗留;看代码时别被它误导以为那是主路径。(inferred:基于全文件内无调用点)
5.5 页码内容从哪来:convert_elements_to_page_content
引用能标"第 N 页",前提是摄取时就按页存过原文。convert_elements_to_page_content(citations.py:70)在文档首次入库时被调用:把带 page_number 的 elements 按页分组,逐页 file_system.save_page_content。没有页号的文档(如纯文本)会被直接跳过(citations.py:77-78),这类文档的引用就只有来源号、没有页码——和 Citation.page_number 允许为 None(citations.py:7)是对上的。
6. 生成回答与引用回填(号怎么翻译回真身)
这节讲第二支核心机制:从 LLM 的结构化输出,到一个每条都能点回真实文档页码的
citations列表。
6.1 结构化响应:ResponseWithCitations
LLM 不是自由发挥,而是被 instructor 约束成一个固定结构:
# 真实源码 dsrag/chat/citations.py:5-12
class Citation(BaseModel):
source_index: int # 来自哪个来源号
page_number: Optional[int] = None # 哪一页(无页码时为 None)
cited_text: str # 支撑该结论的原文片段
class ResponseWithCitations(BaseModel):
response: str # 回答正文
citations: List[Citation] # 一组引用
get_response(instructor_get_response.py:14)是一个统一多家 LLM 的适配层:按模型名分派到 OpenAI / Anthropic / Gemini,各家都用 instructor 强制产出上面这个结构(_handle_*_instructor,instructor_get_response.py:100-118)。对话层因此不关心底层是哪家模型。
6.2 回填:source_index → doc_id → kb_id
模型只会给 source_index(它只认得号)。非流式路径 _get_chat_response(chat.py:572)负责把号翻译回真身,并丢弃翻不出来的幻觉引用:
# 真实源码节选 dsrag/chat/chat.py:619-636
for citation in citations:
citation = citation.model_dump()
if citation["source_index"] not in source_index_to_doc_id:
continue # 模型编了个不存在的号 → 丢
citation["doc_id"] = source_index_to_doc_id[citation["source_index"]] # 号→文档
if citation["doc_id"] in all_doc_ids:
citation["kb_id"] = all_doc_ids[citation["doc_id"]] # 文档→库
else:
continue
formatted_citations.append(citation)
两道校验闭环:号必须在映射里、文档必须属于某个已检索的库,才留下。这是"不让模型伪造出处"的最后一关。最终 interaction 里,citations 每项就有了完整的 {source_index, page_number, cited_text, doc_id, kb_id}。
6.3 组装成 interaction
_get_chat_response 最后打包出统一的 interaction 字典(chat.py:639-651):user_input(内容+时间戳)、model_response(正文+citations+时间戳)、search_queries、relevant_segments。这个形状既写进库,也返回给调用方。