feat(core): cite source documents in the AI answer (#9863)

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<string, any>`
- Add unit test.

<div class='graphite__hidden'>
          <div>🎥 Video uploaded on Graphite:</div>
            <a href="https://app.graphite.dev/media/video/sJGviKxfE3Ap685cl5bj/ec24e664-039e-4fab-bd26-b3312f011daf.mov">
              <img src="https://app.graphite.dev/api/v1/graphite/video/thumbnail/sJGviKxfE3Ap685cl5bj/ec24e664-039e-4fab-bd26-b3312f011daf.mov">
            </a>
          </div>
<video src="https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/sJGviKxfE3Ap685cl5bj/ec24e664-039e-4fab-bd26-b3312f011daf.mov">录屏2025-01-23 10.40.38.mov</video>

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);
 });

packages/backend/server/src/plugins/copilot/prompt/prompts.ts

@@ -947,8 +947,33 @@ const chat: Prompt[] = [
     messages: [
       {
         role: 'system',
-        content:
+        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.
-          "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].`,
       },
     ],
   },

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() {

packages/backend/server/src/plugins/copilot/resolver.ts

@@ -119,7 +119,7 @@ class CreateChatMessageInput implements Omit<SubmittedMessage, 'content'> {
   blobs!: Promise<FileUpload>[] | undefined;
 
   @Field(() => GraphQLJSON, { nullable: true })
-  params!: Record<string, string> | undefined;
+  params!: Record<string, any> | undefined;
 }
 
 enum ChatHistoryOrder {

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),

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
+  params: z.record(z.any()).optional().nullable(),
-    .record(z.union([z.string(), z.array(z.string()), z.record(z.any())]))
-    .optional()
-    .nullable(),
 });
 
 export const PromptMessageSchema = PureMessageSchema.extend({

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<string>;
 
     interface ChatOptions extends AITextActionOptions {
+      // related documents
+      docs?: DocContext[];
       sessionId?: string;
       isRootSession?: boolean;
     }

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<string>;
+  markdown?: Signal<string>;
 }
 
 export interface FileChip extends BaseChip {

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 =
+    const userInput = (markdown ? `${markdown}\n` : '') + text;
-      (markdown ? `${markdown}\n` : '') + `${docsContent}\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,

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) {
+      if (this.chip.markdown) {
-        this.chip.content.value = result.markdown;
+        this.chip.markdown.value = result.markdown;
       } else {
-        this.chip.content = new Signal<string>(result.markdown);
+        this.chip.markdown = new Signal<string>(result.markdown);
       }
       this.updateChip(this.chip, {
         state: 'success',

packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts

@@ -17,7 +17,7 @@ export type TextToTextOptions = {
   sessionId?: string | Promise<string>;
   content?: string;
   attachments?: (string | Blob | File)[];
-  params?: Record<string, string>;
+  params?: Record<string, any>;
   timeout?: number;
   stream?: boolean;
   signal?: AbortSignal;

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,
     });
   });
 

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:
+      content: 'What is AFFiNE AI?',
-        'AFFiNE AI is an assistant with the ability to create well-structured outlines for any given content.\nsummarize',
     });
     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);
   });