From e6f91cced6d5f1bdb9ac63254feaab44aef50c4a Mon Sep 17 00:00:00 2001 From: DarkSky <25152247+darkskygit@users.noreply.github.com> Date: Fri, 27 Jun 2025 23:45:49 +0800 Subject: [PATCH] feat(server): remove context prefetch & integrate context search (#12956) fix AI-173 --- .../src/plugins/copilot/context/service.ts | 54 ++++++++++++------- .../server/src/plugins/copilot/controller.ts | 5 ++ .../src/plugins/copilot/prompt/prompts.ts | 5 ++ .../src/plugins/copilot/providers/provider.ts | 6 ++- .../src/plugins/copilot/providers/types.ts | 1 + .../copilot/tools/doc-semantic-search.ts | 34 +++++++++--- .../components/ai-chat-input/ai-chat-input.ts | 44 ++------------- .../blocksuite/ai/provider/setup-provider.tsx | 1 + .../e2e/utils/chat-panel-utils.ts | 8 +-- 9 files changed, 88 insertions(+), 70 deletions(-) diff --git a/packages/backend/server/src/plugins/copilot/context/service.ts b/packages/backend/server/src/plugins/copilot/context/service.ts index d705cda8ca..e20e619c94 100644 --- a/packages/backend/server/src/plugins/copilot/context/service.ts +++ b/packages/backend/server/src/plugins/copilot/context/service.ts @@ -197,34 +197,52 @@ export class CopilotContextService implements OnApplicationBootstrap { async matchWorkspaceAll( workspaceId: string, content: string, - topK: number = 5, + topK: number, signal?: AbortSignal, - threshold: number = 0.5 + threshold: number = 0.8, + docIds?: string[], + scopedThreshold: number = 0.85 ) { if (!this.embeddingClient) return []; const embedding = await this.embeddingClient.getEmbedding(content, signal); if (!embedding) return []; - const [fileChunks, workspaceChunks] = await Promise.all([ - this.models.copilotWorkspace.matchFileEmbedding( - workspaceId, - embedding, - topK * 2, - threshold - ), - this.models.copilotContext.matchWorkspaceEmbedding( - embedding, - workspaceId, - topK * 2, - threshold - ), - ]); + const [fileChunks, workspaceChunks, scopedWorkspaceChunks] = + await Promise.all([ + this.models.copilotWorkspace.matchFileEmbedding( + workspaceId, + embedding, + topK * 2, + threshold + ), - if (!fileChunks.length && !workspaceChunks.length) return []; + this.models.copilotContext.matchWorkspaceEmbedding( + embedding, + workspaceId, + topK * 2, + threshold + ), + docIds + ? this.models.copilotContext.matchWorkspaceEmbedding( + embedding, + workspaceId, + topK * 2, + scopedThreshold, + docIds + ) + : null, + ]); + + if ( + !fileChunks.length && + !workspaceChunks.length && + !scopedWorkspaceChunks?.length + ) + return []; return await this.embeddingClient.reRank( content, - [...fileChunks, ...workspaceChunks], + [...fileChunks, ...workspaceChunks, ...(scopedWorkspaceChunks || [])], topK, signal ); diff --git a/packages/backend/server/src/plugins/copilot/controller.ts b/packages/backend/server/src/plugins/copilot/controller.ts index 1be3c68760..65797a5bc2 100644 --- a/packages/backend/server/src/plugins/copilot/controller.ts +++ b/packages/backend/server/src/plugins/copilot/controller.ts @@ -257,6 +257,7 @@ export class CopilotController implements BeforeApplicationShutdown { ...session.config.promptConfig, signal: this.getSignal(req), user: user.id, + session: session.config.sessionId, workspace: session.config.workspaceId, reasoning, webSearch, @@ -311,6 +312,7 @@ export class CopilotController implements BeforeApplicationShutdown { ...session.config.promptConfig, signal: this.getSignal(req), user: user.id, + session: session.config.sessionId, workspace: session.config.workspaceId, reasoning, webSearch, @@ -384,6 +386,7 @@ export class CopilotController implements BeforeApplicationShutdown { ...session.config.promptConfig, signal: this.getSignal(req), user: user.id, + session: session.config.sessionId, workspace: session.config.workspaceId, reasoning, webSearch, @@ -463,6 +466,7 @@ export class CopilotController implements BeforeApplicationShutdown { ...session.config.promptConfig, signal: this.getSignal(req), user: user.id, + session: session.config.sessionId, workspace: session.config.workspaceId, }) ).pipe( @@ -586,6 +590,7 @@ export class CopilotController implements BeforeApplicationShutdown { seed: this.parseNumber(params.seed), signal: this.getSignal(req), user: user.id, + session: session.config.sessionId, workspace: session.config.workspaceId, } ) diff --git a/packages/backend/server/src/plugins/copilot/prompt/prompts.ts b/packages/backend/server/src/plugins/copilot/prompt/prompts.ts index 6858b1afd4..1915ead380 100644 --- a/packages/backend/server/src/plugins/copilot/prompt/prompts.ts +++ b/packages/backend/server/src/plugins/copilot/prompt/prompts.ts @@ -1670,6 +1670,11 @@ Your mission is to do your utmost to help users leverage AFFiNE's capabilities f AFFiNE is developed by Toeverything Pte. Ltd., a Singapore-registered company with a diverse international team. The company has also open-sourced BlockSuite and OctoBase to support the creation of tools similar to AFFiNE. The name "AFFiNE" is inspired by the concept of affine transformation, as blocks within AFFiNE can move freely across page, edgeless, and database modes. Currently, the AFFiNE team consists of 25 members and is an engineer-driven open-source company. + +- When searching for information, prioritize searching the user's Workspace information. +- Depending on the complexity of the question and the information returned by the search tools, you can call different tools multiple times to search. + + Today is: {{affine::date}}. User's preferred language is {{affine::language}}. diff --git a/packages/backend/server/src/plugins/copilot/providers/provider.ts b/packages/backend/server/src/plugins/copilot/providers/provider.ts index 726ec9a84f..1bfd3127ce 100644 --- a/packages/backend/server/src/plugins/copilot/providers/provider.ts +++ b/packages/backend/server/src/plugins/copilot/providers/provider.ts @@ -141,7 +141,11 @@ export abstract class CopilotProvider { const context = this.moduleRef.get(CopilotContextService, { strict: false, }); - const searchDocs = buildDocSearchGetter(ac, context); + + const docContext = options.session + ? await context.getBySessionId(options.session) + : null; + const searchDocs = buildDocSearchGetter(ac, context, docContext); tools.doc_semantic_search = createDocSemanticSearchTool( searchDocs.bind(null, options) ); diff --git a/packages/backend/server/src/plugins/copilot/providers/types.ts b/packages/backend/server/src/plugins/copilot/providers/types.ts index 57f09c9c5b..cbe797dd49 100644 --- a/packages/backend/server/src/plugins/copilot/providers/types.ts +++ b/packages/backend/server/src/plugins/copilot/providers/types.ts @@ -161,6 +161,7 @@ export type StreamObject = z.infer; const CopilotProviderOptionsSchema = z.object({ signal: z.instanceof(AbortSignal).optional(), user: z.string().optional(), + session: z.string().optional(), workspace: z.string().optional(), }); diff --git a/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts b/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts index 82fd6e9c64..569c540d1e 100644 --- a/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts +++ b/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts @@ -4,14 +4,20 @@ import { z } from 'zod'; import type { AccessController } from '../../../core/permission'; import type { ChunkSimilarity } from '../../../models'; import type { CopilotContextService } from '../context'; +import type { ContextSession } from '../context/session'; import type { CopilotChatOptions } from '../providers'; import { toolError } from './error'; export const buildDocSearchGetter = ( ac: AccessController, - context: CopilotContextService + context: CopilotContextService, + docContext: ContextSession | null ) => { - const searchDocs = async (options: CopilotChatOptions, query?: string) => { + const searchDocs = async ( + options: CopilotChatOptions, + query?: string, + abortSignal?: AbortSignal + ) => { if (!options || !query?.trim() || !options.user || !options.workspace) { return undefined; } @@ -20,7 +26,11 @@ export const buildDocSearchGetter = ( .workspace(options.workspace) .can('Workspace.Read'); if (!canAccess) return undefined; - const chunks = await context.matchWorkspaceAll(options.workspace, query); + const [chunks, contextChunks] = await Promise.all([ + context.matchWorkspaceAll(options.workspace, query, 10, abortSignal), + docContext?.matchFiles(query, 10, abortSignal) ?? [], + ]); + const docChunks = await ac .user(options.user) .workspace(options.workspace) @@ -29,6 +39,9 @@ export const buildDocSearchGetter = ( 'Doc.Read' ); const fileChunks = chunks.filter(c => 'fileId' in c); + if (contextChunks.length) { + fileChunks.push(...contextChunks); + } if (!docChunks.length && !fileChunks.length) return undefined; return [...fileChunks, ...docChunks]; }; @@ -36,17 +49,24 @@ export const buildDocSearchGetter = ( }; export const createDocSemanticSearchTool = ( - searchDocs: (query: string) => Promise + searchDocs: ( + query: string, + abortSignal?: AbortSignal + ) => Promise ) => { return tool({ description: 'Semantic search for relevant documents in the current workspace', parameters: z.object({ - query: z.string().describe('The query to search for.'), + query: z + .string() + .describe( + 'The query statement to search for, e.g. "What is the capital of France?"' + ), }), - execute: async ({ query }) => { + execute: async ({ query }, options) => { try { - return await searchDocs(query); + return await searchDocs(query, options.abortSignal); } catch (e: any) { return toolError('Doc Semantic Search Failed', e.message); } diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts index 4ac380394f..aa0ab4bcf5 100644 --- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts +++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts @@ -585,7 +585,7 @@ export class AIChatInput extends SignalWatcher( await this._preUpdateMessages(userInput, attachments); const sessionId = await this.createSessionId(); - let contexts = await this._getMatchedContexts(userInput); + let contexts = await this._getMatchedContexts(); if (abortController.signal.aborted) { return; } @@ -685,46 +685,11 @@ export class AIChatInput extends SignalWatcher( } }; - private async _getMatchedContexts(userInput: string) { - const contextId = await this.getContextId(); - const workspaceId = this.host.store.workspace.id; - + private async _getMatchedContexts() { const docContexts = new Map< string, { docId: string; docContent: string } >(); - const fileContexts = new Map< - string, - BlockSuitePresets.AIFileContextOption - >(); - - const { files: matchedFiles = [], docs: matchedDocs = [] } = - (await AIProvider.context?.matchContext( - userInput, - contextId, - workspaceId - )) ?? {}; - - matchedDocs.forEach(doc => { - docContexts.set(doc.docId, { - docId: doc.docId, - docContent: doc.content, - }); - }); - - matchedFiles.forEach(file => { - const context = fileContexts.get(file.fileId); - if (context) { - context.fileContent += `\n${file.content}`; - } else { - fileContexts.set(file.fileId, { - blobId: file.blobId, - fileName: file.name, - fileType: file.mimeType, - fileContent: file.content, - }); - } - }); this.chips.forEach(chip => { if (isDocChip(chip) && !!chip.markdown?.value) { @@ -759,10 +724,7 @@ export class AIChatInput extends SignalWatcher( }; }); - return { - docs, - files: Array.from(fileContexts.values()), - }; + return { docs, files: [] }; } } diff --git a/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx b/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx index 95e4bf5baa..51758b8b5a 100644 --- a/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx +++ b/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx @@ -95,6 +95,7 @@ export function setupAIProvider( client, sessionId, content: input, + timeout: 5 * 60 * 1000, // 5 minutes params: { docs: contexts?.docs, files: contexts?.files, diff --git a/tests/affine-cloud-copilot/e2e/utils/chat-panel-utils.ts b/tests/affine-cloud-copilot/e2e/utils/chat-panel-utils.ts index b54eaaeac2..940a6b7f21 100644 --- a/tests/affine-cloud-copilot/e2e/utils/chat-panel-utils.ts +++ b/tests/affine-cloud-copilot/e2e/utils/chat-panel-utils.ts @@ -165,9 +165,11 @@ export class ChatPanelUtils { const actionList = await message.getByTestId('chat-action-list'); return { message, - content: await message - .locator('chat-content-rich-text editor-host') - .innerText(), + content: ( + await message + .locator('chat-content-rich-text editor-host') + .allInnerTexts() + ).join(' '), actions: { copy: async () => actions.getByTestId('action-copy-button').click(), retry: async () => actions.getByTestId('action-retry-button').click(),