diff --git a/apps/electron/package.json b/apps/electron/package.json index d0e1ee88b..c38cba5f8 100644 --- a/apps/electron/package.json +++ b/apps/electron/package.json @@ -1,6 +1,6 @@ { "name": "@proma/electron", - "version": "0.15.7", + "version": "0.15.8", "description": "Proma next gen ai software with general agents - Electron App", "main": "dist/main.cjs", "author": { @@ -38,6 +38,7 @@ "dist:debug": "bun run scripts/dist.ts --current-arch --verbose" }, "dependencies": { + "@anthropic-ai/tokenizer": "0.0.4", "@anthropic-ai/claude-agent-sdk": "0.3.201", "@anthropic-ai/sdk": "^0.93.0", "@earendil-works/pi-agent-core": "0.80.9", diff --git a/apps/electron/src/main/lib/agent-session-manager.test.ts b/apps/electron/src/main/lib/agent-session-manager.test.ts index d5f531dc5..93fb47e78 100644 --- a/apps/electron/src/main/lib/agent-session-manager.test.ts +++ b/apps/electron/src/main/lib/agent-session-manager.test.ts @@ -2,6 +2,7 @@ import { afterAll, beforeAll, describe, expect, mock, test } from 'bun:test' import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs' import * as os from 'node:os' import { join } from 'node:path' +import { createElectronMock } from './test-helpers/electron-mock' type AgentSessionManager = typeof import('./agent-session-manager') @@ -11,26 +12,9 @@ const originalHome = process.env.HOME const originalPromaDev = process.env.PROMA_DEV const originalClaudeConfigDir = process.env.CLAUDE_CONFIG_DIR -mock.module('electron', () => ({ - app: { - isPackaged: true, - getPath: () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'), - }, - BrowserWindow: class {}, - clipboard: {}, - dialog: {}, - nativeImage: { createFromPath: () => ({}) }, - nativeTheme: {}, - powerMonitor: {}, - powerSaveBlocker: {}, - screen: {}, - shell: {}, - safeStorage: { - isEncryptionAvailable: () => false, - encryptString: (value: string) => Buffer.from(value), - decryptString: (value: Buffer) => value.toString('utf-8'), - }, -})) +mock.module('electron', () => createElectronMock( + () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'), +)) mock.module('node:os', () => ({ ...os, diff --git a/apps/electron/src/main/lib/agent-workspace-manager.test.ts b/apps/electron/src/main/lib/agent-workspace-manager.test.ts index d1012602e..cce596421 100644 --- a/apps/electron/src/main/lib/agent-workspace-manager.test.ts +++ b/apps/electron/src/main/lib/agent-workspace-manager.test.ts @@ -2,6 +2,7 @@ import { afterAll, beforeAll, beforeEach, describe, expect, mock, test } from 'b import { existsSync, mkdirSync, mkdtempSync, readdirSync, rmSync, symlinkSync, writeFileSync } from 'node:fs' import * as os from 'node:os' import { join } from 'node:path' +import { createElectronMock } from './test-helpers/electron-mock' type AgentWorkspaceManager = typeof import('./agent-workspace-manager') type ConfigPathsModule = typeof import('./config-paths') @@ -12,17 +13,9 @@ let tempHome: string const originalHome = process.env.HOME const originalPromaDev = process.env.PROMA_DEV -mock.module('electron', () => ({ - app: { - isPackaged: true, - getPath: () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'), - }, - safeStorage: { - isEncryptionAvailable: () => false, - encryptString: (value: string) => Buffer.from(value), - decryptString: (value: Buffer) => value.toString('utf-8'), - }, -})) +mock.module('electron', () => createElectronMock( + () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'), +)) mock.module('node:os', () => ({ ...os, diff --git a/apps/electron/src/main/lib/channel-runtime-api-key.test.ts b/apps/electron/src/main/lib/channel-runtime-api-key.test.ts index 607aac280..18bdf17cd 100644 --- a/apps/electron/src/main/lib/channel-runtime-api-key.test.ts +++ b/apps/electron/src/main/lib/channel-runtime-api-key.test.ts @@ -3,6 +3,7 @@ import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs' import * as os from 'node:os' import { join } from 'node:path' import { serializeCodexCredentials } from '@proma/shared' +import { createElectronMock } from './test-helpers/electron-mock' type ChannelManagerModule = typeof import('./channel-manager') @@ -11,20 +12,9 @@ let tempHome: string const originalHome = process.env.HOME const originalPromaDev = process.env.PROMA_DEV -mock.module('electron', () => ({ - app: { - isPackaged: true, - getPath: () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'), - }, - safeStorage: { - isEncryptionAvailable: () => false, - encryptString: (value: string) => Buffer.from(value), - decryptString: (value: Buffer) => value.toString('utf-8'), - }, - shell: { - openExternal: async () => undefined, - }, -})) +mock.module('electron', () => createElectronMock( + () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'), +)) mock.module('node:os', () => ({ ...os, diff --git a/apps/electron/src/main/lib/chat-service.ts b/apps/electron/src/main/lib/chat-service.ts index 3f36d5f4a..6ab304156 100644 --- a/apps/electron/src/main/lib/chat-service.ts +++ b/apps/electron/src/main/lib/chat-service.ts @@ -21,11 +21,17 @@ import { streamSSE, fetchTitle, } from '@proma/core' -import type { ImageAttachmentData, ContinuationMessage } from '@proma/core' +import type { ImageAttachmentData, DocumentAttachmentData, ContinuationMessage } from '@proma/core' import { listChannels, resolveChannelRuntimeApiKey } from './channel-manager' import { appendMessage, updateConversationMeta, getConversationMessages } from './conversation-manager' import { readAttachmentAsBase64, isImageAttachment } from './attachment-service' import { extractTextFromAttachment, isDocumentAttachment } from './document-parser' +import { + estimateDocumentTokensConservatively, + estimatePromptTokens, + PROMPT_SAFE_INPUT_TOKENS, + retrieveDocumentContexts, +} from './document-context' import { getFetchFn } from './proxy-fetch' import { getEffectiveProxyUrl } from './proxy-settings-service' import { getEnabledTools } from './chat-tool-registry' @@ -38,6 +44,18 @@ const activeControllers = new Map() /** 最大工具续接轮数（安全上限，防止极端情况下的无限循环） */ const MAX_TOOL_ROUNDS = 999 +const NATIVE_DOCUMENT_MAX_FILE_BYTES = 8 * 1024 * 1024 +const NATIVE_DOCUMENT_MAX_TOTAL_BYTES = 16 * 1024 * 1024 +const NATIVE_DOCUMENT_MAX_TEXT_TOKENS = 80_000 +const NATIVE_DOCUMENT_MAX_TOTAL_TEXT_TOKENS = 120_000 + +interface PreparedDocuments { + nativeDocumentIds: Set + nativeEstimatedTokens: number + retrievalSources: Array<{ filename: string; text: string }> + notices: string[] +} + // ===== 平台相关：图片附件读取器 ===== /** @@ -57,6 +75,57 @@ function getImageAttachmentData(attachments?: FileAttachment[]): ImageAttachment })) } +function getDocumentAttachmentData( + attachments: FileAttachment[] | undefined, + nativeDocumentIds: Set, +): DocumentAttachmentData[] { + return (attachments ?? []).filter((attachment) => nativeDocumentIds.has(attachment.id)).map((attachment) => ({ + filename: attachment.filename, + mediaType: attachment.mediaType, + data: readAttachmentAsBase64(attachment.localPath), + })) +} + +async function prepareDocuments(attachments?: FileAttachment[]): Promise { + const prepared: PreparedDocuments = { + nativeDocumentIds: new Set(), + nativeEstimatedTokens: 0, + retrievalSources: [], + notices: [], + } + if (!attachments || attachments.length === 0) return prepared + + let nativeBytes = 0 + for (const attachment of attachments.filter((item) => isDocumentAttachment(item.mediaType))) { + try { + const text = await extractTextFromAttachment(attachment.localPath) + const normalizedText = text.trim() + const estimatedTokens = normalizedText ? estimateDocumentTokensConservatively(normalizedText) : 0 + const canUseNativePdf = attachment.mediaType === 'application/pdf' + && normalizedText.length > 0 + && attachment.size <= NATIVE_DOCUMENT_MAX_FILE_BYTES + && nativeBytes + attachment.size <= NATIVE_DOCUMENT_MAX_TOTAL_BYTES + && estimatedTokens <= NATIVE_DOCUMENT_MAX_TEXT_TOKENS + && prepared.nativeEstimatedTokens + estimatedTokens <= NATIVE_DOCUMENT_MAX_TOTAL_TEXT_TOKENS + + if (canUseNativePdf) { + prepared.nativeDocumentIds.add(attachment.id) + prepared.nativeEstimatedTokens += estimatedTokens + nativeBytes += attachment.size + } else if (normalizedText) { + prepared.retrievalSources.push({ filename: attachment.filename, text: normalizedText }) + } else { + prepared.notices.push(`[文档 ${attachment.filename} 内容为空]`) + } + } catch (error) { + const errorMsg = error instanceof Error ? error.message : '未知错误' + console.warn(`[聊天服务] 文档提取失败: ${attachment.filename}`, error) + prepared.notices.push(`[文档 ${attachment.filename} 内容提取失败: ${errorMsg}]`) + } + } + return prepared +} + // ===== 文档附件文本提取 ===== /** @@ -71,38 +140,19 @@ function getImageAttachmentData(attachments?: FileAttachment[]): ImageAttachment */ async function enrichMessageWithDocuments( messageText: string, - attachments?: FileAttachment[], + prepared: PreparedDocuments, ): Promise { - if (!attachments || attachments.length === 0) return messageText - - // 筛选出文档类附件（非图片） - const docAttachments = attachments.filter((att) => isDocumentAttachment(att.mediaType)) - if (docAttachments.length === 0) return messageText - - const parts: string[] = [messageText] - - for (const att of docAttachments) { - try { - const text = await extractTextFromAttachment(att.localPath) - if (text.trim()) { - parts.push(`\n\n${text}\n`) - } else { - parts.push(`\n\n[文件内容为空]\n`) - } - } catch (error) { - const errorMsg = error instanceof Error ? error.message : '未知错误' - console.warn(`[聊天服务] 文档提取失败: ${att.filename}`, error) - parts.push(`\n\n[文件内容提取失败: ${errorMsg}]\n`) - } - } - - return parts.join('') + const context = retrieveDocumentContexts(prepared.retrievalSources, messageText) + const sections = [messageText] + if (context.content) sections.push(context.content) + if (prepared.notices.length > 0) sections.push(prepared.notices.join('\n')) + return sections.filter(Boolean).join('\n\n') } /** * 为历史消息列表注入文档附件文本 * - * 遍历历史消息，对包含文档附件的用户消息进行文本增强。 + * 历史附件只保留引用，禁止每一轮重新注入整篇文档。 * 返回新的消息数组（不修改原始消息）。 */ async function enrichHistoryWithDocuments( @@ -115,8 +165,8 @@ async function enrichHistoryWithDocuments( if (msg.role === 'user' && msg.attachments && msg.attachments.length > 0) { const hasDocuments = msg.attachments.some((att) => isDocumentAttachment(att.mediaType)) if (hasDocuments) { - const enrichedContent = await enrichMessageWithDocuments(msg.content, msg.attachments) - enriched.push({ ...msg, content: enrichedContent }) + const names = msg.attachments.filter((att) => isDocumentAttachment(att.mediaType)).map((att) => att.filename) + enriched.push({ ...msg, content: `${msg.content}\n[历史附件引用：${names.join('、')}]` }) continue } } @@ -250,7 +300,21 @@ export async function sendMessage( // 5. 过滤历史并提取文档附件文本 const filteredHistory = filterHistory(fullHistory, contextDividers, contextLength) const enrichedHistory = await enrichHistoryWithDocuments(filteredHistory) - const enrichedUserMessage = await enrichMessageWithDocuments(userMessage, attachments) + const preparedDocuments = await prepareDocuments(attachments) + const enrichedUserMessage = await enrichMessageWithDocuments(userMessage, preparedDocuments) + const estimatedPromptTokens = estimatePromptTokens([ + systemMessage ?? '', + ...enrichedHistory.map((message) => message.content), + enrichedUserMessage, + ]) + preparedDocuments.nativeEstimatedTokens + // 留出模型回复和工具调用空间。超限必须由客户端明确提示，不能让 A2 静默截断。 + if (estimatedPromptTokens > PROMPT_SAFE_INPUT_TOKENS) { + webContents.send(CHAT_IPC_CHANNELS.STREAM_ERROR, { + conversationId, + error: `当前对话预计 ${estimatedPromptTokens.toLocaleString()} tokens，超过快速稳定模式的安全输入预算 ${PROMPT_SAFE_INPUT_TOKENS.toLocaleString()}。请新建会话、缩小文档范围或先生成文档摘要。`, + }) + return + } // 6. 创建 AbortController const controller = new AbortController() @@ -330,6 +394,10 @@ export async function sendMessage( systemMessage: effectiveSystemMessage, attachments, readImageAttachments: getImageAttachmentData, + readDocumentAttachments: (requestAttachments) => getDocumentAttachmentData( + requestAttachments, + preparedDocuments.nativeDocumentIds, + ), thinkingEnabled, tools, continuationMessages: continuationMessages.length > 0 ? continuationMessages : undefined, @@ -406,6 +474,10 @@ export async function sendMessage( systemMessage: effectiveSystemMessage, attachments, readImageAttachments: getImageAttachmentData, + readDocumentAttachments: (requestAttachments) => getDocumentAttachmentData( + requestAttachments, + preparedDocuments.nativeDocumentIds, + ), thinkingEnabled, // 不传 tools，强制模型生成文本回复而非继续调用工具 continuationMessages, diff --git a/apps/electron/src/main/lib/dingtalk-config.test.ts b/apps/electron/src/main/lib/dingtalk-config.test.ts index 3ec9e0ece..a900b6291 100644 --- a/apps/electron/src/main/lib/dingtalk-config.test.ts +++ b/apps/electron/src/main/lib/dingtalk-config.test.ts @@ -3,6 +3,7 @@ import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync import * as os from 'node:os' import { join } from 'node:path' import { createStableDingTalkBotId } from './dingtalk-bot-identity' +import { createElectronMock } from './test-helpers/electron-mock' type DingTalkConfigModule = typeof import('./dingtalk-config') type ConfigPathsModule = typeof import('./config-paths') @@ -13,17 +14,9 @@ let tempHome: string const originalHome = process.env.HOME const originalPromaDev = process.env.PROMA_DEV -mock.module('electron', () => ({ - app: { - isPackaged: true, - getPath: () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'), - }, - safeStorage: { - isEncryptionAvailable: () => false, - encryptString: (value: string) => Buffer.from(value), - decryptString: (value: Buffer) => value.toString('utf-8'), - }, -})) +mock.module('electron', () => createElectronMock( + () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'), +)) mock.module('node:os', () => ({ ...os, diff --git a/apps/electron/src/main/lib/document-context.test.ts b/apps/electron/src/main/lib/document-context.test.ts new file mode 100644 index 000000000..99ed5114f --- /dev/null +++ b/apps/electron/src/main/lib/document-context.test.ts @@ -0,0 +1,55 @@ +import { describe, expect, test } from 'bun:test' +import { + chunkDocument, + DOCUMENT_CONTEXT_MAX_TOKENS, + estimatePromptTokens, + retrieveDocumentContexts, +} from './document-context' + +describe('文档上下文预算', () => { + test('长文档按保守 token 预算重叠切分并检索中文相关内容', () => { + const text = `${'普通内容 '.repeat(4_000)}关键结论：财政政策有效。${'尾部 '.repeat(4_000)}` + const chunks = chunkDocument(text) + const result = retrieveDocumentContexts([{ filename: '财政报告.pdf', text }], '财政政策') + + expect(chunks.length).toBeGreaterThan(1) + expect(chunks.every((chunk) => chunk.tokens <= 1_100)).toBeTrue() + expect(result.content).toContain('财政政策有效') + }) + + test('多份文档共享六万 token 总预算', () => { + const result = retrieveDocumentContexts([ + { filename: '甲.pdf', text: `甲文档 财政政策 ${'研究资料 '.repeat(20_000)}` }, + { filename: '乙.pdf', text: `乙文档 货币政策 ${'研究资料 '.repeat(20_000)}` }, + ], '比较财政政策与货币政策') + + expect(result.includedDocuments).toEqual(['甲.pdf', '乙.pdf']) + expect(result.estimatedTokens).toBeLessThanOrEqual(DOCUMENT_CONTEXT_MAX_TOKENS) + }) + + test('相同文档只注入一次', () => { + const text = `重复文档 ${'内容 '.repeat(2_000)}` + const result = retrieveDocumentContexts([ + { filename: '原件.pdf', text }, + { filename: '副本.pdf', text }, + ], '重复文档') + + expect(result.duplicateDocumentsSkipped).toBe(1) + expect(result.includedDocuments).toEqual(['原件.pdf']) + }) + + test('使用 Anthropic tokenizer 估算完整 prompt', () => { + expect(estimatePromptTokens(['hello world', '财政政策'])).toBeGreaterThan(0) + }) + + test('百万字符文档分块保持线性性能且不调用全文 tokenizer', () => { + const text = '财政政策与货币政策。'.repeat(100_000) + const startedAt = performance.now() + const chunks = chunkDocument(text) + const elapsedMs = performance.now() - startedAt + + expect(chunks.length).toBeGreaterThan(1_000) + expect(chunks.every((chunk) => chunk.tokens <= 1_100)).toBeTrue() + expect(elapsedMs).toBeLessThan(5_000) + }) +}) diff --git a/apps/electron/src/main/lib/document-context.ts b/apps/electron/src/main/lib/document-context.ts new file mode 100644 index 000000000..87855ef83 --- /dev/null +++ b/apps/electron/src/main/lib/document-context.ts @@ -0,0 +1,302 @@ +import { createHash } from 'node:crypto' +import { countTokens } from '@anthropic-ai/tokenizer' + +const CHUNK_MAX_ESTIMATED_TOKENS = 1_100 +const CHUNK_OVERLAP_ESTIMATED_TOKENS = 140 +const MAX_CACHED_DOCUMENTS = 20 + +export const DOCUMENT_CONTEXT_MAX_TOKENS = 60_000 +export const PROMPT_SAFE_INPUT_TOKENS = 180_000 + +export interface DocumentContextSource { + filename: string + text: string +} + +export interface DocumentContextResult { + content: string + estimatedTokens: number + includedDocuments: string[] + duplicateDocumentsSkipped: number +} + +interface DocumentChunk { + index: number + text: string + tokens: number +} + +interface RankedChunk extends DocumentChunk { + documentKey: string + filename: string + score: number +} + +const chunkCache = new Map() + +function normalizeDocument(text: string): string { + return text.replace(/\r\n/g, '\n').replace(/[ \t]+\n/g, '\n').trim() +} + +function documentKey(text: string): string { + return createHash('sha256').update(text).digest('hex') +} + +function escapeXmlAttribute(value: string): string { + return value.replace(/&/g, '&').replace(/"/g, '"').replace(//g, '>') +} + +/** + * 文档全文只做线性、保守估算。官方 tokenizer 仅用于最终入选的最多 60K 上下文， + * 避免对百万字符 PDF 反复分片计数造成 CPU/内存尖峰。 + */ +function conservativeTokenWeight(code: number): number { + if (code <= 0x20) return 0.12 + if ((code >= 0x30 && code <= 0x39) || (code >= 0x41 && code <= 0x5a) || (code >= 0x61 && code <= 0x7a)) { + return 0.34 + } + if (code <= 0x7f) return 0.65 + if ( + (code >= 0x3400 && code <= 0x4dbf) + || (code >= 0x4e00 && code <= 0x9fff) + || (code >= 0xf900 && code <= 0xfaff) + ) { + return 1.85 + } + return 1.35 +} + +export function estimateDocumentTokensConservatively(text: string): number { + let tokens = 0 + for (let index = 0; index < text.length; index++) { + const code = text.charCodeAt(index) + if (code >= 0xd800 && code <= 0xdbff && index + 1 < text.length) { + tokens += 2.5 + index++ + continue + } + tokens += conservativeTokenWeight(code) + } + return Math.ceil(tokens) +} + +function scanChunkEnd(text: string, start: number, budget: number): number { + let tokens = 0 + let end = start + while (end < text.length) { + const code = text.charCodeAt(end) + const isSurrogatePair = code >= 0xd800 && code <= 0xdbff && end + 1 < text.length + const weight = isSurrogatePair ? 2.5 : conservativeTokenWeight(code) + if (end > start && tokens + weight > budget) break + tokens += weight + end += isSurrogatePair ? 2 : 1 + } + return Math.max(start + 1, end) +} + +function preferSentenceBoundary(text: string, start: number, end: number): number { + if (end >= text.length) return end + const minimumBoundary = start + Math.floor((end - start) * 0.65) + const boundary = Math.max( + text.lastIndexOf('\n', end), + text.lastIndexOf('。', end), + text.lastIndexOf('！', end), + text.lastIndexOf('？', end), + text.lastIndexOf('. ', end), + ) + return boundary >= minimumBoundary ? boundary + 1 : end +} + +function findOverlapStart(text: string, chunkStart: number, chunkEnd: number): number { + let tokens = 0 + let start = chunkEnd + while (start > chunkStart + 1) { + const previous = start - 1 + const code = text.charCodeAt(previous) + const isLowSurrogate = code >= 0xdc00 && code <= 0xdfff && previous > chunkStart + const nextStart = isLowSurrogate ? previous - 1 : previous + const weight = isLowSurrogate ? 2.5 : conservativeTokenWeight(code) + if (tokens + weight > CHUNK_OVERLAP_ESTIMATED_TOKENS) break + tokens += weight + start = nextStart + } + return start +} + +export function chunkDocument(text: string): DocumentChunk[] { + const normalized = normalizeDocument(text) + if (!normalized) return [] + + const key = documentKey(normalized) + const cached = chunkCache.get(key) + if (cached) return cached + + const chunks: DocumentChunk[] = [] + let start = 0 + while (start < normalized.length) { + const scannedEnd = scanChunkEnd(normalized, start, CHUNK_MAX_ESTIMATED_TOKENS) + const end = preferSentenceBoundary(normalized, start, scannedEnd) + const chunkText = normalized.slice(start, end).trim() + if (chunkText) { + chunks.push({ + index: chunks.length, + text: chunkText, + tokens: estimateDocumentTokensConservatively(chunkText), + }) + } + if (end >= normalized.length) break + start = findOverlapStart(normalized, start, end) + } + + if (chunkCache.size >= MAX_CACHED_DOCUMENTS) { + const oldest = chunkCache.keys().next().value + if (oldest) chunkCache.delete(oldest) + } + chunkCache.set(key, chunks) + return chunks +} + +function queryTerms(query: string): string[] { + const normalized = query.toLowerCase() + const terms = new Set(normalized.match(/[a-z0-9_-]{2,}/g) ?? []) + for (const sequence of normalized.match(/[\p{Script=Han}]{2,}/gu) ?? []) { + terms.add(sequence) + for (const size of [4, 3, 2]) { + for (let index = 0; index + size <= sequence.length && terms.size < 200; index++) { + terms.add(sequence.slice(index, index + size)) + } + } + } + return [...terms].slice(0, 200) +} + +function countOccurrences(text: string, term: string): number { + let count = 0 + let position = 0 + while (count < 8) { + position = text.indexOf(term, position) + if (position < 0) break + count++ + position += term.length + } + return count +} + +function rankDocumentChunks(source: DocumentContextSource, terms: string[]): RankedChunk[] { + const normalized = normalizeDocument(source.text) + const key = documentKey(normalized) + const filename = source.filename + const filenameLower = filename.toLowerCase() + + return chunkDocument(normalized).map((chunk) => { + const lower = chunk.text.toLowerCase() + let score = 0 + for (const term of terms) { + const occurrences = countOccurrences(lower, term) + if (occurrences > 0) score += 1 + Math.log2(occurrences + 1) + if (filenameLower.includes(term)) score += 2 + } + return { ...chunk, documentKey: key, filename, score } + }).sort((a, b) => b.score - a.score || a.index - b.index) +} + +function renderSelectedChunks(chunks: RankedChunk[]): string { + const byDocument = new Map() + for (const chunk of chunks) { + const current = byDocument.get(chunk.documentKey) ?? [] + current.push(chunk) + byDocument.set(chunk.documentKey, current) + } + + const rendered: string[] = [] + for (const documentChunks of byDocument.values()) { + documentChunks.sort((a, b) => a.index - b.index) + const filename = escapeXmlAttribute(documentChunks[0]?.filename ?? 'document') + const body = documentChunks + .map((chunk) => `\n${chunk.text}\n`) + .join('\n') + rendered.push(`\n${body}\n`) + } + return rendered.join('\n\n') +} + +export function retrieveDocumentContexts( + sources: DocumentContextSource[], + query: string, + maxTokens = DOCUMENT_CONTEXT_MAX_TOKENS, +): DocumentContextResult { + const uniqueSources: Array = [] + const seen = new Set() + let duplicateDocumentsSkipped = 0 + + for (const source of sources) { + const normalized = normalizeDocument(source.text) + if (!normalized) continue + const key = documentKey(normalized) + if (seen.has(key)) { + duplicateDocumentsSkipped++ + continue + } + seen.add(key) + uniqueSources.push({ ...source, text: normalized, key }) + } + + const terms = queryTerms(query) + const rankedByDocument = uniqueSources.map((source) => rankDocumentChunks(source, terms)) + const selected: RankedChunk[] = [] + const selectedIds = new Set() + const usedByDocument = new Map() + const perDocumentLimit = uniqueSources.length > 1 ? Math.floor(maxTokens / 2) : maxTokens + let usedTokens = 0 + + const trySelect = (chunk: RankedChunk): boolean => { + const id = `${chunk.documentKey}:${chunk.index}` + if (selectedIds.has(id)) return false + const documentUsed = usedByDocument.get(chunk.documentKey) ?? 0 + const overhead = 64 + if (documentUsed + chunk.tokens + overhead > perDocumentLimit) return false + if (usedTokens + chunk.tokens + overhead > maxTokens) return false + selected.push(chunk) + selectedIds.add(id) + usedByDocument.set(chunk.documentKey, documentUsed + chunk.tokens + overhead) + usedTokens += chunk.tokens + overhead + return true + } + + for (const ranked of rankedByDocument) { + const best = ranked.find((chunk) => chunk.score > 0) ?? ranked[0] + if (best) trySelect(best) + } + + const remaining = rankedByDocument.flat().sort((a, b) => b.score - a.score || a.index - b.index) + for (const chunk of remaining) trySelect(chunk) + + let content = renderSelectedChunks(selected) + let exactTokens = countTokens(content) + for (let attempt = 0; selected.length > 0 && exactTokens > maxTokens && attempt < 4; attempt++) { + const keepCount = Math.max(0, Math.floor(selected.length * maxTokens / exactTokens * 0.96)) + selected.splice(Math.min(keepCount, selected.length - 1)) + content = renderSelectedChunks(selected) + exactTokens = countTokens(content) + } + while (selected.length > 0 && exactTokens > maxTokens) { + selected.pop() + content = renderSelectedChunks(selected) + exactTokens = countTokens(content) + } + + return { + content, + estimatedTokens: exactTokens, + includedDocuments: [...new Set(selected.map((chunk) => chunk.filename))], + duplicateDocumentsSkipped, + } +} + +export function retrieveDocumentContext(text: string, query: string): string { + return retrieveDocumentContexts([{ filename: 'document', text }], query).content +} + +export function estimatePromptTokens(parts: string[]): number { + return countTokens(parts.filter(Boolean).join('\n')) +} diff --git a/apps/electron/src/main/lib/test-helpers/electron-mock.ts b/apps/electron/src/main/lib/test-helpers/electron-mock.ts new file mode 100644 index 000000000..153f83e95 --- /dev/null +++ b/apps/electron/src/main/lib/test-helpers/electron-mock.ts @@ -0,0 +1,24 @@ +export function createElectronMock(getHome: () => string) { + return { + app: { + isPackaged: true, + getPath: () => getHome(), + }, + BrowserWindow: class {}, + clipboard: {}, + dialog: {}, + nativeImage: { createFromPath: () => ({}) }, + nativeTheme: {}, + powerMonitor: {}, + powerSaveBlocker: {}, + screen: {}, + shell: { + openExternal: async () => undefined, + }, + safeStorage: { + isEncryptionAvailable: () => false, + encryptString: (value: string) => Buffer.from(value), + decryptString: (value: Buffer) => value.toString('utf-8'), + }, + } +} diff --git a/bun.lock b/bun.lock index 0b46a583e..53e86b538 100644 --- a/bun.lock +++ b/bun.lock @@ -36,10 +36,11 @@ }, "apps/electron": { "name": "@proma/electron", - "version": "0.15.6", + "version": "0.15.8", "dependencies": { "@anthropic-ai/claude-agent-sdk": "0.3.201", "@anthropic-ai/sdk": "^0.93.0", + "@anthropic-ai/tokenizer": "0.0.4", "@earendil-works/pi-agent-core": "0.80.9", "@earendil-works/pi-ai": "0.80.9", "@earendil-works/pi-coding-agent": "0.80.9", @@ -158,7 +159,7 @@ }, "packages/core": { "name": "@proma/core", - "version": "0.2.15", + "version": "0.2.16", "dependencies": { "@proma/shared": "workspace:*", "highlight.js": "^11.11.1", @@ -252,6 +253,8 @@ "@anthropic-ai/sdk": ["@anthropic-ai/sdk@0.93.0", "", { "dependencies": { "json-schema-to-ts": "^3.1.1" }, "peerDependencies": { "zod": "^3.25.0 || ^4.0.0" }, "optionalPeers": ["zod"], "bin": { "anthropic-ai-sdk": "bin/cli" } }, "sha512-q9vaSZQVFx6B/gPxetGYfLXSJD5v0sOmh0OpZDq7yCrTSA+Rscvrtyol7JJTW40wEpQB4U1B4JXzxQitbQ3CAA=="], + "@anthropic-ai/tokenizer": ["@anthropic-ai/tokenizer@0.0.4", "", { "dependencies": { "@types/node": "^18.11.18", "tiktoken": "^1.0.10" } }, "sha512-EHRKbxlxlc8W4KCBEseByJ7YwyYCmgu9OyN59H9+IYIGPoKv8tXyQXinkeGDI+cI8Tiuz9wk2jZb/kK7AyvL7g=="], + "@aws-crypto/sha256-browser": ["@aws-crypto/sha256-browser@5.2.0", "", { "dependencies": { "@aws-crypto/sha256-js": "^5.2.0", "@aws-crypto/supports-web-crypto": "^5.2.0", "@aws-crypto/util": "^5.2.0", "@aws-sdk/types": "^3.222.0", "@aws-sdk/util-locate-window": "^3.0.0", "@smithy/util-utf8": "^2.0.0", "tslib": "^2.6.2" } }, "sha512-AXfN/lGotSQwu6HNcEsIASo7kWXZ5HYWvfOmSNKDsEqC4OashTp8alTmaz+F7TC2L083SFv5RdB+qU3Vs1kZqw=="], "@aws-crypto/sha256-js": ["@aws-crypto/sha256-js@5.2.0", "", { "dependencies": { "@aws-crypto/util": "^5.2.0", "@aws-sdk/types": "^3.222.0", "tslib": "^2.6.2" } }, "sha512-FFQQyu7edu4ufvIZ+OadFpHHOt+eSTBaYaki44c+akjg7qZg9oOQeLlk77F6tSYqjDAFClrHJk9tMf0HdVyOvA=="], @@ -2372,6 +2375,8 @@ "through": ["through@2.3.8", "", {}, "sha512-w89qg7PI8wAdvX60bMDP+bFoD5Dvhm9oLheFp5O4a2QF0cSBGsBX4qZmadPMvVqlLJBBci+WqGGOAPvcDeNSVg=="], + "tiktoken": ["tiktoken@1.0.22", "", {}, "sha512-PKvy1rVF1RibfF3JlXBSP0Jrcw2uq3yXdgcEXtKTYn3QJ/cBRBHDnrJ5jHky+MENZ6DIPwNUGWpkVx+7joCpNA=="], + "tiny-typed-emitter": ["tiny-typed-emitter@2.1.0", "", {}, "sha512-qVtvMxeXbVej0cQWKqVSSAHmKZEHAvxdF8HEUBFWts8h+xEo5m/lEiPakuyZ3BnCBjOD8i24kzNOiOLLgsSxhA=="], "tinyexec": ["tinyexec@1.2.2", "", {}, "sha512-M/Q0B2cp4K7kynaT/vnED1j8TlLY+Pp7C6Wl2bl/7u/F0mUVwdyOpwomQb8JpYLitHUssAJRmLZdMCGsrx7i+g=="], @@ -2540,6 +2545,8 @@ "zwitch": ["zwitch@2.0.4", "", {}, "sha512-bXE4cR/kVZhKZX/RjPEflHaKVhUVl85noU3v6b8apfQEc1x4A+zBxjZ4lN8LqGd6WZ3dl98pY4o717VFmoPp+A=="], + "@anthropic-ai/tokenizer/@types/node": ["@types/node@18.19.130", "", { "dependencies": { "undici-types": "~5.26.4" } }, "sha512-GRaXQx6jGfL8sKfaIDD6OupbIHBr9jv7Jnaml9tB7l4v068PAOXqfcujMMo5PhbIs6ggR1XODELqahT2R8v0fg=="], + "@aws-sdk/credential-provider-http/@smithy/node-http-handler": ["@smithy/node-http-handler@4.9.4", "", { "dependencies": { "@smithy/core": "^3.29.2", "@smithy/types": "^4.16.0", "tslib": "^2.6.2" } }, "sha512-BNTop/fSOptmoVk8g+efwHCofFh37g70OWGAFES1TeAAJja1K5aAI8rTE26ETSc5k8IQuWY2kAIoPla01NgYrA=="], "@aws-sdk/credential-provider-sso/@aws-sdk/token-providers": ["@aws-sdk/token-providers@3.1083.0", "", { "dependencies": { "@aws-sdk/core": "^3.975.1", "@aws-sdk/nested-clients": "^3.997.31", "@aws-sdk/types": "^3.974.0", "@smithy/core": "^3.29.2", "@smithy/types": "^4.16.0", "tslib": "^2.6.2" } }, "sha512-s0woKnxuHrExLc5L2ArIH5BMkbonHPtt+5hSBM8oknp9M6QTuUmmAmJ2E0EdzCGONrO+8+ADPqvv6UX0nNcc7A=="], @@ -2844,6 +2851,8 @@ "zip-stream/readable-stream": ["readable-stream@3.6.2", "", { "dependencies": { "inherits": "^2.0.3", "string_decoder": "^1.1.1", "util-deprecate": "^1.0.1" } }, "sha512-9u/sniCrY3D5WdsERHzHE4G2YCXqoG5FTHUiCC4SIbr6XcLZBY05ya9EKjYek9O5xOAwjGq+1JdGBAS7Q9ScoA=="], + "@anthropic-ai/tokenizer/@types/node/undici-types": ["undici-types@5.26.5", "", {}, "sha512-JlCMO+ehdEIKqlFxk6IfVoAUVmgz7cU7zD/h9XZ0qzeosSHmUJVOzSQvvYSYWXkFXC+IfLKSIffhv0sVZup6pA=="], + "@babel/helper-compilation-targets/lru-cache/yallist": ["yallist@3.1.1", "", {}, "sha512-a4UGQaWPH59mOXUYnAG2ewncQS4i4F43Tv3JoAM+s2VDAmS9NsK8GpDMLrCHPksFT7h3K6TOoUNn2pb7RoXx4g=="], "@develar/schema-utils/ajv/json-schema-traverse": ["json-schema-traverse@0.4.1", "", {}, "sha512-xbbCH5dCYU5T8LcEhhuh7HJ88HXuW3qsI3Y0zOZFKfZEHcpWiHU/Jxzk629Brsab/mMiHQti9wMP+845RPe3Vg=="], diff --git a/packages/core/package.json b/packages/core/package.json index ba3b6dc75..3acdecad7 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -1,6 +1,6 @@ { "name": "@proma/core", - "version": "0.2.15", + "version": "0.2.16", "license": "AGPL-3.0-only", "description": "proma core types,storage and core logic with agents", "type": "module", diff --git a/packages/core/src/providers/anthropic-adapter.test.ts b/packages/core/src/providers/anthropic-adapter.test.ts index d8b8ddf6e..230c93aba 100644 --- a/packages/core/src/providers/anthropic-adapter.test.ts +++ b/packages/core/src/providers/anthropic-adapter.test.ts @@ -66,3 +66,46 @@ describe('AnthropicAdapter headers', () => { expect(request.headers['api-key']).toBeUndefined() }) }) + +describe('AnthropicAdapter document blocks', () => { + test('forwards the current PDF as a native Anthropic document block', () => { + const adapter = new AnthropicAdapter('anthropic') + const request = adapter.buildStreamRequest({ + baseUrl: 'https://newapi.example.com', + apiKey: 'test-key', + modelId: 'claude-opus-4-8', + history: [], + userMessage: '分析这份文档', + attachments: [{ + id: 'pdf-1', + filename: '研究报告.pdf', + mediaType: 'application/pdf', + localPath: 'conversation/pdf-1.pdf', + size: 1024, + }], + readImageAttachments: () => [], + readDocumentAttachments: () => [{ + filename: '研究报告.pdf', + mediaType: 'application/pdf', + data: 'JVBERi0xLjQ=', + }], + }) + + const body = JSON.parse(request.body) as { + messages: Array<{ content: Array> }> + } + expect(body.messages[0]?.content).toEqual([ + { + type: 'document', + source: { + type: 'base64', + media_type: 'application/pdf', + data: 'JVBERi0xLjQ=', + }, + name: '研究报告.pdf', + cache_control: { type: 'ephemeral' }, + }, + { type: 'text', text: '分析这份文档' }, + ]) + }) +}) diff --git a/packages/core/src/providers/anthropic-adapter.ts b/packages/core/src/providers/anthropic-adapter.ts index 5cde738b7..3d3f19414 100644 --- a/packages/core/src/providers/anthropic-adapter.ts +++ b/packages/core/src/providers/anthropic-adapter.ts @@ -31,6 +31,7 @@ import type { StreamEvent, TitleRequestInput, ImageAttachmentData, + DocumentAttachmentData, ToolDefinition, ContinuationMessage, } from './types.ts' @@ -42,7 +43,7 @@ import { getPromaUserAgent } from './user-agent.ts' /** Anthropic 内容块（扩展支持 tool_use / tool_result） */ interface AnthropicContentBlock { - type: 'text' | 'image' | 'tool_use' | 'tool_result' | 'thinking' + type: 'text' | 'image' | 'document' | 'tool_use' | 'tool_result' | 'thinking' text?: string source?: { type: 'base64' @@ -61,6 +62,7 @@ interface AnthropicContentBlock { tool_use_id?: string content?: string | AnthropicContentBlock[] is_error?: boolean + cache_control?: { type: 'ephemeral' } } /** Anthropic 消息格式 */ @@ -118,6 +120,19 @@ function buildImageBlocks(imageData: ImageAttachmentData[]): AnthropicContentBlo })) } +function buildDocumentBlocks(documentData: DocumentAttachmentData[]): AnthropicContentBlock[] { + return documentData.map((document) => ({ + type: 'document' as const, + source: { + type: 'base64' as const, + media_type: document.mediaType, + data: document.data, + }, + name: document.filename, + cache_control: { type: 'ephemeral' as const }, + })) +} + /** * 构建包含图片和文本的消息内容 * @@ -126,10 +141,14 @@ function buildImageBlocks(imageData: ImageAttachmentData[]): AnthropicContentBlo function buildMessageContent( text: string, imageData: ImageAttachmentData[], + documentData: DocumentAttachmentData[] = [], ): string | AnthropicContentBlock[] { - if (imageData.length === 0) return text + if (imageData.length === 0 && documentData.length === 0) return text - const content: AnthropicContentBlock[] = buildImageBlocks(imageData) + const content: AnthropicContentBlock[] = [ + ...buildImageBlocks(imageData), + ...buildDocumentBlocks(documentData), + ] if (text) { content.push({ type: 'text', text }) } @@ -152,7 +171,7 @@ function buildMessageContent( function toAnthropicMessages( input: StreamRequestInput, ): AnthropicMessage[] { - const { history, userMessage, attachments, readImageAttachments } = input + const { history, userMessage, attachments, readImageAttachments, readDocumentAttachments } = input // 历史消息转换 const messages: AnthropicMessage[] = history @@ -171,9 +190,10 @@ function toAnthropicMessages( // 当前用户消息 const currentImages = readImageAttachments(attachments) + const currentDocuments = readDocumentAttachments?.(attachments) ?? [] messages.push({ role: 'user', - content: buildMessageContent(userMessage, currentImages), + content: buildMessageContent(userMessage, currentImages, currentDocuments), }) return messages diff --git a/packages/core/src/providers/types.ts b/packages/core/src/providers/types.ts index a49498cfc..493c0f7a4 100644 --- a/packages/core/src/providers/types.ts +++ b/packages/core/src/providers/types.ts @@ -18,6 +18,13 @@ export interface ImageAttachmentData { data: string } +/** 文档附件的 base64 数据（已从磁盘读取） */ +export interface DocumentAttachmentData { + filename: string + mediaType: string + data: string +} + /** * 图片附件读取器 * @@ -26,6 +33,9 @@ export interface ImageAttachmentData { */ export type ImageAttachmentReader = (attachments?: FileAttachment[]) => ImageAttachmentData[] +/** 文档附件读取器；未启用原生文档时返回空数组 */ +export type DocumentAttachmentReader = (attachments?: FileAttachment[]) => DocumentAttachmentData[] + // ===== Tool Use（Function Calling）===== /** 工具参数属性定义 */ @@ -218,6 +228,8 @@ export interface StreamRequestInput { attachments?: FileAttachment[] /** 图片附件读取器（由 Electron 层注入） */ readImageAttachments: ImageAttachmentReader + /** 原生文档读取器（由 Electron 层注入，可选） */ + readDocumentAttachments?: DocumentAttachmentReader /** 是否启用思考模式（各适配器根据供应商 API 自行转换） */ thinkingEnabled?: boolean /** 工具定义列表（可选，启用 function calling） */