feat: add doc copilot context api (#10103)

### What Changed?
- Add graphql APIs.
- Provide context and session service in `AIProvider`.
- Rename the state from `embedding` to `processing`.
- Reafctor front-end session create, update and save logic.

Persist the document selected by the user:
[录屏2025-02-08 11.04.40.mov <span class="graphite__hidden">(uploaded via Graphite)</span> <img class="graphite__hidden" src="https://app.graphite.dev/api/v1/graphite/video/thumbnail/sJGviKxfE3Ap685cl5bj/195a85f2-43c4-4e49-88d9-6b5fc4f235ca.mov" />](https://app.graphite.dev/media/video/sJGviKxfE3Ap685cl5bj/195a85f2-43c4-4e49-88d9-6b5fc4f235ca.mov)

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

@@ -1,6 +1,8 @@
 import { showAILoginRequiredAtom } from '@affine/core/components/affine/auth/ai-login-required';
 import {
+  addContextDocMutation,
   cleanupCopilotSessionMutation,
+  createCopilotContextMutation,
   createCopilotMessageMutation,
   createCopilotSessionMutation,
   forkCopilotSessionMutation,
@@ -9,8 +11,11 @@ import {
   getCopilotSessionsQuery,
   GraphQLError,
   type GraphQLQuery,
+  listContextDocsAndFilesQuery,
+  listContextQuery,
   type QueryOptions,
   type QueryResponse,
+  removeContextDocMutation,
   type RequestOptions,
   updateCopilotSessionMutation,
   UserFriendlyError,
@@ -209,6 +214,74 @@ export class CopilotClient {
     }
   }
 
+  async createContext(workspaceId: string, sessionId: string) {
+    const res = await this.gql({
+      query: createCopilotContextMutation,
+      variables: {
+        workspaceId,
+        sessionId,
+      },
+    });
+    return res.createCopilotContext;
+  }
+
+  async getContextId(workspaceId: string, sessionId: string) {
+    const res = await this.gql({
+      query: listContextQuery,
+      variables: {
+        workspaceId,
+        sessionId,
+      },
+    });
+    return res.currentUser?.copilot?.contexts?.[0]?.id;
+  }
+
+  async addContextDoc(options: OptionsField<typeof addContextDocMutation>) {
+    const res = await this.gql({
+      query: addContextDocMutation,
+      variables: {
+        options,
+      },
+    });
+    return res.addContextDoc;
+  }
+
+  async removeContextDoc(
+    options: OptionsField<typeof removeContextDocMutation>
+  ) {
+    const res = await this.gql({
+      query: removeContextDocMutation,
+      variables: {
+        options,
+      },
+    });
+    return res.removeContextDoc;
+  }
+
+  async addContextFile() {
+    return;
+  }
+
+  async removeContextFile() {
+    return;
+  }
+
+  async getContextDocsAndFiles(
+    workspaceId: string,
+    sessionId: string,
+    contextId: string
+  ) {
+    const res = await this.gql({
+      query: listContextDocsAndFilesQuery,
+      variables: {
+        workspaceId,
+        sessionId,
+        contextId,
+      },
+    });
+    return res.currentUser?.copilot?.contexts?.[0];
+  }
+
   async chatText({
     sessionId,
     messageId,

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

@@ -31,22 +31,29 @@ export type ToImageOptions = TextToTextOptions & {
   seed?: string;
 };
 
-export function createChatSession({
+export async function createChatSession({
   client,
   workspaceId,
   docId,
-  promptName,
+  promptName = 'Chat With AFFiNE AI',
 }: {
   client: CopilotClient;
   workspaceId: string;
   docId: string;
-  promptName: string;
+  promptName?: string;
 }) {
-  return client.createSession({
+  const sessionId = await client.createSession({
     workspaceId,
     docId,
     promptName,
   });
+  // always update the prompt name
+  await updateChatSession({
+    sessionId,
+    client,
+    promptName,
+  });
+  return sessionId;
 }
 
 export function updateChatSession({
@@ -119,7 +126,8 @@ async function createSessionMessage({
   }
   const hasAttachments = attachments && attachments.length > 0;
   const sessionId = await (providedSessionId ??
-    client.createSession({
+    createChatSession({
+      client,
       workspaceId,
       docId,
       promptName: promptName as string,

packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx

@@ -1,12 +1,10 @@
 import { AIProvider } from '@affine/core/blocksuite/presets/ai';
 import { toggleGeneralAIOnboarding } from '@affine/core/components/affine/ai-onboarding/apis';
-import type { AINetworkSearchService } from '@affine/core/modules/ai-button/services/network-search';
 import type { GlobalDialogService } from '@affine/core/modules/dialogs';
 import {
   type getCopilotHistoriesQuery,
   type RequestOptions,
 } from '@affine/graphql';
-import { UnauthorizedError } from '@blocksuite/affine/blocks';
 import { assertExists } from '@blocksuite/affine/global/utils';
 import { z } from 'zod';
 
@@ -39,80 +37,19 @@ const processTypeToPromptName = new Map(
   })
 );
 
-// a single workspace should have only a single chat session
-// user-id:workspace-id:doc-id -> chat session id
-const chatSessions = new Map<
-  string,
-  { getSessionId: Promise<string>; promptName: string }
->();
-
 export function setupAIProvider(
   client: CopilotClient,
-  globalDialogService: GlobalDialogService,
-  networkSearchService: AINetworkSearchService
+  globalDialogService: GlobalDialogService
 ) {
-  function getChatPrompt(options: BlockSuitePresets.ChatOptions) {
-    const { attachments, docs } = options;
-    if (attachments?.length || docs?.length) {
-      return 'Chat With AFFiNE AI';
-    }
-    const { enabled, visible } = networkSearchService;
-    return visible.value && enabled.value
-      ? 'Search With AFFiNE AI'
-      : 'Chat With AFFiNE AI';
-  }
-  async function getChatSessionId(options: BlockSuitePresets.ChatOptions) {
-    const userId = (await AIProvider.userInfo)?.id;
-
-    if (!userId) {
-      throw new UnauthorizedError();
-    }
-
-    const { workspaceId, docId } = options;
-    const storeKey = `${userId}:${workspaceId}:${docId}`;
-    const promptName = getChatPrompt(options);
-    if (!chatSessions.has(storeKey)) {
-      chatSessions.set(storeKey, {
-        getSessionId: createChatSession({
-          client,
-          workspaceId,
-          docId,
-          promptName,
-        }).then(sessionId => {
-          return updateChatSession({
-            sessionId,
-            client,
-            promptName,
-          });
-        }),
-        promptName,
-      });
-    }
-    try {
-      /* oxlint-disable @typescript-eslint/no-non-null-assertion */
-      const { getSessionId, promptName: prevName } =
-        chatSessions.get(storeKey)!;
-      const sessionId = await getSessionId;
-      //update prompt name
-      if (prevName !== promptName) {
-        await updateChatSession({
-          sessionId,
-          client,
-          promptName,
-        });
-        chatSessions.set(storeKey, { getSessionId, promptName });
-      }
-      return sessionId;
-    } catch (err) {
-      // do not cache the error
-      chatSessions.delete(storeKey);
-      throw err;
-    }
-  }
-
   //#region actions
   AIProvider.provide('chat', options => {
-    const sessionId = options.sessionId ?? getChatSessionId(options);
+    const sessionId =
+      options.sessionId ??
+      createChatSession({
+        client,
+        workspaceId: options.workspaceId,
+        docId: options.docId,
+      });
     const { input, docs, ...rest } = options;
     const params = docs?.length
       ? {
@@ -473,6 +410,56 @@ Could you make a new website based on these notes and send back just the html fi
   });
   //#endregion
 
+  AIProvider.provide('session', {
+    createSession: async (
+      workspaceId: string,
+      docId: string,
+      promptName?: string
+    ) => {
+      return createChatSession({
+        client,
+        workspaceId,
+        docId,
+        promptName,
+      });
+    },
+    updateSession: async (sessionId: string, promptName: string) => {
+      return updateChatSession({
+        client,
+        sessionId,
+        promptName,
+      });
+    },
+  });
+
+  AIProvider.provide('context', {
+    createContext: async (workspaceId: string, sessionId: string) => {
+      return client.createContext(workspaceId, sessionId);
+    },
+    getContextId: async (workspaceId: string, sessionId: string) => {
+      return client.getContextId(workspaceId, sessionId);
+    },
+    addContextDoc: async (options: { contextId: string; docId: string }) => {
+      return client.addContextDoc(options);
+    },
+    removeContextDoc: async (options: { contextId: string; docId: string }) => {
+      return client.removeContextDoc(options);
+    },
+    addContextFile: async () => {
+      return client.addContextFile();
+    },
+    removeContextFile: async () => {
+      return client.removeContextFile();
+    },
+    getContextDocsAndFiles: async (
+      workspaceId: string,
+      sessionId: string,
+      contextId: string
+    ) => {
+      return client.getContextDocsAndFiles(workspaceId, sessionId, contextId);
+    },
+  });
+
   AIProvider.provide('histories', {
     actions: async (
       workspaceId: string,

packages/frontend/core/src/components/providers/workspace-side-effects.tsx

@@ -8,7 +8,6 @@ import { SyncAwareness } from '@affine/core/components/affine/awareness';
 import { useRegisterFindInPageCommands } from '@affine/core/components/hooks/affine/use-register-find-in-page-commands';
 import { useRegisterWorkspaceCommands } from '@affine/core/components/hooks/use-register-workspace-commands';
 import { OverCapacityNotification } from '@affine/core/components/over-capacity';
-import { AINetworkSearchService } from '@affine/core/modules/ai-button/services/network-search';
 import {
   EventSourceService,
   FetchService,
@@ -144,7 +143,6 @@ export const WorkspaceSideEffects = () => {
   const graphqlService = useService(GraphQLService);
   const eventSourceService = useService(EventSourceService);
   const fetchService = useService(FetchService);
-  const networkSearchService = useService(AINetworkSearchService);
 
   useEffect(() => {
     const dispose = setupAIProvider(
@@ -153,8 +151,7 @@ export const WorkspaceSideEffects = () => {
         fetchService.fetch,
         eventSourceService.eventSource
       ),
-      globalDialogService,
-      networkSearchService
+      globalDialogService
     );
     return () => {
       dispose();
@@ -164,7 +161,6 @@ export const WorkspaceSideEffects = () => {
     fetchService,
     workspaceDialogService,
     graphqlService,
-    networkSearchService,
     globalDialogService,
   ]);