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>
This commit is contained in:
akumatus
2025-01-24 04:04:00 +00:00
parent 48c26017ae
commit 95cf2e047f
13 changed files with 131 additions and 37 deletions

View File

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

View File

@@ -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].`,
},
],
},

View File

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

View File

@@ -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 {

View File

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

View File

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