From 95cf2e047f2cf3563a5a3a947d7f0004af261e7a Mon Sep 17 00:00:00 2001 From: akumatus Date: Fri, 24 Jan 2025 04:04:00 +0000 Subject: [PATCH] feat(core): cite source documents in the AI answer (#9863) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Support issue [BS-2424](https://linear.app/affine-design/issue/BS-2424). ### What changed? - Add relevant document prompt templates. - Add citation rules in system prompts. - Change message `params` type to `Record` - Add unit test.
🎥 Video uploaded on Graphite:
--- .../server/src/__tests__/copilot.spec.ts | 59 +++++++++++++++++-- .../src/plugins/copilot/prompt/prompts.ts | 29 ++++++++- .../plugins/copilot/providers/perplexity.ts | 4 +- .../server/src/plugins/copilot/resolver.ts | 2 +- .../server/src/plugins/copilot/session.ts | 3 +- .../server/src/plugins/copilot/types.ts | 5 +- .../blocksuite/presets/ai/actions/types.ts | 7 ++- .../presets/ai/chat-panel/chat-context.ts | 2 +- .../presets/ai/chat-panel/chat-panel-input.ts | 25 ++++---- .../ai/chat-panel/components/doc-chip.ts | 6 +- .../block-suite-editor/ai/request.ts | 2 +- .../block-suite-editor/ai/setup-provider.tsx | 15 ++++- .../affine-cloud-copilot/e2e/copilot.spec.ts | 9 ++- 13 files changed, 131 insertions(+), 37 deletions(-) diff --git a/packages/backend/server/src/__tests__/copilot.spec.ts b/packages/backend/server/src/__tests__/copilot.spec.ts index d3ea94c6e3..a0920a1b9b 100644 --- a/packages/backend/server/src/__tests__/copilot.spec.ts +++ b/packages/backend/server/src/__tests__/copilot.spec.ts @@ -629,6 +629,57 @@ test('should revert message correctly', async t => { } }); +test('should handle params correctly in chat session', async t => { + const { prompt, session } = t.context; + + await prompt.set('prompt', 'model', [ + { role: 'system', content: 'hello {{word}}' }, + ]); + + const sessionId = await session.create({ + docId: 'test', + workspaceId: 'test', + userId, + promptName: 'prompt', + }); + + const s = (await session.get(sessionId))!; + + // Case 1: When params is provided directly + { + const directParams = { word: 'direct' }; + const messages = s.finish(directParams); + t.is(messages[0].content, 'hello direct', 'should use provided params'); + } + + // Case 2: When no params provided but last message has params + { + s.push({ + role: 'user', + content: 'test message', + params: { word: 'fromMessage' }, + createdAt: new Date(), + }); + const messages = s.finish({}); + t.is( + messages[0].content, + 'hello fromMessage', + 'should use params from last message' + ); + } + + // Case 3: When neither params provided nor last message has params + { + s.push({ + role: 'user', + content: 'test message without params', + createdAt: new Date(), + }); + const messages = s.finish({}); + t.is(messages[0].content, 'hello ', 'should use empty params'); + } +}); + // ==================== provider ==================== test('should be able to get provider', async t => { @@ -1102,7 +1153,7 @@ test('CitationParser should replace citation placeholders with URLs', t => { 'This is [a] test sentence with [citations [^1]] and [^2] and [3].', `[^1]: {"type":"url","url":"${encodeURIComponent(citations[0])}"}`, `[^2]: {"type":"url","url":"${encodeURIComponent(citations[1])}"}`, - ].join('\n\n'); + ].join('\n'); t.is(result, expected); }); @@ -1145,7 +1196,7 @@ test('CitationParser should replace chunks of citation placeholders with URLs', `[^5]: {"type":"url","url":"${encodeURIComponent(citations[4])}"}`, `[^6]: {"type":"url","url":"${encodeURIComponent(citations[5])}"}`, `[^7]: {"type":"url","url":"${encodeURIComponent(citations[6])}"}`, - ].join('\n\n'); + ].join('\n'); t.is(result, expected); }); @@ -1166,7 +1217,7 @@ test('CitationParser should not replace citation already with URLs', t => { `[^1]: {"type":"url","url":"${encodeURIComponent(citations[0])}"}`, `[^2]: {"type":"url","url":"${encodeURIComponent(citations[1])}"}`, `[^3]: {"type":"url","url":"${encodeURIComponent(citations[2])}"}`, - ].join('\n\n'); + ].join('\n'); t.is(result, expected); }); @@ -1193,6 +1244,6 @@ test('CitationParser should not replace chunks of citation already with URLs', t `[^1]: {"type":"url","url":"${encodeURIComponent(citations[0])}"}`, `[^2]: {"type":"url","url":"${encodeURIComponent(citations[1])}"}`, `[^3]: {"type":"url","url":"${encodeURIComponent(citations[2])}"}`, - ].join('\n\n'); + ].join('\n'); t.is(result, expected); }); diff --git a/packages/backend/server/src/plugins/copilot/prompt/prompts.ts b/packages/backend/server/src/plugins/copilot/prompt/prompts.ts index ebc1effbe0..82886bbc1d 100644 --- a/packages/backend/server/src/plugins/copilot/prompt/prompts.ts +++ b/packages/backend/server/src/plugins/copilot/prompt/prompts.ts @@ -947,8 +947,33 @@ const chat: Prompt[] = [ messages: [ { role: 'system', - content: - "You are AFFiNE AI, a professional and humorous copilot within AFFiNE. You are powered by latest GPT model from OpenAI and AFFiNE. AFFiNE is an open source general purposed productivity tool that contains unified building blocks that users can use on any interfaces, including block-based docs editor, infinite canvas based edgeless graphic mode, or multi-dimensional table with multiple transformable views. Your mission is always to try your very best to assist users to use AFFiNE to write docs, draw diagrams or plan things with these abilities. You always think step-by-step and describe your plan for what to build, using well-structured and clear markdown, written out in great detail. Unless otherwise specified, where list, JSON, or code blocks are required for giving the output. Minimize any other prose so that your responses can be directly used and inserted into the docs. You are able to access to API of AFFiNE to finish your job. You always respect the users' privacy and would not leak their info to anyone else. AFFiNE is made by Toeverything .Pte .Ltd, a company registered in Singapore with a diverse and international team. The company also open sourced blocksuite and octobase for building tools similar to Affine. The name AFFiNE comes from the idea of AFFiNE transform, as blocks in affine can all transform in page, edgeless or database mode. AFFiNE team is now having 25 members, an open source company driven by engineers.", + content: `You are AFFiNE AI, a professional and humorous copilot within AFFiNE. You are powered by latest GPT model from OpenAI and AFFiNE. AFFiNE is an open source general purposed productivity tool that contains unified building blocks that users can use on any interfaces, including block-based docs editor, infinite canvas based edgeless graphic mode, or multi-dimensional table with multiple transformable views. Your mission is always to try your very best to assist users to use AFFiNE to write docs, draw diagrams or plan things with these abilities. You always think step-by-step and describe your plan for what to build, using well-structured and clear markdown, written out in great detail. Unless otherwise specified, where list, JSON, or code blocks are required for giving the output. Minimize any other prose so that your responses can be directly used and inserted into the docs. You are able to access to API of AFFiNE to finish your job. You always respect the users' privacy and would not leak their info to anyone else. AFFiNE is made by Toeverything .Pte .Ltd, a company registered in Singapore with a diverse and international team. The company also open sourced blocksuite and octobase for building tools similar to Affine. The name AFFiNE comes from the idea of AFFiNE transform, as blocks in affine can all transform in page, edgeless or database mode. AFFiNE team is now having 25 members, an open source company driven by engineers. +# Context Documents +The following documents provide relevant context and background information for your reference. +If the provided documents are relevant to the user's query: +- Use them to enrich and support your response +- Cite sources using the citation rules below + +If the documents are not relevant: +- Answer the question directly based on your knowledge +- Do not reference or mention the provided documents + +{{#docs}} +## Document {{index}} +- document_index: {{index}} +- document_id: {{docId}} +- document_content: +{{markdown}} +{{/docs}} + +# Citations Rules: +When referencing information from the provided documents in your response: +1. Use markdown footnote format for citations +2. Add citations immediately after the relevant sentence or paragraph +3. Required format: [^document_index] where document_index is the numerical index of the source document +4. At the end of your response, include the full citation in the format: + [^document_index]:{"type":"doc","docId":"document_id"} +5. Ensure citations adhere strictly to the required format to avoid response errors. Do not add extra spaces in citations like [^ document_index] or [ ^document_index].`, }, ], }, diff --git a/packages/backend/server/src/plugins/copilot/providers/perplexity.ts b/packages/backend/server/src/plugins/copilot/providers/perplexity.ts index 2cda3ed20d..953c68729c 100644 --- a/packages/backend/server/src/plugins/copilot/providers/perplexity.ts +++ b/packages/backend/server/src/plugins/copilot/providers/perplexity.ts @@ -120,7 +120,7 @@ export class CitationParser { } public end() { - return this.flush() + this.getFootnotes(); + return this.flush() + '\n' + this.getFootnotes(); } private flush() { @@ -135,7 +135,7 @@ export class CitationParser { citation )}"}`; }); - return '\n\n' + footnotes.join('\n\n'); + return footnotes.join('\n'); } private getTokenContent() { diff --git a/packages/backend/server/src/plugins/copilot/resolver.ts b/packages/backend/server/src/plugins/copilot/resolver.ts index 51df230923..2b88a721c0 100644 --- a/packages/backend/server/src/plugins/copilot/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/resolver.ts @@ -119,7 +119,7 @@ class CreateChatMessageInput implements Omit { blobs!: Promise[] | undefined; @Field(() => GraphQLJSON, { nullable: true }) - params!: Record | undefined; + params!: Record | undefined; } enum ChatHistoryOrder { diff --git a/packages/backend/server/src/plugins/copilot/session.ts b/packages/backend/server/src/plugins/copilot/session.ts index 16b53f6262..bb126bc747 100644 --- a/packages/backend/server/src/plugins/copilot/session.ts +++ b/packages/backend/server/src/plugins/copilot/session.ts @@ -165,9 +165,10 @@ export class ChatSession implements AsyncDisposable { return finished; } + const lastMessage = messages.at(-1); return [ ...this.state.prompt.finish( - Object.keys(params).length ? params : firstMessage?.params || {}, + Object.keys(params).length ? params : lastMessage?.params || {}, this.config.sessionId ), ...messages.filter(m => m.content?.trim() || m.attachments?.length), diff --git a/packages/backend/server/src/plugins/copilot/types.ts b/packages/backend/server/src/plugins/copilot/types.ts index 3202ea20d9..ec6e2b0711 100644 --- a/packages/backend/server/src/plugins/copilot/types.ts +++ b/packages/backend/server/src/plugins/copilot/types.ts @@ -50,10 +50,7 @@ export const ChatMessageRole = Object.values(AiPromptRole) as [ const PureMessageSchema = z.object({ content: z.string(), attachments: z.array(z.string()).optional().nullable(), - params: z - .record(z.union([z.string(), z.array(z.string()), z.record(z.any())])) - .optional() - .nullable(), + params: z.record(z.any()).optional().nullable(), }); export const PromptMessageSchema = PureMessageSchema.extend({ diff --git a/packages/frontend/core/src/blocksuite/presets/ai/actions/types.ts b/packages/frontend/core/src/blocksuite/presets/ai/actions/types.ts index 5faa4d5fad..2530a05119 100644 --- a/packages/frontend/core/src/blocksuite/presets/ai/actions/types.ts +++ b/packages/frontend/core/src/blocksuite/presets/ai/actions/types.ts @@ -2,6 +2,8 @@ import type { getCopilotHistoriesQuery, RequestOptions } from '@affine/graphql'; import type { EditorHost } from '@blocksuite/affine/block-std'; import type { BlockModel } from '@blocksuite/affine/store'; +import type { DocContext } from '../chat-panel/chat-context'; + export const translateLangs = [ 'English', 'Spanish', @@ -37,7 +39,7 @@ export const imageProcessingTypes = [ ] as const; declare global { - // eslint-disable-next-line @typescript-eslint/no-namespace + // oxlint-disable-next-line @typescript-eslint/no-namespace namespace BlockSuitePresets { type TrackerControl = | 'format-bar' @@ -57,6 +59,7 @@ declare global { } interface AITextActionOptions { + // user input text input?: string; stream?: boolean; attachments?: (string | File | Blob)[]; // blob could only be strings for the moments (url or data urls) @@ -101,6 +104,8 @@ declare global { T['stream'] extends true ? TextStream : Promise; interface ChatOptions extends AITextActionOptions { + // related documents + docs?: DocContext[]; sessionId?: string; isRootSession?: boolean; } diff --git a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-context.ts b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-context.ts index f4f2bdddb2..11a4550d99 100644 --- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-context.ts +++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-context.ts @@ -76,7 +76,7 @@ export interface BaseChip { export interface DocChip extends BaseChip { docId: string; - content?: Signal; + markdown?: Signal; } export interface FileChip extends BaseChip { diff --git a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-input.ts b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-input.ts index 2b47b69049..d468c36b44 100644 --- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-input.ts +++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-input.ts @@ -25,7 +25,7 @@ import { AIProvider } from '../provider'; import { reportResponse } from '../utils/action-reporter'; import { readBlobAsURL } from '../utils/image'; import type { AINetworkSearchConfig } from './chat-config'; -import type { ChatContextValue, ChatMessage } from './chat-context'; +import type { ChatContextValue, ChatMessage, DocContext } from './chat-context'; import { isDocChip } from './components/utils'; const MaximumImageCount = 32; @@ -512,16 +512,11 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) { if (status === 'loading' || status === 'transmitting') return; const { images } = this.chatContextValue; - if (!text && images.length === 0) { + if (!text) { return; } const { doc } = this.host; - const docsContent = chips - .filter(isDocChip) - .map(chip => chip.content?.value || '') - .join('\n'); - this.updateContext({ images: [], status: 'loading', @@ -534,16 +529,14 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) { images?.map(image => readBlobAsURL(image)) ); - const content = - (markdown ? `${markdown}\n` : '') + `${docsContent}\n` + text; - + const userInput = (markdown ? `${markdown}\n` : '') + text; this.updateContext({ items: [ ...this.chatContextValue.items, { id: '', role: 'user', - content: content, + content: userInput, createdAt: new Date().toISOString(), attachments, }, @@ -558,8 +551,16 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) { try { const abortController = new AbortController(); + const docs: DocContext[] = chips + .filter(isDocChip) + .filter(chip => !!chip.markdown?.value && chip.state === 'success') + .map(chip => ({ + docId: chip.docId, + markdown: chip.markdown?.value || '', + })); const stream = AIProvider.actions.chat?.({ - input: content, + input: userInput, + docs: docs, docId: doc.id, attachments: images, workspaceId: doc.workspace.id, diff --git a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/doc-chip.ts b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/doc-chip.ts index 1be6a0b44b..769606cc82 100644 --- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/doc-chip.ts +++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/doc-chip.ts @@ -100,10 +100,10 @@ export class ChatPanelDocChip extends SignalWatcher( doc.load(); } const result = await extractMarkdownFromDoc(doc, this.host.std.provider); - if (this.chip.content) { - this.chip.content.value = result.markdown; + if (this.chip.markdown) { + this.chip.markdown.value = result.markdown; } else { - this.chip.content = new Signal(result.markdown); + this.chip.markdown = new Signal(result.markdown); } this.updateChip(this.chip, { state: 'success', diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts index 4f67651b6a..839fc23a88 100644 --- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts +++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts @@ -17,7 +17,7 @@ export type TextToTextOptions = { sessionId?: string | Promise; content?: string; attachments?: (string | Blob | File)[]; - params?: Record; + params?: Record; timeout?: number; stream?: boolean; signal?: AbortSignal; diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx index 123c82553d..d965ebe791 100644 --- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx +++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx @@ -117,11 +117,22 @@ export function setupAIProvider( const sessionId = options.sessionId ?? getChatSessionId(options.workspaceId, options.docId, options.attachments); + const { input, docs, ...rest } = options; + const params = docs?.length + ? { + docs: docs.map((doc, i) => ({ + docId: doc.docId, + markdown: doc.markdown, + index: i + 1, + })), + } + : undefined; return textToText({ - ...options, + ...rest, client, - content: options.input, + content: input, sessionId, + params, }); }); diff --git a/tests/affine-cloud-copilot/e2e/copilot.spec.ts b/tests/affine-cloud-copilot/e2e/copilot.spec.ts index 1c6a024498..6f2e8a9b80 100644 --- a/tests/affine-cloud-copilot/e2e/copilot.spec.ts +++ b/tests/affine-cloud-copilot/e2e/copilot.spec.ts @@ -773,14 +773,17 @@ test.describe('chat with doc', () => { // oxlint-disable-next-line unicorn/prefer-dom-node-dataset expect(await chip.getAttribute('data-state')).toBe('success'); - await makeChat(page, 'summarize'); + await typeChatSequentially(page, 'What is AFFiNE AI?'); + await page.keyboard.press('Enter'); const history = await collectChat(page); expect(history[0]).toEqual({ name: 'You', - content: - 'AFFiNE AI is an assistant with the ability to create well-structured outlines for any given content.\nsummarize', + content: 'What is AFFiNE AI?', }); expect(history[1].name).toBe('AFFiNE AI'); + expect( + await page.locator('chat-panel affine-footnote-node').count() + ).toBeGreaterThan(0); await clearChat(page); expect((await collectChat(page)).length).toBe(0); });