fix(chat): bound PDF context with retrieved chunks#1240
Open
lzhs1995 wants to merge 3 commits into
Open
Conversation
There was a problem hiding this comment.
Pull request overview
This PR improves Chat 模式下“文档附件（如 PDF）注入提示词”的策略：对抽取出的文档文本做分片缓存，并在发送前按 query 相关性检索片段，控制文档上下文预算，同时在整体提示词过大时给出明确的预检错误，避免模型侧静默截断或失败。
Changes:
- 新增
document-context.ts：实现文档分片、相关片段检索、以及基于字符长度的 token 估算。 - 更新
chat-service.ts：用户当前消息注入“检索到的文档片段”，历史消息仅保留附件引用，并增加 180K token 预检上限错误。 - 新增
document-context.test.ts：覆盖分片与预算上限的基础测试。
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| apps/electron/src/main/lib/document-context.ts | 新增文档分片缓存、query 相关片段检索与 token 粗估逻辑 |
| apps/electron/src/main/lib/document-context.test.ts | 新增对分片与 60K token 预算的单元测试 |
| apps/electron/src/main/lib/chat-service.ts | 使用检索片段注入文档上下文；历史仅保留引用；增加 180K token 预检报错 |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+9
to
+26
| const chunkCache = new Map<string, string[]>() | ||
|
|
||
| function documentKey(text: string): string { | ||
| return createHash('sha256').update(text).digest('hex') | ||
| } | ||
|
|
||
| export function chunkDocument(text: string): string[] { | ||
| const key = documentKey(text) | ||
| const cached = chunkCache.get(key) | ||
| if (cached) return cached | ||
| const normalized = text.replace(/\r\n/g, '\n').trim() | ||
| const chunks: string[] = [] | ||
| for (let start = 0; start < normalized.length; start += CHUNK_CHARS - CHUNK_OVERLAP_CHARS) { | ||
| chunks.push(normalized.slice(start, start + CHUNK_CHARS)) | ||
| } | ||
| chunkCache.set(key, chunks) | ||
| return chunks | ||
| } |
Comment on lines
+32
to
+50
| export function retrieveDocumentContext(text: string, query: string): string { | ||
| const terms = queryTerms(query) | ||
| const ranked = chunkDocument(text).map((chunk, index) => { | ||
| const lower = chunk.toLowerCase() | ||
| let score = 0 | ||
| for (const term of terms) if (lower.includes(term)) score++ | ||
| return { chunk, index, score } | ||
| }).sort((a, b) => b.score - a.score || a.index - b.index) | ||
|
|
||
| const selected: string[] = [] | ||
| let chars = 0 | ||
| for (const item of ranked) { | ||
| if (chars + item.chunk.length > DOCUMENT_CONTEXT_MAX_CHARS) continue | ||
| selected.push(`[片段 ${item.index + 1}]\n${item.chunk}`) | ||
| chars += item.chunk.length | ||
| if (chars >= DOCUMENT_CONTEXT_MAX_CHARS) break | ||
| } | ||
| return selected.join('\n\n') | ||
| } |
Comment on lines
88
to
91
| if (text.trim()) { | ||
| parts.push(`\n<file name="${att.filename}">\n${text}\n</file>`) | ||
| const context = retrieveDocumentContext(text, messageText) | ||
| parts.push(`\n<file name="${att.filename}" mode="retrieved-chunks">\n${context}\n</file>`) | ||
| } else { |
| if (estimatedPromptTokens > 180_000) { | ||
| webContents.send(CHAT_IPC_CHANNELS.STREAM_ERROR, { | ||
| conversationId, | ||
| error: `当前对话预计 ${estimatedPromptTokens.toLocaleString()} tokens，超过 Kiro 安全输入上限 180,000。请新建会话、缩小 PDF 范围或使用文档摘要。`, |
Comment on lines
+5
to
+7
| export const DOCUMENT_CONTEXT_MAX_TOKENS = 60_000 | ||
| // 预留约 2% 给片段编号、XML 包装及 tokenizer 误差。 | ||
| const DOCUMENT_CONTEXT_MAX_CHARS = Math.floor(DOCUMENT_CONTEXT_MAX_TOKENS * 4 * 0.98) |
Author
|
补充本轮可靠性验证：\n\n- PDF 原生 document + 超限 RAG 双路径已完成。\n- 多文档共享 60K RAG 预算，SHA256 去重，完整 prompt 180K 安全预算。\n- 原生 PDF：单份提取文本 ≤80K tokens、合计 ≤120K；单文件 ≤8 MiB、合计 ≤16 MiB。\n- 全量 Bun：414/414 tests passed。\n- 全项目 TypeScript typecheck passed。\n- 修复了 Electron 测试间不完整 mock 污染，原先 404 pass / 2 fail 已清零。\n\n该 PR 仍需与 A2/NewAPI 候选镜像一起完成真实 PDF、连续追问和 60 分钟 soak 后才进入生产灰度。 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
AI-assisted implementation; manually verified with deterministic tests.
Tests