From 022a59b1bd84e016d51d68bf7fe5efb841f5d65c Mon Sep 17 00:00:00 2001 From: SuYao Date: Sun, 25 May 2025 21:08:06 +0800 Subject: [PATCH] feat: enhance citation handling in message export functionality (#6422) * feat: enhance citation handling in message export functionality - Refactored message export functions to include citation content in markdown output. - Introduced a new utility function to extract and format citations from messages. - Updated related imports and adjusted existing markdown generation logic for improved clarity and maintainability. * feat: enhance message export tests to include citation and reasoning content - Added tests to verify inclusion of citation content in markdown output when citation blocks exist. - Ensured proper formatting with double newlines between sections in exported markdown. - Updated existing tests to handle cases with reasoning content and no main text block gracefully. * fix: update citation mapping in export tests for consistency - Modified the citation mapping in export tests to use the index parameter directly, improving clarity and consistency in the generated markdown output. --- src/renderer/src/store/messageBlock.ts | 2 +- .../src/utils/__tests__/export.test.ts | 50 +++++++++++++- src/renderer/src/utils/export.ts | 66 +++++++++---------- src/renderer/src/utils/messageUtils/find.ts | 11 +++- 4 files changed, 92 insertions(+), 37 deletions(-) diff --git a/src/renderer/src/store/messageBlock.ts b/src/renderer/src/store/messageBlock.ts index cf9515fe00..f9b8c34cd5 100644 --- a/src/renderer/src/store/messageBlock.ts +++ b/src/renderer/src/store/messageBlock.ts @@ -81,7 +81,7 @@ const selectBlockEntityById = (state: RootState, blockId: string | undefined) => blockId ? messageBlocksSelectors.selectById(state, blockId) : undefined // Use adapter selector // --- Centralized Citation Formatting Logic --- -const formatCitationsFromBlock = (block: CitationMessageBlock | undefined): Citation[] => { +export const formatCitationsFromBlock = (block: CitationMessageBlock | undefined): Citation[] => { if (!block) return [] let formattedCitations: Citation[] = [] diff --git a/src/renderer/src/utils/__tests__/export.test.ts b/src/renderer/src/utils/__tests__/export.test.ts index 5d308f5a52..e3469c1d8b 100644 --- a/src/renderer/src/utils/__tests__/export.test.ts +++ b/src/renderer/src/utils/__tests__/export.test.ts @@ -24,6 +24,15 @@ vi.mock('@renderer/utils/messageUtils/find', () => ({ // Assuming content exists on ThinkingBlock // Need to cast block to access content if not on base type return (thinkingBlock as any)?.content || '' + }), + getCitationContent: vi.fn((message: Message & { _fullBlocks?: MessageBlock[] }) => { + const citationBlocks = message._fullBlocks?.filter((b) => b.type === MessageBlockType.CITATION) || [] + // Return empty string if no citation blocks, otherwise mock citation content + if (citationBlocks.length === 0) return '' + // Mock citation format: [number] [url](title) + return citationBlocks + .map((_, index) => `[${index + 1}] [https://example${index + 1}.com](Example Citation ${index + 1})`) + .join('\n\n') }) })) @@ -198,6 +207,9 @@ describe('export', () => { const markdown = messageToMarkdown(msg!) expect(markdown).toContain('### 🧑‍💻 User') expect(markdown).toContain('hello user') + // Should have double newlines between sections + const sections = markdown.split('\n\n') + expect(sections.length).toBeGreaterThanOrEqual(3) // title, content, citation (empty) }) it('should format assistant message using main text block', () => { @@ -206,6 +218,9 @@ describe('export', () => { const markdown = messageToMarkdown(msg!) expect(markdown).toContain('### 🤖 Assistant') expect(markdown).toContain('hi assistant') + // Should have double newlines between sections + const sections = markdown.split('\n\n') + expect(sections.length).toBeGreaterThanOrEqual(3) // title, content, citation (empty) }) it('should handle message with no main text block gracefully', () => { @@ -213,7 +228,19 @@ describe('export', () => { mockedMessages.push(msg) const markdown = messageToMarkdown(msg) expect(markdown).toContain('### 🧑‍💻 User') - expect(markdown.trim().endsWith('User')).toBe(true) + // Check that it doesn't fail when no content exists + expect(markdown).toBeDefined() + }) + + it('should include citation content when citation blocks exist', () => { + const msgWithCitation = createMessage({ role: 'assistant', id: 'a_cite' }, [ + { type: MessageBlockType.MAIN_TEXT, content: 'Main content' }, + { type: MessageBlockType.CITATION } + ]) + const markdown = messageToMarkdown(msgWithCitation) + expect(markdown).toContain('### 🤖 Assistant') + expect(markdown).toContain('Main content') + expect(markdown).toContain('[1] [https://example1.com](Example Citation 1)') }) }) @@ -231,7 +258,12 @@ describe('export', () => { const msgWithoutReasoning = createMessage({ role: 'assistant', id: 'a4' }, [ { type: MessageBlockType.MAIN_TEXT, content: 'Simple Answer' } ]) - mockedMessages = [msgWithReasoning, msgWithThinkTag, msgWithoutReasoning] + const msgWithReasoningAndCitation = createMessage({ role: 'assistant', id: 'a5' }, [ + { type: MessageBlockType.MAIN_TEXT, content: 'Answer with citation' }, + { type: MessageBlockType.THINKING, content: 'Some thinking' }, + { type: MessageBlockType.CITATION } + ]) + mockedMessages = [msgWithReasoning, msgWithThinkTag, msgWithoutReasoning, msgWithReasoningAndCitation] }) it('should include reasoning content from thinking block in details section', () => { @@ -243,6 +275,9 @@ describe('export', () => { expect(markdown).toContain('common.reasoning_content') expect(markdown).toContain('Detailed thought process') + // Should have double newlines between sections + const sections = markdown.split('\n\n') + expect(sections.length).toBeGreaterThanOrEqual(3) }) it('should handle tag and replace newlines with
in reasoning', () => { @@ -263,6 +298,17 @@ describe('export', () => { expect(markdown).toContain('Simple Answer') expect(markdown).not.toContain(' { + const msg = mockedMessages.find((m) => m.id === 'a5') + expect(msg).toBeDefined() + const markdown = messageToMarkdownWithReasoning(msg!) + expect(markdown).toContain('### 🤖 Assistant') + expect(markdown).toContain('Answer with citation') + expect(markdown).toContain(' { diff --git a/src/renderer/src/utils/export.ts b/src/renderer/src/utils/export.ts index 219b1ce6e6..44c444bd6c 100644 --- a/src/renderer/src/utils/export.ts +++ b/src/renderer/src/utils/export.ts @@ -8,7 +8,7 @@ import type { Topic } from '@renderer/types' import type { Message } from '@renderer/types/newMessage' import { removeSpecialCharactersForFileName } from '@renderer/utils/file' import { convertMathFormula } from '@renderer/utils/markdown' -import { getMainTextContent, getThinkingContent } from '@renderer/utils/messageUtils/find' +import { getCitationContent, getMainTextContent, getThinkingContent } from '@renderer/utils/messageUtils/find' import { markdownToBlocks } from '@tryfabric/martian' import dayjs from 'dayjs' //TODO: 添加对思考内容的支持 @@ -43,48 +43,48 @@ export function getTitleFromString(str: string, length: number = 80) { return title } -export const messageToMarkdown = (message: Message) => { +const createBaseMarkdown = (message: Message, includeReasoning: boolean = false) => { const { forceDollarMathInMarkdown } = store.getState().settings const roleText = message.role === 'user' ? '🧑‍💻 User' : '🤖 Assistant' const titleSection = `### ${roleText}` + let reasoningSection = '' + + if (includeReasoning) { + let reasoningContent = getThinkingContent(message) + if (reasoningContent) { + if (reasoningContent.startsWith('\n')) { + reasoningContent = reasoningContent.substring(8) + } else if (reasoningContent.startsWith('')) { + reasoningContent = reasoningContent.substring(7) + } + reasoningContent = reasoningContent.replace(/\n/g, '
') + + if (forceDollarMathInMarkdown) { + reasoningContent = convertMathFormula(reasoningContent) + } + reasoningSection = `
+ ${i18n.t('common.reasoning_content')}
+ ${reasoningContent} +
` + } + } + const content = getMainTextContent(message) + const citation = getCitationContent(message) const contentSection = forceDollarMathInMarkdown ? convertMathFormula(content) : content - return [titleSection, '', contentSection].join('\n') + return { titleSection, reasoningSection, contentSection, citation } +} + +export const messageToMarkdown = (message: Message) => { + const { titleSection, contentSection, citation } = createBaseMarkdown(message) + return [titleSection, '', contentSection, citation].join('\n\n') } // 保留接口用于其它导出方法使用 export const messageToMarkdownWithReasoning = (message: Message) => { - const { forceDollarMathInMarkdown } = store.getState().settings - const roleText = message.role === 'user' ? '🧑‍💻 User' : '🤖 Assistant' - const titleSection = `### ${roleText}` - let reasoningContent = getThinkingContent(message) - // 处理思考内容 - let reasoningSection = '' - if (reasoningContent) { - // 移除开头的标记和换行符，并将所有换行符替换为
- if (reasoningContent.startsWith('\n')) { - reasoningContent = reasoningContent.substring(8) - } else if (reasoningContent.startsWith('')) { - reasoningContent = reasoningContent.substring(7) - } - reasoningContent = reasoningContent.replace(/\n/g, '
') - - // 应用数学公式转换（如果启用） - if (forceDollarMathInMarkdown) { - reasoningContent = convertMathFormula(reasoningContent) - } - // 添加思考内容的Markdown格式 - reasoningSection = `
- ${i18n.t('common.reasoning_content')}
- ${reasoningContent} -
` - } - const content = getMainTextContent(message) - - const contentSection = forceDollarMathInMarkdown ? convertMathFormula(content) : content - - return [titleSection, '', reasoningSection + contentSection].join('\n') + const { titleSection, reasoningSection, contentSection, citation } = createBaseMarkdown(message, true) + return [titleSection, '', reasoningSection + contentSection, citation].join('\n\n') } export const messagesToMarkdown = (messages: Message[], exportReasoning?: boolean) => { diff --git a/src/renderer/src/utils/messageUtils/find.ts b/src/renderer/src/utils/messageUtils/find.ts index dd3ae7f92b..a2d804a10a 100644 --- a/src/renderer/src/utils/messageUtils/find.ts +++ b/src/renderer/src/utils/messageUtils/find.ts @@ -1,5 +1,5 @@ import store from '@renderer/store' -import { messageBlocksSelectors } from '@renderer/store/messageBlock' +import { formatCitationsFromBlock, messageBlocksSelectors } from '@renderer/store/messageBlock' import { FileType } from '@renderer/types' import type { CitationMessageBlock, @@ -128,6 +128,15 @@ export const getThinkingContent = (message: Message): string => { return thinkingBlocks.map((block) => block.content).join('\n\n') } +export const getCitationContent = (message: Message): string => { + const citationBlocks = findCitationBlocks(message) + return citationBlocks + .map((block) => formatCitationsFromBlock(block)) + .flat() + .map((citation) => `[${citation.number}] [${citation.url}](${citation.title || citation.url})`) + .join('\n\n') +} + /** * Gets the knowledgeBaseIds array from the *first* MainTextMessageBlock of a message. * Note: Assumes knowledgeBaseIds are only relevant on the first text block, adjust if needed.