feat(server): remove context prefetch & integrate context search (#12956)

fix AI-173

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([
+    const [fileChunks, workspaceChunks, scopedWorkspaceChunks] =
-      this.models.copilotWorkspace.matchFileEmbedding(
+      await Promise.all([
-        workspaceId,
+        this.models.copilotWorkspace.matchFileEmbedding(
-        embedding,
+          workspaceId,
-        topK * 2,
+          embedding,
-        threshold
+          topK * 2,
-      ),
+          threshold
-      this.models.copilotContext.matchWorkspaceEmbedding(
+        ),
-        embedding,
-        workspaceId,
-        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
     );

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

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.
 
 <response_guide>
+<tool_usage_guide>
+- 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.
+</tool_usage_guide>
+
 <real_world_info>
 Today is: {{affine::date}}.
 User's preferred language is {{affine::language}}.

packages/backend/server/src/plugins/copilot/providers/provider.ts

@@ -141,7 +141,11 @@ export abstract class CopilotProvider<C = any> {
             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)
             );

packages/backend/server/src/plugins/copilot/providers/types.ts

@@ -161,6 +161,7 @@ export type StreamObject = z.infer<typeof StreamObjectSchema>;
 const CopilotProviderOptionsSchema = z.object({
   signal: z.instanceof(AbortSignal).optional(),
   user: z.string().optional(),
+  session: z.string().optional(),
   workspace: z.string().optional(),
 });
 

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<ChunkSimilarity[] | undefined>
+  searchDocs: (
+    query: string,
+    abortSignal?: AbortSignal
+  ) => Promise<ChunkSimilarity[] | undefined>
 ) => {
   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);
       }

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) {
+  private async _getMatchedContexts() {
-    const contextId = await this.getContextId();
-    const workspaceId = this.host.store.workspace.id;
-
     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 {
+    return { docs, files: [] };
-      docs,
-      files: Array.from(fileContexts.values()),
-    };
   }
 }
 

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,

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
+      content: (
-        .locator('chat-content-rich-text editor-host')
+        await message
-        .innerText(),
+          .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(),