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 d705cda8c..e20e619c9 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 1be3c6876..65797a5bc 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 6858b1afd..1915ead38 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 726ec9a84..1bfd3127c 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 57f09c9c5..cbe797dd4 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 82fd6e9c6..569c540d1 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 4ac380394..aa0ab4bcf 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 95e4bf5ba..51758b8b5 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 b54eaaeac..940a6b7f2 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(),