From 72f925a2dffdda48023022c7073d04f8bc160587 Mon Sep 17 00:00:00 2001 From: mrsibe Date: Fri, 25 Sep 2026 15:49:41 +0800 Subject: [PATCH] feat(citations): carry structured citations from retrieval into the answer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prompt numbered its sources but threw the provenance away: by the time an answer existed there was no page, block or span to point back at. `SearchResult` already carries a `locator` (#74), so assemble it into a shared `Citation` snapshot — document, page range, block id, char span and the retrieved quote — and persist it under `chat_messages.metadata.citations`. The citation is a snapshot, not a reference: re-indexing or deleting the source later must not erase where the answer was grounded. A defensive reader (`parseCitations`) drops malformed rows individually so one bad entry cannot hide the rest. The renderer's in-memory message never sees the DB row written before streaming, so the persisted metadata also rides along on the `finish` event; otherwise an answer would stay un-citable until the session was reloaded. Refs #69 --- src/main/ipc/chatHandlers.ts | 61 ++----- src/main/services/citations.ts | 90 ++++++++++ src/preload/index.d.ts | 4 +- src/renderer/src/store/chatStore.ts | 13 +- src/shared/types/chat.ts | 7 + src/shared/types/citation.ts | 31 ++++ src/shared/types/index.ts | 3 + src/shared/utils/citations.ts | 73 ++++++++ test/citations.test.ts | 257 ++++++++++++++++++++++++++++ 9 files changed, 493 insertions(+), 46 deletions(-) create mode 100644 src/main/services/citations.ts create mode 100644 src/shared/types/citation.ts create mode 100644 src/shared/utils/citations.ts create mode 100644 test/citations.test.ts diff --git a/src/main/ipc/chatHandlers.ts b/src/main/ipc/chatHandlers.ts index 3563037..855bf55 100644 --- a/src/main/ipc/chatHandlers.ts +++ b/src/main/ipc/chatHandlers.ts @@ -2,52 +2,17 @@ import { ipcMain, IpcMainInvokeEvent } from 'electron' import * as queries from '../db/queries' import { ConnectionManager } from '../models/ConnectionManager' import { SessionAutoSwitchService } from '../services/SessionAutoSwitchService' -import { KnowledgeService, type SearchResult } from '../services/KnowledgeService' +import { KnowledgeService } from '../services/KnowledgeService' +import { buildRAGContext } from '../services/citations' import { validateAndCleanMessages } from '../utils/messageValidator' import Logger from '../../shared/utils/logger' -import type { AnswerSource, RetrievalStatus } from '../../shared/types/chat' +import type { AnswerSource, ChatMessageMetadata, RetrievalStatus } from '../../shared/types/chat' +import type { Citation } from '../../shared/types/citation' import { ChatSchemas, validate } from './validation' // 管理活跃的流式请求 const activeStreams = new Map() -/** - * 构建 RAG 上下文 prompt，并把「这段回答基于哪些段落」一起交出来。 - * - * 之前这里只取 `documentTitle` / `content` / `score` 三个字段， - * `chunkId`、`documentId`、`chunkIndex` 全部被丢掉 —— 于是回答交付之后， - * 界面上再也没有回到原文的路。prompt 文本保持不变，这里只是不再丢弃身份。 - */ -function buildRAGContext(searchResults: SearchResult[]): { - context: string - sources: AnswerSource[] -} { - if (searchResults.length === 0) return { context: '', sources: [] } - - const sources: AnswerSource[] = searchResults.map((result, index) => ({ - index: index + 1, - documentId: result.documentId, - documentTitle: result.documentTitle, - documentType: result.documentType, - chunkId: result.chunkId, - chunkIndex: result.chunkIndex, - content: result.content, - score: result.score - })) - - const contextParts = sources.map( - (source) => `[来源 ${source.index}: ${source.documentTitle}]\n${source.content}` - ) - - const context = `以下是与用户问题相关的背景知识，请参考这些信息来回答： - -${contextParts.join('\n\n---\n\n')} - -请基于以上背景知识回答用户的问题。如果背景知识不足以回答问题，请说明并尽力提供有帮助的回答。` - - return { context, sources } -} - /** * Register chat-related IPC Handlers */ @@ -145,6 +110,7 @@ export function registerChatHandlers( // 于是「没有依据的回答」和「有依据的回答」在界面上完全无法区分。 let retrieval: RetrievalStatus = 'none' let answerSources: AnswerSource[] = [] + let answerCitations: Citation[] = [] try { const embeddingClient = await connectionManager.getEmbeddingClient() @@ -157,9 +123,10 @@ export function registerChatHandlers( }) if (searchResults.length > 0) { - const { context, sources } = buildRAGContext(searchResults) + const { context, sources, citations } = buildRAGContext(searchResults) retrieval = 'used' answerSources = sources + answerCitations = citations Logger.debug( 'ChatHandlers', `RAG: Found ${searchResults.length} relevant chunks for query` @@ -181,11 +148,13 @@ export function registerChatHandlers( Logger.warn('ChatHandlers', 'RAG search failed:', error) } - queries.updateMessageMetadata(assistantMessage.id, { + const answerMetadata: ChatMessageMetadata = { ...(assistantMessage.metadata ?? {}), retrieval, - sources: answerSources - }) + sources: answerSources, + citations: answerCitations + } + queries.updateMessageMetadata(assistantMessage.id, answerMetadata) // 4. 调用 Model Connection 流式生成 const client = await connectionManager.getChatClient() @@ -257,7 +226,11 @@ export function registerChatHandlers( event.sender.send('message-chunk', { messageId: assistantMessage.id, type: 'finish', - metadata: metadata + metadata: metadata, + // The renderer's in-memory message never sees the DB row written + // before streaming, so the persisted provenance rides along here or + // the answer loses its citations until the session is reloaded. + messageMetadata: answerMetadata }) } }, diff --git a/src/main/services/citations.ts b/src/main/services/citations.ts new file mode 100644 index 0000000..49013d5 --- /dev/null +++ b/src/main/services/citations.ts @@ -0,0 +1,90 @@ +import type { Citation } from '../../shared/types/citation' +import type { AnswerSource } from '../../shared/types/chat' +import type { SearchResult } from './KnowledgeService' + +/** + * 检索 → 引用 的组装。 + * + * 检索层（`RetrievedEvidence`）已经带着来源身份、页码区间和块区间；这里只把它 + * 投影成 prompt 与 `chat_messages.metadata` 用的形状，不再回查数据库。保持纯函数 + * 是刻意的 —— #70 的解析/校验也走同一份 `Citation`，两条路径必须能看到同样的 + * 定位信息。 + */ + +/** 一个检索结果 → 一条 citation。`index` 是它在 prompt 里的 1-based 位置。 */ +export function citationFromSearchResult(result: SearchResult, index: number): Citation { + const { blocks, pageStart, pageEnd } = result.locator + const first = blocks[0] + const last = blocks[blocks.length - 1] + + const citation: Citation = { + index, + documentId: result.documentId, + documentTitle: result.documentTitle, + chunkId: result.chunkId, + quote: result.content, + score: result.score + } + + if (result.documentType) citation.documentType = result.documentType + if (pageStart !== null) citation.page = pageStart + if (pageEnd !== null) citation.pageEnd = pageEnd + + // A chunk can cover several blocks. The jump anchor is the first one; the + // char span is the union of the blocks actually covered, which is what the + // reader highlights. + if (first && last) { + citation.blockId = first.blockId + citation.startOffset = first.startOffset + first.startInBlock + citation.endOffset = last.startOffset + last.endInBlock + } + + return citation +} + +export function buildCitations(results: SearchResult[]): Citation[] { + return results.map((result, index) => citationFromSearchResult(result, index + 1)) +} + +/** 检索结果在 prompt / 元数据里共用的「回答基于什么」视图。 */ +export interface RAGContext { + context: string + sources: AnswerSource[] + citations: Citation[] +} + +/** + * 构建 RAG 上下文 prompt，并把「这段回答基于哪些段落」一起交出来。 + * + * 之前这里只取 `documentTitle` / `content` / `score` 三个字段，`chunkId`、 + * `documentId`、`chunkIndex` 全部被丢掉 —— 于是回答交付之后，界面上再也没有回到 + * 原文的路。prompt 文本保持原有结构，这里只是不再丢弃身份，并显式要求模型用 + * `[n]` 标注来源，让回答里的断言可以解析回 citation。 + */ +export function buildRAGContext(searchResults: SearchResult[]): RAGContext { + if (searchResults.length === 0) return { context: '', sources: [], citations: [] } + + const citations = buildCitations(searchResults) + const sources: AnswerSource[] = searchResults.map((result, index) => ({ + index: index + 1, + documentId: result.documentId, + documentTitle: result.documentTitle, + documentType: result.documentType, + chunkId: result.chunkId, + chunkIndex: result.chunkIndex, + content: result.content, + score: result.score + })) + + const contextParts = sources.map( + (source) => `[来源 ${source.index}: ${source.documentTitle}]\n${source.content}` + ) + + const context = `以下是与用户问题相关的背景知识，请参考这些信息来回答： + +${contextParts.join('\n\n---\n\n')} + +请基于以上背景知识回答用户的问题。引用某个来源时，请在相应句子后使用 [n] 标注来源编号（n 为上面的来源编号，例如 [1]）。如果背景知识不足以回答问题，请说明并尽力提供有帮助的回答。` + + return { context, sources, citations } +} diff --git a/src/preload/index.d.ts b/src/preload/index.d.ts index 8d021d0..55546a8 100644 --- a/src/preload/index.d.ts +++ b/src/preload/index.d.ts @@ -1,5 +1,5 @@ import { ElectronAPI } from '@electron-toolkit/preload' -import type { ChatSession, ChatMessage } from '../shared/types/chat' +import type { ChatSession, ChatMessage, ChatMessageMetadata } from '../shared/types/chat' import type { Notebook, Note, @@ -208,6 +208,8 @@ declare global { content?: string reasoningId?: string metadata?: any + /** Persisted `chat_messages.metadata` for the finished answer. */ + messageMetadata?: ChatMessageMetadata }) => void ) => () => void onMessageError: (callback: (data: { messageId: string; error: string }) => void) => () => void diff --git a/src/renderer/src/store/chatStore.ts b/src/renderer/src/store/chatStore.ts index cf36365..b2be935 100644 --- a/src/renderer/src/store/chatStore.ts +++ b/src/renderer/src/store/chatStore.ts @@ -234,7 +234,7 @@ export const useChatStore = create()((set, get) => ({ export function setupChatListeners() { // 监听流式消息片段（AI SDK fullStream 格式） const cleanupChunk = window.api.onMessageChunk((data) => { - const { messageId, type, content } = data + const { messageId, type, content, messageMetadata } = data const store = useChatStore.getState() const cached = store.messageToNotebook[messageId] @@ -306,6 +306,17 @@ export function setupChatListeners() { // 流式传输完成 store.setStreamingMessage(cached.notebookId, null) + // Provenance (retrieval status, sources, citations) is written to the DB + // before the model call, so the in-memory message never saw it. Apply it + // now, or the answer stays un-citable until the session is reloaded. + if (messageMetadata) { + useChatStore.setState((state) => ({ + messages: state.messages.map((msg) => + msg.id === messageId ? { ...msg, metadata: messageMetadata } : msg + ) + })) + } + // 清理缓存 // eslint-disable-next-line @typescript-eslint/no-unused-vars const { [messageId]: _removed, ...rest } = store.messageToNotebook diff --git a/src/shared/types/chat.ts b/src/shared/types/chat.ts index 8f64005..ae530ce 100644 --- a/src/shared/types/chat.ts +++ b/src/shared/types/chat.ts @@ -9,6 +9,7 @@ import type { ChatSession as DBChatSession, ChatMessage as DBChatMessage } from '../../main/db/schema' +import type { Citation } from './citation' /** * 聊天会话接口（完整版） @@ -63,6 +64,12 @@ export type RetrievalStatus = 'used' | 'none' | 'failed' export interface ChatMessageMetadata { sources?: AnswerSource[] retrieval?: RetrievalStatus + /** + * The structured provenance of this answer (#69). `sources` says *what* the + * answer was built from; `citations` says *where* in the source each marker + * points. Both are snapshots written when the answer was produced. + */ + citations?: Citation[] [key: string]: unknown } diff --git a/src/shared/types/citation.ts b/src/shared/types/citation.ts new file mode 100644 index 0000000..709a4c0 --- /dev/null +++ b/src/shared/types/citation.ts @@ -0,0 +1,31 @@ +/** + * 结构化引用 + * + * 一条 citation 是「回答里的一句断言」与「来源里支持它的一处文字」之间的连线。 + * 它刻意是一个**快照**：文字、页码、块 id 在回答生成时一并写入，而不是以后再去 + * 查库。来源文档可以在回答产生之后被重新索引甚至删除，citation 必须在这两件事 + * 之后依然存在 —— 它降级展示，而不是变成一个悬空的 id。 + * + * `quote` 同样是快照：它是检索到的证据本身，不是指向 chunks 表的外键。重新分块 + * 或重新嵌入不能改变一条已经交付的回答里显示的原文。 + */ +export interface Citation { + /** 来源在 prompt 里的 1-based 位置，回答里的 `[n]` 据此解析。 */ + index: number + documentId: string + documentTitle: string + documentType?: string + /** 引用区间起始页；不分页的来源为 `undefined`。 */ + page?: number + /** 引用跨页时的结束页。 */ + pageEnd?: number + /** 跳转应落在的块（区间覆盖的第一个块）。 */ + blockId?: string + chunkId: string + /** 引用区间在 `documents.content` 里的字符偏移。 */ + startOffset?: number + endOffset?: number + /** 检索到的原文，逐字保留。 */ + quote: string + score: number +} diff --git a/src/shared/types/index.ts b/src/shared/types/index.ts index 0f7a1c2..a04d3e3 100644 --- a/src/shared/types/index.ts +++ b/src/shared/types/index.ts @@ -9,6 +9,9 @@ export * from './knowledge' // 导出聊天类型 export * from './chat' +// 导出引用类型 +export * from './citation' + // 导出 Result 错误处理类型 export * from './result' diff --git a/src/shared/utils/citations.ts b/src/shared/utils/citations.ts new file mode 100644 index 0000000..2039160 --- /dev/null +++ b/src/shared/utils/citations.ts @@ -0,0 +1,73 @@ +import type { Citation } from '../types/citation' +import type { ChatMessageMetadata } from '../types/chat' + +/** + * Readers for `chat_messages.metadata.citations`. + * + * The metadata column is an open JSON bag written by whichever version of the app + * produced the message, so a reader can never assume the shape. `parseCitations` + * is the only place that decides what is usable, so a component can render a + * message from a year ago without a guard of its own. + * + * Malformed entries are **dropped individually**: one unreadable citation from an + * older writer is not a reason to hide the citations that did survive. + * + * Dropping is not the same as deleting. A citation whose document no longer + * exists is still a valid, parseable record — the source was real when the answer + * was written. The UI disables the jump; the parser does not pretend it never + * happened. + */ + +const isNonEmptyString = (value: unknown): value is string => + typeof value === 'string' && value.length > 0 + +const toFiniteNumber = (value: unknown): number | undefined => + typeof value === 'number' && Number.isFinite(value) ? value : undefined + +const parseCitation = (value: unknown): Citation | null => { + if (!value || typeof value !== 'object') return null + const candidate = value as Record + + // The four fields a citation cannot be rendered or resolved without. + if (!isNonEmptyString(candidate.documentId)) return null + if (!isNonEmptyString(candidate.documentTitle)) return null + if (!isNonEmptyString(candidate.chunkId)) return null + if (!isNonEmptyString(candidate.quote)) return null + + const citation: Citation = { + index: toFiniteNumber(candidate.index) ?? 0, + documentId: candidate.documentId, + documentTitle: candidate.documentTitle, + chunkId: candidate.chunkId, + quote: candidate.quote, + // An unrecorded score is not "irrelevant"; it is unknown. 0 keeps the field a + // number without claiming anything a reader relies on. + score: toFiniteNumber(candidate.score) ?? 0 + } + + if (isNonEmptyString(candidate.documentType)) citation.documentType = candidate.documentType + if (isNonEmptyString(candidate.blockId)) citation.blockId = candidate.blockId + + const page = toFiniteNumber(candidate.page) + if (page !== undefined) citation.page = page + const pageEnd = toFiniteNumber(candidate.pageEnd) + if (pageEnd !== undefined) citation.pageEnd = pageEnd + const startOffset = toFiniteNumber(candidate.startOffset) + if (startOffset !== undefined) citation.startOffset = startOffset + const endOffset = toFiniteNumber(candidate.endOffset) + if (endOffset !== undefined) citation.endOffset = endOffset + + return citation +} + +/** Every usable citation on a message, ordered by the marker the prompt used. */ +export const parseCitations = (metadata: unknown): Citation[] => { + if (!metadata || typeof metadata !== 'object') return [] + const raw = (metadata as ChatMessageMetadata).citations + if (!Array.isArray(raw)) return [] + + return raw + .map(parseCitation) + .filter((citation): citation is Citation => citation !== null) + .sort((a, b) => a.index - b.index) +} diff --git a/test/citations.test.ts b/test/citations.test.ts new file mode 100644 index 0000000..2388ac8 --- /dev/null +++ b/test/citations.test.ts @@ -0,0 +1,257 @@ +import { test } from 'node:test' +import assert from 'node:assert/strict' +import { + buildCitations, + buildRAGContext, + citationFromSearchResult +} from '../src/main/services/citations.ts' +import { parseCitations } from '../src/shared/utils/citations.ts' +import type { SearchResult } from '../src/main/services/KnowledgeService.ts' + +/** + * A citation is the link between a claim in an answer and the place in the source + * that supports it (#69). It is assembled from the retriever's locator — never by + * querying the database again — and read back defensively, because the metadata + * column is written by whichever version of the app produced the message. + */ + +const searchResult = (over: Partial = {}): SearchResult => ({ + chunkId: 'chunk_1', + documentId: 'doc_1', + documentTitle: 'Attention Is All You Need', + documentType: 'file', + content: 'the retrieved passage', + score: 0.87, + chunkIndex: 4, + locator: { + pageStart: 5, + pageEnd: 5, + blocks: [ + { + blockId: 'blk_doc_1_7', + kind: 'paragraph', + level: null, + page: 5, + bbox: { x: 0.1, y: 0.2, w: 0.8, h: 0.05 }, + text: 'the retrieved passage', + startOffset: 100, + endOffset: 121, + startInBlock: 0, + endInBlock: 21 + } + ] + }, + ...over +}) + +test('a citation carries the page, block and char span of the retrieved chunk', () => { + const citation = citationFromSearchResult(searchResult(), 2) + + assert.deepEqual(citation, { + index: 2, + documentId: 'doc_1', + documentTitle: 'Attention Is All You Need', + documentType: 'file', + page: 5, + pageEnd: 5, + blockId: 'blk_doc_1_7', + chunkId: 'chunk_1', + startOffset: 100, + endOffset: 121, + quote: 'the retrieved passage', + score: 0.87 + }) +}) + +test('a chunk that starts mid-block keeps the exact span, not the block bounds', () => { + const citation = citationFromSearchResult( + searchResult({ + locator: { + pageStart: 3, + pageEnd: 3, + blocks: [ + { + blockId: 'b1', + kind: 'paragraph', + level: null, + page: 3, + bbox: null, + text: 'aaaaaaaaaaaaaaaaaaaa', + startOffset: 50, + endOffset: 70, + startInBlock: 5, + endInBlock: 12 + } + ] + } + }), + 1 + ) + + assert.equal(citation.startOffset, 55) + assert.equal(citation.endOffset, 62) +}) + +test('a chunk spanning several blocks anchors on the first and spans to the last', () => { + const citation = citationFromSearchResult( + searchResult({ + locator: { + pageStart: 5, + pageEnd: 6, + blocks: [ + { + blockId: 'b1', + kind: 'paragraph', + level: null, + page: 5, + bbox: null, + text: 'first', + startOffset: 100, + endOffset: 105, + startInBlock: 0, + endInBlock: 5 + }, + { + blockId: 'b2', + kind: 'paragraph', + level: null, + page: 6, + bbox: null, + text: 'second', + startOffset: 200, + endOffset: 206, + startInBlock: 0, + endInBlock: 6 + } + ] + } + }), + 1 + ) + + assert.equal(citation.blockId, 'b1') + assert.equal(citation.startOffset, 100) + assert.equal(citation.endOffset, 206) + assert.equal(citation.page, 5) + assert.equal(citation.pageEnd, 6) +}) + +test('an unpaginated source produces a citation without page or block fields', () => { + const citation = citationFromSearchResult( + searchResult({ locator: { pageStart: null, pageEnd: null, blocks: [] } }), + 1 + ) + + assert.equal(citation.page, undefined) + assert.equal(citation.pageEnd, undefined) + assert.equal(citation.blockId, undefined) + assert.equal(citation.startOffset, undefined) + assert.equal(citation.endOffset, undefined) +}) + +test('citations are numbered by their position in the prompt', () => { + const citations = buildCitations([searchResult(), searchResult({ chunkId: 'chunk_2' })]) + assert.deepEqual( + citations.map((citation) => citation.index), + [1, 2] + ) + assert.deepEqual( + citations.map((citation) => citation.chunkId), + ['chunk_1', 'chunk_2'] + ) +}) + +test('an answer without retrieval produces no citations', () => { + const { context, sources, citations } = buildRAGContext([]) + assert.equal(context, '') + assert.deepEqual(sources, []) + assert.deepEqual(citations, []) +}) + +test('the prompt asks the model to mark sources so a claim can be resolved', () => { + const { context, sources, citations } = buildRAGContext([searchResult()]) + + assert.match(context, /\[来源 1: Attention Is All You Need\]/) + assert.match(context, /\[n\]/) + assert.equal(sources.length, 1) + assert.equal(citations.length, 1) + assert.equal(citations[0].score, 0.87) +}) + +/** + * The reader half. A citation is a snapshot: it must still parse after the source + * document has been deleted or re-indexed, and one bad row must not hide the good + * ones. + */ + +test('a citation survives a metadata round trip', () => { + const [citation] = parseCitations({ + citations: [citationFromSearchResult(searchResult(), 1)] + }) + + assert.equal(citation.documentId, 'doc_1') + assert.equal(citation.page, 5) + assert.equal(citation.blockId, 'blk_doc_1_7') + assert.equal(citation.startOffset, 100) + assert.equal(citation.endOffset, 121) +}) + +test('a citation whose document no longer exists still parses', () => { + // Deleting the document does not rewrite history: the answer was grounded when + // it was written. The UI disables the jump; the parser keeps the record. + const [citation] = parseCitations({ + citations: [citationFromSearchResult(searchResult(), 1)] + }) + assert.equal(citation.documentTitle, 'Attention Is All You Need') +}) + +test('a citation missing a field it cannot resolve without is dropped', () => { + for (const missing of ['documentId', 'documentTitle', 'chunkId', 'quote']) { + const broken = citationFromSearchResult(searchResult(), 1) as unknown as Record + delete broken[missing] + assert.deepEqual(parseCitations({ citations: [broken] }), [], missing) + } +}) + +test('a malformed citation is dropped without discarding its siblings', () => { + const parsed = parseCitations({ + citations: [ + citationFromSearchResult(searchResult(), 1), + { nope: true }, + citationFromSearchResult(searchResult({ chunkId: 'chunk_2' }), 2) + ] + }) + assert.equal(parsed.length, 2) + assert.deepEqual( + parsed.map((citation) => citation.chunkId), + ['chunk_1', 'chunk_2'] + ) +}) + +test('citations come back ordered by the marker the prompt used', () => { + const parsed = parseCitations({ + citations: [ + citationFromSearchResult(searchResult({ chunkId: 'c3' }), 3), + citationFromSearchResult(searchResult({ chunkId: 'c1' }), 1), + citationFromSearchResult(searchResult({ chunkId: 'c2' }), 2) + ] + }) + assert.deepEqual( + parsed.map((citation) => citation.index), + [1, 2, 3] + ) +}) + +test('metadata that is absent, null or the wrong shape yields no citations', () => { + for (const metadata of [ + undefined, + null, + 'text', + 42, + {}, + { citations: null }, + { citations: 'x' } + ]) { + assert.deepEqual(parseCitations(metadata), [], JSON.stringify(metadata)) + } +})